Improve load speed of the loss log

ui/src/app/api/jobs/[jobID]/loss/route.ts

@@ -5,6 +5,8 @@ import fs from 'fs';
 import { getTrainingFolder } from '@/server/settings';
 
 import sqlite3 from 'sqlite3';
+import zlib from 'zlib';
+import { promisify } from 'util';
 
 export const runtime = 'nodejs';
 
@@ -29,6 +31,31 @@ function closeDb(db: sqlite3.Database) {
   });
 }
 
+// Serialize a metric series to JSON arrays inside SQLite. Building the arrays
+// with json_group_array avoids marshaling one JS object per row, which is the
+// dominant cost on multi-million-point logs. Values are truncated to 7
+// significant digits (float32-level precision) to shrink the wire payload;
+// value_text rows are never numeric (see _coerce_value in logging_aitk.py)
+// so they serialize as null.
+const SERIES_SQL = `
+  SELECT json_group_array(step) AS steps, json_group_array(v) AS vals
+  FROM (
+    SELECT step,
+           CASE WHEN value_real IS NULL THEN NULL
+                ELSE CAST(printf('%.7g', value_real) AS REAL) END AS v
+    FROM metrics
+    WHERE key = ? AND step >= ?
+    ORDER BY step ASC
+  )
+`;
+
+const gzip = promisify(zlib.gzip);
+
+// GET /api/jobs/:id/loss?since_step=N
+// Returns every metric key in one response: {keys, since_step, series: {key:
+// {steps: number[], values: (number|null)[]}}}. Points with step >= since_step
+// are included (default 0 = everything). The body is gzipped when the client
+// accepts it; a 150k-step run with a dozen keys is tens of MB uncompressed.
 export async function GET(request: NextRequest, { params }: { params: { jobID: string } }) {
   // this must be awaited to avoid TS error
   const { jobID } = await params;
@@ -43,58 +70,46 @@ export async function GET(request: NextRequest, { params }: { params: { jobID: s
   try {
     await fs.promises.access(logPath);
   } catch {
-    return NextResponse.json({ keys: [], key: 'loss', points: [] });
+    return NextResponse.json({ keys: [], since_step: 0, series: {} });
   }
 
   const url = new URL(request.url);
-  const key = url.searchParams.get('key') ?? 'loss';
-  const limit = Math.min(Number(url.searchParams.get('limit') ?? 2000), 20000);
-  const sinceStepParam = url.searchParams.get('since_step');
-  const sinceStep = sinceStepParam != null ? Number(sinceStepParam) : null;
-  const stride = Math.max(1, Number(url.searchParams.get('stride') ?? 1));
+  const sinceStepParam = Number(url.searchParams.get('since_step') ?? 0);
+  const sinceStep = Number.isFinite(sinceStepParam) ? sinceStepParam : 0;
 
   const db = openDb(logPath);
 
+  let body: string;
   try {
     const keysRows = await all<{ key: string }>(db, `SELECT key FROM metric_keys ORDER BY key ASC`);
     const keys = keysRows.map((r) => r.key);
 
-    const points = await all<{
-      step: number;
-      wall_time: number;
-      value: number | null;
-      value_text: string | null;
-    }>(
-      db,
-      `
-      SELECT
-        m.step AS step,
-        s.wall_time AS wall_time,
-        m.value_real AS value,
-        m.value_text AS value_text
-      FROM metrics m
-      JOIN steps s ON s.step = m.step
-      WHERE m.key = ?
-        AND (? IS NULL OR m.step > ?)
-        AND (m.step % ?) = 0
-      ORDER BY m.step ASC
-      LIMIT ?
-      `,
-      [key, sinceStep, sinceStep, stride, limit]
-    );
-
-    return NextResponse.json({
-      key,
-      keys,
-      points: points.map((p) => ({
-        step: p.step,
-        wall_time: p.wall_time,
-        value: p.value ?? (p.value_text ? Number(p.value_text) : null),
-      })),
-    });
+    const parts: string[] = [];
+    for (const key of keys) {
+      const rows = await all<{ steps: string; vals: string }>(db, SERIES_SQL, [key, sinceStep]);
+      const steps = rows[0]?.steps ?? '[]';
+      const vals = rows[0]?.vals ?? '[]';
+      parts.push(`${JSON.stringify(key)}:{"steps":${steps},"values":${vals}}`);
+    }
+    body = `{"keys":${JSON.stringify(keys)},"since_step":${sinceStep},"series":{${parts.join(',')}}}`;
   } finally {
     await closeDb(db);
   }
+
+  const headers: Record<string, string> = {
+    'Content-Type': 'application/json',
+    'Cache-Control': 'no-store',
+    Vary: 'Accept-Encoding',
+  };
+  const acceptEncoding = request.headers.get('accept-encoding') ?? '';
+  if (/\bgzip\b/.test(acceptEncoding)) {
+    // Level 1: the payload is mostly digits, so higher levels buy ~5% for 3x
+    // the CPU time.
+    const gz = await gzip(Buffer.from(body), { level: 1 });
+    headers['Content-Encoding'] = 'gzip';
+    return new NextResponse(gz, { headers });
+  }
+  return new NextResponse(body, { headers });
 }
 
 function run(db: sqlite3.Database, sql: string, params: any[] = []) {

ui/src/components/JobLossGraph.tsx

@@ -1,7 +1,7 @@
 'use client';
 
 import { Job } from '@prisma/client';
-import useJobLossLog, { LossPoint } from '@/hooks/useJobLossLog';
+import useJobLossLog, { LossSeries } from '@/hooks/useJobLossLog';
 import { useMemo, useState, useEffect, useRef, useCallback } from 'react';
 import uPlot from 'uplot';
 import 'uplot/dist/uPlot.min.css';
@@ -233,7 +233,16 @@ export default function JobLossGraph({ job }: Props) {
     });
   }, [lossKeys]);
 
-  const activeKeys = useMemo(() => lossKeys.filter(k => enabled[k] !== false), [lossKeys, enabled]);
+  // Keys not yet in `enabled` (data and keys now arrive in the same render,
+  // before the sync effect above has run) take their default here rather than
+  // rendering as enabled and getting pruned a frame later.
+  const activeKeys = useMemo(
+    () =>
+      lossKeys.filter(
+        k => enabled[k] ?? persistedEnabledRef.current?.[k] ?? (k === 'loss/loss' || k === 'val/loss'),
+      ),
+    [lossKeys, enabled],
+  );
 
   // Assign palette slots by position in the (sorted) lossKeys list rather than
   // by hashing the key name — hashing let different keys collide onto the same
@@ -257,11 +266,11 @@ export default function JobLossGraph({ job }: Props) {
     // Union of all steps across active series.
     const stepSet = new Set<number>();
     for (const key of activeKeys) {
-      const pts: LossPoint[] = series[key] ?? [];
-      for (const p of pts) {
-        if (p.value === null || !Number.isFinite(p.value as number)) continue;
-        if (useLogScale && (p.value as number) <= 0) continue;
-        stepSet.add(p.step);
+      const s: LossSeries | undefined = series[key];
+      if (!s) continue;
+      for (let i = 0; i < s.steps.length; i++) {
+        if (useLogScale && s.values[i] <= 0) continue;
+        stepSet.add(s.steps[i]);
       }
     }
     let xs = Array.from(stepSet).sort((a, b) => a - b);
@@ -291,13 +300,13 @@ export default function JobLossGraph({ job }: Props) {
     for (let ki = 0; ki < activeKeys.length; ki++) {
       const key = activeKeys[ki];
       const scaleKey = `y::${key}`;
-      const pts: LossPoint[] = series[key] ?? [];
+      const s: LossSeries = series[key] ?? { steps: [], values: [] };
       const map = new Map<number, number>();
-      for (const p of pts) {
-        if (p.value === null || !Number.isFinite(p.value as number)) continue;
-        if (useLogScale && (p.value as number) <= 0) continue;
-        if (!xsSet.has(p.step)) continue;
-        map.set(p.step, p.value as number);
+      for (let i = 0; i < s.steps.length; i++) {
+        const v = s.values[i];
+        if (useLogScale && v <= 0) continue;
+        if (!xsSet.has(s.steps[i])) continue;
+        map.set(s.steps[i], v);
       }
       const raw: (number | null)[] = xs.map(s => (map.has(s) ? (map.get(s) as number) : null));
       // Metrics logged less often than every step are null at most x positions.

ui/src/hooks/useJobLossLog.tsx

@@ -4,13 +4,20 @@ import { useEffect, useState, useRef, useCallback, useMemo } from 'react';
 import { apiClient } from '@/utils/api';
 import usePollLoop from '@/hooks/usePollLoop';
 
-export interface LossPoint {
-  step: number;
-  wall_time?: number;
-  value: number | null;
+// Columnar series: steps[i] pairs with values[i]. Null / non-finite values are
+// dropped at ingest, so every value here is a finite number.
+export interface LossSeries {
+  steps: number[];
+  values: number[];
 }
 
-type SeriesMap = Record<string, LossPoint[]>;
+type SeriesMap = Record<string, LossSeries>;
+
+interface LossResponse {
+  keys?: string[];
+  since_step?: number;
+  series?: Record<string, { steps: number[]; values: (number | null)[] }>;
+}
 
 export default function useJobLossLog(jobID: string, reloadInterval: null | number = null) {
   const [series, setSeries] = useState<SeriesMap>({});
@@ -20,8 +27,10 @@ export default function useJobLossLog(jobID: string, reloadInterval: null | numb
   const didInitialLoadRef = useRef(false);
   const inFlightRef = useRef(false);
 
-  // track last step per key so polling is incremental per series
-  const lastStepByKeyRef = useRef<Record<string, number | null>>({});
+  // Highest step seen across every key. Polls ask for step >= this (not >) so a
+  // key logged at that step after the previous poll is still picked up; the
+  // per-key merge below drops the one re-sent point.
+  const lastStepRef = useRef<number | null>(null);
 
   const lossKeys = useMemo(() => {
     const base = keys ?? [];
@@ -40,69 +49,60 @@ export default function useJobLossLog(jobID: string, reloadInterval: null | numb
     setStatus(loadStatus);
 
     try {
-      // Step 1: get key list (we can do this by calling endpoint once; it returns keys)
-      // Keep it cheap: limit=1.
-      const first = await apiClient
-        .get(`/api/jobs/${jobID}/loss`, { params: { key: 'loss', limit: 1 } })
-        .then(res => res.data as { keys?: string[] });
+      // One request returns every key. Incremental after the initial load.
+      const params: Record<string, number> = {};
+      if (reloadInterval && didInitialLoadRef.current && lastStepRef.current != null) {
+        params.since_step = lastStepRef.current;
+      }
+      const res = await apiClient.get(`/api/jobs/${jobID}/loss`, { params });
+      const data = res.data as LossResponse;
 
-      const newKeys = first.keys ?? [];
+      const newKeys = data.keys ?? [];
       setKeys(newKeys);
-
       const wantedLossKeys = (newKeys.length ? [...newKeys] : ['loss']).sort();
-
-      // Step 2: fetch each loss key incrementally (since_step per key if polling)
-      const requests = wantedLossKeys.map(k => {
-        const params: Record<string, any> = { key: k };
-
-        if (reloadInterval && lastStepByKeyRef.current[k] != null) {
-          params.since_step = lastStepByKeyRef.current[k];
-        }
-
-        params.limit = 1000000;
-
-        return apiClient
-          .get(`/api/jobs/${jobID}/loss`, { params })
-          .then(res => res.data as { key: string; points?: LossPoint[] });
-      });
-
-      const results = await Promise.all(requests);
+      const incoming = data.series ?? {};
+      const initial = !didInitialLoadRef.current;
 
       setSeries(prev => {
-        const next: SeriesMap = { ...prev };
+        const next: SeriesMap = initial ? {} : { ...prev };
+        let maxStep = lastStepRef.current ?? null;
 
-        for (const r of results) {
-          const k = r.key;
-          const newPoints = (r.points ?? []).filter(p => p.value !== null);
+        for (const k of wantedLossKeys) {
+          const inc = incoming[k];
+          const existing = initial ? undefined : next[k];
+          const prevLast = existing && existing.steps.length ? existing.steps[existing.steps.length - 1] : null;
 
-          if (!didInitialLoadRef.current) {
-            // initial: replace
-            next[k] = newPoints;
-          } else if (newPoints.length) {
-            const existing = next[k] ?? [];
-            const prevLast = existing.length ? existing[existing.length - 1].step : null;
-            const filtered = prevLast == null ? newPoints : newPoints.filter(p => p.step > prevLast);
-            next[k] = filtered.length ? [...existing, ...filtered] : existing;
-          } else {
-            // no new points: keep existing
-            next[k] = next[k] ?? [];
+          const steps: number[] = existing ? existing.steps.slice() : [];
+          const values: number[] = existing ? existing.values.slice() : [];
+          let appended = false;
+          if (inc) {
+            const n = Math.min(inc.steps.length, inc.values.length);
+            for (let i = 0; i < n; i++) {
+              const v = inc.values[i];
+              if (v === null || !Number.isFinite(v)) continue;
+              const s = inc.steps[i];
+              if (prevLast != null && s <= prevLast) continue;
+              steps.push(s);
+              values.push(v as number);
+              appended = true;
+            }
           }
+          if (initial || appended) next[k] = { steps, values };
+          else next[k] = existing ?? { steps: [], values: [] };
 
-          // update last step per key
-          const finalArr = next[k] ?? [];
-          lastStepByKeyRef.current[k] = finalArr.length
-            ? finalArr[finalArr.length - 1].step
-            : (lastStepByKeyRef.current[k] ?? null);
+          const arr = next[k].steps;
+          if (arr.length) {
+            const last = arr[arr.length - 1];
+            if (maxStep == null || last > maxStep) maxStep = last;
+          }
         }
 
         // remove stale keys that no longer exist (rare, but keeps UI clean)
         for (const existingKey of Object.keys(next)) {
-          if (!wantedLossKeys.includes(existingKey)) {
-            delete next[existingKey];
-            delete lastStepByKeyRef.current[existingKey];
-          }
+          if (!wantedLossKeys.includes(existingKey)) delete next[existingKey];
         }
 
+        lastStepRef.current = maxStep;
         return next;
       });
 
@@ -124,13 +124,24 @@ export default function useJobLossLog(jobID: string, reloadInterval: null | numb
       await apiClient.delete(`/api/jobs/${jobID}/loss`, { data: { min_step: minStep, max_step: maxStep } });
       setSeries(prev => {
         const next: SeriesMap = {};
+        let last: number | null = null;
         for (const k of Object.keys(prev)) {
-          const kept = prev[k].filter(p => p.step < minStep || p.step > maxStep);
-          next[k] = kept;
-          // Re-anchor incremental polling to the new tail so a deleted tail
-          // isn't treated as already-fetched.
-          lastStepByKeyRef.current[k] = kept.length ? kept[kept.length - 1].step : null;
+          const steps: number[] = [];
+          const values: number[] = [];
+          const src = prev[k];
+          for (let i = 0; i < src.steps.length; i++) {
+            const s = src.steps[i];
+            if (s < minStep || s > maxStep) {
+              steps.push(s);
+              values.push(src.values[i]);
+            }
+          }
+          next[k] = { steps, values };
+          if (steps.length && (last == null || steps[steps.length - 1] > last)) last = steps[steps.length - 1];
         }
+        // Re-anchor incremental polling to the new tail so a deleted tail
+        // isn't treated as already-fetched.
+        lastStepRef.current = last;
         return next;
       });
     },
@@ -141,7 +152,7 @@ export default function useJobLossLog(jobID: string, reloadInterval: null | numb
   // before the first fetch when jobID changes.
   useEffect(() => {
     didInitialLoadRef.current = false;
-    lastStepByKeyRef.current = {};
+    lastStepRef.current = null;
     setSeries({});
     setKeys([]);
     setStatus('idle');

version.py

@@ -1 +1 @@
-VERSION = "0.13.39"
+VERSION = "0.13.40"