fix(db): chunk trace-replay uploads (#570)

cquil11 · web-flow · commit f4346d6976e3 · 2026-07-13T15:15:10.000-05:00
Upload large agentic trace values in 4 MiB pieces and assemble them atomically in PostgreSQL. Extend the workflow timeout so recovered sweeps retain enough headroom.

中文：将大型 agentic trace 数据拆分为 4 MiB 分块上传，并在 PostgreSQL 中以事务方式原子组装。同时延长工作流超时时间，为恢复完整 sweep 预留足够时间。
diff --git a/.github/workflows/ingest-agentic-results.yml b/.github/workflows/ingest-agentic-results.yml
@@ -94,7 +94,9 @@ jobs:
     name: Ingest agentic benchmark results
     # Blob-heavy: uploading trace-replay sidecars for a ~20-point sweep takes
     # far longer than a fixed-seq-len ingest.
-    timeout-minutes: 60
+    # High-concurrency rows also precompute request timelines from GiB-scale
+    # artifacts, so retain enough headroom for a complete recovered sweep.
+    timeout-minutes: 180
     runs-on: blacksmith-4vcpu-ubuntu-2404
     permissions:
       contents: read
diff --git a/packages/db/src/etl/trace-replay-ingest.test.ts b/packages/db/src/etl/trace-replay-ingest.test.ts
@@ -0,0 +1,57 @@
+import type postgres from 'postgres';
+import { describe, expect, it } from 'vitest';
+
+import {
+  TRACE_REPLAY_UPLOAD_CHUNK_BYTES,
+  uploadTraceReplayPayloadChunks,
+} from './trace-replay-ingest';
+
+interface SqlCall {
+  text: string;
+  values: unknown[];
+}
+
+function mockTransactionSql(): { sql: postgres.TransactionSql; calls: SqlCall[] } {
+  const calls: SqlCall[] = [];
+  const sql = ((strings: TemplateStringsArray, ...values: unknown[]) => {
+    calls.push({ text: strings.join('?'), values });
+    return Promise.resolve([]);
+  }) as unknown as postgres.TransactionSql;
+  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.
+    const measuredPayloads = [
+      ['profile_export_jsonl_gz', Buffer.alloc(22_992_290)],
+      ['server_metrics_json_gz', Buffer.alloc(49_891_135)],
+      ['request_timeline', Buffer.alloc(22_146_655)],
+    ] as const;
+    const { sql, calls } = mockTransactionSql();
+
+    let expectedParts = 0;
+    for (const [field, payload] of measuredPayloads) {
+      const parts = await uploadTraceReplayPayloadChunks(sql, field, payload);
+      const expected = Math.ceil(payload.length / TRACE_REPLAY_UPLOAD_CHUNK_BYTES);
+      expect(parts).toBe(expected);
+      expectedParts += expected;
+    }
+
+    expect(calls).toHaveLength(expectedParts);
+    expect(calls.every((call) => call.text.includes('trace_replay_upload_parts'))).toBe(true);
+    const chunks = calls.flatMap((call) => call.values.filter(Buffer.isBuffer));
+    expect(chunks).toHaveLength(expectedParts);
+    expect(Math.max(...chunks.map((chunk) => chunk.length))).toBe(TRACE_REPLAY_UPLOAD_CHUNK_BYTES);
+    expect(chunks.reduce((total, chunk) => total + chunk.length, 0)).toBe(
+      measuredPayloads.reduce((total, [, payload]) => total + payload.length, 0),
+    );
+  });
+
+  it('does not issue a query for a missing payload', async () => {
+    const { sql, calls } = mockTransactionSql();
+
+    await expect(uploadTraceReplayPayloadChunks(sql, 'chart_series', null)).resolves.toBe(0);
+    expect(calls).toHaveLength(0);
+  });
+});
diff --git a/packages/db/src/etl/trace-replay-ingest.ts b/packages/db/src/etl/trace-replay-ingest.ts
@@ -19,6 +19,21 @@ import type { ServerMetricsContext } from './server-metrics-adapters';
 
 type Sql = ReturnType<typeof postgres>;
 
+/**
+ * Keep each postgres.js Bind message comfortably below the large payload that
+ * can stall through Neon's proxy. The final row is assembled server-side in a
+ * transaction, so callers still get all-or-nothing trace persistence.
+ */
+export const TRACE_REPLAY_UPLOAD_CHUNK_BYTES = 4 * 1024 * 1024;
+
+type TraceReplayUploadField =
+  | 'profile_export_jsonl_gz'
+  | 'server_metrics_csv'
+  | 'server_metrics_json_gz'
+  | 'aggregate_stats'
+  | 'chart_series'
+  | 'request_timeline';
+
 export interface TraceReplayIngestOptions {
   metricsContext?: ServerMetricsContext;
   progressLabel?: string;
@@ -38,6 +53,34 @@ function elapsed(startMs: number): string {
   return `${((Date.now() - startMs) / 1000).toFixed(1)}s`;
 }
 
+function jsonBuffer(value: unknown | null): Buffer | null {
+  if (value === null) return null;
+  return Buffer.from(JSON.stringify(structuredClone(value)), 'utf8');
+}
+
+/**
+ * Upload one trace-replay value as bounded binary parameters. The caller must
+ * create `pg_temp.trace_replay_upload_parts` in the same transaction first.
+ */
+export async function uploadTraceReplayPayloadChunks(
+  sql: postgres.TransactionSql,
+  field: TraceReplayUploadField,
+  payload: Buffer | null,
+): Promise<number> {
+  if (!payload) return 0;
+
+  let part = 0;
+  for (let offset = 0; offset < payload.length; offset += TRACE_REPLAY_UPLOAD_CHUNK_BYTES) {
+    const chunk = payload.subarray(offset, offset + TRACE_REPLAY_UPLOAD_CHUNK_BYTES);
+    await sql`
+      insert into pg_temp.trace_replay_upload_parts (field, part, data)
+      values (${field}, ${part}, ${chunk})
+    `;
+    part += 1;
+  }
+  return part;
+}
+
 /**
  * Persist the per-point trace files and link them to `benchmarkResultIds`.
  *
@@ -119,86 +162,144 @@ export async function insertTraceReplay(
       `timeline_requests=${requestTimeline?.requests.length ?? 0} (${elapsed(computeStart)})`,
   );
 
+  const aggregateStatsJson = jsonBuffer(aggregateStats);
+  const chartSeriesJson = jsonBuffer(chartSeries);
+  const requestTimelineJson = jsonBuffer(requestTimeline);
+
   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<typeof sql.json>[0])},
-      ${chartSeries === null ? null : sql.json(structuredClone(chartSeries) as unknown as Parameters<typeof sql.json>[0])},
-      ${requestTimeline === null ? null : sql.json(structuredClone(requestTimeline) as unknown as Parameters<typeof sql.json>[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)
+  log(`uploading trace_replay payloads in ${formatBytes(TRACE_REPLAY_UPLOAD_CHUNK_BYTES)} chunks`);
+  await sql.begin(async (tx) => {
+    await tx`
+      create temporary table trace_replay_upload_parts (
+        field text not null,
+        part integer not null,
+        data bytea not null,
+        primary key (field, part)
+      ) on commit drop
+    `;
+
+    const payloads: [TraceReplayUploadField, Buffer | null][] = [
+      ['profile_export_jsonl_gz', profileGz],
+      ['server_metrics_csv', serverMetricsCsv],
+      ['server_metrics_json_gz', metricsJsonGz],
+      ['aggregate_stats', aggregateStatsJson],
+      ['chart_series', chartSeriesJson],
+      ['request_timeline', requestTimelineJson],
+    ];
+    for (const [field, payload] of payloads) {
+      const uploadStart = Date.now();
+      const parts = await uploadTraceReplayPayloadChunks(tx, field, payload);
+      log(
+        `uploaded ${field}=${formatBytes(payload?.length)} in ${parts} part(s) ` +
+          `(${elapsed(uploadStart)})`,
+      );
+    }
+
+    log('assembling 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 (
+        (
+          select string_agg(data, ''::bytea order by part)
+          from pg_temp.trace_replay_upload_parts
+          where field = 'profile_export_jsonl_gz'
+        ),
+        ${profileSize},
+        (
+          select string_agg(data, ''::bytea order by part)
+          from pg_temp.trace_replay_upload_parts
+          where field = 'server_metrics_csv'
+        ),
+        ${csvSize},
+        (
+          select string_agg(data, ''::bytea order by part)
+          from pg_temp.trace_replay_upload_parts
+          where field = 'server_metrics_json_gz'
+        ),
+        ${metricsJsonSize},
+        (
+          select convert_from(string_agg(data, ''::bytea order by part), 'UTF8')::jsonb
+          from pg_temp.trace_replay_upload_parts
+          where field = 'aggregate_stats'
+        ),
+        (
+          select convert_from(string_agg(data, ''::bytea order by part), 'UTF8')::jsonb
+          from pg_temp.trace_replay_upload_parts
+          where field = 'chart_series'
+        ),
+        (
+          select convert_from(string_agg(data, ''::bytea order by part), 'UTF8')::jsonb
+          from pg_temp.trace_replay_upload_parts
+          where field = 'request_timeline'
         )
-        where id = any(${sql.array(unlinked.map((r) => r.id))}::bigint[])
-      `;
-      log('updated cache-hit metrics from chart series');
+      )
+      returning id
+    `;
+    log(`assembled trace_replay_id=${traceReplayId}`);
+
+    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 (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');
+      }
     }
-  }
+  });
+  log(`inserted trace_replay payload (${elapsed(insertStart)})`);
 }
