From c959b3740be9caf2065859cfdd450640b8adb090 Mon Sep 17 00:00:00 2001 From: Aditya Sanjeev Date: Fri, 2 Oct 2026 20:33:20 -0700 Subject: [PATCH] Run several website model jobs at once, one per GPU worker Today one global lock makes every job wait, even when other GPUs are idle. GPU_WORKER_PARALLEL=N (default 1, capped by the number of workers, ignored while workers are disabled or the older EPAI_REMOTE_ENABLED single-host mode is on) lets up to N jobs run at once. With the default, behavior is unchanged. Safety: - Worker reservation: workers are health-checked without being reserved and the winner is reserved atomically afterwards. Previously the check reserved every idle worker, so with two simultaneous jobs the second saw "no worker free" and would have run on the web host's GPU. Reservations are re-read after the health checks, so a worker another job took meanwhile is waited for rather than mistaken for "nothing to wait for". The winner is committed only if, under the lock, it is still unreserved and not in cooldown, so a worker another job just failed on is never handed on a stale health score. A worker released during the health checks is reconsidered before waiting or falling back, so a job is not sent to the web host while a worker has just become free. - A job that finds every usable worker busy with our own jobs waits (bounded by GPU_WORKER_MAX_WAIT_SECONDS, cancellable) instead of overflowing to the web host. It uses the web host only when no worker is up, or after the wait. - The web host's own GPU runs at most one model at a time however many jobs are in flight (a lock of its own). This includes MedIA-Agentic, which runs on the web host's GPU directly and used to be covered only by the old global lock. Cancellation is checked after the lock is acquired as well as while waiting. Not taken when workers are disabled, so that path is unchanged. - A session never runs against itself: a repeat request for a session whose job is queued or running waits for it (as it did behind the old lock) without using a job slot, instead of overlapping it in the same workspace. - The older EPAI_REMOTE_ENABLED mode keeps jobs strictly serial. It is read with the same yes/no parser the ePAI runner uses, and there is now a single parser for every environment flag, so a spelling such as "y" cannot mean yes to one setting and no to another. - The SSH connection master per worker is created under a lock. - The queue cap (INFERENCE_MAX_PENDING) defaults to 3 + N (4 as before at 1). Verified with a stress test (8 jobs over 3 workers: no double booking), unit tests for waiting/cancel/limits and for every review finding (each fails on the code without its fix), and a real run of 3 ePAI jobs with N=2 on bdmap2 and bdmap4: jobs 1 and 2 ran side by side, job 3 took the first free worker; 92 s total against about 140 s one at a time. --- flask-server/.env.example | 5 + flask-server/api/api_blueprint.py | 10 +- flask-server/deploy/GPU_WORKERS.md | 52 ++- flask-server/services/auto_segmentor.py | 149 ++++-- flask-server/services/gpu_workers.py | 115 ++++- flask-server/tests/unit/test_gpu_workers.py | 494 ++++++++++++++++++++ 6 files changed, 766 insertions(+), 59 deletions(-) diff --git a/flask-server/.env.example b/flask-server/.env.example index 02fa68188..2ef35879f 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 e7773b93d..9eab280ef 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 3424c9f08..a2f6c6c8e 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 2e1e4920a..2e950c90c 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 6b1f31474..29ba8478d 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 a210f43b0..d16e54db6 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")