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/__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/_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..484ead01 100644 --- a/workbench/_api/routes/causal_mediation.py +++ b/workbench/_api/routes/causal_mediation.py @@ -7,11 +7,9 @@ from pydantic import BaseModel, Field from ..auth import require_user_email -from ..data_models import NDIFResponse +from ..sse import stream from ..state import AppState, get_state -from nnsightful.types import LogitLensData - router = APIRouter() @@ -27,12 +25,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 +163,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 +192,23 @@ 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} - - -@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.", - ) + # 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) - 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, - ) + 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, + ) - return {"data": data} + # `_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/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..88bc076f 100644 --- a/workbench/_api/routes/lens.py +++ b/workbench/_api/routes/lens.py @@ -2,11 +2,12 @@ from enum import Enum import torch as t -from fastapi import APIRouter, Depends, HTTPException +from fastapi import APIRouter, Depends from pydantic import BaseModel -from ..auth import require_user_email, user_has_model_access -from ..data_models import NDIFResponse, Token +from ..auth import require_model_access, require_user_email +from ..data_models import Token +from ..sse import stream from ..state import AppState, get_state ############ LINE ############ @@ -33,10 +34,6 @@ class Line(BaseModel): data: list[Point] -class LensLineResponse(NDIFResponse): - data: list[Line] | None = None - - router = APIRouter() @@ -63,11 +60,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 +79,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 +111,16 @@ 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``).""" + require_model_access(state, user_email, req.model) + return stream(line(req, state), 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 +139,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 +195,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 +214,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 +265,12 @@ 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``).""" + 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/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..6265604b 100644 --- a/workbench/_api/routes/models.py +++ b/workbench/_api/routes/models.py @@ -8,8 +8,9 @@ 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 ..auth import get_user_email, require_model_access, require_user_email +from ..data_models import Token, ModelHeat +from ..sse import stream from ..telemetry import TelemetryClient, RequestStatus from ..state import AppState, get_state @@ -123,23 +124,83 @@ 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``). + """ + 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=denied.detail, + ) + raise + + 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 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", + ) + + return stream(result, finish) + + 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 +212,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 +225,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 +249,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 +277,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 +304,20 @@ def generate(req: Completion, state: AppState): new_token_ids = model.generator.output[0].save() if state.remote: - return tracer.backend.job_id - - return values_V, indices_V, new_token_ids - + return backend -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 +347,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..4f0007dd --- /dev/null +++ b/workbench/_api/sse.py @@ -0,0 +1,166 @@ +"""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 json +from typing import Any, AsyncIterator, Callable + +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" + +# Given the dict of saved values NDIF returns, produce what the client should +# 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 +# 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)) + + +def stream(result: Any, process: ProcessFn) -> StreamingResponse: + """The SSE response for a run, however that run turned out. + + ``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. + + ``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. + """ + 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 stream it — 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. + + The tool's ``to_data_obj`` is what turns NDIF's dict of saved values into the + 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: + # 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 + # 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 stream(backend, lambda saves: tool.to_data_obj(**saves["results"])) + + +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 value_frames(value: Any) -> AsyncIterator[str]: + """A one-frame stream: for local execution, which has nothing to report.""" + yield sse_event("data", jsonify(value)) 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/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" ] } 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..84452fee 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 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. */ -// 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; }