Skip to content

Commit 0567840

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. The deltas the aggregator hands over now carry 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. BucketDelta held exactly those two sums and had no other purpose, so it collapses to a plain ResourceSlot per bucket. 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: 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, and dropping a typmod does not rewrite the table. duration_seconds goes too. It only ever existed so readers could reconstitute the product, nothing consulted it on its own, and it counted kernel-seconds rather than wall-clock, so keeping it would leave a column that invites the same misreading. 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 0567840

12 files changed

Lines changed: 618 additions & 202 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/__init__.py‎

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
from .types import (
2-
BucketDelta,
32
DomainFactorResult,
43
DomainFairShareData,
54
DomainFairShareSearchResult,
@@ -59,6 +58,5 @@
5958
"UserUsageBucketKey",
6059
"ProjectUsageBucketKey",
6160
"DomainUsageBucketKey",
62-
"BucketDelta",
6361
"UsageBucketAggregationResult",
6462
)

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

Lines changed: 9 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -416,23 +416,15 @@ class DomainUsageBucketKey:
416416

417417

418418
@dataclass
419-
class BucketDelta:
420-
"""Separated resource amount and duration for a usage bucket.
419+
class UsageBucketAggregationResult:
420+
"""Resource-seconds to add to each bucket, from one observation tick.
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+
Each value is ``sum(amount_k * duration_k)`` over the slices folded into that
423+
bucket. It must be accumulated as a sum of per-slice products: summing the
424+
amounts and the durations separately and multiplying afterwards gives a cross
425+
product inflated by the number of slices.
426426
"""
427427

428-
slots: ResourceSlot = field(default_factory=ResourceSlot)
429-
duration_seconds: int = 0
430-
431-
432-
@dataclass
433-
class UsageBucketAggregationResult:
434-
"""Result of aggregating kernel usage into daily buckets."""
435-
436-
user_usage_deltas: dict[UserUsageBucketKey, BucketDelta] = field(default_factory=dict)
437-
project_usage_deltas: dict[ProjectUsageBucketKey, BucketDelta] = field(default_factory=dict)
438-
domain_usage_deltas: dict[DomainUsageBucketKey, BucketDelta] = field(default_factory=dict)
428+
user_usage_deltas: dict[UserUsageBucketKey, ResourceSlot] = field(default_factory=dict)
429+
project_usage_deltas: dict[ProjectUsageBucketKey, ResourceSlot] = field(default_factory=dict)
430+
domain_usage_deltas: dict[DomainUsageBucketKey, ResourceSlot] = field(default_factory=dict)

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

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

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

Lines changed: 10 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -442,10 +442,15 @@ 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.
449454
450455
One entry per (bucket_id, slot_name). ``bucket_type`` is a discriminator
451456
indicating which parent table (domain/project/user_usage_buckets) owns
@@ -457,10 +462,7 @@ class UsageBucketEntryRow(Base): # type: ignore[misc]
457462
bucket_id: Mapped[uuid.UUID] = mapped_column("bucket_id", GUID(), nullable=False)
458463
bucket_type: Mapped[str] = mapped_column("bucket_type", sa.String(length=16), nullable=False)
459464
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-
)
463-
duration_seconds: Mapped[int] = mapped_column("duration_seconds", sa.Integer(), nullable=False)
465+
resource_usage: Mapped[Decimal] = mapped_column("resource_usage", sa.Numeric(), nullable=False)
464466
capacity: Mapped[Decimal] = mapped_column(
465467
"capacity", sa.Numeric(precision=24, scale=6), nullable=False
466468
)

‎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)