Optimize agentic trace ingestion without changing payloads

packages/db/src/etl/compute-aggregate-stats.ts

@@ -61,7 +61,7 @@ export function mergeProfileStatsUpgrade(
 }
 
 /** Metric subtrees we extract via stream-parse on oversized server blobs. */
-const TARGET_METRIC_KEYS = new Set([
+export const AGGREGATE_SERVER_METRIC_KEYS = new Set([
   'vllm:kv_cache_usage_perc',
   'vllm:gpu_cache_usage_perc',
   'vllm:prefix_cache_hits',
@@ -78,10 +78,36 @@ const TARGET_METRIC_KEYS = new Set([
 async function streamExtractServer(
   buffer: Buffer,
 ): Promise<{ kvCacheUtil: number[]; prefixCacheHitRate: number[] }> {
-  const collected = await streamCollectKeys<unknown>(buffer, 'metrics', TARGET_METRIC_KEYS);
+  const collected = await streamCollectKeys<unknown>(
+    buffer,
+    'metrics',
+    AGGREGATE_SERVER_METRIC_KEYS,
+  );
   return extractServerMetricSamples(JSON.stringify({ metrics: collected }));
 }
 
+/**
+ * Add server-derived distributions to profile stats using an already parsed
+ * profiling metric map. Ingest uses this to share one server JSON parse with
+ * chart-series generation; the output shape and ordering match
+ * `computeAggregateStats()` exactly.
+ */
+export function withServerMetricAggregateStats(
+  profileStats: AggregateStats,
+  metrics: Record<string, unknown>,
+): AggregateStats {
+  try {
+    const server = extractServerMetricSamples(JSON.stringify({ metrics }));
+    return {
+      ...profileStats,
+      kvCacheUtil: percentilesOf(server.kvCacheUtil),
+      prefixCacheHitRate: percentilesOf(server.prefixCacheHitRate),
+    };
+  } catch {
+    return profileStats;
+  }
+}
+
 /**
  * Compute the full versioned stats bundle from a (profile, server-metrics)
  * blob pair. Either blob may be null (e.g. only the server file existed) —

packages/db/src/etl/compute-chart-series.ts

@@ -156,18 +156,18 @@ interface RawSeries {
   timeslices?: RawSlice[];
 }
 
-interface RawMetric {
+export interface RawMetric {
   series?: RawSeries[];
 }
 
-type MetricsMap = Record<string, RawMetric>;
+export type MetricsMap = Record<string, RawMetric>;
 
 /**
  * The set of metric subtrees the chart consumes. Includes both vllm:* and
  * sglang:* names so the stream-parse fallback collects whichever framework
  * the blob was emitted by — `buildSeriesFromMetrics` then picks per metric.
  */
-const CHART_METRIC_KEYS = new Set([
+export const CHART_METRIC_KEYS = new Set([
   // vLLM
   'vllm:kv_cache_usage_perc',
   'vllm:gpu_cache_usage_perc',
@@ -257,6 +257,19 @@ export async function computeChartSeries(
   return buildSeriesFromMetrics(metrics, context);
 }
 
+/**
+ * Build the chart payload from already parsed phase maps. This is the same
+ * merge + projection used by `computeChartSeries()`, exposed so ingest can
+ * share one server JSON parse with aggregate-stat computation.
+ */
+export function computeChartSeriesFromMetricPhases(
+  profiling: MetricsMap,
+  warmup: MetricsMap,
+  context: ServerMetricsContext = {},
+): ChartSeries {
+  return buildSeriesFromMetrics(mergePhaseMetrics(profiling, warmup), context);
+}
+
 /**
  * Aggregate one timeslice field across all series of a metric, indexed by
  * `start_ns`. Multi-engine vllm deployments report one series per engine —

packages/db/src/etl/compute-trace-derived.test.ts

@@ -0,0 +1,142 @@
+import { gzipSync } from 'node:zlib';
+
+import { describe, expect, it } from 'vitest';
+
+import { computeAggregateStats } from './compute-aggregate-stats.js';
+import { computeChartSeries } from './compute-chart-series.js';
+import { computeRequestTimeline } from './compute-request-timeline.js';
+import { computeTraceDerivedPayloads } from './compute-trace-derived.js';
+
+function makeProfileBlob(): Buffer {
+  return gzipSync(
+    Buffer.from(
+      [
+        {
+          metadata: {
+            conversation_id: 'conv-1',
+            turn_index: 0,
+            worker_id: 'worker-1',
+            agent_depth: 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' },
+            time_to_first_token: { value: 20, unit: 'ms' },
+            inter_token_latency: { value: 5, unit: 'ms' },
+          },
+        },
+        {
+          metadata: {
+            conversation_id: 'conv-1',
+            turn_index: 1,
+            worker_id: 'worker-1',
+            agent_depth: 0,
+            benchmark_phase: 'profiling',
+            credit_issued_ns: 6_000,
+            request_start_ns: 7_000,
+            request_end_ns: 10_000,
+          },
+          metrics: {
+            input_sequence_length: { value: 256, unit: 'tokens' },
+            output_sequence_length: { value: 32, unit: 'tokens' },
+            time_to_first_token: { value: 30, unit: 'ms' },
+            inter_token_latency: { value: 6, unit: 'ms' },
+          },
+        },
+      ]
+        .map((record) => JSON.stringify(record))
+        .join('\n'),
+    ),
+  );
+}
+
+function metric(
+  values: { start_ns: number; end_ns: number; avg?: number; rate?: number }[],
+  labels: Record<string, string> = {},
+) {
+  return { series: [{ endpoint_url: 'worker.test:8000', labels, timeslices: values }] };
+}
+
+function makeServerBlob(): Buffer {
+  return gzipSync(
+    Buffer.from(
+      JSON.stringify({
+        warmup_metrics: {
+          'vllm:kv_cache_usage_perc': metric([{ start_ns: 0, end_ns: 1e9, avg: 0.1 }]),
+          'vllm:prompt_tokens': metric([{ start_ns: 0, end_ns: 1e9, rate: 100 }]),
+        },
+        metrics: {
+          'vllm:kv_cache_usage_perc': metric([
+            { start_ns: 10e9, end_ns: 11e9, avg: 0.4 },
+            { start_ns: 11e9, end_ns: 12e9, avg: 0.6 },
+          ]),
+          'vllm:prefix_cache_hits': metric([{ start_ns: 10e9, end_ns: 11e9, rate: 80 }]),
+          'vllm:prefix_cache_queries': metric([{ start_ns: 10e9, end_ns: 11e9, rate: 100 }]),
+          'vllm:gpu_prefix_cache_hits': metric([{ start_ns: 10e9, end_ns: 11e9, rate: 80 }]),
+          'vllm:gpu_prefix_cache_queries': metric([{ start_ns: 10e9, end_ns: 11e9, rate: 100 }]),
+          'vllm:num_requests_running': metric([{ start_ns: 10e9, end_ns: 11e9, avg: 3 }]),
+          'vllm:num_requests_waiting': metric([{ start_ns: 10e9, end_ns: 11e9, avg: 2 }]),
+          'vllm:prompt_tokens': metric([{ start_ns: 10e9, end_ns: 11e9, rate: 900 }]),
+          'vllm:generation_tokens': metric([{ start_ns: 10e9, end_ns: 11e9, rate: 450 }]),
+        },
+      }),
+    ),
+  );
+}
+
+describe('computeTraceDerivedPayloads', () => {
+  it('is byte-for-byte JSON equivalent to the independent upload computations', async () => {
+    const profileBlob = makeProfileBlob();
+    const serverBlob = makeServerBlob();
+    const context = { framework: 'dynamo-vllm', disagg: true } as const;
+
+    const [aggregateStats, chartSeries, requestTimeline] = await Promise.all([
+      computeAggregateStats({ profileBlob, serverBlob }),
+      computeChartSeries(serverBlob, context),
+      Promise.resolve(computeRequestTimeline(profileBlob)),
+    ]);
+    const optimized = await computeTraceDerivedPayloads(profileBlob, serverBlob, context);
+
+    expect(Buffer.from(JSON.stringify(optimized.aggregateStats))).toEqual(
+      Buffer.from(JSON.stringify(aggregateStats)),
+    );
+    expect(Buffer.from(JSON.stringify(optimized.chartSeries))).toEqual(
+      Buffer.from(JSON.stringify(chartSeries)),
+    );
+    expect(Buffer.from(JSON.stringify(optimized.requestTimeline))).toEqual(
+      Buffer.from(JSON.stringify(requestTimeline)),
+    );
+  });
+
+  it('produces the same payloads through the oversized streaming path', async () => {
+    const profileBlob = makeProfileBlob();
+    const serverBlob = makeServerBlob();
+    const bounded = await computeTraceDerivedPayloads(profileBlob, serverBlob);
+    const streamed = await computeTraceDerivedPayloads(
+      profileBlob,
+      serverBlob,
+      {},
+      {
+        maxInMemoryBytes: 1,
+      },
+    );
+
+    expect(streamed).toEqual(bounded);
+  });
+
+  it('preserves independent malformed-input fallbacks', async () => {
+    const profileBlob = makeProfileBlob();
+    const malformedServer = Buffer.from('not-gzip');
+    const optimized = await computeTraceDerivedPayloads(profileBlob, malformedServer);
+
+    expect(optimized.aggregateStats).toEqual(
+      await computeAggregateStats({ profileBlob, serverBlob: malformedServer }),
+    );
+    expect(optimized.chartSeries).toBeNull();
+    expect(optimized.requestTimeline).toEqual(computeRequestTimeline(profileBlob));
+  });
+});

packages/db/src/etl/compute-trace-derived.ts

@@ -0,0 +1,95 @@
+/**
+ * Compute every derived trace-replay payload while parsing server metrics only
+ * once. The standalone aggregate/chart helpers remain the compatibility path
+ * for backfills; ingest uses this coordinator to avoid repeated GiB-scale
+ * decompression and JSON tokenization.
+ */
+
+import {
+  AGGREGATE_SERVER_METRIC_KEYS,
+  computeAggregateStats,
+  withServerMetricAggregateStats,
+  type AggregateStats,
+} from './compute-aggregate-stats.js';
+import {
+  CHART_METRIC_KEYS,
+  computeChartSeriesFromMetricPhases,
+  type ChartSeries,
+  type MetricsMap,
+  type RawMetric,
+} from './compute-chart-series.js';
+import { computeRequestTimeline, type RequestTimeline } from './compute-request-timeline.js';
+import { collectMetricPhases } from './gzip-json-stream.js';
+import type { ServerMetricsContext } from './server-metrics-adapters.js';
+
+export interface TraceDerivedPayloads {
+  aggregateStats: AggregateStats;
+  chartSeries: ChartSeries | null;
+  requestTimeline: RequestTimeline | null;
+}
+
+export interface TraceDerivedComputeOptions {
+  /** Override the bounded fast-path threshold, primarily for streaming tests. */
+  maxInMemoryBytes?: number;
+}
+
+const DERIVED_SERVER_METRIC_KEYS = new Set([...CHART_METRIC_KEYS, ...AGGREGATE_SERVER_METRIC_KEYS]);
+
+function selectMetrics(metrics: MetricsMap, wanted: ReadonlySet<string>): MetricsMap {
+  const selected: MetricsMap = {};
+  for (const [name, metric] of Object.entries(metrics)) {
+    if (wanted.has(name)) selected[name] = metric;
+  }
+  return selected;
+}
+
+/**
+ * Produce the same three JSON values previously computed independently in
+ * `insertTraceReplay()`. Malformed inputs retain the old failure isolation:
+ * profile stats/timeline can succeed without server metrics, and a chart
+ * projection failure does not discard aggregate stats.
+ */
+export async function computeTraceDerivedPayloads(
+  profileBlob: Buffer | null,
+  serverBlob: Buffer | null,
+  metricsContext: ServerMetricsContext = {},
+  options: TraceDerivedComputeOptions = {},
+): Promise<TraceDerivedPayloads> {
+  const profileStatsPromise = computeAggregateStats({ profileBlob, serverBlob: null });
+  const requestTimeline = computeRequestTimeline(profileBlob);
+
+  const phases = serverBlob
+    ? await collectMetricPhases<RawMetric>(
+        serverBlob,
+        DERIVED_SERVER_METRIC_KEYS,
+        options.maxInMemoryBytes,
+      ).catch(() => null)
+    : null;
+  let aggregateStats = await profileStatsPromise;
+  let chartSeries: ChartSeries | null = null;
+
+  if (phases) {
+    const aggregateMetrics = phases.complete
+      ? phases.metrics
+      : selectMetrics(phases.metrics, AGGREGATE_SERVER_METRIC_KEYS);
+    aggregateStats = withServerMetricAggregateStats(aggregateStats, aggregateMetrics);
+
+    // The historical in-memory path builds chart timing metadata from every
+    // metric, while its oversized streaming fallback retains only chart keys.
+    // Preserve that distinction exactly even though the shared streaming pass
+    // also collects the aggregate-only GPU prefix-cache aliases.
+    const profiling = phases.complete
+      ? phases.metrics
+      : selectMetrics(phases.metrics, CHART_METRIC_KEYS);
+    const warmup = phases.complete
+      ? phases.warmupMetrics
+      : selectMetrics(phases.warmupMetrics, CHART_METRIC_KEYS);
+    try {
+      chartSeries = computeChartSeriesFromMetricPhases(profiling, warmup, metricsContext);
+    } catch {
+      chartSeries = null;
+    }
+  }
+
+  return { aggregateStats, chartSeries, requestTimeline };
+}

packages/db/src/etl/gzip-json-stream.test.ts

@@ -2,7 +2,11 @@ import { gzipSync } from 'node:zlib';
 
 import { describe, expect, it } from 'vitest';
 
-import { gunzipJsonWithinLimit, streamCollectKeys } from './gzip-json-stream.js';
+import {
+  collectMetricPhases,
+  gunzipJsonWithinLimit,
+  streamCollectKeys,
+} from './gzip-json-stream.js';
 
 describe('gunzipJsonWithinLimit', () => {
   const json = JSON.stringify({ metrics: { value: 1 } });
@@ -61,3 +65,47 @@ describe('streamCollectKeys', () => {
     ).rejects.toThrow();
   });
 });
+
+describe('collectMetricPhases', () => {
+  const blob = gzipSync(
+    JSON.stringify({
+      metadata: { ignored: true },
+      metrics: {
+        wanted: { series: [{ timeslices: [{ start_ns: 1, rate: 2 }] }] },
+        ignored: { series: [{ timeslices: [{ start_ns: 3, rate: 4 }] }] },
+      },
+      warmup_metrics: {
+        wanted: { series: [{ timeslices: [{ start_ns: 0, rate: 1 }] }] },
+        ignored: { series: [] },
+      },
+    }),
+  );
+
+  it('retains the complete phase maps on the bounded fast path', async () => {
+    const phases = await collectMetricPhases(blob, new Set(['wanted']));
+
+    expect(phases.complete).toBe(true);
+    expect(Object.keys(phases.metrics)).toEqual(['wanted', 'ignored']);
+    expect(Object.keys(phases.warmupMetrics)).toEqual(['wanted', 'ignored']);
+  });
+
+  it('collects both filtered phase maps from one streaming parse', async () => {
+    const phases = await collectMetricPhases(blob, new Set(['wanted']), 1);
+
+    expect(phases).toEqual({
+      metrics: {
+        wanted: { series: [{ timeslices: [{ start_ns: 1, rate: 2 }] }] },
+      },
+      warmupMetrics: {
+        wanted: { series: [{ timeslices: [{ start_ns: 0, rate: 1 }] }] },
+      },
+      complete: false,
+    });
+  });
+
+  it('rejects malformed gzip input on both paths', async () => {
+    await expect(
+      collectMetricPhases(Buffer.from('not gzip'), new Set(['wanted']), 1),
+    ).rejects.toThrow();
+  });
+});