utilities/configurations/scheduler_cascade.py

0001 """Scheduler cascade resolution utilities.
0002
0003 This module provides functions to resolve the effective scheduler configuration
0004 for a given stage/branch/workflow, following the cascade precedence:
0005
0006     objective-level default → workflow-level default → branch-level default → stage-level override
0007
0008 The cascade allows users to set sensible defaults at higher levels and override
0009 them at lower levels only when needed.
0010 """
0011
0012 from typing import Optional, Dict, Any
0013 from pydantic import BaseModel
0014
0015 from aid2e.utilities.configurations.scheduler_config import SchedulerConfiguration
0016
0017
0018 def resolve_scheduler_cascade(
0019     stage_scheduler: Optional[SchedulerConfiguration] = None,
0020     branch_scheduler: Optional[SchedulerConfiguration] = None,
0021     workflow_scheduler: Optional[SchedulerConfiguration] = None,
0022     objective_scheduler: Optional[SchedulerConfiguration] = None,
0023     global_scheduler: Optional[SchedulerConfiguration] = None,
0024 ) -> Optional[SchedulerConfiguration]:
0025     """
0026     Resolve the effective scheduler configuration using cascade precedence.
0027
0028     Cascade order (highest to lowest priority):
0029     1. Stage-level scheduler (stage override)
0030     2. Branch-level scheduler (branch default)
0031     3. Workflow-level scheduler (workflow default)
0032     4. Objective-level scheduler (objective default)
0033     5. Global scheduler (global default)
0034
0035     Args:
0036         stage_scheduler: Scheduler at stage level (highest priority)
0037         branch_scheduler: Scheduler at branch level
0038         workflow_scheduler: Scheduler at workflow level
0039         objective_scheduler: Scheduler at objective level
0040         global_scheduler: Global scheduler (lowest priority)
0041
0042     Returns:
0043         The first non-None scheduler in the cascade, or None if all are None
0044
0045     Example:
0046         >>> stage_sched = SchedulerConfiguration(runner_type="SlurmRunner", parameters={})
0047         >>> branch_sched = SchedulerConfiguration(runner_type="JobLibRunner", parameters={})
0048         >>> effective = resolve_scheduler_cascade(stage_sched, branch_sched)
0049         >>> assert effective == stage_sched  # Stage overrides branch
0050
0051         >>> effective = resolve_scheduler_cascade(None, branch_sched)
0052         >>> assert effective == branch_sched  # Branch used if stage is None
0053     """
0054     # Check in order of precedence
0055     if stage_scheduler is not None:
0056         return stage_scheduler
0057     if branch_scheduler is not None:
0058         return branch_scheduler
0059     if workflow_scheduler is not None:
0060         return workflow_scheduler
0061     if objective_scheduler is not None:
0062         return objective_scheduler
0063     if global_scheduler is not None:
0064         return global_scheduler
0065     return None
0066
0067
0068 def create_scheduler_context(
0069     objective_scheduler: Optional[SchedulerConfiguration] = None,
0070     workflow_scheduler: Optional[SchedulerConfiguration] = None,
0071     branch_scheduler: Optional[SchedulerConfiguration] = None,
0072     stage_scheduler: Optional[SchedulerConfiguration] = None,
0073     global_scheduler: Optional[SchedulerConfiguration] = None,
0074 ) -> Dict[str, Any]:
0075     """
0076     Create a context dictionary with scheduler information for logging/debugging.
0077
0078     Args:
0079         objective_scheduler: Scheduler at objective level
0080         workflow_scheduler: Scheduler at workflow level
0081         branch_scheduler: Scheduler at branch level
0082         stage_scheduler: Scheduler at stage level
0083         global_scheduler: Global scheduler
0084
0085     Returns:
0086         Dictionary with scheduler cascade information for each level
0087
0088     Example:
0089         >>> context = create_scheduler_context(
0090         ...     workflow_scheduler=SchedulerConfiguration(...),
0091         ...     branch_scheduler=SchedulerConfiguration(...),
0092         ... )
0093         >>> print(context["effective_scheduler"])  # Will show the effective one
0094     """
0095     effective = resolve_scheduler_cascade(
0096         stage_scheduler, branch_scheduler, workflow_scheduler, objective_scheduler, global_scheduler
0097     )
0098
0099     return {
0100         "cascade_levels": {
0101             "stage": stage_scheduler.runner_type if stage_scheduler else None,
0102             "branch": branch_scheduler.runner_type if branch_scheduler else None,
0103             "workflow": workflow_scheduler.runner_type if workflow_scheduler else None,
0104             "objective": objective_scheduler.runner_type if objective_scheduler else None,
0105             "global": global_scheduler.runner_type if global_scheduler else None,
0106         },
0107         "effective_scheduler": effective.runner_type if effective else None,
0108         "source": _get_cascade_source(stage_scheduler, branch_scheduler, workflow_scheduler, objective_scheduler, global_scheduler),
0109     }
0110
0111
0112 def _get_cascade_source(
0113     stage_scheduler, branch_scheduler, workflow_scheduler, objective_scheduler, global_scheduler
0114 ) -> str:
0115     """Determine which level provided the effective scheduler."""
0116     if stage_scheduler is not None:
0117         return "stage"
0118     if branch_scheduler is not None:
0119         return "branch"
0120     if workflow_scheduler is not None:
0121         return "workflow"
0122     if objective_scheduler is not None:
0123         return "objective"
0124     if global_scheduler is not None:
0125         return "global"
0126     return "none"