fix(data): cut athena stats sync cost

Compute all stat dimensions and both grains in one GROUPING SETS query so each sync pass scans the source table once instead of eight times. Hourly passes now only recompute the current ISO week; a daily full pass refreshes the whole display window. Restarted daemons resume the hourly cadence from the last completed sync instead of immediately re-running a pass, and the workgroup kills any query scanning more than 2 TB.
2026-07-21 02:06:00 +00:00 · 2026-07-10 06:36:26 -05:00
parent 60c88336b9
commit 08a2a623dc
5 changed files with 122 additions and 77 deletions
@@ -67,6 +67,10 @@ const athenaWorkgroup = new aws.athena.Workgroup("LakeAthenaWorkgroup", {
  configuration: {
    enforceWorkgroupConfiguration: true,
    publishCloudwatchMetricsEnabled: true,
+    // Athena bills $5/TB scanned; kill any query that would scan more than 2 TB
+    // so a regression cannot silently burn money. Stats sync full passes scan
+    // ~250 GB as of 2026-07.
+    bytesScannedCutoffPerQuery: 2 * 1024 ** 4,
    resultConfiguration: {
      outputLocation: $interpolate`s3://${athenaResultsBucket.bucket}/`,
    },
@@ -14,7 +14,10 @@ import { normalizeCountry, normalizeTier, type StatBaseAggregate } from "./stat"

 export type StatDimension = "model" | "provider" | "geo" | "geo_model"

-export function buildStatsQuery(periodStart: Date, periodEnd: Date, dimension: StatDimension) {
+// All stat dimensions and both grains are computed in one query via GROUPING SETS so
+// the source table is scanned once per sync pass; separate queries per dimension (and
+// the previous weekly/daily UNION ALL) each re-scanned the same events.
+export function buildStatsQuery(periodStart: Date, periodEnd: Date) {
  const periodStartValue = sqlString(periodStart.toISOString())
  const periodEndValue = sqlString(periodEnd.toISOString())
  const periodStartDateValue = sqlString(periodStart.toISOString().slice(0, 10))
@@ -22,23 +25,6 @@ export function buildStatsQuery(periodStart: Date, periodEnd: Date, dimension: S
  const sourceTable = [Resource.InferenceEvent.catalog, Resource.InferenceEvent.database, Resource.InferenceEvent.table]
    .map(sqlIdentifier)
    .join(".")
-  const dimensionSql = (() => {
-    if (dimension === "model")
-      return {
-        select: "provider, model, COALESCE(MAX(NULLIF(provider_model, '')), '') AS provider_model",
-        groupBy: "provider, model",
-      }
-    if (dimension === "provider") return { select: "provider", groupBy: "provider" }
-    if (dimension === "geo_model")
-      return {
-        select: "provider, model, country, COALESCE(MAX(NULLIF(continent, '')), '') AS continent",
-        groupBy: "provider, model, country",
-      }
-    return {
-      select: "'all' AS provider, 'all' AS model, country, COALESCE(MAX(NULLIF(continent, '')), '') AS continent",
-      groupBy: "country",
-    }
-  })()
  const aggregateColumns = `
    COUNT(DISTINCT session) AS sessions,
    COUNT(*) AS requests,
@@ -135,34 +121,41 @@ WITH normalized AS (
    COALESCE(cost_total_microcents, cost_total * 1000000) AS cost_total_microcents
  FROM normalized
  WHERE lower(model) NOT IN (${[...EXCLUDED_MODELS].map(sqlString).join(", ")})
-), weekly AS (
+), periods AS (
  SELECT
    concat(CAST(year_of_week(event_time) AS varchar), '-W', lpad(CAST(week(event_time) AS varchar), 2, '0')) AS week_key,
+    substr(to_iso8601(date_trunc('day', event_time)), 1, 10) AS day_key,
    *
  FROM filtered
-), daily AS (
-  SELECT substr(to_iso8601(date_trunc('day', event_time)), 1, 10) AS day_key, *
-  FROM filtered
 )
 SELECT
-  'week' AS grain,
-  week_key AS period_key,
+  CASE WHEN grouping(week_key) = 0 THEN 'week' ELSE 'day' END AS grain,
+  COALESCE(week_key, day_key) AS period_key,
  ${sqlString(Resource.StatsSyncConfig.dataset)} AS dataset,
+  CASE
+    WHEN grouping(country) = 0 AND grouping(model) = 0 THEN 'geo_model'
+    WHEN grouping(country) = 0 THEN 'geo'
+    WHEN grouping(model) = 0 THEN 'model'
+    ELSE 'provider'
+  END AS dimension,
  tier,
-  ${dimensionSql.select},
+  CASE WHEN grouping(provider) = 0 THEN provider ELSE 'all' END AS provider,
+  CASE WHEN grouping(model) = 0 THEN model WHEN grouping(country) = 0 THEN 'all' END AS model,
+  CASE WHEN grouping(model) = 0 AND grouping(country) = 1 THEN COALESCE(MAX(NULLIF(provider_model, '')), '') END AS provider_model,
+  CASE WHEN grouping(country) = 0 THEN country END AS country,
+  CASE WHEN grouping(country) = 0 THEN COALESCE(MAX(NULLIF(continent, '')), '') END AS continent,
  ${aggregateColumns}
-FROM weekly
-GROUP BY week_key, tier, ${dimensionSql.groupBy}
-UNION ALL
-SELECT
-  'day' AS grain,
-  day_key AS period_key,
-  ${sqlString(Resource.StatsSyncConfig.dataset)} AS dataset,
-  tier,
-  ${dimensionSql.select},
-  ${aggregateColumns}
-FROM daily
-GROUP BY day_key, tier, ${dimensionSql.groupBy}
+FROM periods
+GROUP BY GROUPING SETS (
+  (week_key, tier, provider, model),
+  (week_key, tier, provider),
+  (week_key, tier, country),
+  (week_key, tier, provider, model, country),
+  (day_key, tier, provider, model),
+  (day_key, tier, provider),
+  (day_key, tier, country),
+  (day_key, tier, provider, model, country)
+)
 ORDER BY grain, period_key, total_tokens DESC
 `
 }
@@ -1,4 +1,4 @@
-import { and, asc, eq, inArray, or } from "drizzle-orm"
+import { and, asc, eq, inArray, max, or } from "drizzle-orm"
 import { Effect, Layer } from "effect"
 import * as Context from "effect/Context"
 import { DatabaseError, DrizzleClient } from "../database"
@@ -43,6 +43,7 @@ export type ModelStatMetric = {
 export declare namespace ModelStatRepo {
  export interface Service {
    readonly listDaily: () => Effect.Effect<ModelStatMetric[], DatabaseError>
+    readonly lastSyncedAt: () => Effect.Effect<Date | null, DatabaseError>
    readonly upsert: (rows: ModelStatRow[]) => Effect.Effect<void, DatabaseError>
    readonly deleteRetiredDimensions: (rows: ModelStatRow[]) => Effect.Effect<void, DatabaseError>
  }
@@ -111,6 +112,14 @@ export class ModelStatRepo extends Context.Service<ModelStatRepo, ModelStatRepo.
        })
      })

+      const lastSyncedAt = Effect.fn("ModelStatRepo.lastSyncedAt")(function* () {
+        const result = yield* Effect.tryPromise({
+          try: () => db.select({ value: max(modelStat.updated_at) }).from(modelStat),
+          catch: (cause) => DatabaseError.make({ cause }),
+        })
+        return result[0]?.value ?? null
+      })
+
      const upsert = Effect.fn("ModelStatRepo.upsert")(function* (rows: ModelStatRow[]) {
        yield* Effect.forEach(
          chunks(rows, UPSERT_CHUNK_SIZE),
@@ -192,7 +201,7 @@ export class ModelStatRepo extends Context.Service<ModelStatRepo, ModelStatRepo.
        })
      })

-      return ModelStatRepo.of({ listDaily, upsert, deleteRetiredDimensions })
+      return ModelStatRepo.of({ listDaily, lastSyncedAt, upsert, deleteRetiredDimensions })
    }),
  )
 }
@@ -12,24 +12,22 @@ const DATALAKE_INGESTION_LAG_MS = 5 * 60_000
 const STATS_DATA_START_MS = new Date("2026-05-28T00:00:00.000Z").getTime()
 const WEEK_MS = 7 * 86_400_000
 const DISPLAY_WINDOW_MS = 56 * 86_400_000
+// Anchor incremental passes to the ISO week containing this lookback, so the pass
+// after a week boundary still recomputes the previous week's final aggregates even
+// if the boundary pass itself failed.
+const INCREMENTAL_LOOKBACK_MS = 2 * 3_600_000

 export type SyncStatsResult = { ok: true; rows: number; startedAt: string; periodStart: string; periodEnd: string }
 export type SyncStatsError = AthenaQueryError | AthenaQueryTimeoutError | DatabaseError

-export const syncStats: () => Effect.Effect<
+export const syncStats: (options?: { full?: boolean }) => Effect.Effect<
  SyncStatsResult,
  SyncStatsError,
  Athena | ModelStatRepo | ProviderStatRepo | GeoStatRepo
-> = Effect.fn("StatSync.sync")(function* () {
+> = Effect.fn("StatSync.sync")(function* (options?: { full?: boolean }) {
  const startedAt = yield* DateTime.nowAsDate
  const periodEnd = new Date(Math.floor((startedAt.getTime() - DATALAKE_INGESTION_LAG_MS) / 60_000) * 60_000)
-  // May 27 was partial, so keep Athena stats anchored at the first complete day.
-  const periodStart = new Date(
-    Math.max(
-      Math.min(startOfIsoWeek(periodEnd).getTime() - WEEK_MS, periodEnd.getTime() - DISPLAY_WINDOW_MS),
-      STATS_DATA_START_MS,
-    ),
-  )
+  const periodStart = options?.full ? fullPeriodStart(periodEnd) : incrementalPeriodStart(periodEnd)
  const athena = yield* Athena
  const modelStats = yield* ModelStatRepo
  const providerStats = yield* ProviderStatRepo
@@ -37,26 +35,16 @@ export const syncStats: () => Effect.Effect<

  yield* logRuntimeCheck()

-  const [modelAggregates, providerAggregates, geoAggregates, geoModelAggregates] = yield* Effect.all(
-    [
-      athena
-        .query(buildStatsQuery(periodStart, periodEnd, "model"))
-        .pipe(Effect.map((rows) => rows.flatMap(toModelAggregate))),
-      athena
-        .query(buildStatsQuery(periodStart, periodEnd, "provider"))
-        .pipe(Effect.map((rows) => rows.flatMap(toProviderAggregate))),
-      athena
-        .query(buildStatsQuery(periodStart, periodEnd, "geo"))
-        .pipe(Effect.map((rows) => rows.flatMap(toGeoAggregate))),
-      athena
-        .query(buildStatsQuery(periodStart, periodEnd, "geo_model"))
-        .pipe(Effect.map((rows) => rows.flatMap(toGeoAggregate))),
-    ],
-    { concurrency: "unbounded" },
+  const rows = yield* athena.query(buildStatsQuery(periodStart, periodEnd))
+  const modelRows = modelRowsFromAggregates(
+    rows.filter((row) => row.dimension === "model").flatMap(toModelAggregate),
+  )
+  const providerRows = providerRowsFromAggregates(
+    rows.filter((row) => row.dimension === "provider").flatMap(toProviderAggregate),
+  )
+  const geoRows = geoRowsFromAggregates(
+    rows.filter((row) => row.dimension === "geo" || row.dimension === "geo_model").flatMap(toGeoAggregate),
  )
-  const modelRows = modelRowsFromAggregates(modelAggregates)
-  const providerRows = providerRowsFromAggregates(providerAggregates)
-  const geoRows = geoRowsFromAggregates([...geoAggregates, ...geoModelAggregates])

  yield* Effect.all([modelStats.upsert(modelRows), providerStats.upsert(providerRows), geoStats.upsert(geoRows)], {
    concurrency: "unbounded",
@@ -92,6 +80,29 @@ export const syncStats: () => Effect.Effect<
  }
 })

+// May 27 was partial, so keep Athena stats anchored at the first complete day.
+function fullPeriodStart(periodEnd: Date) {
+  return new Date(
+    Math.max(
+      Math.min(startOfIsoWeek(periodEnd).getTime() - WEEK_MS, periodEnd.getTime() - DISPLAY_WINDOW_MS),
+      STATS_DATA_START_MS,
+    ),
+  )
+}
+
+// Events are append-only, so completed periods never change once synced; hourly
+// passes only recompute the periods the current ISO week can still touch. The daily
+// full pass refreshes the whole display window (normalization changes, retired
+// dimension cleanup).
+function incrementalPeriodStart(periodEnd: Date) {
+  return new Date(
+    Math.max(
+      startOfIsoWeek(new Date(periodEnd.getTime() - INCREMENTAL_LOOKBACK_MS)).getTime(),
+      STATS_DATA_START_MS,
+    ),
+  )
+}
+
 function logRuntimeCheck() {
  return Effect.logInfo(
    `athena stats runtime check ${JSON.stringify({
@@ -1,21 +1,49 @@
 import * as NodeRuntime from "@effect/platform-node/NodeRuntime"
 import { Athena } from "@opencode-ai/stats-core/athena"
+import { ModelStatRepo } from "@opencode-ai/stats-core/domain/model"
 import { layer as statsLayer } from "@opencode-ai/stats-core/runtime"
 import { syncStats } from "@opencode-ai/stats-core/stat-sync"
-import { Cause, Effect, Layer, Schedule } from "effect"
+import { Cause, Duration, Effect, Layer, Schedule } from "effect"

 const SYNC_INTERVAL = "1 hour"
+const SYNC_INTERVAL_MS = 3_600_000

 const runtimeLayer = Layer.mergeAll(statsLayer, Athena.layer)
-const syncPass = syncStats().pipe(
-  Effect.catchCause((cause) =>
-    Effect.logWarning(`stats sync failed ${JSON.stringify({ cause: Cause.pretty(cause) })}`),
-  ),
-)
-const daemon = Effect.logInfo("stats sync daemon started").pipe(
-  Effect.andThen(syncPass.pipe(Effect.repeat(Schedule.fixed(SYNC_INTERVAL)))),
-  Effect.forkScoped,
-)
+
+const daemon = Effect.gen(function* () {
+  yield* Effect.logInfo("stats sync daemon started")
+  yield* initialDelay()
+
+  // One full pass per UTC day (including the first pass after boot) refreshes the
+  // whole display window; every other pass only recomputes the current ISO week.
+  let lastFullDay = ""
+  const pass = Effect.gen(function* () {
+    const today = new Date().toISOString().slice(0, 10)
+    const full = lastFullDay !== today
+    yield* syncStats({ full })
+    if (full) lastFullDay = today
+  }).pipe(
+    Effect.catchCause((cause) =>
+      Effect.logWarning(`stats sync failed ${JSON.stringify({ cause: Cause.pretty(cause) })}`),
+    ),
+  )
+  yield* pass.pipe(Effect.repeat(Schedule.fixed(SYNC_INTERVAL)))
+}).pipe(Effect.forkScoped)
+
+// A restarted daemon must not immediately re-run the expensive Athena pass; resume
+// the hourly cadence from the last completed sync instead. This caps the Athena
+// spend of a crash loop at one pass per interval.
+const initialDelay = Effect.fnUntraced(function* () {
+  const modelStats = yield* ModelStatRepo
+  const lastSynced = yield* modelStats.lastSyncedAt().pipe(Effect.catchCause(() => Effect.succeed(null)))
+  if (!lastSynced) return
+  const delayMs = Math.min(SYNC_INTERVAL_MS - (Date.now() - lastSynced.getTime()), SYNC_INTERVAL_MS)
+  if (delayMs <= 0) return
+  yield* Effect.logInfo(
+    `stats sync delaying first pass ${JSON.stringify({ lastSyncedAt: lastSynced.toISOString(), delayMs })}`,
+  )
+  yield* Effect.sleep(Duration.millis(delayMs))
+})

 NodeRuntime.runMain(Layer.launch(Layer.effectDiscard(daemon).pipe(Layer.provide(runtimeLayer))), {
  disableErrorReporting: true,