diff --git a/packages/db/src/etl/benchmark-ingest.ts b/packages/db/src/etl/benchmark-ingest.ts index 2a2382c8..c9808b61 100644 --- a/packages/db/src/etl/benchmark-ingest.ts +++ b/packages/db/src/etl/benchmark-ingest.ts @@ -101,32 +101,36 @@ export async function insertServerLog( ): Promise { if (benchmarkResultIds.length === 0) return; - // Only link rows that don't already have a server log - const unlinked = await sql<{ id: number }[]>` - select id from benchmark_results - where id = any(${sql.array(benchmarkResultIds)}::bigint[]) - and server_log_id is null - `; - if (unlinked.length === 0) return; - - const [{ id: logId }] = await sql<{ id: number }[]>` - insert into server_logs (server_log) values (${serverLog}) - returning id - `; - // Derive the KV-cache pool size (tokens) from the log's authoritative - // "GPU KV cache size: N tokens" line(s) and stash it on the result's metrics - // JSON, mirroring how trace-replay-ingest derives cache-hit rates. The - // scraped vllm:cache_config_info metric can't reconstruct this for MLA models. - const kvCachePoolTokens = kvCachePoolTokensFromServerLog(serverLog); - await sql` - update benchmark_results - set server_log_id = ${logId}${ - kvCachePoolTokens === null - ? sql`` - : sql`, metrics = jsonb_set(metrics, '{kv_cache_pool_tokens}', to_jsonb(${kvCachePoolTokens}::bigint))` - } - where id = any(${sql.array(unlinked.map((r) => r.id))}::bigint[]) - `; + await sql.begin(async (tx) => { + // Lock before the check so concurrent config workers cannot create an + // unreferenced duplicate server_logs row for the same benchmark result. + const unlinked = await tx<{ id: number }[]>` + select id from benchmark_results + where id = any(${tx.array(benchmarkResultIds)}::bigint[]) + and server_log_id is null + for update + `; + if (unlinked.length === 0) return; + + const [{ id: logId }] = await tx<{ id: number }[]>` + insert into server_logs (server_log) values (${serverLog}) + returning id + `; + // Derive the KV-cache pool size (tokens) from the log's authoritative + // "GPU KV cache size: N tokens" line(s) and stash it on the result's metrics + // JSON, mirroring how trace-replay-ingest derives cache-hit rates. The + // scraped vllm:cache_config_info metric can't reconstruct this for MLA models. + const kvCachePoolTokens = kvCachePoolTokensFromServerLog(serverLog); + await tx` + update benchmark_results + set server_log_id = ${logId}${ + kvCachePoolTokens === null + ? tx`` + : tx`, metrics = jsonb_set(metrics, '{kv_cache_pool_tokens}', to_jsonb(${kvCachePoolTokens}::bigint))` + } + where id = any(${tx.array(unlinked.map((r) => r.id))}::bigint[]) + `; + }); } /** diff --git a/packages/db/src/etl/trace-replay-ingest.ts b/packages/db/src/etl/trace-replay-ingest.ts index 1c739b7d..0ef8b588 100644 --- a/packages/db/src/etl/trace-replay-ingest.ts +++ b/packages/db/src/etl/trace-replay-ingest.ts @@ -12,9 +12,9 @@ import { gzipSync } from 'node:zlib'; import type postgres from 'postgres'; -import { computeAggregateStats } from './compute-aggregate-stats.js'; -import { computeChartSeries } from './compute-chart-series.js'; -import { computeRequestTimeline } from './compute-request-timeline.js'; +import { computeAggregateStats, type AggregateStats } from './compute-aggregate-stats'; +import { computeChartSeries, type ChartSeries } from './compute-chart-series'; +import { computeRequestTimeline, type RequestTimeline } from './compute-request-timeline'; import type { ServerMetricsContext } from './server-metrics-adapters'; type Sql = ReturnType; @@ -24,6 +24,20 @@ export interface TraceReplayIngestOptions { progressLabel?: string; } +export interface PreparedTraceReplay { + profileGz: Buffer | null; + profileSize: number | null; + serverMetricsCsv: Buffer | null; + csvSize: number | null; + metricsJsonGz: Buffer | null; + metricsJsonSize: number | null; + aggregateStats: AggregateStats; + chartSeries: ChartSeries | null; + requestTimeline: RequestTimeline | null; + gzipMs: number; + computeMs: number; +} + function formatBytes(bytes: number | null | undefined): string { if (bytes === null || bytes === undefined) return 'none'; if (bytes < 1024) return `${bytes} B`; @@ -38,6 +52,52 @@ function elapsed(startMs: number): string { return `${((Date.now() - startMs) / 1000).toFixed(1)}s`; } +export async function findUnlinkedTraceReplayBenchmarkIds( + sql: Sql, + benchmarkResultIds: number[], +): Promise { + if (benchmarkResultIds.length === 0) return []; + const rows = await sql<{ id: number }[]>` + select id from benchmark_results + where id = any(${sql.array(benchmarkResultIds)}::bigint[]) + and trace_replay_id is null + `; + return rows.map((row) => row.id); +} + +export async function prepareTraceReplay( + profileExportJsonl: Buffer | null, + serverMetricsCsv: Buffer | null, + serverMetricsJson: Buffer | null, + metricsContext: ServerMetricsContext = {}, +): Promise { + const gzipStart = Date.now(); + const profileGz = profileExportJsonl ? gzipSync(profileExportJsonl) : null; + const metricsJsonGz = serverMetricsJson ? gzipSync(serverMetricsJson) : null; + const gzipMs = Date.now() - gzipStart; + + const computeStart = Date.now(); + const [aggregateStats, chartSeries, requestTimeline] = await Promise.all([ + computeAggregateStats({ profileBlob: profileGz, serverBlob: metricsJsonGz }), + computeChartSeries(metricsJsonGz, metricsContext), + Promise.resolve(computeRequestTimeline(profileGz)), + ]); + + return { + profileGz, + profileSize: profileExportJsonl?.length ?? null, + serverMetricsCsv, + csvSize: serverMetricsCsv?.length ?? null, + metricsJsonGz, + metricsJsonSize: serverMetricsJson?.length ?? null, + aggregateStats, + chartSeries, + requestTimeline, + gzipMs, + computeMs: Date.now() - computeStart, + }; +} + /** * Persist the per-point trace files and link them to `benchmarkResultIds`. * @@ -72,133 +132,146 @@ export async function insertTraceReplay( if (benchmarkResultIds.length === 0) return; if (!profileExportJsonl && !serverMetricsCsv && !serverMetricsJson) return; - // Only link rows that don't already point at a trace_replay row — keeps - // re-ingest from inserting duplicate sibling blobs. const linkStart = Date.now(); log(`checking ${benchmarkResultIds.length} benchmark row(s) for existing links`); - const unlinked = await sql<{ id: number }[]>` - select id from benchmark_results - where id = any(${sql.array(benchmarkResultIds)}::bigint[]) - and trace_replay_id is null - `; - log(`found ${unlinked.length} unlinked row(s) (${elapsed(linkStart)})`); - if (unlinked.length === 0) { + const unlinkedIds = await findUnlinkedTraceReplayBenchmarkIds(sql, benchmarkResultIds); + log(`found ${unlinkedIds.length} unlinked row(s) (${elapsed(linkStart)})`); + if (unlinkedIds.length === 0) { log('skipping blob insert; all benchmark rows already linked'); return; } - const gzipStart = Date.now(); log( `compressing profile=${formatBytes(profileExportJsonl?.length)}, ` + `server_csv=${formatBytes(serverMetricsCsv?.length)}, ` + `server_json=${formatBytes(serverMetricsJson?.length)}`, ); - const profileGz = profileExportJsonl ? gzipSync(profileExportJsonl) : null; - const profileSize = profileExportJsonl ? profileExportJsonl.length : null; - const csvSize = serverMetricsCsv ? serverMetricsCsv.length : null; - const metricsJsonGz = serverMetricsJson ? gzipSync(serverMetricsJson) : null; - const metricsJsonSize = serverMetricsJson ? serverMetricsJson.length : null; + const prepared = await prepareTraceReplay( + profileExportJsonl, + serverMetricsCsv, + serverMetricsJson, + metricsContext, + ); log( - `compressed profile=${formatBytes(profileGz?.length)}, ` + - `server_json=${formatBytes(metricsJsonGz?.length)} (${elapsed(gzipStart)})`, + `compressed profile=${formatBytes(prepared.profileGz?.length)}, ` + + `server_json=${formatBytes(prepared.metricsJsonGz?.length)} ` + + `(${(prepared.gzipMs / 1000).toFixed(1)}s)`, ); - - // Pre-compute aggregate stats + chart-ready time-series + per-request - // timeline so the detail page doesn't have to re-parse these blobs on - // every request. Each helper tolerates a null blob and falls back to - // a streaming parser for oversized server_metrics blobs. - const computeStart = Date.now(); - log('computing aggregate stats, chart series, and request timeline'); - const [aggregateStats, chartSeries, requestTimeline] = await Promise.all([ - computeAggregateStats({ profileBlob: profileGz, serverBlob: metricsJsonGz }), - computeChartSeries(metricsJsonGz, metricsContext), - Promise.resolve(computeRequestTimeline(profileGz)), - ]); log( - `computed derived JSON: chart_windows=${chartSeries?.timeslicesCount ?? 0}, ` + - `timeline_requests=${requestTimeline?.requests.length ?? 0} (${elapsed(computeStart)})`, + `computed derived JSON: chart_windows=${prepared.chartSeries?.timeslicesCount ?? 0}, ` + + `timeline_requests=${prepared.requestTimeline?.requests.length ?? 0} ` + + `(${(prepared.computeMs / 1000).toFixed(1)}s)`, ); - const insertStart = Date.now(); - log('inserting trace_replay blob row'); - const [{ id: traceReplayId }] = await sql<{ id: number }[]>` - insert into agentic_trace_replay ( - profile_export_jsonl_gz, - profile_export_uncompressed_size, - server_metrics_csv, - server_metrics_csv_size, - server_metrics_json_gz, - server_metrics_json_uncompressed_size, - aggregate_stats, - chart_series, - request_timeline - ) - values ( - ${profileGz}, - ${profileSize}, - ${serverMetricsCsv}, - ${csvSize}, - ${metricsJsonGz}, - ${metricsJsonSize}, - ${sql.json(structuredClone(aggregateStats) as unknown as Parameters[0])}, - ${chartSeries === null ? null : sql.json(structuredClone(chartSeries) as unknown as Parameters[0])}, - ${requestTimeline === null ? null : sql.json(structuredClone(requestTimeline) as unknown as Parameters[0])} - ) - returning id - `; - log(`inserted trace_replay_id=${traceReplayId} (${elapsed(insertStart)})`); - - const updateStart = Date.now(); - log(`linking trace_replay_id=${traceReplayId} to ${unlinked.length} benchmark row(s)`); - await sql` - update benchmark_results - set trace_replay_id = ${traceReplayId} - where id = any(${sql.array(unlinked.map((r) => r.id))}::bigint[]) - `; - log(`linked benchmark rows (${elapsed(updateStart)})`); - - // Derive lifetime GPU + CPU cache hit rates from chart_series. SGLang - // runs don't populate these in the harness JSON; vLLM runs do but only - // for GPU. We always recompute to keep the derivation consistent with - // what the detail-page charts plot — overwriting any pre-existing value. - // - // Source label naming differs by framework / cache topology: - // SGLang hicache: 'cache hit (HBM)' + 'cache hit (CPU offload)' - // SGLang older: 'cache hit' (no tier breakdown) - // vLLM LMCache: 'local_cache_hit' + 'external_kv_transfer' (+ 'local_compute' for miss) - // vLLM single: falls back to prefixCacheHitsTps total (= local cache only) - if (chartSeries && chartSeries.prefillTps.length > 0) { - const sumPrompts = chartSeries.prefillTps.reduce((s, p) => s + p.value, 0); - if (sumPrompts > 0) { - const sumOf = (name: string): number => - (chartSeries.promptTokensBySource[name] ?? []).reduce((s, p) => s + p.value, 0); - // CPU-offload hits: SGLang hicache + vLLM LMCache external transfer. - const cpuHits = sumOf('cache hit (CPU offload)') + sumOf('external_kv_transfer'); - // GPU/HBM hits from source breakdown, summed across known aliases. - const hbmFromBreakdown = - sumOf('cache hit (HBM)') + sumOf('cache hit') + sumOf('local_cache_hit'); - // If the source breakdown has any GPU entry, use it. Otherwise fall back - // to total prefixCacheHitsTps sum (single-source vLLM path with no - // by_source metric — equals the lone cache counter's lifetime). - const gpuHits = - hbmFromBreakdown > 0 - ? hbmFromBreakdown - : chartSeries.prefixCacheHitsTps.reduce((s, p) => s + p.value, 0); - const gpuRate = gpuHits / sumPrompts; - const cpuRate = cpuHits > 0 ? cpuHits / sumPrompts : null; - await sql` - update benchmark_results - set metrics = jsonb_set( - case when ${cpuRate}::numeric is not null - then jsonb_set(metrics, '{server_cpu_cache_hit_rate}', to_jsonb(${cpuRate}::numeric)) - else metrics - end, - '{server_gpu_cache_hit_rate}', - to_jsonb(${gpuRate}::numeric) - ) - where id = any(${sql.array(unlinked.map((r) => r.id))}::bigint[]) - `; - log('updated cache-hit metrics from chart series'); + await insertPreparedTraceReplay(sql, benchmarkResultIds, prepared, options); +} + +export async function insertPreparedTraceReplay( + sql: Sql, + benchmarkResultIds: number[], + prepared: PreparedTraceReplay, + options: TraceReplayIngestOptions = {}, +): Promise { + const { progressLabel } = options; + const log = (message: string): void => { + if (progressLabel) console.log(` trace_replay ${progressLabel}: ${message}`); + }; + + await sql.begin(async (tx) => { + // Row locks make the check-and-link atomic when multiple config workers + // converge on the same idempotency key. + const unlinked = await tx<{ id: number }[]>` + select id from benchmark_results + where id = any(${tx.array(benchmarkResultIds)}::bigint[]) + and trace_replay_id is null + for update + `; + if (unlinked.length === 0) { + log('skipping blob insert; all benchmark rows already linked'); + return; } - } + + const insertStart = Date.now(); + log('inserting trace_replay blob row'); + const [{ id: traceReplayId }] = await tx<{ id: number }[]>` + insert into agentic_trace_replay ( + profile_export_jsonl_gz, + profile_export_uncompressed_size, + server_metrics_csv, + server_metrics_csv_size, + server_metrics_json_gz, + server_metrics_json_uncompressed_size, + aggregate_stats, + chart_series, + request_timeline + ) + values ( + ${prepared.profileGz}, + ${prepared.profileSize}, + ${prepared.serverMetricsCsv}, + ${prepared.csvSize}, + ${prepared.metricsJsonGz}, + ${prepared.metricsJsonSize}, + ${tx.json(structuredClone(prepared.aggregateStats) as unknown as Parameters[0])}, + ${prepared.chartSeries === null ? null : tx.json(structuredClone(prepared.chartSeries) as unknown as Parameters[0])}, + ${prepared.requestTimeline === null ? null : tx.json(structuredClone(prepared.requestTimeline) as unknown as Parameters[0])} + ) + returning id + `; + log(`inserted trace_replay_id=${traceReplayId} (${elapsed(insertStart)})`); + + const updateStart = Date.now(); + log(`linking trace_replay_id=${traceReplayId} to ${unlinked.length} benchmark row(s)`); + await tx` + update benchmark_results + set trace_replay_id = ${traceReplayId} + where id = any(${tx.array(unlinked.map((r) => r.id))}::bigint[]) + `; + log(`linked benchmark rows (${elapsed(updateStart)})`); + + // Derive lifetime GPU + CPU cache hit rates from chart_series. SGLang + // runs don't populate these in the harness JSON; vLLM runs do but only + // for GPU. We always recompute to keep the derivation consistent with + // what the detail-page charts plot — overwriting any pre-existing value. + // + // Source label naming differs by framework / cache topology: + // SGLang hicache: 'cache hit (HBM)' + 'cache hit (CPU offload)' + // SGLang older: 'cache hit' (no tier breakdown) + // vLLM LMCache: 'local_cache_hit' + 'external_kv_transfer' (+ 'local_compute' for miss) + // vLLM single: falls back to prefixCacheHitsTps total (= local cache only) + if (prepared.chartSeries && prepared.chartSeries.prefillTps.length > 0) { + const sumPrompts = prepared.chartSeries.prefillTps.reduce((s, p) => s + p.value, 0); + if (sumPrompts > 0) { + const sumOf = (name: string): number => + (prepared.chartSeries?.promptTokensBySource[name] ?? []).reduce((s, p) => s + p.value, 0); + // CPU-offload hits: SGLang hicache + vLLM LMCache external transfer. + const cpuHits = sumOf('cache hit (CPU offload)') + sumOf('external_kv_transfer'); + // GPU/HBM hits from source breakdown, summed across known aliases. + const hbmFromBreakdown = + sumOf('cache hit (HBM)') + sumOf('cache hit') + sumOf('local_cache_hit'); + // If the source breakdown has any GPU entry, use it. Otherwise fall back + // to total prefixCacheHitsTps sum (single-source vLLM path with no + // by_source metric — equals the lone cache counter's lifetime). + const gpuHits = + hbmFromBreakdown > 0 + ? hbmFromBreakdown + : prepared.chartSeries.prefixCacheHitsTps.reduce((s, p) => s + p.value, 0); + const gpuRate = gpuHits / sumPrompts; + const cpuRate = cpuHits > 0 ? cpuHits / sumPrompts : null; + await tx` + update benchmark_results + set metrics = jsonb_set( + case when ${cpuRate}::numeric is not null + then jsonb_set(metrics, '{server_cpu_cache_hit_rate}', to_jsonb(${cpuRate}::numeric)) + else metrics + end, + '{server_gpu_cache_hit_rate}', + to_jsonb(${gpuRate}::numeric) + ) + where id = any(${tx.array(unlinked.map((r) => r.id))}::bigint[]) + `; + log('updated cache-hit metrics from chart series'); + } + } + }); } diff --git a/packages/db/src/etl/trace-replay-worker-bootstrap.mjs b/packages/db/src/etl/trace-replay-worker-bootstrap.mjs new file mode 100644 index 00000000..099b48d5 --- /dev/null +++ b/packages/db/src/etl/trace-replay-worker-bootstrap.mjs @@ -0,0 +1,3 @@ +import { tsImport } from 'tsx/esm/api'; + +await tsImport('./trace-replay-worker.ts', import.meta.url); diff --git a/packages/db/src/etl/trace-replay-worker-pool.ts b/packages/db/src/etl/trace-replay-worker-pool.ts new file mode 100644 index 00000000..21084074 --- /dev/null +++ b/packages/db/src/etl/trace-replay-worker-pool.ts @@ -0,0 +1,103 @@ +import { Worker } from 'node:worker_threads'; + +import type { PreparedTraceReplay } from './trace-replay-ingest'; +import type { ServerMetricsContext } from './server-metrics-adapters'; +import type { TraceReplayWorkerRequest, TraceReplayWorkerResponse } from './trace-replay-worker'; + +export interface TraceReplayWork { + profilePath: string | null; + serverMetricsCsvPath: string | null; + serverMetricsJsonPath: string | null; + metricsContext: ServerMetricsContext; +} + +interface QueuedWork extends TraceReplayWork { + id: number; + resolve: (prepared: PreparedTraceReplay) => void; + reject: (error: Error) => void; +} + +interface WorkerSlot { + worker: Worker; + current: QueuedWork | null; +} + +function hydratePrepared(prepared: TraceReplayWorkerResponse & { ok: true }): PreparedTraceReplay { + return { + ...prepared.prepared, + profileGz: prepared.prepared.profileGz ? Buffer.from(prepared.prepared.profileGz) : null, + serverMetricsCsv: prepared.prepared.serverMetricsCsv + ? Buffer.from(prepared.prepared.serverMetricsCsv) + : null, + metricsJsonGz: prepared.prepared.metricsJsonGz + ? Buffer.from(prepared.prepared.metricsJsonGz) + : null, + }; +} + +export class TraceReplayWorkerPool { + readonly #slots: WorkerSlot[] = []; + readonly #queue: QueuedWork[] = []; + #nextId = 1; + #closed = false; + + constructor(size: number) { + if (!Number.isInteger(size) || size < 1) throw new Error(`Invalid worker pool size: ${size}`); + for (let index = 0; index < size; index++) this.#slots.push(this.#createSlot()); + } + + prepare(work: TraceReplayWork): Promise { + if (this.#closed) return Promise.reject(new Error('Trace replay worker pool is closed')); + return new Promise((resolve, reject) => { + this.#queue.push({ ...work, id: this.#nextId++, resolve, reject }); + this.#dispatch(); + }); + } + + async close(): Promise { + this.#closed = true; + await Promise.all(this.#slots.map((slot) => slot.worker.terminate())); + } + + #createSlot(): WorkerSlot { + const slot: WorkerSlot = { + worker: new Worker(new URL('trace-replay-worker-bootstrap.mjs', import.meta.url)), + current: null, + }; + slot.worker.on('message', (response: TraceReplayWorkerResponse) => { + const current = slot.current; + if (!current || response.id !== current.id) return; + slot.current = null; + if (response.ok) current.resolve(hydratePrepared(response)); + else current.reject(new Error(response.error)); + this.#dispatch(); + }); + slot.worker.on('error', (error) => { + slot.current?.reject(error instanceof Error ? error : new Error(String(error))); + slot.current = null; + if (!this.#closed) { + const index = this.#slots.indexOf(slot); + this.#slots[index] = this.#createSlot(); + this.#dispatch(); + } + }); + return slot; + } + + #dispatch(): void { + for (const slot of this.#slots) { + if (slot.current || this.#queue.length === 0) continue; + const work = this.#queue.shift()!; + slot.current = work; + const request: TraceReplayWorkerRequest = { + id: work.id, + profilePath: work.profilePath, + serverMetricsCsvPath: work.serverMetricsCsvPath, + serverMetricsJsonPath: work.serverMetricsJsonPath, + metricsContext: work.metricsContext, + }; + // oxlint-disable-next-line unicorn/require-post-message-target-origin -- Node Worker has no targetOrigin parameter. + slot.worker.postMessage(request); + } + } +} diff --git a/packages/db/src/etl/trace-replay-worker.ts b/packages/db/src/etl/trace-replay-worker.ts new file mode 100644 index 00000000..c11dc44e --- /dev/null +++ b/packages/db/src/etl/trace-replay-worker.ts @@ -0,0 +1,46 @@ +import fs from 'node:fs'; +import { parentPort } from 'node:worker_threads'; + +import { prepareTraceReplay } from './trace-replay-ingest'; +import type { ServerMetricsContext } from './server-metrics-adapters'; + +export interface TraceReplayWorkerRequest { + id: number; + profilePath: string | null; + serverMetricsCsvPath: string | null; + serverMetricsJsonPath: string | null; + metricsContext: ServerMetricsContext; +} + +export type TraceReplayWorkerResponse = + | { id: number; ok: true; prepared: Awaited> } + | { id: number; ok: false; error: string }; + +const port = parentPort; +if (!port) throw new Error('trace-replay-worker must run in a worker thread'); + +port.on('message', async (request: TraceReplayWorkerRequest) => { + try { + const prepared = await prepareTraceReplay( + request.profilePath ? fs.readFileSync(request.profilePath) : null, + request.serverMetricsCsvPath ? fs.readFileSync(request.serverMetricsCsvPath) : null, + request.serverMetricsJsonPath ? fs.readFileSync(request.serverMetricsJsonPath) : null, + request.metricsContext, + ); + const response = { + id: request.id, + ok: true, + prepared, + } satisfies TraceReplayWorkerResponse; + // oxlint-disable-next-line unicorn/require-post-message-target-origin -- Node MessagePort has no targetOrigin parameter. + port.postMessage(response); + } catch (error) { + const response = { + id: request.id, + ok: false, + error: error instanceof Error ? (error.stack ?? error.message) : String(error), + } satisfies TraceReplayWorkerResponse; + // oxlint-disable-next-line unicorn/require-post-message-target-origin -- Node MessagePort has no targetOrigin parameter. + port.postMessage(response); + } +}); diff --git a/packages/db/src/ingest-ci-run.ts b/packages/db/src/ingest-ci-run.ts index cada82d6..ff138c30 100644 --- a/packages/db/src/ingest-ci-run.ts +++ b/packages/db/src/ingest-ci-run.ts @@ -50,7 +50,11 @@ import { bulkUpsertAvailability, insertServerLog, } from './etl/benchmark-ingest'; -import { insertTraceReplay } from './etl/trace-replay-ingest'; +import { + findUnlinkedTraceReplayBenchmarkIds, + insertPreparedTraceReplay, +} from './etl/trace-replay-ingest'; +import { TraceReplayWorkerPool } from './etl/trace-replay-worker-pool'; import { discoverTraceReplayArtifacts } from './etl/trace-artifact-discovery'; import { datasetSlugFromBenchmarkRow } from './etl/dataset-provenance'; import { mapAggEvalRow, mapEvalRow } from './etl/eval-mapper'; @@ -223,6 +227,21 @@ function findJsonFiles(dir: string): string[] { .map((f) => path.join(dir, f)); } +async function forEachConcurrent( + items: T[], + concurrency: number, + worker: (item: T, index: number) => Promise, +): Promise { + let nextIndex = 0; + const runners = Array.from({ length: Math.min(concurrency, items.length) }, async () => { + while (nextIndex < items.length) { + const index = nextIndex++; + await worker(items[index]!, index); + } + }); + await Promise.all(runners); +} + // ── Main ──────────────────────────────────────────────────────────────────── async function main(): Promise { @@ -428,156 +447,180 @@ async function main(): Promise { const allBmkFiles = [...bmkFiles, ...allBmkDirs.flatMap((d) => findJsonFiles(d))]; console.log(` Found ${allBmkFiles.length} benchmark JSON file(s)`); - for (const [fileIndex, file] of allBmkFiles.entries()) { - const fileStart = Date.now(); - const relativeFile = path.relative(artifactsDir, file); - console.log( - ` [${fileIndex + 1}/${allBmkFiles.length}] ${relativeFile} (${formatBytes(fileSize(file))})`, - ); - const data = readJson(file); - if (!data) { - console.log(` skipped unreadable JSON (${elapsed(fileStart)})`); - continue; - } + const configuredConcurrency = Number.parseInt(process.env.INGEST_CONFIG_CONCURRENCY ?? '', 10); + const configConcurrency = Number.isFinite(configuredConcurrency) + ? Math.max(1, Math.min(4, configuredConcurrency)) + : Math.max(1, Math.min(3, os.availableParallelism())); + console.log(` Processing up to ${configConcurrency} benchmark configs concurrently`); + const traceReplayPool = + traceReplayPaths.size > 0 ? new TraceReplayWorkerPool(configConcurrency) : null; - const rawRows: Record[] = Array.isArray(data) - ? data - : [data as Record]; - console.log(` raw rows: ${rawRows.length}`); + try { + await forEachConcurrent(allBmkFiles, configConcurrency, async (file, fileIndex) => { + const fileStart = Date.now(); + const relativeFile = path.relative(artifactsDir, file); + console.log( + ` [${fileIndex + 1}/${allBmkFiles.length}] ${relativeFile} (${formatBytes(fileSize(file))})`, + ); + const data = readJson(file); + if (!data) { + console.log(` skipped unreadable JSON (${elapsed(fileStart)})`); + return; + } - for (const rawRow of rawRows) { - if (!rawRow || typeof rawRow !== 'object') continue; - const datasetSlug = datasetSlugFromBenchmarkRow(rawRow); - if (datasetSlug) datasetSlugs.add(datasetSlug); - } + const rawRows: Record[] = Array.isArray(data) + ? data + : [data as Record]; + console.log(` raw rows: ${rawRows.length}`); - const rows = rawRows - .filter((r) => typeof r === 'object' && r !== null) - .map((r) => mapBenchmarkRow(r, tracker)) - .filter((r): r is NonNullable => r !== null); + for (const rawRow of rawRows) { + if (!rawRow || typeof rawRow !== 'object') continue; + const datasetSlug = datasetSlugFromBenchmarkRow(rawRow); + if (datasetSlug) datasetSlugs.add(datasetSlug); + } - console.log(` mapped rows: ${rows.length}`); - if (rows.length === 0) { - console.log(` skipped; no mappable rows (${elapsed(fileStart)})`); - continue; - } + const rows = rawRows + .filter((r) => typeof r === 'object' && r !== null) + .map((r) => mapBenchmarkRow(r, tracker)) + .filter((r): r is NonNullable => r !== null); - const toInsert = []; - for (const row of rows) { - try { - const configId = await getOrCreateConfig(row.config); - toInsert.push({ ...row, configId }); - } catch (error: any) { - tracker.recordDbError(`config for ${path.basename(file)}`, error); + console.log(` mapped rows: ${rows.length}`); + if (rows.length === 0) { + console.log(` skipped; no mappable rows (${elapsed(fileStart)})`); + return; } - } - console.log(` rows with resolved configs: ${toInsert.length}`); - - if (toInsert.length > 0) { - try { - const insertStart = Date.now(); - const { newCount, dupCount, insertedIds } = await bulkIngestBenchmarkRows( - sql, - toInsert, - workflowRunId, - date, - ); - console.log( - ` benchmark rows: +${newCount} new, ${dupCount} dup, ` + - `${insertedIds.length} id(s) (${elapsed(insertStart)})`, - ); - totalNewBmk += newCount; - totalDupBmk += dupCount; - - // Build availability only after successful insert - for (const r of toInsert) { - availRows.push({ - model: r.config.model, - isl: r.isl, - osl: r.osl, - precision: r.config.precision, - hardware: r.config.hardware, - framework: r.config.framework, - specMethod: r.config.specMethod, - disagg: r.config.disagg, - benchmarkType: r.benchmarkType, - }); + + const toInsert = []; + for (const row of rows) { + try { + const configId = await getOrCreateConfig(row.config); + toInsert.push({ ...row, configId }); + } catch (error: any) { + tracker.recordDbError(`config for ${path.basename(file)}`, error); } + } + console.log(` rows with resolved configs: ${toInsert.length}`); + + if (toInsert.length > 0) { + try { + const insertStart = Date.now(); + const { newCount, dupCount, insertedIds } = await bulkIngestBenchmarkRows( + sql, + toInsert, + workflowRunId, + date, + ); + console.log( + ` benchmark rows: +${newCount} new, ${dupCount} dup, ` + + `${insertedIds.length} id(s) (${elapsed(insertStart)})`, + ); + totalNewBmk += newCount; + totalDupBmk += dupCount; + + // Build availability only after successful insert + for (const r of toInsert) { + availRows.push({ + model: r.config.model, + isl: r.isl, + osl: r.osl, + precision: r.config.precision, + hardware: r.config.hardware, + framework: r.config.framework, + specMethod: r.config.specMethod, + disagg: r.config.disagg, + benchmarkType: r.benchmarkType, + }); + } - const parentDir = path.basename(path.dirname(file)); - if (parentDir.startsWith('bmk_') && insertedIds.length > 0) { - // Single-turn artifacts are `bmk_` paired with - // `server_logs_`. Agentic artifacts are `bmk_agentic_` - // but the server log is still `server_logs_` (no `agentic_` - // prefix), so fall back to the fully-stripped suffix — otherwise - // agentic rows never get their server log (and KV-pool size) linked. - const configKey = parentDir.replace(/^bmk_/u, ''); - const logPath = - serverLogPaths.get(configKey) ?? - serverLogPaths.get(stripBmkAndAgenticPrefix(parentDir)); - if (logPath) { - try { - const serverLogStart = Date.now(); - console.log( - ` server_log ${path.basename(logPath)} (${formatBytes(fileSize(logPath))})`, - ); - const serverLog = fs.readFileSync(logPath, 'utf8').replaceAll('\u0000', ''); - await insertServerLog(sql, insertedIds, serverLog); - console.log(` server_log linked (${elapsed(serverLogStart)})`); - } catch (error: any) { - tracker.recordDbError(`server_log for ${configKey}`, error); + const parentDir = path.basename(path.dirname(file)); + if (parentDir.startsWith('bmk_') && insertedIds.length > 0) { + // Single-turn artifacts are `bmk_` paired with + // `server_logs_`. Agentic artifacts are `bmk_agentic_` + // but the server log is still `server_logs_` (no `agentic_` + // prefix), so fall back to the fully-stripped suffix — otherwise + // agentic rows never get their server log (and KV-pool size) linked. + const configKey = parentDir.replace(/^bmk_/u, ''); + const logPath = + serverLogPaths.get(configKey) ?? + serverLogPaths.get(stripBmkAndAgenticPrefix(parentDir)); + if (logPath) { + try { + const serverLogStart = Date.now(); + console.log( + ` server_log ${path.basename(logPath)} (${formatBytes(fileSize(logPath))})`, + ); + const serverLog = fs.readFileSync(logPath, 'utf8').replaceAll('\u0000', ''); + await insertServerLog(sql, insertedIds, serverLog); + console.log(` server_log linked (${elapsed(serverLogStart)})`); + } catch (error: any) { + tracker.recordDbError(`server_log for ${configKey}`, error); + } } } - } - // Trace-replay sibling lookup for agentic points only. The aiperf - // harness emits `agentic_/trace_replay/...` next to the - // `bmk_agentic_` artifact we just ingested. - if (parentDir.startsWith('bmk_agentic_') && insertedIds.length > 0) { - const suffix = stripBmkAndAgenticPrefix(parentDir); - const concMatch = path.basename(file).match(/_conc(?\d+)\.json$/u); - const trace = - (concMatch?.groups?.conc - ? traceReplayPaths.get(`${suffix}|${concMatch.groups.conc}`) - : undefined) ?? traceReplayPaths.get(suffix); - if (trace) { - try { - const traceStart = Date.now(); - console.log( - ` trace_replay ${suffix}: ` + - `profile=${formatBytes(fileSize(trace.profileJsonl))}, ` + - `server_csv=${formatBytes(fileSize(trace.serverMetricsCsv))}, ` + - `server_json=${formatBytes(fileSize(trace.serverMetricsJson))}`, - ); - const profile = trace.profileJsonl ? fs.readFileSync(trace.profileJsonl) : null; - const metrics = trace.serverMetricsCsv - ? fs.readFileSync(trace.serverMetricsCsv) - : null; - const metricsJson = trace.serverMetricsJson - ? fs.readFileSync(trace.serverMetricsJson) - : null; - await insertTraceReplay(sql, insertedIds, profile, metrics, metricsJson, { - metricsContext: { - framework: toInsert[0]?.config.framework, - disagg: toInsert[0]?.config.disagg, - }, - progressLabel: suffix, - }); - totalTraceReplayLinked += insertedIds.length; - console.log(` trace_replay ${suffix}: done (${elapsed(traceStart)})`); - } catch (error: any) { - tracker.recordDbError(`trace_replay for ${suffix}`, error); + // Trace-replay sibling lookup for agentic points only. The aiperf + // harness emits `agentic_/trace_replay/...` next to the + // `bmk_agentic_` artifact we just ingested. + if (parentDir.startsWith('bmk_agentic_') && insertedIds.length > 0) { + const suffix = stripBmkAndAgenticPrefix(parentDir); + const concMatch = path.basename(file).match(/_conc(?\d+)\.json$/u); + const trace = + (concMatch?.groups?.conc + ? traceReplayPaths.get(`${suffix}|${concMatch.groups.conc}`) + : undefined) ?? traceReplayPaths.get(suffix); + if (trace) { + try { + const traceStart = Date.now(); + console.log( + ` trace_replay ${suffix}: ` + + `profile=${formatBytes(fileSize(trace.profileJsonl))}, ` + + `server_csv=${formatBytes(fileSize(trace.serverMetricsCsv))}, ` + + `server_json=${formatBytes(fileSize(trace.serverMetricsJson))}`, + ); + const options = { + metricsContext: { + framework: toInsert[0]?.config.framework, + disagg: toInsert[0]?.config.disagg, + }, + progressLabel: suffix, + }; + const unlinkedIds = await findUnlinkedTraceReplayBenchmarkIds(sql, insertedIds); + if (unlinkedIds.length === 0) { + console.log( + ` trace_replay ${suffix}: skipping; benchmark rows already linked`, + ); + } else { + const prepared = await traceReplayPool!.prepare({ + profilePath: trace.profileJsonl ?? null, + serverMetricsCsvPath: trace.serverMetricsCsv ?? null, + serverMetricsJsonPath: trace.serverMetricsJson ?? null, + metricsContext: options.metricsContext, + }); + console.log( + ` trace_replay ${suffix}: prepared in ` + + `${((prepared.gzipMs + prepared.computeMs) / 1000).toFixed(1)}s ` + + `(worker CPU)`, + ); + await insertPreparedTraceReplay(sql, insertedIds, prepared, options); + } + totalTraceReplayLinked += insertedIds.length; + console.log(` trace_replay ${suffix}: done (${elapsed(traceStart)})`); + } catch (error: any) { + tracker.recordDbError(`trace_replay for ${suffix}`, error); + } + } else { + console.log(` trace_replay ${suffix}: missing sibling artifact`); + tracker.skips.traceReplayMissing++; } - } else { - console.log(` trace_replay ${suffix}: missing sibling artifact`); - tracker.skips.traceReplayMissing++; } + } catch (error: any) { + tracker.recordDbError(path.basename(file), error); } - } catch (error: any) { - tracker.recordDbError(path.basename(file), error); } - } - console.log(` finished ${relativeFile} (${elapsed(fileStart)})`); + console.log(` finished ${relativeFile} (${elapsed(fileStart)})`); + }); + } finally { + await traceReplayPool?.close(); } console.log(` Benchmarks: +${totalNewBmk} new, ${totalDupBmk} dup`); if (totalTraceReplayLinked > 0 || tracker.skips.traceReplayMissing > 0) {