Schedulers API

Schedulers for local and distributed execution backends.

Current built-in scheduler: - JobLibScheduler: Local parallel execution using joblib.

Registries mirror the optimizer layout for consistency, providing register, get, and list_registered helpers in aid2e.schedulers._registry.

BaseScheduler

Bases: ABC

Define the common scheduler interface.

Schedulers execute workflow stages on different backends (local, SLURM, PanDA, etc.). They handle job submission, monitoring, retries, and artifact collection.

Source code in src/aid2e/schedulers/base.py

class BaseScheduler(ABC):
    """Define the common scheduler interface.

    Schedulers execute workflow stages on different backends (local, SLURM, PanDA, etc.).
    They handle job submission, monitoring, retries, and artifact collection.
    """

    def __init__(self, config: Optional[BaseModel] = None) -> None:
        """Initialize the scheduler with executor-specific configuration.

        Args:
            config: Executor-specific config (e.g., JobLibRunnerConfig, SlurmRunnerConfig).
        """

        self.config = config or {}

    @abstractmethod
    def run_stage(
        self,
        stage_name: str,
        job_definitions: List[Dict[str, Any]],
        parallelism_policy: Optional[Dict[str, Any]] = None,
        working_dir: Optional[str] = None,
    ) -> StageExecutionResult:
        """Execute all jobs in a stage respecting parallelism constraints.

        Args:
            stage_name: Name of the stage being executed.
            job_definitions: Job dictionaries with command, payload, outputs, etc.
            parallelism_policy: Parallelism settings (max_concurrent, retry_max, timeout_sec).
            working_dir: Working directory for job execution.

        Returns:
            StageExecutionResult describing job outcomes and collected artifacts.
        """


    @abstractmethod
    def check_status(self, job_id: str) -> JobStatus:
        """Check the status of a previously submitted job.

        Args:
            job_id: Unique job identifier returned from ``run_stage``.

        Returns:
            JobStatus with current state and metrics (if available).
        """

    @abstractmethod
    def cancel_job(self, job_id: str) -> bool:
        """Cancel a job if it is still running.

        Args:
            job_id: Unique job identifier.

        Returns:
            True if the job was cancelled, False otherwise.
        """

    @abstractmethod
    def submit_stage(
        self,
        stage_name: str,
        job_definitions: List[Dict[str, Any]],
        parallelism_policy: Optional[Dict[str, Any]] = None,
        working_dir: Optional[str] = None,
    ) -> str:
        """Submit a stage for asynchronous execution.

        Unlike `run_stage`, which may block until completion and return a
        `StageExecutionResult`, `submit_stage` should schedule the stage and
        return immediately with a `stage_id` that can be used to poll status
        and retrieve results later.

        Returns:
            A unique `stage_id` string that identifies the submitted stage.
        """

    @abstractmethod
    def check_stage_status(self, stage_id: str) -> StageStatus:
        """Check the current status of an asynchronously submitted stage.

        Args:
            stage_id: The identifier returned by `submit_stage`.

        Returns:
            A `StageStatus` object summarizing progress and (optionally)
            per-job statuses.
        """

    @abstractmethod
    def get_stage_results(self, stage_id: str) -> StageExecutionResult:
        """Retrieve final execution results for a completed stage.

        This should block or raise an informative error if the stage is not
        yet finished, depending on the scheduler implementation's semantics.

        Args:
            stage_id: The identifier returned by `submit_stage`.

        Returns:
            A `StageExecutionResult` containing artifact collection and
            per-job statuses for the stage.
        """

    def shutdown(self) -> None:
        """Clean up scheduler resources (optional for implementations)."""

        return None

__init__(config=None)

Initialize the scheduler with executor-specific configuration.

Parameters:

Name	Type	Description	Default
config	Optional[BaseModel]	Executor-specific config (e.g., JobLibRunnerConfig, SlurmRunnerConfig).	None

Source code in src/aid2e/schedulers/base.py

def __init__(self, config: Optional[BaseModel] = None) -> None:
    """Initialize the scheduler with executor-specific configuration.

    Args:
        config: Executor-specific config (e.g., JobLibRunnerConfig, SlurmRunnerConfig).
    """

    self.config = config or {}

cancel_job(job_id) abstractmethod

Cancel a job if it is still running.

Parameters:

Name	Type	Description	Default
job_id	str	Unique job identifier.	required

Returns:

Type	Description
bool	True if the job was cancelled, False otherwise.

Source code in src/aid2e/schedulers/base.py

@abstractmethod
def cancel_job(self, job_id: str) -> bool:
    """Cancel a job if it is still running.

    Args:
        job_id: Unique job identifier.

    Returns:
        True if the job was cancelled, False otherwise.
    """

check_stage_status(stage_id) abstractmethod

Check the current status of an asynchronously submitted stage.

Parameters:

Name	Type	Description	Default
stage_id	str	The identifier returned by submit_stage.	required

Returns:

Type	Description
StageStatus	A StageStatus object summarizing progress and (optionally)
StageStatus	per-job statuses.

Source code in src/aid2e/schedulers/base.py

@abstractmethod
def check_stage_status(self, stage_id: str) -> StageStatus:
    """Check the current status of an asynchronously submitted stage.

    Args:
        stage_id: The identifier returned by `submit_stage`.

    Returns:
        A `StageStatus` object summarizing progress and (optionally)
        per-job statuses.
    """

check_status(job_id) abstractmethod

Check the status of a previously submitted job.

Parameters:

Name	Type	Description	Default
job_id	str	Unique job identifier returned from run_stage.	required

Returns:

Type	Description
JobStatus	JobStatus with current state and metrics (if available).

Source code in src/aid2e/schedulers/base.py

@abstractmethod
def check_status(self, job_id: str) -> JobStatus:
    """Check the status of a previously submitted job.

    Args:
        job_id: Unique job identifier returned from ``run_stage``.

    Returns:
        JobStatus with current state and metrics (if available).
    """

get_stage_results(stage_id) abstractmethod

Retrieve final execution results for a completed stage.

This should block or raise an informative error if the stage is not yet finished, depending on the scheduler implementation's semantics.

Parameters:

Name	Type	Description	Default
stage_id	str	The identifier returned by submit_stage.	required

Returns:

Type	Description
StageExecutionResult	A StageExecutionResult containing artifact collection and
StageExecutionResult	per-job statuses for the stage.

Source code in src/aid2e/schedulers/base.py

@abstractmethod
def get_stage_results(self, stage_id: str) -> StageExecutionResult:
    """Retrieve final execution results for a completed stage.

    This should block or raise an informative error if the stage is not
    yet finished, depending on the scheduler implementation's semantics.

    Args:
        stage_id: The identifier returned by `submit_stage`.

    Returns:
        A `StageExecutionResult` containing artifact collection and
        per-job statuses for the stage.
    """

run_stage(stage_name, job_definitions, parallelism_policy=None, working_dir=None) abstractmethod

Execute all jobs in a stage respecting parallelism constraints.

Parameters:

Name	Type	Description	Default
stage_name	str	Name of the stage being executed.	required
job_definitions	List[Dict[str, Any]]	Job dictionaries with command, payload, outputs, etc.	required
parallelism_policy	Optional[Dict[str, Any]]	Parallelism settings (max_concurrent, retry_max, timeout_sec).	None
working_dir	Optional[str]	Working directory for job execution.	None

Returns:

Type	Description
StageExecutionResult	StageExecutionResult describing job outcomes and collected artifacts.

Source code in src/aid2e/schedulers/base.py

@abstractmethod
def run_stage(
    self,
    stage_name: str,
    job_definitions: List[Dict[str, Any]],
    parallelism_policy: Optional[Dict[str, Any]] = None,
    working_dir: Optional[str] = None,
) -> StageExecutionResult:
    """Execute all jobs in a stage respecting parallelism constraints.

    Args:
        stage_name: Name of the stage being executed.
        job_definitions: Job dictionaries with command, payload, outputs, etc.
        parallelism_policy: Parallelism settings (max_concurrent, retry_max, timeout_sec).
        working_dir: Working directory for job execution.

    Returns:
        StageExecutionResult describing job outcomes and collected artifacts.
    """

shutdown()

Clean up scheduler resources (optional for implementations).

Source code in src/aid2e/schedulers/base.py

def shutdown(self) -> None:
    """Clean up scheduler resources (optional for implementations)."""

    return None

submit_stage(stage_name, job_definitions, parallelism_policy=None, working_dir=None) abstractmethod

Submit a stage for asynchronous execution.

Unlike run_stage, which may block until completion and return a StageExecutionResult, submit_stage should schedule the stage and return immediately with a stage_id that can be used to poll status and retrieve results later.

Returns:

Type	Description
str	A unique stage_id string that identifies the submitted stage.

Source code in src/aid2e/schedulers/base.py

@abstractmethod
def submit_stage(
    self,
    stage_name: str,
    job_definitions: List[Dict[str, Any]],
    parallelism_policy: Optional[Dict[str, Any]] = None,
    working_dir: Optional[str] = None,
) -> str:
    """Submit a stage for asynchronous execution.

    Unlike `run_stage`, which may block until completion and return a
    `StageExecutionResult`, `submit_stage` should schedule the stage and
    return immediately with a `stage_id` that can be used to poll status
    and retrieve results later.

    Returns:
        A unique `stage_id` string that identifies the submitted stage.
    """

JobStatus

Bases: BaseModel

Represent status information for a single job.

Parameters:

Name	Type	Description	Default
job_id		Unique job identifier.	required
status		Current status ("queued", "running", "completed", "failed", "cancelled").	required
return_code		Exit code when completed or failed.	required
stdout		Standard output from the job.	required
stderr		Standard error from the job.	required
outputs		Optional output data (e.g., objectives, results from Python callables).	required
metrics		Optional metrics (e.g., runtime, memory usage).	required

Source code in src/aid2e/schedulers/base.py

class JobStatus(BaseModel):
    """Represent status information for a single job.

    Args:
        job_id: Unique job identifier.
        status: Current status ("queued", "running", "completed", "failed", "cancelled").
        return_code: Exit code when completed or failed.
        stdout: Standard output from the job.
        stderr: Standard error from the job.
        outputs: Optional output data (e.g., objectives, results from Python callables).
        metrics: Optional metrics (e.g., runtime, memory usage).
    """

    job_id: str
    status: str
    return_code: Optional[int] = None
    stdout: Optional[str] = None
    stderr: Optional[str] = None
    outputs: Optional[Dict[str, Any]] = None
    metrics: Optional[Dict[str, Any]] = None

StageExecutionResult

Bases: BaseModel

Capture the result of executing all jobs in a stage.

Parameters:

Name	Type	Description	Default
stage_name		Name of the executed stage.	required
job_statuses		Status for each job in the stage.	required
artifacts		Output artifacts collected from the stage (path -> content).	required
success		Whether all jobs completed successfully.	required
error_message		Optional error message if the stage failed.	required

Source code in src/aid2e/schedulers/base.py

class StageExecutionResult(BaseModel):
    """Capture the result of executing all jobs in a stage.

    Args:
        stage_name: Name of the executed stage.
        job_statuses: Status for each job in the stage.
        artifacts: Output artifacts collected from the stage (path -> content).
        success: Whether all jobs completed successfully.
        error_message: Optional error message if the stage failed.
    """

    stage_name: str
    job_statuses: List[JobStatus]
    artifacts: Dict[str, Any]
    success: bool
    error_message: Optional[str] = None

__getattr__(name)

Lazy-load JobLib subpackage when accessed.

Source code in src/aid2e/schedulers/__init__.py

def __getattr__(name: str):
    """Lazy-load JobLib subpackage when accessed."""
    if name == "JobLib":
        from aid2e.schedulers import JobLib as _joblib
        return _joblib
    if name == "JobLibRunnerConfig":
        from aid2e.schedulers.JobLib import JobLibRunnerConfig
        return JobLibRunnerConfig
    if name == "JobLibScheduler":
        from aid2e.schedulers.JobLib import JobLibScheduler
        return JobLibScheduler
    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

get(name)

Retrieve a scheduler class by name (with lazy loading).

Parameters:

Name	Type	Description	Default
name	str	Identifier that was used during registration.	required

Returns:

Type	Description
Type[BaseScheduler]	Scheduler class implementing BaseScheduler.

Raises:

Type	Description
KeyError	If the scheduler is not registered.

Source code in src/aid2e/schedulers/_registry.py

def get(name: str) -> Type[BaseScheduler]:
    """Retrieve a scheduler class by name (with lazy loading).

    Args:
        name: Identifier that was used during registration.

    Returns:
        Scheduler class implementing BaseScheduler.

    Raises:
        KeyError: If the scheduler is not registered.
    """

    name_key = name.lower()

    # First check if already loaded
    if name_key in _scheduler_registry:
        return _scheduler_registry[name_key]

    # Try lazy loader
    if name_key in _scheduler_loaders:
        scheduler_class = _scheduler_loaders[name_key]()
        _scheduler_registry[name_key] = scheduler_class
        return scheduler_class

    available = list(_scheduler_registry.keys()) + list(_scheduler_loaders.keys())
    raise KeyError(f"Scheduler '{name}' not registered. Available: {available}")

is_registered(name)

Return True if the scheduler name is registered or can be lazy-loaded.

Source code in src/aid2e/schedulers/_registry.py

def is_registered(name: str) -> bool:
    """Return True if the scheduler name is registered or can be lazy-loaded."""

    name_key = name.lower()
    return name_key in _scheduler_registry or name_key in _scheduler_loaders

list_registered()

Return a copy of the registered schedulers mapping.

Loads all lazy-registered schedulers.

Source code in src/aid2e/schedulers/_registry.py

def list_registered() -> Dict[str, Type[BaseScheduler]]:
    """Return a copy of the registered schedulers mapping.

    Loads all lazy-registered schedulers.
    """

    # Load all lazy schedulers
    for name in list(_scheduler_loaders.keys()):
        if name not in _scheduler_registry:
            try:
                _scheduler_registry[name] = _scheduler_loaders[name]()
            except Exception:
                pass  # Skip if fails to load

    return _scheduler_registry.copy()

register(name, scheduler_class)

Register a scheduler implementation.

Parameters:

Name	Type	Description	Default
name	str	Identifier used to retrieve the scheduler (case-insensitive).	required
scheduler_class	Type[BaseScheduler]	Scheduler class that implements BaseScheduler.	required

Raises:

Type	Description
ValueError	If name is already registered or scheduler_class is invalid.

Source code in src/aid2e/schedulers/_registry.py

def register(name: str, scheduler_class: Type[BaseScheduler]) -> None:
    """Register a scheduler implementation.

    Args:
        name: Identifier used to retrieve the scheduler (case-insensitive).
        scheduler_class: Scheduler class that implements BaseScheduler.

    Raises:
        ValueError: If name is already registered or scheduler_class is invalid.
    """

    name_key = name.lower()
    if name_key in _scheduler_registry:
        raise ValueError(f"Scheduler '{name}' already registered")
    if not issubclass(scheduler_class, BaseScheduler):
        raise ValueError("Scheduler class must inherit from BaseScheduler")

    _scheduler_registry[name_key] = scheduler_class