schedulers/PanDAiDDS/config.py

0001 """Pydantic configuration model for the PanDAiDDS scheduler.
0002
0003 This module defines a Pydantic model that mirrors the expected PanDAiDDS
0004 attributes used when submitting jobs to the PanDA system. The fields are
0005 inferred from the example `panda_attrs` dictionary used by the EPIC/PanDA
0006 integration.
0007
0008 The model is registered with the scheduler configuration registry so the
0009 framework can discover and validate PanDA runner configurations.
0010 """
0011
0012 import os
0013 import getpass
0014 from typing import List, Optional, Any
0015 from pydantic import BaseModel, Field, field_validator, model_validator
0016
0017 from aid2e.utilities.configurations.scheduler_registry import register as register_runner_config
0018
0019
0020 class PanDAiDDSRunnerConfig(BaseModel):
0021     """Configuration for the PanDAiDDS scheduler.
0022
0023     Auto-generation behavior:
0024     - name: Auto-generates as 'user.<username>.aid2e_job' if not provided
0025             Override username via PANDA_USERNAME environment variable
0026     - source_dir: Auto-sets to project root directory if not provided
0027                   Override via PANDA_SOURCE_DIR environment variable
0028     - init_env: Auto-sets to 'source setup_aid2e.sh' if not provided
0029                 If provided as string, appends ' && source setup_aid2e.sh'
0030
0031     Fields:
0032     - name: str (auto-generated)
0033     - init_env: any (auto-set to source setup script, or prepended if string)
0034     - cloud: str
0035     - queue: str
0036     - working_group: str
0037     - task_type: str
0038     - source_dir: str (auto-set to project root)
0039     - source_dir_parent_level: int
0040     - exclude_source_files: List[str] (includes .venv, venv, .git)
0041     - max_walltime: int (seconds)
0042     - core_count: int
0043     - total_memory: int (MB)
0044     - enable_separate_log: bool
0045     - job_dir: Optional[str]
0046     """
0047
0048     name: Optional[str] = Field(
0049         default=None,
0050         description=(
0051             "PanDA job name, must start with 'user.<username>'. "
0052             "If not provided, will be auto-generated from username. "
0053             "Set PANDA_USERNAME environment variable to override system username."
0054         ),
0055     )
0056     init_env: Optional[Any] = Field(
0057         default=None,
0058         description=(
0059             "Initialization environment (callable, dict, or other) to prepare remote jobs. "
0060             "If not provided, defaults to 'source setup_aid2e.sh'. "
0061             "If a string is provided, 'source setup_aid2e.sh' will be appended to it."
0062         ),
0063     )
0064     post_script: Optional[Any] = Field(
0065         default="rm -fr .local .venv src examples tests docs",
0066         description=(
0067             "Post-execution script (callable, dict, or other) to clean up after remote jobs. "
0068             "If not provided, defaults to 'rm -fr .local .venv src examples tests docs'. "
0069             "If a string is provided, it will be executed after the job completes."
0070         ),
0071     )
0072
0073     @model_validator(mode='after')
0074     def validate_and_set_defaults(self) -> 'PanDAiDDSRunnerConfig':
0075         """Validate or generate PanDA job name and set source_dir defaults.
0076
0077         The name must start with 'user.<username>'. If not provided, it will be
0078         auto-generated using the system username (or PANDA_USERNAME env var).
0079
0080         The source_dir defaults to the project root directory. Can be overridden
0081         via PANDA_SOURCE_DIR environment variable.
0082
0083         The init_env defaults to sourcing setup_aid2e.sh. If a string value is
0084         already provided, the setup script will be appended to it.
0085
0086         Returns:
0087             Self with validated/generated name, source_dir, and init_env.
0088
0089         Raises:
0090             ValueError: If the provided name doesn't start with 'user.'.
0091         """
0092         # Validate and generate name
0093         if self.name is not None and self.name != "":
0094             if not self.name.startswith("user."):
0095                 raise ValueError(
0096                     f"PanDA job name must start with 'user.<username>', got: {self.name}"
0097                 )
0098         else:
0099             # Auto-generate name from username
0100             # Check environment variable first, then fall back to system username
0101             username = os.environ.get("PANDA_USERNAME") or getpass.getuser()
0102             self.name = f"user.{username}.aid2e_job"
0103
0104         # Set source_dir to project root if not provided
0105         if self.source_dir is None:
0106             # Check environment variable first
0107             env_source = os.environ.get("PANDA_SOURCE_DIR")
0108             if env_source:
0109                 self.source_dir = env_source
0110             else:
0111                 # Default to project root directory
0112                 # Navigate from this config file: .../src/aid2e/schedulers/PanDAiDDS/config.py
0113                 # Go up to project root: ../../../.. from this file
0114                 config_file_dir = os.path.dirname(os.path.abspath(__file__))
0115                 project_root = os.path.abspath(os.path.join(config_file_dir, "..", "..", "..", ".."))
0116                 self.source_dir = project_root
0117
0118         # Set init_env to source setup_aid2e.sh, or append to existing
0119         if self.init_env is None:
0120             self.init_env = "source setup_aid2e.sh; bash install_aid2e_dependencies.sh; "
0121         else:
0122             # If init_env is already set, append the setup script after it
0123             if isinstance(self.init_env, str):
0124                 self.init_env = f"source setup_aid2e.sh && bash install_aid2e_dependencies.sh && {self.init_env}"
0125             # Note: If init_env is a callable or other type, leave it as-is
0126
0127         return self
0128
0129     cloud: Optional[str] = Field(
0130         default=None,
0131         description="Cloud/region for the PanDA submission (e.g. 'US')",
0132     )
0133     queue: Optional[str] = Field(
0134         default=None,
0135         description="PanDA queue name to submit jobs to",
0136     )
0137     working_group: Optional[str] = Field(
0138         default="AID2E",
0139         description="Working group for the PanDA submission (e.g. 'wg_epic', 'AID2E')",
0140     )
0141     task_type: Optional[str] = Field(
0142         default="AID2E",
0143         description="Task type for PanDA job classification (e.g. 'test', 'prod', 'analysis', 'AID2E')",
0144     )
0145     source_dir: str = Field(
0146         default=None,
0147         description=(
0148             "Directory whose contents should be uploaded to PanDA for remote jobs. "
0149             "If not provided, defaults to project root directory. "
0150             "Set PANDA_SOURCE_DIR environment variable to override."
0151         ),
0152     )
0153     source_dir_parent_level: int = Field(
0154         default=1,
0155         description="How many parent levels above the source_dir should be included when uploading",
0156     )
0157     exclude_source_files: List[str] = Field(
0158         default_factory=lambda: [
0159             r"(^|/)\.[^/]+",
0160             r"^doc",
0161             r"^DTLZ2",
0162             r".*json$",
0163             r".*log$",
0164             "work",
0165             "log",
0166             "OUTDIR",
0167             "calibrations",
0168             "fieldmaps",
0169             "gdml",
0170             "EICrecon-drich-mobo",
0171             "eic-software",
0172             "epic-geom-drich-mobo",
0173             "irt",
0174             "share",
0175             "back*",
0176             "__pycache__",
0177             ".venv",
0178             "venv",
0179             ".git",
0180             ".local",
0181         ],
0182         description="Filename patterns to exclude when packaging source_dir",
0183     )
0184     max_walltime: Optional[int] = Field(
0185         default=None,
0186         description="Maximum walltime in seconds for a PanDA job",
0187     )
0188     core_count: int = Field(
0189         default=1,
0190         description="Number of CPU cores requested per job",
0191     )
0192     total_memory: int = Field(
0193         default=4000,
0194         description="Total memory in MB requested per job",
0195     )
0196     enable_separate_log: bool = Field(
0197         default=True,
0198         description="Whether to enable separate log files for each remote job",
0199     )
0200     job_dir: Optional[str] = Field(
0201         default=None,
0202         description="Remote job working directory (if applicable)",
0203     )
0204
0205
0206 # Register with the runner-config registry for lookup by runner_type
0207 register_runner_config("PanDAiDDSRunner", PanDAiDDSRunnerConfig)