File indexing completed on 2026-08-12 08:24:56
0001 """Base scheduler abstract class for job execution.
0002
0003 All schedulers (JobLib, SLURM, PanDA, etc.) inherit from BaseScheduler
0004 and implement the core interface: run_stage(), check_status(), cancel_job().
0005
0006 Project: AID2E v0.0.0 - AI assisted Detector Design for EIC
0007 Homepage: https://aid2e.github.io/AID2E-framework
0008 Repository: https://github.com/aid2e/AID2E-framework.git
0009 """
0010
0011 from abc import ABC, abstractmethod
0012 from typing import Dict, Any, List, Optional
0013 from pydantic import BaseModel
0014
0015
0016 class JobStatus(BaseModel):
0017 """Represent status information for a single job.
0018
0019 Args:
0020 job_id: Unique job identifier.
0021 status: Current status ("queued", "running", "completed", "failed", "cancelled").
0022 return_code: Exit code when completed or failed.
0023 stdout: Standard output from the job.
0024 stderr: Standard error from the job.
0025 outputs: Optional output data (e.g., objectives, results from Python callables).
0026 metrics: Optional metrics (e.g., runtime, memory usage).
0027 """
0028
0029 job_id: str
0030 status: str
0031 return_code: Optional[int] = None
0032 stdout: Optional[str] = None
0033 stderr: Optional[str] = None
0034 outputs: Optional[Dict[str, Any]] = None
0035 metrics: Optional[Dict[str, Any]] = None
0036
0037
0038 class StageExecutionResult(BaseModel):
0039 """Capture the result of executing all jobs in a stage.
0040
0041 Args:
0042 stage_name: Name of the executed stage.
0043 job_statuses: Status for each job in the stage.
0044 artifacts: Output artifacts collected from the stage (path -> content).
0045 success: Whether all jobs completed successfully.
0046 error_message: Optional error message if the stage failed.
0047 """
0048
0049 stage_name: str
0050 job_statuses: List[JobStatus]
0051 artifacts: Dict[str, Any]
0052 success: bool
0053 error_message: Optional[str] = None
0054
0055
0056 class StageStatus(BaseModel):
0057 """Lightweight status summary for a submitted stage.
0058
0059 This model is intended for polling asynchronous stage submissions. It
0060 provides progress information and optional per-job statuses when available.
0061 """
0062
0063 stage_id: str
0064 status: str
0065 completed_jobs: int = 0
0066 total_jobs: Optional[int] = None
0067 progress: Optional[float] = None
0068 job_statuses: Optional[List[JobStatus]] = None
0069
0070
0071 class BaseScheduler(ABC):
0072 """Define the common scheduler interface.
0073
0074 Schedulers execute workflow stages on different backends (local, SLURM, PanDA, etc.).
0075 They handle job submission, monitoring, retries, and artifact collection.
0076 """
0077
0078 def __init__(self, config: Optional[BaseModel] = None) -> None:
0079 """Initialize the scheduler with executor-specific configuration.
0080
0081 Args:
0082 config: Executor-specific config (e.g., JobLibRunnerConfig, SlurmRunnerConfig).
0083 """
0084
0085 self.config = config or {}
0086
0087 @abstractmethod
0088 def run_stage(
0089 self,
0090 stage_name: str,
0091 job_definitions: List[Dict[str, Any]],
0092 parallelism_policy: Optional[Dict[str, Any]] = None,
0093 working_dir: Optional[str] = None,
0094 ) -> StageExecutionResult:
0095 """Execute all jobs in a stage respecting parallelism constraints.
0096
0097 Args:
0098 stage_name: Name of the stage being executed.
0099 job_definitions: Job dictionaries with command, payload, outputs, etc.
0100 parallelism_policy: Parallelism settings (max_concurrent, retry_max, timeout_sec).
0101 working_dir: Working directory for job execution.
0102
0103 Returns:
0104 StageExecutionResult describing job outcomes and collected artifacts.
0105 """
0106
0107
0108 @abstractmethod
0109 def check_status(self, job_id: str) -> JobStatus:
0110 """Check the status of a previously submitted job.
0111
0112 Args:
0113 job_id: Unique job identifier returned from ``run_stage``.
0114
0115 Returns:
0116 JobStatus with current state and metrics (if available).
0117 """
0118
0119 @abstractmethod
0120 def cancel_job(self, job_id: str) -> bool:
0121 """Cancel a job if it is still running.
0122
0123 Args:
0124 job_id: Unique job identifier.
0125
0126 Returns:
0127 True if the job was cancelled, False otherwise.
0128 """
0129
0130 @abstractmethod
0131 def submit_stage(
0132 self,
0133 stage_name: str,
0134 job_definitions: List[Dict[str, Any]],
0135 parallelism_policy: Optional[Dict[str, Any]] = None,
0136 working_dir: Optional[str] = None,
0137 ) -> str:
0138 """Submit a stage for asynchronous execution.
0139
0140 Unlike `run_stage`, which may block until completion and return a
0141 `StageExecutionResult`, `submit_stage` should schedule the stage and
0142 return immediately with a `stage_id` that can be used to poll status
0143 and retrieve results later.
0144
0145 Returns:
0146 A unique `stage_id` string that identifies the submitted stage.
0147 """
0148
0149 @abstractmethod
0150 def check_stage_status(self, stage_id: str) -> StageStatus:
0151 """Check the current status of an asynchronously submitted stage.
0152
0153 Args:
0154 stage_id: The identifier returned by `submit_stage`.
0155
0156 Returns:
0157 A `StageStatus` object summarizing progress and (optionally)
0158 per-job statuses.
0159 """
0160
0161 @abstractmethod
0162 def get_stage_results(self, stage_id: str) -> StageExecutionResult:
0163 """Retrieve final execution results for a completed stage.
0164
0165 This should block or raise an informative error if the stage is not
0166 yet finished, depending on the scheduler implementation's semantics.
0167
0168 Args:
0169 stage_id: The identifier returned by `submit_stage`.
0170
0171 Returns:
0172 A `StageExecutionResult` containing artifact collection and
0173 per-job statuses for the stage.
0174 """
0175
0176 def shutdown(self) -> None:
0177 """Clean up scheduler resources (optional for implementations)."""
0178
0179 return None
