From d868b538be07186fc325dbb7fa1ed39ea8dd2713 Mon Sep 17 00:00:00 2001 From: JadenFK Date: Thu, 13 Aug 2026 14:36:21 -0400 Subject: [PATCH 1/5] Run against a 0.8 NDIF nnterp's nnsight-0.8 branch carries the port, and nnsightful's tools speak nnterp's vocabulary rather than nnsight's, so neither repo needed a source change -- but the pickled request carries the model wrapper, whose class is nnterp.StandardizedTransformer, and NDIF ships neither library. Register it by value alongside nnsightful; without it the server cannot read the payload and reports ModuleNotFoundError dressed as a corrupt-payload error. Also trust better-sqlite3 so bun builds its native binding -- without it drizzle-kit push cannot open the local SQLite database at all. Co-Authored-By: Claude Opus 5 --- workbench/_api/routes/__init__.py | 12 ++++++++++++ workbench/_web/bun.lock | 1 + workbench/_web/package.json | 1 + 3 files changed, 14 insertions(+) diff --git a/workbench/_api/routes/__init__.py b/workbench/_api/routes/__init__.py index 32a7057d..c0f2fc75 100644 --- a/workbench/_api/routes/__init__.py +++ b/workbench/_api/routes/__init__.py @@ -8,7 +8,19 @@ from nnsight import ndif import nnsightful +import nnterp + +# Ship both libraries' source with each request, so the server can rebuild them +# without having them installed. NDIF's own requirements carry neither. +# +# nnsightful holds the tools the traced block calls. nnterp holds the class of +# the model the block is written against -- every tool talks to a +# StandardizedTransformer (`model.layers_output`, `model.project_on_vocab`), and +# the wrapper is part of the pickled request, so without this the server fails to +# read the payload at all: "ModuleNotFoundError: No module named 'nnterp'", +# surfaced to the user as a corrupt-payload error rather than a missing import. ndif.register(nnsightful) +ndif.register(nnterp) __all__ = [ "lens", diff --git a/workbench/_web/bun.lock b/workbench/_web/bun.lock index ec97becd..7ddbf43e 100644 --- a/workbench/_web/bun.lock +++ b/workbench/_web/bun.lock @@ -84,6 +84,7 @@ }, }, "trustedDependencies": [ + "better-sqlite3", "nnsightful", ], "packages": { diff --git a/workbench/_web/package.json b/workbench/_web/package.json index 28ffc17e..f798a29b 100644 --- a/workbench/_web/package.json +++ b/workbench/_web/package.json @@ -93,6 +93,7 @@ "typescript": "^5.9.2" }, "trustedDependencies": [ + "better-sqlite3", "nnsightful" ] } From dd99cdb1bb2e8148a106761980af0b378cbf0e4f Mon Sep 17 00:00:00 2001 From: JadenFK Date: Thu, 13 Aug 2026 14:48:05 -0400 Subject: [PATCH 2/5] One streamed request per run, instead of start / poll / collect A tool run took three legs: POST /start for a job id, the browser polling NDIF's /response/{id} until COMPLETED, then POST /results/{id} to collect and shape it. It is now one POST that stays open and answers in Server-Sent Events -- each NDIF status as it lands, then the finished payload. nnsight 0.8 ships AsyncRemoteBackend, which submits on the trace's exit and then hands back the raw status updates instead of consuming them. That is the whole mechanism: _api/sse.py turns those updates into `status` frames, calls the tool's to_data_obj when the saved values arrive, and emits one `data` frame. There is no custom backend subclass -- an earlier attempt at this (feat/sse-remote-backend, against 0.7) needed 102 lines of one because 0.7 had no async path. What this buys, beyond the round trips: * The browser no longer talks to NDIF at all, so config.ts has no NDIF URL and needs none -- one origin, and no second host to reach through a tunnel or an ingress. The hardcoded localhost:5001 goes with it; NDIF's API has been on 8001 since long before 0.8. * Status is pushed rather than sampled, so QUEUED position and the RUNNING->COMPLETED transition land when they happen rather than up to a second later. * A collect step could find its model gone -- causal_mediation carried a 503 for exactly that window. One connection, no window. What it costs is that a run lives and dies with its connection: polling a job id survived a reload, and this does not. NDIF still finishes the job, and the result is still written by the caller's mutation, so what is lost is the ability to rejoin a run in progress. Runs are seconds to a couple of minutes. The routes lost their response_model, so the payload is encoded by hand -- through jsonable_encoder, because a payload is often a plain dict with pydantic models nested inside it (a generation's `completion` is a list of Token). Long silences are covered by a comment frame every 15s. A cold 70B deploy can sit quiet for minutes, and an idle connection is what proxies reap; the warmup path in deployApi used to poll on a 20-minute ceiling and now just holds the stream open. Errors split by when they happen. Before the stream opens -- a 403 for a model the caller cannot use -- they stay ordinary HTTP failures. After, they can only be an `error` frame, because the status line is long gone. Verified against the self-hosted 0.8 NDIF: the logit lens streams QUEUED -> DISPATCHED -> RUNNING -> COMPLETED and returns Paris for "The Eiffel Tower is in the city of", and generation returns its completion. Typecheck and lint show the same errors as before the change, none of them new. Co-Authored-By: Claude Opus 5 --- workbench/_api/routes/activation_patching.py | 55 +-- workbench/_api/routes/causal_mediation.py | 86 ++--- workbench/_api/routes/j_lens.py | 51 +-- workbench/_api/routes/lens.py | 164 ++++----- workbench/_api/routes/logit_lens.py | 52 +-- workbench/_api/routes/models.py | 324 ++++++------------ workbench/_api/sse.py | 212 ++++++++++++ workbench/_api/state.py | 30 +- .../_web/src/lib/api/activationPatchingApi.ts | 7 +- workbench/_web/src/lib/api/chartApi.ts | 16 +- workbench/_web/src/lib/api/deployApi.ts | 111 ++---- workbench/_web/src/lib/api/jlensApi.ts | 9 +- workbench/_web/src/lib/api/lensApi.ts | 9 +- workbench/_web/src/lib/api/modelsApi.ts | 14 +- workbench/_web/src/lib/api/patchLensApi.ts | 12 +- workbench/_web/src/lib/config.ts | 38 +- workbench/_web/src/lib/runAndStream.ts | 143 ++++++++ workbench/_web/src/lib/startAndPoll.ts | 107 ------ .../_web/src/stores/useModelDeployment.ts | 21 +- workbench/_web/src/types/deployment.ts | 1 - 20 files changed, 683 insertions(+), 779 deletions(-) create mode 100644 workbench/_api/sse.py create mode 100644 workbench/_web/src/lib/runAndStream.ts delete mode 100644 workbench/_web/src/lib/startAndPoll.ts diff --git a/workbench/_api/routes/activation_patching.py b/workbench/_api/routes/activation_patching.py index 4eb70297..321d13cc 100644 --- a/workbench/_api/routes/activation_patching.py +++ b/workbench/_api/routes/activation_patching.py @@ -1,19 +1,17 @@ -from fastapi import APIRouter, Request, Depends from typing import List, Union -from pydantic import BaseModel -from ..data_models import NDIFResponse +from fastapi import APIRouter, Depends +from pydantic import BaseModel +from nnsightful.tools.activation_patching import activation_patching -from ..state import AppState from ..auth import require_user_email -from ..state import get_state - -from nnsightful.types import ActivationPatchingData -from nnsightful.tools.activation_patching import activation_patching +from ..sse import stream_tool +from ..state import AppState, get_state router = APIRouter() + class ActivationPatchingRequest(BaseModel): model_name: str src_prompt: str @@ -23,48 +21,21 @@ class ActivationPatchingRequest(BaseModel): tgt_freeze: List[int] = [] token_ids: List[int] -class ActivationPatchingResponse(NDIFResponse): - data: ActivationPatchingData | None = None - -@router.post("/start", response_model=ActivationPatchingResponse) -async def start_activation_patching( +@router.post("/run") +async def run_activation_patching( request: ActivationPatchingRequest, state: AppState = Depends(get_state), user_email: str = Depends(require_user_email), ): - model = state[request.model_name] - backend = state.make_backend(model=model) - - output = activation_patching._run( - model, + """Run activation patching, streaming status until the data lands (see ``sse``).""" + return stream_tool( + state, + activation_patching, + state[request.model_name], request.src_prompt, request.tgt_prompt, request.src_pos, request.tgt_pos, request.tgt_freeze, - remote=state.remote, - backend=backend, - non_blocking=state.remote, - raw=False, ) - - if not backend.blocking: - return {"job_id": output} - - return {"data": activation_patching.to_data_obj(**output)} - - -@router.post("/results/{job_id}", response_model=ActivationPatchingResponse) -async def collect_results( - job_id: str, - request: ActivationPatchingRequest, - state: AppState = Depends(get_state), - user_email: str = Depends(require_user_email), -): - backend = state.make_backend(job_id=job_id) - results = backend()['results'] - - data = activation_patching.to_data_obj(**results) - - return {"data": data} \ No newline at end of file diff --git a/workbench/_api/routes/causal_mediation.py b/workbench/_api/routes/causal_mediation.py index 26fc96b7..a1536201 100644 --- a/workbench/_api/routes/causal_mediation.py +++ b/workbench/_api/routes/causal_mediation.py @@ -4,14 +4,13 @@ import torch from fastapi import APIRouter, Depends, HTTPException +from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field from ..auth import require_user_email -from ..data_models import NDIFResponse +from ..sse import HEADERS, MEDIA_TYPE, stream_backend, stream_value from ..state import AppState, get_state -from nnsightful.types import LogitLensData - router = APIRouter() @@ -27,12 +26,6 @@ class CausalMediationRequest(BaseModel): include_entropy: bool = True -class CausalMediationResponse(NDIFResponse): - """Identical shape to LogitLensResponse so the frontend can reuse the - existing logit-lens transform/renderer.""" - data: LogitLensData | None = None - - def _format_lens( logits: torch.Tensor, tokenizer, @@ -171,20 +164,22 @@ def _run_causal_mediation( logits = torch.cat(per_layer_logits, dim=0).save() if remote and backend is not None: - return {"job_id": backend.job_id} + # Nothing to return: the values land later, on the backend's stream. + return None return {"logits": logits} -@router.post("/start", response_model=CausalMediationResponse) -async def start_causal_mediation( +@router.post("/run") +async def run_causal_mediation( req: CausalMediationRequest, state: AppState = Depends(get_state), user_email: str = Depends(require_user_email), ): + """Patch one residual across prompts and lens the result, streamed (see ``sse``).""" model = state[req.model] _validate_indices(req, model) - backend = state.make_backend(model=model) + backend = state.make_backend(model) raw = _run_causal_mediation( model, @@ -198,53 +193,28 @@ async def start_causal_mediation( backend=backend, ) - if "job_id" in raw: - return {"job_id": raw["job_id"]} - - input_tokens = _decode_input_tokens(model.tokenizer, req.tgt_prompt) - data = _format_lens( - raw["logits"], - tokenizer=model.tokenizer, - model_name=req.model, - input_tokens=input_tokens, - n_layers=model.num_layers, - top_k=req.topk, - include_entropy=req.include_entropy, - ) - return {"data": data} + # Read off the model now rather than when the values land: one connection + # holds the request open, so unlike the old collect step there is no window + # in which NDIF could stop serving this model and leave `state[...]` raising. + tokenizer = model.tokenizer + input_tokens = _decode_input_tokens(tokenizer, req.tgt_prompt) + def process(saves: dict): + return _format_lens( + saves["logits"], + tokenizer=tokenizer, + model_name=req.model, + input_tokens=input_tokens, + n_layers=model.num_layers, + top_k=req.topk, + include_entropy=req.include_entropy, + ) -@router.post("/results/{job_id}", response_model=CausalMediationResponse) -async def collect_causal_mediation( - job_id: str, - req: CausalMediationRequest, - state: AppState = Depends(get_state), - user_email: str = Depends(require_user_email), -): - backend = state.make_backend(job_id=job_id) - results = backend() - - # The model can be deregistered from the catalog (NDIF stopped serving it) - # between /start and /results; state[...] raises KeyError in that case. - # Surface a clear 503 instead of an opaque 500. - try: - model = state[req.model] - except KeyError: - raise HTTPException( - status_code=503, - detail=f"Model {req.model} is no longer available; please re-run.", + if not state.remote: + return StreamingResponse( + stream_value(process(raw)), media_type=MEDIA_TYPE, headers=HEADERS ) - tokenizer = model.tokenizer - input_tokens = _decode_input_tokens(tokenizer, req.tgt_prompt) - data = _format_lens( - results["logits"], - tokenizer=tokenizer, - model_name=req.model, - input_tokens=input_tokens, - n_layers=model.num_layers, - top_k=req.topk, - include_entropy=req.include_entropy, + return StreamingResponse( + stream_backend(backend, process), media_type=MEDIA_TYPE, headers=HEADERS ) - - return {"data": data} diff --git a/workbench/_api/routes/j_lens.py b/workbench/_api/routes/j_lens.py index cc588ef2..a99d04cf 100644 --- a/workbench/_api/routes/j_lens.py +++ b/workbench/_api/routes/j_lens.py @@ -1,15 +1,15 @@ from fastapi import APIRouter, Depends from pydantic import BaseModel -from ..state import AppState, get_state -from ..auth import require_user_email - -from ..data_models import NDIFResponse -from nnsightful.types import JLensData from nnsightful.tools.j_lens import j_lens +from ..auth import require_user_email +from ..sse import stream_tool +from ..state import AppState, get_state + router = APIRouter() + class JLensRequest(BaseModel): model: str prompt: str @@ -17,38 +17,17 @@ class JLensRequest(BaseModel): include_entropy: bool = True # Whether to include entropy data -class JLensResponse(NDIFResponse): - data: JLensData | None = None - - -@router.post("/start", response_model=JLensResponse) -async def start_j_lens( +@router.post("/run") +async def run_j_lens( req: JLensRequest, state: AppState = Depends(get_state), user_email: str = Depends(require_user_email), ): - model = state[req.model] - backend = state.make_backend(model=model) - - output = j_lens._run(model, req.prompt, remote=state.remote, backend=backend, non_blocking=state.remote, raw=False, top_k=req.topk) - - if not backend.blocking: - return {"job_id": output} - - - return {"data": j_lens.to_data_obj(**output)} - - -@router.post("/results/{job_id}", response_model=JLensResponse) -async def collect_j_lens( - job_id: str, - req: JLensRequest, - state: AppState = Depends(get_state), - user_email: str = Depends(require_user_email), -): - backend = state.make_backend(job_id=job_id) - results = backend()['results'] - - data = j_lens.to_data_obj(**results) - - return {"data": data} + """Run the Jacobian lens, streaming status until the data lands (see ``sse``).""" + return stream_tool( + state, + j_lens, + state[req.model], + req.prompt, + top_k=req.topk, + ) diff --git a/workbench/_api/routes/lens.py b/workbench/_api/routes/lens.py index adf02d69..f5fcc82d 100644 --- a/workbench/_api/routes/lens.py +++ b/workbench/_api/routes/lens.py @@ -3,10 +3,12 @@ import torch as t from fastapi import APIRouter, Depends, HTTPException +from fastapi.responses import StreamingResponse from pydantic import BaseModel from ..auth import require_user_email, user_has_model_access -from ..data_models import NDIFResponse, Token +from ..data_models import Token +from ..sse import HEADERS, MEDIA_TYPE, stream_backend, stream_value from ..state import AppState, get_state ############ LINE ############ @@ -33,11 +35,31 @@ class Line(BaseModel): data: list[Point] -class LensLineResponse(NDIFResponse): - data: list[Line] | None = None +router = APIRouter() -router = APIRouter() +def _stream(state: AppState, user_email: str, *, model: str, run, process): + """Access-check, run, and stream — the shape both lens v1 routes share. + + Kept local rather than shared with ``models._stream_trace``: this path logs no + telemetry, and lens v1 is hidden and on its way out (see CLAUDE.md), so the + two should be able to diverge without one dragging the other. + """ + if state.remote and not user_has_model_access(user_email, model, state): + raise HTTPException( + status_code=403, detail=f"User does not have access to {model}" + ) + + result = run() + + if not state.remote: + return StreamingResponse( + stream_value(process(result)), media_type=MEDIA_TYPE, headers=HEADERS + ) + + return StreamingResponse( + stream_backend(result, process), media_type=MEDIA_TYPE, headers=HEADERS + ) def line(req: LensLineRequest, state: AppState) -> list[t.Tensor]: @@ -63,11 +85,9 @@ def _compute_rank(logits): elif req.stat == LensStatistic.RANK: _compute_func = _compute_rank - with model.trace( - req.prompt, - remote=state.remote, - backend=state.make_backend(model=model), - ) as tracer: + backend = state.make_backend(model) + + with model.trace(req.prompt, remote=state.remote, backend=backend): results = [] for layer in model.model.layers: hidden_BLD = layer.output @@ -84,22 +104,18 @@ def _compute_rank(logits): results.save() if state.remote: - return tracer.backend.job_id - - return results + return backend - -def get_remote_line(user_email: str, job_id: str, state: AppState): - backend = state.make_backend(job_id=job_id) - results = backend() - return results["results"] + return {"results": results} def process_line_results( - results: list[t.Tensor], + saves: dict, req: LensLineRequest, state: AppState, ): + """Turn the trace's saved values into the chart's lines.""" + results = saves["results"] tok = state[req.model].tokenizer target_token_strs = tok.batch_decode(req.token.target_ids) @@ -120,43 +136,21 @@ def process_line_results( return lines -@router.post("/start-line", response_model=LensLineResponse) -async def start_line( - req: LensLineRequest, - state: AppState = Depends(get_state), - user_email: str = Depends(require_user_email) -): - - if state.remote: - if not user_has_model_access(user_email, req.model, state): - message = f"User does not have access to {req.model}" - raise HTTPException(status_code=403, detail=message) - - try: - result = line(req, state) - except Exception as e: - raise e - - if state.remote: - return {"job_id": result} - - return {"data": process_line_results(result, req, state)} - - -@router.post("/results-line/{job_id}", response_model=LensLineResponse) -async def collect_line( - job_id: str, +@router.post("/run-line") +async def run_line( req: LensLineRequest, state: AppState = Depends(get_state), user_email: str = Depends(require_user_email) ): + """Legacy lens v1 line, streamed (see ``sse``).""" + return _stream( + state, + user_email, + model=req.model, + run=lambda: line(req, state), + process=lambda saves: process_line_results(saves, req, state), + ) - try: - results = get_remote_line(user_email, job_id, state) - except Exception as e: - raise e - - return {"data": process_line_results(results, req, state)} ############ GRID ############ @@ -175,10 +169,6 @@ class GridRow(BaseModel): right_axis_label: str | None = None -class GridLensResponse(NDIFResponse): - data: list[GridRow] | None = None - - def heatmap( req: GridLensRequest, state: AppState ) -> tuple[list[t.Tensor], list[t.Tensor]]: @@ -235,11 +225,9 @@ def _compute_entropy(hs_decoded, logits): elif req.stat == LensStatistic.ENTROPY: _compute_func = _compute_entropy - with model.trace( - req.prompt, - remote=state.remote, - backend=state.make_backend(model=model), - ) as tracer: + backend = state.make_backend(model) + + with model.trace(req.prompt, remote=state.remote, backend=backend): hs_decoded = [] for layer in model.model.layers[:-1]: @@ -256,26 +244,18 @@ def _compute_entropy(hs_decoded, logits): pred_ids.save() if state.remote: - return tracer.backend.job_id + return backend - return stats, pred_ids + return {"stats": stats, "pred_ids": pred_ids} -def get_remote_heatmap( - user_email: str, - job_id: str, - state: AppState -) -> tuple[list[t.Tensor], list[t.Tensor]]: - backend = state.make_backend(job_id=job_id) - results = backend() - return results["stats"], results["pred_ids"] def process_grid_results( - stats: list[t.Tensor], - pred_ids: list[t.Tensor], + saves: dict, lens_request: GridLensRequest, state: AppState, ): + stats, pred_ids = saves["stats"], saves["pred_ids"] tok = state[lens_request.model].tokenizer input_strs = tok.batch_decode(tok.encode(lens_request.prompt)) @@ -315,39 +295,17 @@ def process_grid_results( return rows -@router.post("/start-grid", response_model=GridLensResponse) -async def get_grid( +@router.post("/run-grid") +async def run_grid( req: GridLensRequest, state: AppState = Depends(get_state), user_email: str = Depends(require_user_email) ): - if state.remote: - if not user_has_model_access(user_email, req.model, state): - message = f"User does not have access to {req.model}" - raise HTTPException(status_code=403, detail=message) - - try: - result = heatmap(req, state) - except Exception as e: - raise e - - if state.remote: - return {"job_id": result} - - probs, pred_ids = result - return {"data": process_grid_results(probs, pred_ids, req, state)} - - -@router.post("/results-grid/{job_id}", response_model=GridLensResponse) -async def collect_grid( - job_id: str, - lens_request: GridLensRequest, - state: AppState = Depends(get_state), - user_email: str = Depends(require_user_email) -): - try: - probs, pred_ids = get_remote_heatmap(user_email, job_id, state) - except Exception as e: - raise e - - return {"data": process_grid_results(probs, pred_ids, lens_request, state)} + """Legacy lens v1 grid, streamed (see ``sse``).""" + return _stream( + state, + user_email, + model=req.model, + run=lambda: heatmap(req, state), + process=lambda saves: process_grid_results(saves, req, state), + ) diff --git a/workbench/_api/routes/logit_lens.py b/workbench/_api/routes/logit_lens.py index f9279a3b..5cd0bc33 100644 --- a/workbench/_api/routes/logit_lens.py +++ b/workbench/_api/routes/logit_lens.py @@ -1,15 +1,15 @@ from fastapi import APIRouter, Depends from pydantic import BaseModel -from ..state import AppState, get_state -from ..auth import require_user_email - -from ..data_models import NDIFResponse -from nnsightful.types import LogitLensData from nnsightful.tools.logit_lens import logit_lens +from ..auth import require_user_email +from ..sse import stream_tool +from ..state import AppState, get_state + router = APIRouter() + class LogitLensRequest(BaseModel): model: str prompt: str @@ -17,38 +17,18 @@ class LogitLensRequest(BaseModel): include_entropy: bool = True # Whether to include entropy data -class LogitLensResponse(NDIFResponse): - data: LogitLensData | None = None - - -@router.post("/start", response_model=LogitLensResponse) -async def start_logit_lens( +@router.post("/run") +async def run_logit_lens( req: LogitLensRequest, state: AppState = Depends(get_state), user_email: str = Depends(require_user_email), ): - model = state[req.model] - backend = state.make_backend(model=model) - - output = logit_lens._run(model, req.prompt, remote=state.remote, backend=backend, non_blocking=state.remote, raw=False, top_k=req.topk) - - if not backend.blocking: - return {"job_id": output} - - - return {"data": logit_lens.to_data_obj(**output)} - - -@router.post("/results/{job_id}", response_model=LogitLensResponse) -async def collect_logit_lens( - job_id: str, - req: LogitLensRequest, - state: AppState = Depends(get_state), - user_email: str = Depends(require_user_email), -): - backend = state.make_backend(job_id=job_id) - results = backend()['results'] - - data = logit_lens.to_data_obj(**results) - - return {"data": data} \ No newline at end of file + """Run the logit lens, streaming status until the data lands (see ``sse``).""" + return stream_tool( + state, + logit_lens, + state[req.model], + req.prompt, + top_k=req.topk, + include_entropy=req.include_entropy, + ) diff --git a/workbench/_api/routes/models.py b/workbench/_api/routes/models.py index cd97f59a..0d704b36 100644 --- a/workbench/_api/routes/models.py +++ b/workbench/_api/routes/models.py @@ -4,12 +4,14 @@ import requests import torch as t from fastapi import APIRouter, Depends, HTTPException +from fastapi.responses import StreamingResponse from pydantic import BaseModel from nnsightful.tools.j_lens import j_lens from ..auth import get_user_email, require_user_email, user_has_model_access -from ..data_models import NDIFResponse, Token, ModelHeat +from ..data_models import Token, ModelHeat +from ..sse import HEADERS, MEDIA_TYPE, stream_backend, stream_value from ..telemetry import TelemetryClient, RequestStatus from ..state import AppState, get_state @@ -123,23 +125,86 @@ async def get_models( return models +def _stream_trace( + state: AppState, + user_email: str, + *, + model: str, + method: str, + run, + process, +): + """Access-check a trace, run it, and stream it — what both model routes do. + + ``run`` returns the backend to stream (remote) or the saved values themselves + (local); ``process`` turns saved values into the client's payload. Telemetry + brackets the whole thing. + + The access check raises rather than streaming a failure, because nothing has + been sent yet: a 403 is still a 403. Once ``run`` has submitted, every later + failure reaches the client as an `error` frame instead (see ``sse``). + """ + if state.remote and not user_has_model_access(user_email, model, state): + message = f"User does not have access to {model}" + TelemetryClient.log_request( + RequestStatus.ERROR, user_email, method=method, type="NEXT_TOKEN", msg=message, + ) + raise HTTPException(status_code=403, detail=message) + + TelemetryClient.log_request( + RequestStatus.STARTED, user_email, method=method, type="NEXT_TOKEN", + ) + + try: + result = run() + except Exception as error: + TelemetryClient.log_request( + RequestStatus.ERROR, user_email, method=method, type="NEXT_TOKEN", msg=str(error), + ) + raise + + def finish(saves: dict): + data = process(saves) + TelemetryClient.log_request( + RequestStatus.COMPLETE, user_email, method=method, type="NEXT_TOKEN", + ) + return data + + if not state.remote: + return StreamingResponse( + stream_value(finish(result)), media_type=MEDIA_TYPE, headers=HEADERS + ) + + # No job id to log: it belonged to the poll-and-collect flow, and the async + # backend never surfaces one. If telemetry is switched back on and the + # correlation matters, take it from the first status update. + TelemetryClient.log_request( + RequestStatus.READY, user_email, method=method, type="NEXT_TOKEN", + ) + return StreamingResponse( + stream_backend(result, finish), media_type=MEDIA_TYPE, headers=HEADERS + ) + + class LensCompletion(BaseModel): model: str prompt: str token: Token -def prediction( - req: LensCompletion, state: AppState -) -> tuple[t.Tensor, t.Tensor] | str: +def prediction(req: LensCompletion, state: AppState): + """Trace the model for the next-token distribution at the requested position. + + Returns the backend to stream when remote, and the saved values themselves + when local — the two things a route can go on to do. Either way what reaches + ``process_prediction`` is the same dict, keyed as it is saved here, because + that is how NDIF hands the values back. + """ model = state[req.model] idx = req.token.idx + backend = state.make_backend(model) - with model.trace( - req.prompt, - remote=state.remote, - backend=state.make_backend(model=model), - ) as tracer: + with model.trace(req.prompt, remote=state.remote, backend=backend): logits_BLV = model.logits # Get logits for the correct index @@ -151,17 +216,10 @@ def prediction( values_LV = values_LV_indices_LV[0].save() indices_LV = values_LV_indices_LV[1].save() - if state.remote: - return tracer.backend.job_id - - return values_LV, indices_LV + if state.remote: + return backend -def get_remote_prediction( - job_id: str, state: AppState -) -> tuple[t.Tensor, t.Tensor]: - backend = state.make_backend(job_id=job_id) - results = backend() - return results["values_LV"], results["indices_LV"] + return {"values_LV": values_LV, "indices_LV": indices_LV} class Prediction(BaseModel): @@ -171,16 +229,9 @@ class Prediction(BaseModel): texts: list[str] -class PredictionResponse(NDIFResponse): - data: Prediction | None = None - - -def process_prediction( - values_LV: t.Tensor, - indices_LV: t.Tensor, - req: LensCompletion, - state: AppState, -): +def process_prediction(saves: dict, req: LensCompletion, state: AppState): + """Turn the trace's saved values into the client's `Prediction`.""" + values_LV, indices_LV = saves["values_LV"], saves["indices_LV"] tok = state[req.model].tokenizer idxs = [req.token.idx] @@ -202,90 +253,22 @@ def process_prediction( return prediction -@router.post("/start-prediction", response_model=PredictionResponse) -async def start_prediction( - prediction_request: LensCompletion, - state: AppState = Depends(get_state), - user_email: str = Depends(require_user_email) -): - if state.remote: - if not user_has_model_access(user_email, prediction_request.model, state): - message = f"User does not have access to {prediction_request.model}" - TelemetryClient.log_request( - RequestStatus.ERROR, - user_email, - method="PREDICTION", - type="NEXT_TOKEN", - msg=message, - ) - raise HTTPException(status_code=403, detail=message) - - TelemetryClient.log_request( - RequestStatus.STARTED, - user_email, - method="PREDICTION", - type="NEXT_TOKEN", - ) - - try: - result = prediction(prediction_request, state) - except Exception as e: - TelemetryClient.log_request( - RequestStatus.ERROR, - user_email, - method="PREDICTION", - type="NEXT_TOKEN", - msg=str(e), - ) - raise e - - if state.remote: - TelemetryClient.log_request( - RequestStatus.READY, - user_email, - method="PREDICTION", - type="NEXT_TOKEN", - job_id=result - ) - return {"job_id": result} - - values_LV, indices_LV = result - data = process_prediction(values_LV, indices_LV, prediction_request, state) - return {"data": data} - - -@router.post("/results-prediction/{job_id}", response_model=PredictionResponse) -async def results_prediction( - job_id: str, +@router.post("/run-prediction") +async def run_prediction( prediction_request: LensCompletion, state: AppState = Depends(get_state), user_email: str = Depends(require_user_email) ): - - try: - values_LV, indices_LV = get_remote_prediction(job_id, state) - data = process_prediction(values_LV, indices_LV, prediction_request, state) - except Exception as e: - TelemetryClient.log_request( - RequestStatus.ERROR, - user_email, - job_id=job_id, - method="PREDICTION", - type="NEXT_TOKEN", - msg=str(e), - ) - raise e - - TelemetryClient.log_request( - RequestStatus.COMPLETE, + """Next-token distribution at one position, streamed (see ``sse``).""" + return _stream_trace( + state, user_email, - job_id=job_id, + model=prediction_request.model, method="PREDICTION", - type="NEXT_TOKEN", + run=lambda: prediction(prediction_request, state), + process=lambda saves: process_prediction(saves, prediction_request, state), ) - return {"data": data} - class Completion(BaseModel): prompt: str @@ -298,18 +281,20 @@ class Generation(BaseModel): last_token_prediction: Prediction -class GenerationResponse(NDIFResponse): - data: Generation | None = None - - def generate(req: Completion, state: AppState): + """Generate a completion, saving the last step's distribution. + + Returns the backend to stream when remote and the saved values when local, + the same way :func:`prediction` does. + """ model = state[req.model] last_iter = req.max_new_tokens - 1 + backend = state.make_backend(model) with model.generate( req.prompt, max_new_tokens=req.max_new_tokens, remote=state.remote, - backend=state.make_backend(model=model), + backend=backend, ) as tracer: with tracer.iter[last_iter]: @@ -323,26 +308,20 @@ def generate(req: Completion, state: AppState): new_token_ids = model.generator.output[0].save() if state.remote: - return tracer.backend.job_id + return backend - return values_V, indices_V, new_token_ids - - -def get_remote_generate( - job_id: str, state: AppState -) -> tuple[t.Tensor, t.Tensor, t.Tensor]: - backend = state.make_backend(job_id=job_id) - results = backend() - return results["values_V"], results["indices_V"], results["new_token_ids"] + return { + "values_V": values_V, + "indices_V": indices_V, + "new_token_ids": new_token_ids, + } -def process_generation_results( - values_V: t.Tensor, - indices_V: t.Tensor, - new_token_ids: t.Tensor, - req: Completion, - state: AppState, -): +def process_generation_results(saves: dict, req: Completion, state: AppState): + """Turn the trace's saved values into the client's `Generation`.""" + values_V = saves["values_V"] + indices_V = saves["indices_V"] + new_token_ids = saves["new_token_ids"] tok = state[req.model].tokenizer new_token_text = tok.batch_decode(new_token_ids) @@ -372,93 +351,18 @@ def process_generation_results( } -@router.post("/start-generate", response_model=GenerationResponse) -async def start_generate( - req: Completion, - state: AppState = Depends(get_state), - user_email: str = Depends(require_user_email) -): - - if state.remote: - if not user_has_model_access(user_email, req.model, state): - message = f"User does not have access to {req.model}" - TelemetryClient.log_request( - RequestStatus.ERROR, - user_email, - method="GENERATE", - type="NEXT_TOKEN", - msg=message, - ) - raise HTTPException(status_code=403, detail=message) - - TelemetryClient.log_request( - RequestStatus.STARTED, - user_email, - method="GENERATE", - type="NEXT_TOKEN", - ) - - try: - result = generate(req, state) - except Exception as e: - TelemetryClient.log_request( - RequestStatus.ERROR, - user_email, - method="GENERATE", - type="NEXT_TOKEN", - msg=str(e), - ) - raise e - - if state.remote: - TelemetryClient.log_request( - RequestStatus.READY, - user_email, - method="GENERATE", - type="NEXT_TOKEN", - job_id=result - ) - return {"job_id": result} - - else: - values_V, indices_V, new_token_ids = result - - data = process_generation_results( - values_V, indices_V, new_token_ids, req, state - ) - return {"data": data} - - -@router.post("/results-generate/{job_id}", response_model=GenerationResponse) -async def results_generate( - job_id: str, +@router.post("/run-generate") +async def run_generate( req: Completion, state: AppState = Depends(get_state), user_email: str = Depends(require_user_email) ): - - try: - values_V, indices_V, new_token_ids = get_remote_generate(job_id, state) - data = process_generation_results( - values_V, indices_V, new_token_ids, req, state - ) - except Exception as e: - TelemetryClient.log_request( - RequestStatus.ERROR, - user_email, - job_id=job_id, - method="GENERATE", - type="NEXT_TOKEN", - msg=str(e), - ) - raise e - - TelemetryClient.log_request( - RequestStatus.COMPLETE, + """Generate a completion, streamed (see ``sse``).""" + return _stream_trace( + state, user_email, - job_id=job_id, + model=req.model, method="GENERATE", - type="NEXT_TOKEN", + run=lambda: generate(req, state), + process=lambda saves: process_generation_results(saves, req, state), ) - - return {"data": data} diff --git a/workbench/_api/sse.py b/workbench/_api/sse.py new file mode 100644 index 00000000..9c9f72a6 --- /dev/null +++ b/workbench/_api/sse.py @@ -0,0 +1,212 @@ +"""Server-Sent Events helpers shared by the tool routes. + +Every model-touching route is one POST that stays open: the browser gets each +NDIF status update as it lands and then the finished payload, over a single +connection. That replaces a three-legged flow — POST /start for a job id, poll +NDIF directly until COMPLETED, POST /results/{job_id} — and with it the browser's +need to reach NDIF at all. NDIF is now only ever spoken to from this process. + +The event vocabulary is small and every route emits the same one: + + status a raw nnsight ResponseModel, minus `data` (RECEIVED, QUEUED, RUNNING…) + data the finished payload, JSON-encoded. Exactly one, and it ends the stream + error {"error": "..."}. Also terminal + +Once the stream is open a failure has to be an `error` frame rather than an HTTP +status, because the headers went out when the stream opened and the status line +is long gone. A route may still fail the ordinary way *before* it starts +streaming — a 403 for a model the caller cannot use is still a 403 — and the +client handles both; what it must never see is a request that returns 200 and +then goes quiet. + +Local execution (``REMOTE=false``) is streamed too, as a single `data` frame. It +has nothing to report, but giving it the same shape keeps the development mode +off its own path through the UI. +""" + +from __future__ import annotations + +import asyncio +import inspect +import json +from typing import Any, AsyncIterator, Callable, Union + +from fastapi.encoders import jsonable_encoder +from fastapi.responses import StreamingResponse +from nnsight.schema.response import ResponseModel, Status +from pydantic import BaseModel + +MEDIA_TYPE = "text/event-stream" + +# Given the dict of saved values NDIF returns, produce what the client should +# get. Sync or async — a route that has to await something (a tokenizer call, a +# second request) returns the awaitable and `stream_backend` awaits it. +ProcessFn = Callable[[dict], Union[BaseModel, dict, list, Any]] + +# Sent to whatever sits in front of this app. SSE only works if nothing between +# here and the browser buffers the response: nginx (and the ingress in front of +# the preview deployments) buffers proxied responses by default, which holds +# every frame until the stream closes and turns live status into one burst at +# the end. +HEADERS = { + "Cache-Control": "no-cache", + "Connection": "keep-alive", + "X-Accel-Buffering": "no", +} + + +def sse_event(event: str, data: str) -> str: + """Format one SSE frame.""" + return f"event: {event}\ndata: {data}\n\n" + + +def _jsonify(payload: Any) -> str: + """JSON-encode a payload that may be, or contain, pydantic models. + + ``jsonable_encoder`` is what FastAPI applied itself when these routes still + declared a ``response_model``; a frame is written by hand now, so it has to be + applied here. Not just ``model_dump_json`` on the top level: a payload is + often a plain dict or list with models *inside* it — a generation's + ``completion`` is a list of ``Token`` — which plain ``json.dumps`` refuses. + """ + return json.dumps(jsonable_encoder(payload)) + + +# How long a stream may go quiet before it sends a comment frame. +# +# A job can sit QUEUED behind someone else's for minutes, or spend them loading a +# cold 70B, without NDIF having anything new to say — and an idle connection is +# exactly what a proxy reaps. Well under the usual 60s idle timeouts. +HEARTBEAT_SECONDS = 15.0 + + +async def _with_heartbeat(frames: AsyncIterator[str]) -> AsyncIterator[str]: + """Pass frames through, filling any silence with SSE comments. + + A comment (a frame starting ``:``) is defined to be ignored by every SSE + client, so this is invisible to the browser and to `runAndStream`; all it does + is keep bytes moving so nothing in between decides the connection is dead. + """ + iterator = frames.__aiter__() + pending = asyncio.ensure_future(iterator.__anext__()) + try: + while True: + try: + # Shielded: a timeout must not cancel the receive we are waiting + # on, only stop waiting on it for now. + yield await asyncio.wait_for( + asyncio.shield(pending), HEARTBEAT_SECONDS + ) + except asyncio.TimeoutError: + yield ": keepalive\n\n" + continue + except StopAsyncIteration: + return + pending = asyncio.ensure_future(iterator.__anext__()) + finally: + # The client hung up (or we are done): stop waiting on the socket. + pending.cancel() + + +def stream_backend(backend, process: ProcessFn) -> AsyncIterator[str]: + """Drive an ``AsyncRemoteBackend`` and yield SSE frames for what it reports. + + nnsight's async backend yields raw ``ResponseModel`` updates and then, once the + job completes, the downloaded dict of saved values as its final item — so the + type of the yielded object, not a status check, is what says "this is the + result". Deliberately the *raw* stream: it neither renders nnsight's terminal + display nor raises on a server-side error, which is what lets an ERROR become + an `error` frame here rather than a traceback out of a half-written response. + + ``process`` shapes the saved values into the payload the client wants — for a + tool, that is its ``to_data_obj``. + """ + return _with_heartbeat(_backend_frames(backend, process)) + + +async def _backend_frames(backend, process: ProcessFn) -> AsyncIterator[str]: + try: + failure = None + + async for update in backend: + if isinstance(update, ResponseModel): + # Forward the status as-is. `data` is dropped: on COMPLETED it is + # the object-store URL, which is this process's business and often + # signed for a host the browser cannot reach anyway. + if update.status == Status.ERROR: + failure = update.description + yield sse_event("status", update.model_dump_json(exclude={"data"})) + continue + + # Not a status: the saved values, which only arrive after COMPLETED. + result = process(update) + if inspect.isawaitable(result): + result = await result + yield sse_event("data", _jsonify(result)) + + if failure is not None: + yield sse_event("error", json.dumps({"error": failure})) + except Exception as error: + # Includes anything `process` raised. The stream has to close cleanly + # either way, so the exception becomes the last frame. + yield sse_event("error", json.dumps({"error": str(error)})) + + +def stream_tool(state, tool, model, *args: Any, **kwargs: Any) -> StreamingResponse: + """Run an nnsightful tool and return its progress as an SSE response. + + The whole of what a tool route does. Three of them (logit lens, j-lens, + activation patching) differ only in which tool and which arguments, so they + say exactly that and nothing else. + + Remote and local end up at the same event vocabulary, deliberately: local + execution has no status to report, so it yields a single `data` frame and the + client cannot tell the difference. That is the point — ``REMOTE=false`` is a + development mode, and it should not need its own path through the UI. + + The tool's ``to_data_obj`` is what turns NDIF's dict of saved values into the + payload; it runs here, when the values land, rather than in a separate + collect request. + + That dict is keyed by the *name of the variable the tool saved*, which for + every nnsightful tool is ``results`` — so unwrapping that key is what turns + NDIF's reply into the tool's own arguments. It is a real coupling to the + tool's internals, and it is the one the previous collect routes had too + (``backend()["results"]``). + """ + if not state.remote: + return StreamingResponse( + stream_value(tool(model, *args, remote=False, **kwargs)), + media_type=MEDIA_TYPE, + headers=HEADERS, + ) + + backend = state.make_backend(model) + # Submits on the trace's exit and returns; the awaiting happens in the + # stream. `non_blocking` keeps the tool from reaching for a result that, on + # this path, arrives long after the block's frame is gone. + tool._run( + model, + *args, + remote=True, + backend=backend, + non_blocking=True, + raw=False, + **kwargs, + ) + + return StreamingResponse( + stream_backend(backend, lambda saves: tool.to_data_obj(**saves["results"])), + media_type=MEDIA_TYPE, + headers=HEADERS, + ) + + +async def stream_value(value: Any) -> AsyncIterator[str]: + """A one-frame stream: for local execution, which has nothing to report.""" + yield sse_event("data", _jsonify(value)) + + +async def stream_error(message: str) -> AsyncIterator[str]: + """A one-frame stream carrying a failure the route caught before any work.""" + yield sse_event("error", json.dumps({"error": message})) diff --git a/workbench/_api/state.py b/workbench/_api/state.py index ce892306..10651424 100644 --- a/workbench/_api/state.py +++ b/workbench/_api/state.py @@ -7,7 +7,7 @@ from nnsight import CONFIG from nnterp import StandardizedTransformer -from nnsight.intervention.backends.remote import RemoteBackend +from nnsight.intervention.backends.remote import AsyncRemoteBackend from .data_models import ModelHeat @@ -226,26 +226,28 @@ def get_all_model_list(self) -> list[dict]: """ return self._metadata.all_dumped() - def make_backend(self, model: StandardizedTransformer | None = None, job_id: str | None = None): + def make_backend(self, model: StandardizedTransformer): """Create an nnsight backend for the current deployment mode. - Returns a ``RemoteBackend`` when ``self.remote`` is True, otherwise - ``None`` (local execution uses the in-process model directly). + Returns an ``AsyncRemoteBackend`` when ``self.remote`` is True, otherwise + ``None`` — local execution runs the model in this process and has no + backend and nothing to stream. + + The async backend submits the moment the trace block exits, exactly as the + blocking one does; what differs is that it then hands back the status + updates instead of consuming them, so a route can forward each to the + browser (see ``sse.stream_backend``). Nothing here holds a job id: a + request now lives for one HTTP connection rather than being started, + polled and collected across three. Args: - model: Loaded wrapper; its model key is forwarded to NDIF when - starting a new remote job. - job_id: Existing NDIF job ID for polling results. + model: Loaded wrapper; its model key tells NDIF what to run on. """ - if self.remote: - return RemoteBackend( - job_id=job_id, - blocking=False, - model_key=model.to_model_key() if model is not None else None, - ) - else: + if not self.remote: return None + return AsyncRemoteBackend(model.to_model_key()) + def __getitem__(self, model_name: str): """Alias for ``get_model`` — enables ``state[model_name]`` in handlers.""" return self.get_model(model_name) diff --git a/workbench/_web/src/lib/api/activationPatchingApi.ts b/workbench/_web/src/lib/api/activationPatchingApi.ts index 8e875646..b9dd9d51 100644 --- a/workbench/_web/src/lib/api/activationPatchingApi.ts +++ b/workbench/_web/src/lib/api/activationPatchingApi.ts @@ -12,7 +12,7 @@ import { } from "@/types/activationPatching"; import { queryKeys } from "../queryKeys"; import { toast } from "sonner"; -import { startAndPoll } from "../startAndPoll"; +import { runAndStream } from "../runAndStream"; import { createUserHeadersAction } from "@/actions/auth"; /** @@ -42,10 +42,9 @@ const getActivationPatching = async ( token_ids: [], // Backend will use src_pred and clean_pred from results }; - return await startAndPoll( - config.endpoints.startActivationPatching, + return await runAndStream( + config.endpoints.runActivationPatching, apiRequest, - config.endpoints.resultsActivationPatching, headers, ); }; diff --git a/workbench/_web/src/lib/api/chartApi.ts b/workbench/_web/src/lib/api/chartApi.ts index fa32d6a9..2e6e8439 100644 --- a/workbench/_web/src/lib/api/chartApi.ts +++ b/workbench/_web/src/lib/api/chartApi.ts @@ -21,7 +21,7 @@ import { useCapture } from "@/components/providers/CaptureProvider"; import { Line, HeatmapRow, ChartView } from "@/types/charts"; import { queryKeys } from "../queryKeys"; import { toast } from "sonner"; -import { startAndPoll } from "../startAndPoll"; +import { runAndStream } from "../runAndStream"; import { useHeatmapView, useLineView } from "@/components/charts/ViewProvider"; import { createUserHeadersAction } from "@/actions/auth"; @@ -36,12 +36,7 @@ const getLensLine = async (lensRequest: { completion: LensConfigData; chartId: s token: lensRequest.completion.token, }; - return await startAndPoll( - config.endpoints.startLensLine, - lineRequest, - config.endpoints.resultsLensLine, - headers, - ); + return await runAndStream(config.endpoints.runLensLine, lineRequest, headers); }; export const useLensLine = () => { @@ -123,12 +118,7 @@ const getLensGrid = async (lensRequest: { completion: LensConfigData; chartId: s prompt: lensRequest.completion.prompt, }; - return await startAndPoll( - config.endpoints.startLensGrid, - gridRequest, - config.endpoints.resultsLensGrid, - headers, - ); + return await runAndStream(config.endpoints.runLensGrid, gridRequest, headers); }; export const useLensGrid = () => { diff --git a/workbench/_web/src/lib/api/deployApi.ts b/workbench/_web/src/lib/api/deployApi.ts index cb5154bc..811ad7c3 100644 --- a/workbench/_web/src/lib/api/deployApi.ts +++ b/workbench/_web/src/lib/api/deployApi.ts @@ -1,98 +1,41 @@ import config from "@/lib/config"; import { createUserHeadersAction } from "@/actions/auth"; +import { runAndStream } from "@/lib/runAndStream"; /** * Cold-model deployment ("warmup") API — isolated from normal generation/tool - * execution. It reuses the `/models/start-generate` endpoint with a tiny - * throwaway prompt purely to make NDIF deploy the model; the response is never - * surfaced as a generation result or stored as history. + * execution. It reuses the generation endpoint with a tiny throwaway prompt + * purely to make NDIF deploy the model; the response is never surfaced as a + * generation result or stored as history. * - * "Deployed" is signalled when the NDIF job reaches COMPLETED — i.e. the tiny - * warmup generation actually ran end-to-end, which proves the model is loaded - * and serving. (RUNNING is a weaker signal: NDIF can report it at dispatch - * time, before the replica finishes loading the weights, so we don't treat it - * as deployed.) Unlike normal tool requests (`startAndPoll`, 60s hard - * timeout), cold deploys can take minutes, so the poll uses a generous ceiling - * and never fails early on slowness alone. + * "Deployed" means the warmup generation actually returned — i.e. the model ran + * a forward pass end to end, which proves it is loaded and serving. NDIF's + * RUNNING is a weaker signal: it can be reported at dispatch time, before the + * replica has finished loading its weights. + * + * This used to POST for a job id and then poll NDIF from the browser on a + * 20-minute ceiling. It is now one streamed request like every other, so there + * is no ceiling: the connection *is* the wait, and the server keeps it alive + * through the long silences a cold load produces (see `sse.HEARTBEAT_SECONDS`). + * What that costs is that a reload abandons the wait — though not the + * deployment, which NDIF carries on with regardless. */ -// Generous safety ceiling — cold deploys are slow, but we don't want a truly -// stuck job to poll forever. ~20 minutes. -const DEPLOY_POLL_TIMEOUT_MS = 20 * 60 * 1000; -const DEPLOY_POLL_INTERVAL_MS = 2000; - export class DeploymentError extends Error {} -/** Fire the warmup request. Returns the NDIF job id (remote) or null (local, - * where the model is effectively already available). */ -export async function submitWarmup(model: string): Promise { +/** Warm a model up, resolving once it has provably run. */ +export async function deployModel(model: string): Promise { const headers = await createUserHeadersAction(); - const resp = await fetch(config.getApiUrl(config.endpoints.startGenerate), { - method: "POST", - // Match startAndPoll/getModels: carry the oauth2-proxy session cookie - // cross-origin. Without this, the preview env returns a non-job - // response and the warmup silently "succeeds" with no job_id. - credentials: "include", - headers: { "Content-Type": "application/json", ...headers }, - body: JSON.stringify({ model, prompt: "Hello", max_new_tokens: 1 }), - }); - if (!resp.ok) { - throw new DeploymentError(`Failed to start deployment (HTTP ${resp.status})`); - } - const data = (await resp.json()) as { job_id?: string | null; data?: unknown }; - if (data.job_id) return data.job_id; - // A local (non-remote) backend returns a synchronous result with no - // job_id — the model is genuinely available immediately. `!= null` so a - // `{ data: null }` response (no real result) falls through to the throw. - if (data.data != null) return null; - // 200 OK but neither a job id nor a local result means the request never - // reached the model backend (e.g. an auth gateway answered instead). - // Surface it as a failure rather than a false "deployed". - throw new DeploymentError("Deployment did not start (no job id returned)"); -} - -/** - * Poll the NDIF job until the warmup generation has COMPLETED — the point at - * which the model has provably executed a request and is therefore loaded and - * hot. Rejects on NDIF error statuses or the safety-ceiling timeout. - * `onStatus` reports the raw NDIF status for UI. - */ -export async function pollUntilDeployed( - jobId: string, - onStatus?: (status: string) => void, - signal?: AbortSignal, -): Promise { - const startedAt = Date.now(); - while (true) { - if (signal?.aborted) throw new DeploymentError("Deployment cancelled"); - if (Date.now() - startedAt > DEPLOY_POLL_TIMEOUT_MS) { - throw new DeploymentError("Deployment timed out"); - } - - let resp: Response; - try { - resp = await fetch(config.ndifStatusUrl(jobId), { signal }); - } catch { - // fetch rejects (vs. returning non-ok) on network failure — DNS, - // refused connection, etc. Surface it as a status-check problem - // rather than the generic "Deployment failed". - throw new DeploymentError("Couldn't reach NDIF to check deployment status"); - } - if (!resp.ok) throw new DeploymentError("Deployment status check failed"); - const data = (await resp.json()) as { status?: string; description?: string }; - const status = data.status; - if (status) onStatus?.(status); - - // Only COMPLETED proves the model actually ran the warmup forward pass. - // RUNNING is intentionally NOT treated as deployed: NDIF can report it - // while the replica is still loading weights, which would flip the UI - // to "ready" before the model is genuinely servable. - if (status === "COMPLETED") return; - - if (status === "ERROR" || status === "NNSIGHT_ERROR") { - throw new DeploymentError("Deployment failed on the backend"); - } - await new Promise((r) => setTimeout(r, DEPLOY_POLL_INTERVAL_MS)); + try { + await runAndStream( + config.endpoints.runGenerate, + { model, prompt: "Hello", max_new_tokens: 1 }, + headers, + ); + } catch (error) { + throw new DeploymentError( + error instanceof Error ? error.message : "Deployment failed", + ); } } diff --git a/workbench/_web/src/lib/api/jlensApi.ts b/workbench/_web/src/lib/api/jlensApi.ts index ff773a35..684ef833 100644 --- a/workbench/_web/src/lib/api/jlensApi.ts +++ b/workbench/_web/src/lib/api/jlensApi.ts @@ -8,7 +8,7 @@ import { setChartData } from "@/lib/queries/chartQueries"; import { JLensConfigData, JLensData } from "@/types/jlens"; import { queryKeys } from "../queryKeys"; import { toast } from "sonner"; -import { startAndPoll } from "../startAndPoll"; +import { runAndStream } from "../runAndStream"; import { createUserHeadersAction } from "@/actions/auth"; /** @@ -33,12 +33,7 @@ const getJLens = async (lensRequest: JLensRequest): Promise => { include_entropy: lensRequest.completion.includeEntropy ?? true, }; - return await startAndPoll( - config.endpoints.startJLens, - request, - config.endpoints.resultsJLens, - headers, - ); + return await runAndStream(config.endpoints.runJLens, request, headers); }; /** diff --git a/workbench/_web/src/lib/api/lensApi.ts b/workbench/_web/src/lib/api/lensApi.ts index 87eeab6f..4a03cbc4 100644 --- a/workbench/_web/src/lib/api/lensApi.ts +++ b/workbench/_web/src/lib/api/lensApi.ts @@ -8,7 +8,7 @@ import { setChartData } from "@/lib/queries/chartQueries"; import { Lens2ConfigData, Lens2Data } from "@/types/lens2"; import { queryKeys } from "../queryKeys"; import { toast } from "sonner"; -import { startAndPoll } from "../startAndPoll"; +import { runAndStream } from "../runAndStream"; import { createUserHeadersAction } from "@/actions/auth"; /** @@ -33,12 +33,7 @@ const getLens2 = async (lensRequest: Lens2Request): Promise => { include_entropy: lensRequest.completion.includeEntropy ?? true, }; - return await startAndPoll( - config.endpoints.startLens2, - request, - config.endpoints.resultsLens2, - headers, - ); + return await runAndStream(config.endpoints.runLens2, request, headers); }; /** diff --git a/workbench/_web/src/lib/api/modelsApi.ts b/workbench/_web/src/lib/api/modelsApi.ts index 54c9995a..516e84bc 100644 --- a/workbench/_web/src/lib/api/modelsApi.ts +++ b/workbench/_web/src/lib/api/modelsApi.ts @@ -2,7 +2,7 @@ import { useCallback, useEffect } from "react"; import config from "@/lib/config"; import type { LensConfigData } from "@/types/lens"; import type { Model, Token } from "@/types/models"; -import { startAndPoll } from "../startAndPoll"; +import { runAndStream } from "../runAndStream"; import { useMutation, useQuery, useQueryClient } from "@tanstack/react-query"; import { toast } from "sonner"; import { useWorkspace } from "@/stores/useWorkspace"; @@ -19,12 +19,7 @@ interface Prediction { const getPrediction = async (request: LensConfigData): Promise => { const headers = await createUserHeadersAction(); - return await startAndPoll( - config.endpoints.startPrediction, - request, - config.endpoints.resultsPrediction, - headers, - ); + return await runAndStream(config.endpoints.runPrediction, request, headers); }; export const usePrediction = () => { @@ -49,10 +44,9 @@ export interface GenerationResponse { const generate = async (request: Completion): Promise => { const headers = await createUserHeadersAction(); - return await startAndPoll( - config.endpoints.startGenerate, + return await runAndStream( + config.endpoints.runGenerate, request, - config.endpoints.resultsGenerate, headers, ); }; diff --git a/workbench/_web/src/lib/api/patchLensApi.ts b/workbench/_web/src/lib/api/patchLensApi.ts index d27c6766..07f306e2 100644 --- a/workbench/_web/src/lib/api/patchLensApi.ts +++ b/workbench/_web/src/lib/api/patchLensApi.ts @@ -6,7 +6,7 @@ import config from "@/lib/config"; import { useMutation, useQueryClient } from "@tanstack/react-query"; -import { startAndPoll } from "../startAndPoll"; +import { runAndStream } from "../runAndStream"; import { createUserHeadersAction } from "@/actions/auth"; import { setChartData, getChartById } from "@/lib/queries/chartQueries"; import { createLensRun, updateLensRunIntervention } from "@/lib/queries/lensRunQueries"; @@ -82,15 +82,14 @@ const runLogitLens = async ( includeEntropy: boolean, headers: Record, ): Promise => { - return await startAndPoll( - config.endpoints.startLens2, + return await runAndStream( + config.endpoints.runLens2, { model, prompt, topk, include_entropy: includeEntropy, }, - config.endpoints.resultsLens2, headers, ); }; @@ -237,10 +236,9 @@ export const usePatchLensIntervention = () => { include_entropy: includeEntropy, }; - const result = await startAndPoll( - config.endpoints.startCausalMediation, + const result = await runAndStream( + config.endpoints.runCausalMediation, body, - config.endpoints.resultsCausalMediation, headers, ); diff --git a/workbench/_web/src/lib/config.ts b/workbench/_web/src/lib/config.ts index 6db5d0df..678537cf 100644 --- a/workbench/_web/src/lib/config.ts +++ b/workbench/_web/src/lib/config.ts @@ -1,40 +1,26 @@ // Configuration for the application +// Every model-touching endpoint is one streaming POST (see runAndStream.ts). +// There is deliberately no NDIF URL here: the browser talks only to this app's +// own backend, which is the only thing that speaks to NDIF. const config = { backendUrl: process.env.NEXT_PUBLIC_BACKEND_URL || "http://localhost:8000", - ndifUrl: - process.env.NEXT_PUBLIC_LOCAL_NDIF === "true" - ? "http://localhost:5001" - : "https://api.ndif.us", endpoints: { - startLensLine: "/lens/start-line", - resultsLensLine: (jobId: string) => `/lens/results-line/${jobId}`, + // Lens v1 — hidden, do not extend (see CLAUDE.md). + runLensLine: "/lens/run-line", + runLensGrid: "/lens/run-grid", - startLensGrid: "/lens/start-grid", - resultsLensGrid: (jobId: string) => `/lens/results-grid/${jobId}`, + runLens2: "/logit_lens/run", + runJLens: "/j_lens/run", + runCausalMediation: "/causal_mediation/run", + runActivationPatching: "/activation_patching/run", - startLens2: "/logit_lens/start", - resultsLens2: (jobId: string) => `/logit_lens/results/${jobId}`, - - startJLens: "/j_lens/start", - resultsJLens: (jobId: string) => `/j_lens/results/${jobId}`, - - startCausalMediation: "/causal_mediation/start", - resultsCausalMediation: (jobId: string) => `/causal_mediation/results/${jobId}`, - - startActivationPatching: "/activation_patching/start", - resultsActivationPatching: (jobId: string) => `/activation_patching/results/${jobId}`, - - startPrediction: "/models/start-prediction", - resultsPrediction: (jobId: string) => `/models/results-prediction/${jobId}`, - - startGenerate: "/models/start-generate", - resultsGenerate: (jobId: string) => `/models/results-generate/${jobId}`, + runPrediction: "/models/run-prediction", + runGenerate: "/models/run-generate", models: "/models/", }, getApiUrl: (endpoint: string) => `${config.backendUrl}${endpoint}`, - ndifStatusUrl: (jobId: string) => `${config.ndifUrl}/response/${jobId}`, } as const; export default config; diff --git a/workbench/_web/src/lib/runAndStream.ts b/workbench/_web/src/lib/runAndStream.ts new file mode 100644 index 00000000..52617190 --- /dev/null +++ b/workbench/_web/src/lib/runAndStream.ts @@ -0,0 +1,143 @@ +import config from "./config"; +import { useWorkspace } from "@/stores/useWorkspace"; + +/** + * One POST that stays open for the life of a job. + * + * Replaces the three-legged flow this used to take — POST /start for a job id, + * poll NDIF's /response/{id} from the browser until COMPLETED, POST /results/{id} + * — with a single request whose response is a stream of Server-Sent Events. The + * backend forwards each NDIF status as it lands and then the finished payload. + * + * The browser no longer talks to NDIF at all, which is why there is no NDIF URL + * in `config` any more: the only origin this app calls is its own backend. + * + * The tradeoff is that a run now lives and dies with its connection. Polling a + * job id survived a reload; this does not. Runs are seconds to a couple of + * minutes, and the result is written to the workspace by the caller's mutation + * either way, so what is lost is the ability to *rejoin* a run in progress. + */ + +type SSEEvent = { event: string; data: string }; + +/** Parse a fetch body into SSE events: blank-line-separated `event:`/`data:` blocks. */ +async function* parseSSE(body: ReadableStream): AsyncGenerator { + const reader = body.getReader(); + const decoder = new TextDecoder(); + let buffer = ""; + + try { + while (true) { + const { done, value } = await reader.read(); + if (done) break; + buffer += decoder.decode(value, { stream: true }); + + let separator: number; + while ((separator = buffer.indexOf("\n\n")) !== -1) { + const frame = buffer.slice(0, separator); + buffer = buffer.slice(separator + 2); + + let eventName = "message"; + const dataLines: string[] = []; + for (const line of frame.split("\n")) { + if (line.startsWith("event:")) { + eventName = line.slice(6).trim(); + } else if (line.startsWith("data:")) { + // One leading space after the colon is part of the + // framing, not the payload. + dataLines.push(line.slice(5).replace(/^ /, "")); + } + } + if (dataLines.length === 0) continue; + yield { event: eventName, data: dataLines.join("\n") }; + } + } + } finally { + reader.releaseLock(); + } +} + +/** The message from a failed response, preferring FastAPI's `detail`. */ +async function failureMessage(response: Response): Promise { + try { + const body = await response.json(); + if (typeof body?.detail === "string") return body.detail; + } catch { + /* not JSON; fall through to the status line */ + } + return `Request failed: ${response.status} ${response.statusText}`; +} + +/** + * POST to a streaming endpoint and resolve with its final payload. + * + * Pushes every status the backend reports into `useWorkspace.jobStatus`, so the + * header pill tracks the job. Throws on an `error` frame, on a failure before + * the stream opened (a 403 for a model the user can't reach), or if the stream + * ends without delivering data. + */ +export async function runAndStream( + endpoint: string, + body: unknown, + headers?: Record, +): Promise { + const { setJobStatus } = useWorkspace.getState(); + + const response = await fetch(config.getApiUrl(endpoint), { + method: "POST", + // See modelsApi.ts: send oauth2-proxy cookies cross-origin. + credentials: "include", + headers: { + "Content-Type": "application/json", + Accept: "text/event-stream", + ...headers, + }, + body: JSON.stringify(body), + }); + + if (!response.ok || !response.body) { + setJobStatus("Error"); + throw new Error(await failureMessage(response)); + } + + let data: T | null = null; + let failure: string | null = null; + + for await (const frame of parseSSE(response.body)) { + if (frame.event === "status") { + try { + const status = JSON.parse(frame.data); + // QUEUED carries its position in the description, which the pill + // renders — mirroring what the old poll loop displayed. + if (status?.status === "QUEUED") { + const position = status?.description?.match(/\d+/)?.[0]; + setJobStatus(position ? `QUEUED: ${position}` : "QUEUED"); + } else if (status?.status) { + setJobStatus(status.status); + } + } catch { + /* a malformed status frame is not worth failing the run over */ + } + } else if (frame.event === "data") { + data = JSON.parse(frame.data) as T; + setJobStatus("Idle"); + } else if (frame.event === "error") { + try { + failure = (JSON.parse(frame.data) as { error?: string }).error ?? frame.data; + } catch { + failure = frame.data; + } + } + } + + if (failure !== null) { + setJobStatus("Error"); + throw new Error(failure); + } + if (data === null) { + // The connection closed early — a dropped stream, or a proxy that cut it. + setJobStatus("Error"); + throw new Error("The run ended before returning a result"); + } + return data; +} diff --git a/workbench/_web/src/lib/startAndPoll.ts b/workbench/_web/src/lib/startAndPoll.ts deleted file mode 100644 index 22cdf3e9..00000000 --- a/workbench/_web/src/lib/startAndPoll.ts +++ /dev/null @@ -1,107 +0,0 @@ -import config from "./config"; -import { useWorkspace } from "@/stores/useWorkspace"; - -const POLL_TIMEOUT_MS = 300000; -const POLL_INTERVAL_MS = 1000; - -async function awaitNDIFJob(jobId: string): Promise { - const startedAt = Date.now(); - const { setJobStatus } = useWorkspace.getState(); - while (true) { - if (Date.now() - startedAt > POLL_TIMEOUT_MS) { - setJobStatus("timeout"); - throw new Error("Timed out waiting for job to complete"); - } - - const pollResp = await fetch(config.ndifStatusUrl(jobId)); - if (!pollResp.ok) throw new Error("Polling failed"); - const data = await pollResp.json(); - const status = data?.status as string | undefined; - - if (status === "COMPLETED") { - setJobStatus("Idle"); - return; - } - - if (status === "ERROR" || status === "NNSIGHT_ERROR") { - setJobStatus("Error"); - console.error(data); - throw new Error("Job failed"); - } - - if (status === "QUEUED") { - const match = data?.description?.match(/\d+/); - const num = match ? parseInt(match[0], 10) : null; - setJobStatus(num !== null ? `${status}: ${num}` : status); - } else if (status) { - setJobStatus(status); - } - - // For non-terminal statuses, wait and try again - await new Promise((r) => setTimeout(r, POLL_INTERVAL_MS)); - } -} - -type JobStartResponse = { job_id: string | null } & { data?: T } & Record; - -async function startJob( - url: string, - body: unknown, - headers?: Record, -): Promise> { - const response = await fetch(url, { - method: "POST", - // See modelsApi.ts: send oauth2-proxy cookies cross-origin. - credentials: "include", - headers: { - "Content-Type": "application/json", - ...headers, - }, - body: JSON.stringify(body), - }); - if (!response.ok) throw new Error("Failed to start job"); - return await response.json(); -} - -async function fetchResults( - url: string, - body: unknown, - headers?: Record, -): Promise { - const resp = await fetch(url, { - method: "POST", - credentials: "include", - headers: { - "Content-Type": "application/json", - ...headers, - }, - body: JSON.stringify(body), - }); - if (!resp.ok) throw new Error("Failed to fetch results"); - return resp.json() as Promise; -} - -export async function startAndPoll( - startEndpoint: string, - body: unknown, - resultsEndpoint: (jobId: string) => string, - headers?: Record, -): Promise { - const startUrl = config.getApiUrl(startEndpoint); - const response = await startJob(startUrl, body, headers); - const jobId = response?.job_id ?? null; - if (jobId) { - await awaitNDIFJob(jobId); - - const resultsUrl = config.getApiUrl(resultsEndpoint(jobId)); - const results = await fetchResults(resultsUrl, body, headers); - if (results && typeof results === "object" && "data" in results) { - return (results as { data: T }).data; - } - return results as T; - } - if ("data" in response) { - return (response as { data: T | null }).data as T; - } - return response as unknown as T; -} diff --git a/workbench/_web/src/stores/useModelDeployment.ts b/workbench/_web/src/stores/useModelDeployment.ts index 329f2152..a6a6c6ea 100644 --- a/workbench/_web/src/stores/useModelDeployment.ts +++ b/workbench/_web/src/stores/useModelDeployment.ts @@ -3,7 +3,7 @@ import { create } from "zustand"; import { toast } from "sonner"; import type { DeploymentPhase, DeploymentState } from "@/types/deployment"; -import { submitWarmup, pollUntilDeployed, DeploymentError } from "@/lib/api/deployApi"; +import { deployModel, DeploymentError } from "@/lib/api/deployApi"; /** Deploy toasts read a touch lighter than the default — a translucent popover * surface with a subtle blur instead of a fully opaque panel. The blur earns @@ -59,20 +59,13 @@ export const useModelDeployment = create()((set, get) => { }); const run = async (model: string) => { - setPhase(model, { phase: "submitting", error: undefined, jobId: undefined }); + setPhase(model, { phase: "submitting", error: undefined }); try { - const jobId = await submitWarmup(model); - if (!jobId) { - // Local backend (or non-remote): the model is effectively - // available immediately. - setPhase(model, { phase: "ready" }); - toast.success(`${model.split("/").pop()} is now available`, { - style: DEPLOY_TOAST_STYLE, - }); - return; - } - setPhase(model, { phase: "deploying", jobId }); - await pollUntilDeployed(jobId); + // One request that stays open until the warmup generation returns. + // A local (non-remote) backend answers it immediately; a cold remote + // one takes as long as the load does. + setPhase(model, { phase: "deploying" }); + await deployModel(model); setPhase(model, { phase: "ready" }); toast.success(`${model.split("/").pop()} is now available`, { description: "The model is deployed and ready to run.", diff --git a/workbench/_web/src/types/deployment.ts b/workbench/_web/src/types/deployment.ts index 86c0267b..7e96fec7 100644 --- a/workbench/_web/src/types/deployment.ts +++ b/workbench/_web/src/types/deployment.ts @@ -17,7 +17,6 @@ export interface DeploymentState { model: string; phase: DeploymentPhase; /** NDIF job id of the in-flight warmup request, once started. */ - jobId?: string; /** Human-readable error when phase === "error". */ error?: string; } From 2e764caa8d5abb887f88c59539fe93b54d017de2 Mon Sep 17 00:00:00 2001 From: JadenFK Date: Thu, 13 Aug 2026 15:23:58 -0400 Subject: [PATCH 3/5] sse: drop the await that never fires process is called with the dict of saved values and hands back the payload; every one of the seven is a plain def or lambda, so the isawaitable branch was never taken. It came across from the earlier 0.7 attempt at this, and I wrote a comment justifying it -- that a route might await a tokenizer call or a second request -- which described nothing that exists. The type said Union[BaseModel, dict, list, Any], which is Any with decoration. Co-Authored-By: Claude Opus 5 --- workbench/_api/sse.py | 15 +++++---------- 1 file changed, 5 insertions(+), 10 deletions(-) diff --git a/workbench/_api/sse.py b/workbench/_api/sse.py index 9c9f72a6..f7ba577b 100644 --- a/workbench/_api/sse.py +++ b/workbench/_api/sse.py @@ -27,21 +27,19 @@ from __future__ import annotations import asyncio -import inspect import json -from typing import Any, AsyncIterator, Callable, Union +from typing import Any, AsyncIterator, Callable from fastapi.encoders import jsonable_encoder from fastapi.responses import StreamingResponse from nnsight.schema.response import ResponseModel, Status -from pydantic import BaseModel MEDIA_TYPE = "text/event-stream" # Given the dict of saved values NDIF returns, produce what the client should -# get. Sync or async — a route that has to await something (a tokenizer call, a -# second request) returns the awaitable and `stream_backend` awaits it. -ProcessFn = Callable[[dict], Union[BaseModel, dict, list, Any]] +# get. Called on the event loop, so it must not block for long -- every one of +# these is arithmetic over tensors that are already in memory. +ProcessFn = Callable[[dict], Any] # Sent to whatever sits in front of this app. SSE only works if nothing between # here and the browser buffers the response: nginx (and the ingress in front of @@ -139,10 +137,7 @@ async def _backend_frames(backend, process: ProcessFn) -> AsyncIterator[str]: continue # Not a status: the saved values, which only arrive after COMPLETED. - result = process(update) - if inspect.isawaitable(result): - result = await result - yield sse_event("data", _jsonify(result)) + yield sse_event("data", _jsonify(process(update))) if failure is not None: yield sse_event("error", json.dumps({"error": failure})) From 1b20cabb136bac7a7f2e1bc3b7ce6ef14395cd74 Mon Sep 17 00:00:00 2001 From: JadenFK Date: Thu, 13 Aug 2026 15:31:26 -0400 Subject: [PATCH 4/5] sse: no heartbeat Simpler without it, and the one deployment it would have protected is the nginx-ingress preview path (60s proxy-read-timeout, not overridden). Prod does not run behind Modal, so the 300s function ceiling that a heartbeat could not have fixed anyway is moot. If an idle stream does start getting cut, the annotation is the better lever than a comment frame here; the docstring says where. Co-Authored-By: Claude Opus 5 --- workbench/_api/sse.py | 49 ++++--------------------- workbench/_web/src/lib/api/deployApi.ts | 10 ++--- 2 files changed, 12 insertions(+), 47 deletions(-) diff --git a/workbench/_api/sse.py b/workbench/_api/sse.py index f7ba577b..838c8439 100644 --- a/workbench/_api/sse.py +++ b/workbench/_api/sse.py @@ -26,7 +26,6 @@ from __future__ import annotations -import asyncio import json from typing import Any, AsyncIterator, Callable @@ -70,43 +69,7 @@ def _jsonify(payload: Any) -> str: return json.dumps(jsonable_encoder(payload)) -# How long a stream may go quiet before it sends a comment frame. -# -# A job can sit QUEUED behind someone else's for minutes, or spend them loading a -# cold 70B, without NDIF having anything new to say — and an idle connection is -# exactly what a proxy reaps. Well under the usual 60s idle timeouts. -HEARTBEAT_SECONDS = 15.0 - - -async def _with_heartbeat(frames: AsyncIterator[str]) -> AsyncIterator[str]: - """Pass frames through, filling any silence with SSE comments. - - A comment (a frame starting ``:``) is defined to be ignored by every SSE - client, so this is invisible to the browser and to `runAndStream`; all it does - is keep bytes moving so nothing in between decides the connection is dead. - """ - iterator = frames.__aiter__() - pending = asyncio.ensure_future(iterator.__anext__()) - try: - while True: - try: - # Shielded: a timeout must not cancel the receive we are waiting - # on, only stop waiting on it for now. - yield await asyncio.wait_for( - asyncio.shield(pending), HEARTBEAT_SECONDS - ) - except asyncio.TimeoutError: - yield ": keepalive\n\n" - continue - except StopAsyncIteration: - return - pending = asyncio.ensure_future(iterator.__anext__()) - finally: - # The client hung up (or we are done): stop waiting on the socket. - pending.cancel() - - -def stream_backend(backend, process: ProcessFn) -> AsyncIterator[str]: +async def stream_backend(backend, process: ProcessFn) -> AsyncIterator[str]: """Drive an ``AsyncRemoteBackend`` and yield SSE frames for what it reports. nnsight's async backend yields raw ``ResponseModel`` updates and then, once the @@ -118,11 +81,13 @@ def stream_backend(backend, process: ProcessFn) -> AsyncIterator[str]: ``process`` shapes the saved values into the payload the client wants — for a tool, that is its ``to_data_obj``. - """ - return _with_heartbeat(_backend_frames(backend, process)) - -async def _backend_frames(backend, process: ProcessFn) -> AsyncIterator[str]: + Nothing is sent to fill the silences between updates. A job can sit QUEUED or + loading for minutes with nothing to report, and the nginx ingress in front of + the preview deployments will cut a connection idle for 60s + (``proxy-read-timeout``, not overridden in ``deploy/preview/values.yaml``). If + that starts biting, the fix is that annotation or a comment frame here. + """ try: failure = None diff --git a/workbench/_web/src/lib/api/deployApi.ts b/workbench/_web/src/lib/api/deployApi.ts index 811ad7c3..84452fee 100644 --- a/workbench/_web/src/lib/api/deployApi.ts +++ b/workbench/_web/src/lib/api/deployApi.ts @@ -14,11 +14,11 @@ import { runAndStream } from "@/lib/runAndStream"; * replica has finished loading its weights. * * This used to POST for a job id and then poll NDIF from the browser on a - * 20-minute ceiling. It is now one streamed request like every other, so there - * is no ceiling: the connection *is* the wait, and the server keeps it alive - * through the long silences a cold load produces (see `sse.HEARTBEAT_SECONDS`). - * What that costs is that a reload abandons the wait — though not the - * deployment, which NDIF carries on with regardless. + * 20-minute ceiling. It is now one streamed request like every other, so the + * connection *is* the wait. Two things that costs: a reload abandons the wait + * (though not the deployment, which NDIF carries on with), and a cold load can + * go minutes without a status update, which is longer than some proxies leave + * an idle connection alone. */ export class DeploymentError extends Error {} From 06fb6c611f65a39076526c7cd9631862065d99ba Mon Sep 17 00:00:00 2001 From: JadenFK Date: Thu, 13 Aug 2026 16:26:40 -0400 Subject: [PATCH 5/5] sse: one streamer for both shapes, and plainer names Three notes from review. No leading underscores on what this change added: _jsonify -> jsonify, _stream_trace -> stream_trace, and lens's _stream is gone entirely. The underscored names still in these files (_refresh_catalog, _format_lens, _run_causal_mediation) predate it and are left alone. stream_error was never called. It came across from the earlier 0.7 attempt at this, like the await that went in the last commit. And the local-vs-remote branch was written out four times -- in stream_tool and in all three routes that do not use it -- each wrapping the same StreamingResponse with the same media type and headers. That is now one `stream(result, process)`, which dispatches on whether it was handed an AsyncRemoteBackend or the saved values themselves. Dispatching on the object rather than on state.remote keeps it in step with what the trace actually did, since that flag is what decided the shape. Routes no longer import MEDIA_TYPE, HEADERS, StreamingResponse or the two frame generators. The 403 was also written twice, so it moves to auth.require_model_access next to the predicate it wraps; models.py still logs the denial, which is the only thing it did differently. A lens v1 route is now its access check and one line. Same four statuses and the same Paris from hakone; the frontend is untouched. Co-Authored-By: Claude Opus 5 --- workbench/_api/auth.py | 14 +++ workbench/_api/routes/causal_mediation.py | 14 +-- workbench/_api/routes/lens.py | 49 ++------ workbench/_api/routes/models.py | 40 +++---- workbench/_api/sse.py | 136 +++++++++++----------- 5 files changed, 108 insertions(+), 145 deletions(-) diff --git a/workbench/_api/auth.py b/workbench/_api/auth.py index 512462b8..f91ccaf7 100644 --- a/workbench/_api/auth.py +++ b/workbench/_api/auth.py @@ -46,3 +46,17 @@ def user_has_model_access(user_email: str, model_name: str, state: "AppState") - return True + +def require_model_access(state: "AppState", user_email: str, model_name: str) -> None: + """Refuse a caller who cannot use this model, before anything runs. + + A route calls this while it can still fail the ordinary way: once it starts + streaming, the status line is gone and a refusal can only be an `error` frame + (see ``sse``). Local deployments gate nothing -- there is no catalog to be + outside of. + """ + if state.remote and not user_has_model_access(user_email, model_name, state): + raise HTTPException( + status_code=403, detail=f"User does not have access to {model_name}" + ) + diff --git a/workbench/_api/routes/causal_mediation.py b/workbench/_api/routes/causal_mediation.py index a1536201..484ead01 100644 --- a/workbench/_api/routes/causal_mediation.py +++ b/workbench/_api/routes/causal_mediation.py @@ -4,11 +4,10 @@ import torch from fastapi import APIRouter, Depends, HTTPException -from fastapi.responses import StreamingResponse from pydantic import BaseModel, Field from ..auth import require_user_email -from ..sse import HEADERS, MEDIA_TYPE, stream_backend, stream_value +from ..sse import stream from ..state import AppState, get_state router = APIRouter() @@ -210,11 +209,6 @@ def process(saves: dict): include_entropy=req.include_entropy, ) - if not state.remote: - return StreamingResponse( - stream_value(process(raw)), media_type=MEDIA_TYPE, headers=HEADERS - ) - - return StreamingResponse( - stream_backend(backend, process), media_type=MEDIA_TYPE, headers=HEADERS - ) + # `_run_causal_mediation` returns None when remote -- the values come off the + # backend's stream -- and the saved values themselves when local. + return stream(backend if state.remote else raw, process) diff --git a/workbench/_api/routes/lens.py b/workbench/_api/routes/lens.py index f5fcc82d..88bc076f 100644 --- a/workbench/_api/routes/lens.py +++ b/workbench/_api/routes/lens.py @@ -2,13 +2,12 @@ from enum import Enum import torch as t -from fastapi import APIRouter, Depends, HTTPException -from fastapi.responses import StreamingResponse +from fastapi import APIRouter, Depends from pydantic import BaseModel -from ..auth import require_user_email, user_has_model_access +from ..auth import require_model_access, require_user_email from ..data_models import Token -from ..sse import HEADERS, MEDIA_TYPE, stream_backend, stream_value +from ..sse import stream from ..state import AppState, get_state ############ LINE ############ @@ -38,30 +37,6 @@ class Line(BaseModel): router = APIRouter() -def _stream(state: AppState, user_email: str, *, model: str, run, process): - """Access-check, run, and stream — the shape both lens v1 routes share. - - Kept local rather than shared with ``models._stream_trace``: this path logs no - telemetry, and lens v1 is hidden and on its way out (see CLAUDE.md), so the - two should be able to diverge without one dragging the other. - """ - if state.remote and not user_has_model_access(user_email, model, state): - raise HTTPException( - status_code=403, detail=f"User does not have access to {model}" - ) - - result = run() - - if not state.remote: - return StreamingResponse( - stream_value(process(result)), media_type=MEDIA_TYPE, headers=HEADERS - ) - - return StreamingResponse( - stream_backend(result, process), media_type=MEDIA_TYPE, headers=HEADERS - ) - - def line(req: LensLineRequest, state: AppState) -> list[t.Tensor]: model = state[req.model] idx = req.token.idx @@ -143,13 +118,8 @@ async def run_line( user_email: str = Depends(require_user_email) ): """Legacy lens v1 line, streamed (see ``sse``).""" - return _stream( - state, - user_email, - model=req.model, - run=lambda: line(req, state), - process=lambda saves: process_line_results(saves, req, state), - ) + require_model_access(state, user_email, req.model) + return stream(line(req, state), lambda saves: process_line_results(saves, req, state)) ############ GRID ############ @@ -302,10 +272,5 @@ async def run_grid( user_email: str = Depends(require_user_email) ): """Legacy lens v1 grid, streamed (see ``sse``).""" - return _stream( - state, - user_email, - model=req.model, - run=lambda: heatmap(req, state), - process=lambda saves: process_grid_results(saves, req, state), - ) + require_model_access(state, user_email, req.model) + return stream(heatmap(req, state), lambda saves: process_grid_results(saves, req, state)) diff --git a/workbench/_api/routes/models.py b/workbench/_api/routes/models.py index 0d704b36..6265604b 100644 --- a/workbench/_api/routes/models.py +++ b/workbench/_api/routes/models.py @@ -4,14 +4,13 @@ import requests import torch as t from fastapi import APIRouter, Depends, HTTPException -from fastapi.responses import StreamingResponse from pydantic import BaseModel from nnsightful.tools.j_lens import j_lens -from ..auth import get_user_email, require_user_email, user_has_model_access +from ..auth import get_user_email, require_model_access, require_user_email from ..data_models import Token, ModelHeat -from ..sse import HEADERS, MEDIA_TYPE, stream_backend, stream_value +from ..sse import stream from ..telemetry import TelemetryClient, RequestStatus from ..state import AppState, get_state @@ -125,7 +124,7 @@ async def get_models( return models -def _stream_trace( +def stream_trace( state: AppState, user_email: str, *, @@ -144,12 +143,14 @@ def _stream_trace( been sent yet: a 403 is still a 403. Once ``run`` has submitted, every later failure reaches the client as an `error` frame instead (see ``sse``). """ - if state.remote and not user_has_model_access(user_email, model, state): - message = f"User does not have access to {model}" + try: + require_model_access(state, user_email, model) + except HTTPException as denied: TelemetryClient.log_request( - RequestStatus.ERROR, user_email, method=method, type="NEXT_TOKEN", msg=message, + RequestStatus.ERROR, user_email, method=method, type="NEXT_TOKEN", + msg=denied.detail, ) - raise HTTPException(status_code=403, detail=message) + raise TelemetryClient.log_request( RequestStatus.STARTED, user_email, method=method, type="NEXT_TOKEN", @@ -170,20 +171,15 @@ def finish(saves: dict): ) return data - if not state.remote: - return StreamingResponse( - stream_value(finish(result)), media_type=MEDIA_TYPE, headers=HEADERS + if state.remote: + # No job id to log: it belonged to the poll-and-collect flow, and the + # async backend never surfaces one. If telemetry is switched back on and + # the correlation matters, take it from the first status update. + TelemetryClient.log_request( + RequestStatus.READY, user_email, method=method, type="NEXT_TOKEN", ) - # No job id to log: it belonged to the poll-and-collect flow, and the async - # backend never surfaces one. If telemetry is switched back on and the - # correlation matters, take it from the first status update. - TelemetryClient.log_request( - RequestStatus.READY, user_email, method=method, type="NEXT_TOKEN", - ) - return StreamingResponse( - stream_backend(result, finish), media_type=MEDIA_TYPE, headers=HEADERS - ) + return stream(result, finish) class LensCompletion(BaseModel): @@ -260,7 +256,7 @@ async def run_prediction( user_email: str = Depends(require_user_email) ): """Next-token distribution at one position, streamed (see ``sse``).""" - return _stream_trace( + return stream_trace( state, user_email, model=prediction_request.model, @@ -358,7 +354,7 @@ async def run_generate( user_email: str = Depends(require_user_email) ): """Generate a completion, streamed (see ``sse``).""" - return _stream_trace( + return stream_trace( state, user_email, model=req.model, diff --git a/workbench/_api/sse.py b/workbench/_api/sse.py index 838c8439..4f0007dd 100644 --- a/workbench/_api/sse.py +++ b/workbench/_api/sse.py @@ -31,6 +31,7 @@ from fastapi.encoders import jsonable_encoder from fastapi.responses import StreamingResponse +from nnsight.intervention.backends.remote import AsyncRemoteBackend from nnsight.schema.response import ResponseModel, Status MEDIA_TYPE = "text/event-stream" @@ -57,7 +58,7 @@ def sse_event(event: str, data: str) -> str: return f"event: {event}\ndata: {data}\n\n" -def _jsonify(payload: Any) -> str: +def jsonify(payload: Any) -> str: """JSON-encode a payload that may be, or contain, pydantic models. ``jsonable_encoder`` is what FastAPI applied itself when these routes still @@ -69,77 +70,44 @@ def _jsonify(payload: Any) -> str: return json.dumps(jsonable_encoder(payload)) -async def stream_backend(backend, process: ProcessFn) -> AsyncIterator[str]: - """Drive an ``AsyncRemoteBackend`` and yield SSE frames for what it reports. +def stream(result: Any, process: ProcessFn) -> StreamingResponse: + """The SSE response for a run, however that run turned out. - nnsight's async backend yields raw ``ResponseModel`` updates and then, once the - job completes, the downloaded dict of saved values as its final item — so the - type of the yielded object, not a status check, is what says "this is the - result". Deliberately the *raw* stream: it neither renders nnsight's terminal - display nor raises on a server-side error, which is what lets an ERROR become - an `error` frame here rather than a traceback out of a half-written response. - - ``process`` shapes the saved values into the payload the client wants — for a - tool, that is its ``to_data_obj``. + ``result`` is what tracing produced, and it comes in two shapes: an + ``AsyncRemoteBackend`` to follow, or — when the model ran in this process — + the saved values themselves. Dispatching on the object rather than on + ``state.remote`` keeps this in step with what the trace actually did, since + that flag is what decided the shape in the first place. - Nothing is sent to fill the silences between updates. A job can sit QUEUED or - loading for minutes with nothing to report, and the nginx ingress in front of - the preview deployments will cut a connection idle for 60s - (``proxy-read-timeout``, not overridden in ``deploy/preview/values.yaml``). If - that starts biting, the fix is that annotation or a comment frame here. + ``process`` turns saved values into the client's payload and runs either way: + a local run has no status to report, so it is one `data` frame and the client + cannot tell the difference. That is the point — ``REMOTE=false`` is a + development mode, and it should not need its own path through the UI. """ - try: - failure = None - - async for update in backend: - if isinstance(update, ResponseModel): - # Forward the status as-is. `data` is dropped: on COMPLETED it is - # the object-store URL, which is this process's business and often - # signed for a host the browser cannot reach anyway. - if update.status == Status.ERROR: - failure = update.description - yield sse_event("status", update.model_dump_json(exclude={"data"})) - continue - - # Not a status: the saved values, which only arrive after COMPLETED. - yield sse_event("data", _jsonify(process(update))) - - if failure is not None: - yield sse_event("error", json.dumps({"error": failure})) - except Exception as error: - # Includes anything `process` raised. The stream has to close cleanly - # either way, so the exception becomes the last frame. - yield sse_event("error", json.dumps({"error": str(error)})) + frames = ( + backend_frames(result, process) + if isinstance(result, AsyncRemoteBackend) + else value_frames(process(result)) + ) + return StreamingResponse(frames, media_type=MEDIA_TYPE, headers=HEADERS) def stream_tool(state, tool, model, *args: Any, **kwargs: Any) -> StreamingResponse: - """Run an nnsightful tool and return its progress as an SSE response. - - The whole of what a tool route does. Three of them (logit lens, j-lens, - activation patching) differ only in which tool and which arguments, so they - say exactly that and nothing else. + """Run an nnsightful tool and stream it — the whole of what a tool route does. - Remote and local end up at the same event vocabulary, deliberately: local - execution has no status to report, so it yields a single `data` frame and the - client cannot tell the difference. That is the point — ``REMOTE=false`` is a - development mode, and it should not need its own path through the UI. + Three of them (logit lens, j-lens, activation patching) differ only in which + tool and which arguments, so they say exactly that and nothing else. The tool's ``to_data_obj`` is what turns NDIF's dict of saved values into the - payload; it runs here, when the values land, rather than in a separate - collect request. - - That dict is keyed by the *name of the variable the tool saved*, which for - every nnsightful tool is ``results`` — so unwrapping that key is what turns - NDIF's reply into the tool's own arguments. It is a real coupling to the - tool's internals, and it is the one the previous collect routes had too + payload. That dict is keyed by the *name of the variable the tool saved*, + which for every nnsightful tool is ``results`` — so unwrapping that key is + what turns NDIF's reply into the tool's own arguments. It is a real coupling + to the tool's internals, and it is the one the previous collect routes had too (``backend()["results"]``). """ if not state.remote: - return StreamingResponse( - stream_value(tool(model, *args, remote=False, **kwargs)), - media_type=MEDIA_TYPE, - headers=HEADERS, - ) + # The tool has already shaped this one: `__call__` ends in to_data_obj. + return stream(tool(model, *args, remote=False, **kwargs), lambda data: data) backend = state.make_backend(model) # Submits on the trace's exit and returns; the awaiting happens in the @@ -155,18 +123,44 @@ def stream_tool(state, tool, model, *args: Any, **kwargs: Any) -> StreamingRespo **kwargs, ) - return StreamingResponse( - stream_backend(backend, lambda saves: tool.to_data_obj(**saves["results"])), - media_type=MEDIA_TYPE, - headers=HEADERS, - ) + return stream(backend, lambda saves: tool.to_data_obj(**saves["results"])) -async def stream_value(value: Any) -> AsyncIterator[str]: - """A one-frame stream: for local execution, which has nothing to report.""" - yield sse_event("data", _jsonify(value)) +async def backend_frames(backend, process: ProcessFn) -> AsyncIterator[str]: + """Drive an ``AsyncRemoteBackend``, yielding SSE frames for what it reports. + + nnsight's async backend yields raw ``ResponseModel`` updates and then, once the + job completes, the downloaded dict of saved values as its final item — so the + type of the yielded object, not a status check, is what says "this is the + result". Deliberately the *raw* stream: it neither renders nnsight's terminal + display nor raises on a server-side error, which is what lets an ERROR become + an `error` frame here rather than a traceback out of a half-written response. + + Nothing is sent to fill the silences between updates. A job can sit QUEUED or + loading for minutes with nothing to report, and the nginx ingress in front of + the preview deployments will cut a connection idle for 60s + (``proxy-read-timeout``, not overridden in ``deploy/preview/values.yaml``). If + that starts biting, the fix is that annotation or a comment frame here. + """ + try: + async for update in backend: + if not isinstance(update, ResponseModel): + # The saved values, which only arrive after COMPLETED. + yield sse_event("data", jsonify(process(update))) + continue + + # Forward the status as-is. `data` is dropped: on COMPLETED it is the + # object-store URL, which is this process's business and often signed + # for a host the browser cannot reach anyway. + yield sse_event("status", update.model_dump_json(exclude={"data"})) + if update.status == Status.ERROR: + yield sse_event("error", json.dumps({"error": update.description})) + except Exception as error: + # Includes anything `process` raised. The stream has to close cleanly + # either way, so the exception becomes the last frame. + yield sse_event("error", json.dumps({"error": str(error)})) -async def stream_error(message: str) -> AsyncIterator[str]: - """A one-frame stream carrying a failure the route caught before any work.""" - yield sse_event("error", json.dumps({"error": message})) +async def value_frames(value: Any) -> AsyncIterator[str]: + """A one-frame stream: for local execution, which has nothing to report.""" + yield sse_event("data", jsonify(value))