Skip to content

Commit 0f9b02b

Browse files
jopemachineclaude
jopemachine
and
claude
committed
fix(BA-6927): stop cross-multiplying usage bucket amounts and durations
FairShareAggregator accumulated the raw resource amounts and the slice durations of a bucket separately and multiplied them afterwards, yielding (sum amount_k) * (sum duration_k) instead of sum(amount_k * duration_k). The cross product inflates a bucket by the number of kernel slices folded into it, so a user running 4 GPU kernels saw 4x their real fGPU-seconds. The factor compounds up the hierarchy, since each level aggregates the kernels of everything beneath it. BucketDelta now carries the per-slice products already summed, which is the only shape that fixes this: (sum amount, sum duration) does not determine sum(amount * duration), so the pairing has to survive into the delta. usage_bucket_entries.amount becomes resource_usage, the name kernel_usage_ records, the three bucket tables and the GraphQL and REST responses already use for this quantity, so one vocabulary runs the whole way through. The column also drops its precision limit and holds the product directly rather than a factor readers had to multiply back out. No fixed precision is defensible here -- a domain-level daily mem bucket runs to ~1e18 byte-seconds on a large cluster -- while PostgreSQL's unconstrained numeric has no ceiling. Renaming a column and dropping a typmod are both metadata-only; neither rewrites the table. duration_seconds is no longer read anywhere; it stays as provenance. The read paths previously summed the stored figure without combining it with duration at all, returning something 300x smaller than the resource-seconds the fair share calculator assumes (it divides by capacity * lookback_days * 86400), so scheduling priorities were computed on the wrong scale -- a second defect independent of the cross product. Existing buckets cannot be corrected in place, in either the JSONB mirror or the normalized entries, so the migration rebuilds both from kernel_usage_records, which stores per-slice products and was never affected. The oldest retained day is excluded because retention purges kernel records by period_end and may have truncated it; buckets older than the kernel record retention window keep their inflated values rather than being silently zeroed. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
1 parent 1b14eef commit 0f9b02b

11 files changed

Lines changed: 637 additions & 125 deletions

File tree

‎changes/12965.fix.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1 @@
1+
Fix user, project and domain usage buckets being over-reported by the number of concurrently running kernels, and rebuild the affected buckets from the recorded per-kernel usage.

‎src/ai/backend/manager/data/fair_share/types.py‎

Lines changed: 10 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -417,15 +417,19 @@ class DomainUsageBucketKey:
417417

418418
@dataclass
419419
class BucketDelta:
420-
"""Separated resource amount and duration for a usage bucket.
420+
"""Accumulated resource usage and observation duration for a usage bucket.
421421
422-
Stores raw resource amounts and duration separately instead of
423-
pre-multiplied resource-seconds. The product ``amount * duration_seconds``
424-
is computed at SQL query time where PostgreSQL auto-extends NUMERIC precision,
425-
eliminating overflow risk for large memory values.
422+
``resource_usage`` is in resource-seconds and must be accumulated as
423+
``sum(amount_k * duration_k)``. Summing amounts and durations separately
424+
and multiplying afterwards gives a cross product inflated by the number of
425+
slices in the bucket.
426+
427+
``duration_seconds`` is the total observed duration of those slices. The
428+
two divide -- never multiply -- into the mean amount that
429+
``usage_bucket_entries`` stores.
426430
"""
427431

428-
slots: ResourceSlot = field(default_factory=ResourceSlot)
432+
resource_usage: ResourceSlot = field(default_factory=ResourceSlot)
429433
duration_seconds: int = 0
430434

431435

‎src/ai/backend/manager/models/alembic/versions/c4a91d7e05b2_rebuild_inflated_usage_buckets.py‎

Lines changed: 210 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,210 @@
1+
"""rebuild inflated usage buckets
2+
3+
The aggregator summed amounts and durations separately and multiplied them
4+
afterwards, so both the JSONB ``resource_usage`` mirror and the normalized
5+
entries hold a cross product. Neither can be corrected in place, so both are
6+
rebuilt from ``kernel_usage_records``, which was never affected.
7+
8+
``usage_bucket_entries.amount`` becomes ``resource_usage``, matching what the
9+
JSONB mirror on ``kernel_usage_records`` and the three bucket tables already call
10+
this quantity, and drops its precision limit. It now holds the product directly
11+
rather than a factor readers had to multiply back out: a domain-level daily mem
12+
bucket runs past any fixed precision, and unconstrained NUMERIC has no ceiling.
13+
Renaming a column and dropping a typmod are both metadata-only; neither rewrites
14+
the table.
15+
16+
Revision ID: c4a91d7e05b2
17+
Revises: 3f9a1c7b2e04
18+
Create Date: 2026-07-20 00:00:00.000000
19+
20+
"""
21+
22+
from datetime import date, timedelta
23+
24+
import sqlalchemy as sa
25+
from alembic import op
26+
27+
# revision identifiers, used by Alembic.
28+
revision = "c4a91d7e05b2"
29+
down_revision = "3f9a1c7b2e04"
30+
# Part of: NEXT_RELEASE_VERSION
31+
branch_labels = None
32+
depends_on = None
33+
34+
35+
# Bucket tables keyed by the columns that identify the owning entity. Every key
36+
# column has the same name in kernel_usage_records, so the join is by name.
37+
_BUCKET_LEVELS: list[tuple[str, str, list[str]]] = [
38+
("user_usage_buckets", "user", ["user_uuid", "project_id", "resource_group"]),
39+
("project_usage_buckets", "project", ["project_id", "resource_group"]),
40+
("domain_usage_buckets", "domain", ["domain_name", "resource_group"]),
41+
]
42+
43+
44+
def upgrade() -> None:
45+
op.alter_column(
46+
"usage_bucket_entries",
47+
"amount",
48+
new_column_name="resource_usage",
49+
existing_type=sa.Numeric(precision=24, scale=6),
50+
type_=sa.Numeric(),
51+
existing_nullable=False,
52+
)
53+
54+
conn = op.get_bind()
55+
covered = _covered_date_range(conn)
56+
if covered is None:
57+
# No usage records to rebuild from (fresh install, or everything purged).
58+
return
59+
rebuild_from, rebuild_to = covered
60+
for table_name, bucket_type, key_columns in _BUCKET_LEVELS:
61+
_rebuild_buckets(conn, table_name, bucket_type, key_columns, rebuild_from, rebuild_to)
62+
63+
64+
def downgrade() -> None:
65+
# The rebuilt values are the correct ones and the inflated originals cannot
66+
# be reconstructed, so only the column definition is reverted. Values that
67+
# exceed the restored precision will fail the cast, which is the honest
68+
# outcome: they do not fit the old column.
69+
op.alter_column(
70+
"usage_bucket_entries",
71+
"resource_usage",
72+
new_column_name="amount",
73+
existing_type=sa.Numeric(),
74+
type_=sa.Numeric(precision=24, scale=6),
75+
existing_nullable=False,
76+
)
77+
78+
79+
def _covered_date_range(conn: sa.engine.Connection) -> tuple[date, date] | None:
80+
"""Return the date range that kernel_usage_records can faithfully rebuild.
81+
82+
The oldest retained day is excluded because retention purges by ``period_end``
83+
and may have truncated it. Buckets outside the range keep their inflated
84+
values rather than being zeroed: they are unrecoverable, and zeroing them
85+
would destroy the only usage history left.
86+
"""
87+
row = conn.execute(
88+
sa.text(
89+
"SELECT min((period_start AT TIME ZONE 'UTC')::date) AS min_date, "
90+
" max((period_start AT TIME ZONE 'UTC')::date) AS max_date "
91+
"FROM kernel_usage_records"
92+
)
93+
).one()
94+
if row.min_date is None or row.max_date is None:
95+
return None
96+
rebuild_from = row.min_date + timedelta(days=1)
97+
if rebuild_from > row.max_date:
98+
return None
99+
return rebuild_from, row.max_date
100+
101+
102+
def _rebuild_buckets(
103+
conn: sa.engine.Connection,
104+
table_name: str,
105+
bucket_type: str,
106+
key_columns: list[str],
107+
rebuild_from: date,
108+
rebuild_to: date,
109+
) -> None:
110+
"""Recompute one bucket level's entries and JSONB from kernel_usage_records."""
111+
key_list = ", ".join(key_columns)
112+
join_on = " AND ".join(f"b.{col} = agg.{col}" for col in key_columns)
113+
params = {"rebuild_from": rebuild_from, "rebuild_to": rebuild_to}
114+
115+
# Per-slot resource-seconds, plus the duration they span. The duration is
116+
# the same for every slot of a bucket, so it is aggregated separately.
117+
agg_cte = f"""
118+
WITH agg AS (
119+
SELECT {key_list},
120+
(period_start AT TIME ZONE 'UTC')::date AS period_date,
121+
kv.key AS slot_name,
122+
SUM(kv.value::numeric) AS resource_seconds
123+
FROM kernel_usage_records,
124+
LATERAL jsonb_each_text(resource_usage) AS kv
125+
WHERE (period_start AT TIME ZONE 'UTC')::date
126+
BETWEEN :rebuild_from AND :rebuild_to
127+
GROUP BY {key_list}, period_date, kv.key
128+
),
129+
agg_duration AS (
130+
SELECT {key_list},
131+
(period_start AT TIME ZONE 'UTC')::date AS period_date,
132+
SUM(EXTRACT(EPOCH FROM (period_end - period_start)))::bigint
133+
AS duration_seconds
134+
FROM kernel_usage_records
135+
WHERE (period_start AT TIME ZONE 'UTC')::date
136+
BETWEEN :rebuild_from AND :rebuild_to
137+
GROUP BY {key_list}, period_date
138+
)
139+
"""
140+
141+
# capacity is refilled by the next observation tick, so dropping it is safe.
142+
conn.execute(
143+
sa.text(
144+
f"""
145+
DELETE FROM usage_bucket_entries e
146+
USING {table_name} b
147+
WHERE e.bucket_id = b.id
148+
AND e.bucket_type = :bucket_type
149+
AND b.period_start BETWEEN :rebuild_from AND :rebuild_to
150+
"""
151+
),
152+
{**params, "bucket_type": bucket_type},
153+
)
154+
conn.execute(
155+
sa.text(
156+
f"""
157+
{agg_cte}
158+
INSERT INTO usage_bucket_entries
159+
(bucket_id, bucket_type, slot_name, resource_usage,
160+
duration_seconds, capacity)
161+
SELECT b.id, :bucket_type, agg.slot_name, agg.resource_seconds,
162+
agg_duration.duration_seconds, 0
163+
FROM agg
164+
JOIN agg_duration USING ({key_list}, period_date)
165+
JOIN {table_name} b
166+
ON {join_on}
167+
AND b.period_start = agg.period_date
168+
ON CONFLICT (bucket_id, slot_name) DO UPDATE
169+
SET resource_usage = EXCLUDED.resource_usage,
170+
duration_seconds = EXCLUDED.duration_seconds
171+
"""
172+
),
173+
{**params, "bucket_type": bucket_type},
174+
)
175+
176+
# Buckets with no matching kernel records collapse to an empty slot map.
177+
conn.execute(
178+
sa.text(
179+
f"""
180+
{agg_cte},
181+
per_bucket AS (
182+
SELECT {key_list}, period_date,
183+
jsonb_object_agg(slot_name, resource_seconds) AS usage
184+
FROM agg
185+
GROUP BY {key_list}, period_date
186+
)
187+
UPDATE {table_name} b
188+
SET resource_usage = COALESCE(agg.usage, '{{}}'::jsonb)
189+
FROM per_bucket agg
190+
WHERE {join_on}
191+
AND b.period_start = agg.period_date
192+
AND b.period_start BETWEEN :rebuild_from AND :rebuild_to
193+
"""
194+
),
195+
params,
196+
)
197+
conn.execute(
198+
sa.text(
199+
f"""
200+
UPDATE {table_name} b
201+
SET resource_usage = '{{}}'::jsonb
202+
WHERE b.period_start BETWEEN :rebuild_from AND :rebuild_to
203+
AND NOT EXISTS (
204+
SELECT 1 FROM usage_bucket_entries e
205+
WHERE e.bucket_id = b.id AND e.bucket_type = :bucket_type
206+
)
207+
"""
208+
),
209+
{**params, "bucket_type": bucket_type},
210+
)

‎src/ai/backend/manager/models/resource_usage_history/row.py‎

Lines changed: 13 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -442,10 +442,18 @@ class UserUsageBucketRow(LifecycleTimestampsMixin, Base): # type: ignore[misc]
442442
class UsageBucketEntryRow(Base): # type: ignore[misc]
443443
"""Per-slot normalized entry for usage bucket aggregation (Phase 3).
444444
445-
Stores amount and duration separately instead of pre-multiplied resource-seconds,
446-
eliminating overflow risk for large memory values.
447-
The product ``amount * duration_seconds`` is computed at SQL query time
448-
where PostgreSQL auto-extends NUMERIC precision.
445+
``resource_usage`` accumulates a kernel's ``occupied_slots`` integrated
446+
over the time it held them. It records an allocation, not a measurement: a
447+
kernel holding a GPU idle for an hour counts the same as one saturating it.
448+
That is what fair share wants -- holding a resource denies it to others -- but
449+
it means "usage" would overstate what this column knows.
450+
451+
Declared as unconstrained NUMERIC on purpose: a domain-level daily mem bucket
452+
runs to ~1e18 byte-seconds on a large cluster, past any fixed precision worth
453+
writing down, and PostgreSQL's unconstrained numeric has no such ceiling.
454+
455+
``duration_seconds`` records the total observed duration. Nothing reads it
456+
today -- it is kept as provenance for the figure beside it.
449457
450458
One entry per (bucket_id, slot_name). ``bucket_type`` is a discriminator
451459
indicating which parent table (domain/project/user_usage_buckets) owns
@@ -457,9 +465,7 @@ class UsageBucketEntryRow(Base): # type: ignore[misc]
457465
bucket_id: Mapped[uuid.UUID] = mapped_column("bucket_id", GUID(), nullable=False)
458466
bucket_type: Mapped[str] = mapped_column("bucket_type", sa.String(length=16), nullable=False)
459467
slot_name: Mapped[str] = mapped_column("slot_name", sa.String(length=64), nullable=False)
460-
amount: Mapped[Decimal] = mapped_column(
461-
"amount", sa.Numeric(precision=24, scale=6), nullable=False
462-
)
468+
resource_usage: Mapped[Decimal] = mapped_column("resource_usage", sa.Numeric(), nullable=False)
463469
duration_seconds: Mapped[int] = mapped_column("duration_seconds", sa.Integer(), nullable=False)
464470
capacity: Mapped[Decimal] = mapped_column(
465471
"capacity", sa.Numeric(precision=24, scale=6), nullable=False

‎src/ai/backend/manager/repositories/fair_share/db_source/db_source.py‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1585,7 +1585,7 @@ async def _fetch_raw_usage_buckets(
15851585
UserUsageBucketRow.project_id,
15861586
UserUsageBucketRow.period_start,
15871587
ube.c.slot_name,
1588-
ube.c.amount,
1588+
ube.c.resource_usage.label("resource_usage"),
15891589
)
15901590
.select_from(
15911591
sa.join(
@@ -1612,7 +1612,7 @@ async def _fetch_raw_usage_buckets(
16121612
ProjectUsageBucketRow.project_id,
16131613
ProjectUsageBucketRow.period_start,
16141614
ube.c.slot_name,
1615-
ube.c.amount,
1615+
ube.c.resource_usage.label("resource_usage"),
16161616
)
16171617
.select_from(
16181618
sa.join(
@@ -1639,7 +1639,7 @@ async def _fetch_raw_usage_buckets(
16391639
DomainUsageBucketRow.domain_name,
16401640
DomainUsageBucketRow.period_start,
16411641
ube.c.slot_name,
1642-
ube.c.amount,
1642+
ube.c.resource_usage.label("resource_usage"),
16431643
)
16441644
.select_from(
16451645
sa.join(
@@ -1668,7 +1668,7 @@ async def _fetch_raw_usage_buckets(
16681668
user_buckets[key] = {}
16691669
if row.period_start not in user_buckets[key]:
16701670
user_buckets[key][row.period_start] = ResourceSlot()
1671-
user_buckets[key][row.period_start][row.slot_name] = Decimal(str(row.amount))
1671+
user_buckets[key][row.period_start][row.slot_name] = Decimal(str(row.resource_usage))
16721672

16731673
project_buckets: dict[uuid.UUID, dict[date, ResourceSlot]] = {}
16741674
for row in project_rows:
@@ -1677,7 +1677,7 @@ async def _fetch_raw_usage_buckets(
16771677
if row.period_start not in project_buckets[row.project_id]:
16781678
project_buckets[row.project_id][row.period_start] = ResourceSlot()
16791679
project_buckets[row.project_id][row.period_start][row.slot_name] = Decimal(
1680-
str(row.amount)
1680+
str(row.resource_usage)
16811681
)
16821682

16831683
domain_buckets: dict[str, dict[date, ResourceSlot]] = {}
@@ -1687,7 +1687,7 @@ async def _fetch_raw_usage_buckets(
16871687
if row.period_start not in domain_buckets[row.domain_name]:
16881688
domain_buckets[row.domain_name][row.period_start] = ResourceSlot()
16891689
domain_buckets[row.domain_name][row.period_start][row.slot_name] = Decimal(
1690-
str(row.amount)
1690+
str(row.resource_usage)
16911691
)
16921692

16931693
return RawUsageBucketsByLevel(

0 commit comments

Comments
 (0)