From f7ccf222bf8776220257935a1c5df7241960db04 Mon Sep 17 00:00:00 2001 From: wangliwei Date: Tue, 30 Jun 2026 14:51:37 +0800 Subject: [PATCH] feat(runtime): support mapping-based scheduler configs Cause: environments currently accept scheduler strings or ready-made scheduler objects, which makes policy-driven runtime configuration harder to express consistently. Change: let resolve_scheduler() accept mapping configs via type/name/kind, normalize common load-aware aliases, and pass through LoadAwareScheduler options such as strategy and max_concurrent. Validation: python3 -m compileall src/sage/runtime/scheduler.py src/tests/test_runtime_local_consolidation.py; PYTHONPATH=/private/tmp/sage-testdeps:src python3 -m pytest src/tests/test_runtime_local_consolidation.py -q -k scheduler --- src/sage/runtime/scheduler.py | 57 ++++++++++++++++++- src/tests/test_runtime_local_consolidation.py | 35 ++++++++++++ 2 files changed, 90 insertions(+), 2 deletions(-) diff --git a/src/sage/runtime/scheduler.py b/src/sage/runtime/scheduler.py index 66bd995d7..9163189e9 100644 --- a/src/sage/runtime/scheduler.py +++ b/src/sage/runtime/scheduler.py @@ -3,6 +3,7 @@ import os import time from abc import ABC, abstractmethod +from collections.abc import Mapping from dataclasses import dataclass, field from typing import Any @@ -251,6 +252,55 @@ def _has_scheduler_interface(value: Any) -> bool: return callable(getattr(value, "make_decision", None)) +def _normalize_scheduler_name(raw_value: Any) -> str: + normalized = str(raw_value or "").strip().lower().replace("-", "_").replace(" ", "_") + if not normalized: + raise ValueError("scheduler name must be a non-empty string.") + return normalized + + +def _coerce_positive_int(raw_value: Any, *, field_name: str) -> int: + try: + value = int(raw_value) + except (TypeError, ValueError) as exc: + raise ValueError(f"{field_name} must be a positive integer.") from exc + if value <= 0: + raise ValueError(f"{field_name} must be a positive integer.") + return value + + +def _resolve_scheduler_from_mapping(*, scheduler: Mapping[str, Any], platform: str): + scheduler_type = ( + scheduler.get("type") + or scheduler.get("name") + or scheduler.get("kind") + ) + if scheduler_type is None: + raise ValueError( + "scheduler mapping must include one of: 'type', 'name', or 'kind'." + ) + + scheduler_name = _normalize_scheduler_name(scheduler_type) + if scheduler_name == "fifo": + return FIFOScheduler(platform=platform) + if scheduler_name in {"load_aware", "loadaware"}: + strategy = str(scheduler.get("strategy") or "balanced").strip() or "balanced" + max_concurrent = _coerce_positive_int( + scheduler.get("max_concurrent", 10), + field_name="scheduler.max_concurrent", + ) + return LoadAwareScheduler( + platform=platform, + max_concurrent=max_concurrent, + strategy=strategy, + ) + + raise ValueError( + "Unknown scheduler type: " + f"{scheduler_type}. Available options: 'fifo', 'load_aware'" + ) + + def create_default_scheduler(*, platform: str): return FIFOScheduler(platform=platform) @@ -260,7 +310,7 @@ def resolve_scheduler(*, scheduler: Any, platform: str): return create_default_scheduler(platform=platform) if isinstance(scheduler, str): - scheduler_lower = scheduler.lower() + scheduler_lower = _normalize_scheduler_name(scheduler) if scheduler_lower == "fifo": return FIFOScheduler(platform=platform) if scheduler_lower in {"load_aware", "loadaware"}: @@ -269,11 +319,14 @@ def resolve_scheduler(*, scheduler: Any, platform: str): f"Unknown scheduler type: {scheduler}. Available options: 'fifo', 'load_aware'" ) + if isinstance(scheduler, Mapping): + return _resolve_scheduler_from_mapping(scheduler=scheduler, platform=platform) + if _has_scheduler_interface(scheduler): return scheduler raise TypeError( - "scheduler must be None, str, or an object implementing make_decision(), " + "scheduler must be None, str, mapping, or an object implementing make_decision(), " f"got {type(scheduler)}" ) diff --git a/src/tests/test_runtime_local_consolidation.py b/src/tests/test_runtime_local_consolidation.py index 2d4bd1330..6703d7b71 100644 --- a/src/tests/test_runtime_local_consolidation.py +++ b/src/tests/test_runtime_local_consolidation.py @@ -257,11 +257,46 @@ def test_packet_and_stop_signal_are_main_repo_owned() -> None: def test_scheduler_resolution_uses_in_tree_implementations() -> None: fifo = resolve_scheduler(scheduler="fifo", platform="local") load_aware = resolve_scheduler(scheduler="load_aware", platform="local") + configured_load_aware = resolve_scheduler( + scheduler={ + "type": "load-aware", + "strategy": "spread", + "max_concurrent": 4, + }, + platform="local", + ) assert isinstance(fifo, FIFOScheduler) assert isinstance(load_aware, LoadAwareScheduler) + assert isinstance(configured_load_aware, LoadAwareScheduler) assert fifo.get_metrics()["scheduler_type"] == "FIFO" assert load_aware.get_metrics()["scheduler_type"] == "LoadAware" + assert configured_load_aware.strategy == "spread" + assert configured_load_aware.max_concurrent == 4 + + +def test_scheduler_resolution_accepts_scheduler_mapping_aliases() -> None: + scheduler = resolve_scheduler( + scheduler={ + "name": "load aware", + "strategy": "pack", + "max_concurrent": 2, + }, + platform="local", + ) + + assert isinstance(scheduler, LoadAwareScheduler) + assert scheduler.strategy == "pack" + assert scheduler.max_concurrent == 2 + + +def test_scheduler_resolution_rejects_scheduler_mapping_without_type() -> None: + try: + resolve_scheduler(scheduler={"strategy": "spread"}, platform="local") + except ValueError as exc: + assert "scheduler mapping must include one of" in str(exc) + else: # pragma: no cover - defensive guard + raise AssertionError("resolve_scheduler() should reject mapping configs without type") def test_local_environment_batch_submit_runs_without_kernel_dependency() -> None: