refactor: non-blocking remote job dispatch (HARN-4)

fusionlab/agent_tools.py

@@ -35,8 +35,8 @@
 
 from fusionlab import mast, virtual
 from fusionlab.api import devices, simulate_point
-from fusionlab.api_replay import _shot, db_search, replay_shot, shots
-from fusionlab.api_virtual import EditIn, VirtualIn, virtual_gate, virtual_shot
+from fusionlab.api_replay import _replay_bundle, _shot, db_search, shots
+from fusionlab.api_virtual import EditIn, VirtualIn, _virtual_shot_body, virtual_gate
 from fusionlab.physics import DEVICES
 
 SEARCH_LIMIT = 200   # the tool's own cap; /db/search itself allows 1000 — the model works page-size
@@ -117,7 +117,7 @@ def load_shot(shot_id: int, H: float = 1.0) -> ToolResult:
     """A cached shot's replay summary — the /replay/{id} internals. Time arrays stay out of the model's
     context: the UI draws the traces after the driver opens the shot."""
     try:
-        r = replay_shot(int(shot_id), H=float(H))
+        r = _replay_bundle(int(shot_id), H=float(H))
     except HTTPException as e:
         cached = mast.cached_shots()
         return _err(f"{e.detail} Replay covers only the shots cached under data/shots ({len(cached)}: {cached}); "
@@ -156,7 +156,7 @@ def limit_distances(shot_id: int) -> ToolResult:
     Answers 'how close is shot X to the beta limit over time'. Distances, never forecasts: 1.0 is
     at the limit, not a prediction of a disruption."""
     try:
-        r = replay_shot(int(shot_id))
+        r = _replay_bundle(int(shot_id))
     except HTTPException as e:
         cached = mast.cached_shots()
         return _err(f"{e.detail} Replay covers only the shots cached under data/shots ({len(cached)}: {cached}); "
@@ -288,7 +288,7 @@ def run_whatif(shot_id: int, p_nbi_MW: float | None = None, p_nbi_scale: float |
     if not edit:
         return _err("no edit given: pass at least one of p_nbi_MW/p_nbi_scale, n_scale, ip_scale, bt_scale, nbi_shift_s")
 
-    r = virtual_shot(int(shot_id), VirtualIn(edit=EditIn(**edit), closure=closure))
+    r = _virtual_shot_body(int(shot_id), VirtualIn(edit=EditIn(**edit), closure=closure))
     apply = {"sliders": {k: _fmt(v, 3 if k == "nbi_shift_s" else 2) for k, v in edit.items()},
              "events": ["fusionlab:tab"]}
     proposal = {"tool": "run_whatif",

fusionlab/api.py

@@ -83,17 +83,18 @@ def compute_cancel(task: str):
 
 
 @app.post("/compute/benchmark")
-def compute_benchmark(profile: str = "quick"):
+async def compute_benchmark(profile: str = "quick"):
     """The measured-timings benchmark through the selected provider (the worker's `benchmark` task,
     with the local fallback) — the one named task that reports section progress. POST: it occupies
     the compute slot for its whole run, so it must not be triggerable by navigation or prefetch
-    (same reasoning as the USD export)."""
+    (same reasoning as the USD export). The body is awaited on the compute pool (HARN-4): the run
+    — remote or local — never pins a shared request-threadpool token."""
     from fusionlab.bench import PROFILES
     from fusionlab.compute import status
 
     if profile not in PROFILES:
         raise HTTPException(422, f"profile must be one of {sorted(PROFILES)}")
-    out = status.run_task("benchmark", {"profile": profile})
+    out = await status.run_task_async("benchmark", {"profile": profile})
     return {**out, "computed_on": status.provenance("benchmark").get("host")}
 
 

fusionlab/api_replay.py

@@ -8,7 +8,7 @@
 import os
 import re
 import threading
-from functools import lru_cache
+from functools import lru_cache, partial
 from pathlib import Path
 
 import numpy as np
@@ -19,7 +19,12 @@
 from fusionlab import mast
 from fusionlab.compute import warm_up as compute_warm_up
 from fusionlab.compute.status import capabilities as compute_capabilities
-from fusionlab.compute.status import provenance, run_task
+from fusionlab.compute.status import (
+    provenance,
+    run_compute_async,
+    run_task,
+    run_task_async,
+)
 from fusionlab.physics import (
     BETA_N_LIMIT,
     LIMIT_NAMES,
@@ -149,14 +154,13 @@ def _replay_bundle(shot_id: int, H: float = 1.0) -> dict:
 
 
 @router.get("/replay/{shot_id}")
-def replay_shot(shot_id: int, H: float = 1.0):
+async def replay_shot(shot_id: int, H: float = 1.0):
     """Measured traces, model traces on the same time base, limits, and the geometry to draw the shot."""
-    return _replay_bundle(shot_id, H)
+    return await run_compute_async(partial(_replay_bundle, shot_id, H))
 
 
-@router.get("/replay/{shot_id}/psi/{i}")
-def replay_psi(shot_id: int, i: int):
-    """Normalised poloidal flux (0 axis, 1 separatrix) on the EFIT grid for one time slice, rows = Z."""
+def _replay_psi_slice(shot_id: int, i: int) -> dict:
+    """One psi slice's response body — the blocking part the async route below runs on the compute pool."""
     s = _shot(shot_id)
     if "psi_n" not in s:
         raise HTTPException(404, f"shot {shot_id}'s bundle has no equilibrium reconstruction (psi)")
@@ -172,6 +176,12 @@ def replay_psi(shot_id: int, i: int):
     return out
 
 
+@router.get("/replay/{shot_id}/psi/{i}")
+async def replay_psi(shot_id: int, i: int):
+    """Normalised poloidal flux (0 axis, 1 separatrix) on the EFIT grid for one time slice, rows = Z."""
+    return await run_compute_async(partial(_replay_psi_slice, shot_id, i))
+
+
 @router.post("/replay/{shot_id}/usd")
 def replay_usd(shot_id: int, fmt: str = "usdc"):
     """OpenUSD export (Omniverse-compatible): vessel, PF coils and the plasma boundary time-sampled over the shot.
@@ -212,9 +222,9 @@ def _lines(shot_id: int):
     return pts, q, qc["summary"], out["device"], out["psi_n_start"], where
 
 
-@router.get("/replay/{shot_id}/fieldlines/{i}")
-def replay_fieldlines(shot_id: int, i: int):
-    """Field lines of the EFIT reconstruction at one time slice (axisymmetric; no islands or 3D fields)."""
+def _replay_fieldlines_slice(shot_id: int, i: int) -> dict:
+    """One field-lines response body — the blocking part (a Warp trace through run_task) the async
+    route below runs on the compute pool."""
     s = _shot(shot_id)
     if not 0 <= i < s["t_s"].size:
         raise HTTPException(404, "time index out of range")
@@ -231,6 +241,12 @@ def replay_fieldlines(shot_id: int, i: int):
             "shot_check": {k: summary[k] for k in ("n_compared", "median_rel_err", "p95_rel_err", "psi_n_drift_max")}}
 
 
+@router.get("/replay/{shot_id}/fieldlines/{i}")
+async def replay_fieldlines(shot_id: int, i: int):
+    """Field lines of the EFIT reconstruction at one time slice (axisymmetric; no islands or 3D fields)."""
+    return await run_compute_async(partial(_replay_fieldlines_slice, shot_id, i))
+
+
 @lru_cache(maxsize=1)   # process-lifetime cache: a refetched or retrained artifact on disk is not seen until restart
 def _eq_model():
     """The equilibrium surrogate's holdout shots. None when no trained model is on disk."""
@@ -318,19 +334,19 @@ def start_warm_up():
 
 
 @router.get("/eq_surrogate")
-def eq_surrogate_metrics():
+async def eq_surrogate_metrics():
     """Holdout metrics of the equilibrium surrogate, as written by scripts/train_eq_surrogate.py."""
     try:
-        return run_task("model_metrics", {"model": "eq_surrogate"})
+        return await run_task_async("model_metrics", {"model": "eq_surrogate"})
     except FileNotFoundError:
         raise HTTPException(404, "no metrics: run scripts/train_eq_surrogate.py") from None
 
 
 @router.get("/surrogate")
-def surrogate_metrics():
+async def surrogate_metrics():
     """Holdout errors of the learned correction, as written by `python -m fusionlab.surrogate`."""
     try:
-        return run_task("model_metrics", {"model": "surrogate"})
+        return await run_task_async("model_metrics", {"model": "surrogate"})
     except FileNotFoundError:
         raise HTTPException(404, "no trained model: run `make train`") from None
 
@@ -600,17 +616,24 @@ def _export(bundles: list[dict], fmt: str, H: float):
     raise HTTPException(422, 'fmt must be "csv" or "json"')
 
 
+def _replay_bundles(shot_ids: list[int], H: float) -> list[dict]:
+    """Full replay bundles in request order — one blocking call for the compute pool."""
+    return [_replay_bundle(i, H) for i in shot_ids]
+
+
 @router.get("/export/{shot_id}")
-def export_shot(shot_id: int, H: float = 1.0, fmt: str = "csv"):
+async def export_shot(shot_id: int, H: float = 1.0, fmt: str = "csv"):
     """Take-away export of one shot's derived quantities: every measured trace, model trace and
     summary number the replay serves, as a tidy CSV (one row per shot/slice/quantity, each row
     labelled with its provenance and unit) or as JSON (the full replay bundle + the provenance map).
     The numbers are the replay bundle's own, never a recomputation."""
-    return _export([_replay_bundle(shot_id, H)], fmt, H)
+    bundle = await run_compute_async(partial(_replay_bundle, shot_id, H))
+    return _export([bundle], fmt, H)
 
 
 @router.get("/export")
-def export_compare(shots: str, H: float = 1.0, fmt: str = "csv"):
+async def export_compare(shots: str, H: float = 1.0, fmt: str = "csv"):
     """Compare-variant export: the same tidy table over several shots (the compare view's cap),
     a `shot_id` column so any spreadsheet or notebook can group them. Order follows the request."""
-    return _export([_replay_bundle(i, H) for i in _export_ids(shots)], fmt, H)
+    bundles = await run_compute_async(partial(_replay_bundles, _export_ids(shots), H))
+    return _export(bundles, fmt, H)

fusionlab/api_virtual.py

@@ -2,12 +2,15 @@
 
 from __future__ import annotations
 
+from functools import partial
+
 import numpy as np
 from fastapi import APIRouter, HTTPException
 from pydantic import BaseModel, Field
 
 from fusionlab import virtual
 from fusionlab.api_replay import ATTRIBUTION, _j, _shot
+from fusionlab.compute.status import run_compute_async
 
 router = APIRouter()
 
@@ -58,7 +61,15 @@ def virtual_gate():
 
 
 @router.post("/virtual/{shot_id}")
-def virtual_shot(shot_id: int, body: VirtualIn):
+async def virtual_shot(shot_id: int, body: VirtualIn):
+    """A what-if anchored on a cached real shot — the body runs on the compute pool (HARN-4):
+    the learned closure reaches the surrogate through run_task, whose remote wait must not pin a
+    shared request-threadpool token."""
+    return await run_compute_async(partial(_virtual_shot_body, shot_id, body))
+
+
+def _virtual_shot_body(shot_id: int, body: VirtualIn) -> dict:
+    """The virtual-shot response body (the blocking part the async route above runs on the pool)."""
     s = _shot(shot_id)
     g = virtual.gate()
     edit = virtual.Edit(**body.edit.model_dump())

fusionlab/compute/ssh.py

@@ -4,15 +4,17 @@
 FusionLab never sees or stores a credential — opens a local port-forward to the worker, and talks
 plain HTTP over the tunnel:
 
-    FastAPI router (sync def, threadpool) ── SshProvider ── SSH (asyncssh) ── fusionlab worker
+    FastAPI router (async def) ── compute pool ── SshProvider ── SSH (asyncssh) ── fusionlab worker
                                                 │                on the user's host (binds 127.0.0.1)
                                                 └─ local forward 127.0.0.1:<port> → worker
 
 Wiring, with the SSH part as small as possible:
 
 - ``_SshSession`` owns the only async work (connect, forward, remote autostart) on a private event
-  loop in a thread. The routers are sync ``def`` (threadpool), so the blocking protocol methods
-  never touch the API's event loop; PR-1's calling convention is unchanged.
+  loop in a thread. The API's async routes await compute bodies on the compute pool
+  (``status.run_compute_async`` — HARN-4), so the blocking protocol methods and a remote job's
+  whole ``result()`` wait run on the pool's own threads: the shared request threadpool every sync
+  endpoint (``/health``) depends on stays free no matter how many jobs are in flight.
 - ``SshProvider`` implements the PR-1 protocol: ``submit`` queues the job and returns at once,
   ``result`` waits for the terminal state and decodes the FLF1 payload, ``stream`` relays the
   worker's SSE events to async consumers (PR 3's progress chip). Every failure — unreachable host,

fusionlab/compute/status.py

@@ -23,9 +23,10 @@
 import logging
 import threading
 import time
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Callable
 from concurrent.futures import ThreadPoolExecutor
 from concurrent.futures import TimeoutError as ProbeTimeout
+from functools import partial
 from typing import Any
 
 from fusionlab.compute import get_provider, provider_name
@@ -187,6 +188,34 @@ def run_task(task: str, inputs: dict) -> Any:
     return out
 
 
+# ---------------------------------------------------------------- non-blocking dispatch (HARN-4)
+# Async routes await their compute bodies on this pool instead of the shared request threadpool:
+# a remote job's result() wait (up to DEFAULT_RESULT_TIMEOUT_S) used to pin one of the ~40 worker
+# tokens every sync endpoint shares, and a few long traces starved /health. One plain pool serves
+# whichever loop awaits it (uvicorn's, TestClient portals') — threads and queues need no event loop.
+COMPUTE_SLOTS = 8   # concurrent compute bodies; extra requests queue here, never on shared endpoints
+_task_pool = ThreadPoolExecutor(max_workers=COMPUTE_SLOTS, thread_name_prefix="fusionlab-compute-task")
+
+
+async def run_compute_async(fn: Callable[[], Any]) -> Any:
+    """One blocking compute body, awaited off the shared request threadpool (HARN-4).
+
+    A sync router has no choice but to hold its worker token for the whole response — the token a
+    remote result() poll keeps for its whole life is one /health and /db/search never get back.
+    Bodies awaited here run on the compute pool, so however many jobs wait, shared endpoints keep
+    answering. A cancelled request cannot stop a running thread: the body runs to completion and
+    its caches and provenance land as if the client had waited — the same semantics a sync
+    router's abandoned thread had.
+    """
+    return await asyncio.get_running_loop().run_in_executor(_task_pool, fn)
+
+
+async def run_task_async(task: str, inputs: dict) -> Any:
+    """run_task for async routes — the same provider selection, local fallback and provenance,
+    awaited off the shared request threadpool."""
+    return await run_compute_async(partial(run_task, task, inputs))
+
+
 async def progress_stream(task: str, *, wait_s: float | None = None,
                           grace_s: float | None = None) -> AsyncIterator[str]:
     """Server-sent events for the job of `task`: the provider's stream relayed frame-for-frame in the

tests/test_api.py

@@ -322,7 +322,9 @@ def slow_fetch(shot_id):
     out = []
 
     def hit():
-        out.append(api_replay.replay_shot(uncached)["meta"]["shot_id"])
+        # the async route shells out to this body on the compute pool; the concurrency under test
+        # lives here, so three threads call the body directly (HARN-4)
+        out.append(api_replay._replay_bundle(uncached)["meta"]["shot_id"])
 
     threads = [threading.Thread(target=hit) for _ in range(3)]
     for t in threads: