diff --git a/config.py b/config.py index 9f8ffa4..d99603a 100644 --- a/config.py +++ b/config.py @@ -21,7 +21,22 @@ class Instance(BaseModel): class Scale(BaseModel): minimum_jobs: int = 11 maximum_jobs: int = 11 - adaptive: bool = False + worker_recovery_timeout_seconds: float = 600.0 + worker_recovery_check_interval_seconds: float = 10.0 + + @model_validator(mode="after") + def validate_limits(self): + if self.minimum_jobs < 1: + raise ValueError("minimum_jobs must be at least 1") + if self.maximum_jobs < self.minimum_jobs: + raise ValueError("maximum_jobs must be >= minimum_jobs") + if self.worker_recovery_timeout_seconds <= 0: + raise ValueError("worker_recovery_timeout_seconds must be positive") + if self.worker_recovery_check_interval_seconds <= 0: + raise ValueError( + "worker_recovery_check_interval_seconds must be positive" + ) + return self instance: Instance = Instance() scale: Scale = Scale() diff --git a/config.template.yaml b/config.template.yaml index bc9bfba..80a8755 100644 --- a/config.template.yaml +++ b/config.template.yaml @@ -25,7 +25,8 @@ executor: # scale: # minimum_jobs: 11 # maximum_jobs: 11 -# adaptive: false +# worker_recovery_timeout_seconds: 600 +# worker_recovery_check_interval_seconds: 10 inputs: specz: diff --git a/packages/executor.py b/packages/executor.py index 9d5a8ab..8e8b2ca 100644 --- a/packages/executor.py +++ b/packages/executor.py @@ -46,7 +46,7 @@ def get_executor(executor_config: Dict[str, Any], logs_dir: str | None = None): - Local: start a LocalCluster(**args). - SLURM: start SLURMCluster with n_workers = minimum_jobs * processes. Each job will use the provided cores, processes, and memory. - Optionally enable adapt() when scale.adaptive is true. + Always enable adaptive scaling after submitting the minimum jobs. Args: executor_config: Dictionary with: @@ -58,7 +58,6 @@ def get_executor(executor_config: Dict[str, Any], logs_dir: str | None = None): - scale: * minimum_jobs (int) * maximum_jobs (int) - * adaptive (bool, default false) logs_dir: Directory to store SLURM job logs (if provided). Returns: @@ -105,8 +104,6 @@ def get_executor(executor_config: Dict[str, Any], logs_dir: str | None = None): min_jobs = int(scale_cfg.get("minimum_jobs", 0)) max_jobs = int(scale_cfg.get("maximum_jobs", 0) or 0) - adaptive = bool(scale_cfg.get("adaptive", False)) - # n_workers in SLURMCluster = total worker processes n_workers_init = min_jobs * processes @@ -128,18 +125,14 @@ def get_executor(executor_config: Dict[str, Any], logs_dir: str | None = None): cluster = SLURMCluster(n_workers=n_workers_init, **instance_cfg) logger.info("SLURMCluster started with instance args=%s", instance_cfg) - # Adaptive SLURM scaling can thrash on this pipeline because its task - # graph alternates between driver-heavy and worker-heavy stages. Keep a - # stable allocation unless adaptive scaling is explicitly requested. - if adaptive and max_jobs > 0: - cluster.adapt(minimum_jobs=min_jobs, maximum_jobs=max_jobs) - logger.info( - "Adaptive scaling enabled: minimum_jobs=%d maximum_jobs=%d", - min_jobs, - max_jobs, - ) - else: - logger.info("Fixed SLURM allocation enabled: jobs=%d", min_jobs) + # Keep adaptive scaling enabled even when min == max. In that case it + # acts as a fixed-size allocation that can replace jobs lost mid-run. + cluster.adapt(minimum_jobs=min_jobs, maximum_jobs=max_jobs) + logger.info( + "Adaptive scaling enabled: minimum_jobs=%d maximum_jobs=%d", + min_jobs, + max_jobs, + ) return cluster diff --git a/packages/worker_health.py b/packages/worker_health.py new file mode 100644 index 0000000..a8b3bf1 --- /dev/null +++ b/packages/worker_health.py @@ -0,0 +1,121 @@ +"""Lightweight monitoring for the minimum Dask worker floor.""" + +from __future__ import annotations + +import logging +import threading +import time +from collections.abc import Callable +from typing import Any + + +class WorkerFloorMonitor: + """Abort a run when the cluster stays below its configured worker floor.""" + + def __init__( + self, + client: Any, + minimum_workers: int, + recovery_timeout_seconds: float = 600.0, + check_interval_seconds: float = 10.0, + logger: logging.LoggerAdapter | None = None, + on_timeout: Callable[[], None] | None = None, + ) -> None: + self.client = client + self.minimum_workers = minimum_workers + self.recovery_timeout_seconds = recovery_timeout_seconds + self.check_interval_seconds = check_interval_seconds + self.logger = logger or logging.LoggerAdapter( + logging.getLogger("crc.worker_health"), {"phase": "resources"} + ) + self.on_timeout = on_timeout + self._below_since: float | None = None + self._timed_out = False + self._query_failure_logged = False + self._stop_event = threading.Event() + self._thread: threading.Thread | None = None + + @property + def timed_out(self) -> bool: + """Whether the worker floor failed to recover within the timeout.""" + return self._timed_out + + def check_once(self, now: float | None = None) -> None: + """Evaluate one scheduler snapshot.""" + if self._timed_out or self.minimum_workers <= 0: + return + now = time.monotonic() if now is None else now + try: + current_workers = len(self.client.scheduler_info().get("workers", {})) + except Exception as exc: + if not self._query_failure_logged: + self.logger.warning( + "Could not query worker count; worker-floor timer unchanged: %s", + exc, + ) + self._query_failure_logged = True + return + + self._query_failure_logged = False + if current_workers >= self.minimum_workers: + if self._below_since is not None: + self.logger.info( + "Worker floor recovered: current=%d minimum=%d after %.1fs.", + current_workers, + self.minimum_workers, + now - self._below_since, + ) + self._below_since = None + return + + if self._below_since is None: + self._below_since = now + self.logger.warning( + "Worker count below configured floor: current=%d minimum=%d; " + "waiting up to %.1fs for adaptive recovery.", + current_workers, + self.minimum_workers, + self.recovery_timeout_seconds, + ) + return + + elapsed = now - self._below_since + if elapsed < self.recovery_timeout_seconds: + return + + self._timed_out = True + self.logger.error( + "Worker floor did not recover: current=%d minimum=%d elapsed=%.1fs; " + "aborting pipeline.", + current_workers, + self.minimum_workers, + elapsed, + ) + if self.on_timeout is not None: + try: + self.on_timeout() + except Exception: + self.logger.exception("Failed to close the degraded Dask cluster.") + + def start(self) -> None: + """Start monitoring in a daemon thread.""" + if self._thread is not None or self.minimum_workers <= 0: + return + self._thread = threading.Thread( + target=self._run, + name="crc-worker-floor-monitor", + daemon=True, + ) + self._thread.start() + + def _run(self) -> None: + while not self._stop_event.is_set() and not self._timed_out: + self.check_once() + self._stop_event.wait(self.check_interval_seconds) + + def stop(self) -> None: + """Stop monitoring without triggering the recovery callback.""" + self._stop_event.set() + if self._thread is not None and self._thread is not threading.current_thread(): + self._thread.join(timeout=max(1.0, self.check_interval_seconds * 2)) + diff --git a/scripts/crd-run.py b/scripts/crd-run.py index cb61f08..daa7d6e 100644 --- a/scripts/crd-run.py +++ b/scripts/crd-run.py @@ -51,6 +51,7 @@ from product_handle import save_dataframe from resource_usage import ResourceUsageMonitor from specz import prepare_catalog +from worker_health import WorkerFloorMonitor from utils import ( configure_exception_hook, configure_warning_handler, @@ -66,6 +67,7 @@ __all__ = ["main"] _resource_usage_monitor: ResourceUsageMonitor | None = None +_worker_floor_monitor: WorkerFloorMonitor | None = None # ----------------------- @@ -565,7 +567,7 @@ def main( config_path: str, cwd: str = ".", base_dir_override: str | None = None ) -> None: """Run the CRC pipeline end-to-end.""" - global _resource_usage_monitor + global _resource_usage_monitor, _worker_floor_monitor delete_temp_files = True # set True to aggressively clean intermediates @@ -770,7 +772,7 @@ def main( _resource_usage_monitor = ResourceUsageMonitor(client) _resource_usage_monitor.start() - # Ensure the minimum number of workers start within 10 seconds + # Do not start pipeline work until the configured worker floor is ready. exec_args = config.get("executor", {}).get("args", {}) or {} instance_cfg = exec_args.get("instance", {}) or {} scale_cfg = exec_args.get("scale", {}) or {} @@ -778,26 +780,52 @@ def main( procs = int(instance_cfg.get("processes", 1) or 1) min_jobs = scale_cfg.get("minimum_jobs") min_workers = 0 if min_jobs is None else int(min_jobs) * procs + recovery_timeout = float( + scale_cfg.get("worker_recovery_timeout_seconds", 600.0) + ) + recovery_interval = float( + scale_cfg.get("worker_recovery_check_interval_seconds", 10.0) + ) + + def _abort_degraded_cluster() -> None: + try: + client.close(timeout=5) + finally: + cluster.close() if min_workers > 0: log_init.info( - "Waiting up to 10s for minimum_workers=%d to start...", min_workers + "Waiting up to %.1fs for minimum_workers=%d to start...", + recovery_timeout, + min_workers, ) try: - client.wait_for_workers(min_workers, timeout=10) - except Exception: - current_workers = len(client.scheduler_info().get("workers", {})) - log_init.warning( - "Timeout: waited 10 seconds but minimum_workers=%d did not start " - "(current=%d). Proceeding anyway.", + client.wait_for_workers(min_workers, timeout=recovery_timeout) + except Exception as exc: + try: + current_workers = len(client.scheduler_info().get("workers", {})) + except Exception: + current_workers = -1 + log_init.error( + "Minimum worker floor was not reached during startup: " + "minimum_workers=%d current=%d timeout=%.1fs. Aborting.", min_workers, current_workers, + recovery_timeout, ) + try: + _abort_degraded_cluster() + except Exception: + log_init.exception("Failed to close cluster after startup timeout.") + raise RuntimeError( + "Dask minimum worker floor was not reached during startup." + ) from exc else: current_workers = len(client.scheduler_info().get("workers", {})) log_init.info( - "Confirmed: minimum_workers=%d started within 10s (current=%d).", + "Confirmed: minimum_workers=%d started within %.1fs (current=%d).", min_workers, + recovery_timeout, current_workers, ) @@ -807,6 +835,24 @@ def main( current_workers, ) + if min_workers > 0: + _worker_floor_monitor = WorkerFloorMonitor( + client, + minimum_workers=min_workers, + recovery_timeout_seconds=recovery_timeout, + check_interval_seconds=recovery_interval, + logger=_phase_logger(base_logger, "resources"), + on_timeout=_abort_degraded_cluster, + ) + _worker_floor_monitor.start() + log_init.info( + "Worker-floor monitor enabled: minimum_workers=%d " + "recovery_timeout=%.1fs check_interval=%.1fs.", + min_workers, + recovery_timeout, + recovery_interval, + ) + log_init.info("END init: pipeline bootstrap") # Dask perf report (global) @@ -1496,6 +1542,8 @@ def _submit_pair(executor: ThreadPoolExecutor) -> bool: base_logger.error( "Unknown combine_mode: %s", combine_mode, extra={"phase": "crossmatch"} ) + if _worker_floor_monitor is not None: + _worker_floor_monitor.stop() client.close() cluster.close() return @@ -2074,6 +2122,11 @@ def _submit_pair(executor: ThreadPoolExecutor) -> bool: # Update process info try: + if _worker_floor_monitor is not None and _worker_floor_monitor.timed_out: + raise RuntimeError( + "Dask worker count stayed below the configured minimum beyond " + "the recovery timeout." + ) update_process_info( process_info, process_info_path, @@ -2208,6 +2261,8 @@ def _submit_pair(executor: ThreadPoolExecutor) -> bool: log_cons.warning("Could not delete temp_dir %s: %s", temp_dir, e) log_cons.info("END consolidation: export complete") + if _worker_floor_monitor is not None: + _worker_floor_monitor.stop() client.close() cluster.close() @@ -2271,6 +2326,8 @@ def _pick_next_process_dir(root: str) -> str: _resource_usage_monitor.report( logging.LoggerAdapter(lg, {"phase": "resources"}) ) + if _worker_floor_monitor is not None: + _worker_floor_monitor.stop() msg = f"Pipeline {'completed successfully' if ok else 'terminated with errors'} in {dur:.2f} seconds. (run dir: {base_dir})" if lg.handlers: logging.LoggerAdapter(lg, {"phase": "consolidation"}).info(msg) diff --git a/tests/test_executor_scaling.py b/tests/test_executor_scaling.py index 7dad6cc..31f4218 100644 --- a/tests/test_executor_scaling.py +++ b/tests/test_executor_scaling.py @@ -15,7 +15,7 @@ def adapt(self, **kwargs): self.adapt_calls.append(kwargs) -def test_slurm_executor_uses_fixed_allocation_by_default(monkeypatch): +def test_slurm_executor_starts_minimum_workers_and_always_enables_adapt(monkeypatch): monkeypatch.setattr(executor, "SLURMCluster", FakeSlurmCluster) cluster = executor.get_executor( @@ -29,10 +29,10 @@ def test_slurm_executor_uses_fixed_allocation_by_default(monkeypatch): ) assert cluster.kwargs["n_workers"] == 11 - assert cluster.adapt_calls == [] + assert cluster.adapt_calls == [{"minimum_jobs": 11, "maximum_jobs": 11}] -def test_slurm_executor_enables_adaptive_scaling_only_explicitly(monkeypatch): +def test_slurm_executor_supports_different_adaptive_limits(monkeypatch): monkeypatch.setattr(executor, "SLURMCluster", FakeSlurmCluster) cluster = executor.get_executor( @@ -43,7 +43,6 @@ def test_slurm_executor_enables_adaptive_scaling_only_explicitly(monkeypatch): "scale": { "minimum_jobs": 7, "maximum_jobs": 15, - "adaptive": True, }, }, } diff --git a/tests/test_worker_health.py b/tests/test_worker_health.py new file mode 100644 index 0000000..59d12a5 --- /dev/null +++ b/tests/test_worker_health.py @@ -0,0 +1,55 @@ +import logging + +from worker_health import WorkerFloorMonitor + + +class SnapshotClient: + def __init__(self, worker_counts): + self.worker_counts = iter(worker_counts) + + def scheduler_info(self): + count = next(self.worker_counts) + return {"workers": {f"worker-{index}": {} for index in range(count)}} + + +def _monitor(worker_counts, on_timeout=None): + return WorkerFloorMonitor( + SnapshotClient(worker_counts), + minimum_workers=3, + recovery_timeout_seconds=600, + logger=logging.LoggerAdapter(logging.getLogger("test"), {}), + on_timeout=on_timeout, + ) + + +def test_worker_floor_timer_resets_after_recovery(): + monitor = _monitor([2, 3, 2]) + + monitor.check_once(now=0) + monitor.check_once(now=100) + monitor.check_once(now=200) + + assert monitor._below_since == 200 + assert not monitor.timed_out + + +def test_worker_floor_does_not_abort_before_timeout(): + called = [] + monitor = _monitor([2, 2], on_timeout=lambda: called.append(True)) + + monitor.check_once(now=0) + monitor.check_once(now=599) + + assert not monitor.timed_out + assert called == [] + + +def test_worker_floor_aborts_after_timeout(): + called = [] + monitor = _monitor([2, 2], on_timeout=lambda: called.append(True)) + + monitor.check_once(now=0) + monitor.check_once(now=600) + + assert monitor.timed_out + assert called == [True]