Skip to content

Split submit_request from handle_response; add async submit + get_response#648

Merged
JadenFiotto-Kaufman merged 1 commit into
devfrom
feat/async-submit
Apr 21, 2026
Merged

Split submit_request from handle_response; add async submit + get_response#648
JadenFiotto-Kaufman merged 1 commit into
devfrom
feat/async-submit

Conversation

@JadenFiotto-Kaufman
Copy link
Copy Markdown
Member

Summary

Refactor `RemoteBackend` so a subclass (or an async caller that wants streaming control) can submit a request without paying for `handle_response`'s side effects, and can use `httpx.AsyncClient` instead of the sync path.

Motivated by the workbench SSE endpoint flow: the FastAPI route wants to `async yield` the initial `ResponseModel` to the browser before anything touches the terminal-display logic or raises on error, and it wants to run the POST inside the event loop without a threadpool.

Changes

  • `submit_request` no longer calls `handle_response`. It's now a pure POST that returns the initial `ResponseModel` (still records `self.job_id` as natural bookkeeping). Callers are responsible for dispatching the response when they're ready.
  • Add `async_submit_request` — mirrors `submit_request` via `httpx.AsyncClient`.
  • Add `async_get_response` — mirrors `get_response` for symmetry in the non-blocking poll path.
  • Factor the shared `status_code`/`ResponseModel` parsing into a small `_parse_submit_response` helper used by both submit variants.
  • Update the three internal callers (`blocking_request`, `async_request`, `non_blocking_request`) to call `handle_response` explicitly on the initial response, preserving today's behaviour. `async_request` now also uses `async_submit_request` instead of the sync one.

Test plan

  • Unit / integration tests in nnsight pass on `dev`.
  • Blocking remote trace (existing `with model.trace(..., remote=True)`) still shows the status display and completes.
  • Async remote trace (`tracer.asynchronous=True`) still completes end-to-end; no thread blocking on the POST.
  • Non-blocking submit + poll round-trip still works (submit → job_id → poll to COMPLETED).
  • Workbench SSE endpoint (companion PR) exercises the new `async_submit_request` end-to-end against NDIF — verified locally against gpt2 on api.ndif.us.

Companion

Workbench side: ndif-team/workbench#112 — its `StreamingRemoteBackend` now calls `super().async_submit_request()` instead of duplicating the POST.

submit_request now does only the POST and returns the initial
ResponseModel (with self.job_id recorded). Dispatching the response
to handle_response is left to the caller so that async / streaming
consumers can yield the initial response before any side effects
(display updates, RemoteException raises) run.

Add async_submit_request and async_get_response using
httpx.AsyncClient so callers inside an event loop don't need a
threadpool or a duplicated POST path. The new parsing helper
_parse_submit_response is shared between the sync and async submits.

The three internal callers (blocking_request, async_request,
non_blocking_request) now invoke handle_response explicitly on the
initial response to preserve today's behaviour.
JadenFiotto-Kaufman added a commit to ndif-team/workbench that referenced this pull request Apr 20, 2026
The duplicated _async_submit in StreamingRemoteBackend existed only
because the parent RemoteBackend didn't expose an async submit. With
ndif-team/nnsight#648 merged, we can delete it and call
super().async_submit_request() directly — which also means the
job_id bookkeeping (self.job_id = response.id) is handled by the
parent, not open-coded here.

Depends on ndif-team/nnsight#648.
@JadenFiotto-Kaufman JadenFiotto-Kaufman merged commit 92fd355 into dev Apr 21, 2026
1 check passed
@JadenFiotto-Kaufman JadenFiotto-Kaufman deleted the feat/async-submit branch April 21, 2026 22:10
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant