From e30d53d43189f2788072f6eb772632bf67dad793 Mon Sep 17 00:00:00 2001 From: Obvious Date: Sun, 11 Oct 2026 19:47:55 +0000 Subject: [PATCH] refactor: non-blocking remote job dispatch (HARN-4) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Async compute routes now await their bodies on a dedicated bounded pool (status.run_compute_async / run_task_async, 8 slots) instead of holding a shared request-threadpool token for the whole response — a remote job's result() wait used to pin one of ~40 worker tokens for up to 1800 s, so a few concurrent traces could starve /health and every other sync endpoint. Routes converted by extracting their sync bodies (_replay_bundle-style names) and shelling to the pool: /replay/{id}, /replay/{id}/psi/{i}, /replay/{id}/fieldlines/{i}, /eq_surrogate, /surrogate, /export/{id}, /export, POST /compute/benchmark, POST /virtual/{id}. Bounded sync routes (/compute probe, /compute/cancel, USD export) deliberately stay sync. Provider selection, depth-1 queue, local fallback, RemoteComputeError typing, provenance labels, and HARN-1/HARN-2 progress + cancel routes are unchanged; agent tools now call the sync bodies directly. New tests: run_task_async semantics (remote success + fallback), router- level remote coverage incl. fallback labeling (audit gap 1), /health responsiveness under 4 parallel slow-remote replays (fails on the old sync dispatch), and cancel/progress SSE through the awaited dispatch. Co-authored-by: Tapan Ganatma Nakkina --- fusionlab/agent_tools.py | 10 +- fusionlab/api.py | 7 +- fusionlab/api_replay.py | 59 +++++-- fusionlab/api_virtual.py | 13 +- fusionlab/compute/ssh.py | 8 +- fusionlab/compute/status.py | 31 +++- tests/test_api.py | 4 +- tests/test_nonblocking_compute.py | 269 ++++++++++++++++++++++++++++++ 8 files changed, 369 insertions(+), 32 deletions(-) create mode 100644 tests/test_nonblocking_compute.py diff --git a/fusionlab/agent_tools.py b/fusionlab/agent_tools.py index 9f13534..1b8b583 100644 --- a/fusionlab/agent_tools.py +++ b/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", diff --git a/fusionlab/api.py b/fusionlab/api.py index d69cd44..f246e28 100644 --- a/fusionlab/api.py +++ b/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")} diff --git a/fusionlab/api_replay.py b/fusionlab/api_replay.py index 1b084ba..3dd79fa 100644 --- a/fusionlab/api_replay.py +++ b/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) diff --git a/fusionlab/api_virtual.py b/fusionlab/api_virtual.py index d73a625..49589e0 100644 --- a/fusionlab/api_virtual.py +++ b/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()) diff --git a/fusionlab/compute/ssh.py b/fusionlab/compute/ssh.py index 282fd8a..354d1b2 100644 --- a/fusionlab/compute/ssh.py +++ b/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: → 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, diff --git a/fusionlab/compute/status.py b/fusionlab/compute/status.py index 50fc4ad..982f98f 100644 --- a/fusionlab/compute/status.py +++ b/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 diff --git a/tests/test_api.py b/tests/test_api.py index 4e5915c..01cfb75 100644 --- a/tests/test_api.py +++ b/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: diff --git a/tests/test_nonblocking_compute.py b/tests/test_nonblocking_compute.py new file mode 100644 index 0000000..70b3511 --- /dev/null +++ b/tests/test_nonblocking_compute.py @@ -0,0 +1,269 @@ +"""HARN-4: async routes await compute on the dedicated pool — the shared request threadpool stays free. + +Three layers, mirroring the issue's acceptance criteria: + +- ``run_task_async`` — the awaited entry point keeps run_task's provider selection, local fallback + and provenance exactly. +- Router-level remote coverage (the audit's gap 1): a configured remote through real routes, + including the fallback labels the responses and provenance carry. +- The concurrency property: parallel replays against a slow remote never starve ``/health``. + +No SSH: the provider is a fake whose ``result()`` takes its time (the exact wait a sync router +used to pin a threadpool token for). +""" + +import asyncio +import threading +import time + +import anyio +import httpx +import numpy as np +import pytest +from fastapi.testclient import TestClient + +from fusionlab import api_replay, compute +from fusionlab.api import app +from fusionlab.compute import status +from fusionlab.compute.provider import JobHandle +from fusionlab.compute.ssh import RemoteComputeError + +client = TestClient(app) + + +# ---------------------------------------------------------------- fakes and helpers +def shaped_payload(task: str, inputs: dict, psi_shape: tuple[int, ...]) -> dict | np.ndarray: + """Task-shaped results the replay path's consumers accept (the bodies index and scale arrays, + they never inspect provenance of the fake's values).""" + if task == "model_metrics": + return {"block_size": 100, "held_out_blocks": [], "test_shot_ids": []} + if task == "run_surrogate_shot": + return np.ones(np.shape(inputs["shot"]["t_s"])) + if task == "run_eq_surrogate": + return np.zeros(psi_shape) + if task == "run_surrogate": + return np.ones(np.shape(inputs["features"])[:1]) + raise AssertionError(f"unexpected task {task!r}") + + +class SlowRemote: + """A provider-shaped remote: submit queues instantly, result() takes delay_s to land. + + ``stream`` matches the worker's SSE vocabulary (queued/running/terminal); with ``fail`` every + job dies the mid-job way run_task's fallback is built on; ``cancel`` makes result() raise the + same cancelled error HARN-2's route expects. + """ + + name, host = "ssh", "slowsrvr" + + def __init__(self, delay_s: float = 0.0, fail: str | None = None, + psi_shape: tuple[int, ...] = (1, 1, 1)) -> None: + self.delay_s, self.fail, self.psi_shape = delay_s, fail, psi_shape + self._inputs: dict[str, dict] = {} + self._cancelled = threading.Event() + self._done = threading.Event() + self.cancel_calls = 0 + + def capabilities(self) -> dict: + return {"provider": self.name, "host": self.host, "surrogate_available": True, + "eq_surrogate_available": True, "version_mismatch": None, "remote": {}} + + def health(self) -> None: + return None + + def submit(self, task: str, inputs: dict) -> JobHandle: + self._inputs[task] = inputs + return JobHandle(job_id=f"j-{task}", task=task, provider=self.name) + + async def stream(self, job: JobHandle): + yield {"event": "queued", "task": job.task} + yield {"event": "running", "task": job.task} + while not self._done.is_set(): + await asyncio.sleep(0.02) + if self._cancelled.is_set(): + yield {"event": "cancelled", "task": job.task} + elif self.fail: + yield {"event": "error", "error": self.fail} + else: + yield {"event": "completed", "job_id": job.job_id, "task": job.task} + + def cancel(self, job: JobHandle) -> None: + self.cancel_calls += 1 + self._cancelled.set() + + def result(self, job: JobHandle): + deadline = time.monotonic() + self.delay_s + while time.monotonic() < deadline and not self._cancelled.is_set(): + time.sleep(0.01) + self._done.set() + if self._cancelled.is_set(): + raise RemoteComputeError("job", "remote job was cancelled") + if self.fail: + raise RemoteComputeError("job", self.fail) + return shaped_payload(job.task, self._inputs[job.task], self.psi_shape) + + +@pytest.fixture() +def clean(monkeypatch): + """No selected provider, no status state leaking between tests (same as test_compute_status), + plus the replay-path caches: a stub answer must never be served to another test.""" + monkeypatch.setenv("FUSIONLAB_COMPUTE_PROVIDER", "local") + monkeypatch.setattr(compute, "_provider", None) + status._reset_for_tests() + api_replay._eq_psi.cache_clear() + api_replay._eq_model.cache_clear() + yield + status._reset_for_tests() + api_replay._eq_psi.cache_clear() + api_replay._eq_model.cache_clear() + + +# ---------------------------------------------------------------- run_task_async +def test_run_task_async_matches_run_task_on_remote_success(clean): + """Remote success through the awaited entry point: same payload, same provenance as run_task.""" + stub = SlowRemote(delay_s=0.05) + compute._provider = stub + out = asyncio.run(status.run_task_async("model_metrics", {"model": "surrogate"})) + assert out == {"block_size": 100, "held_out_blocks": [], "test_shot_ids": []} + assert status.provenance("model_metrics") == {"provider": "ssh", "host": "slowsrvr"} + + +def test_run_task_async_falls_back_and_records(clean): + """Remote failure through the awaited entry point: the local answer comes back, the fallback + is recorded with the remote's reason, and provenance flips to local.""" + stub = SlowRemote(fail="worker died mid-job") + compute._provider = stub + out = asyncio.run(status.run_task_async("model_metrics", {"model": "surrogate"})) + assert "block_size" in out # the local metrics file answered, not the fake + fb = status.fallback_state() + assert fb["fallback"] is True and "mid-job" in fb["reason"] + assert status.provenance("model_metrics") == {"provider": "local", "host": None} + + +# ---------------------------------------------------------------- router-level remote (audit gap 1) +def test_router_serves_the_remote_answer_with_remote_provenance(clean): + """A configured remote through a real route: the response body is the remote's own answer and + the provenance names it — run_task-through-router, exercised for the first time here.""" + stub = SlowRemote(delay_s=0.05) + compute._provider = stub + r = client.get("/surrogate") + assert r.status_code == 200 + assert r.json() == {"block_size": 100, "held_out_blocks": [], "test_shot_ids": []} + assert status.provenance("model_metrics") == {"provider": "ssh", "host": "slowsrvr"} + + +def test_replay_route_labels_remote_provenance(clean): + """The replay bundle says where its learned correction ran: the remote host, with no fallback.""" + stub = SlowRemote(delay_s=0.02) + compute._provider = stub + body = client.get("/replay/30166").json() + assert body["summary"]["learned_computed_on"] == "slowsrvr" + assert status.fallback_state()["fallback"] is False + + +def test_replay_route_falls_back_with_honest_labels(clean): + """A dying remote mid-replay: the bundle is still served, built locally — learned_computed_on + is None (no host), the fallback reason is recorded, and the provenance map says local.""" + stub = SlowRemote(fail="worker died mid-job") + compute._provider = stub + body = client.get("/replay/30166").json() + assert body["summary"]["learned_computed_on"] is None + fb = status.fallback_state() + assert fb["fallback"] is True and "mid-job" in fb["reason"] + assert status.provenance("run_surrogate_shot")["provider"] == "local" + + +# ---------------------------------------------------------------- the concurrency property (HARN-4) +async def _health_stays_responsive_while_replays_wait() -> None: + """Four parallel replays against a slow remote; a watcher polls /health throughout. + + The shared default limiter is shrunk to 2 tokens first: pre-HARN-4 (sync routers), each replay + held one of those 2 tokens for its whole ~2 s of remote waiting and /health starved past 1 s; + post-HARN-4 the replays wait on the compute pool and /health keeps its token. + """ + anyio.to_thread.current_default_thread_limiter().total_tokens = 2 + compute._provider = SlowRemote(delay_s=1.0, psi_shape=api_replay._shot(30166)["psi_n"].shape) + + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as c: + + async def one_replay() -> None: + r = await c.get("/replay/30166") + assert r.status_code == 200 + + async def watch_health(stop: asyncio.Event) -> list[float]: + latencies = [] + while not stop.is_set(): + t0 = time.monotonic() + r = await c.get("/health") + latencies.append(time.monotonic() - t0) + assert r.json() == {"ok": True} + await asyncio.sleep(0.05) + return latencies + + stop = asyncio.Event() + watcher = asyncio.create_task(watch_health(stop)) + await asyncio.gather(*(one_replay() for _ in range(4))) + stop.set() + latencies = await watcher + assert latencies and max(latencies) < 1.0, max(latencies) + + +def test_health_stays_responsive_while_replays_wait(clean): + asyncio.run(_health_stays_responsive_while_replays_wait()) + + +# ---------------------------------------------------------------- HARN-1/HARN-2 through the new dispatch +async def _cancel_stops_an_awaited_remote_job() -> None: + """HARN-2's route must reach a job launched through run_task_async (an async route in flight), + and the request that was waiting on it must come back as the local fallback answer.""" + compute._provider = SlowRemote(delay_s=3.0) + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as c: + pending = asyncio.create_task(c.get("/surrogate")) + for _ in range(250): + if status.active_job("model_metrics") is not None: + break + await asyncio.sleep(0.02) + assert status.active_job("model_metrics") is not None + cancel_r = await c.post("/compute/cancel?task=model_metrics") + out = await pending + stub = compute._provider + assert cancel_r.status_code == 200 and cancel_r.json()["cancelling"] is True + assert stub.cancel_calls == 1 # the remote was told, not just abandoned client-side + assert out.status_code == 200 and "block_size" in out.json() # the fallback answered locally + fb = status.fallback_state() + assert fb["fallback"] is True and "cancel" in fb["reason"] + assert status.provenance("model_metrics")["provider"] == "local" + + +def test_cancel_stops_an_awaited_remote_job(clean): + asyncio.run(_cancel_stops_an_awaited_remote_job()) + + +async def _progress_streams_an_awaited_remote_job() -> None: + """HARN-1's SSE route must relay a job launched through run_task_async, labelled with the + remote provider it ran on.""" + compute._provider = SlowRemote(delay_s=1.0) + transport = httpx.ASGITransport(app=app) + async with httpx.AsyncClient(transport=transport, base_url="http://test") as c: + pending = asyncio.create_task(c.get("/surrogate")) + for _ in range(250): + if status.active_job("model_metrics") is not None: + break + await asyncio.sleep(0.02) + body = "" + async with c.stream("GET", "/compute/progress?task=model_metrics") as r: + assert r.status_code == 200 + async for chunk in r.aiter_raw(): + body += chunk.decode() + if "event: running" in body: + break # the live prefix is the assertion; leaving early cancels the stream + out = await pending + assert "event: queued" in body and "event: running" in body + assert '"provider": "ssh"' in body # the relay names where the job actually ran + assert out.status_code == 200 # the compute request itself completes remotely, untouched + + +def test_progress_streams_an_awaited_remote_job(clean): + asyncio.run(_progress_streams_an_awaited_remote_job())