diff --git a/flask-server/.env.example b/flask-server/.env.example index 02fa6818..2ef35879 100644 --- a/flask-server/.env.example +++ b/flask-server/.env.example @@ -63,6 +63,11 @@ CANCERVERSE_LOWRES_PATH=/home/visitor/cancerverse_lowres # GPU_WORKER_MAX_CPU_LOAD=0.75 # GPU_WORKER_COOLDOWN_SECONDS=600 # GPU_WORKER_SSH_MULTIPLEX=true +# Jobs allowed to run at once, one per worker (capped by the number of workers; +# 1 = strictly one at a time). A job waits up to MAX_WAIT for a busy worker +# before it may use this host's GPU. +# GPU_WORKER_PARALLEL=1 +# GPU_WORKER_MAX_WAIT_SECONDS=300 # Keep BLAS/OpenMP from multiplying CPU threads inside each model subprocess. # Set model-specific values only after measuring the host; these are safe # process-wide defaults for a shared web server. diff --git a/flask-server/api/api_blueprint.py b/flask-server/api/api_blueprint.py index e7773b93..9eab280e 100644 --- a/flask-server/api/api_blueprint.py +++ b/flask-server/api/api_blueprint.py @@ -1,7 +1,7 @@ from flask import Blueprint, send_file, make_response, request, jsonify, Response, current_app, stream_with_context from werkzeug.utils import secure_filename from services.session_manager import SessionManager, generate_uuid -from services.auto_segmentor import run_auto_segmentation, cancel_session, cancel_all_inference +from services.auto_segmentor import run_auto_segmentation, cancel_session, cancel_all_inference, max_parallel_jobs from services.mesh_generation import ( bake_case_meshes, generate_mesh_manifest, @@ -2166,12 +2166,12 @@ def download_segmentation_zip(id): # services.auto_segmentor serializes the actual model execution, but an # unbounded number of accepted requests would still create one blocked Python # thread and retain one uploaded volume per request. Keep the cap configurable -# so it can be sized to the host's RAM and GPU; the default includes the job -# currently running plus a small waiting queue. +# so it can be sized to the host's RAM and GPU; the default is the jobs that can +# run at once plus a small waiting queue (4 with one job at a time, as always). try: - _INFERENCE_MAX_PENDING = max(1, int(os.getenv("INFERENCE_MAX_PENDING", "4"))) + _INFERENCE_MAX_PENDING = max(1, int(os.getenv("INFERENCE_MAX_PENDING", str(3 + max_parallel_jobs())))) except (TypeError, ValueError): - _INFERENCE_MAX_PENDING = 4 + _INFERENCE_MAX_PENDING = 3 + max_parallel_jobs() _INFERENCE_PENDING_SLOTS = threading.BoundedSemaphore(_INFERENCE_MAX_PENDING) diff --git a/flask-server/deploy/GPU_WORKERS.md b/flask-server/deploy/GPU_WORKERS.md index 3424c9f0..a2f6c6c8 100644 --- a/flask-server/deploy/GPU_WORKERS.md +++ b/flask-server/deploy/GPU_WORKERS.md @@ -68,8 +68,8 @@ download over the 100 Gbit/s LAN, typically 1-2 s. Transfers abort after 120 s without progress (hard cap 15 min each). Model run time is the same as on bdmap1 (same GB10 hardware). -Jobs still run one at a time (the same global lock as local runs), so -throughput is unchanged; the gain is that bdmap1's GPU stays free. +Jobs run one at a time unless `GPU_WORKER_PARALLEL` is raised (see Parallel +jobs); the main gain is that bdmap1's GPU stays free. If a worker runs its own warm predictor on the same port as `EPAI_WARM_URL` / `LESIONSEG_WARM_URL`, the forwarded URL (127.0.0.1) makes jobs on that worker @@ -77,6 +77,48 @@ use it automatically; otherwise they use the cold path, as bdmap1 does when its warm predictor is not running. Starting warm predictors on the workers makes jobs faster than today. +## Assumptions + +- **Each worker has one GPU** (GB10), and jobs are placed per worker, not per + GPU. Model commands carry the web host's `CUDA_VISIBLE_DEVICES` choice + (default 0) to the worker unchanged. A multi-GPU worker would need per-GPU + placement first. +- **One gunicorn process.** Worker reservations live in that process's memory, + which matches the documented deployment (one process, eight threads). +- **No strict first-come-first-served order** between waiting jobs, as with the + single lock before. + +## Parallel jobs + +By default one model job runs at a time, exactly as before. Set +`GPU_WORKER_PARALLEL=N` (in `.env`, with the workers enabled) to let up to N jobs +run at once, one per worker. It is capped at the number of workers in +`GPU_WORKER_HOSTS`, and ignored (1) while workers are disabled. Recommended: the +number of dependable workers (for example 2 with bdmap2 and bdmap4), leaving +the others, such as bdmap3, as spares that take over when one is down. + +How it stays safe: +- **No double booking.** Workers are health-checked without being reserved, and + the winner is reserved atomically, so two jobs never get the same worker and + one job never sees "no worker free" just because another is checking. +- **Waiting beats overflowing.** If every usable worker is busy with our own + jobs, a job waits (up to `GPU_WORKER_MAX_WAIT_SECONDS`, 300; cancel works while + waiting) instead of running on bdmap1. bdmap1 is used only when no worker is + up at all, or after that wait. Jobs beyond N wait in the normal queue, and + the queue size (`INFERENCE_MAX_PENDING`) defaults to `3 + N`. +- **bdmap1's own GPU still runs at most one model at a time**, however many + jobs are in flight: fallback jobs queue on a lock of their own. +- **A session never runs against itself.** A repeat request for a session whose + job is still queued or running (a retry, a double click) waits for it, as it + did behind the old single lock, and does not use up a job slot while waiting. +- **The older single-host ePAI mode (`EPAI_REMOTE_ENABLED`) keeps jobs strictly + serial**, because it sends every ePAI job to one fixed GPU outside the pool. +- Placement is still by health and load, so a hot, throttled or busy worker is + skipped, and a failing one cools down. + +To roll back to one job at a time, unset `GPU_WORKER_PARALLEL` (or set it to 1) +and reload. + ## One-time setup (per worker) 1. Passwordless SSH from bdmap1 as `visitor`: @@ -107,6 +149,12 @@ Restart gunicorn with the usual deploy procedure. To roll back, set ## Maintenance +- **Reload or restart gunicorn only when no job is running.** Worker reservations + and the web host GPU lock live in the process; a new process starts with none. + (The health check still keeps it off a worker that is busy, but a job in + flight is lost when its process exits anyway.) Check that no model process + runs on bdmap1 and no worker has a session folder or GPU process first. + - After installing or updating a model, conda env or weights on bdmap1, run `flask-server/scripts/sync_gpu_workers.sh` (all hosts from `.env`). `flask-server/scripts` itself is synced automatically with every job. diff --git a/flask-server/services/auto_segmentor.py b/flask-server/services/auto_segmentor.py index 2e1e4920..2e950c90 100644 --- a/flask-server/services/auto_segmentor.py +++ b/flask-server/services/auto_segmentor.py @@ -1,3 +1,4 @@ +import contextlib import os import uuid import signal @@ -14,8 +15,89 @@ # Load environment variables load_dotenv() -# Only one model inference runs at a time to avoid GPU OOM -_gpu_lock = threading.Lock() +def _is_truthy(value) -> bool: + return gpu_workers._truthy(value) + + +def max_parallel_jobs() -> int: + """How many model jobs may run at once. + + 1 (strictly one at a time, as before) unless GPU workers are enabled and + GPU_WORKER_PARALLEL asks for more; never more than there are workers, since + each worker runs one job at a time, and 1 while the older EPAI_REMOTE_ENABLED + mode is on. Read once at import. + """ + if not gpu_workers.enabled(): + return 1 + if _is_truthy(os.getenv("EPAI_REMOTE_ENABLED", "false")): + # The older explicit single-host ePAI mode sends every ePAI job to one + # fixed GPU outside the worker pool; only serial jobs are safe there. + return 1 + try: + wanted = int(os.getenv("GPU_WORKER_PARALLEL", "1")) + except ValueError: + wanted = 1 + return max(1, min(wanted, len(gpu_workers.hosts()))) + + +# A job holds one of these slots for its whole run, so concurrent requests queue +# instead of OOM-ing a GPU. One slot = the old single global lock. +_job_slots = threading.BoundedSemaphore(max_parallel_jobs()) + +# bdmap1's own GPU runs at most one model at a time, however many jobs are in +# flight: it is the machine that serves the website (shared CPU/GPU memory), so +# fallback jobs queue here rather than pile onto it. +_local_gpu_lock = threading.Lock() + + +_session_locks = {} # session id -> [lock, waiters+holder count] + + +@contextlib.contextmanager +def _session_exclusive(session_id): + """A session never runs against itself. + + A repeat request (a retry, a double click) for a session whose job is still + queued or running used to serialize behind it on the single global lock. + With several jobs in flight it would overlap it, two runs writing the same + workspace, outputs, process entry and job state. It waits here instead, and + it does so before taking a job slot, so a waiting duplicate costs no capacity. + """ + if not session_id: + yield + return + with _session_procs_lock: + entry = _session_locks.setdefault(session_id, [threading.Lock(), 0]) + entry[1] += 1 + try: + with entry[0]: + yield + finally: + with _session_procs_lock: + entry[1] -= 1 + if entry[1] == 0: + _session_locks.pop(session_id, None) + + +def _current_session_cancelled() -> bool: + sid = getattr(_thread_session, "sid", None) + return bool(sid) and sid in _cancelled_sessions + + +@contextlib.contextmanager +def _local_gpu_slot(cancelled): + """Hold the web host's GPU for one local model run; cancellable while waiting.""" + while not _local_gpu_lock.acquire(timeout=1.0): + if cancelled(): + raise RuntimeError("Inference cancelled") + try: + # Also covers a cancel that landed while waiting but just as the lock + # came free (no process existed to kill), and a lock that was free. + if cancelled(): + raise RuntimeError("Inference cancelled") + yield + finally: + _local_gpu_lock.release() # ── Per-session subprocess tracking (for user-initiated cancel) ── # Each session's worker thread binds its session id thread-locally; _tracked_run @@ -118,17 +200,23 @@ def _cancelled(): if _session_procs.get(sid) is registered[-1]: _session_procs.pop(sid, None) - proc = subprocess.Popen(cmd, **kwargs) - if sid: - with _session_procs_lock: - _session_procs[sid] = proc - try: - stdout, stderr = proc.communicate() - finally: + is_model_cmd = bool(remote_dir) and bool(kwargs.get("shell")) and isinstance(cmd, str) + if is_model_cmd and gpu_workers.enabled(): + local_slot = _local_gpu_slot(_current_session_cancelled) + else: + local_slot = contextlib.nullcontext() + with local_slot: + proc = subprocess.Popen(cmd, **kwargs) if sid: with _session_procs_lock: - if _session_procs.get(sid) is proc: - _session_procs.pop(sid, None) + _session_procs[sid] = proc + try: + stdout, stderr = proc.communicate() + finally: + if sid: + with _session_procs_lock: + if _session_procs.get(sid) is proc: + _session_procs.pop(sid, None) if retry_of and proc.returncode == 0: # Local succeeded where the worker failed: the worker is at fault. gpu_workers.mark_failed(retry_of) @@ -290,23 +378,24 @@ def _resolve_conda_activate_path(): def run_auto_segmentation(input_path, session_dir, model, session_id=None, on_start=None): """Run one model; see _run_auto_segmentation. Always clears per-job state.""" token = object() - try: - return _run_auto_segmentation(input_path, session_dir, model, session_id, on_start, token) - finally: - _thread_session.remote_session_dir = None - if session_id: - with _session_procs_lock: - # Runs after _gpu_lock is released, so a newer run of the same - # session may already be active; only clear our own state. - if _active_sessions.get(session_id) is token: - del _active_sessions[session_id] - _cancelled_sessions.discard(session_id) + with _session_exclusive(session_id): + try: + return _run_auto_segmentation(input_path, session_dir, model, session_id, on_start, token) + finally: + _thread_session.remote_session_dir = None + if session_id: + with _session_procs_lock: + # Only clear our own state. + if _active_sessions.get(session_id) is token: + del _active_sessions[session_id] + _cancelled_sessions.discard(session_id) def _run_auto_segmentation(input_path, session_dir, model, session_id=None, on_start=None, token=None): """ Dispatch to the appropriate model inference function. - Serialized via _gpu_lock so concurrent requests queue instead of OOM-ing. + Limited by _job_slots (one at a time unless GPU workers allow more), so + concurrent requests queue instead of OOM-ing. Returns the output directory path on success, raises on failure. session_id: binds this worker thread so _tracked_run/cancel_session can @@ -315,7 +404,7 @@ def _run_auto_segmentation(input_path, session_dir, model, session_id=None, on_s run (used when the user cancelled while the job was still queued) and makes this function return None. """ - with _gpu_lock: + with _job_slots: if session_id: # Before on_start: a cancel landing after its status check is # still recorded. A flag left from an earlier run is dropped. @@ -1199,10 +1288,6 @@ def _run_lesionsegmenter_inference(input_path: str, session_dir: str, conda_path return output_ct_dir -def _is_truthy(value: str) -> bool: - return str(value or "").strip().lower() in {"1", "true", "yes", "y", "on"} - - def _run_checked_process(cmd: list[str], error_prefix: str): process = _tracked_run(cmd, text=True, capture_output=True) if process.returncode != 0: @@ -1517,7 +1602,13 @@ def _run_media_agentic_inference( print(f"[INFO] Running MedIA-Agentic {model_type} inference...") print(f"[INFO] Command: {' '.join(cmd)}") - result = subprocess.run(cmd, capture_output=True, text=True) + # This model runs on the web host's own GPU (it is not routed to a worker), + # so with several jobs in flight it queues on the same lock as every other + # local model run. Only taken when workers are enabled (the disabled path is + # unchanged: one job at a time already). + gpu_slot = _local_gpu_slot(_current_session_cancelled) if gpu_workers.enabled() else contextlib.nullcontext() + with gpu_slot: + result = subprocess.run(cmd, capture_output=True, text=True) if result.returncode != 0: print(f"[ERROR] MedIA-Agentic inference failed: {result.stderr}") diff --git a/flask-server/services/gpu_workers.py b/flask-server/services/gpu_workers.py index 6b1f3147..29ba8478 100644 --- a/flask-server/services/gpu_workers.py +++ b/flask-server/services/gpu_workers.py @@ -33,6 +33,11 @@ SSH connections are reused per worker, so the per-job overhead is about the time to copy the session's files over the LAN. +Several jobs can run at once (GPU_WORKER_PARALLEL, default 1), one per worker: +workers are reserved atomically after their health check, a job waits for a +busy worker rather than overflowing to the web host, and the web host's own GPU +still runs at most one model at a time (see auto_segmentor._local_gpu_slot). + Disabled unless GPU_WORKERS_ENABLED is true, so deploying this changes nothing. """ from __future__ import annotations @@ -109,7 +114,9 @@ class RemoteResultLost(RemoteRunFailed): def _truthy(value) -> bool: - return str(value or "").strip().lower() in ("1", "true", "yes", "on") + """The one yes/no parser for environment flags (auto_segmentor uses it too), + so a value like "y" can never mean yes to one setting and no to another.""" + return str(value or "").strip().lower() in ("1", "true", "yes", "y", "on") def enabled() -> bool: @@ -145,6 +152,10 @@ def _mux_opts(host: str) -> list[str]: return ["-o", "ControlMaster=no", "-o", f"ControlPath={_control_path(host)}"] +_master_locks: dict[str, threading.Lock] = {} +_master_locks_guard = threading.Lock() + + def _ensure_master(host: str) -> None: """Start the shared SSH connection for a worker if it is not running. @@ -154,6 +165,13 @@ def _ensure_master(host: str) -> None: """ if not _truthy(os.getenv("GPU_WORKER_SSH_MULTIPLEX", "true")): return + with _master_locks_guard: + lock = _master_locks.setdefault(host, threading.Lock()) + with lock: # two jobs must not both create (or remove) the same socket + _ensure_master_locked(host) + + +def _ensure_master_locked(host: str) -> None: target = f"{_user()}@{host}" path = _control_path(host) check = subprocess.run(["ssh", "-o", f"ControlPath={path}", "-O", "check", target], @@ -199,6 +217,7 @@ def _run(cmd, timeout=None, **kwargs) -> subprocess.CompletedProcess: # ── Host selection ────────────────────────────────────────────────────────── _busy_lock = threading.Lock() +_busy_cv = threading.Condition(_busy_lock) # notified whenever a worker is released _busy_hosts: set[str] = set() @@ -321,26 +340,19 @@ def mark_failed(host: str, secs: float | None = None) -> None: _cooldown_until[host] = max(_cooldown_until.get(host, 0), time.time() + secs) -def acquire_host(exclude=()) -> str | None: - """Reserve the best healthy worker. Caller must release. +def _max_wait_seconds() -> float: + return _env_float("GPU_WORKER_MAX_WAIT_SECONDS", 300) - Every candidate is health-checked in parallel (a dead host costs at most - the SSH connect timeout, not one timeout per host), then the least loaded - is chosen; ties keep the configured order. - """ - now = time.time() - with _busy_lock: - candidates = [h for h in hosts() if h not in exclude and h not in _busy_hosts - and _cooldown_until.get(h, 0) <= now] - _busy_hosts.update(candidates) - if not candidates: - return None + +def _score_all(candidates) -> dict: + """Health-check candidates in parallel (a dead host costs at most the SSH + connect timeout, not one timeout per host). host -> score or None.""" scores = {} def check(h): try: scores[h] = _host_score(h) - except Exception as e: # never leave a host reserved forever + except Exception as e: print(f"[gpu_workers] health check on {h} failed: {e}") scores[h] = None @@ -349,17 +361,74 @@ def check(h): t.start() for t in threads: t.join(timeout=30) - healthy = [(scores[h], i, h) for i, h in enumerate(candidates) if scores.get(h) is not None] - best = min(healthy)[2] if healthy else None - for h in candidates: - if h != best: - release_host(h) - return best + return scores + + +def acquire_host(exclude=(), cancelled=None) -> str | None: + """Reserve the best healthy worker nobody holds. Caller must release. + + Safe with several jobs at once: + * health checks run on the unreserved workers WITHOUT reserving them, so + a concurrent job does not see them as taken while they are only being + checked; the winner is reserved atomically afterwards, and a job that + loses that race simply looks again; + * the least loaded healthy worker wins, ties keep the configured order; + * the winner is committed only if, under the lock, it is still unreserved + and not in cooldown (its score may be a moment stale); + * if no worker is usable right now but some are only busy with our own + jobs, wait for one to be released (bounded by GPU_WORKER_MAX_WAIT_SECONDS, + cancellable) instead of overflowing to the web host. Nothing to wait + for (all down, cooling down or excluded) returns None at once. + """ + deadline = time.monotonic() + _max_wait_seconds() + announced = False + while True: + if cancelled is not None and cancelled(): + raise WorkerCancelled("cancelled while waiting for a GPU worker") + now = time.time() + with _busy_lock: + eligible = [h for h in hosts() if h not in exclude and _cooldown_until.get(h, 0) <= now] + free = [h for h in eligible if h not in _busy_hosts] + holding = [h for h in eligible if h in _busy_hosts] + if free: + scores = _score_all(free) + ranked = sorted((scores[h], i, h) for i, h in enumerate(free) if scores.get(h) is not None) + for _, _, h in ranked: + with _busy_lock: + # Re-check everything the choice depended on, under the lock + # mark_failed uses: another job may have reserved this worker, + # failed on it and put it in cooldown since we scored it. + if h not in _busy_hosts and _cooldown_until.get(h, 0) <= time.time(): + _busy_hosts.add(h) + return h + if ranked: + continue # healthy workers were all taken meanwhile: look again + with _busy_lock: + # Reservations can change while we were health-checking: another job + # may have taken or released a worker. Reconsider newly free workers + # before waiting or falling back; their health has not been checked. + now = time.time() + eligible = [h for h in hosts() if h not in exclude and _cooldown_until.get(h, 0) <= now] + if any(h not in _busy_hosts and h not in free for h in eligible): + continue + holding = [h for h in eligible if h in _busy_hosts] + if not holding: + return None # nobody will free up: the caller falls back + remaining = deadline - time.monotonic() + if remaining <= 0: + return None + if not announced: + print(f"[gpu_workers] all usable workers are busy with other jobs; waiting up to {remaining:.0f}s") + announced = True + with _busy_cv: + if not any(h not in _busy_hosts for h in holding): # none released since we looked + _busy_cv.wait(timeout=min(remaining, 5.0)) def release_host(host: str) -> None: - with _busy_lock: + with _busy_cv: _busy_hosts.discard(host) + _busy_cv.notify_all() # ── Remote execution ──────────────────────────────────────────────────────── @@ -679,7 +748,7 @@ def run(cmd: str, session_dir: str, cwd: str | None, popen, while True: if cancelled(): raise WorkerCancelled("cancelled before dispatch") - host = acquire_host(exclude=tried) + host = acquire_host(exclude=tried, cancelled=cancelled) if host is None: raise WorkerUnavailable(f"no idle GPU worker (tried: {tried or 'none idle'})") tried.append(host) diff --git a/flask-server/tests/unit/test_gpu_workers.py b/flask-server/tests/unit/test_gpu_workers.py index a210f43b..d16e54db 100644 --- a/flask-server/tests/unit/test_gpu_workers.py +++ b/flask-server/tests/unit/test_gpu_workers.py @@ -553,3 +553,497 @@ def test_finished_run_does_not_clear_a_newer_runs_cancel(monkeypatch): finally: auto_segmentor._active_sessions.pop("sid-race", None) auto_segmentor._cancelled_sessions.discard("sid-race") + + +# ── Parallel jobs ─────────────────────────────────────────────────────────── +import threading +import time + + +def _score_everyone_equal(monkeypatch): + monkeypatch.setattr(gpu_workers, "_host_score", lambda h: (0.0, 0.0, 0)) + + +def test_concurrent_jobs_never_get_the_same_worker(monkeypatch): + # 8 threads fight over 3 workers, 40 rounds each: no worker is ever held by + # two jobs at once, and every job eventually gets one. + monkeypatch.setenv("GPU_WORKER_HOSTS", "w1,w2,w3") + monkeypatch.setenv("GPU_WORKER_MAX_WAIT_SECONDS", "30") + _score_everyone_equal(monkeypatch) + holders, problems, done = {}, [], [] + guard = threading.Lock() + + def job(): + for _ in range(40): + h = gpu_workers.acquire_host() + if h is None: + problems.append("got None while workers exist") + return + with guard: + if h in holders: + problems.append(f"{h} held twice") + holders[h] = True + time.sleep(0.001) + with guard: + del holders[h] + gpu_workers.release_host(h) + done.append(1) + + threads = [threading.Thread(target=job) for _ in range(8)] + [t.start() for t in threads] + [t.join(60) for t in threads] + assert problems == [] and len(done) == 8 and not gpu_workers._busy_hosts + + +def test_two_simultaneous_jobs_use_two_workers(monkeypatch): + monkeypatch.setenv("GPU_WORKER_HOSTS", "w1,w2") + _score_everyone_equal(monkeypatch) + got, barrier = [], threading.Barrier(2) + + def job(): + barrier.wait() + got.append(gpu_workers.acquire_host()) + + threads = [threading.Thread(target=job) for _ in range(2)] + [t.start() for t in threads] + [t.join(10) for t in threads] + assert sorted(got) == ["w1", "w2"] # neither saw "no free worker" and fell back + for h in got: + gpu_workers.release_host(h) + + +def test_job_waits_for_a_busy_worker_instead_of_overflowing(monkeypatch): + monkeypatch.setenv("GPU_WORKER_HOSTS", "w1") + monkeypatch.setenv("GPU_WORKER_MAX_WAIT_SECONDS", "30") + _score_everyone_equal(monkeypatch) + gpu_workers._busy_hosts.add("w1") # another job holds it + got = [] + t = threading.Thread(target=lambda: got.append(gpu_workers.acquire_host())) + t.start() + time.sleep(0.3) + assert got == [] # still waiting, not overflowed to the web host + gpu_workers.release_host("w1") + t.join(10) + assert got == ["w1"] + gpu_workers.release_host("w1") + + +def test_wait_is_bounded(monkeypatch): + monkeypatch.setenv("GPU_WORKER_HOSTS", "w1") + monkeypatch.setenv("GPU_WORKER_MAX_WAIT_SECONDS", "0.4") + _score_everyone_equal(monkeypatch) + gpu_workers._busy_hosts.add("w1") + start = time.monotonic() + assert gpu_workers.acquire_host() is None # then the caller falls back locally + assert 0.3 < time.monotonic() - start < 3 + + +def test_nothing_to_wait_for_returns_at_once(monkeypatch): + monkeypatch.setenv("GPU_WORKER_HOSTS", "w1,w2") + monkeypatch.setenv("GPU_WORKER_MAX_WAIT_SECONDS", "60") + monkeypatch.setattr(gpu_workers, "_host_score", lambda h: None) # every worker down + start = time.monotonic() + assert gpu_workers.acquire_host() is None + assert time.monotonic() - start < 2 + + +def test_wait_for_a_worker_is_cancellable(monkeypatch): + monkeypatch.setenv("GPU_WORKER_HOSTS", "w1") + monkeypatch.setenv("GPU_WORKER_MAX_WAIT_SECONDS", "60") + _score_everyone_equal(monkeypatch) + gpu_workers._busy_hosts.add("w1") + flag, errors = [], [] + + def job(): + try: + gpu_workers.acquire_host(cancelled=lambda: bool(flag)) + except gpu_workers.WorkerCancelled as e: + errors.append(e) + + t = threading.Thread(target=job) + t.start() + time.sleep(0.2) + flag.append(True) + gpu_workers.release_host("w1") # wakes the waiter, which sees the cancel + t.join(10) + assert len(errors) == 1 + + +def test_unhealthy_free_worker_waits_for_the_busy_healthy_one(monkeypatch): + monkeypatch.setenv("GPU_WORKER_HOSTS", "w1,w2") + monkeypatch.setenv("GPU_WORKER_MAX_WAIT_SECONDS", "30") + monkeypatch.setattr(gpu_workers, "_host_score", lambda h: (0.0, 0.0, 0) if h == "w1" else None) + gpu_workers._busy_hosts.add("w1") # healthy but busy; w2 is down + got = [] + t = threading.Thread(target=lambda: got.append(gpu_workers.acquire_host())) + t.start() + time.sleep(0.3) + assert got == [] + gpu_workers.release_host("w1") + t.join(10) + assert got == ["w1"] + gpu_workers.release_host("w1") + + +@pytest.mark.parametrize("other_busy", [False, True]) +def test_worker_released_during_health_checks_is_reconsidered(monkeypatch, other_busy): + monkeypatch.setenv("GPU_WORKER_HOSTS", "w1,w2,w3" if other_busy else "w1,w2") + monkeypatch.setenv("GPU_WORKER_MAX_WAIT_SECONDS", "0") + gpu_workers._busy_hosts.add("w1") + if other_busy: + gpu_workers._busy_hosts.add("w3") + checked = [] + + def score(host): + checked.append(host) + if host == "w2": + # A healthy worker finishes while the only free spare is checked. + gpu_workers.release_host("w1") + return None + return (0.0, 0.0, 0) + + monkeypatch.setattr(gpu_workers, "_host_score", score) + try: + assert gpu_workers.acquire_host() == "w1" + assert "w1" in checked + assert "w3" not in checked + finally: + gpu_workers.release_host("w1") + gpu_workers.release_host("w3") + + +def test_max_parallel_jobs_is_one_unless_asked_and_capped_by_workers(monkeypatch): + monkeypatch.setenv("GPU_WORKER_PARALLEL", "3") + assert auto_segmentor.max_parallel_jobs() == 1 # workers disabled: strictly serial + _enable(monkeypatch, "w1,w2") + assert auto_segmentor.max_parallel_jobs() == 2 # capped by the number of workers + monkeypatch.setenv("GPU_WORKER_PARALLEL", "1") + assert auto_segmentor.max_parallel_jobs() == 1 + monkeypatch.delenv("GPU_WORKER_PARALLEL") + assert auto_segmentor.max_parallel_jobs() == 1 # default: no change on deploy + for bad in ("abc", "0", "-4", ""): + monkeypatch.setenv("GPU_WORKER_PARALLEL", bad) + assert auto_segmentor.max_parallel_jobs() == 1 + + +def _run_jobs_blocked(monkeypatch, slots, jobs): + """Start `jobs` model runs that block inside the model; return the peak + number inside at once.""" + monkeypatch.setattr(auto_segmentor, "_job_slots", threading.BoundedSemaphore(slots)) + inside, peak, release = [], [], threading.Event() + lock = threading.Lock() + + def fake_epai(**kw): + with lock: + inside.append(1) + peak.append(len(inside)) + release.wait(10) + with lock: + inside.pop() + return "out" + + monkeypatch.setattr(auto_segmentor, "_run_epai_inference", fake_epai) + monkeypatch.setattr(auto_segmentor, "_resolve_conda_activate_path", lambda: "") + threads = [threading.Thread(target=auto_segmentor.run_auto_segmentation, + args=("in", "/s", "ePAI"), kwargs={"session_id": f"sess-{i}"}) + for i in range(jobs)] + [t.start() for t in threads] + time.sleep(0.4) + seen = max(peak) if peak else 0 + release.set() + [t.join(10) for t in threads] + return seen + + +def test_job_slots_bound_how_many_models_run_at_once(monkeypatch): + assert _run_jobs_blocked(monkeypatch, slots=1, jobs=3) == 1 # exactly the old behavior + assert _run_jobs_blocked(monkeypatch, slots=2, jobs=3) == 2 # the third queues + + +def _local_fallback_setup(monkeypatch, popen): + _enable(monkeypatch) + monkeypatch.setattr(gpu_workers, "run", lambda *a, **k: (_ for _ in ()).throw(gpu_workers.WorkerUnavailable("down"))) + monkeypatch.setattr(auto_segmentor.subprocess, "Popen", popen) + + +def test_web_host_gpu_runs_one_model_at_a_time_even_with_parallel_jobs(monkeypatch): + active, peak = [], [] + guard = threading.Lock() + + class SlowProc(_FakeProc): + def communicate(self, timeout=None): + with guard: + active.append(1) + peak.append(len(active)) + time.sleep(0.2) + with guard: + active.pop() + return "", "" + + _local_fallback_setup(monkeypatch, lambda cmd, **k: SlowProc()) + + def job(i): + auto_segmentor._thread_session.remote_session_dir = "/home/visitor/tmp/session-1" + auto_segmentor._thread_session.sid = f"s{i}" + auto_segmentor._tracked_run("model", shell=True, check=True) + + threads = [threading.Thread(target=job, args=(i,)) for i in range(3)] + [t.start() for t in threads] + [t.join(20) for t in threads] + assert max(peak) == 1 and len(peak) == 3 # all three ran, one at a time + + +def test_waiting_for_the_web_host_gpu_is_cancellable(monkeypatch): + _local_fallback_setup(monkeypatch, lambda cmd, **k: pytest.fail("ran while the web GPU was busy")) + auto_segmentor._cancelled_sessions.add("s-cancel") + assert auto_segmentor._local_gpu_lock.acquire(timeout=1) # another job holds it + try: + errors = [] + + def job(): + auto_segmentor._thread_session.remote_session_dir = "/home/visitor/tmp/session-1" + auto_segmentor._thread_session.sid = "s-cancel" + try: + auto_segmentor._tracked_run("model", shell=True, check=True) + except RuntimeError as e: + errors.append(str(e)) + + t = threading.Thread(target=job) + t.start() + t.join(10) + assert errors == ["Inference cancelled"] + finally: + auto_segmentor._local_gpu_lock.release() + auto_segmentor._cancelled_sessions.discard("s-cancel") + + +def test_disabled_path_never_touches_the_web_gpu_lock(monkeypatch): + # With workers disabled the old behavior is byte-for-byte: no extra lock. + monkeypatch.setattr(auto_segmentor.subprocess, "Popen", lambda cmd, **k: _FakeProc()) + assert auto_segmentor._local_gpu_lock.acquire(timeout=1) + try: + out = [] + t = threading.Thread(target=lambda: out.append( + auto_segmentor._tracked_run("model", shell=True, check=True).returncode)) + t.start() + t.join(5) + assert out == [0] # would hang if the lock were taken + finally: + auto_segmentor._local_gpu_lock.release() + + +# ── Review round: findings against the parallel-jobs change ───────────────── + +def test_job_waits_when_another_job_reserves_a_worker_during_its_health_check(monkeypatch): + # Finding: the busy-worker snapshot predates the health checks. Here job B + # reserves w2 while job A is still checking it, so A must wait for it + # instead of concluding "nothing to wait for" and falling back. + monkeypatch.setenv("GPU_WORKER_HOSTS", "w1,w2") + monkeypatch.setenv("GPU_WORKER_MAX_WAIT_SECONDS", "30") + state = {"healthy": False, "b_took_it": False} + + def score(h): + if h == "w2" and not state["b_took_it"]: + state["b_took_it"] = True + gpu_workers._busy_hosts.add("w2") # job B takes it mid-check, once + return (0.0, 0.0, 0) if state["healthy"] and h == "w2" else None + + monkeypatch.setattr(gpu_workers, "_host_score", score) + got = [] + start = time.monotonic() + t = threading.Thread(target=lambda: got.append((gpu_workers.acquire_host(), time.monotonic() - start))) + t.start() + time.sleep(0.4) + assert got == [] # still waiting, did not return None / fall back + state["healthy"] = True + gpu_workers.release_host("w2") + t.join(10) + assert got and got[0][0] == "w2" and got[0][1] >= 0.35 + gpu_workers.release_host("w2") + + +def test_cancel_is_noticed_even_when_the_web_gpu_lock_is_free(): + # Finding: cancellation was only checked when the lock wait timed out. + with pytest.raises(RuntimeError, match="cancelled"): + with auto_segmentor._local_gpu_slot(lambda: True): + pytest.fail("ran a model for a cancelled job") + assert auto_segmentor._local_gpu_lock.acquire(blocking=False) # lock was released + auto_segmentor._local_gpu_lock.release() + + +def test_cancel_landing_just_as_the_web_gpu_lock_frees_is_noticed(): + assert auto_segmentor._local_gpu_lock.acquire(timeout=1) # another job holds it + flag, outcome = [], [] + + def waiter(): + try: + with auto_segmentor._local_gpu_slot(lambda: bool(flag)): + outcome.append("ran") + except RuntimeError: + outcome.append("cancelled") + + t = threading.Thread(target=waiter) + t.start() + time.sleep(0.2) + flag.append(True) # user cancels while it waits... + auto_segmentor._local_gpu_lock.release() # ...and the lock frees at once + t.join(10) + assert outcome == ["cancelled"] + + +def _media_agentic_probe(monkeypatch, tmp_path, enabled): + """Run the MedIA-Agentic runner with a fake subprocess; return whether the + web host's GPU lock was held while the model command ran.""" + if enabled: + _enable(monkeypatch) + held = [] + + def fake_run(cmd, **kw): + held.append(auto_segmentor._local_gpu_lock.locked()) + return subprocess.CompletedProcess(cmd, 1, "", "stop here") + + ct = tmp_path / "ct.nii.gz" + ct.write_bytes(b"x") + monkeypatch.setattr(subprocess, "run", fake_run) + with pytest.raises(RuntimeError, match="MedIA-Agentic inference failed"): + auto_segmentor._run_media_agentic_inference(str(ct), str(tmp_path / "s"), "organs") + return held + + +def test_media_agentic_queues_on_the_web_gpu_lock_when_workers_are_enabled(monkeypatch, tmp_path): + # Finding: it launches a GPU model directly, bypassing the lock, so it could + # run alongside a local fallback job or another MedIA-Agentic job. + assert _media_agentic_probe(monkeypatch, tmp_path, enabled=True) == [True] + + +def test_media_agentic_is_unchanged_when_workers_are_disabled(monkeypatch, tmp_path): + assert _media_agentic_probe(monkeypatch, tmp_path, enabled=False) == [False] + + +# ── Review round 2 ────────────────────────────────────────────────────────── + +def _run_sessions_blocked(monkeypatch, slots, session_ids): + """Start one model run per entry of session_ids (repeats = the same session + submitted again); return the peak number running at once and per-session peaks.""" + monkeypatch.setattr(auto_segmentor, "_job_slots", threading.BoundedSemaphore(slots)) + inside, peaks, release = {}, {"all": 0}, threading.Event() + lock = threading.Lock() + + def fake_epai(session_dir, **kw): + sid = auto_segmentor._thread_session.sid + with lock: + inside[sid] = inside.get(sid, 0) + 1 + peaks["all"] = max(peaks["all"], sum(inside.values())) + peaks[sid] = max(peaks.get(sid, 0), inside[sid]) + release.wait(10) + with lock: + inside[sid] -= 1 + return "out" + + monkeypatch.setattr(auto_segmentor, "_run_epai_inference", fake_epai) + monkeypatch.setattr(auto_segmentor, "_resolve_conda_activate_path", lambda: "") + threads = [threading.Thread(target=auto_segmentor.run_auto_segmentation, + args=("in", "/s", "ePAI"), kwargs={"session_id": sid}) + for sid in session_ids] + [t.start() for t in threads] + time.sleep(0.5) + snapshot = dict(peaks) + release.set() + [t.join(10) for t in threads] + return snapshot, peaks + + +def test_a_repeated_session_never_overlaps_itself(monkeypatch): + # Finding: with parallel jobs, a retry / double click for the same session + # ran alongside the first, both writing the same workspace and outputs. + before_release, final = _run_sessions_blocked(monkeypatch, slots=2, session_ids=["dup", "dup"]) + assert before_release["dup"] == 1 # the second waits for the first + assert final["dup"] == 1 # ... and still runs afterwards (no run lost) + + +def test_a_waiting_duplicate_does_not_use_up_a_job_slot(monkeypatch): + # dup, dup, other with 2 slots: the duplicate waits on its session, not on + # a slot, so the other session still runs in parallel with the first. + before_release, _ = _run_sessions_blocked(monkeypatch, slots=2, session_ids=["dup", "dup", "other"]) + assert before_release["all"] == 2 and before_release["dup"] == 1 and before_release["other"] == 1 + + +def test_session_locks_do_not_leak(monkeypatch): + _run_sessions_blocked(monkeypatch, slots=2, session_ids=["a", "a", "b"]) + assert auto_segmentor._session_locks == {} + + +def test_legacy_single_host_epai_mode_keeps_jobs_serial(monkeypatch): + # Finding: EPAI_REMOTE_ENABLED sends every ePAI job to one fixed GPU outside + # the worker pool, which parallel jobs would overload. + _enable(monkeypatch, "w1,w2") + monkeypatch.setenv("GPU_WORKER_PARALLEL", "2") + assert auto_segmentor.max_parallel_jobs() == 2 + monkeypatch.setenv("EPAI_REMOTE_ENABLED", "true") + assert auto_segmentor.max_parallel_jobs() == 1 + monkeypatch.setenv("EPAI_REMOTE_ENABLED", "false") + assert auto_segmentor.max_parallel_jobs() == 2 + + +# ── Review round 3 ────────────────────────────────────────────────────────── + +@pytest.mark.parametrize("spelling", ["1", "true", "TRUE", "yes", "y", "Y", " on ", "On"]) +def test_every_spelling_that_enables_legacy_epai_mode_forces_serial_jobs(monkeypatch, spelling): + # Finding: the ePAI runner accepts "y" but the serial-mode guard used a + # parser that did not, so EPAI_REMOTE_ENABLED=y plus parallel jobs would have + # sent two ePAI jobs to the same fixed GPU. + _enable(monkeypatch, "w1,w2") + monkeypatch.setenv("GPU_WORKER_PARALLEL", "2") + monkeypatch.setenv("EPAI_REMOTE_ENABLED", spelling) + assert auto_segmentor._is_truthy(spelling) is True # the runner really treats it as on + assert auto_segmentor.max_parallel_jobs() == 1 + + +@pytest.mark.parametrize("spelling", ["", "0", "false", "no", "n", "off", "maybe", "2"]) +def test_other_spellings_leave_parallel_jobs_alone(monkeypatch, spelling): + _enable(monkeypatch, "w1,w2") + monkeypatch.setenv("GPU_WORKER_PARALLEL", "2") + monkeypatch.setenv("EPAI_REMOTE_ENABLED", spelling) + assert auto_segmentor.max_parallel_jobs() == 2 + + +def test_there_is_only_one_yes_no_parser(): + # The ePAI runner and the worker settings must never disagree about a value. + assert auto_segmentor._is_truthy("y") and gpu_workers._truthy("y") + for value in ["1", "true", "yes", "y", "on", "Y", " TRUE ", "0", "false", "no", "n", "off", "", None, "t", "2"]: + assert auto_segmentor._is_truthy(value) == gpu_workers._truthy(value), value + + +# ── Review round 4 ────────────────────────────────────────────────────────── + +def _score_while_another_job_fails_on(monkeypatch, victim, scores): + """Health scores where, as `victim` is being checked, another job reserves + it, fails on it (cooldown) and releases it: the scorer's result is stale.""" + state = {"done": False} + + def score(h): + if h == victim and not state["done"]: + state["done"] = True + gpu_workers._busy_hosts.add(victim) # job B reserves it... + gpu_workers.mark_failed(victim) # ...fails on it... + gpu_workers.release_host(victim) # ...and releases it + return scores[h] + + monkeypatch.setattr(gpu_workers, "_host_score", score) + + +def test_a_worker_that_failed_meanwhile_is_not_reserved_on_a_stale_score(monkeypatch): + # Finding: the commit step checked busy but not cooldown, so a worker another + # job had just put in a 600 s cooldown was handed straight to a new job. + monkeypatch.setenv("GPU_WORKER_HOSTS", "w1") + _score_while_another_job_fails_on(monkeypatch, "w1", {"w1": (0.0, 0.0, 0)}) + assert gpu_workers.acquire_host() is None # nothing usable: the caller falls back + assert "w1" not in gpu_workers._busy_hosts + + +def test_another_healthy_worker_is_used_when_the_best_one_just_failed(monkeypatch): + monkeypatch.setenv("GPU_WORKER_HOSTS", "w1,w2") + _score_while_another_job_fails_on(monkeypatch, "w1", {"w1": (0.0, 0.0, 0), "w2": (5.0, 0.0, 0)}) + assert gpu_workers.acquire_host() == "w2" # w1 ranked first but is cooling down now + assert gpu_workers._busy_hosts == {"w2"} + gpu_workers.release_host("w2")