From feef4a9cbc7c1d523adccaea0b9e5f62131bb2b6 Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Wed, 5 Aug 2026 10:02:13 -0500 Subject: [PATCH 1/2] perf(db): streamline trace metric parsing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the dual token fan-out with one packed-token consumer that only assembles selected metric values. Add regression coverage for nested lookalikes, scalar and compound values, duplicate keys, malformed input, and both metric phases.\n\n中文：优化 trace 指标解析。用单一 packed-token 消费器替代双路 token 分发，仅组装需要的指标值；补充嵌套同名字段、标量与复合值、重复键、异常输入及两个指标阶段的回归测试。 --- packages/db/src/etl/gzip-json-stream.test.ts | 48 ++++++ packages/db/src/etl/gzip-json-stream.ts | 156 ++++++++++++++----- 2 files changed, 163 insertions(+), 41 deletions(-) diff --git a/packages/db/src/etl/gzip-json-stream.test.ts b/packages/db/src/etl/gzip-json-stream.test.ts index 30aa9fea..b0ba1326 100644 --- a/packages/db/src/etl/gzip-json-stream.test.ts +++ b/packages/db/src/etl/gzip-json-stream.test.ts @@ -103,6 +103,54 @@ describe('collectMetricPhases', () => { }); }); + it('ignores nested phase lookalikes and preserves selected scalar and compound values', async () => { + const adversarial = gzipSync( + JSON.stringify({ + metadata: { + metrics: { wanted: { nested: 'must not be selected' } }, + warmup_metrics: { wanted: { nested: 'must not be selected' } }, + }, + metrics: { + wanted: [1, { nested: true }, null], + scalar: 42, + ignored: { very: { deeply: ['nested', 'value'] } }, + }, + warmup_metrics: { + wanted: false, + scalar: 'warmup', + ignored: [1, 2, 3], + }, + }), + ); + + await expect( + collectMetricPhases(adversarial, new Set(['wanted', 'scalar']), 1), + ).resolves.toEqual({ + metrics: { + wanted: [1, { nested: true }, null], + scalar: 42, + }, + warmupMetrics: { + wanted: false, + scalar: 'warmup', + }, + complete: false, + }); + }); + + it('matches JSON semantics for duplicate phase and metric keys', async () => { + const duplicateKeys = gzipSync( + '{"metrics":{"wanted":1,"wanted":2},"metrics":{"wanted":3},' + + '"warmup_metrics":{"wanted":4,"wanted":5}}', + ); + + await expect(collectMetricPhases(duplicateKeys, new Set(['wanted']), 1)).resolves.toEqual({ + metrics: { wanted: 3 }, + warmupMetrics: { wanted: 5 }, + complete: false, + }); + }); + it('rejects malformed gzip input on both paths', async () => { await expect( collectMetricPhases(Buffer.from('not gzip'), new Set(['wanted']), 1), diff --git a/packages/db/src/etl/gzip-json-stream.ts b/packages/db/src/etl/gzip-json-stream.ts index 48723f57..f6d7e6af 100644 --- a/packages/db/src/etl/gzip-json-stream.ts +++ b/packages/db/src/etl/gzip-json-stream.ts @@ -7,13 +7,14 @@ * stream-json pipeline collects only the top-level subtrees callers need. */ -import { PassThrough, Readable } from 'node:stream'; -import { pipeline } from 'node:stream/promises'; +import { Readable } from 'node:stream'; import { createGunzip, gunzipSync } from 'node:zlib'; import { chain } from 'stream-chain'; import { parser } from 'stream-json'; +import Assembler from 'stream-json/assembler.js'; +import type { Token } from 'stream-json/parser.js'; import { pick } from 'stream-json/filters/pick.js'; import { streamObject } from 'stream-json/streamers/stream-object.js'; @@ -82,24 +83,122 @@ export interface MetricPhaseMaps { complete: boolean; } -async function collectTokenBranch( - input: PassThrough, - filter: 'metrics' | 'warmup_metrics', +function isValueStart(token: Token): boolean { + return ( + token.name === 'startObject' || + token.name === 'startArray' || + token.name === 'stringValue' || + token.name === 'numberValue' || + token.name === 'nullValue' || + token.name === 'trueValue' || + token.name === 'falseValue' + ); +} + +function updateDepth(depth: number, token: Token): number { + if (token.name === 'startObject' || token.name === 'startArray') return depth + 1; + if (token.name === 'endObject' || token.name === 'endArray') return depth - 1; + return depth; +} + +/** + * Consume the parser token stream once and assemble only selected direct + * children of the two metric phase objects. Unselected metric subtrees still + * have to be tokenized so the JSON can be validated, but they are never copied + * to secondary streams or materialized as JavaScript objects. + */ +async function streamCollectMetricPhases( + buffer: Buffer, wanted: ReadonlySet, -): Promise> { - const collected: Record = {}; - const output = chain([input, pick({ filter }), streamObject()]); - for await (const chunk of output) { - const { key, value } = chunk as { key: string; value: T }; - if (wanted.has(key)) collected[key] = value; +): Promise> { + let metrics: Record = {}; + let warmupMetrics: Record = {}; + let depth = 0; + let phase: 'metrics' | 'warmup_metrics' | null = null; + let topLevelKey: string | null = null; + let metricKey: string | null = null; + let valueAssembler: Assembler | null = null; + let valueTarget: Record | null = null; + let valueKey: string | null = null; + + // Packed-only values avoid emitting start/chunk/end tokens in addition to + // their complete value token. This materially reduces the token count for + // metric-heavy multi-GiB documents. + const tokens = chain([ + Readable.from(buffer), + createGunzip(), + parser({ + packKeys: true, + packStrings: true, + packNumbers: true, + streamKeys: false, + streamStrings: false, + streamNumbers: false, + }), + ]); + + for await (const rawToken of tokens) { + const token = rawToken as Token; + const depthBefore = depth; + + if (valueAssembler) { + valueAssembler.consume(token); + depth = updateDepth(depth, token); + if (valueAssembler.done) { + valueTarget![valueKey!] = valueAssembler.current as T; + valueAssembler = null; + valueTarget = null; + valueKey = null; + } + continue; + } + + if (token.name === 'keyValue') { + if (depthBefore === 1) { + topLevelKey = token.value; + } else if (depthBefore === 2 && phase) { + metricKey = token.value; + } + } else if (isValueStart(token)) { + if (depthBefore === 1 && topLevelKey !== null) { + if (token.name === 'startObject' && topLevelKey === 'metrics') { + metrics = {}; + phase = 'metrics'; + } else if (token.name === 'startObject' && topLevelKey === 'warmup_metrics') { + warmupMetrics = {}; + phase = 'warmup_metrics'; + } + topLevelKey = null; + } else if (depthBefore === 2 && phase && metricKey !== null) { + if (wanted.has(metricKey)) { + valueTarget = phase === 'metrics' ? metrics : warmupMetrics; + valueKey = metricKey; + valueAssembler = new Assembler(); + valueAssembler.consume(token); + if (valueAssembler.done) { + valueTarget[valueKey] = valueAssembler.current as T; + valueAssembler = null; + valueTarget = null; + valueKey = null; + } + } + metricKey = null; + } + } + + depth = updateDepth(depth, token); + if (phase && depth === 1 && (token.name === 'endObject' || token.name === 'endArray')) { + phase = null; + metricKey = null; + } } - return collected; + + return { metrics, warmupMetrics, complete: false }; } /** - * Gunzip and parse both server-metric phase blocks once. Large documents fan - * the parser's token stream out to two lightweight selectors, avoiding one - * complete decompression + JSON tokenization pass per phase. + * Gunzip and parse both server-metric phase blocks once. Large documents use a + * single token consumer which materializes only the selected metric values. */ export async function collectMetricPhases( buffer: Buffer, @@ -119,30 +218,5 @@ export async function collectMetricPhases( }; } - // Attach every branch before starting the source pipeline so no parser - // tokens can be missed. PassThrough backpressure keeps the two consumers in - // lockstep without buffering the full document. - const profilingInput = new PassThrough({ objectMode: true }); - const warmupInput = new PassThrough({ objectMode: true }); - const tokenTee = new PassThrough({ objectMode: true }); - tokenTee.pipe(profilingInput); - tokenTee.pipe(warmupInput); - - const profiling = collectTokenBranch(profilingInput, 'metrics', wanted); - const warmup = collectTokenBranch(warmupInput, 'warmup_metrics', wanted); - const tokens = chain([Readable.from(buffer), createGunzip(), parser()]); - - try { - const [, metrics, warmupMetrics] = await Promise.all([ - pipeline(tokens, tokenTee), - profiling, - warmup, - ]); - return { metrics, warmupMetrics, complete: false }; - } catch (error) { - tokenTee.destroy(); - profilingInput.destroy(); - warmupInput.destroy(); - throw error; - } + return await streamCollectMetricPhases(buffer, wanted); } From f4e45187561f25c3a798b81923ab0f5e55e91dc7 Mon Sep 17 00:00:00 2001 From: Cam Quilici Date: Wed, 5 Aug 2026 10:18:06 -0500 Subject: [PATCH 2/2] perf(db): parallelize trace replay preparation MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Split CPU-heavy trace preparation from atomic persistence, run it in a vCPU-scaled bounded worker pool, cap concurrent uploads, and recheck links under row locks. Add worker transfer, queue-bound, failure-recovery, and concurrent-ingest regression coverage.\n\n中文：并行化 trace replay 预处理。将 CPU 密集型预处理与原子化持久化拆分，使用按 vCPU 数量伸缩的有界 worker 池执行，限制并发上传，并在行锁内重新检查关联状态；补充 worker 数据传输、队列内存上限、失败恢复及并发摄取回归测试。 --- packages/db/src/etl/async-semaphore.ts | 38 ++ .../db/src/etl/trace-replay-ingest.test.ts | 71 ++++ packages/db/src/etl/trace-replay-ingest.ts | 333 +++++++++++------- .../src/etl/trace-replay-worker-pool.test.ts | 186 ++++++++++ .../db/src/etl/trace-replay-worker-pool.ts | 181 ++++++++++ .../src/etl/trace-replay-worker-protocol.ts | 29 ++ .../db/src/etl/trace-replay-worker-smoke.ts | 108 ++++++ packages/db/src/etl/trace-replay-worker.ts | 63 ++++ packages/db/src/ingest-ci-run.ts | 93 ++++- 9 files changed, 967 insertions(+), 135 deletions(-) create mode 100644 packages/db/src/etl/async-semaphore.ts create mode 100644 packages/db/src/etl/trace-replay-worker-pool.test.ts create mode 100644 packages/db/src/etl/trace-replay-worker-pool.ts create mode 100644 packages/db/src/etl/trace-replay-worker-protocol.ts create mode 100644 packages/db/src/etl/trace-replay-worker-smoke.ts create mode 100644 packages/db/src/etl/trace-replay-worker.ts diff --git a/packages/db/src/etl/async-semaphore.ts b/packages/db/src/etl/async-semaphore.ts new file mode 100644 index 00000000..5d4a8344 --- /dev/null +++ b/packages/db/src/etl/async-semaphore.ts @@ -0,0 +1,38 @@ +/** Minimal FIFO semaphore for bounding concurrent asynchronous operations. */ +export class AsyncSemaphore { + private active = 0; + private readonly waiters: (() => void)[] = []; + private readonly limit: number; + + constructor(limit: number) { + if (!Number.isSafeInteger(limit) || limit < 1) { + throw new Error(`Semaphore limit must be a positive integer, received ${limit}`); + } + this.limit = limit; + } + + async run(operation: () => Promise): Promise { + await this.acquire(); + try { + return await operation(); + } finally { + this.release(); + } + } + + private async acquire(): Promise { + if (this.active < this.limit) { + this.active += 1; + return; + } + await new Promise((resolve) => { + this.waiters.push(resolve); + }); + this.active += 1; + } + + private release(): void { + this.active -= 1; + this.waiters.shift()?.(); + } +} diff --git a/packages/db/src/etl/trace-replay-ingest.test.ts b/packages/db/src/etl/trace-replay-ingest.test.ts index fd0b3fd3..a77d87aa 100644 --- a/packages/db/src/etl/trace-replay-ingest.test.ts +++ b/packages/db/src/etl/trace-replay-ingest.test.ts @@ -9,6 +9,8 @@ import { describe, expect, it } from 'vitest'; import { TRACE_REPLAY_UPLOAD_CHUNK_BYTES, gzipTraceReplayInput, + persistPreparedTraceReplay, + type PreparedTraceReplay, uploadTraceReplayPayloadChunks, } from './trace-replay-ingest'; @@ -26,6 +28,45 @@ function mockTransactionSql(): { sql: postgres.TransactionSql; calls: SqlCall[] return { sql, calls }; } +function preparedFixture(): PreparedTraceReplay { + return { + profileGz: Buffer.from('profile'), + profileSize: 70, + serverMetricsCsv: Buffer.from('csv'), + serverMetricsCsvSize: 30, + serverMetricsJsonGz: Buffer.from('metrics'), + serverMetricsJsonSize: 700, + aggregateStatsJson: Buffer.from('{"version":1}'), + chartSeriesJson: Buffer.from('{"version":12}'), + requestTimelineJson: Buffer.from('{"version":1}'), + chartWindows: 2, + timelineRequests: 3, + compressionMs: 10, + computeMs: 20, + cacheHitRates: null, + }; +} + +function mockSqlWithTransaction(lockedRows: { id: number }[]): { + sql: Parameters[0]; + calls: SqlCall[]; +} { + const calls: SqlCall[] = []; + const execute = (strings: TemplateStringsArray, ...values: unknown[]) => { + const text = strings.join('?'); + calls.push({ text, values }); + if (text.includes('for update')) return Promise.resolve(lockedRows); + if (text.includes('insert into agentic_trace_replay')) return Promise.resolve([{ id: 123 }]); + return Promise.resolve([]); + }; + const tx = Object.assign(execute, { array: (values: unknown[]) => values }); + const sql = Object.assign(execute, { + array: (values: unknown[]) => values, + begin: (operation: (transaction: typeof tx) => Promise) => operation(tx), + }) as unknown as Parameters[0]; + return { sql, calls }; +} + describe('uploadTraceReplayPayloadChunks', () => { it('bounds every Bind payload for the measured 90 MiB staging row', async () => { // Exact payload sizes from InferenceX run 29181694248, item 8/9. @@ -80,3 +121,33 @@ describe('gzipTraceReplayInput', () => { } }); }); + +describe('persistPreparedTraceReplay', () => { + it('rechecks links under a row lock and avoids creating an orphan after a concurrent ingest', async () => { + const { sql, calls } = mockSqlWithTransaction([]); + + await expect(persistPreparedTraceReplay(sql, [41], preparedFixture())).resolves.toBe(0); + expect(calls).toHaveLength(1); + expect(calls[0].text).toContain('for update'); + expect(calls.some((call) => call.text.includes('insert into agentic_trace_replay'))).toBe( + false, + ); + }); + + it('uploads every payload and links exactly the rows locked by the transaction', async () => { + const { sql, calls } = mockSqlWithTransaction([{ id: 41 }]); + + await expect(persistPreparedTraceReplay(sql, [41], preparedFixture())).resolves.toBe(1); + const lockIndex = calls.findIndex((call) => call.text.includes('for update')); + const blobIndex = calls.findIndex((call) => + call.text.includes('insert into agentic_trace_replay'), + ); + const linkCall = calls.find((call) => call.text.includes('set trace_replay_id')); + expect(lockIndex).toBe(0); + expect(blobIndex).toBeGreaterThan(lockIndex); + expect( + calls.filter((call) => call.text.includes('trace_replay_upload_parts (field, part, data)')), + ).toHaveLength(6); + expect(linkCall?.values.some((value) => Array.isArray(value) && value.includes(41))).toBe(true); + }); +}); diff --git a/packages/db/src/etl/trace-replay-ingest.ts b/packages/db/src/etl/trace-replay-ingest.ts index 5e6fa65e..bd8c0403 100644 --- a/packages/db/src/etl/trace-replay-ingest.ts +++ b/packages/db/src/etl/trace-replay-ingest.ts @@ -46,6 +46,23 @@ export interface TraceReplayIngestOptions { progressLabel?: string; } +export interface PreparedTraceReplay { + profileGz: Buffer | null; + profileSize: number | null; + serverMetricsCsv: Buffer | null; + serverMetricsCsvSize: number | null; + serverMetricsJsonGz: Buffer | null; + serverMetricsJsonSize: number | null; + aggregateStatsJson: Buffer | null; + chartSeriesJson: Buffer | null; + requestTimelineJson: Buffer | null; + chartWindows: number; + timelineRequests: number; + compressionMs: number; + computeMs: number; + cacheHitRates: { gpu: number; cpu: number | null } | null; +} + function formatBytes(bytes: number | null | undefined): string { if (bytes === null || bytes === undefined) return 'none'; if (bytes < 1024) return `${bytes} B`; @@ -65,6 +82,30 @@ function jsonBuffer(value: unknown | null): Buffer | null { return Buffer.from(JSON.stringify(structuredClone(value)), 'utf8'); } +function cacheHitRatesFromChartSeries( + chartSeries: Awaited>['chartSeries'], +): PreparedTraceReplay['cacheHitRates'] { + if (!chartSeries || chartSeries.prefillTps.length === 0) return null; + const sumPrompts = chartSeries.prefillTps.reduce((sum, point) => sum + point.value, 0); + if (!(sumPrompts > 0)) return null; + + const sumOf = (name: string): number => + (chartSeries.promptTokensBySource[name] ?? []).reduce((sum, point) => sum + point.value, 0); + // Preserve the historical source aliases exactly: SGLang hicache reports + // HBM/CPU labels while vLLM LMCache reports local/external transfer labels. + const cpuHits = sumOf('cache hit (CPU offload)') + sumOf('external_kv_transfer'); + const hbmFromBreakdown = sumOf('cache hit (HBM)') + sumOf('cache hit') + sumOf('local_cache_hit'); + const gpuHits = + hbmFromBreakdown > 0 + ? hbmFromBreakdown + : chartSeries.prefixCacheHitsTps.reduce((sum, point) => sum + point.value, 0); + + return { + gpu: gpuHits / sumPrompts, + cpu: cpuHits > 0 ? cpuHits / sumPrompts : null, + }; +} + /** * Gzip a trace input without materializing file-backed GiB-scale exports in * Node's heap. Buffer inputs remain supported for callers and unit tests. @@ -118,98 +159,108 @@ export async function uploadTraceReplayPayloadChunks( return part; } -/** - * Persist the per-point trace files and link them to `benchmarkResultIds`. - * - * @param sql Active `postgres` connection. - * @param benchmarkResultIds DB ids of the benchmark_results rows produced by - * the same `bmk_agentic_` artifact whose - * sibling `agentic_` directory holds these - * trace files. - * @param profileExportJsonl Raw bytes or a file path for `profile_export.jsonl`. - * Gzipped before storage; file paths stream from disk. - * @param serverMetricsCsv Raw bytes or a file path for `server_metrics_export.csv`. - * Stored as-is. - * @param serverMetricsJson Raw bytes or a file path for `server_metrics_export.json` — - * per-scrape time-series of every Prometheus metric. - * Optional, streamed and gzipped before storage (~42x ratio). - * @param options Canonical framework/disagg context plus optional - * progress label for CI logs. - */ -export async function insertTraceReplay( +/** Resolve benchmark rows which still need a trace sidecar before scheduling CPU work. */ +export async function findUnlinkedTraceReplayIds( sql: Sql, benchmarkResultIds: number[], - profileExportJsonl: TraceReplayInput, - serverMetricsCsv: TraceReplayInput, - serverMetricsJson: TraceReplayInput = null, - options: TraceReplayIngestOptions = {}, -): Promise { - const { metricsContext = {}, progressLabel } = options; - const log = (message: string): void => { - if (progressLabel) console.log(` trace_replay ${progressLabel}: ${message}`); - }; - - 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 }[]>` +): 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 `; - log(`found ${unlinked.length} unlinked row(s) (${elapsed(linkStart)})`); - if (unlinked.length === 0) { - log('skipping blob insert; all benchmark rows already linked'); - return; - } + return rows.map((row) => row.id); +} - const gzipStart = Date.now(); - log('reading and compressing trace inputs'); +/** + * Perform the file IO and CPU-heavy derivation independently of Postgres so + * callers can execute this stage in worker threads. + */ +export async function prepareTraceReplay( + profileExportJsonl: TraceReplayInput, + serverMetricsCsv: TraceReplayInput, + serverMetricsJson: TraceReplayInput = null, + metricsContext: ServerMetricsContext = {}, +): Promise { + const compressionStart = Date.now(); const [profile, csv, metricsJson] = await Promise.all([ gzipTraceReplayInput(profileExportJsonl), readTraceReplayInput(serverMetricsCsv), gzipTraceReplayInput(serverMetricsJson), ]); - const profileGz = profile.data; - const profileSize = profile.sourceSize; - const serverMetricsCsvData = csv.data; - const csvSize = csv.sourceSize; - const metricsJsonGz = metricsJson.data; - const metricsJsonSize = metricsJson.sourceSize; - log( - `compressed profile=${formatBytes(profileSize)} -> ${formatBytes(profileGz?.length)}, ` + - `server_csv=${formatBytes(csvSize)}, ` + - `server_json=${formatBytes(metricsJsonSize)} -> ${formatBytes(metricsJsonGz?.length)} ` + - `(${elapsed(gzipStart)})`, - ); + const compressionMs = Date.now() - compressionStart; - // 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 computeTraceDerivedPayloads( - profileGz, - metricsJsonGz, + profile.data, + metricsJson.data, metricsContext, ); - log( - `computed derived JSON: chart_windows=${chartSeries?.timeslicesCount ?? 0}, ` + - `timeline_requests=${requestTimeline?.requests.length ?? 0} (${elapsed(computeStart)})`, - ); + const computeMs = Date.now() - computeStart; + + return { + profileGz: profile.data, + profileSize: profile.sourceSize, + serverMetricsCsv: csv.data, + serverMetricsCsvSize: csv.sourceSize, + serverMetricsJsonGz: metricsJson.data, + serverMetricsJsonSize: metricsJson.sourceSize, + aggregateStatsJson: jsonBuffer(aggregateStats), + chartSeriesJson: jsonBuffer(chartSeries), + requestTimelineJson: jsonBuffer(requestTimeline), + chartWindows: chartSeries?.timeslicesCount ?? 0, + timelineRequests: requestTimeline?.requests.length ?? 0, + compressionMs, + computeMs, + cacheHitRates: cacheHitRatesFromChartSeries(chartSeries), + }; +} + +/** + * Persist a prepared trace atomically. Rows are locked and rechecked here so + * concurrent ingest attempts cannot both attach different trace sidecars. + */ +export async function persistPreparedTraceReplay( + sql: Sql, + benchmarkResultIds: number[], + prepared: PreparedTraceReplay, + options: Pick = {}, +): Promise { + const { progressLabel } = options; + const log = (message: string): void => { + if (progressLabel) console.log(` trace_replay ${progressLabel}: ${message}`); + }; + if (benchmarkResultIds.length === 0) return 0; - const aggregateStatsJson = jsonBuffer(aggregateStats); - const chartSeriesJson = jsonBuffer(chartSeries); - const requestTimelineJson = jsonBuffer(requestTimeline); + const { + profileGz, + profileSize, + serverMetricsCsv, + serverMetricsCsvSize, + serverMetricsJsonGz, + serverMetricsJsonSize, + aggregateStatsJson, + chartSeriesJson, + requestTimelineJson, + cacheHitRates, + } = prepared; + let linkedCount = 0; const insertStart = Date.now(); log(`uploading trace_replay payloads in ${formatBytes(TRACE_REPLAY_UPLOAD_CHUNK_BYTES)} chunks`); await sql.begin(async (tx) => { + 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; rows were linked by another ingest'); + return; + } + await tx` create temporary table trace_replay_upload_parts ( field text not null, @@ -221,8 +272,8 @@ export async function insertTraceReplay( const payloads: [TraceReplayUploadField, Buffer | null][] = [ ['profile_export_jsonl_gz', profileGz], - ['server_metrics_csv', serverMetricsCsvData], - ['server_metrics_json_gz', metricsJsonGz], + ['server_metrics_csv', serverMetricsCsv], + ['server_metrics_json_gz', serverMetricsJsonGz], ['aggregate_stats', aggregateStatsJson], ['chart_series', chartSeriesJson], ['request_timeline', requestTimelineJson], @@ -261,13 +312,13 @@ export async function insertTraceReplay( from pg_temp.trace_replay_upload_parts where field = 'server_metrics_csv' ), - ${csvSize}, + ${serverMetricsCsvSize}, ( select string_agg(data, ''::bytea order by part) from pg_temp.trace_replay_upload_parts where field = 'server_metrics_json_gz' ), - ${metricsJsonSize}, + ${serverMetricsJsonSize}, ( select convert_from(string_agg(data, ''::bytea order by part), 'UTF8')::jsonb from pg_temp.trace_replay_upload_parts @@ -293,54 +344,98 @@ export async function insertTraceReplay( await tx` update benchmark_results set trace_replay_id = ${traceReplayId} - where id = any(${tx.array(unlinked.map((r) => r.id))}::bigint[]) + where id = any(${tx.array(unlinked.map((row) => row.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 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'); - } + if (cacheHitRates) { + await tx` + update benchmark_results + set metrics = jsonb_set( + case when ${cacheHitRates.cpu}::numeric is not null + then jsonb_set( + metrics, + '{server_cpu_cache_hit_rate}', + to_jsonb(${cacheHitRates.cpu}::numeric) + ) + else metrics + end, + '{server_gpu_cache_hit_rate}', + to_jsonb(${cacheHitRates.gpu}::numeric) + ) + where id = any(${tx.array(unlinked.map((row) => row.id))}::bigint[]) + `; + log('updated cache-hit metrics from chart series'); } + linkedCount = unlinked.length; }); - log(`inserted trace_replay payload (${elapsed(insertStart)})`); + if (linkedCount > 0) log(`inserted trace_replay payload (${elapsed(insertStart)})`); + return linkedCount; +} + +/** + * Persist the per-point trace files and link them to `benchmarkResultIds`. + * + * @param sql Active `postgres` connection. + * @param benchmarkResultIds DB ids of the benchmark_results rows produced by + * the same `bmk_agentic_` artifact whose + * sibling `agentic_` directory holds these + * trace files. + * @param profileExportJsonl Raw bytes or a file path for `profile_export.jsonl`. + * Gzipped before storage; file paths stream from disk. + * @param serverMetricsCsv Raw bytes or a file path for `server_metrics_export.csv`. + * Stored as-is. + * @param serverMetricsJson Raw bytes or a file path for `server_metrics_export.json` — + * per-scrape time-series of every Prometheus metric. + * Optional, streamed and gzipped before storage (~42x ratio). + * @param options Canonical framework/disagg context plus optional + * progress label for CI logs. + */ +export async function insertTraceReplay( + sql: Sql, + benchmarkResultIds: number[], + profileExportJsonl: TraceReplayInput, + serverMetricsCsv: TraceReplayInput, + serverMetricsJson: TraceReplayInput = null, + options: TraceReplayIngestOptions = {}, +): Promise { + const { metricsContext = {}, progressLabel } = options; + const log = (message: string): void => { + if (progressLabel) console.log(` trace_replay ${progressLabel}: ${message}`); + }; + + 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 unlinkedIds = await findUnlinkedTraceReplayIds(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; + } + + log('reading and compressing trace inputs'); + const prepared = await prepareTraceReplay( + profileExportJsonl, + serverMetricsCsv, + serverMetricsJson, + metricsContext, + ); + log( + `compressed profile=${formatBytes(prepared.profileSize)} -> ${formatBytes(prepared.profileGz?.length)}, ` + + `server_csv=${formatBytes(prepared.serverMetricsCsvSize)}, ` + + `server_json=${formatBytes(prepared.serverMetricsJsonSize)} -> ` + + `${formatBytes(prepared.serverMetricsJsonGz?.length)} ` + + `(${(prepared.compressionMs / 1000).toFixed(1)}s)`, + ); + log( + `computed derived JSON: chart_windows=${prepared.chartWindows}, ` + + `timeline_requests=${prepared.timelineRequests} ` + + `(${(prepared.computeMs / 1000).toFixed(1)}s)`, + ); + await persistPreparedTraceReplay(sql, unlinkedIds, prepared, { progressLabel }); } diff --git a/packages/db/src/etl/trace-replay-worker-pool.test.ts b/packages/db/src/etl/trace-replay-worker-pool.test.ts new file mode 100644 index 00000000..e85ef970 --- /dev/null +++ b/packages/db/src/etl/trace-replay-worker-pool.test.ts @@ -0,0 +1,186 @@ +import { execFile } from 'node:child_process'; +import { createHash } from 'node:crypto'; +import { mkdtemp, rm, writeFile } from 'node:fs/promises'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { fileURLToPath } from 'node:url'; +import { promisify } from 'node:util'; +import { gunzipSync } from 'node:zlib'; + +import { afterEach, describe, expect, it } from 'vitest'; + +import { prepareTraceReplay, type PreparedTraceReplay } from './trace-replay-ingest'; +import { AsyncSemaphore } from './async-semaphore'; +import { resolveTraceReplayWorkerCount } from './trace-replay-worker-pool'; + +const tempDirs: string[] = []; +const execFileAsync = promisify(execFile); + +afterEach(async () => { + await Promise.all(tempDirs.splice(0).map((dir) => rm(dir, { recursive: true, force: true }))); +}); + +function metric(rate: number) { + return { + series: [ + { + endpoint_url: 'worker.test:8000', + labels: {}, + timeslices: [{ start_ns: 1e9, end_ns: 2e9, rate }], + }, + ], + }; +} + +async function traceFixture() { + const dir = await mkdtemp(join(tmpdir(), 'trace-worker-test-')); + tempDirs.push(dir); + const profile = join(dir, 'profile_export.jsonl'); + const csv = join(dir, 'server_metrics_export.csv'); + const metrics = join(dir, 'server_metrics_export.json'); + const profileRaw = Buffer.from( + JSON.stringify({ + metadata: { + conversation_id: 'conv-1', + turn_index: 0, + benchmark_phase: 'profiling', + credit_issued_ns: 1_000, + request_start_ns: 2_000, + request_end_ns: 5_000, + }, + metrics: { + input_sequence_length: { value: 128, unit: 'tokens' }, + output_sequence_length: { value: 64, unit: 'tokens' }, + }, + }), + ); + const csvRaw = Buffer.from('timestamp,value\n1,2\n'); + const metricsRaw = Buffer.from( + JSON.stringify({ + warmup_metrics: { 'vllm:prompt_tokens': metric(10) }, + metrics: { + 'vllm:prompt_tokens': metric(100), + 'vllm:generation_tokens': metric(50), + }, + }), + ); + await Promise.all([ + writeFile(profile, profileRaw), + writeFile(csv, csvRaw), + writeFile(metrics, metricsRaw), + ]); + return { profile, csv, metrics, profileRaw, csvRaw, metricsRaw }; +} + +async function waitUntil(predicate: () => boolean, timeoutMs = 5_000): Promise { + const deadline = Date.now() + timeoutMs; + while (!predicate()) { + if (Date.now() >= deadline) throw new Error('Timed out waiting for condition'); + await new Promise((resolve) => { + setTimeout(resolve, 10); + }); + } +} + +function sha256(buffer: Buffer | null): string | null { + return buffer ? createHash('sha256').update(buffer).digest('hex') : null; +} + +function fingerprint(prepared: PreparedTraceReplay) { + return { + profileGz: sha256(prepared.profileGz), + profileRaw: prepared.profileGz ? sha256(gunzipSync(prepared.profileGz)) : null, + profileSize: prepared.profileSize, + serverMetricsCsv: sha256(prepared.serverMetricsCsv), + serverMetricsCsvSize: prepared.serverMetricsCsvSize, + serverMetricsJsonGz: sha256(prepared.serverMetricsJsonGz), + serverMetricsJsonRaw: prepared.serverMetricsJsonGz + ? sha256(gunzipSync(prepared.serverMetricsJsonGz)) + : null, + serverMetricsJsonSize: prepared.serverMetricsJsonSize, + aggregateStatsJson: sha256(prepared.aggregateStatsJson), + chartSeriesJson: sha256(prepared.chartSeriesJson), + requestTimelineJson: sha256(prepared.requestTimelineJson), + chartWindows: prepared.chartWindows, + timelineRequests: prepared.timelineRequests, + cacheHitRates: prepared.cacheHitRates, + }; +} + +describe('resolveTraceReplayWorkerCount', () => { + it.each([ + [1, 1], + [4, 1], + [8, 2], + [16, 4], + [64, 4], + ])('maps %i vCPUs to %i bounded workers', (vcpus, expected) => { + expect(resolveTraceReplayWorkerCount(vcpus, undefined)).toBe(expected); + }); + + it('supports a bounded explicit override and ignores invalid values', () => { + expect(resolveTraceReplayWorkerCount(16, '6')).toBe(6); + expect(resolveTraceReplayWorkerCount(16, '99')).toBe(8); + expect(resolveTraceReplayWorkerCount(8, '0')).toBe(2); + expect(resolveTraceReplayWorkerCount(8, 'invalid')).toBe(2); + }); +}); + +describe('AsyncSemaphore', () => { + it('never exceeds its configured concurrency', async () => { + const semaphore = new AsyncSemaphore(2); + let active = 0; + let maximum = 0; + let release!: () => void; + const gate = new Promise((resolve) => { + release = resolve; + }); + const operations = Array.from({ length: 5 }, () => + semaphore.run(async () => { + active += 1; + maximum = Math.max(maximum, active); + await gate; + active -= 1; + }), + ); + + await waitUntil(() => active === 2); + expect(maximum).toBe(2); + release(); + await Promise.all(operations); + expect(maximum).toBe(2); + }); +}); + +describe('TraceReplayWorkerPool', () => { + it('transfers byte-identical payloads, bounds buffers, and recovers after a failed job', async () => { + const fixture = await traceFixture(); + const expected = await prepareTraceReplay(fixture.profile, fixture.csv, fixture.metrics); + const smokeScript = fileURLToPath(new URL('trace-replay-worker-smoke.ts', import.meta.url)); + const { stdout } = await execFileAsync( + 'bun', + [smokeScript, fixture.profile, fixture.csv, fixture.metrics], + { maxBuffer: 1024 * 1024 }, + ); + const result = JSON.parse(stdout) as { + consumersBeforeRelease: number; + totalConsumers: number; + direct: ReturnType; + results: ReturnType[]; + rejectedMissingInput: boolean; + recoveredTimelineRequests: number; + }; + + expect(result.consumersBeforeRelease).toBe(2); + expect(result.totalConsumers).toBe(3); + expect(result.results).toEqual([result.direct, result.direct, result.direct]); + const expectedFingerprint = fingerprint(expected); + expect(result.direct).toEqual({ + ...expectedFingerprint, + profileGz: result.direct.profileGz, + serverMetricsJsonGz: result.direct.serverMetricsJsonGz, + }); + expect(result.rejectedMissingInput).toBe(true); + expect(result.recoveredTimelineRequests).toBe(1); + }, 15_000); +}); 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..ee122411 --- /dev/null +++ b/packages/db/src/etl/trace-replay-worker-pool.ts @@ -0,0 +1,181 @@ +import { availableParallelism } from 'node:os'; +import { Worker } from 'node:worker_threads'; + +import type { PreparedTraceReplay } from './trace-replay-ingest'; +import type { + PreparedTraceReplayWire, + TraceReplayWorkerJob, + TraceReplayWorkerRequest, + TraceReplayWorkerResponse, +} from './trace-replay-worker-protocol'; + +const DEFAULT_MAX_WORKERS = 4; +const OVERRIDE_MAX_WORKERS = 8; + +/** + * Reserve roughly three quarters of the vCPUs for Bun, gzip, database IO, and + * runner overhead. Four GiB-scale parser workers are enough to saturate the + * current 16-vCPU ingest runner without multiplying memory use excessively. + */ +export function resolveTraceReplayWorkerCount( + vcpus = availableParallelism(), + override = process.env.INGEST_TRACE_WORKERS, +): number { + const normalizedVcpus = Math.max(1, Math.floor(vcpus)); + if (override !== undefined) { + const parsed = Number(override); + if (Number.isSafeInteger(parsed) && parsed > 0) { + return Math.min(parsed, normalizedVcpus, OVERRIDE_MAX_WORKERS); + } + } + return Math.max(1, Math.min(DEFAULT_MAX_WORKERS, Math.floor(normalizedVcpus / 4))); +} + +interface PendingTask { + id: number; + job: TraceReplayWorkerJob; + consume: (prepared: PreparedTraceReplay) => Promise; + resolve: (value: T) => void; + reject: (error: Error) => void; +} + +interface WorkerSlot { + worker: Worker; + task: PendingTask | null; +} + +const BUFFER_FIELDS = [ + 'profileGz', + 'serverMetricsCsv', + 'serverMetricsJsonGz', + 'aggregateStatsJson', + 'chartSeriesJson', + 'requestTimelineJson', +] as const; + +function fromWire(wire: PreparedTraceReplayWire): PreparedTraceReplay { + const prepared = { ...wire } as unknown as PreparedTraceReplay; + for (const field of BUFFER_FIELDS) { + const value = wire[field]; + prepared[field] = value + ? Buffer.from(value.buffer as ArrayBuffer, value.byteOffset, value.byteLength) + : null; + } + return prepared; +} + +/** + * A fixed worker pool whose slot remains reserved until `consume` finishes. + * This bounds prepared-buffer memory to the worker count even when Postgres is + * slower than computation, instead of allowing completed GiB-scale jobs to + * accumulate in an unbounded upload queue. + */ +export class TraceReplayWorkerPool { + private readonly queue: PendingTask[] = []; + private readonly slots: WorkerSlot[] = []; + private nextTaskId = 1; + private closing = false; + readonly size: number; + + constructor(size: number) { + if (!Number.isSafeInteger(size) || size < 1) { + throw new Error(`Worker count must be a positive integer, received ${size}`); + } + this.size = size; + } + + run( + job: TraceReplayWorkerJob, + consume: (prepared: PreparedTraceReplay) => Promise, + ): Promise { + if (this.closing) return Promise.reject(new Error('Trace replay worker pool is closing')); + this.ensureWorkers(); + return new Promise((resolve, reject) => { + this.queue.push({ + id: this.nextTaskId++, + job, + consume: consume as (prepared: PreparedTraceReplay) => Promise, + resolve: resolve as (value: unknown) => void, + reject, + }); + this.dispatch(); + }); + } + + async close(): Promise { + this.closing = true; + const closeError = new Error('Trace replay worker pool closed before queued work completed'); + for (const task of this.queue.splice(0)) task.reject(closeError); + await Promise.all(this.slots.map((slot) => slot.worker.terminate())); + this.slots.length = 0; + } + + private ensureWorkers(): void { + while (this.slots.length < this.size) { + const worker = new Worker(new URL('trace-replay-worker.ts', import.meta.url)); + const slot: WorkerSlot = { worker, task: null }; + worker.on('message', (response: TraceReplayWorkerResponse) => { + void this.handleResponse(slot, response); + }); + worker.on('error', (error) => + this.handleWorkerFailure(slot, error instanceof Error ? error : new Error(String(error))), + ); + worker.on('exit', (code) => { + if (!this.closing && code !== 0) { + this.handleWorkerFailure(slot, new Error(`Trace replay worker exited with code ${code}`)); + } + }); + this.slots.push(slot); + } + } + + private dispatch(): void { + for (const slot of this.slots) { + if (slot.task) continue; + const task = this.queue.shift(); + if (!task) return; + slot.task = task; + const request: TraceReplayWorkerRequest = { id: task.id, job: task.job }; + // oxlint-disable-next-line unicorn/require-post-message-target-origin -- node:worker_threads has no targetOrigin parameter + slot.worker.postMessage(request); + } + } + + private async handleResponse( + slot: WorkerSlot, + response: TraceReplayWorkerResponse, + ): Promise { + const task = slot.task; + if (!task || response.id !== task.id) { + this.handleWorkerFailure(slot, new Error(`Unexpected trace worker response ${response.id}`)); + return; + } + + try { + if (!response.ok) { + const error = new Error(response.error.message); + error.stack = response.error.stack ?? error.stack; + throw error; + } + task.resolve(await task.consume(fromWire(response.prepared))); + } catch (error) { + task.reject(error instanceof Error ? error : new Error(String(error))); + } finally { + slot.task = null; + this.dispatch(); + } + } + + private handleWorkerFailure(slot: WorkerSlot, error: Error): void { + if (!this.slots.includes(slot)) return; + slot.task?.reject(error); + slot.task = null; + const index = this.slots.indexOf(slot); + this.slots.splice(index, 1); + void slot.worker.terminate(); + if (!this.closing) { + this.ensureWorkers(); + this.dispatch(); + } + } +} diff --git a/packages/db/src/etl/trace-replay-worker-protocol.ts b/packages/db/src/etl/trace-replay-worker-protocol.ts new file mode 100644 index 00000000..d098134a --- /dev/null +++ b/packages/db/src/etl/trace-replay-worker-protocol.ts @@ -0,0 +1,29 @@ +import type { PreparedTraceReplay } from './trace-replay-ingest'; +import type { ServerMetricsContext } from './server-metrics-adapters'; + +export interface TraceReplayWorkerJob { + profileExportJsonl: string | null; + serverMetricsCsv: string | null; + serverMetricsJson: string | null; + metricsContext: ServerMetricsContext; +} + +export interface TraceReplayWorkerRequest { + id: number; + job: TraceReplayWorkerJob; +} + +type BufferField = + | 'profileGz' + | 'serverMetricsCsv' + | 'serverMetricsJsonGz' + | 'aggregateStatsJson' + | 'chartSeriesJson' + | 'requestTimelineJson'; + +export type PreparedTraceReplayWire = Omit & + Record; + +export type TraceReplayWorkerResponse = + | { id: number; ok: true; prepared: PreparedTraceReplayWire } + | { id: number; ok: false; error: { message: string; stack?: string } }; diff --git a/packages/db/src/etl/trace-replay-worker-smoke.ts b/packages/db/src/etl/trace-replay-worker-smoke.ts new file mode 100644 index 00000000..33d25efb --- /dev/null +++ b/packages/db/src/etl/trace-replay-worker-smoke.ts @@ -0,0 +1,108 @@ +import { createHash } from 'node:crypto'; +import { tmpdir } from 'node:os'; +import { join } from 'node:path'; +import { gunzipSync } from 'node:zlib'; + +import { prepareTraceReplay, type PreparedTraceReplay } from './trace-replay-ingest'; +import { TraceReplayWorkerPool } from './trace-replay-worker-pool'; + +const [profileExportJsonl, serverMetricsCsv, serverMetricsJson] = process.argv.slice(2); +if (!profileExportJsonl || !serverMetricsCsv || !serverMetricsJson) { + throw new Error('profile, csv, and metrics paths are required'); +} + +function sha256(buffer: Buffer | null): string | null { + return buffer ? createHash('sha256').update(buffer).digest('hex') : null; +} + +function fingerprint(prepared: PreparedTraceReplay) { + return { + profileGz: sha256(prepared.profileGz), + profileRaw: prepared.profileGz ? sha256(gunzipSync(prepared.profileGz)) : null, + profileSize: prepared.profileSize, + serverMetricsCsv: sha256(prepared.serverMetricsCsv), + serverMetricsCsvSize: prepared.serverMetricsCsvSize, + serverMetricsJsonGz: sha256(prepared.serverMetricsJsonGz), + serverMetricsJsonRaw: prepared.serverMetricsJsonGz + ? sha256(gunzipSync(prepared.serverMetricsJsonGz)) + : null, + serverMetricsJsonSize: prepared.serverMetricsJsonSize, + aggregateStatsJson: sha256(prepared.aggregateStatsJson), + chartSeriesJson: sha256(prepared.chartSeriesJson), + requestTimelineJson: sha256(prepared.requestTimelineJson), + chartWindows: prepared.chartWindows, + timelineRequests: prepared.timelineRequests, + cacheHitRates: prepared.cacheHitRates, + }; +} + +async function delay(milliseconds: number): Promise { + await new Promise((resolve) => { + setTimeout(resolve, milliseconds); + }); +} + +async function waitUntil(predicate: () => boolean): Promise { + const deadline = Date.now() + 10 * 60_000; + while (!predicate()) { + if (Date.now() >= deadline) throw new Error('Timed out waiting for worker consumers'); + await delay(10); + } +} + +const pool = new TraceReplayWorkerPool(2); +const job = { + profileExportJsonl, + serverMetricsCsv, + serverMetricsJson, + metricsContext: {}, +}; +let consumers = 0; +let release!: () => void; +const gate = new Promise((resolve) => { + release = resolve; +}); + +try { + const direct = fingerprint( + await prepareTraceReplay(profileExportJsonl, serverMetricsCsv, serverMetricsJson), + ); + const tasks = Array.from({ length: 3 }, () => + pool.run(job, async (prepared) => { + consumers += 1; + await gate; + return fingerprint(prepared); + }), + ); + await waitUntil(() => consumers === 2); + await delay(50); + const consumersBeforeRelease = consumers; + release(); + const results = await Promise.all(tasks); + + let rejectedMissingInput = false; + try { + await pool.run( + { ...job, profileExportJsonl: join(tmpdir(), 'missing-trace-worker-profile.jsonl') }, + () => Promise.resolve(undefined), + ); + } catch { + rejectedMissingInput = true; + } + const recoveredTimelineRequests = await pool.run(job, (prepared) => + Promise.resolve(prepared.timelineRequests), + ); + + console.log( + JSON.stringify({ + consumersBeforeRelease, + totalConsumers: consumers, + direct, + results, + rejectedMissingInput, + recoveredTimelineRequests, + }), + ); +} finally { + await pool.close(); +} 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..1835d0c8 --- /dev/null +++ b/packages/db/src/etl/trace-replay-worker.ts @@ -0,0 +1,63 @@ +import { parentPort } from 'node:worker_threads'; + +import { prepareTraceReplay, type PreparedTraceReplay } from './trace-replay-ingest'; +import type { + PreparedTraceReplayWire, + TraceReplayWorkerRequest, + TraceReplayWorkerResponse, +} from './trace-replay-worker-protocol'; + +const port = parentPort; +if (!port) throw new Error('trace-replay-worker must run in a worker thread'); + +const BUFFER_FIELDS = [ + 'profileGz', + 'serverMetricsCsv', + 'serverMetricsJsonGz', + 'aggregateStatsJson', + 'chartSeriesJson', + 'requestTimelineJson', +] as const; + +function toWire(prepared: PreparedTraceReplay): { + prepared: PreparedTraceReplayWire; + transfer: ArrayBuffer[]; +} { + const wire = { ...prepared } as unknown as PreparedTraceReplayWire; + const transfer: ArrayBuffer[] = []; + for (const field of BUFFER_FIELDS) { + const buffer = prepared[field]; + if (!buffer) continue; + const transferable = + buffer.buffer instanceof ArrayBuffer && + buffer.byteOffset === 0 && + buffer.byteLength === buffer.buffer.byteLength + ? new Uint8Array(buffer.buffer) + : Uint8Array.from(buffer); + wire[field] = transferable; + transfer.push(transferable.buffer); + } + return { prepared: wire, transfer }; +} + +port.on('message', async ({ id, job }: TraceReplayWorkerRequest) => { + try { + const result = await prepareTraceReplay( + job.profileExportJsonl, + job.serverMetricsCsv, + job.serverMetricsJson, + job.metricsContext, + ); + const { prepared, transfer } = toWire(result); + const response: TraceReplayWorkerResponse = { id, ok: true, prepared }; + port.postMessage(response, transfer); + } catch (error) { + const normalized = error instanceof Error ? error : new Error(String(error)); + const response: TraceReplayWorkerResponse = { + id, + ok: false, + error: { message: normalized.message, stack: normalized.stack }, + }; + port.postMessage(response); + } +}); diff --git a/packages/db/src/ingest-ci-run.ts b/packages/db/src/ingest-ci-run.ts index 85f03469..f5b17844 100644 --- a/packages/db/src/ingest-ci-run.ts +++ b/packages/db/src/ingest-ci-run.ts @@ -17,6 +17,7 @@ * INGEST_RUN_ID — (CI mode) Workflow run ID * INGEST_ARTIFACTS_PATH — (CI mode) Local path to pre-downloaded artifacts * INGEST_REPO — (CI mode) Source repo slug (owner/name) + * INGEST_TRACE_WORKERS — Optional trace worker override (default scales with vCPUs, max 4) * reused-ingest-metadata/reuse_source_run.json overrides reused rows to the * original source sweep run, so public links point at the real benchmark run. */ @@ -50,7 +51,12 @@ import { bulkUpsertAvailability, insertServerLog, } from './etl/benchmark-ingest'; -import { insertTraceReplay } from './etl/trace-replay-ingest'; +import { findUnlinkedTraceReplayIds, persistPreparedTraceReplay } from './etl/trace-replay-ingest'; +import { + resolveTraceReplayWorkerCount, + TraceReplayWorkerPool, +} from './etl/trace-replay-worker-pool'; +import { AsyncSemaphore } from './etl/async-semaphore'; import { discoverTraceReplayArtifacts } from './etl/trace-artifact-discovery'; import { datasetSlugFromBenchmarkRow } from './etl/dataset-provenance'; import { mapAggEvalRow, mapEvalRow } from './etl/eval-mapper'; @@ -424,6 +430,16 @@ async function main(): Promise { if (traceReplayPaths.size > 0) { console.log(` Found ${traceReplayPaths.size} trace_replay sibling artifact(s)`); } + const traceWorkerCount = resolveTraceReplayWorkerCount(); + const traceWorkerPool = new TraceReplayWorkerPool(traceWorkerCount); + const traceUploadLimiter = new AsyncSemaphore(Math.min(2, traceWorkerCount)); + const traceTasks: Promise[] = []; + if (traceReplayPaths.size > 0) { + console.log( + ` Trace preparation: ${traceWorkerCount} worker(s) across ${os.availableParallelism()} vCPU(s), ` + + `${Math.min(2, traceWorkerCount)} concurrent upload(s)`, + ); + } const allBmkFiles = [...bmkFiles, ...allBmkDirs.flatMap((d) => findJsonFiles(d))]; console.log(` Found ${allBmkFiles.length} benchmark JSON file(s)`); @@ -565,22 +581,62 @@ async function main(): Promise { `server_csv=${formatBytes(fileSize(trace.serverMetricsCsv))}, ` + `server_json=${formatBytes(fileSize(trace.serverMetricsJson))}`, ); - await insertTraceReplay( - sql, - insertedIds, - trace.profileJsonl, - trace.serverMetricsCsv, - trace.serverMetricsJson, - { - metricsContext: { - framework: toInsert[0]?.config.framework, - disagg: toInsert[0]?.config.disagg, - }, - progressLabel: suffix, - }, + const linkStart = Date.now(); + console.log( + ` trace_replay ${suffix}: checking ${insertedIds.length} benchmark row(s) for existing links`, + ); + const unlinkedIds = await findUnlinkedTraceReplayIds(sql, insertedIds); + console.log( + ` trace_replay ${suffix}: found ${unlinkedIds.length} unlinked row(s) (${elapsed(linkStart)})`, ); - totalTraceReplayLinked += insertedIds.length; - console.log(` trace_replay ${suffix}: done (${elapsed(traceStart)})`); + if (unlinkedIds.length === 0) { + console.log( + ` trace_replay ${suffix}: skipping blob insert; all benchmark rows already linked`, + ); + console.log(` trace_replay ${suffix}: done (${elapsed(traceStart)})`); + } else { + console.log(` trace_replay ${suffix}: queued for worker preparation`); + const task = traceWorkerPool + .run( + { + profileExportJsonl: trace.profileJsonl, + serverMetricsCsv: trace.serverMetricsCsv, + serverMetricsJson: trace.serverMetricsJson, + metricsContext: { + framework: toInsert[0]?.config.framework, + disagg: toInsert[0]?.config.disagg, + }, + }, + // oxlint-disable-next-line no-loop-func -- each callback closes over this iteration's block-scoped trace metadata + async (prepared) => { + console.log( + ` trace_replay ${suffix}: compressed ` + + `profile=${formatBytes(prepared.profileSize)} -> ${formatBytes(prepared.profileGz?.length)}, ` + + `server_csv=${formatBytes(prepared.serverMetricsCsvSize)}, ` + + `server_json=${formatBytes(prepared.serverMetricsJsonSize)} -> ` + + `${formatBytes(prepared.serverMetricsJsonGz?.length)} ` + + `(${(prepared.compressionMs / 1000).toFixed(1)}s)`, + ); + console.log( + ` trace_replay ${suffix}: computed derived JSON: ` + + `chart_windows=${prepared.chartWindows}, ` + + `timeline_requests=${prepared.timelineRequests} ` + + `(${(prepared.computeMs / 1000).toFixed(1)}s)`, + ); + const linked = await traceUploadLimiter.run(() => + persistPreparedTraceReplay(sql, unlinkedIds, prepared, { + progressLabel: suffix, + }), + ); + totalTraceReplayLinked += linked; + console.log(` trace_replay ${suffix}: done (${elapsed(traceStart)})`); + }, + ) + .catch((error: any) => { + tracker.recordDbError(`trace_replay for ${suffix}`, error); + }); + traceTasks.push(task); + } } catch (error: any) { tracker.recordDbError(`trace_replay for ${suffix}`, error); } @@ -595,6 +651,11 @@ async function main(): Promise { } console.log(` finished ${relativeFile} (${elapsed(fileStart)})`); } + if (traceTasks.length > 0) { + console.log(` Waiting for ${traceTasks.length} queued trace replay task(s)`); + await Promise.all(traceTasks); + } + await traceWorkerPool.close(); console.log(` Benchmarks: +${totalNewBmk} new, ${totalDupBmk} dup`); if (totalTraceReplayLinked > 0 || tracker.skips.traceReplayMissing > 0) { console.log(