Skip to content

Commit

Permalink
improvement: tracing perf, hide charts and skip markdown (#3282)
Browse files Browse the repository at this point in the history 
Building off @Light2Dark's great work on the tracing panel. Had a few
tweaks for perf and show in the **Labs** section of the settings

cc @Light2Dark for review
  • Loading branch information
@mscolnick
mscolnick authored Dec 23, 2024
1 parent 0116550 commit 0555fa2
Show file tree
Hide file tree
Showing 4 changed files with 287 additions and 78 deletions.
26 changes: 25 additions & 1 deletion frontend/src/components/app-config/user-config-form.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1096,7 +1096,7 @@ export const UserConfigForm: React.FC = () => {
render={({ field }) => (
<div className="flex flex-col gap-y-1">
<FormItem className={formItemClasses}>
<FormLabel className="font-normal">Chat Sidebar</FormLabel>
<FormLabel className="font-normal">Chat sidebar</FormLabel>
<FormControl>
<Checkbox
data-testid="chat-sidebar-checkbox"
Expand All @@ -1112,6 +1112,30 @@ export const UserConfigForm: React.FC = () => {
</div>
)}
/>
<FormField
control={form.control}
name="experimental.tracing"
render={({ field }) => (
<div className="flex flex-col gap-y-1">
<FormItem className={formItemClasses}>
<FormLabel className="font-normal">
Tracing sidebar
</FormLabel>
<FormControl>
<Checkbox
data-testid="tracing-sidebar-checkbox"
checked={field.value === true}
onCheckedChange={field.onChange}
/>
</FormControl>
</FormItem>
<FormDescription>
Enable experimental tracing sidebar to debug cell execution
times and dependencies.
</FormDescription>
</div>
)}
/>
</SettingGroup>
);
}
Expand Down
201 changes: 128 additions & 73 deletions frontend/src/components/tracing/tracing.tsx
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,7 @@ import {
type Run,
useRunsActions,
} from "@/core/cells/runs";
import { useAtomValue } from "jotai";
import { atom, useAtomValue, useSetAtom } from "jotai";
import { CellLink } from "../editor/links/cell-link";
import { formatLogTimestamp } from "@/core/cells/logs";
import { useCellIds } from "@/core/cells/cells";
Expand All @@ -35,9 +35,13 @@ import { ClearButton } from "../buttons/clear-button";
import { cn } from "@/utils/cn";
import { PanelEmptyState } from "../editor/chrome/panels/empty-state";
import { type ResolvedTheme, useTheme } from "@/theme/useTheme";
import useResizeObserver from "use-resize-observer";

const expandedRunsAtom = atom<Map<RunId, boolean>>(new Map<RunId, boolean>());

export const Tracing: React.FC = () => {
const { runIds: newestToOldestRunIds, runMap } = useAtomValue(runsAtom);
const expandedRuns = useAtomValue(expandedRunsAtom);
const { clearRuns } = useRunsActions();

const { theme } = useTheme();
Expand Down Expand Up @@ -82,13 +86,15 @@ export const Tracing: React.FC = () => {
</div>

<div className="flex flex-col gap-3">
{newestToOldestRunIds.map((runId: RunId) => {
{newestToOldestRunIds.map((runId: RunId, index: number) => {
const run = runMap.get(runId);
if (run) {
return (
<TraceBlock
key={run.runId}
run={run}
isExpanded={expandedRuns.get(run.runId)}
isMostRecentRun={index === 0}
chartPosition={chartPosition}
theme={theme}
/>
Expand All @@ -101,24 +107,24 @@ export const Tracing: React.FC = () => {
};

const LazyVega = React.lazy(() =>
import("react-vega").then((m) => ({ default: m.Vega })),
import("react-vega").then((m) => ({ default: m.VegaLite })),
);
interface ChartProps {
className?: string;
width: number;
height: number;
vegaSpec: VisualizationSpec;
signalListeners: SignalListeners;
theme: ResolvedTheme;
}

const Chart: React.FC<ChartProps> = (props: ChartProps) => {
const { ref, width = 300 } = useResizeObserver<HTMLDivElement>();
return (
<div className={props.className}>
<div className={props.className} ref={ref}>
<LazyVega
spec={props.vegaSpec}
theme={props.theme === "dark" ? "dark" : undefined}
width={props.width}
width={width - 50}
height={props.height}
signalListeners={props.signalListeners}
actions={false}
Expand All @@ -134,24 +140,64 @@ interface VegaHoverCellSignal {

const TraceBlock: React.FC<{
run: Run;
/**
* undefined means the user hasn't clicked on this run yet
*/
isExpanded: boolean | undefined;
isMostRecentRun: boolean;
chartPosition: ChartPosition;
theme: ResolvedTheme;
}> = ({ run, chartPosition, theme }) => {
const [collapsed, setCollapsed] = useState(false);
const [hoveredCellId, setHoveredCellId] = useState<CellId | null>();

// To send signals to Vega from React, we bind a hidden input element
const hiddenInputRef = useRef<HTMLInputElement>(null);
const dispatchHoverEvent = (cellId: CellId | null) => {
// dispatch input event to trigger vega's param to update
if (hiddenInputRef.current) {
hiddenInputRef.current.value = String(cellId);
hiddenInputRef.current.dispatchEvent(
new Event("input", { bubbles: true }),
);
}
}> = ({ run, isMostRecentRun, chartPosition, isExpanded, theme }) => {
const setExpandedRuns = useSetAtom(expandedRunsAtom);
// We prefer the user's last click, but if they haven't clicked on this run,
// we expand the most recent run by default, otherwise we collapse it.
isExpanded = isExpanded ?? isMostRecentRun;

const onToggleExpanded = () => {
setExpandedRuns((prev) => {
const newMap = new Map(prev);
newMap.set(run.runId, !isExpanded);
return newMap;
});
};

const Icon = isExpanded ? ChevronDown : ChevronRight;
const chevron = <Icon height={16} className="inline" />;

const traceTitle = (
<span className="text-sm cursor-pointer" onClick={onToggleExpanded}>
Run - {formatLogTimestamp(run.runStartTime)}
{chevron}
</span>
);

if (!isExpanded) {
return (
<div key={run.runId} className="flex flex-col">
<pre className="font-mono font-semibold">{traceTitle}</pre>
</div>
);
}

return (
<TraceBlockBody
key={run.runId}
run={run}
chartPosition={chartPosition}
theme={theme}
title={traceTitle}
/>
);
};

const TraceBlockBody: React.FC<{
run: Run;
chartPosition: ChartPosition;
theme: ResolvedTheme;
title: React.ReactNode;
}> = ({ run, chartPosition, theme, title }) => {
const [hoveredCellId, setHoveredCellId] = useState<CellId | null>();

const handleVegaSignal = {
[VEGA_HOVER_SIGNAL]: (name: string, value: unknown) => {
const signalValue = value as VegaHoverCellSignal;
Expand All @@ -160,19 +206,6 @@ const TraceBlock: React.FC<{
},
};

const Icon = collapsed ? ChevronRight : ChevronDown;
const chevron = <Icon height={16} className="inline" />;

const TraceTitle = (
<span
className="text-sm cursor-pointer"
onClick={() => setCollapsed(!collapsed)}
>
Run - {formatLogTimestamp(run.runStartTime)}
{chevron}
</span>
);

const cellIds = useCellIds();

const chartValues: ChartValues[] = run.cellRuns.map((cellRun) => {
Expand All @@ -197,41 +230,26 @@ const TraceBlock: React.FC<{
),
).spec;

const TraceRows = (
<div className="text-xs mt-0.5 ml-3 flex flex-col gap-0.5">
<input
type="text"
id={hiddenInputElementId}
defaultValue={hoveredCellId || ""}
hidden={true}
ref={hiddenInputRef}
/>
{run.cellRuns.map((cellRun) => (
<TraceRow
key={cellRun.cellId}
cellRun={cellRun}
hovered={cellRun.cellId === hoveredCellId}
dispatchHoverEvent={dispatchHoverEvent}
/>
))}
</div>
const traceRows = (
<TraceRows
run={run}
hoveredCellId={hoveredCellId}
hiddenInputElementId={hiddenInputElementId}
/>
);

if (chartPosition === "above") {
return (
<div key={run.runId} className="flex flex-col">
<pre className="font-mono font-semibold">
{TraceTitle}
{!collapsed && (
<Chart
vegaSpec={vegaSpec}
width={320}
height={120}
signalListeners={handleVegaSignal}
theme={theme}
/>
)}
{!collapsed && TraceRows}
{title}
<Chart
vegaSpec={vegaSpec}
height={120}
signalListeners={handleVegaSignal}
theme={theme}
/>
{traceRows}
</pre>
</div>
);
Expand All @@ -240,19 +258,56 @@ const TraceBlock: React.FC<{
return (
<div key={run.runId} className="flex flex-row">
<pre className="font-mono font-semibold">
{TraceTitle}
{!collapsed && TraceRows}
{title}
{traceRows}
</pre>
{!collapsed && (
<Chart
className="-mt-0.5"
vegaSpec={vegaSpec}
width={240}
height={100}
signalListeners={handleVegaSignal}
theme={theme}
<Chart
className="-mt-0.5 flex-1"
vegaSpec={vegaSpec}
height={100}
signalListeners={handleVegaSignal}
theme={theme}
/>
</div>
);
};

const TraceRows = (props: {
run: Run;
hoveredCellId: CellId | null | undefined;
hiddenInputElementId: string;
}) => {
const { run, hoveredCellId, hiddenInputElementId } = props;

// To send signals to Vega from React, we bind a hidden input element
const hiddenInputRef = useRef<HTMLInputElement>(null);
const dispatchHoverEvent = (cellId: CellId | null) => {
// dispatch input event to trigger vega's param to update
if (hiddenInputRef.current) {
hiddenInputRef.current.value = String(cellId);
hiddenInputRef.current.dispatchEvent(
new Event("input", { bubbles: true }),
);
}
};

return (
<div className="text-xs mt-0.5 ml-3 flex flex-col gap-0.5">
<input
type="text"
id={hiddenInputElementId}
defaultValue={hoveredCellId || ""}
hidden={true}
ref={hiddenInputRef}
/>
{run.cellRuns.map((cellRun) => (
<TraceRow
key={cellRun.cellId}
cellRun={cellRun}
hovered={cellRun.cellId === hoveredCellId}
dispatchHoverEvent={dispatchHoverEvent}
/>
)}
))}
</div>
);
};
Expand Down
Loading

0 comments on commit 0555fa2

Please sign in to comment.