feat: add adaptive worker recovery and minimum worker safeguards

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()

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:

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
 

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))
+

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,34 +772,60 @@ 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 {}
 
     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)

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,
                 },
             },
         }