Add /forward_pass route + bundle transformer-explainer in preview - #117
Open
jon-bell wants to merge 18 commits into
Open
Add /forward_pass route + bundle transformer-explainer in preview#117jon-bell wants to merge 18 commits into
jon-bell wants to merge 18 commits into
Conversation
Wires up the new nnsightful forward_pass tool as a workbench API route
and serves the transformer-explainer Svelte SPA from the same Next.js
host so PR preview deploys exercise the whole flow end to end.
Backend:
- New routes/forward_pass.py: /start, /results/{job_id}, /status/{job_id}.
The status endpoint proxies NDIF so the browser never sees NDIF_API_KEY.
- state.py: ModelMetadata gains n_heads, n_kv_heads, d_model, d_head,
vocab_size, positional_kind, arch_kind. Auto-derived from
AutoConfig.from_pretrained() the same way n_layers is. Surfaces on the
existing /models response without further changes.
- main.py: register the new route, allow http://localhost:5173 / :4173
origins in dev mode (vite default + preview), add GZipMiddleware so
the larger forward_pass payloads compress before egress.
Frontend:
- next.config.js: rewrite /transformer-explainer{,/} to
/transformer-explainer/index.html so the bundled SvelteKit SPA's
entrypoint resolves via Next.js public/ serving.
- workbench/_web/Dockerfile: new te-builder stage clones
ndif-team/transformer-explainer at TE_REF, runs `npm run build` with
VITE_WORKBENCH_API set to the preview API host, copies the static
output into public/transformer-explainer/.
CI:
- preview-deploy.yml: pass VITE_WORKBENCH_API, VITE_USER_EMAIL, TE_REPO,
and TE_REF as build args to the web image. TE_REF defaults to
forward-pass; overridable via the repo variable
TRANSFORMER_EXPLAINER_REF for future bumps.
Deps:
- pyproject.toml + uv.lock: pin nnsightful at jon-bell/nnsightful@ea17d7b3
(forward-pass branch) to pick up the new forward_pass tool. Will switch
to upstream once merged.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Contributor
|
The latest updates on your projects. Learn more about Vercel for GitHub.
|
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR adds a forward-pass inference API with enriched model metadata, integrates the transformer-explainer Svelte app into the web container, and updates build configuration to support both features end-to-end. ChangesForward Pass API and Transformer-Explainer Integration
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
The te-builder stage set NODE_ENV=production before `npm install`, which makes npm skip devDependencies — including vite, sveltekit, and svelte-check. `npm run build` then ran `vite build` → command not found → exit 127 → preview build failure. Split into install (NODE_ENV unset) then build (NODE_ENV=production) so the production svelte.config.js base-path branch still kicks in for the bundle without starving the install of build tools. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (1)
workbench/_web/Dockerfile (1)
16-17: ⚡ Quick winAllow immutable
TE_REFvalues (commit SHAs), not only branch/tag checkout mode.Line 17 uses
git clone --branch, which is branch/tag-oriented and makes immutable SHA pinning awkward for reproducible previews.Suggested refactor
ARG TE_REPO=https://github.com/ndif-team/transformer-explainer.git ARG TE_REF=forward-pass -RUN git clone --depth 1 --branch ${TE_REF} ${TE_REPO} te +RUN git init te \ + && cd te \ + && git remote add origin ${TE_REPO} \ + && git fetch --depth 1 origin ${TE_REF} \ + && git checkout --detach FETCH_HEAD🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@workbench/_web/Dockerfile` around lines 16 - 17, The Dockerfile currently uses git clone --branch ${TE_REF} ${TE_REPO} te which fails to support immutable commit SHAs; change the RUN to first clone the repo (without --branch) into te, then explicitly fetch/checkout ${TE_REF} so both branch/tag names and commit SHAs work (use git -C te fetch --depth=1 origin ${TE_REF} if needed and git -C te checkout ${TE_REF}); reference TE_REF and TE_REPO and the git clone line when making the change.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@workbench/_api/routes/forward_pass.py`:
- Around line 99-101: In the except block that catches requests.RequestException
(in workbench/_api/routes/forward_pass.py where it logs "forward_pass status
proxy failed for {job_id}: {exc}"), re-raise the HTTPException using exception
chaining so the original traceback is preserved: change the current "raise
HTTPException(status_code=502, detail='NDIF status unreachable')" to "raise
HTTPException(status_code=502, detail='NDIF status unreachable') from exc".
Ensure this change is applied in the same except requests.RequestException as
exc handler that references job_id and logger.
- Around line 65-66: state.make_backend(job_id=job_id) can return None in local
mode, so avoid calling backend() blindly; modify the forward_pass handler to
check if backend is None after backend = state.make_backend(job_id=job_id) and
handle that case (e.g., fetch results from the local state store or return an
explicit error/empty response) instead of invoking backend(); ensure you update
the results assignment (currently results = backend()["results"]) to use the
safe alternative when backend is None and keep using the same symbols
(state.make_backend, backend, results) so the logic is clear and robust in local
mode.
- Line 38: Validate that the requested model key exists in the state before
indexing into it: check if req.model is present in the state mapping (e.g.,
using "if req.model not in state") and handle the missing-model case by
returning or raising the appropriate error/HTTP response instead of letting a
KeyError propagate; update the code around the model assignment (the line "model
= state[req.model]") to perform this guard and use the same error/response
pattern your route uses for other missing resources.
---
Nitpick comments:
In `@workbench/_web/Dockerfile`:
- Around line 16-17: The Dockerfile currently uses git clone --branch ${TE_REF}
${TE_REPO} te which fails to support immutable commit SHAs; change the RUN to
first clone the repo (without --branch) into te, then explicitly fetch/checkout
${TE_REF} so both branch/tag names and commit SHAs work (use git -C te fetch
--depth=1 origin ${TE_REF} if needed and git -C te checkout ${TE_REF});
reference TE_REF and TE_REPO and the git clone line when making the change.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro Plus
Run ID: 2c02d14e-4f08-4d27-88f6-8a8d133b9b94
⛔ Files ignored due to path filters (1)
uv.lockis excluded by!**/*.lock
📒 Files selected for processing (8)
.github/workflows/preview-deploy.ymlpyproject.tomlworkbench/_api/main.pyworkbench/_api/routes/__init__.pyworkbench/_api/routes/forward_pass.pyworkbench/_api/state.pyworkbench/_web/Dockerfileworkbench/_web/next.config.js
| state: AppState = Depends(get_state), | ||
| user_email: str = Depends(require_user_email), | ||
| ): | ||
| model = state[req.model] |
There was a problem hiding this comment.
Add validation for model existence.
Accessing state[req.model] will raise a KeyError if the model name is not loaded. Validate the model exists before accessing it.
🛡️ Proposed fix
+ if req.model not in state.models:
+ raise HTTPException(status_code=404, detail=f"Model '{req.model}' not found")
model = state[req.model]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@workbench/_api/routes/forward_pass.py` at line 38, Validate that the
requested model key exists in the state before indexing into it: check if
req.model is present in the state mapping (e.g., using "if req.model not in
state") and handle the missing-model case by returning or raising the
appropriate error/HTTP response instead of letting a KeyError propagate; update
the code around the model assignment (the line "model = state[req.model]") to
perform this guard and use the same error/response pattern your route uses for
other missing resources.
Comment on lines
+65
to
+66
| backend = state.make_backend(job_id=job_id) | ||
| results = backend()["results"] |
There was a problem hiding this comment.
Handle backend being None in local mode.
In local mode, state.make_backend(job_id=job_id) returns None (line 186 in state.py), so calling backend() will raise TypeError: 'NoneType' object is not callable.
🛡️ Proposed fix
backend = state.make_backend(job_id=job_id)
+ if backend is None:
+ raise HTTPException(status_code=400, detail="Results endpoint requires remote mode")
results = backend()["results"]📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| backend = state.make_backend(job_id=job_id) | |
| results = backend()["results"] | |
| backend = state.make_backend(job_id=job_id) | |
| if backend is None: | |
| raise HTTPException(status_code=400, detail="Results endpoint requires remote mode") | |
| results = backend()["results"] |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@workbench/_api/routes/forward_pass.py` around lines 65 - 66,
state.make_backend(job_id=job_id) can return None in local mode, so avoid
calling backend() blindly; modify the forward_pass handler to check if backend
is None after backend = state.make_backend(job_id=job_id) and handle that case
(e.g., fetch results from the local state store or return an explicit
error/empty response) instead of invoking backend(); ensure you update the
results assignment (currently results = backend()["results"]) to use the safe
alternative when backend is None and keep using the same symbols
(state.make_backend, backend, results) so the logic is clear and robust in local
mode.
Comment on lines
+99
to
+101
| except requests.RequestException as exc: | ||
| logger.warning(f"forward_pass status proxy failed for {job_id}: {exc}") | ||
| raise HTTPException(status_code=502, detail="NDIF status unreachable") |
There was a problem hiding this comment.
Use raise ... from exc for better exception context.
When re-raising within an exception handler, use raise ... from exc to preserve the original traceback.
♻️ Proposed fix
except requests.RequestException as exc:
logger.warning(f"forward_pass status proxy failed for {job_id}: {exc}")
- raise HTTPException(status_code=502, detail="NDIF status unreachable")
+ raise HTTPException(status_code=502, detail="NDIF status unreachable") from exc🧰 Tools
🪛 Ruff (0.15.15)
[warning] 101-101: Within an except clause, raise exceptions with raise ... from err or raise ... from None to distinguish them from errors in exception handling
(B904)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@workbench/_api/routes/forward_pass.py` around lines 99 - 101, In the except
block that catches requests.RequestException (in
workbench/_api/routes/forward_pass.py where it logs "forward_pass status proxy
failed for {job_id}: {exc}"), re-raise the HTTPException using exception
chaining so the original traceback is preserved: change the current "raise
HTTPException(status_code=502, detail='NDIF status unreachable')" to "raise
HTTPException(status_code=502, detail='NDIF status unreachable') from exc".
Ensure this change is applied in the same except requests.RequestException as
exc handler that references job_id and logger.
Contributor
|
🚀 Preview deployed
|
…mer-explainer/ The previous build emitted /_app/... asset URLs instead of /transformer-explainer/_app/..., so every CSS/JS chunk 404'd when the SPA was served from the workbench /transformer-explainer/ subpath. Pair with the TE fix that switches the base resolution to an explicit BASE_PATH env var (d009601 on ndif-team/transformer-explainer); set it to /transformer-explainer in the te-builder stage as both a build arg and a runtime env so svelte.config.js sees it. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The git clone of ndif-team/transformer-explainer was layer-cached based solely on the ARG values, so a TE branch update never re-ran the clone and the bundle kept the pre-fix base-path bug. ADD the per-ref commit JSON from GitHub's API before the clone — Docker invalidates the layer whenever the URL content changes, which forces a fresh clone (and a fresh `npm install && vite build`) on every TE HEAD move while preserving the cache for unrelated CI runs. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
SvelteKit adapter-static emits relative asset URLs (./_app/immutable/...) which only resolve correctly when the document URL ends in a slash. Previously we *rewrote* the slashless path to serve index.html under the slashless URL, so the browser treated the last segment as a file and hoisted every CSS/JS reference up to host root (/_app/...), 404ing everything. Switch to a 307 redirect that adds the trailing slash; the rewrite is kept for the canonical /transformer-explainer/ form and serves the prerendered HTML. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
|
The latest updates on your projects. Learn more about Argos notifications ↗︎
|
…ault) The TE bundle now emits absolute /transformer-explainer/_app/... URLs (9756deb on transformer-explainer@forward-pass), so we no longer need the browser document URL to end in a slash. The previous redirect fought with Next.js's built-in trailingSlash:false handling and produced ERR_TOO_MANY_REDIRECTS. Keep the rewrite for both slash and slashless forms. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Pulls in the ModelSelector dropdown z-index + click-outside fix from ndif-team/transformer-explainer@60714c0. The TE_REF cache-bust ADD will re-clone the branch, no other changes required. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Switching from a branch name to a 40-char commit SHA so each preview build is tied to a known TE commit. Bumping the value busts the docker layer cache (the SHA shows up in both the cache-bust ADD URL and the git clone command), which removes any ambiguity about whether the deployed bundle has the latest fixes. Initial pin: a4f3b3e029a8a0c36e47c16094908e7dddbdf508 (60714c0 ModelSelector z-index/click-outside fix + a4f3b3e new Playwright dropdown spec) The git clone path branches on regex: bare SHAs need fetch+checkout because `git clone --branch` doesn't accept SHAs, but human-readable refs (branches/tags) still use the simple `--branch` form for speed. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The deployed bundle was hitting http://localhost:8000 because something in the te-builder cache chain produced a build where the env var hadn't been applied. Two-pronged hardening: 1. Drop the localhost:8000 default on the ARG. Missing build-args now fail the build with a clear message instead of silently producing a broken bundle. 2. Write the URL into .env.production in the cloned TE source before the npm install layer. .env files are loaded by vite at build time AND become part of the source-layer cache key — so changing them does invalidate downstream layers regardless of how the cache is scoped. Belt-and-braces with the existing ENV directives. Also bumps TE_REF to 0d0c3a8 (diagnostic logging of the API base), which gives users a console line that confirms what the bundle is configured with — and a loud warning if the fallback ever fires again. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The previous attempt required a new build-arg (VITE_WORKBENCH_API), but pull_request_target runs the workflow YAML from main — not from the PR HEAD — so my workflow updates on this branch never took effect. The build then failed at the new guard because main's workflow never passed the arg. Reuse NEXT_PUBLIC_BACKEND_URL (which main's workflow already passes for the Next.js build) to derive VITE_WORKBENCH_API. The te-builder still fails loudly if neither is set, and the explicit arg still wins if/when the workflow is updated. Same .env.production / no-default-fallback hardening from df929a6 carries over. This also explains the original mixed-content bug: the deployed previews were running main's workflow YAML, which never passed VITE_WORKBENCH_API, so the Dockerfile's old default `http://localhost:8000` shipped in every bundle. With this change, the same API host the Next.js app already talks to gets baked into the TE bundle as well. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Picks up ndif-team/transformer-explainer@93067a3 — TE's fetch now sends the oauth2_proxy session cookie cross-origin so the API stops 302'ing to the auth flow on /models/. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Picks up jon-bell/nnsightful@75b4b8c — forward_pass tool no longer records a pydantic_core call into the NDIF trace graph, fixing remote runs that were returning "Module pydantic_core._pydantic_core is not whitelisted". Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Picks up jon-bell/nnsightful@5a2b3e7 — _detect_arch now returns a plain dict, so the trace context's closure never captures any Pydantic instances. NDIF's dill-based serializer can pickle the frame cleanly without dragging in pydantic_core. The previous bump (75b4b8c) only fixed the model_dump call but left the ForwardPassArch instance in scope, where it was still being captured silently. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Three coordinated bumps:
- workbench _api/state.py: detect model_type=="gptj" → arch_kind="gptj",
treat presence of cfg.rotary_dim as RoPE (GPT-J's flavor of partial
RoPE doesn't set rope_theta/rope_parameters). The /models response
now correctly tags EleutherAI/gpt-j-6b as {arch_kind:"gptj",
positional_kind:"rope"}.
- nnsightful → 075d074: forward_pass tool gains a GPT-J branch
(out_proj / single ln_1 / partial RoPE via rotate_every_two) and a
multi-GPU device-mismatch fix on the RoPE cos/sin. Pytest now has 17
cases covering GPT-2, Llama (full RoPE), and GPT-J (partial RoPE).
- TE → 22b1402: ModelSelector defaults to GPT-J when available, falls
through to GPT-2 then any allowed model. New purple gptj badge.
Local verification before push:
- nnsightful pytest 17/17 green
- TE playwright 7/7 + 1 skip green (Llama spec skips in local mode,
same as before)
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Picks up ndif-team/transformer-explainer 166bf6d which: * replaces the hardcoded gpt-j > gpt-2 > first-allowed preference in +page.svelte's onMount with a runtime probe loop. The page now iterates /models smallest-first and adopts the first model whose forward pass actually completes. Fixes the case where the previous default (gpt-j-6b) is HOT on NDIF but failing with "Module nnsight.intervention.batching is not whitelisted" — instead of dead-ending the page, the probe advances and lands on a working model (typically Llama-3.1-8B in current NDIF state). * adds a comprehensive Llama-3.1-8B forward-pass E2E spec (24-token prompt, GQA shape, RoPE row-sum invariant, causal-mask zero pattern, K/V storage at n_kv_heads) that runs against real NDIF in ~15s. * generalizes the page-load smoke (renamed gpt-2 → bootstrap) and the whitespace error-states test so neither pins openai-community/gpt2 as a required HOT deployment. Full TE Playwright suite green locally (8 tests, 1.4 min) against the live workbench backend + real NDIF. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Picks up ndif-team/transformer-explainer eec90b6 which adds a long- prompt test case (≥128 tokens, observed ~147) to the Llama-3.1-8B spec. Uses a Buffer-based extraction path to work around V8's max string size (the decoded payload is ~600+ MB at this length). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
* nnsightful 075d074 → 3204cba: forward_pass now ships only `scores` per layer. The masked + softmax views are derived client-side. Halves the wire payload (S=103: 387 MB → 176 MB decoded). * transformer-explainer 22b1402 → 7599b64: adds deriveAttention.ts with the matching client-side derivation, Llama E2E spec updated to drop the Buffer-extraction hack and use plain .json() again. Adds retry-on-transient-NDIF-error to both the page-load bootstrap loop and the test helper, since NDIF's "is not whitelisted" race is firing for ~80% of jobs against pinned Llama deployments today. Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Main-side changes since we branched: * Model catalog + LRU eviction rewrite (Patch Lens work): AppState split ModelMetadata / fetch_metadata into a new workbench/_api/metadata.py module and rebuilt AppState around a catalog / pinned set / _active_models LRU. Our forward-pass changes to state.py (arch_kind/positional_kind derivation and the extra ModelMetadata fields) were dropped in favor of main's structure — the same fields now live on metadata.ModelMetadata and are populated in metadata.fetch_model_metadata. Route/handler behavior is unchanged because state[repo] / state.make_backend still work as before. * CI: build cache moved from GHA cache to GHCR registry-backed cache (#121). Preserved that in preview-deploy.yml alongside our VITE_WORKBENCH_API / TE_REPO / TE_REF build-args for the transformer-explainer bundle. * CI: api host now uses the pr-<n>-api.<domain> sibling-subdomain layout (#122) for wildcard TLS — our TE_REF still references steps.meta.outputs.api_host so it picks up the new format automatically. * Backend: pyproject pins transformers>=5.11.0 for OLMo3/Gemma support; we kept our jon-bell/nnsightful@3204cba pin (payload-trim fork) rather than main's AdamBelfki3/nnsightful, since the transformer-explainer bundle depends on the wire format that fork produces. * Routes: kept both forward_pass and causal_mediation registered. Deleted workbench/_api/_metadata_cache.json in this merge commit so next startup re-fetches with the new schema (the checked-in cache predates our arch_kind / positional_kind / d_model fields; loading it gives every model default-zero arch info and the TE runtime probe falls back to alphabetical order). Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Wires up the new nnsightful forward_pass tool as a workbench API route and serves the transformer-explainer Svelte SPA from the same Next.js host so PR preview deploys exercise the whole flow end to end.
Backend:
Frontend:
npm run buildwith VITE_WORKBENCH_API set to the preview API host, copies the static output into public/transformer-explainer/.CI:
Deps:
Summary by CodeRabbit
/transformer-explainerroute/transformer-explainerand/transformer-explainer/load correctlynnsightfuldependency