Source code for autojob.settings

"""Global settings and configuration for ``autojob``."""

import logging
from logging import handlers
import os
from pathlib import Path
from typing import Any
from typing import ClassVar

from pydantic import Field
from pydantic import field_validator
from pydantic import model_validator
from pydantic_settings import BaseSettings
from pydantic_settings import PydanticBaseSettingsSource
from pydantic_settings import SettingsConfigDict
from pydantic_settings.sources import TomlConfigSettingsSource

logger = logging.getLogger(__name__)

AUTOJOB_HOME = (
    Path()
    .home()
    .resolve()
    .joinpath(
        ".config",
        "autojob",
    )
)
_DEFAULT_CONFIG_FILE = AUTOJOB_HOME.joinpath("config.toml")


[docs] class AutojobSettings(BaseSettings): """Settings for ``autojob``.""" # Input Files INPUT_ATOMS_FILE: str = Field( default="in.traj", description="the default filename to use for the input Atoms file", ) INPUTS_FILE: str = Field( default="inputs.json", description="the default filename to use for the input parameters file", ) # TODO: Replace with CalculationInputs.calculation_script attribute TASK_SCRIPT: str = Field( default="run.py", description="the default filename to use for the python script", ) # TODO: Replace with TaskInputs.task_script attribute SLURM_SCRIPT: str = Field( default="vasp.sh", description="the default filename to use for the SLURM script", ) DEFAULT_TASK_SCRIPT_FILE: str = Field( default="run.sh", description="the default filename to use for the task script", ) DEFAULT_CALCULATION_SCRIPT_FILE: str = Field( default="run.py", description="the default filename to use for the calculation script", ) # Output Files OUTPUT_ATOMS_FILE: str = Field( default="final.traj", description="the default filename to use for the output Atoms file", ) SCHEDULER_STATS_FILE: str = Field( default="job_stats.json", description="the default filename to use for the job stats file", ) ARCHIVE_FILE: str = Field( default="task.json", description="the default filename to use to archive harvested tasks", ) # Metadata Files TASK_METADATA_FILE: str = Field( default="job.json", description="the default filename to use to store task metadata", ) TASK_GROUP_METADATA_FILE: str = Field( default="calculation.json", description="the default filename to use to store task group metadata", ) STUDY_METADATA_FILE: str = Field( default="study.json", description="the default filename to use to store study metadata", ) STUDY_GROUP_METADATA_FILE: str = Field( default="study_group.json", description="the default filename to use to store study group " "metadata", ) # Workflow Files RECORD_FILE: str = Field( default="record.txt", description="the default filename to use to store the study record", ) WORKFLOW_FILE: str = Field( default="workflow.json", description="the default filename to use to store study workflow data", ) PARAMETRIZATION_FILE: str = Field( default="parametrizations.json", description="the default filename to use to store study " "parametrization data", ) # Template Files TEMPLATE_DIR: Path | None = Field( default=None, description="If not None, specifies the directory to use to load " "templates", ) TASK_SCRIPT_TEMPLATE: str = Field( default="run.sh.j2", description="the name of the default scheduler script template to use", ) CALCULATION_SCRIPT_TEMPLATE: str = Field( default="run.py.j2", description="the name of the default task script template to use", ) SCHEDULER: str = Field( default="slurm", description="the name of the scheduler to use", ) # Logging & General Behaviour LOG_FILE: Path | None = Field( default=None, description="The filename for the log file. Note that this variable " "is mainly for storing state for application-like use. If you are " "using autojob as a library, you may be better served " "configuring handlers.", ) LOG_LEVEL: int = Field( default=logging.DEBUG, description="The default log level.", ) STRICT_MODE: bool = Field( default=True, description="Sets the default behaviour of data retrieval functions " "and methods. If True, such functions will raise errors when they " "fail. Otherwise, failure will pass with a log message only. This " "may be useful if you are harvesting the results of incomplete tasks.", ) LEGACY_MODE: bool = Field( default=True, description="Sets the format study group, study, task group, and task " "IDs. If True, said IDs will be 10-character alphanumerics prefixed " "by single letters (g, s, c, j, respectively). IDs will be UUID4s " "otherwise.", ) DEFAULT_TASK: str = Field( default="task", description="The name of the default task type to use." ) # VASP Settings VASP_KEEP_DOS: bool = Field( default=False, description="Whether or not to store the complete DOS from VASP " "calculations", ) model_config: ClassVar[SettingsConfigDict] = SettingsConfigDict( env_prefix="AUTOJOB_", case_sensitive=True, env_ignore_empty=True )
[docs] @field_validator("LOG_LEVEL", mode="plain") @classmethod def validate_log_level(cls, v: Any) -> int: """Validate the global log level.""" if isinstance(v, int): return v try: level: int = getattr(logging, v) return level except AttributeError as err: msg = f"{v} not a valid logging level" raise ValueError(msg) from err except TypeError as err: msg = f"Unable to convert {v} into a logging level" raise ValueError(msg) from err
[docs] @model_validator(mode="after") def configure_logging(self) -> "AutojobSettings": """Configure logging based on user settings.""" if self.LOG_FILE: fh = handlers.RotatingFileHandler( self.LOG_FILE, encoding="utf-8", maxBytes=int(1e6), backupCount=3, ) log_format = ( "%(asctime)s - %(name)s::%(funcName)s::%(lineno)s - " "%(levelname)s - %(message)s " ) formatter = logging.Formatter(log_format) fh.setFormatter(formatter) fh.setLevel(level=self.LOG_LEVEL) logger.addHandler(fh) return self
@classmethod def settings_customise_sources( cls, settings_cls: type[BaseSettings], init_settings: PydanticBaseSettingsSource, env_settings: PydanticBaseSettingsSource, dotenv_settings: PydanticBaseSettingsSource, file_secret_settings: PydanticBaseSettingsSource, ) -> tuple[PydanticBaseSettingsSource, ...]: """Add a TOML configuration file to the settings sources. Note that the name of the TOML file can be set via environment variables. Otherwise, the default filename is used. """ toml_file = os.environ.get("AUTOJOB_CONFIG_FILE", _DEFAULT_CONFIG_FILE) logger.info("Configuration file will be read from %s", toml_file) return ( init_settings, env_settings, dotenv_settings, TomlConfigSettingsSource(settings_cls, toml_file), file_secret_settings, )