From ff521aaedbf5e4945d8a2b6f22a26f0eab394ce9 Mon Sep 17 00:00:00 2001 From: Myles Scolnick Date: Fri, 3 Jan 2025 15:48:15 -0500 Subject: [PATCH] fix: small tracing perf (#3336) - Changes the tracing runs from a map to list for O(1) lookup - Only runs the tracing reducer when the feature flag is enabled --- frontend/src/components/tracing/tracing.tsx | 26 +-- .../src/core/cells/__tests__/runs.test.ts | 57 ++++--- frontend/src/core/cells/runs.ts | 159 ++++++++++-------- .../src/core/websocket/useMarimoWebSocket.tsx | 5 +- 4 files changed, 140 insertions(+), 107 deletions(-) diff --git a/frontend/src/components/tracing/tracing.tsx b/frontend/src/components/tracing/tracing.tsx index 317e1484735..0bbebc43fea 100644 --- a/frontend/src/components/tracing/tracing.tsx +++ b/frontend/src/components/tracing/tracing.tsx @@ -209,17 +209,19 @@ const TraceBlockBody: React.FC<{ const cellIds = useCellIds(); - const chartValues: ChartValues[] = run.cellRuns.map((cellRun) => { - const elapsedTime = cellRun.elapsedTime ?? 0; - return { - cell: cellRun.cellId, - cellNum: cellIds.inOrderIds.indexOf(cellRun.cellId), - startTimestamp: formatChartTime(cellRun.startTime), - endTimestamp: formatChartTime(cellRun.startTime + elapsedTime), - elapsedTime: formatElapsedTime(elapsedTime * 1000), - status: cellRun.status, - }; - }); + const chartValues: ChartValues[] = [...run.cellRuns.values()].map( + (cellRun) => { + const elapsedTime = cellRun.elapsedTime ?? 0; + return { + cell: cellRun.cellId, + cellNum: cellIds.inOrderIds.indexOf(cellRun.cellId), + startTimestamp: formatChartTime(cellRun.startTime), + endTimestamp: formatChartTime(cellRun.startTime + elapsedTime), + elapsedTime: formatElapsedTime(elapsedTime * 1000), + status: cellRun.status, + }; + }, + ); const hiddenInputElementId = `hiddenInputElement-${run.runId}`; const vegaSpec = compile( @@ -301,7 +303,7 @@ const TraceRows = (props: { hidden={true} ref={hiddenInputRef} /> - {run.cellRuns.map((cellRun) => ( + {[...run.cellRuns.values()].map((cellRun) => ( (map: Map | undefined): T { + invariant(map, "Map is undefined"); + return map.values().next().value; +} + describe("RunsState Reducer", () => { let state: RunsState; @@ -53,15 +59,18 @@ describe("RunsState Reducer", () => { expect(nextState.runMap.get(runId)).toEqual({ runId, runStartTime: timestamp, - cellRuns: [ - { + cellRuns: new Map([ + [ cellId, - code: code.slice(0, MAX_CODE_LENGTH), - elapsedTime: 0, - startTime: timestamp, - status: "queued", - }, - ], + { + cellId, + code: code.slice(0, MAX_CODE_LENGTH), + elapsedTime: 0, + startTime: timestamp, + status: "queued", + }, + ], + ]), }); }); @@ -124,8 +133,10 @@ describe("RunsState Reducer", () => { }); expect(updatedState.runIds).toEqual([runId]); - expect(updatedState.runMap.get(runId)?.cellRuns[0].status).toBe("running"); - expect(updatedState.runMap.get(runId)?.cellRuns[0].startTime).toBe( + expect(first(updatedState.runMap.get(runId)?.cellRuns).status).toBe( + "running", + ); + expect(first(updatedState.runMap.get(runId)?.cellRuns).startTime).toBe( runStartTimestamp, ); @@ -143,11 +154,15 @@ describe("RunsState Reducer", () => { }); expect(successState.runIds).toEqual([runId]); - expect(successState.runMap.get(runId)?.cellRuns[0].status).toBe("success"); - expect(successState.runMap.get(runId)?.cellRuns[0].startTime).toBe( + expect(first(successState.runMap.get(runId)?.cellRuns).status).toBe( + "success", + ); + expect(first(successState.runMap.get(runId)?.cellRuns).startTime).toBe( runStartTimestamp, ); - expect(successState.runMap.get(runId)?.cellRuns[0].elapsedTime).toBe(5000); + expect(first(successState.runMap.get(runId)?.cellRuns).elapsedTime).toBe( + 5000, + ); }); it("should limit the number of runs to MAX_RUNS", () => { @@ -191,7 +206,9 @@ describe("RunsState Reducer", () => { }, }); - expect(nextState.runMap.get(runId)?.cellRuns[0].code).toBe(truncatedCode); + expect(first(nextState.runMap.get(runId)?.cellRuns).code).toBe( + truncatedCode, + ); }); it("should update the run status to error when stderr occurs", () => { @@ -216,8 +233,8 @@ describe("RunsState Reducer", () => { }); expect(errorState.runIds).toEqual([runId]); - expect(errorState.runMap.get(runId)?.cellRuns[0].status).toBe("error"); - expect(errorState.runMap.get(runId)?.cellRuns[0].elapsedTime).toBe( + expect(first(errorState.runMap.get(runId)?.cellRuns).status).toBe("error"); + expect(first(errorState.runMap.get(runId)?.cellRuns).elapsedTime).toBe( errorTimestamp - timestamp, ); }); @@ -244,8 +261,8 @@ describe("RunsState Reducer", () => { }); expect(errorState.runIds).toEqual([runId]); - expect(errorState.runMap.get(runId)?.cellRuns[0].status).toBe("error"); - expect(errorState.runMap.get(runId)?.cellRuns[0].elapsedTime).toBe( + expect(first(errorState.runMap.get(runId)?.cellRuns).status).toBe("error"); + expect(first(errorState.runMap.get(runId)?.cellRuns).elapsedTime).toBe( errorTimestamp - timestamp, ); }); @@ -280,7 +297,7 @@ describe("RunsState Reducer", () => { }); expect(finalState.runIds).toEqual([runId]); - expect(finalState.runMap.get(runId)?.cellRuns[0].status).toBe("error"); + expect(first(finalState.runMap.get(runId)?.cellRuns).status).toBe("error"); }); it("should order runs from newest to oldest", () => { @@ -412,7 +429,7 @@ describe("RunsState Reducer", () => { expect(nextState.runIds).toEqual([runId]); expect(nextState.runMap.size).toBe(1); - expect(nextState.runMap.get(runId)?.cellRuns.length).toBe(2); + expect(nextState.runMap.get(runId)?.cellRuns.size).toBe(2); }); }); diff --git a/frontend/src/core/cells/runs.ts b/frontend/src/core/cells/runs.ts index 537b5d26ae7..73611c3e963 100644 --- a/frontend/src/core/cells/runs.ts +++ b/frontend/src/core/cells/runs.ts @@ -15,7 +15,7 @@ export interface CellRun { export interface Run { runId: RunId; - cellRuns: CellRun[]; + cellRuns: Map; runStartTime: number; } @@ -49,91 +49,102 @@ const { if (!runId) { return state; } - let runIds: RunId[]; - let run = state.runMap.get(runId); - if (run) { - runIds = state.runIds; - // Shallow copy the run to avoid mutating the existing run - run = { ...run }; - } else { - // If it is a brand new run and the cell code is "pure markdown", - // we don't want to show the trace since it's not helpful. - // This spams the tracing because we re-render pure markdown on keystrokes. - if (isPureMarkdown(code)) { - return state; - } + const existingRun = state.runMap.get(runId); + // If it is a brand new run and the cell code is "pure markdown", + // we don't want to show the trace since it's not helpful. + // This spams the tracing because we re-render pure markdown on keystrokes. + if (!existingRun && isPureMarkdown(code)) { + return state; + } - run = { - runId: runId, - cellRuns: [], + // We determine if the cell operation errored by looking at the output + const erroredOutput = + cellOperation.output && + (cellOperation.output.channel === "marimo-error" || + cellOperation.output.channel === "stderr"); + + let status: CellRun["status"] = erroredOutput + ? "error" + : cellOperation.status === "queued" + ? "queued" + : cellOperation.status === "running" + ? "running" + : "success"; + + // Create new run if needed + if (!existingRun) { + const newRun: Run = { + runId, + cellRuns: new Map([ + [ + cellOperation.cell_id as CellId, + { + cellId: cellOperation.cell_id as CellId, + code: code.slice(0, MAX_CODE_LENGTH), + elapsedTime: 0, + status: status, + startTime: cellOperation.timestamp, + }, + ], + ]), runStartTime: cellOperation.timestamp, }; - runIds = [runId, ...state.runIds]; + // Manage run history size + const runIds = [runId, ...state.runIds]; + const nextRunMap = new Map(state.runMap); if (runIds.length > MAX_RUNS) { const oldestRunId = runIds.pop(); if (oldestRunId) { - state.runMap.delete(oldestRunId); + nextRunMap.delete(oldestRunId); } } + + nextRunMap.set(runId, newRun); + return { + runIds, + runMap: nextRunMap, + }; } - // We determine if the cell operation errored by looking at the output - const erroredOutput = - cellOperation.output && - (cellOperation.output.channel === "marimo-error" || - cellOperation.output.channel === "stderr"); + // Update existing run + const nextCellRuns = new Map(existingRun.cellRuns); + const existingCellRun = nextCellRuns.get(cellOperation.cell_id as CellId); - const nextRuns: CellRun[] = []; - let found = false; - for (const existingCellRun of run.cellRuns) { - if (existingCellRun.cellId === cellOperation.cell_id) { - const hasErroredPreviously = existingCellRun.status === "error"; - let status: CellRun["status"]; - let startTime = existingCellRun.startTime; - - if (hasErroredPreviously || erroredOutput) { - status = "error"; - } else if (cellOperation.status === "queued") { - status = "queued"; - } else if (cellOperation.status === "running") { - status = "running"; - startTime = cellOperation.timestamp; - } else { - status = "success"; - } + // Early return if nothing changed + if ( + existingCellRun && + !erroredOutput && + cellOperation.status === "queued" + ) { + return state; + } - let elapsedTime: number | undefined = undefined; - if (status === "success" || status === "error") { - elapsedTime = cellOperation.timestamp - existingCellRun.startTime; - } + if (existingCellRun) { + const hasErroredPreviously = existingCellRun.status === "error"; - nextRuns.push({ - ...existingCellRun, - startTime: startTime, - elapsedTime: elapsedTime, - status: status, - }); - found = true; - } else { - nextRuns.push(existingCellRun); - } - } - if (!found) { - let status: CellRun["status"]; - - if (erroredOutput) { - status = "error"; - } else if (cellOperation.status === "queued") { - status = "queued"; - } else if (cellOperation.status === "running") { - status = "running"; - } else { - status = "success"; - } + // Compute new status and timing + status = hasErroredPreviously || erroredOutput ? "error" : status; + + const startTime = + cellOperation.status === "running" + ? cellOperation.timestamp + : existingCellRun.startTime; + + const elapsedTime = + status === "success" || status === "error" + ? cellOperation.timestamp - existingCellRun.startTime + : undefined; - nextRuns.push({ + nextCellRuns.set(cellOperation.cell_id as CellId, { + ...existingCellRun, + startTime, + elapsedTime, + status, + }); + } else { + nextCellRuns.set(cellOperation.cell_id as CellId, { cellId: cellOperation.cell_id as CellId, code: code.slice(0, MAX_CODE_LENGTH), elapsedTime: 0, @@ -142,14 +153,14 @@ const { }); } - run.cellRuns = nextRuns; - const nextRunMap = new Map(state.runMap); - nextRunMap.set(runId, run); + nextRunMap.set(runId, { + ...existingRun, + cellRuns: nextCellRuns, + }); return { ...state, - runIds: runIds, runMap: nextRunMap, }; }, diff --git a/frontend/src/core/websocket/useMarimoWebSocket.tsx b/frontend/src/core/websocket/useMarimoWebSocket.tsx index 6136221e92d..8b1d2c730cf 100644 --- a/frontend/src/core/websocket/useMarimoWebSocket.tsx +++ b/frontend/src/core/websocket/useMarimoWebSocket.tsx @@ -41,6 +41,7 @@ import { capabilitiesAtom } from "../config/capabilities"; import { UI_ELEMENT_REGISTRY } from "../dom/uiregistry"; import { reloadSafe } from "@/utils/reload-safe"; import { useRunsActions } from "../cells/runs"; +import { getFeatureFlag } from "../config/feature-flag"; /** * WebSocket that connects to the Marimo kernel and handles incoming messages. @@ -115,7 +116,9 @@ export function useMarimoWebSocket(opts: { case "cell-op": { handleCellOperation(msg.data, handleCellMessage); const cellData = getNotebook().cellData[msg.data.cell_id as CellId]; - addCellOperation({ cellOperation: msg.data, code: cellData.code }); + if (getFeatureFlag("tracing")) { + addCellOperation({ cellOperation: msg.data, code: cellData.code }); + } return; }