From 5b1ba79a1ca9d146246f4d772e14b3477ccae7b8 Mon Sep 17 00:00:00 2001 From: Herman Brunborg Date: Sun, 9 Aug 2026 04:40:50 -0700 Subject: [PATCH] Replace orchestrator pools during execution takeover - Add live-run discovery and WebSocket takeover for compatible Slurm pools - Redirect adopted workers to successor coordinators and re-exec them - Expose takeover errors and cover the workflow with tests --- src/furu/__init__.py | 2 + src/furu/dag.py | 8 +- src/furu/execution/execution_coordinator.py | 190 +++++- src/furu/execution/server.py | 70 +- src/furu/execution/takeover.py | 359 ++++++++++ src/furu/utils.py | 7 + src/furu/worker/_cli.py | 100 ++- src/furu/worker/backends/slurm/backend.py | 136 +++- src/furu/worker/backends/slurm/pool.py | 31 + src/furu/worker/endpoint.py | 33 + src/furu/worker/loop.py | 17 +- tests/test_slurm_backend.py | 685 +++++++++++--------- tests/test_takeover.py | 562 ++++++++++++++++ 13 files changed, 1828 insertions(+), 372 deletions(-) create mode 100644 src/furu/execution/takeover.py create mode 100644 src/furu/worker/endpoint.py create mode 100644 tests/test_takeover.py diff --git a/src/furu/__init__.py b/src/furu/__init__.py index bfa1aafc..7073e611 100644 --- a/src/furu/__init__.py +++ b/src/furu/__init__.py @@ -6,6 +6,7 @@ from furu.core import Missing, Spec from furu.dependencies import dependency from furu.diff import diff +from furu.execution.execution_coordinator import FuruReplacedError from furu.execution.load_or_create import create, load_existing from furu.logging import get_logger from furu.migration.steps import ( @@ -39,6 +40,7 @@ __all__ = [ "Added", "Codec", + "FuruReplacedError", "GiB", "Metadata", "MigrationStep", diff --git a/src/furu/dag.py b/src/furu/dag.py index 975893bd..0d931d91 100644 --- a/src/furu/dag.py +++ b/src/furu/dag.py @@ -14,6 +14,10 @@ from furu.execution.execution_coordinator import ExecutionCoordinator +class RunningObjectError(RuntimeError): + """An object is currently being computed by another process.""" + + @dataclass(eq=False) class DagNode: obj: Spec @@ -49,7 +53,9 @@ def _add_to_dag(coordinator: ExecutionCoordinator, objs: Sequence[Spec]) -> None continue case "running": # TODO: handle already-running objects as external dependencies. - raise RuntimeError(f"cannot add running object to DAG: {obj.object_id}") + raise RunningObjectError( + f"cannot add running object to DAG: {obj.object_id}" + ) case "missing" | "failed": pass case "stale": diff --git a/src/furu/execution/execution_coordinator.py b/src/furu/execution/execution_coordinator.py index 21d77874..ff8164c5 100644 --- a/src/furu/execution/execution_coordinator.py +++ b/src/furu/execution/execution_coordinator.py @@ -1,6 +1,7 @@ from __future__ import annotations import hashlib +import os import threading import time from collections.abc import Iterator, Sequence @@ -13,7 +14,13 @@ from furu.config import get_config from furu.core import Spec -from furu.dag import DagNode, _add_to_dag, _update_dag_blocking_dependencies +from furu.dag import ( + DagNode, + RunningObjectError, + _add_to_dag, + _update_dag_blocking_dependencies, +) +from furu.execution.takeover import AdoptedPool, ClaimResponse, PoolInventory from furu.logging import ( _scoped_component, _scoped_log_files, @@ -34,11 +41,19 @@ ) if TYPE_CHECKING: - from furu.worker.backends.protocol import WorkerBackend + from furu.worker.backends.protocol import WorkerBackend, WorkerPool logger = get_logger() +# After a takeover, aborted old-run workers release their compute locks as +# they unwind; a hard-killed worker's lock only goes stale after ~35s. +_POST_TAKEOVER_LOCK_WAIT_S = 45.0 + + +class FuruReplacedError(RuntimeError): + """This run surrendered its workers to a successor run and stopped.""" + @dataclass(frozen=True, slots=True) class RunningJob: @@ -68,6 +83,8 @@ class ExecutionCoordinator: done: threading.Event = field(default_factory=threading.Event) finish_error: str | None = None submit_provenance: SubmitProvenance | None = None + pools: list[WorkerPool] = field(default_factory=list) + replaced_by: str | None = None def _failed_counts(self) -> tuple[int, int]: failed_retry = sum( @@ -100,20 +117,25 @@ def run[ObjsT: Sequence[Spec]]( if max_retries_per_object is None: max_retries_per_object = get_config().worker.max_retries_per_object coordinator = cls(max_retries_per_object=max_retries_per_object) - _add_to_dag(coordinator, objs) digest = hashlib.blake2s(digest_size=16) for obj in objs: digest.update(obj.object_id.encode("utf-8")) digest.update(b"\0") coordinator.executor_id = digest.hexdigest() - if not coordinator.nodes_by_id: - with coordinator.log_context(), coordinator.lock: - logger.info( - "all objects already exist; no execution coordinator work to run" - ) - coordinator._maybe_finish_locked() - return objs + # A takeover must happen before DAG construction: the old run's + # workers are still computing objects this DAG may contain. + replace_selector = os.environ.get("FURU_REPLACE_ORCHESTRATOR") or None + + if replace_selector is None: + _add_to_dag(coordinator, objs) + if not coordinator.nodes_by_id: + with coordinator.log_context(), coordinator.lock: + logger.info( + "all objects already exist; no execution coordinator work to run" + ) + coordinator._maybe_finish_locked() + return objs # One capture (and at most one snapshot build) for the whole batch; # every job carries this same frozen submit half. @@ -129,39 +151,79 @@ def run[ObjsT: Sequence[Spec]]( with coordinator.log_context(): logger.info( - "starting exec=%s · %d ready · %d blocked", + "starting exec=%s — replace with FURU_REPLACE_ORCHESTRATOR=auto", coordinator.executor_id[:5], - len(coordinator.ready), - len(coordinator.blocked), extra=log_detail( executor_id=coordinator.executor_id, executor_dir=coordinator.executor_dir, ), ) - pools = [] try: with execution_coordinator_server( coordinator, bind_host=bind_host, port=port ) as server: logger.info("server listening on %s", server.server_url) - for backend in worker_backends: - pool = backend.start_pool( + adoptions: dict[int, AdoptedPool] = {} + if replace_selector is not None: + from furu.execution.takeover import perform_takeover + + adoptions = perform_takeover( + selector=replace_selector, coordinator=coordinator, - bound_port=server.bound_port, - auth_token=server.auth_token, - executor_dir=coordinator.executor_dir, - provenance=coordinator.submit_provenance, + server=server, + worker_backends=worker_backends, ) - pools.append(pool) + _add_to_dag_after_takeover(coordinator, objs) + if not coordinator.nodes_by_id: + with coordinator.lock: + logger.info( + "all objects already exist; " + "shutting down adopted workers" + ) + coordinator._maybe_finish_locked() + logger.info( + "%d ready · %d blocked", + len(coordinator.ready), + len(coordinator.blocked), + ) + for index, backend in enumerate(worker_backends): + if (adopt := adoptions.get(index)) is not None: + from furu.worker.backends.slurm.backend import ( + SlurmWorkerBackend, + ) + + assert isinstance(backend, SlurmWorkerBackend) + pool = backend.start_pool( + coordinator=coordinator, + bound_port=server.bound_port, + auth_token=server.auth_token, + executor_dir=coordinator.executor_dir, + provenance=coordinator.submit_provenance, + adopt=adopt, + ) + else: + pool = backend.start_pool( + coordinator=coordinator, + bound_port=server.bound_port, + auth_token=server.auth_token, + executor_dir=coordinator.executor_dir, + provenance=coordinator.submit_provenance, + ) + coordinator.pools.append(pool) logger.info("pool started · %s", type(backend).__name__) coordinator.done.wait() finally: - if pools: - with ThreadPoolExecutor(max_workers=len(pools)) as executor: + if coordinator.pools: + with ThreadPoolExecutor( + max_workers=len(coordinator.pools) + ) as executor: stop_futures = [ - executor.submit(pool.stop, timeout=5) for pool in pools + executor.submit(pool.stop, timeout=5) + for pool in coordinator.pools ] - for pool, future in zip(pools, stop_futures, strict=True): + for pool, future in zip( + coordinator.pools, stop_futures, strict=True + ): if (exc := future.exception()) is not None: logger.error( "pool stop failed · %s · %s", @@ -399,7 +461,61 @@ def job_result(self, object_id: str, request: JobResult) -> None: self._maybe_finish_locked() self.lock.notify_all() + def takeover_inventory(self) -> list[PoolInventory]: + from furu.worker.backends.slurm.pool import SlurmWorkerPool + + return [ + pool.takeover_inventory() + for pool in self.pools + if isinstance(pool, SlurmWorkerPool) + ] + + def surrender_pools(self, adopt: list[str]) -> ClaimResponse: + """Hand the named pools' jobs to a successor; the rest are cancelled + by this run's ordinary shutdown path.""" + from furu.worker.backends.slurm.pool import SlurmWorkerPool + + pools_by_id = { + pool.pool_id: pool + for pool in self.pools + if isinstance(pool, SlurmWorkerPool) + } + if unknown := set(adopt) - set(pools_by_id): + raise ValueError( + f"cannot surrender unknown pools: {', '.join(sorted(unknown))}" + ) + adopted: list[PoolInventory] = [] + for pool_id in adopt: + pool = pools_by_id[pool_id] + pool.surrender() + adopted.append(pool.takeover_inventory()) + with self.log_context(): + logger.info( + "surrendered %d of %d slurm pools", + len(adopted), + len(pools_by_id), + extra=log_detail(adopted=",".join(adopt)), + ) + return ClaimResponse( + adopted=adopted, + cancelled=[pool_id for pool_id in pools_by_id if pool_id not in set(adopt)], + ) + + def replaced(self, successor_executor_id: str) -> None: + with self.log_context(), self.lock: + if self.done.is_set(): + return + self.replaced_by = successor_executor_id + self.finish_error = ( + f"orchestrator replaced by successor {successor_executor_id}" + ) + logger.info("furu execution coordinator finished: %s", self.finish_error) + self.done.set() + self.lock.notify_all() + def raise_for_failure(self) -> None: + if self.replaced_by is not None: + raise FuruReplacedError(self.finish_error) if self.finish_error is not None: raise RuntimeError(self.finish_error) @@ -448,3 +564,27 @@ def _maybe_finish_locked(self) -> None: else: logger.info("furu execution coordinator finished successfully") self.done.set() + + +def _add_to_dag_after_takeover( + coordinator: ExecutionCoordinator, objs: Sequence[Spec] +) -> None: + """Build the DAG once the aborted old-run workers' compute locks clear. + + Probes on a throwaway coordinator because a failed ``_add_to_dag`` leaves + partially classified nodes behind; the real build only runs clean. + """ + deadline = time.monotonic() + _POST_TAKEOVER_LOCK_WAIT_S + while True: + probe = ExecutionCoordinator( + max_retries_per_object=coordinator.max_retries_per_object + ) + try: + _add_to_dag(probe, objs) + break + except RunningObjectError: + if time.monotonic() >= deadline: + raise + logger.info("waiting for adopted workers to release compute locks") + time.sleep(1.0) + _add_to_dag(coordinator, objs) diff --git a/src/furu/execution/server.py b/src/furu/execution/server.py index 6461d32d..187a4f0f 100644 --- a/src/furu/execution/server.py +++ b/src/furu/execution/server.py @@ -10,6 +10,13 @@ from websockets.sync.server import ServerConnection, basic_auth, serve from furu.execution.execution_coordinator import ExecutionCoordinator +from furu.execution.takeover import ( + TAKEOVER_PATH, + ClaimRequest, + PoolsRequest, + PoolsResponse, + register_live_run, +) from furu.logging import get_logger, log_detail from furu.worker.protocol import HelloMessage, job_result_adapter @@ -58,6 +65,48 @@ def _serve_worker( coordinator.worker_lost(worker) +def _serve_takeover( + coordinator: ExecutionCoordinator, + connection: ServerConnection, + busy: threading.Lock, +) -> None: + """Two-message handshake with a successor run, on one connection. + + Inventory and claim stay separate so the destructive step is explicit and + the successor can prepare (venv, config, endpoint files) between them; a + drop before the claim aborts the takeover with nothing surrendered. + """ + with coordinator.log_context(): + if not busy.acquire(blocking=False): + logger.warning("rejecting takeover connection: one already in progress") + connection.close(1013, "takeover already in progress") + return + try: + request = PoolsRequest.model_validate_json(connection.recv(timeout=10.0)) + logger.info( + "takeover requested by successor %s", request.successor_executor_id + ) + connection.send( + PoolsResponse( + executor_id=coordinator.executor_id, + executor_dir=coordinator.executor_dir, + pools=coordinator.takeover_inventory(), + ).model_dump_json() + ) + try: + claim = ClaimRequest.model_validate_json(connection.recv()) + except ConnectionClosed: + logger.warning( + "successor disconnected before claiming; keeping all pools" + ) + return + response = coordinator.surrender_pools(claim.adopt) + connection.send(response.model_dump_json()) + coordinator.replaced(request.successor_executor_id) + finally: + busy.release() + + @contextmanager def execution_coordinator_server( coordinator: ExecutionCoordinator, *, bind_host: str, port: int @@ -65,12 +114,19 @@ def execution_coordinator_server( auth_token = token_urlsafe(32) connections: set[ServerConnection] = set() connections_changed = threading.Condition() + takeover_busy = threading.Lock() def handler(connection: ServerConnection) -> None: with connections_changed: connections.add(connection) try: - _serve_worker(coordinator, connection) + if ( + connection.request is not None + and connection.request.path == TAKEOVER_PATH + ): + _serve_takeover(coordinator, connection, takeover_busy) + else: + _serve_worker(coordinator, connection) finally: with connections_changed: connections.discard(connection) @@ -90,11 +146,17 @@ def handler(connection: ServerConnection) -> None: ) thread.start() try: - yield ExecutionCoordinatorServer( - bound_host=bound_host, + with register_live_run( + executor_id=coordinator.executor_id, + executor_dir=coordinator.executor_dir, bound_port=bound_port, auth_token=auth_token, - ) + ): + yield ExecutionCoordinatorServer( + bound_host=bound_host, + bound_port=bound_port, + auth_token=auth_token, + ) finally: coordinator.fail("execution coordinator server closed before the run finished") server.shutdown() diff --git a/src/furu/execution/takeover.py b/src/furu/execution/takeover.py new file mode 100644 index 00000000..3d6cd37b --- /dev/null +++ b/src/furu/execution/takeover.py @@ -0,0 +1,359 @@ +"""Orchestrator takeover: let a new run inherit an old run's Slurm workers. + +Three cooperating pieces (see also ``furu.worker.endpoint``): + +1. every coordinator registers itself in ``/live/`` while alive, + with liveness backed by a heartbeat lock rather than trusting cleanup; +2. a successor holds a two-message handshake (inventory, then claim) on a + single ``/takeover`` WebSocket connection to the old coordinator; +3. between the two messages the successor atomically rewrites the matched + pools' endpoint files, so any worker that wakes up after the old server + stops already finds the new coordinates. +""" + +from __future__ import annotations + +import os +import socket +import subprocess +from collections.abc import Iterator +from contextlib import contextmanager +from dataclasses import dataclass +from datetime import UTC, datetime +from pathlib import Path +from typing import TYPE_CHECKING, Literal + +from pydantic import BaseModel, ConfigDict, ValidationError +from websockets.headers import build_authorization_basic +from websockets.sync.client import connect + +from furu.config import get_config +from furu.locking import is_active_lock, lock, read_text_or_none, unlink_if_exists +from furu.logging import get_logger +from furu.utils import atomic_replace_private_file +from furu.worker.endpoint import read_worker_endpoint, write_worker_endpoint + +if TYPE_CHECKING: + from furu.execution.execution_coordinator import ExecutionCoordinator + from furu.execution.server import ExecutionCoordinatorServer + from furu.worker.backends.protocol import WorkerBackend + +logger = get_logger() + +TAKEOVER_PATH = "/takeover" +_INVENTORY_TIMEOUT_S = 30.0 +# The successor extracts a snapshot and builds a venv between inventory and +# claim, and the old side joins its scale threads before answering. +_CLAIM_TIMEOUT_S = 300.0 +_SLURM_COMMAND_TIMEOUT_S = 60.0 + + +class PoolJob(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + job_id: str + state: str + + +class PoolInventory(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + pool_id: str + fingerprint: str + endpoint_file: Path + jobs: list[PoolJob] + + +class PoolsRequest(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + type: Literal["pools_request"] = "pools_request" + successor_executor_id: str + + +class PoolsResponse(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + type: Literal["pools"] = "pools" + executor_id: str + executor_dir: Path + pools: list[PoolInventory] + + +class ClaimRequest(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + type: Literal["claim"] = "claim" + adopt: list[str] + + +class ClaimResponse(BaseModel): + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + type: Literal["claimed"] = "claimed" + adopted: list[PoolInventory] + cancelled: list[str] + + +@dataclass(frozen=True, slots=True) +class AdoptedPool: + """What a successor pool inherits: identity, indirection file, jobs.""" + + pool_id: str + endpoint_file: Path + job_ids: tuple[str, ...] + + +# -------------------------------------------------------------------------- +# Live-run registry + + +class LiveRunEntry(BaseModel): + """Self-sufficient handle to a live coordinator: one read gives a + successor everything. Contains the auth token: never log it verbatim.""" + + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + executor_id: str + server_url: str + auth_token: str + executor_dir: Path + pid: int + host: str + started_at: str + + +def _live_registry_dir() -> Path: + return get_config().run_directories.executions / "live" + + +def _live_entry_path(executor_id: str) -> Path: + return _live_registry_dir() / f"{executor_id}.json" + + +def _live_lock_path(executor_id: str) -> Path: + return _live_registry_dir() / f"{executor_id}.lock" + + +@contextmanager +def register_live_run( + *, + executor_id: str, + executor_dir: Path, + bound_port: int, + auth_token: str, +) -> Iterator[None]: + """Advertise this coordinator while it runs. + + The heartbeat lock, not the entry file, is the liveness signal: a + SIGKILL'd coordinator's entry goes stale with its lock and is ignored + (and garbage-collected) by the next run that looks. + """ + connect_host = get_config().worker.connect_host or socket.getfqdn() + entry = LiveRunEntry( + executor_id=executor_id, + server_url=f"ws://{connect_host}:{bound_port}", + auth_token=auth_token, + executor_dir=executor_dir, + pid=os.getpid(), + host=socket.gethostname(), + started_at=datetime.now(UTC).isoformat(), + ) + _live_registry_dir().mkdir(parents=True, exist_ok=True) + entry_path = _live_entry_path(executor_id) + with lock(_live_lock_path(executor_id)): + atomic_replace_private_file(entry_path, entry.model_dump_json(indent=2) + "\n") + try: + yield + finally: + unlink_if_exists(entry_path) + + +def _discover_live_run(selector: str, *, exclude_executor_id: str) -> LiveRunEntry: + live_dir = _live_registry_dir() + entries: list[LiveRunEntry] = [] + for entry_path in sorted(live_dir.glob("*.json")) if live_dir.is_dir() else []: + raw = read_text_or_none(entry_path) + if raw is None: + continue + try: + entry = LiveRunEntry.model_validate_json(raw) + except ValidationError: + logger.warning("ignoring malformed live-run entry at %s", entry_path) + continue + if entry.executor_id == exclude_executor_id: + continue + if not is_active_lock(_live_lock_path(entry.executor_id)): + logger.info("garbage-collecting stale live-run entry %s", entry.executor_id) + unlink_if_exists(entry_path) + continue + entries.append(entry) + if selector != "auto": + entries = [entry for entry in entries if entry.executor_id.startswith(selector)] + match entries: + case []: + raise RuntimeError( + f"FURU_REPLACE_ORCHESTRATOR={selector}: no live run to replace" + ) + case [entry]: + return entry + case _: + raise RuntimeError( + f"FURU_REPLACE_ORCHESTRATOR={selector} matches several live runs: " + + ", ".join(entry.executor_id for entry in entries) + + "; disambiguate with an executor-id prefix" + ) + + +# -------------------------------------------------------------------------- +# Successor side + + +def perform_takeover( + *, + selector: str, + coordinator: ExecutionCoordinator, + server: ExecutionCoordinatorServer, + worker_backends: tuple[WorkerBackend, ...], +) -> dict[int, AdoptedPool]: + """Discover, match, rewrite, claim, signal. Returns adoptions keyed by + the matched backend's index in ``worker_backends``.""" + from furu.worker.backends.slurm.backend import SlurmWorkerBackend + + provenance = coordinator.submit_provenance + assert provenance is not None + + entry = _discover_live_run(selector, exclude_executor_id=coordinator.executor_id) + logger.info("taking over coordinator %s at %s", entry.executor_id, entry.server_url) + with connect( + entry.server_url + TAKEOVER_PATH, + additional_headers={ + "Authorization": build_authorization_basic("furu", entry.auth_token) + }, + max_size=None, + ) as connection: + connection.send( + PoolsRequest( + successor_executor_id=coordinator.executor_id + ).model_dump_json() + ) + inventory = PoolsResponse.model_validate_json( + connection.recv(timeout=_INVENTORY_TIMEOUT_S) + ) + + unmatched = {pool.pool_id: pool for pool in inventory.pools} + matched: list[tuple[int, SlurmWorkerBackend, PoolInventory]] = [] + for index, backend in enumerate(worker_backends): + if not isinstance(backend, SlurmWorkerBackend): + continue + fingerprint = backend.fingerprint() + pool = next( + (p for p in unmatched.values() if p.fingerprint == fingerprint), None + ) + if pool is None: + continue + del unmatched[pool.pool_id] + matched.append((index, backend, pool)) + if not matched: + raise RuntimeError( + "FURU_REPLACE_ORCHESTRATOR: no pool of the live run " + f"{inventory.executor_id} matches this run's backends; " + "old pool fingerprints: " + + ( + ", ".join( + f"{pool.pool_id}={pool.fingerprint}" for pool in inventory.pools + ) + or "" + ) + + "; new backend fingerprints: " + + ( + ", ".join( + backend.fingerprint() + for backend in worker_backends + if isinstance(backend, SlurmWorkerBackend) + ) + or "" + ) + + ". Unset FURU_REPLACE_ORCHESTRATOR to start cold." + ) + if unmatched: + logger.info( + "no backend matches old pools %s; the old run will cancel them", + ", ".join(sorted(unmatched)), + ) + + # Rewrite endpoint files before claiming: from the moment the old + # server can shut down, any worker that wakes up must already find + # the new coordinates. + for _, backend, pool in matched: + old_endpoint = read_worker_endpoint(pool.endpoint_file) + write_worker_endpoint( + pool.endpoint_file, + backend.takeover_endpoint( + generation=old_endpoint.generation + 1, + bound_port=server.bound_port, + auth_token=server.auth_token, + executor_dir=coordinator.executor_dir, + provenance=provenance, + ), + ) + + connection.send( + ClaimRequest( + adopt=[pool.pool_id for _, _, pool in matched] + ).model_dump_json() + ) + response = ClaimResponse.model_validate_json( + connection.recv(timeout=_CLAIM_TIMEOUT_S) + ) + + claimed = {pool.pool_id: pool for pool in response.adopted} + adoptions: dict[int, AdoptedPool] = {} + signal_job_ids: list[str] = [] + for index, _, pool in matched: + final = claimed.get(pool.pool_id, pool) + adoptions[index] = AdoptedPool( + pool_id=pool.pool_id, + endpoint_file=pool.endpoint_file, + job_ids=tuple(job.job_id for job in final.jobs), + ) + signal_job_ids.extend( + job.job_id for job in final.jobs if job.state == "RUNNING" + ) + logger.info( + "adopted pool %s · %d jobs · old worker logs remain under %s", + pool.pool_id, + len(final.jobs), + inventory.executor_dir, + ) + if response.cancelled: + logger.info( + "old run cancels its unadopted pools: %s", ", ".join(response.cancelled) + ) + _signal_adopted_workers(signal_job_ids) + return adoptions + + +def _signal_adopted_workers(job_ids: list[str]) -> None: + """Interrupt busy workers without killing their allocations; each unwinds, + re-reads its endpoint file, and re-execs into the new snapshot. On failure + a busy worker still redirects at its next touch of the dead old socket.""" + if not job_ids: + return + try: + result = subprocess.run( + ["scancel", "--signal=USR1", "--batch", *job_ids], + check=False, + capture_output=True, + text=True, + timeout=_SLURM_COMMAND_TIMEOUT_S, + ) + except subprocess.TimeoutExpired: + logger.warning("scancel --signal=USR1 timed out for %s", ",".join(job_ids)) + return + if result.returncode != 0: + logger.warning( + "scancel --signal=USR1 failed for %s: %s", + ",".join(job_ids), + result.stderr.strip(), + ) diff --git a/src/furu/utils.py b/src/furu/utils.py index 4cb0af28..72989032 100644 --- a/src/furu/utils.py +++ b/src/furu/utils.py @@ -124,6 +124,13 @@ def write_private_file(path: Path, contents: str, *, mode: int) -> None: path.chmod(mode) +def atomic_replace_private_file(path: Path, contents: str) -> None: + """Atomically replace ``path`` with ``contents``, 0600 from the first byte.""" + tmp_path = nfs_safe_unique_name(path, name="tmp") + write_private_file(tmp_path, contents, mode=0o600) + tmp_path.rename(path) + + def format_duration(seconds: float) -> str: """Compact human duration for log lines: 850ms, 3.2s, 2m05s, 1h05m.""" if seconds < 1: diff --git a/src/furu/worker/_cli.py b/src/furu/worker/_cli.py index 893b53f9..634fabad 100644 --- a/src/furu/worker/_cli.py +++ b/src/furu/worker/_cli.py @@ -1,23 +1,37 @@ import argparse +import os +import signal +import sys from collections.abc import Sequence from pathlib import Path +from types import FrameType +from websockets.exceptions import WebSocketException + +from furu.config import _WORKER_JSON_CONFIG_FILE_ENV_VAR +from furu.logging import get_logger from furu.resources import ResourceRequest -from furu.worker.loop import worker_loop +from furu.worker.endpoint import read_worker_endpoint +from furu.worker.loop import WorkerPreempted, worker_loop + +logger = get_logger("worker.cli") + + +def _raise_preempted(signum: int, frame: FrameType | None) -> None: + raise WorkerPreempted def main(argv: Sequence[str] | None = None) -> int: parser = argparse.ArgumentParser(allow_abbrev=False) parser.add_argument( - "--server-url", - required=True, - help="execution coordinator WebSocket URL (ws://host:port)", - ) - parser.add_argument( - "--auth-token-file", + "--endpoint-file", required=True, type=Path, - help="path to a file containing the execution coordinator auth token", + help=( + "JSON file naming the execution coordinator URL, auth token, and " + "project; re-read after a disconnect so a successor coordinator " + "can redirect this worker" + ), ) parser.add_argument( "--resource-cpus", @@ -55,18 +69,64 @@ def main(argv: Sequence[str] | None = None) -> int: ) args = parser.parse_args(argv) - worker_loop( - server_url=args.server_url, - auth_token=args.auth_token_file.read_text(encoding="utf-8").rstrip(), - resource_request=ResourceRequest( - cpus=args.resource_cpus, - gpus=args.resource_gpus, - memory_gib=args.resource_memory_gib, - ), - idle_timeout=args.idle_timeout, - component=args.component, - backend=args.backend, - ) + endpoint = read_worker_endpoint(args.endpoint_file) + try: + signal.signal(signal.SIGUSR1, _raise_preempted) + except ValueError: + # Not the main thread (e.g. a test runner); preemption then only + # arrives via the connection-closed path. + pass + + error: BaseException | None = None + try: + exit_reason = worker_loop( + server_url=endpoint.server_url, + auth_token=endpoint.auth_token, + resource_request=ResourceRequest( + cpus=args.resource_cpus, + gpus=args.resource_gpus, + memory_gib=args.resource_memory_gib, + ), + idle_timeout=args.idle_timeout, + component=args.component, + backend=args.backend, + ) + except WorkerPreempted: + logger.info("preempted; abandoning in-flight work") + exit_reason = "disconnected" + except (OSError, WebSocketException) as exc: + error = exc + exit_reason = "disconnected" + + if exit_reason == "disconnected": + current = read_worker_endpoint(args.endpoint_file) + if current.generation > endpoint.generation: + logger.info( + "endpoint generation %d -> %d; re-exec against %s", + endpoint.generation, + current.generation, + current.server_url, + ) + os.environ[_WORKER_JSON_CONFIG_FILE_ENV_VAR] = current.config_file + # The submit-side venv belongs to the old project; uv would warn + # before selecting the new --project venv. + os.environ.pop("VIRTUAL_ENV", None) + os.execvp( + "uv", + [ + "uv", + "run", + "--frozen", + "--project", + current.project_root, + "python", + "-m", + "furu.worker._cli", + *(sys.argv[1:] if argv is None else argv), + ], + ) + if error is not None: + raise error return 0 diff --git a/src/furu/worker/backends/slurm/backend.py b/src/furu/worker/backends/slurm/backend.py index 6e30e18f..dc2b3620 100644 --- a/src/furu/worker/backends/slurm/backend.py +++ b/src/furu/worker/backends/slurm/backend.py @@ -1,5 +1,7 @@ from __future__ import annotations +import hashlib +import json import os import secrets import shlex @@ -14,12 +16,14 @@ _WORKER_JSON_CONFIG_FILE_ENV_VAR, get_config, ) +from furu.execution.takeover import AdoptedPool from furu.provenance import EnvironmentIdentity, SubmitProvenance from furu.resources import ResourceRequest from furu.snapshot import extract_snapshot from furu.utils import write_private_file from furu.worker.backends.slurm.pool import SlurmWorkerPool from furu.worker.backends.slurm.resources import SlurmResources +from furu.worker.endpoint import WorkerEndpoint, write_worker_endpoint if TYPE_CHECKING: from furu.execution.execution_coordinator import ExecutionCoordinator @@ -27,6 +31,15 @@ type SlurmExport = Literal["NIL", "ALL"] | tuple[str, ...] | None +def _endpoint_field_lookup(field_name: str) -> str: + """Shell fragment resolving one endpoint-file field at script runtime.""" + return ( + '"$(python3 -c ' + "'import json,sys; print(json.load(open(sys.argv[1]))[sys.argv[2]])' " + f'"$furu_endpoint_file" {field_name})"' + ) + + @dataclass(frozen=True, slots=True) class SlurmWorkerBackend: max_workers: int @@ -45,20 +58,35 @@ class SlurmWorkerBackend: export: SlurmExport = None use_job_arrays: bool = True - def start_pool( - self, - *, - coordinator: ExecutionCoordinator, - bound_port: int, - auth_token: str, - executor_dir: Path, - provenance: SubmitProvenance, - ) -> SlurmWorkerPool: + def fingerprint(self) -> str: + """Hash of exactly the fields that change what a worker *is*. + + Submitter-side concerns (max_workers, poll interval, idle timeout, + connect host/port) and the snapshot are deliberately excluded: an old + pool whose fingerprint matches can serve this backend's jobs. + """ + payload = json.dumps( + { + "sbatch_args": self.resources.to_sbatch_args(), + "job_name": self.job_name, + "export": self.export, + "use_job_arrays": self.use_job_arrays, + "pre_worker_commands": self.pre_worker_commands, + }, + sort_keys=True, + separators=(",", ":"), + ) + return "sha256:" + hashlib.sha256(payload.encode("utf-8")).hexdigest() + + def _server_url(self, bound_port: int) -> str: connect_port = ( bound_port if self.worker_connect_port is None else self.worker_connect_port ) - server_url = f"ws://{self.worker_connect_host}:{connect_port}" + return f"ws://{self.worker_connect_host}:{connect_port}" + def _prepare_project(self, provenance: SubmitProvenance) -> tuple[Path, Path]: + """Resolve (chdir, project_root), extracting the snapshot and building + its venv so workers never race to create it.""" chdir = Path.cwd().resolve() project_root = Path(EnvironmentIdentity.capture().project_root) if provenance.snapshot_id is not None: @@ -78,12 +106,9 @@ def start_pool( env={k: v for k, v in os.environ.items() if k != "VIRTUAL_ENV"}, check=True, ) - worker_dir = executor_dir.resolve() / "workers" - worker_dir.mkdir(parents=True, exist_ok=True) - - token_file = worker_dir / f"worker-{secrets.token_hex(16)}.token" - write_private_file(token_file, auth_token, mode=0o600) + return chdir, project_root + def _write_worker_config(self, worker_dir: Path) -> Path: # Workers may run from a different directory (the extracted snapshot), # so pin any relative data directories to the submit-side anchor. config = get_config() @@ -96,6 +121,63 @@ def start_pool( config.model_dump_json(indent=2) + "\n", mode=0o600, ) + return config_file + + def takeover_endpoint( + self, + *, + generation: int, + bound_port: int, + auth_token: str, + executor_dir: Path, + provenance: SubmitProvenance, + ) -> WorkerEndpoint: + """Endpoint content pointing adopted workers at this run, with the + snapshot venv built and this run's worker config written.""" + _, project_root = self._prepare_project(provenance) + worker_dir = executor_dir.resolve() / "workers" + worker_dir.mkdir(parents=True, exist_ok=True) + return WorkerEndpoint( + generation=generation, + server_url=self._server_url(bound_port), + auth_token=auth_token, + project_root=str(project_root), + config_file=str(self._write_worker_config(worker_dir)), + ) + + def start_pool( + self, + *, + coordinator: ExecutionCoordinator, + bound_port: int, + auth_token: str, + executor_dir: Path, + provenance: SubmitProvenance, + adopt: AdoptedPool | None = None, + ) -> SlurmWorkerPool: + chdir, project_root = self._prepare_project(provenance) + worker_dir = executor_dir.resolve() / "workers" + worker_dir.mkdir(parents=True, exist_ok=True) + + if adopt is None: + pool_id = secrets.token_hex(8) + endpoint_file = worker_dir / f"endpoint-{pool_id}.json" + write_worker_endpoint( + endpoint_file, + WorkerEndpoint( + generation=1, + server_url=self._server_url(bound_port), + auth_token=auth_token, + project_root=str(project_root), + config_file=str(self._write_worker_config(worker_dir)), + ), + ) + else: + # The takeover already rewrote the adopted pool's endpoint file to + # point here; the pool's queued job scripts bake that path, so + # this run's new submissions must keep using the same file. + pool_id = adopt.pool_id + endpoint_file = adopt.endpoint_file resource_request = ResourceRequest( cpus=self.resources.cpus_per_worker, @@ -110,7 +192,7 @@ def start_pool( scripts_dir = worker_dir / "scripts" scripts_dir.mkdir(parents=True, exist_ok=True) - script_path = scripts_dir / f"worker-{secrets.token_hex(16)}.sh" + script_path = scripts_dir / f"worker-{pool_id}.sh" if self.use_job_arrays: component_line = 'furu_worker_component="slurm-worker-${SLURM_ARRAY_JOB_ID}a${SLURM_ARRAY_TASK_ID}"\n' else: @@ -122,8 +204,11 @@ def start_pool( "#!/bin/bash\n" "set -euo pipefail\n" "\n" - "export " - f"{_WORKER_JSON_CONFIG_FILE_ENV_VAR}={shlex.quote(str(config_file))}\n" + # The endpoint file is the one runtime indirection between this + # (copied-at-submit) script and its coordinator: resolving the + # config, project, and server here rather than at submit time + # is what lets a successor run redirect already-queued jobs. + f"furu_endpoint_file={shlex.quote(str(endpoint_file))}\n" "\n" f"{component_line}" "\n" @@ -133,13 +218,16 @@ def start_pool( # and makes uv warn before it selects the snapshot's .venv. "unset VIRTUAL_ENV\n" "\n" + "export " + f"{_WORKER_JSON_CONFIG_FILE_ENV_VAR}=" + f"{_endpoint_field_lookup('config_file')}\n" + "\n" # --frozen forbids silent lock updates on the node; --project # pins the environment regardless of --chdir. "exec uv run --frozen " - f"--project {shlex.quote(str(project_root))} \\\n" + f"--project {_endpoint_field_lookup('project_root')} \\\n" " python -m furu.worker._cli \\\n" - f" --server-url {shlex.quote(server_url)} \\\n" - f" --auth-token-file {shlex.quote(str(token_file))} \\\n" + ' --endpoint-file "$furu_endpoint_file" \\\n' ' --component "${furu_worker_component}" \\\n' " --backend slurm \\\n" f" --idle-timeout {self.worker_idle_timeout} \\\n" @@ -188,7 +276,11 @@ def start_pool( target=lambda: pool_holder[0]._scale_loop(), name="furu-slurm-worker-pool-scale", ), - _job_ids=[], + _job_ids=list(adopt.job_ids) if adopt is not None else [], + _pool_id=pool_id, + _fingerprint=self.fingerprint(), + _endpoint_file=endpoint_file, + _surrendered=threading.Event(), ) pool_holder.append(pool) pool._scale_thread.start() diff --git a/src/furu/worker/backends/slurm/pool.py b/src/furu/worker/backends/slurm/pool.py index 5c56abbf..9221f3bd 100644 --- a/src/furu/worker/backends/slurm/pool.py +++ b/src/furu/worker/backends/slurm/pool.py @@ -8,6 +8,7 @@ from pathlib import Path from typing import TYPE_CHECKING +from furu.execution.takeover import PoolInventory, PoolJob from furu.logging import _scoped_component, get_logger from furu.resources import ResourceRequest @@ -48,11 +49,41 @@ class SlurmWorkerPool: _use_job_arrays: bool _scale_thread: threading.Thread _job_ids: list[str] + _pool_id: str + _fingerprint: str + _endpoint_file: Path + _surrendered: threading.Event + + @property + def pool_id(self) -> str: + return self._pool_id + + def takeover_inventory(self) -> PoolInventory: + states = self._active_job_states() or {} + return PoolInventory( + pool_id=self._pool_id, + fingerprint=self._fingerprint, + endpoint_file=self._endpoint_file, + jobs=[ + PoolJob(job_id=job_id, state=states.get(job_id, "UNKNOWN")) + for job_id in list(self._job_ids) + ], + ) + + def surrender(self) -> None: + """Hand this pool's jobs to a successor: stop scaling now (so the job + list is final) and make ``stop()`` leave the jobs alone.""" + self._surrendered.set() + self._stop_event.set() + self._scale_thread.join(timeout=_SLURM_COMMAND_TIMEOUT_S) def stop(self, *, timeout: float) -> None: with _scoped_component("slurm"): self._stop_event.set() self._scale_thread.join(timeout=timeout) + if self._surrendered.is_set(): + # The jobs belong to the successor run now. + return deadline = time.monotonic() + timeout while time.monotonic() < deadline: diff --git a/src/furu/worker/endpoint.py b/src/furu/worker/endpoint.py new file mode 100644 index 00000000..0cad1091 --- /dev/null +++ b/src/furu/worker/endpoint.py @@ -0,0 +1,33 @@ +from __future__ import annotations + +from pathlib import Path + +from pydantic import BaseModel, ConfigDict + +from furu.utils import atomic_replace_private_file + + +class WorkerEndpoint(BaseModel): + """Who is my coordinator, and what code do I run. + + The worker's sbatch script bakes only this file's *path*; everything that + a takeover must swap together — URL, token, project, config — lives here + and changes under one atomic rename, so a worker can never observe them + out of sync. Contains the auth token: never log it verbatim. + """ + + model_config = ConfigDict(extra="forbid", frozen=True, strict=True) + + generation: int # monotonic; a worker re-execs only if it grew + server_url: str + auth_token: str + project_root: str + config_file: str + + +def read_worker_endpoint(path: Path) -> WorkerEndpoint: + return WorkerEndpoint.model_validate_json(path.read_text(encoding="utf-8")) + + +def write_worker_endpoint(path: Path, endpoint: WorkerEndpoint) -> None: + atomic_replace_private_file(path, endpoint.model_dump_json(indent=2) + "\n") diff --git a/src/furu/worker/loop.py b/src/furu/worker/loop.py index 6b13d4c1..eff15dd8 100644 --- a/src/furu/worker/loop.py +++ b/src/furu/worker/loop.py @@ -2,7 +2,7 @@ import time import traceback -from typing import assert_never +from typing import Literal, assert_never from websockets.exceptions import ConnectionClosed from websockets.headers import build_authorization_basic @@ -20,6 +20,15 @@ logger = get_logger("worker.loop") +class WorkerPreempted(BaseException): + """Raised by the SIGUSR1 handler: abandon in-flight work immediately. + + A BaseException so the fault barriers that convert crashes into + JobFailedResults cannot swallow it; unwinding still releases compute + locks and retires the warm child via the ``finally`` blocks it passes. + """ + + def _run_job( job: protocol.Job, child_slot: ChildSlot ) -> tuple[protocol.JobResult, str | None]: @@ -62,7 +71,7 @@ def worker_loop( idle_timeout: float | None, component: str, backend: str, -) -> None: +) -> Literal["idle", "disconnected"]: worker_backend_token = _worker_backend.set(backend) with _scoped_component(component): child_slot = ChildSlot() @@ -91,10 +100,10 @@ def worker_loop( "no work for %s; worker exiting", format_duration(idle_timeout), ) - return + return "idle" except ConnectionClosed: logger.info("server closed the connection; worker exiting") - return + return "disconnected" job = protocol.Job.model_validate_json(message) task_started_at = time.monotonic() diff --git a/tests/test_slurm_backend.py b/tests/test_slurm_backend.py index 93b80497..6b1afbb8 100644 --- a/tests/test_slurm_backend.py +++ b/tests/test_slurm_backend.py @@ -3,12 +3,12 @@ import json import logging import os -import shlex import shutil import stat import subprocess import sys import textwrap +import threading from collections.abc import Callable from pathlib import Path from typing import Any @@ -24,6 +24,7 @@ get_config, ) from furu.execution.execution_coordinator import ExecutionCoordinator +from furu.execution.takeover import AdoptedPool from furu.provenance import ( EnvironmentIdentity, GitIdentity, @@ -42,6 +43,8 @@ MemoryPerNode, SlurmResources, ) +from furu.worker.endpoint import read_worker_endpoint +from furu.worker.loop import WorkerPreempted class _StubCoordinator(ExecutionCoordinator): @@ -74,7 +77,13 @@ def start(self) -> None: def join(self, timeout: float | None = None) -> None: pass - monkeypatch.setattr(slurm_backend_module.threading, "Thread", NoopThread) + class ThreadingShim: + Thread = NoopThread + Event = threading.Event + + # Patch only the backend module's view of ``threading``: other components + # (e.g. the locking heartbeat) still need real threads. + monkeypatch.setattr(slurm_backend_module, "threading", ThreadingShim) def _submit_provenance() -> SubmitProvenance: @@ -100,13 +109,45 @@ def _submit_provenance() -> SubmitProvenance: ) -def test_worker_cli_reads_auth_token_file( +def _write_endpoint(path: Path, **overrides: object) -> None: + data: dict[str, object] = { + "generation": 1, + "server_url": "ws://execution-coordinator.test:1234", + "auth_token": "secret", + "project_root": "/proj/one", + "config_file": "/cfg/one.json", + } + data.update(overrides) + path.write_text(json.dumps(data)) + + +def _worker_cli_args(endpoint_file: Path, *extra: str) -> list[str]: + return [ + "--endpoint-file", + str(endpoint_file), + "--resource-cpus", + "1", + "--resource-gpus", + "0", + "--resource-memory-gib", + "0", + "--idle-timeout", + "60", + "--component", + "test-worker", + "--backend", + "slurm", + *extra, + ] + + +def test_worker_cli_reads_endpoint_file( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: calls: list[tuple[str, str, ResourceRequest, float | None]] = [] - token_file = tmp_path / "worker.token" - token_file.write_text("secret\n\n") + endpoint_file = tmp_path / "endpoint.json" + _write_endpoint(endpoint_file) def worker_loop( *, @@ -116,48 +157,27 @@ def worker_loop( idle_timeout: float | None, component: str, backend: str, - ) -> None: + ) -> str: calls.append((server_url, auth_token, resource_request, idle_timeout)) + return "idle" monkeypatch.setattr(_cli, "worker_loop", worker_loop) - assert ( - _cli.main( - [ - "--server-url", - "http://execution-coordinator.test", - "--auth-token-file", - str(token_file), - "--resource-cpus", - "1", - "--resource-gpus", - "0", - "--resource-memory-gib", - "0", - "--idle-timeout", - "60", - "--component", - "test-worker", - "--backend", - "slurm", - ] - ) - == 0 - ) + assert _cli.main(_worker_cli_args(endpoint_file)) == 0 assert calls == [ - ("http://execution-coordinator.test", "secret", ResourceRequest(), 60.0) + ("ws://execution-coordinator.test:1234", "secret", ResourceRequest(), 60.0) ] - assert token_file.exists() + assert endpoint_file.exists() -def test_worker_cli_reads_resource_request( +def test_worker_cli_reads_resource_request_and_idle_timeout( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: calls: list[tuple[ResourceRequest, float | None]] = [] - token_file = tmp_path / "worker.token" - token_file.write_text("secret") + endpoint_file = tmp_path / "endpoint.json" + _write_endpoint(endpoint_file) def worker_loop( *, @@ -167,18 +187,17 @@ def worker_loop( idle_timeout: float | None, component: str, backend: str, - ) -> None: + ) -> str: calls.append((resource_request, idle_timeout)) + return "idle" monkeypatch.setattr(_cli, "worker_loop", worker_loop) assert ( _cli.main( [ - "--server-url", - "http://execution-coordinator.test", - "--auth-token-file", - str(token_file), + "--endpoint-file", + str(endpoint_file), "--resource-cpus", "4", "--resource-gpus", @@ -186,7 +205,7 @@ def worker_loop( "--resource-memory-gib", "16", "--idle-timeout", - "30", + "0.25", "--component", "test-worker", "--backend", @@ -196,16 +215,16 @@ def worker_loop( == 0 ) - assert calls == [(ResourceRequest(cpus=4, gpus=1, memory_gib=16), 30.0)] + assert calls == [(ResourceRequest(cpus=4, gpus=1, memory_gib=16), 0.25)] -def test_worker_cli_reads_idle_timeout( +def test_worker_cli_reads_component_override( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: - calls: list[float | None] = [] - token_file = tmp_path / "worker.token" - token_file.write_text("secret") + captured: list[str] = [] + endpoint_file = tmp_path / "endpoint.json" + _write_endpoint(endpoint_file) def worker_loop( *, @@ -215,294 +234,174 @@ def worker_loop( idle_timeout: float | None, component: str, backend: str, - ) -> None: - calls.append(idle_timeout) + ) -> str: + captured.append(component) + return "idle" monkeypatch.setattr(_cli, "worker_loop", worker_loop) - assert ( - _cli.main( - [ - "--server-url", - "http://execution-coordinator.test", - "--auth-token-file", - str(token_file), - "--resource-cpus", - "4", - "--resource-gpus", - "1", - "--resource-memory-gib", - "0", - "--idle-timeout", - "0.25", - "--component", - "test-worker", - "--backend", - "slurm", - ] - ) - == 0 - ) + args = _worker_cli_args(endpoint_file) + args[args.index("--component") + 1] = "worker-a" + assert _cli.main(args) == 0 - assert calls == [0.25] + assert captured == ["worker-a"] -def _run_worker_cli_capturing_component( +@pytest.mark.parametrize( + "missing", + ["--endpoint-file", "--component", "--idle-timeout", "--resource-cpus"], +) +def test_worker_cli_requires_arguments( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch, - token_file: Path, - extra_args: list[str], -) -> str: - captured: list[str] = [] + missing: str, +) -> None: + endpoint_file = tmp_path / "endpoint.json" + _write_endpoint(endpoint_file) - def worker_loop( - *, - server_url: str, - auth_token: str, - resource_request: ResourceRequest, - idle_timeout: float | None, - component: str, - backend: str, - ) -> None: - captured.append(component) + def worker_loop(**kwargs: object) -> str: + raise AssertionError("worker_loop should not be called") monkeypatch.setattr(_cli, "worker_loop", worker_loop) - assert ( - _cli.main( - [ - "--server-url", - "http://execution-coordinator.test", - "--auth-token-file", - str(token_file), - "--resource-cpus", - "1", - "--resource-gpus", - "0", - "--resource-memory-gib", - "0", - "--idle-timeout", - "60", - "--backend", - "slurm", - *extra_args, - ] - ) - == 0 - ) - (component,) = captured - return component + args = _worker_cli_args(endpoint_file) + index = args.index(missing) + del args[index : index + 2] + with pytest.raises(SystemExit) as exc_info: + _cli.main(args) -def test_worker_cli_reads_component_override( - tmp_path: Path, - monkeypatch: pytest.MonkeyPatch, -) -> None: - token_file = tmp_path / "worker.token" - token_file.write_text("secret") - - component = _run_worker_cli_capturing_component( - monkeypatch, token_file, ["--component", "worker-a"] - ) - - assert component == "worker-a" + assert exc_info.value.code == 2 -def test_worker_cli_requires_component( +@pytest.mark.parametrize( + "rejected", + [ + ("--auth-token", "secret"), + ("--auth-token-file", "/tmp/worker.token"), + ("--server-url", "ws://execution-coordinator.test:1"), + ], +) +def test_worker_cli_rejects_direct_connection_arguments( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, + rejected: tuple[str, str], ) -> None: - token_file = tmp_path / "worker.token" - token_file.write_text("secret") + endpoint_file = tmp_path / "endpoint.json" + _write_endpoint(endpoint_file) - def worker_loop( - *, - server_url: str, - auth_token: str, - resource_request: ResourceRequest, - idle_timeout: float | None, - component: str, - backend: str, - ) -> None: + def worker_loop(**kwargs: object) -> str: raise AssertionError("worker_loop should not be called") monkeypatch.setattr(_cli, "worker_loop", worker_loop) with pytest.raises(SystemExit) as exc_info: - _cli.main( - [ - "--server-url", - "http://execution-coordinator.test", - "--auth-token-file", - str(token_file), - "--resource-cpus", - "1", - "--resource-gpus", - "0", - "--resource-memory-gib", - "0", - "--idle-timeout", - "60", - ] - ) + _cli.main(_worker_cli_args(endpoint_file, *rejected)) assert exc_info.value.code == 2 -def test_worker_cli_requires_resource_request( +def _capture_execvp(monkeypatch: pytest.MonkeyPatch) -> list[tuple[str, list[str]]]: + execs: list[tuple[str, list[str]]] = [] + monkeypatch.setattr( + _cli.os, "execvp", lambda file, argv: execs.append((file, list(argv))) + ) + return execs + + +def test_worker_cli_reexecs_when_endpoint_generation_grows( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: - calls: list[ResourceRequest] = [] - token_file = tmp_path / "worker.token" - token_file.write_text("secret") - - def worker_loop( - *, - server_url: str, - auth_token: str, - resource_request: ResourceRequest, - idle_timeout: float | None, - ) -> None: - calls.append(resource_request) - - monkeypatch.setattr(_cli, "worker_loop", worker_loop) - - with pytest.raises(SystemExit) as exc_info: - _cli.main( - [ - "--server-url", - "http://execution-coordinator.test", - "--auth-token-file", - str(token_file), - "--idle-timeout", - "60", - "--component", - "test-worker", - ] + endpoint_file = tmp_path / "endpoint.json" + _write_endpoint(endpoint_file) + + def worker_loop(**kwargs: object) -> str: + # A takeover rewrites the endpoint file while the worker is connected. + _write_endpoint( + endpoint_file, + generation=2, + server_url="ws://successor.test:1", + project_root="/proj/two", + config_file="/cfg/two.json", ) + return "disconnected" - assert exc_info.value.code == 2 - assert calls == [] + monkeypatch.setattr(_cli, "worker_loop", worker_loop) + execs = _capture_execvp(monkeypatch) + monkeypatch.setenv("VIRTUAL_ENV", "/old/venv") + monkeypatch.setenv(_WORKER_JSON_CONFIG_FILE_ENV_VAR, "/cfg/one.json") + args = _worker_cli_args(endpoint_file) + assert _cli.main(args) == 0 -def test_worker_cli_requires_auth_token_file(monkeypatch: pytest.MonkeyPatch) -> None: - calls: list[tuple[str, str]] = [] + ((file, argv),) = execs + assert file == "uv" + assert argv[:8] == [ + "uv", + "run", + "--frozen", + "--project", + "/proj/two", + "python", + "-m", + "furu.worker._cli", + ] + assert argv[8:] == args + assert os.environ[_WORKER_JSON_CONFIG_FILE_ENV_VAR] == "/cfg/two.json" + assert "VIRTUAL_ENV" not in os.environ - def worker_loop( - *, - server_url: str, - auth_token: str, - resource_request: ResourceRequest, - idle_timeout: float | None, - ) -> None: - calls.append((server_url, auth_token)) - monkeypatch.setattr(_cli, "worker_loop", worker_loop) +def test_worker_cli_exits_without_reexec_when_generation_unchanged( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + endpoint_file = tmp_path / "endpoint.json" + _write_endpoint(endpoint_file) - with pytest.raises(SystemExit) as exc_info: - _cli.main( - [ - "--server-url", - "http://execution-coordinator.test", - "--resource-cpus", - "1", - "--resource-gpus", - "0", - "--resource-memory-gib", - "0", - "--idle-timeout", - "60", - "--component", - "test-worker", - ] - ) + monkeypatch.setattr(_cli, "worker_loop", lambda **kwargs: "disconnected") + execs = _capture_execvp(monkeypatch) - assert exc_info.value.code == 2 - assert calls == [] + assert _cli.main(_worker_cli_args(endpoint_file)) == 0 + assert execs == [] -def test_worker_cli_requires_idle_timeout( +def test_worker_cli_reexecs_after_preemption( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: - calls: list[float | None] = [] - token_file = tmp_path / "worker.token" - token_file.write_text("secret") + endpoint_file = tmp_path / "endpoint.json" + _write_endpoint(endpoint_file) - def worker_loop( - *, - server_url: str, - auth_token: str, - resource_request: ResourceRequest, - idle_timeout: float | None, - ) -> None: - calls.append(idle_timeout) + def worker_loop(**kwargs: object) -> str: + _write_endpoint(endpoint_file, generation=2, project_root="/proj/two") + raise WorkerPreempted monkeypatch.setattr(_cli, "worker_loop", worker_loop) + execs = _capture_execvp(monkeypatch) - with pytest.raises(SystemExit) as exc_info: - _cli.main( - [ - "--server-url", - "http://execution-coordinator.test", - "--auth-token-file", - str(token_file), - "--resource-cpus", - "1", - "--resource-gpus", - "0", - "--resource-memory-gib", - "0", - ] - ) + assert _cli.main(_worker_cli_args(endpoint_file)) == 0 - assert exc_info.value.code == 2 - assert calls == [] + ((_, argv),) = execs + assert argv[3:5] == ["--project", "/proj/two"] -def test_worker_cli_rejects_auth_token_argument( +def test_worker_cli_reraises_connect_failure_when_endpoint_unchanged( tmp_path: Path, monkeypatch: pytest.MonkeyPatch, ) -> None: - calls: list[tuple[str, str]] = [] - token_file = tmp_path / "worker.token" - token_file.write_text("secret") + endpoint_file = tmp_path / "endpoint.json" + _write_endpoint(endpoint_file) - def worker_loop( - *, - server_url: str, - auth_token: str, - resource_request: ResourceRequest, - idle_timeout: float | None, - ) -> None: - calls.append((server_url, auth_token)) + def worker_loop(**kwargs: object) -> str: + raise OSError("connection refused") monkeypatch.setattr(_cli, "worker_loop", worker_loop) + execs = _capture_execvp(monkeypatch) - with pytest.raises(SystemExit) as exc_info: - _cli.main( - [ - "--server-url", - "http://execution-coordinator.test", - "--auth-token-file", - str(token_file), - "--resource-cpus", - "1", - "--resource-gpus", - "0", - "--resource-memory-gib", - "0", - "--idle-timeout", - "60", - "--component", - "test-worker", - "--auth-token", - "secret", - ] - ) - - assert exc_info.value.code == 2 - assert calls == [] + with pytest.raises(OSError, match="connection refused"): + _cli.main(_worker_cli_args(endpoint_file)) + assert execs == [] def test_slurm_backend_submits_workers_with_required_sbatch_options( @@ -569,17 +468,19 @@ def test_slurm_backend_submits_workers_with_required_sbatch_options( script_path = Path(argv[-1]) script = script_path.read_text() - assert "--auth-token-file" in script - assert "--auth-token " not in script + assert "--auth-token" not in script + assert "--server-url" not in script assert "secret-token" not in script project_root = EnvironmentIdentity.capture().project_root assert script.index('echo "Hello" > /tmp/hey') < script.index("exec uv run") assert script.index("unset VIRTUAL_ENV") < script.index("exec uv run") - assert f"exec uv run --frozen --project {project_root}" in script + # The project is resolved from the endpoint file at script runtime, so a + # takeover can redirect queued jobs into a new snapshot. + assert 'exec uv run --frozen --project "$(python3 -c' in script assert "python -m furu.worker._cli" in script assert sys.executable not in script assert "--backend slurm" in script - assert "--server-url ws://execution-coordinator.cluster:1234" in script + assert '--endpoint-file "$furu_endpoint_file"' in script assert "SLURM_ARRAY_TASK_ID" in script assert "SLURM_ARRAY_JOB_ID" in script assert ( @@ -595,12 +496,18 @@ def test_slurm_backend_submits_workers_with_required_sbatch_options( assert "FURU_DIRECTORIES__EXECUTIONS" not in script assert not (worker_dir / "secrets").exists() - token_files = sorted(worker_dir.glob("worker-*.token")) - assert len(token_files) == 1 - for token_file in token_files: - assert _mode(token_file) == 0o600 - assert token_file.read_text() == "secret-token" - assert str(token_file) in script + assert sorted(worker_dir.glob("worker-*.token")) == [] + + endpoint_files = sorted(worker_dir.glob("endpoint-*.json")) + assert len(endpoint_files) == 1 + (endpoint_file,) = endpoint_files + assert _mode(endpoint_file) == 0o600 + assert f"furu_endpoint_file={endpoint_file}" in script + endpoint = read_worker_endpoint(endpoint_file) + assert endpoint.generation == 1 + assert endpoint.server_url == "ws://execution-coordinator.cluster:1234" + assert endpoint.auth_token == "secret-token" + assert endpoint.project_root == project_root config_files = sorted(worker_dir.glob("worker-*.config.json")) assert len(config_files) == 1 @@ -610,15 +517,14 @@ def test_slurm_backend_submits_workers_with_required_sbatch_options( _Config.model_validate_json(config_file.read_text(encoding="utf-8")) == get_config() ) - assert f"export {_WORKER_JSON_CONFIG_FILE_ENV_VAR}={config_file}" in script + assert endpoint.config_file == str(config_file) + # The config env var is resolved from the endpoint file at script runtime. + assert f'export {_WORKER_JSON_CONFIG_FILE_ENV_VAR}="$(python3 -c' in script assert not sbatch_records[0]["has_execution_coordinator_environment"] assert "secret-token" not in record_file.read_text() - assert all(token_file.exists() for token_file in token_files) - assert all(config_file.exists() for config_file in config_files) - @pytest.mark.skipif(shutil.which("bash") is None, reason="requires bash") @pytest.mark.parametrize( @@ -722,6 +628,62 @@ def test_slurm_array_worker_component_label_derivation_under_bash( assert result.stdout == "slurm-worker-100a7" +@pytest.mark.skipif( + shutil.which("bash") is None or shutil.which("python3") is None, + reason="requires bash and python3", +) +def test_slurm_worker_script_resolves_endpoint_fields_under_bash( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _disable_slurm_pool_scale_thread(monkeypatch) + backend = SlurmWorkerBackend( + max_workers=1, + resources=SlurmResources(cpus_per_worker=1), + worker_connect_host="execution-coordinator.cluster", + ) + pool = backend.start_pool( + coordinator=_StubCoordinator(), + bound_port=1234, + auth_token="secret-token", + executor_dir=tmp_path / "executor", + provenance=_submit_provenance(), + ) + endpoint = read_worker_endpoint(pool._endpoint_file) + lines = pool._script_path.read_text().splitlines() + endpoint_line = next( + line for line in lines if line.startswith("furu_endpoint_file=") + ) + export_line = next( + line + for line in lines + if line.startswith(f"export {_WORKER_JSON_CONFIG_FILE_ENV_VAR}=") + ) + project_lookup = ( + next(line for line in lines if "--project" in line) + .removeprefix("exec uv run --frozen --project ") + .removesuffix(" \\") + ) + script = ( + "set -euo pipefail\n" + f"{endpoint_line}\n" + f"{export_line}\n" + f'printf "%s\\n" "${_WORKER_JSON_CONFIG_FILE_ENV_VAR}"\n' + f"printf '%s' {project_lookup}\n" + ) + + result = subprocess.run( + ["bash", "-c", script], + capture_output=True, + text=True, + check=True, + ) + + config_file, project_root = result.stdout.split("\n") + assert config_file == endpoint.config_file + assert project_root == endpoint.project_root + + @pytest.mark.parametrize( ("export", "expected_args"), [ @@ -1092,9 +1054,9 @@ def test_slurm_backend_builds_server_url_from_worker_connect_host( sbatch_records = [record for record in records if record["executable"] == "sbatch"] assert len(sbatch_records) == 1 - script_path = Path(sbatch_records[0]["argv"][-1]) - script = script_path.read_text() - assert "--server-url ws://execution-coordinator.cluster:4321" in script + endpoint = read_worker_endpoint(pool._endpoint_file) + assert endpoint.server_url == "ws://execution-coordinator.cluster:4321" + script = Path(sbatch_records[0]["argv"][-1]).read_text() assert f"--idle-timeout {get_config().worker.idle_timeout_seconds}" in script @@ -1125,9 +1087,9 @@ def test_slurm_backend_worker_connect_port_overrides_bound_port( sbatch_records = [record for record in records if record["executable"] == "sbatch"] assert len(sbatch_records) == 1 - script = Path(sbatch_records[0]["argv"][-1]).read_text() - assert "--server-url ws://execution-coordinator.cluster:9000" in script - assert ":4321" not in script + endpoint = read_worker_endpoint(pool._endpoint_file) + assert endpoint.server_url == "ws://execution-coordinator.cluster:9000" + assert ":4321" not in endpoint.server_url def test_slurm_backend_worker_connect_host_defaults_to_config() -> None: @@ -1779,6 +1741,10 @@ def _install_fake_slurm( with open(record_file, "a", encoding="utf-8") as file: file.write(json.dumps({"executable": "scancel", "argv": sys.argv[1:]}) + "\\n") + if any(arg.startswith("--signal") for arg in sys.argv[1:]): + # Signal delivery leaves the job running. + sys.exit(0) + cancelled_jobs = set(sys.argv[1:]) with open(active_file, encoding="utf-8") as file: active_jobs = [ @@ -1876,8 +1842,9 @@ def test_slurm_backend_runs_workers_from_the_extracted_snapshot( ).resolve() assert (code_dir / "pyproject.toml").is_file() assert f"--chdir={code_dir}" in pool._sbatch_base_args + endpoint = read_worker_endpoint(pool._endpoint_file) + assert endpoint.project_root == str(code_dir) script = pool._script_path.read_text() - assert f"--project {shlex.quote(str(code_dir))}" in script assert str(repo) not in script # The venv is built once at submit so workers never race to create it. assert uv_commands == [["uv", "sync", "--frozen", "--project", str(code_dir)]] @@ -1913,3 +1880,129 @@ def test_slurm_backend_pins_relative_data_directories_for_workers( written = _Config.model_validate_json(config_file.read_text(encoding="utf-8")) assert written.directories.objects == work_dir / "furu-data" / "objects" assert written.directories.snapshots == work_dir / "furu-data" / "snapshots" + + +def test_slurm_backend_fingerprint_covers_worker_identity_only() -> None: + def backend(**overrides: Any) -> SlurmWorkerBackend: + kwargs: dict[str, Any] = { + "max_workers": 2, + "resources": SlurmResources(cpus_per_worker=4, partition="debug"), + "worker_connect_host": "execution-coordinator.cluster", + } + kwargs.update(overrides) + return SlurmWorkerBackend(**kwargs) + + base = backend().fingerprint() + # Submitter-side concerns do not change what a worker is. + assert backend(max_workers=9).fingerprint() == base + assert backend(poll_interval=99.0).fingerprint() == base + assert backend(worker_idle_timeout=1.0).fingerprint() == base + assert ( + backend(worker_connect_host="other", worker_connect_port=1).fingerprint() + == base + ) + # Anything reaching sbatch or the worker runtime does. + assert backend(resources=SlurmResources(cpus_per_worker=8)).fingerprint() != base + assert backend(job_name="other").fingerprint() != base + assert backend(use_job_arrays=False).fingerprint() != base + assert backend(export="ALL").fingerprint() != base + assert backend(pre_worker_commands=("module load cuda",)).fingerprint() != base + + +def test_slurm_backend_start_pool_adopts_inherited_jobs( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _disable_slurm_pool_scale_thread(monkeypatch) + backend = SlurmWorkerBackend( + max_workers=4, + resources=SlurmResources(cpus_per_worker=1), + worker_connect_host="execution-coordinator.cluster", + ) + adopted_endpoint = tmp_path / "old-executor" / "workers" / "endpoint-cafe0123.json" + + pool = backend.start_pool( + coordinator=_StubCoordinator(), + bound_port=1234, + auth_token="secret-token", + executor_dir=tmp_path / "executor", + provenance=_submit_provenance(), + adopt=AdoptedPool( + pool_id="cafe0123", + endpoint_file=adopted_endpoint, + job_ids=("100_0", "100_1"), + ), + ) + + assert pool._job_ids == ["100_0", "100_1"] + assert pool.pool_id == "cafe0123" + assert pool._endpoint_file == adopted_endpoint + # The takeover already rewrote the adopted endpoint file; the pool must + # not write a competing one of its own. + assert list((tmp_path / "executor" / "workers").glob("endpoint-*.json")) == [] + script = pool._script_path.read_text() + assert f"furu_endpoint_file={adopted_endpoint}" in script + + +def test_slurm_worker_pool_stop_leaves_jobs_alone_after_surrender( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _disable_slurm_pool_scale_thread(monkeypatch) + record_file, active_file = _install_fake_slurm(tmp_path, monkeypatch) + backend = SlurmWorkerBackend( + max_workers=2, + resources=SlurmResources(cpus_per_worker=1), + worker_connect_host="execution-coordinator.cluster", + poll_interval=0, + ) + pool = backend.start_pool( + coordinator=_StubCoordinator(2), + bound_port=1234, + auth_token="secret-token", + executor_dir=tmp_path / "executor", + provenance=_submit_provenance(), + ) + pool._scale_once() + assert pool._job_ids == ["100_0", "100_1"] + + pool.surrender() + pool.stop(timeout=0) + + assert active_file.read_text() != "" + assert not any( + record["executable"] == "scancel" for record in _read_records(record_file) + ) + + +def test_slurm_worker_pool_reports_takeover_inventory( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + _disable_slurm_pool_scale_thread(monkeypatch) + _record_file, active_file = _install_fake_slurm(tmp_path, monkeypatch) + backend = SlurmWorkerBackend( + max_workers=2, + resources=SlurmResources(cpus_per_worker=1), + worker_connect_host="execution-coordinator.cluster", + poll_interval=0, + ) + pool = backend.start_pool( + coordinator=_StubCoordinator(2), + bound_port=1234, + auth_token="secret-token", + executor_dir=tmp_path / "executor", + provenance=_submit_provenance(), + ) + pool._scale_once() + active_file.write_text("100_0 RUNNING\n100_1 PENDING\n") + + inventory = pool.takeover_inventory() + + assert inventory.pool_id == pool.pool_id + assert inventory.fingerprint == backend.fingerprint() + assert inventory.endpoint_file == pool._endpoint_file + assert [(job.job_id, job.state) for job in inventory.jobs] == [ + ("100_0", "RUNNING"), + ("100_1", "PENDING"), + ] diff --git a/tests/test_takeover.py b/tests/test_takeover.py new file mode 100644 index 00000000..2576fe2f --- /dev/null +++ b/tests/test_takeover.py @@ -0,0 +1,562 @@ +from __future__ import annotations + +import os +import shutil +import stat +import subprocess +import sys +import time +from collections.abc import Iterator +from contextlib import contextmanager +from pathlib import Path + +import pytest +from test_slurm_backend import ( + _disable_slurm_pool_scale_thread, + _install_fake_slurm, + _read_records, + _submit_provenance, +) +from websockets.exceptions import ConnectionClosed +from websockets.headers import build_authorization_basic +from websockets.sync.client import ClientConnection, connect + +import furu +from furu.config import _Config, get_config +from furu.dag import _add_to_dag +from furu.execution.execution_coordinator import ( + ExecutionCoordinator, + FuruReplacedError, +) +from furu.execution.server import ( + ExecutionCoordinatorServer, + execution_coordinator_server, +) +from furu.execution.takeover import ( + TAKEOVER_PATH, + AdoptedPool, + ClaimRequest, + ClaimResponse, + LiveRunEntry, + PoolsRequest, + PoolsResponse, + _discover_live_run, + _live_registry_dir, + perform_takeover, + register_live_run, +) +from furu.provenance import ( + EnvironmentIdentity, + GitIdentity, + SubmitContext, + SubmitProvenance, +) +from furu.testing import override_config +from furu.worker.backends.local import LocalThreadWorkerBackend +from furu.worker.backends.slurm.backend import SlurmWorkerBackend +from furu.worker.backends.slurm.pool import SlurmWorkerPool +from furu.worker.backends.slurm.resources import SlurmResources +from furu.worker.endpoint import ( + WorkerEndpoint, + read_worker_endpoint, + write_worker_endpoint, +) + +OLD_EXECUTOR_ID = "0123456789abcdef" * 2 +NEW_EXECUTOR_ID = "fedcba9876543210" * 2 + + +def _coordinator(executor_id: str) -> ExecutionCoordinator: + coordinator = ExecutionCoordinator(max_retries_per_object=0) + coordinator.executor_id = executor_id + return coordinator + + +@contextmanager +def _loopback_connect_host() -> Iterator[None]: + """The live registry advertises the configured connect host; point it at + the loopback interface these test servers actually bind.""" + data = get_config().model_dump() + data["worker"]["connect_host"] = "127.0.0.1" + with override_config(_Config.model_validate(data)): + yield + + +# --------------------------------------------------------------------------- +# Live-run registry + + +def test_register_live_run_is_discoverable_while_held() -> None: + with register_live_run( + executor_id=OLD_EXECUTOR_ID, + executor_dir=Path("/somewhere/executions") / OLD_EXECUTOR_ID, + bound_port=4567, + auth_token="registry-token", + ): + entry = _discover_live_run("auto", exclude_executor_id=NEW_EXECUTOR_ID) + assert entry.executor_id == OLD_EXECUTOR_ID + assert entry.auth_token == "registry-token" + assert entry.server_url.endswith(":4567") + entry_path = _live_registry_dir() / f"{OLD_EXECUTOR_ID}.json" + assert stat.S_IMODE(entry_path.stat().st_mode) == 0o600 + + # A run never discovers itself. + with pytest.raises(RuntimeError, match="no live run"): + _discover_live_run("auto", exclude_executor_id=OLD_EXECUTOR_ID) + + with pytest.raises(RuntimeError, match="no live run"): + _discover_live_run("auto", exclude_executor_id=NEW_EXECUTOR_ID) + + +def test_discover_garbage_collects_entries_without_active_heartbeat() -> None: + live_dir = _live_registry_dir() + live_dir.mkdir(parents=True, exist_ok=True) + stale_path = live_dir / f"{OLD_EXECUTOR_ID}.json" + stale_path.write_text( + LiveRunEntry( + executor_id=OLD_EXECUTOR_ID, + server_url="ws://gone:1", + auth_token="stale", + executor_dir=Path("/gone"), + pid=1, + host="gone", + started_at="2026-01-01T00:00:00+00:00", + ).model_dump_json() + ) + + with pytest.raises(RuntimeError, match="no live run"): + _discover_live_run("auto", exclude_executor_id=NEW_EXECUTOR_ID) + assert not stale_path.exists() + + +def test_discover_requires_disambiguation_and_accepts_prefix() -> None: + other_id = "aaaa" + OLD_EXECUTOR_ID[4:] + with ( + register_live_run( + executor_id=OLD_EXECUTOR_ID, + executor_dir=Path("/a"), + bound_port=1, + auth_token="a", + ), + register_live_run( + executor_id=other_id, + executor_dir=Path("/b"), + bound_port=2, + auth_token="b", + ), + ): + with pytest.raises(RuntimeError, match="matches several live runs"): + _discover_live_run("auto", exclude_executor_id=NEW_EXECUTOR_ID) + entry = _discover_live_run("aaaa", exclude_executor_id=NEW_EXECUTOR_ID) + assert entry.executor_id == other_id + + +# --------------------------------------------------------------------------- +# Takeover handshake (old-coordinator side) + + +@contextmanager +def _old_run_with_pool( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, + *, + job_lines: str = "100_0 RUNNING\n100_1 PENDING\n", +) -> Iterator[ + tuple[ExecutionCoordinator, ExecutionCoordinatorServer, SlurmWorkerPool, Path] +]: + _disable_slurm_pool_scale_thread(monkeypatch) + record_file, active_file = _install_fake_slurm(tmp_path, monkeypatch) + coordinator = _coordinator(OLD_EXECUTOR_ID) + backend = SlurmWorkerBackend( + max_workers=2, + resources=SlurmResources(cpus_per_worker=1), + worker_connect_host="127.0.0.1", + ) + with ( + _loopback_connect_host(), + execution_coordinator_server( + coordinator, bind_host="127.0.0.1", port=0 + ) as server, + ): + pool = backend.start_pool( + coordinator=coordinator, + bound_port=server.bound_port, + auth_token=server.auth_token, + executor_dir=coordinator.executor_dir, + provenance=_submit_provenance(), + ) + coordinator.pools.append(pool) + pool._job_ids[:] = ["100_0", "100_1"] + active_file.write_text(job_lines) + yield coordinator, server, pool, record_file + + +def _takeover_connection(server: ExecutionCoordinatorServer) -> ClientConnection: + return connect( + f"ws://127.0.0.1:{server.bound_port}{TAKEOVER_PATH}", + additional_headers={ + "Authorization": build_authorization_basic("furu", server.auth_token) + }, + ) + + +def test_takeover_handshake_surrenders_claimed_pools( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + with _old_run_with_pool(tmp_path, monkeypatch) as ( + coordinator, + server, + pool, + record_file, + ): + with _takeover_connection(server) as connection: + connection.send( + PoolsRequest(successor_executor_id=NEW_EXECUTOR_ID).model_dump_json() + ) + inventory = PoolsResponse.model_validate_json(connection.recv(timeout=5)) + assert inventory.executor_id == OLD_EXECUTOR_ID + (pool_inventory,) = inventory.pools + assert pool_inventory.pool_id == pool.pool_id + assert [(job.job_id, job.state) for job in pool_inventory.jobs] == [ + ("100_0", "RUNNING"), + ("100_1", "PENDING"), + ] + + connection.send(ClaimRequest(adopt=[pool.pool_id]).model_dump_json()) + response = ClaimResponse.model_validate_json(connection.recv(timeout=5)) + assert response.cancelled == [] + (adopted,) = response.adopted + assert [job.job_id for job in adopted.jobs] == ["100_0", "100_1"] + + assert coordinator.done.wait(timeout=5) + with pytest.raises(FuruReplacedError, match=NEW_EXECUTOR_ID): + coordinator.raise_for_failure() + + # The ordinary shutdown path must leave the surrendered jobs alone. + pool.stop(timeout=0) + assert not any( + record["executable"] == "scancel" for record in _read_records(record_file) + ) + + +def test_takeover_handshake_rejects_unknown_pool_ids( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + with _old_run_with_pool(tmp_path, monkeypatch) as (coordinator, server, pool, _): + with _takeover_connection(server) as connection: + connection.send( + PoolsRequest(successor_executor_id=NEW_EXECUTOR_ID).model_dump_json() + ) + connection.recv(timeout=5) + connection.send(ClaimRequest(adopt=["deadbeef"]).model_dump_json()) + with pytest.raises(ConnectionClosed): + connection.recv(timeout=5) + + assert not coordinator.done.wait(timeout=0.2) + assert not pool._surrendered.is_set() + + +def test_successor_drop_before_claim_keeps_all_pools( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + with _old_run_with_pool(tmp_path, monkeypatch) as (coordinator, server, pool, _): + with _takeover_connection(server) as connection: + connection.send( + PoolsRequest(successor_executor_id=NEW_EXECUTOR_ID).model_dump_json() + ) + connection.recv(timeout=5) + + assert not coordinator.done.wait(timeout=0.5) + assert not pool._surrendered.is_set() + assert coordinator.replaced_by is None + + +def test_second_takeover_connection_is_rejected_as_busy( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + with ( + _old_run_with_pool(tmp_path, monkeypatch) as (_, server, _, _), + _takeover_connection(server) as first, + ): + first.send( + PoolsRequest(successor_executor_id=NEW_EXECUTOR_ID).model_dump_json() + ) + first.recv(timeout=5) # the handshake is now in progress + + with _takeover_connection(server) as second: + with pytest.raises(ConnectionClosed) as exc_info: + second.recv(timeout=5) + assert exc_info.value.rcvd is not None + assert exc_info.value.rcvd.code == 1013 + + +# --------------------------------------------------------------------------- +# perform_takeover (successor side) + + +def test_perform_takeover_adopts_matching_pool_and_signals_running_jobs( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + with _old_run_with_pool(tmp_path, monkeypatch) as ( + old_coordinator, + _old_server, + old_pool, + record_file, + ): + old_endpoint_file = old_pool._endpoint_file + assert read_worker_endpoint(old_endpoint_file).generation == 1 + + new_coordinator = _coordinator(NEW_EXECUTOR_ID) + new_coordinator.submit_provenance = _submit_provenance() + matching_backend = SlurmWorkerBackend( + max_workers=5, # deliberately different: excluded from the fingerprint + resources=SlurmResources(cpus_per_worker=1), + worker_connect_host="127.0.0.1", + ) + unmatched_backend = SlurmWorkerBackend( + max_workers=1, + resources=SlurmResources(cpus_per_worker=32), + worker_connect_host="127.0.0.1", + ) + + with execution_coordinator_server( + new_coordinator, bind_host="127.0.0.1", port=0 + ) as new_server: + adoptions = perform_takeover( + selector="auto", + coordinator=new_coordinator, + server=new_server, + worker_backends=(unmatched_backend, matching_backend), + ) + + assert adoptions == { + 1: AdoptedPool( + pool_id=old_pool.pool_id, + endpoint_file=old_endpoint_file, + job_ids=("100_0", "100_1"), + ) + } + + # The endpoint file now points at the successor, atomically. + endpoint = read_worker_endpoint(old_endpoint_file) + assert endpoint.generation == 2 + assert endpoint.server_url == f"ws://127.0.0.1:{new_server.bound_port}" + assert endpoint.auth_token == new_server.auth_token + assert Path(endpoint.config_file).is_relative_to( + new_coordinator.executor_dir + ) + + # Only the RUNNING job is signalled; PENDING jobs need nothing. + signal_records = [ + record + for record in _read_records(record_file) + if record["executable"] == "scancel" + ] + assert signal_records == [ + { + "executable": "scancel", + "argv": ["--signal=USR1", "--batch", "100_0"], + } + ] + + assert old_coordinator.done.wait(timeout=5) + with pytest.raises(FuruReplacedError, match=NEW_EXECUTOR_ID): + old_coordinator.raise_for_failure() + + # The successor's pool starts seeded with the inherited jobs and + # keeps submitting new workers against the same endpoint file. + new_pool = matching_backend.start_pool( + coordinator=new_coordinator, + bound_port=new_server.bound_port, + auth_token=new_server.auth_token, + executor_dir=new_coordinator.executor_dir, + provenance=new_coordinator.submit_provenance, + adopt=adoptions[1], + ) + assert new_pool._job_ids == ["100_0", "100_1"] + assert new_pool.pool_id == old_pool.pool_id + assert new_pool._endpoint_file == old_endpoint_file + + +def test_perform_takeover_errors_when_no_pool_matches( + tmp_path: Path, + monkeypatch: pytest.MonkeyPatch, +) -> None: + with _old_run_with_pool(tmp_path, monkeypatch) as (old_coordinator, _, pool, _): + new_coordinator = _coordinator(NEW_EXECUTOR_ID) + new_coordinator.submit_provenance = _submit_provenance() + mismatched_backend = SlurmWorkerBackend( + max_workers=1, + resources=SlurmResources(cpus_per_worker=32), + worker_connect_host="127.0.0.1", + ) + + with ( + execution_coordinator_server( + new_coordinator, bind_host="127.0.0.1", port=0 + ) as new_server, + pytest.raises(RuntimeError, match="no pool of the live run"), + ): + perform_takeover( + selector="auto", + coordinator=new_coordinator, + server=new_server, + worker_backends=(mismatched_backend,), + ) + + # The old run is unaffected: nothing was claimed. + assert not old_coordinator.done.wait(timeout=0.5) + assert not pool._surrendered.is_set() + assert read_worker_endpoint(pool._endpoint_file).generation == 1 + + +# --------------------------------------------------------------------------- +# ExecutionCoordinator.run wiring + + +class _TakeoverProbeSpec(furu.Spec): + value: int = 0 + + def create(self) -> int: + return self.value + + +def test_run_errors_when_replace_is_requested_but_nothing_is_live( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("FURU_REPLACE_ORCHESTRATOR", "auto") + + with pytest.raises(RuntimeError, match="no live run to replace"): + ExecutionCoordinator.run( + [_TakeoverProbeSpec(value=1)], + worker_backends=(LocalThreadWorkerBackend(),), + ) + + +def test_run_treats_empty_replace_selector_as_unset( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setenv("FURU_REPLACE_ORCHESTRATOR", "") + + (obj,) = ExecutionCoordinator.run( + [_TakeoverProbeSpec(value=2)], + worker_backends=(LocalThreadWorkerBackend(),), + ) + assert furu.create(obj) == 2 + + +# --------------------------------------------------------------------------- +# End to end: a real worker process redirects to a successor coordinator + + +@pytest.mark.skipif(shutil.which("uv") is None, reason="requires uv") +def test_worker_process_redirects_to_successor_and_completes_work( + tmp_path: Path, +) -> None: + """An idle worker of the old run must, once the endpoint file is rewritten + and the old server goes away, re-exec and complete the new run's work.""" + config = get_config() + config = config.model_copy(update={"directories": config.directories.anchored()}) + config_file = tmp_path / "worker.config.json" + config_file.write_text(config.model_dump_json(indent=2)) + project_root = EnvironmentIdentity.capture().project_root + + provenance = SubmitProvenance( + git=GitIdentity( + commit="0" * 40, + branch=None, + remote=None, + repo_root=".", + dirty=False, + diff_stats=None, + ), + # Real environment identity so the worker-side lock-hash check passes. + environment=EnvironmentIdentity.capture(), + snapshot_id=None, + submitted=SubmitContext.capture(), + ) + + old_coordinator = _coordinator(OLD_EXECUTOR_ID) + obj = _TakeoverProbeSpec(value=7) + new_coordinator = _coordinator(NEW_EXECUTOR_ID) + new_coordinator.submit_provenance = provenance + _add_to_dag(new_coordinator, [obj]) + + worker_env = dict(os.environ) + worker_env["_FURU_WORKER_JSON_CONFIG_FILE"] = str(config_file) + # The worker resolves this test module's Spec class by qualified name. + worker_env["PYTHONPATH"] = os.pathsep.join( + [str(Path(__file__).resolve().parent), worker_env.get("PYTHONPATH", "")] + ) + worker_env.pop("VIRTUAL_ENV", None) + + endpoint_file = tmp_path / "endpoint.json" + worker: subprocess.Popen[bytes] | None = None + try: + with execution_coordinator_server( + new_coordinator, bind_host="127.0.0.1", port=0 + ) as new_server: + with execution_coordinator_server( + old_coordinator, bind_host="127.0.0.1", port=0 + ) as old_server: + write_worker_endpoint( + endpoint_file, + WorkerEndpoint( + generation=1, + server_url=f"ws://127.0.0.1:{old_server.bound_port}", + auth_token=old_server.auth_token, + project_root=project_root, + config_file=str(config_file), + ), + ) + worker = subprocess.Popen( + [ + sys.executable, + "-m", + "furu.worker._cli", + "--endpoint-file", + str(endpoint_file), + "--resource-cpus", + "1", + "--resource-gpus", + "0", + "--resource-memory-gib", + "0", + "--idle-timeout", + "120", + "--component", + "e2e-worker", + "--backend", + "slurm", + ], + env=worker_env, + ) + # Give the worker a moment to read generation 1 and connect; + # if it is slower than this it starts against the rewritten + # endpoint instead — the (equally valid) queued-job path. + time.sleep(2.0) + write_worker_endpoint( + endpoint_file, + WorkerEndpoint( + generation=2, + server_url=f"ws://127.0.0.1:{new_server.bound_port}", + auth_token=new_server.auth_token, + project_root=project_root, + config_file=str(config_file), + ), + ) + # The old server is gone; the worker re-reads the endpoint file, + # re-execs, connects to the successor, and executes its job. + assert new_coordinator.done.wait(timeout=120) + new_coordinator.raise_for_failure() + assert worker.wait(timeout=60) == 0 + finally: + if worker is not None and worker.poll() is None: + worker.kill() + + assert furu.load_existing([obj]) == [7]