Add loss graph to the ui

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

@@ -0,0 +1,98 @@
+import { NextRequest, NextResponse } from 'next/server';
+import { PrismaClient } from '@prisma/client';
+import path from 'path';
+import fs from 'fs';
+import { getTrainingFolder } from '@/server/settings';
+
+import sqlite3 from 'sqlite3';
+
+export const runtime = 'nodejs';
+
+const prisma = new PrismaClient();
+
+function openDb(filename: string) {
+  const db = new sqlite3.Database(filename);
+  db.configure('busyTimeout', 30_000);
+  return db;
+}
+
+function all<T = any>(db: sqlite3.Database, sql: string, params: any[] = []) {
+  return new Promise<T[]>((resolve, reject) => {
+    db.all(sql, params, (err, rows) => {
+      if (err) reject(err);
+      else resolve(rows as T[]);
+    });
+  });
+}
+
+function closeDb(db: sqlite3.Database) {
+  return new Promise<void>((resolve, reject) => {
+    db.close((err) => (err ? reject(err) : resolve()));
+  });
+}
+
+export async function GET(request: NextRequest, { params }: { params: { jobID: string } }) {
+  // this must be awaited to avoid TS error
+  const { jobID } = await params;
+
+  const job = await prisma.job.findUnique({ where: { id: jobID } });
+  if (!job) return NextResponse.json({ error: 'Job not found' }, { status: 404 });
+
+  const trainingFolder = await getTrainingFolder();
+  const jobFolder = path.join(trainingFolder, job.name);
+  const logPath = path.join(jobFolder, 'loss_log.db');
+
+  if (!fs.existsSync(logPath)) {
+    return NextResponse.json({ keys: [], key: 'loss', points: [] });
+  }
+
+  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 db = openDb(logPath);
+
+  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),
+      })),
+    });
+  } finally {
+    await closeDb(db);
+  }
+}

ui/src/app/jobs/[jobID]/page.tsx

@@ -10,9 +10,10 @@ import JobOverview from '@/components/JobOverview';
 import { redirect } from 'next/navigation';
 import JobActionBar from '@/components/JobActionBar';
 import JobConfigViewer from '@/components/JobConfigViewer';
+import JobLossGraph from '@/components/JobLossGraph';
 import { Job } from '@prisma/client';
 
-type PageKey = 'overview' | 'samples' | 'config';
+type PageKey = 'overview' | 'samples' | 'config' | 'loss_log';
 
 interface Page {
   name: string;
@@ -36,6 +37,12 @@ const pages: Page[] = [
     menuItem: SampleImagesMenu,
     mainCss: 'pt-24',
   },
+  {
+    name: 'Loss Graph',
+    value: 'loss_log',
+    component: JobLossGraph,
+    mainCss: 'pt-24',
+  },
   {
     name: 'Config File',
     value: 'config',

ui/src/app/jobs/new/jobConfig.ts

@@ -92,6 +92,10 @@ export const defaultJobConfig: JobConfig = {
           switch_boundary_every: 1,
           loss_type: 'mse',
         },
+        logging: {
+          log_every: 1,
+          use_ui_logger: true,
+        },
         model: {
           name_or_path: 'ostris/Flex.1-alpha',
           quantize: true,
@@ -187,5 +191,13 @@ export const migrateJobConfig = (jobConfig: JobConfig): JobConfig => {
       false) as boolean;
     delete jobConfig.config.process[0].model.auto_memory;
   }
+
+  if (!('logging' in jobConfig.config.process[0])) {
+    //@ts-ignore
+    jobConfig.config.process[0].logging = {
+      log_every: 1,
+      use_ui_logger: true,
+    };
+  }
   return jobConfig;
 };