diff --git a/src/prime_rl/dashboard/README.md b/src/prime_rl/dashboard/README.md index c072c5dd37..74442e7b1d 100644 --- a/src/prime_rl/dashboard/README.md +++ b/src/prime_rl/dashboard/README.md @@ -15,6 +15,14 @@ cluster head node, a laptop against a mounted outputs dir: uv sync --extra dashboard && uv run dashboard [output_dir ...] ``` +The trace viewer's **Messages** mode keeps structured `message.content` and +`trace.tools` visibly separate. **Rendered** decodes each selected branch's +recorded post-renderer `token_ids` as one sequence, retaining special tokens; +it never reconstructs a chat template. The recorded IDs remain the source of +truth, and the viewer reports when IDs, the renderer model, or its tokenizer +are unavailable. Text, advantage, logprob, mask, and content signals apply to +both views; Rendered with Text uses the exact full-sequence decode. + This package is fully AI-generated and maintained by agents - it is not meant to be read or edited by humans. Change it by asking an agent, and verify through the browser smoke tests. The integration suite covers it end to end (`tests/integration/dashboard_smoke.py` runs after every integration test). diff --git a/src/prime_rl/dashboard/server.py b/src/prime_rl/dashboard/server.py index 3ea8c9f1a1..d2df83f1d2 100755 --- a/src/prime_rl/dashboard/server.py +++ b/src/prime_rl/dashboard/server.py @@ -767,6 +767,56 @@ def decode_pieces(model: str, ids: list[int]) -> list[str] | None: return pieces +def trace_node_paths(trace: dict) -> list[list[int]]: + """Root-to-leaf node indexes in the same order as the trace viewer.""" + nodes = trace.get("nodes") or [] + has_child = {node.get("parent") for node in nodes if isinstance(node, dict) and isinstance(node.get("parent"), int)} + paths = [] + for leaf in (index for index in range(len(nodes)) if index not in has_child): + path = [] + seen = set() + index = leaf + while isinstance(index, int) and 0 <= index < len(nodes) and index not in seen: + seen.add(index) + path.append(index) + parent = nodes[index].get("parent") if isinstance(nodes[index], dict) else None + index = parent if isinstance(parent, int) else None + paths.append(list(reversed(path))) + return paths + + +def rendered_token_text(trace: dict, model: str | None) -> dict: + """Decode recorded post-renderer IDs as full branch sequences.""" + nodes = trace.get("nodes") or [] + paths = trace_node_paths(trace) + if not any(isinstance(node, dict) and node.get("token_ids") for node in nodes): + return {"status": "missing_token_ids", "model": model, "paths": []} + if not model: + return {"status": "missing_model", "model": None, "paths": []} + tokenizer = get_tokenizer(model) + if tokenizer is None: + return {"status": "tokenizer_unavailable", "model": model, "paths": []} + + def decode_path(path: list[int]) -> dict: + ids = [token_id for index in path for token_id in (nodes[index].get("token_ids") or [])] + try: + text = tokenizer.decode(ids, skip_special_tokens=False) + except Exception: + text = None + return {"nodes": path, "token_count": len(ids), "text": text} + + rendered_paths = [decode_path(path) for path in paths] + all_nodes = list(range(len(nodes))) + all_nodes_rendered = decode_path(all_nodes) + status = "ok" if all(path["text"] is not None for path in rendered_paths + [all_nodes_rendered]) else "decode_error" + return { + "status": status, + "model": model, + "paths": rendered_paths, + "all_nodes": all_nodes_rendered, + } + + @app.get("/api/runs/{run}/rollouts/{step}/{kind}/{subset}/series") def episode_series(run: str, step: int, kind: str, subset: str, etag: str | None = None, after: int = 0) -> dict: """Per-episode series over a traces file (x = episode order): reward, shape, and the @@ -797,7 +847,15 @@ def value(s: dict, key: str): @app.get("/api/runs/{run}/rollouts/{step}/{kind}/{subset}/{line}") -def get_episode(run: str, step: int, kind: str, subset: str, line: int, tokens: bool = False) -> dict: +def get_episode( + run: str, + step: int, + kind: str, + subset: str, + line: int, + tokens: bool = False, + rendered: bool = False, +) -> dict: path = traces_path(run, step, kind, subset) offsets = line_offsets(path) if not 0 <= line < len(offsets): @@ -805,17 +863,18 @@ def get_episode(run: str, step: int, kind: str, subset: str, line: int, tokens: with path.open("rb") as f: f.seek(offsets[line]) rec = orjson.loads(f.readline()) - if not tokens: + if not tokens and not rendered: return rec fallback_model = model_name(main_config(get_run_dir(run))[1]) for trace in rec.get("traces") or []: client = ((trace.get("agent") or {}).get("config") or {}).get("client") or {} model = client.get("renderer_model_name") or fallback_model - if not model: - continue - for node in trace.get("nodes") or []: - if node.get("token_ids"): - node["token_strs"] = decode_pieces(model, node["token_ids"]) + if tokens and model: + for node in trace.get("nodes") or []: + if node.get("token_ids"): + node["token_strs"] = decode_pieces(model, node["token_ids"]) + if rendered: + trace["rendered_tokens"] = rendered_token_text(trace, model) return rec diff --git a/src/prime_rl/dashboard/static/app.js b/src/prime_rl/dashboard/static/app.js index ce668ae815..89d6c4cf48 100644 --- a/src/prime_rl/dashboard/static/app.js +++ b/src/prime_rl/dashboard/static/app.js @@ -45,6 +45,7 @@ const state = { errorsOnly: prefs.traceErrorsOnly ?? false, sort: (prefs.traceSort ?? "line:asc").split(":")[0], order: (prefs.traceSort ?? "line:asc").split(":")[1], + viewMode: prefs.tokenSignal === "rendered" ? "rendered" : (prefs.traceViewMode ?? "messages"), }, }; @@ -1947,6 +1948,8 @@ let currentEpisode = null; let currentLine = null; let currentTraceIdx = 0; let currentBranchIdx = 0; +let episodeOpenVersion = 0; +let episodeEnrichmentVersion = 0; const COPY_SVG = `` + @@ -2051,42 +2054,61 @@ async function modalStep(delta) { } } -function fetchEpisode(line, withTokens) { +function fetchEpisode(line, withTokens, withRendered = false) { const traces = state.traces; - const qs = withTokens ? "?tokens=true" : ""; + const params = new URLSearchParams(); + if (withTokens) params.set("tokens", "true"); + if (withRendered) params.set("rendered", "true"); + const qs = params.size ? `?${params}` : ""; return api(`/api/runs/${encodeURIComponent(state.run)}/rollouts/${traces.step}/${traces.kind}/${traces.subset}/${line}${qs}`); } /* token strings multiply the payload of a big episode, so they are fetched only - while a token signal is selected — the plain view ships the raw record */ + for token signals or the rendered-token view — the plain view ships the raw record */ async function ensureTokens() { - if (!currentEpisode || currentEpisode._hasTokens || !$("#token-signal").value) return; + if (!currentEpisode) return; + const wantsPieces = !!$("#token-signal").value; + const wantsRendered = state.traces.viewMode === "rendered"; + if ((!wantsPieces || currentEpisode._hasTokens) && (!wantsRendered || currentEpisode._hasRendered)) return; const line = currentLine; - const episode = await fetchEpisode(line, true); - if (line !== currentLine) return; - episode._hasTokens = true; + const withTokens = wantsPieces || !!currentEpisode._hasTokens; + const withRendered = wantsRendered || !!currentEpisode._hasRendered; + const requestVersion = ++episodeEnrichmentVersion; + const episode = await fetchEpisode(line, withTokens, withRendered); + if (line !== currentLine || requestVersion !== episodeEnrichmentVersion) return; + episode._hasTokens = withTokens; + episode._hasRendered = withRendered; currentEpisode = episode; } async function openEpisode(line) { + const requestVersion = ++episodeOpenVersion; + episodeEnrichmentVersion++; $("#trace-modal").hidden = false; $("#drawer-backdrop").hidden = false; currentLine = line; + currentEpisode = null; renderModalStep(); renderRolloutList(); $("#tm-messages").innerHTML = `
loading episode…
`; $("#tm-meta").innerHTML = ""; const withTokens = !!$("#token-signal").value; - const episode = await fetchEpisode(line, withTokens); - if (line !== currentLine) return; // user already moved to another rollout + const withRendered = state.traces.viewMode === "rendered"; + const episode = await fetchEpisode(line, withTokens, withRendered); + if (line !== currentLine || requestVersion !== episodeOpenVersion) return; episode._hasTokens = withTokens; + episode._hasRendered = withRendered; currentEpisode = episode; currentTraceIdx = 0; currentBranchIdx = 0; renderEpisode(); + await ensureTokens(); + if (line === currentLine && requestVersion === episodeOpenVersion) renderEpisode(); } function closeDrawer() { + episodeOpenVersion++; + episodeEnrichmentVersion++; $("#trace-modal").hidden = true; $("#drawer-backdrop").hidden = true; currentEpisode = null; @@ -2145,7 +2167,7 @@ function renderTokenNode(node, signal, maxAbsAdv) { const logprobAt = alignedSignal(node, node.logprobs); const advantageAt = alignedSignal(node, node.advantages); const spans = ids.map((id, i) => { - const text = strs ? strs[i] : ` ${id} `; + const text = strs?.[i] ?? ` ${id} `; const logprob = logprobAt(i), advantage = advantageAt(i); let bg = ""; if (signal === "advantage" && advantage != null && maxAbsAdv > 0) { @@ -2205,6 +2227,95 @@ function reasoningBlock(content) { ); } +function normalizedTools(tools) { + if (tools == null) return []; + return Array.isArray(tools) ? tools : [tools]; +} + +function toolParts(tool, index) { + if (!tool || typeof tool !== "object" || Array.isArray(tool)) + return { name: `tool ${index + 1}`, description: "Malformed tool definition", parameters: tool }; + const value = tool.function && typeof tool.function === "object" ? tool.function : tool; + return { + name: typeof value.name === "string" && value.name ? value.name : `tool ${index + 1}`, + description: typeof value.description === "string" ? value.description : "No description recorded.", + parameters: value.parameters ?? value.input_schema ?? value.schema ?? null, + }; +} + +function toolDefinitionsHtml(trace) { + const tools = normalizedTools(trace.tools); + if (!tools.length) return ""; + const names = tools.map((tool, i) => toolParts(tool, i).name); + const body = tools.map((tool, i) => { + const parts = toolParts(tool, i); + const schema = parts.parameters == null ? "No parameters/schema recorded." : + typeof parts.parameters === "string" ? parts.parameters : JSON.stringify(parts.parameters, null, 2); + return ( + `
${esc(parts.name)}` + + `${preview(parts.description, 140)}` + + `` + + `` + + `
${esc(parts.description)}
` + + `
Parameters / JSON schema` + + `
` + + `
${esc(schema)}
` + ); + }).join(""); + return ( + `
Tool definitions` + + `${tools.length} tool${tools.length === 1 ? "" : "s"}` + + `${esc(names.join(", "))}` + + `` + + `` + + body + `
` + ); +} + +function renderedTokensHtml(trace, branches) { + const rendered = trace.rendered_tokens; + const errors = errorBannersHtml(episodeErrors(currentEpisode, trace)); + if (!rendered) return emptyState("rendered text not loaded", "select this view again to load recorded token IDs") + errors; + const signal = $("#token-signal").value; + const path = currentPath(trace, branches); + const tokenCount = path.reduce((count, index) => count + (trace.nodes[index]?.token_ids?.length || 0), 0); + const unavailable = { + missing_token_ids: ["no recorded token IDs", "This trace cannot provide post-renderer text because its nodes have no token_ids."], + missing_model: ["tokenizer model unavailable", "Neither renderer_model_name nor the run model was recorded."], + tokenizer_unavailable: ["tokenizer unavailable", `Could not load the recorded renderer tokenizer${rendered.model ? ` (${rendered.model})` : ""}. Token IDs remain authoritative.`], + decode_error: ["recorded tokens could not be decoded", "The tokenizer was found, but it could not decode this recorded sequence."], + }; + const selected = currentBranchIdx === -1 ? rendered.all_nodes : rendered.paths?.[currentBranchIdx]; + if (signal && tokenCount) { + let maxAbsAdv = 0; + for (const node of trace.nodes || []) + for (const advantage of node.advantages || []) maxAbsAdv = Math.max(maxAbsAdv, Math.abs(advantage)); + const body = path.map((index) => renderTokenNode(trace.nodes[index], signal, maxAbsAdv)).join(""); + return ( + `
Rendered tokens/text` + + `${fmtCompact(tokenCount)} tokens` + + (selected?.text != null ? `${preview(selected.text, 180)}` : ``) + + (selected?.text != null ? `` : "") + + `` + + `` + + `
${body}
` + errors + ); + } + if (selected?.text == null) { + const [title, detail] = unavailable[rendered.status] ?? ["rendered text unavailable", "The recorded token sequence could not be decoded."]; + return emptyState(title, detail) + errors; + } + return ( + `
Rendered tokens/text` + + `${fmtCompact(selected.token_count)} tokens` + + `${preview(selected.text, 180)}` + + `` + + `` + + `` + + `
${esc(selected.text)}
` + errors + ); +} + let entriesObserver = null; function episodeErrors(ep, trace) { @@ -2243,9 +2354,15 @@ function renderMessages(ep, trace, branches) { container.innerHTML = emptyState("no traces", "this episode carries no trace data") + errorsHtml; return; } + if (state.traces.viewMode === "rendered") { + container.innerHTML = renderedTokensHtml(trace, branches); + return; + } const signal = $("#token-signal").value; const path = currentPath(trace, branches); const concatenated = currentBranchIdx === -1; + const toolsHtml = toolDefinitionsHtml(trace); + const systemPosition = path.findIndex((idx) => trace.nodes[idx]?.message?.role === "system"); let maxAbsAdv = 0; for (const node of trace.nodes || []) for (const a of node.advantages || []) maxAbsAdv = Math.max(maxAbsAdv, Math.abs(a)); @@ -2266,7 +2383,7 @@ function renderMessages(ep, trace, branches) { const reasoning = node.message?.reasoning_content ?? node.message?.reasoning; if (reasoning) subs.push(reasoningBlock(reasoning)); const toolCalls = (node.message?.tool_calls || []).map(toolCallHtml); - return ( + const messageHtml = `
` + `${String(i + 1).padStart(2, "0")}` + `${esc(role)}` + @@ -2277,14 +2394,15 @@ function renderMessages(ep, trace, branches) { subs.join("") + (body ? `
${body}
` : "") + toolCalls.join("") + - `
` - ); + ``; + return messageHtml + (i === systemPosition ? toolsHtml : ""); }; // long traces render in chunks as the reader scrolls — a 1MB episode with // hundreds of turns paints the first screen immediately const CHUNK = 30; let rendered = Math.min(path.length, CHUNK); container.innerHTML = + (systemPosition === -1 ? toolsHtml : "") + path.slice(0, rendered).map(entryHtml).join("") + (rendered < path.length ? `
scroll for ${path.length - rendered} more entries
` : "") + errorsHtml; @@ -2449,6 +2567,7 @@ function renderEpisode() { `` : ""; $("#tm-tabs-row").hidden = traceTabs.hidden && branchTabs.hidden; + setActive("#trace-view-mode", "mode", state.traces.viewMode); renderRolloutList(); renderMessages(ep, trace, branches); renderMeta(ep, trace, branches); @@ -2882,6 +3001,15 @@ $("#token-signal").addEventListener("change", async () => { renderEpisode(); savePrefs(); }); +$("#trace-view-mode").addEventListener("click", async (e) => { + const btn = e.target.closest("[data-mode]"); + if (!btn || btn.dataset.mode === state.traces.viewMode) return; + state.traces.viewMode = btn.dataset.mode; + setActive("#trace-view-mode", "mode", state.traces.viewMode); + await ensureTokens(); + renderEpisode(); + savePrefs(); +}); $("#tm-trace-tabs").addEventListener("click", (e) => { const btn = e.target.closest("[data-trace]"); if (btn) { currentTraceIdx = +btn.dataset.trace; currentBranchIdx = 0; renderEpisode(); } @@ -2895,18 +3023,41 @@ $("#tm-list").addEventListener("click", (e) => { if (item) openEpisode(+item.dataset.line); }); $("#tm-collapse").addEventListener("click", () => - document.querySelectorAll("#tm-messages details.entry").forEach((d) => (d.open = false)) + document.querySelectorAll("#tm-messages details").forEach((d) => (d.open = false)) ); $("#tm-expand").addEventListener("click", () => document.querySelectorAll("#tm-messages details").forEach((d) => (d.open = true)) ); $("#tm-messages").addEventListener("click", (e) => { - const btn = e.target.closest("[data-copy]"); + const btn = e.target.closest("[data-copy], [data-copy-tool], [data-copy-schema], [data-copy-tools], [data-copy-rendered]"); if (!btn) return; e.preventDefault(); e.stopPropagation(); - const node = currentEpisode?.traces?.[currentTraceIdx]?.nodes?.[+btn.dataset.copy]; - if (node) copyText(messageText(node.message), btn); + const trace = currentEpisode?.traces?.[currentTraceIdx]; + if (!trace) return; + if (btn.dataset.copy != null) { + const node = trace.nodes?.[+btn.dataset.copy]; + if (node) copyText(messageText(node.message), btn); + return; + } + const tools = normalizedTools(trace.tools); + if (btn.hasAttribute("data-copy-tools")) return copyText(JSON.stringify(tools, null, 2), btn); + if (btn.dataset.copyTool != null) return copyText(JSON.stringify(tools[+btn.dataset.copyTool], null, 2), btn); + if (btn.dataset.copySchema != null) { + const schema = toolParts(tools[+btn.dataset.copySchema], +btn.dataset.copySchema).parameters; + return copyText(typeof schema === "string" ? schema : JSON.stringify(schema, null, 2), btn); + } + if (btn.dataset.copyRendered) { + const rendered = trace.rendered_tokens; + const selected = currentBranchIdx === -1 ? rendered?.all_nodes : rendered?.paths?.[currentBranchIdx]; + if (btn.dataset.copyRendered === "text") { + if (selected?.text != null) copyText(selected.text, btn); + return; + } + const path = currentPath(trace, traceBranches(trace)); + const ids = path.flatMap((index) => trace.nodes?.[index]?.token_ids || []); + return copyText(JSON.stringify(ids), btn); + } }); $("#tm-meta").addEventListener("click", (e) => { const btn = e.target.closest("[data-copytext]"); @@ -2938,6 +3089,7 @@ function savePrefs() { collapsedSections: [...state.metrics.collapsedSections], traceErrorsOnly: state.traces.errorsOnly, traceSort: `${state.traces.sort}:${state.traces.order}`, + traceViewMode: state.traces.viewMode, logView: state.logs.view, logComponents: state.logs.components ? [...state.logs.components] : null, logLevel: state.logs.level, @@ -3000,7 +3152,7 @@ document.addEventListener("visibilitychange", () => { renderLogLevel(); $("#log-search").value = prefs.logSearch ?? ""; $("#config-search").value = prefs.configSearch ?? ""; - $("#token-signal").value = prefs.tokenSignal ?? ""; + $("#token-signal").value = prefs.tokenSignal === "rendered" ? "" : (prefs.tokenSignal ?? ""); for (const sel of ["#run-select", "#trace-env", "#trace-sort", "#tm-env", "#tm-sort", "#attempt-select", "#token-signal"]) dressSelect($(sel)); syncTraceFilterControls(); diff --git a/src/prime_rl/dashboard/static/index.html b/src/prime_rl/dashboard/static/index.html index 02c1ec1def..0460b0dd13 100644 --- a/src/prime_rl/dashboard/static/index.html +++ b/src/prime_rl/dashboard/static/index.html @@ -190,12 +190,16 @@
Trace Viewer
- + + + + + diff --git a/src/prime_rl/dashboard/static/style.css b/src/prime_rl/dashboard/static/style.css index 5fd2534335..e6a6ccb6a3 100644 --- a/src/prime_rl/dashboard/static/style.css +++ b/src/prime_rl/dashboard/static/style.css @@ -788,6 +788,43 @@ details.sub > summary { details.sub .sub-name { font-weight: 500; color: var(--grey-6); } details.sub > .entry-body { border-top: 1px solid var(--hairline); } +/* tool definitions are model context, but deliberately sit outside numbered messages */ +details.tool-definitions, details.rendered-transcript { + border: 1px solid var(--grey-3); + margin-bottom: 12px; + background: rgba(183, 166, 250, 0.025); +} +details.tool-definitions > summary, details.rendered-transcript > summary, +details.tool-definition > summary { + display: flex; + align-items: center; + gap: 10px; + min-width: 0; + padding: 7px 12px; + background: var(--grey-2); +} +details.tool-definitions[open] > summary, details.rendered-transcript[open] > summary, +details.tool-definition[open] > summary { border-bottom: 1px solid var(--hairline); } +details.tool-definitions[open] > summary .entry-chev, +details.rendered-transcript[open] > summary .entry-chev, +details.tool-definition[open] > summary .entry-chev { transform: rotate(90deg); } +.context-label, .tool-def-name { color: var(--chart-2); font-size: 10px; font-weight: 500; letter-spacing: var(--track-caps); text-transform: uppercase; } +details.tool-definition { margin: 8px 12px; border: 1px solid var(--hairline); } +.tool-description { padding: 10px 12px; color: var(--grey-5); font-size: 12px; white-space: pre-wrap; word-break: break-word; } +.schema-head { display: flex; align-items: center; padding: 6px 12px; border-top: 1px solid var(--hairline); color: var(--grey-4); font-size: 10px; letter-spacing: var(--track-caps); text-transform: uppercase; } +.schema-head .icon-btn { margin-left: auto; } +.tool-schema, .rendered-text { + margin: 0; + padding: 12px 14px; + overflow: auto; + color: var(--grey-5); + font: 400 11px/1.55 var(--font-mono); + tab-size: 2; +} +.tool-schema { max-height: 380px; border-top: 1px solid var(--hairline); white-space: pre; } +.rendered-text { white-space: pre-wrap; overflow-wrap: anywhere; } +#trace-view-mode button { white-space: nowrap; } + /* right pane: overview metadata */ .tm-reward-big { color: var(--positive); } .tm-reward-big.neg { color: var(--negative); }