diff --git a/ui/src/app/api/jobs/[jobID]/loss/route.ts b/ui/src/app/api/jobs/[jobID]/loss/route.ts index 6ad0e0b..2fad591 100644 --- a/ui/src/app/api/jobs/[jobID]/loss/route.ts +++ b/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 = { + '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[] = []) { diff --git a/ui/src/components/JobLossGraph.tsx b/ui/src/components/JobLossGraph.tsx index cda523f..61eb555 100644 --- a/ui/src/components/JobLossGraph.tsx +++ b/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(); 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(); - 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. diff --git a/ui/src/hooks/useJobLossLog.tsx b/ui/src/hooks/useJobLossLog.tsx index fb80575..7fedc35 100644 --- a/ui/src/hooks/useJobLossLog.tsx +++ b/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; +type SeriesMap = Record; + +interface LossResponse { + keys?: string[]; + since_step?: number; + series?: Record; +} export default function useJobLossLog(jobID: string, reloadInterval: null | number = null) { const [series, setSeries] = useState({}); @@ -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>({}); + // 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(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 = {}; + 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 = { 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'); diff --git a/version.py b/version.py index f8b754a..edef351 100644 --- a/version.py +++ b/version.py @@ -1 +1 @@ -VERSION = "0.13.39" +VERSION = "0.13.40"