From 2d41a41328db603db02376d7006bd40cb2ba0ecb Mon Sep 17 00:00:00 2001 From: Gyubong Date: Mon, 20 Jul 2026 13:09:02 +0900 Subject: [PATCH 1/2] 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.slots becomes resource_usage -- it no longer holds an allocation but the allocation already integrated over the slice, which is what every other carrier of this quantity is called (the resource_usage column on kernel_usage_records and on all three bucket tables). Naming it resource_seconds would have put it next to duration_seconds and read as a second duration. The three repository call sites no longer re-multiply by the accumulated duration. usage_bucket_entries.amount widens to 32 digits: it stores a daily total per entity, and a domain-level mem bucket on a large cluster exceeds the previous 1e18 ceiling. 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 resource-seconds 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) --- changes/12965.fix.md | 1 + .../backend/manager/data/fair_share/types.py | 15 +- ...1d7e05b2_rebuild_inflated_usage_buckets.py | 185 +++++++++ .../models/resource_usage_history/row.py | 9 +- .../db_source/db_source.py | 37 +- .../scheduler/fair_share/aggregator.py | 14 +- .../test_usage_bucket_entries.py | 53 +-- .../test_aggregator_bucket_aggregation.py | 373 ++++++++++++++++-- 8 files changed, 582 insertions(+), 105 deletions(-) create mode 100644 changes/12965.fix.md create mode 100644 src/ai/backend/manager/models/alembic/versions/c4a91d7e05b2_rebuild_inflated_usage_buckets.py diff --git a/changes/12965.fix.md b/changes/12965.fix.md new file mode 100644 index 000000000000..4e1835190a80 --- /dev/null +++ b/changes/12965.fix.md @@ -0,0 +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. diff --git a/src/ai/backend/manager/data/fair_share/types.py b/src/ai/backend/manager/data/fair_share/types.py index 0463692d2124..d709c54aeee4 100644 --- a/src/ai/backend/manager/data/fair_share/types.py +++ b/src/ai/backend/manager/data/fair_share/types.py @@ -417,15 +417,18 @@ class DomainUsageBucketKey: @dataclass class BucketDelta: - """Separated resource amount and duration for a usage bucket. + """Accumulated resource usage and observation duration for a usage bucket. - Stores raw resource amounts and duration separately instead of - pre-multiplied resource-seconds. The product ``amount * duration_seconds`` - is computed at SQL query time where PostgreSQL auto-extends NUMERIC precision, - eliminating overflow risk for large memory values. + ``resource_usage`` is in resource-seconds and must be accumulated as + ``sum(amount_k * duration_k)``. Summing amounts and durations separately + and multiplying afterwards gives a cross product inflated by the number of + slices in the bucket. + + ``duration_seconds`` is for reporting only; both fields are sums, so they + are never multiplied with each other. """ - slots: ResourceSlot = field(default_factory=ResourceSlot) + resource_usage: ResourceSlot = field(default_factory=ResourceSlot) duration_seconds: int = 0 diff --git a/src/ai/backend/manager/models/alembic/versions/c4a91d7e05b2_rebuild_inflated_usage_buckets.py b/src/ai/backend/manager/models/alembic/versions/c4a91d7e05b2_rebuild_inflated_usage_buckets.py new file mode 100644 index 000000000000..b776086274be --- /dev/null +++ b/src/ai/backend/manager/models/alembic/versions/c4a91d7e05b2_rebuild_inflated_usage_buckets.py @@ -0,0 +1,185 @@ +"""rebuild inflated usage buckets + +The aggregator summed amounts and durations separately and multiplied them +afterwards, so both the JSONB ``resource_usage`` mirror and the normalized +entries hold a cross product. Neither can be corrected in place, so both are +rebuilt from ``kernel_usage_records``, which was never affected. + +``usage_bucket_entries.amount`` also widens, because a domain-level daily mem +bucket can exceed the previous 24-digit ceiling. + +Revision ID: c4a91d7e05b2 +Revises: 3f9a1c7b2e04 +Create Date: 2026-07-20 00:00:00.000000 + +""" + +from datetime import date, timedelta + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "c4a91d7e05b2" +down_revision = "3f9a1c7b2e04" +# Part of: NEXT_RELEASE_VERSION +branch_labels = None +depends_on = None + + +# Bucket tables keyed by the columns that identify the owning entity. Every key +# column has the same name in kernel_usage_records, so the join is by name. +_BUCKET_LEVELS: list[tuple[str, str, list[str]]] = [ + ("user_usage_buckets", "user", ["user_uuid", "project_id", "resource_group"]), + ("project_usage_buckets", "project", ["project_id", "resource_group"]), + ("domain_usage_buckets", "domain", ["domain_name", "resource_group"]), +] + + +def upgrade() -> None: + op.alter_column( + "usage_bucket_entries", + "amount", + existing_type=sa.Numeric(precision=24, scale=6), + type_=sa.Numeric(precision=32, scale=6), + existing_nullable=False, + ) + + conn = op.get_bind() + covered = _covered_date_range(conn) + if covered is None: + # No usage records to rebuild from (fresh install, or everything purged). + return + rebuild_from, rebuild_to = covered + for table_name, bucket_type, key_columns in _BUCKET_LEVELS: + _rebuild_buckets(conn, table_name, bucket_type, key_columns, rebuild_from, rebuild_to) + + +def downgrade() -> None: + # Keeps the rebuilt values; only the widened precision is reverted. + op.alter_column( + "usage_bucket_entries", + "amount", + existing_type=sa.Numeric(precision=32, scale=6), + type_=sa.Numeric(precision=24, scale=6), + existing_nullable=False, + ) + + +def _covered_date_range(conn: sa.engine.Connection) -> tuple[date, date] | None: + """Return the date range that kernel_usage_records can faithfully rebuild. + + The oldest retained day is excluded because retention purges by ``period_end`` + and may have truncated it. Buckets outside the range keep their inflated + values rather than being zeroed: they are unrecoverable, and zeroing them + would destroy the only usage history left. + """ + row = conn.execute( + sa.text( + "SELECT min((period_start AT TIME ZONE 'UTC')::date) AS min_date, " + " max((period_start AT TIME ZONE 'UTC')::date) AS max_date " + "FROM kernel_usage_records" + ) + ).one() + if row.min_date is None or row.max_date is None: + return None + rebuild_from = row.min_date + timedelta(days=1) + if rebuild_from > row.max_date: + return None + return rebuild_from, row.max_date + + +def _rebuild_buckets( + conn: sa.engine.Connection, + table_name: str, + bucket_type: str, + key_columns: list[str], + rebuild_from: date, + rebuild_to: date, +) -> None: + """Recompute one bucket level's entries and JSONB from kernel_usage_records.""" + key_list = ", ".join(key_columns) + join_on = " AND ".join(f"b.{col} = agg.{col}" for col in key_columns) + params = {"rebuild_from": rebuild_from, "rebuild_to": rebuild_to} + + # Per-slot resource-seconds, summed over every kernel slice of the day. + agg_cte = f""" + WITH agg AS ( + SELECT {key_list}, + (period_start AT TIME ZONE 'UTC')::date AS period_date, + kv.key AS slot_name, + SUM(kv.value::numeric) AS resource_seconds + FROM kernel_usage_records, + LATERAL jsonb_each_text(resource_usage) AS kv + WHERE (period_start AT TIME ZONE 'UTC')::date + BETWEEN :rebuild_from AND :rebuild_to + GROUP BY {key_list}, period_date, kv.key + ) + """ + + # capacity is refilled by the next observation tick, so dropping it is safe. + conn.execute( + sa.text( + f""" + DELETE FROM usage_bucket_entries e + USING {table_name} b + WHERE e.bucket_id = b.id + AND e.bucket_type = :bucket_type + AND b.period_start BETWEEN :rebuild_from AND :rebuild_to + """ + ), + {**params, "bucket_type": bucket_type}, + ) + conn.execute( + sa.text( + f""" + {agg_cte} + INSERT INTO usage_bucket_entries + (bucket_id, bucket_type, slot_name, amount, duration_seconds, capacity) + SELECT b.id, :bucket_type, agg.slot_name, agg.resource_seconds, 0, 0 + FROM agg + JOIN {table_name} b + ON {join_on} + AND b.period_start = agg.period_date + ON CONFLICT (bucket_id, slot_name) DO UPDATE + SET amount = EXCLUDED.amount + """ + ), + {**params, "bucket_type": bucket_type}, + ) + + # Buckets with no matching kernel records collapse to an empty slot map. + conn.execute( + sa.text( + f""" + {agg_cte}, + per_bucket AS ( + SELECT {key_list}, period_date, + jsonb_object_agg(slot_name, resource_seconds) AS usage + FROM agg + GROUP BY {key_list}, period_date + ) + UPDATE {table_name} b + SET resource_usage = COALESCE(agg.usage, '{{}}'::jsonb) + FROM per_bucket agg + WHERE {join_on} + AND b.period_start = agg.period_date + AND b.period_start BETWEEN :rebuild_from AND :rebuild_to + """ + ), + params, + ) + conn.execute( + sa.text( + f""" + UPDATE {table_name} b + SET resource_usage = '{{}}'::jsonb + WHERE b.period_start BETWEEN :rebuild_from AND :rebuild_to + AND NOT EXISTS ( + SELECT 1 FROM usage_bucket_entries e + WHERE e.bucket_id = b.id AND e.bucket_type = :bucket_type + ) + """ + ), + {**params, "bucket_type": bucket_type}, + ) diff --git a/src/ai/backend/manager/models/resource_usage_history/row.py b/src/ai/backend/manager/models/resource_usage_history/row.py index 6d98822adc16..32ca808ba8a8 100644 --- a/src/ai/backend/manager/models/resource_usage_history/row.py +++ b/src/ai/backend/manager/models/resource_usage_history/row.py @@ -442,10 +442,9 @@ class UserUsageBucketRow(LifecycleTimestampsMixin, Base): # type: ignore[misc] class UsageBucketEntryRow(Base): # type: ignore[misc] """Per-slot normalized entry for usage bucket aggregation (Phase 3). - Stores amount and duration separately instead of pre-multiplied resource-seconds, - eliminating overflow risk for large memory values. - The product ``amount * duration_seconds`` is computed at SQL query time - where PostgreSQL auto-extends NUMERIC precision. + ``amount`` holds ``sum(amount_k * duration_k)`` resource-seconds and is what + the read paths sum. ``duration_seconds`` is for reporting only; both columns + are sums, so they are never multiplied together. One entry per (bucket_id, slot_name). ``bucket_type`` is a discriminator indicating which parent table (domain/project/user_usage_buckets) owns @@ -458,7 +457,7 @@ class UsageBucketEntryRow(Base): # type: ignore[misc] bucket_type: Mapped[str] = mapped_column("bucket_type", sa.String(length=16), nullable=False) slot_name: Mapped[str] = mapped_column("slot_name", sa.String(length=64), nullable=False) amount: Mapped[Decimal] = mapped_column( - "amount", sa.Numeric(precision=24, scale=6), nullable=False + "amount", sa.Numeric(precision=32, scale=6), nullable=False ) duration_seconds: Mapped[int] = mapped_column("duration_seconds", sa.Integer(), nullable=False) capacity: Mapped[Decimal] = mapped_column( diff --git a/src/ai/backend/manager/repositories/resource_usage_history/db_source/db_source.py b/src/ai/backend/manager/repositories/resource_usage_history/db_source/db_source.py index b6fe4bc91d08..f5a23d0e3380 100644 --- a/src/ai/backend/manager/repositories/resource_usage_history/db_source/db_source.py +++ b/src/ai/backend/manager/repositories/resource_usage_history/db_source/db_source.py @@ -398,7 +398,7 @@ async def _fetch_aggregated_usage_by_user( UserUsageBucketRow.user_uuid, UserUsageBucketRow.project_id, ube.c.slot_name, - sa.func.sum(ube.c.amount).label("total_amount"), + sa.func.sum(ube.c.amount).label("total_resource_usage"), ) .select_from( sa.join( @@ -429,7 +429,7 @@ async def _fetch_aggregated_usage_by_user( key = (row.user_uuid, row.project_id) if key not in aggregated: aggregated[key] = ResourceSlot() - aggregated[key][row.slot_name] = row.total_amount + aggregated[key][row.slot_name] = row.total_resource_usage return aggregated async def get_aggregated_usage_by_project( @@ -448,7 +448,7 @@ async def get_aggregated_usage_by_project( sa.select( ProjectUsageBucketRow.project_id, ube.c.slot_name, - sa.func.sum(ube.c.amount).label("total_amount"), + sa.func.sum(ube.c.amount).label("total_resource_usage"), ) .select_from( sa.join( @@ -477,7 +477,7 @@ async def get_aggregated_usage_by_project( for row in rows: if row.project_id not in aggregated: aggregated[row.project_id] = ResourceSlot() - aggregated[row.project_id][row.slot_name] = row.total_amount + aggregated[row.project_id][row.slot_name] = row.total_resource_usage return aggregated async def get_aggregated_usage_by_domain( @@ -496,7 +496,7 @@ async def get_aggregated_usage_by_domain( sa.select( DomainUsageBucketRow.domain_name, ube.c.slot_name, - sa.func.sum(ube.c.amount).label("total_amount"), + sa.func.sum(ube.c.amount).label("total_resource_usage"), ) .select_from( sa.join( @@ -525,7 +525,7 @@ async def get_aggregated_usage_by_domain( for row in rows: if row.domain_name not in aggregated: aggregated[row.domain_name] = ResourceSlot() - aggregated[row.domain_name][row.slot_name] = row.total_amount + aggregated[row.domain_name][row.slot_name] = row.total_resource_usage return aggregated # ==================== Bucket Delta Updates ==================== @@ -587,11 +587,7 @@ async def _increment_user_usage_buckets( for key, bucket_delta in deltas.items(): lookup_key = (key.user_uuid, key.project_id, key.resource_group, key.period_date) existing_usage = existing.get(lookup_key, ResourceSlot()) - # JSONB stores resource-seconds (amount * seconds) for legacy compatibility - resource_seconds = self._calculate_resource_seconds( - bucket_delta.slots, bucket_delta.duration_seconds - ) - new_usage = existing_usage + resource_seconds + new_usage = existing_usage + bucket_delta.resource_usage # Upsert with merged usage (JSONB) stmt = ( @@ -684,10 +680,7 @@ async def _increment_project_usage_buckets( for key, bucket_delta in deltas.items(): lookup_key = (key.project_id, key.resource_group, key.period_date) existing_usage = existing.get(lookup_key, ResourceSlot()) - resource_seconds = self._calculate_resource_seconds( - bucket_delta.slots, bucket_delta.duration_seconds - ) - new_usage = existing_usage + resource_seconds + new_usage = existing_usage + bucket_delta.resource_usage # Upsert with merged usage (JSONB) stmt = ( @@ -771,10 +764,7 @@ async def _increment_domain_usage_buckets( for key, bucket_delta in deltas.items(): lookup_key = (key.domain_name, key.resource_group, key.period_date) existing_usage = existing.get(lookup_key, ResourceSlot()) - resource_seconds = self._calculate_resource_seconds( - bucket_delta.slots, bucket_delta.duration_seconds - ) - new_usage = existing_usage + resource_seconds + new_usage = existing_usage + bucket_delta.resource_usage # Upsert with merged usage (JSONB) stmt = ( @@ -851,18 +841,13 @@ async def _upsert_bucket_entries( ) -> None: """Upsert normalized usage_bucket_entries for a bucket. - For each slot in the delta, insert or update an entry row. - ``amount`` stores the raw resource amount (not pre-multiplied) - and ``duration_seconds`` stores the actual observation duration. - The product ``amount * duration_seconds`` is computed at SQL query - time where PostgreSQL auto-extends NUMERIC precision, eliminating - overflow risk for large memory values. + Both columns accumulate independently and are never multiplied together. ``capacity`` is set to 0 here; it is updated separately during fair share factor calculation when the cluster capacity is known. """ entry_table = UsageBucketEntryRow.__table__ - for slot_name, value in bucket_delta.slots.items(): + for slot_name, value in bucket_delta.resource_usage.items(): stmt = ( pg_insert(entry_table) .values( diff --git a/src/ai/backend/manager/sokovan/scheduler/fair_share/aggregator.py b/src/ai/backend/manager/sokovan/scheduler/fair_share/aggregator.py index 0386eddba819..7e38fb80a768 100644 --- a/src/ai/backend/manager/sokovan/scheduler/fair_share/aggregator.py +++ b/src/ai/backend/manager/sokovan/scheduler/fair_share/aggregator.py @@ -223,10 +223,8 @@ def _add_to_bucket_deltas( ) -> None: """Add resource usage to bucket deltas for a day. - Accumulates raw resource amounts and duration separately. - Slots are accumulated additively (sum of ``raw_slots`` across all - slices within the same bucket key) while ``duration_seconds`` tracks - total observation time. + The segment is converted to resource-seconds before accumulation; see + ``BucketDelta`` for why it cannot be multiplied afterwards. Args: spec: Original spec (for entity identifiers) @@ -237,6 +235,8 @@ def _add_to_bucket_deltas( project_deltas: Project deltas to update (mutated) domain_deltas: Domain deltas to update (mutated) """ + segment_usage = self._calculate_resource_seconds(raw_slots, segment_seconds) + # User bucket key user_key = UserUsageBucketKey( user_uuid=spec.user_uuid, @@ -248,7 +248,7 @@ def _add_to_bucket_deltas( ) ud = user_deltas[user_key] user_deltas[user_key] = BucketDelta( - slots=ud.slots + raw_slots, + resource_usage=ud.resource_usage + segment_usage, duration_seconds=ud.duration_seconds + segment_seconds, ) @@ -262,7 +262,7 @@ def _add_to_bucket_deltas( ) pd = project_deltas[project_key] project_deltas[project_key] = BucketDelta( - slots=pd.slots + raw_slots, + resource_usage=pd.resource_usage + segment_usage, duration_seconds=pd.duration_seconds + segment_seconds, ) @@ -275,7 +275,7 @@ def _add_to_bucket_deltas( ) dd = domain_deltas[domain_key] domain_deltas[domain_key] = BucketDelta( - slots=dd.slots + raw_slots, + resource_usage=dd.resource_usage + segment_usage, duration_seconds=dd.duration_seconds + segment_seconds, ) diff --git a/tests/unit/manager/repositories/resource_usage_history/test_usage_bucket_entries.py b/tests/unit/manager/repositories/resource_usage_history/test_usage_bucket_entries.py index af827906e294..be0bea122bd1 100644 --- a/tests/unit/manager/repositories/resource_usage_history/test_usage_bucket_entries.py +++ b/tests/unit/manager/repositories/resource_usage_history/test_usage_bucket_entries.py @@ -1,7 +1,8 @@ """Tests for UsageBucketEntryRow and normalized bucket entry operations. -Phase 3 (BA-4308): Verifies that usage bucket entries are correctly created, -upserted, and aggregated via the normalized usage_bucket_entries table. +Verifies that usage bucket entries are correctly created, upserted, and +aggregated via the normalized usage_bucket_entries table. Entries store +resource-seconds; the read paths sum that column directly. """ from __future__ import annotations @@ -124,8 +125,9 @@ async def test_increment_domain_buckets_creates_entries( test_domain_name: str, ) -> None: """Verify that incrementing domain buckets also writes normalized entries.""" - raw_slots = ResourceSlot({"cpu": Decimal("2"), "mem": Decimal("4096000")}) - duration = 300 # 5-minute slice + # 5-minute slice of 2 CPU / 4096000 mem + resource_usage = ResourceSlot({"cpu": Decimal("600"), "mem": Decimal("1228800000")}) + duration = 300 period = date(2024, 1, 15) result = UsageBucketAggregationResult( @@ -137,13 +139,13 @@ async def test_increment_domain_buckets_creates_entries( resource_group="default", resource_group_id=RESOURCE_GROUP_ID, period_date=period, - ): BucketDelta(slots=raw_slots, duration_seconds=duration), + ): BucketDelta(resource_usage=resource_usage, duration_seconds=duration), }, ) await db_source.increment_usage_buckets(result) - # Verify entries were created with separated amount/duration + # Verify entries were created with resource-seconds and duration async with db_with_cleanup.begin_readonly_session() as db_sess: entry_rows = ( ( @@ -161,8 +163,8 @@ async def test_increment_domain_buckets_creates_entries( slot_map = {e.slot_name: e for e in entry_rows} assert "cpu" in slot_map assert "mem" in slot_map - assert slot_map["cpu"].amount == Decimal("2") - assert slot_map["mem"].amount == Decimal("4096000") + assert slot_map["cpu"].amount == Decimal("600") + assert slot_map["mem"].amount == Decimal("1228800000") assert slot_map["cpu"].duration_seconds == 300 assert slot_map["mem"].duration_seconds == 300 @@ -181,20 +183,20 @@ async def test_increment_accumulates_entries( period_date=period, ) - # First increment: 2 CPUs for 300 seconds + # First increment: 2 CPUs for 300 seconds -> 600 CPU-seconds result1 = UsageBucketAggregationResult( user_usage_deltas={}, project_usage_deltas={}, domain_usage_deltas={ key: BucketDelta( - slots=ResourceSlot({"cpu": Decimal("2")}), + resource_usage=ResourceSlot({"cpu": Decimal("600")}), duration_seconds=300, ), }, ) await db_source.increment_usage_buckets(result1) - # Second increment: 3 CPUs for 300 seconds + # Second increment: 3 CPUs for 300 seconds -> 900 CPU-seconds replacement_resource_group_id = ResourceGroupID(uuid.uuid4()) replacement_key = DomainUsageBucketKey( domain_name=test_domain_name, @@ -207,14 +209,14 @@ async def test_increment_accumulates_entries( project_usage_deltas={}, domain_usage_deltas={ replacement_key: BucketDelta( - slots=ResourceSlot({"cpu": Decimal("3")}), + resource_usage=ResourceSlot({"cpu": Decimal("900")}), duration_seconds=300, ), }, ) await db_source.increment_usage_buckets(result2) - # Verify accumulated: amount = 2 + 3 = 5, duration = 300 + 300 = 600 + # Verify accumulated: resource_usage = 600 + 900 = 1500, duration = 600 async with db_with_cleanup.begin_readonly_session() as db_sess: entry_rows = ( ( @@ -230,7 +232,7 @@ async def test_increment_accumulates_entries( assert len(entry_rows) == 1 assert entry_rows[0].slot_name == "cpu" - assert entry_rows[0].amount == Decimal("5") + assert entry_rows[0].amount == Decimal("1500") assert entry_rows[0].duration_seconds == 600 stored_resource_group_id = await db_sess.scalar( @@ -251,8 +253,9 @@ async def test_increment_user_buckets_creates_entries( """Verify that incrementing user buckets also writes normalized entries.""" user_uuid = uuid.uuid4() project_id = uuid.uuid4() - raw_slots = ResourceSlot({"cpu": Decimal("3"), "cuda.device": Decimal("2")}) - duration = 300 # 5-minute slice + # 5-minute slice of 3 CPU / 2 cuda.device + resource_usage = ResourceSlot({"cpu": Decimal("900"), "cuda.device": Decimal("600")}) + duration = 300 period = date(2024, 1, 15) result = UsageBucketAggregationResult( @@ -264,7 +267,7 @@ async def test_increment_user_buckets_creates_entries( resource_group="default", resource_group_id=RESOURCE_GROUP_ID, period_date=period, - ): BucketDelta(slots=raw_slots, duration_seconds=duration), + ): BucketDelta(resource_usage=resource_usage, duration_seconds=duration), }, project_usage_deltas={}, domain_usage_deltas={}, @@ -272,7 +275,7 @@ async def test_increment_user_buckets_creates_entries( await db_source.increment_usage_buckets(result) - # Verify entries were created with separated amount/duration + # Verify entries were created with resource-seconds and duration async with db_with_cleanup.begin_readonly_session() as db_sess: entry_rows = ( ( @@ -288,8 +291,8 @@ async def test_increment_user_buckets_creates_entries( assert len(entry_rows) == 2 slot_map = {e.slot_name: e for e in entry_rows} - assert slot_map["cpu"].amount == Decimal("3") - assert slot_map["cuda.device"].amount == Decimal("2") + assert slot_map["cpu"].amount == Decimal("900") + assert slot_map["cuda.device"].amount == Decimal("600") assert slot_map["cpu"].duration_seconds == 300 async def test_aggregated_usage_reads_from_entries( @@ -302,7 +305,7 @@ async def test_aggregated_usage_reads_from_entries( period1 = date(2024, 1, 15) period2 = date(2024, 1, 16) - # Insert two domain buckets with entries (raw amount, not resource-seconds) + # Insert two domain buckets with entries (resource-seconds) result = UsageBucketAggregationResult( user_usage_deltas={}, project_usage_deltas={}, @@ -313,7 +316,7 @@ async def test_aggregated_usage_reads_from_entries( resource_group_id=RESOURCE_GROUP_ID, period_date=period1, ): BucketDelta( - slots=ResourceSlot({"cpu": Decimal("2")}), + resource_usage=ResourceSlot({"cpu": Decimal("600")}), duration_seconds=300, ), DomainUsageBucketKey( @@ -322,7 +325,7 @@ async def test_aggregated_usage_reads_from_entries( resource_group_id=RESOURCE_GROUP_ID, period_date=period2, ): BucketDelta( - slots=ResourceSlot({"cpu": Decimal("3")}), + resource_usage=ResourceSlot({"cpu": Decimal("900")}), duration_seconds=300, ), }, @@ -337,5 +340,5 @@ async def test_aggregated_usage_reads_from_entries( ) assert test_domain_name in aggregated - # 2 + 3 = 5 (raw amounts summed across buckets) - assert aggregated[test_domain_name]["cpu"] == Decimal("5") + # 600 + 900 = 1500 CPU-seconds summed across buckets + assert aggregated[test_domain_name]["cpu"] == Decimal("1500") diff --git a/tests/unit/manager/sokovan/scheduler/fair_share/test_aggregator_bucket_aggregation.py b/tests/unit/manager/sokovan/scheduler/fair_share/test_aggregator_bucket_aggregation.py index 430d2cb1d762..e2b8eb9d9c82 100644 --- a/tests/unit/manager/sokovan/scheduler/fair_share/test_aggregator_bucket_aggregation.py +++ b/tests/unit/manager/sokovan/scheduler/fair_share/test_aggregator_bucket_aggregation.py @@ -3,12 +3,13 @@ Verifies that kernel usage specs are correctly split by day boundaries and aggregated into user/project/domain buckets. -Phase 3 (BA-4308): BucketDelta stores raw amount and duration separately. +BucketDelta stores resource-seconds and the observed duration separately. """ from __future__ import annotations -from datetime import UTC, date, datetime +from dataclasses import dataclass +from datetime import UTC, date, datetime, timedelta from decimal import Decimal from uuid import UUID, uuid4 @@ -30,6 +31,13 @@ RESOURCE_GROUP_ID = ResourceGroupID(uuid4()) +_USER_A = UUID("11111111-1111-4111-8111-111111111111") +_USER_B = UUID("22222222-2222-4222-8222-222222222222") +_USER_C = UUID("33333333-3333-4333-8333-333333333333") +_PROJECT_1 = UUID("aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa") +_PROJECT_2 = UUID("bbbbbbbb-bbbb-4bbb-8bbb-bbbbbbbbbbbb") +_TICK_DAY = date(2026, 7, 17) + def make_datetime( year: int, month: int, day: int, hour: int, minute: int, second: int = 0 @@ -164,7 +172,8 @@ def test_spec_empty_range(self, aggregator: FairShareAggregator) -> None: class TestAggregateKernelUsageToBuckets: """Tests for aggregate_kernel_usage_to_buckets method. - BucketDelta stores raw slots and duration_seconds separately. + BucketDelta stores resource-seconds (amount * duration summed per slice) + and the total observed duration separately. """ def test_single_spec_single_day(self, aggregator: FairShareAggregator) -> None: @@ -200,7 +209,8 @@ def test_single_spec_single_day(self, aggregator: FairShareAggregator) -> None: ) assert user_key in result.user_usage_deltas delta = result.user_usage_deltas[user_key] - assert delta.slots["cpu"] == Decimal("2") + # 2 CPU for 300s -> 600 CPU-seconds + assert delta.resource_usage["cpu"] == Decimal("600") assert delta.duration_seconds == 300 # Project bucket @@ -213,7 +223,7 @@ def test_single_spec_single_day(self, aggregator: FairShareAggregator) -> None: period_date=date(2024, 1, 15), ) assert project_key in result.project_usage_deltas - assert result.project_usage_deltas[project_key].slots["cpu"] == Decimal("2") + assert result.project_usage_deltas[project_key].resource_usage["cpu"] == Decimal("600") # Domain bucket assert len(result.domain_usage_deltas) == 1 @@ -224,7 +234,7 @@ def test_single_spec_single_day(self, aggregator: FairShareAggregator) -> None: period_date=date(2024, 1, 15), ) assert domain_key in result.domain_usage_deltas - assert result.domain_usage_deltas[domain_key].slots["cpu"] == Decimal("2") + assert result.domain_usage_deltas[domain_key].resource_usage["cpu"] == Decimal("600") def test_multiple_specs_same_user_same_day_aggregated( self, aggregator: FairShareAggregator @@ -255,12 +265,11 @@ def test_multiple_specs_same_user_same_day_aggregated( result = aggregator.aggregate_kernel_usage_to_buckets(specs) - # Should have only one user bucket with summed raw slots and duration assert len(result.user_usage_deltas) == 1 user_key = list(result.user_usage_deltas.keys())[0] delta = result.user_usage_deltas[user_key] - # Raw slots accumulate: 2 + 2 = 4 - assert delta.slots["cpu"] == Decimal("4") + # Per slice: 2*300 + 2*300 = 1200, not (2+2) * (300+300) = 2400 + assert delta.resource_usage["cpu"] == Decimal("1200") # Durations accumulate: 300 + 300 = 600 assert delta.duration_seconds == 600 @@ -308,11 +317,11 @@ def test_spec_crossing_midnight_creates_two_buckets( assert day1_key in result.user_usage_deltas assert day2_key in result.user_usage_deltas - # Day 1: 3 minutes = 180s, raw slots = 2 CPU - assert result.user_usage_deltas[day1_key].slots["cpu"] == Decimal("2") + # Day 1: 2 CPU for 180s -> 360 CPU-seconds + assert result.user_usage_deltas[day1_key].resource_usage["cpu"] == Decimal("360") assert result.user_usage_deltas[day1_key].duration_seconds == 180 - # Day 2: 3 minutes = 180s, raw slots = 2 CPU - assert result.user_usage_deltas[day2_key].slots["cpu"] == Decimal("2") + # Day 2: 2 CPU for 180s -> 360 CPU-seconds + assert result.user_usage_deltas[day2_key].resource_usage["cpu"] == Decimal("360") assert result.user_usage_deltas[day2_key].duration_seconds == 180 @@ -436,16 +445,14 @@ def test_backlogged_usage_split_correctly(self, aggregator: FairShareAggregator) period_date=date(2024, 1, 16), ) - # Day 1: 5 specs, raw_slots=2 each, accumulated = 2*5 = 10 - # duration: 240 + 300 + 300 + 300 + 300 = 1440 seconds + # 5 specs at 2 CPU, durations 240 + 300*4 = 1440s d1 = result.user_usage_deltas[day1_key] - assert d1.slots["cpu"] == Decimal("10") + assert d1.resource_usage["cpu"] == Decimal("2880") assert d1.duration_seconds == 1440 - # Day 2: 3 specs, raw_slots=2 each, accumulated = 2*3 = 6 - # duration: 300 + 300 + 180 = 780 seconds + # 3 specs at 2 CPU, durations 300 + 300 + 180 = 780s d2 = result.user_usage_deltas[day2_key] - assert d2.slots["cpu"] == Decimal("6") + assert d2.resource_usage["cpu"] == Decimal("1560") assert d2.duration_seconds == 780 def test_backlogged_usage_with_midnight_crossing_spec( @@ -540,17 +547,14 @@ def test_backlogged_usage_with_midnight_crossing_spec( period_date=date(2024, 1, 16), ) - # Day 1: specs contribute raw_slots=2 each - # 4 complete specs + crossing spec day1 part = 5 contributions = 2*5 = 10 - # duration: 240 + 300 + 300 + 300 + 300(crossing) = 1440s + # 4 complete specs + the crossing spec's day1 part, all at 2 CPU: 1440s d1 = result.user_usage_deltas[day1_key] - assert d1.slots["cpu"] == Decimal("10") + assert d1.resource_usage["cpu"] == Decimal("2880") assert d1.duration_seconds == 1440 - # Day 2: crossing spec day2 part + 2 specs = 3 contributions = 2*3 = 6 - # duration: 300(crossing) + 300 + 180 = 780s + # The crossing spec's day2 part + 2 specs, all at 2 CPU: 780s d2 = result.user_usage_deltas[day2_key] - assert d2.slots["cpu"] == Decimal("6") + assert d2.resource_usage["cpu"] == Decimal("1560") assert d2.duration_seconds == 780 def test_backlogged_multiple_users(self, aggregator: FairShareAggregator) -> None: @@ -597,10 +601,9 @@ def test_backlogged_multiple_users(self, aggregator: FairShareAggregator) -> Non resource_group_id=specs[0].resource_group_id, period_date=date(2024, 1, 15), ) - # User1: raw=2 for 300s, User2: raw=2 for 300s (day1 part) - # Accumulated slots: 2 + 2 = 4 + # User1 and User2 each 2 CPU for 300s: 2*300 + 2*300 = 1200 pd1 = result.project_usage_deltas[project_day1_key] - assert pd1.slots["cpu"] == Decimal("4") + assert pd1.resource_usage["cpu"] == Decimal("1200") assert pd1.duration_seconds == 600 # 300 + 300 project_day2_key = ProjectUsageBucketKey( @@ -610,9 +613,9 @@ def test_backlogged_multiple_users(self, aggregator: FairShareAggregator) -> Non resource_group_id=specs[0].resource_group_id, period_date=date(2024, 1, 16), ) - # User2 only: raw=2 for 300s (day2 part) + # User2 only: 2 CPU for 300s (day2 part) -> 600 CPU-seconds pd2 = result.project_usage_deltas[project_day2_key] - assert pd2.slots["cpu"] == Decimal("2") + assert pd2.resource_usage["cpu"] == Decimal("600") assert pd2.duration_seconds == 300 @@ -649,7 +652,7 @@ def test_spec_exactly_at_midnight_boundary(self, aggregator: FairShareAggregator user_key = list(result.user_usage_deltas.keys())[0] assert user_key.period_date == date(2024, 1, 15) delta = result.user_usage_deltas[user_key] - assert delta.slots["cpu"] == Decimal("2") + assert delta.resource_usage["cpu"] == Decimal("600") assert delta.duration_seconds == 300 def test_spec_starting_exactly_at_midnight(self, aggregator: FairShareAggregator) -> None: @@ -696,8 +699,306 @@ def test_multiple_resource_types(self, aggregator: FairShareAggregator) -> None: assert len(result.user_usage_deltas) == 2 for _key, delta in result.user_usage_deltas.items(): - # Each day gets raw slots + 120 seconds - assert delta.slots["cpu"] == Decimal("2") - assert delta.slots["mem"] == Decimal("4096") - assert delta.slots["cuda.shares"] == Decimal("1") + assert delta.resource_usage["cpu"] == Decimal("240") + assert delta.resource_usage["mem"] == Decimal("491520") + assert delta.resource_usage["cuda.shares"] == Decimal("120") assert delta.duration_seconds == 120 + + +@dataclass(frozen=True) +class _ConcurrentKernelsCase: + """One tick observing ``kernel_count`` kernels of the same user.""" + + kernel_count: int + expected_resource_usage: Decimal + expected_duration_seconds: int + + +class TestConcurrentKernelsNotCrossMultiplied: + """Regression: a bucket must not scale with the square of the kernel count. + + Summing amounts and durations separately and multiplying afterwards + inflates a bucket by exactly the number of kernels folded into it. + """ + + @pytest.mark.parametrize( + "case", + [ + _ConcurrentKernelsCase( + kernel_count=1, + expected_resource_usage=Decimal("300"), + expected_duration_seconds=300, + ), + _ConcurrentKernelsCase( + kernel_count=2, + expected_resource_usage=Decimal("600"), + expected_duration_seconds=600, + ), + _ConcurrentKernelsCase( + kernel_count=4, + expected_resource_usage=Decimal("1200"), + expected_duration_seconds=1200, + ), + _ConcurrentKernelsCase( + kernel_count=10, + expected_resource_usage=Decimal("3000"), + expected_duration_seconds=3000, + ), + ], + ids=lambda case: f"{case.kernel_count}-kernels", + ) + def test_one_tick_of_concurrent_kernels( + self, + aggregator: FairShareAggregator, + case: _ConcurrentKernelsCase, + ) -> None: + """N kernels at 1 fGPU for one 300s slice total N*300 fGPU-seconds.""" + user_uuid = uuid4() + project_id = uuid4() + raw_slots = ResourceSlot({"cuda.shares": Decimal("1")}) + + specs = [ + make_spec( + period_start=make_datetime(2026, 7, 17, 10, 0, 0), + period_end=make_datetime(2026, 7, 17, 10, 5, 0), + resource_usage=ResourceSlot({"cuda.shares": Decimal("300")}), + occupied_slots=raw_slots, + user_uuid=user_uuid, + project_id=project_id, + ) + for _ in range(case.kernel_count) + ] + + result = aggregator.aggregate_kernel_usage_to_buckets(specs) + + user_key = UserUsageBucketKey( + user_uuid=user_uuid, + project_id=project_id, + domain_name="default", + resource_group="default", + resource_group_id=RESOURCE_GROUP_ID, + period_date=date(2026, 7, 17), + ) + delta = result.user_usage_deltas[user_key] + assert delta.resource_usage["cuda.shares"] == case.expected_resource_usage + assert delta.duration_seconds == case.expected_duration_seconds + + domain_key = DomainUsageBucketKey( + domain_name="default", + resource_group="default", + resource_group_id=RESOURCE_GROUP_ID, + period_date=date(2026, 7, 17), + ) + domain_delta = result.domain_usage_deltas[domain_key] + assert domain_delta.resource_usage["cuda.shares"] == case.expected_resource_usage + + def test_full_day_of_four_kernels(self, aggregator: FairShareAggregator) -> None: + """A full day of 4 kernels sums to 4 * 86400, not 4 * 4 * 86400.""" + user_uuid = uuid4() + project_id = uuid4() + raw_slots = ResourceSlot({"cuda.shares": Decimal("1")}) + + specs = [ + make_spec( + period_start=datetime(2026, 7, 17, tzinfo=UTC).replace( + hour=slice_index // 12, minute=(slice_index % 12) * 5 + ), + period_end=datetime(2026, 7, 17, tzinfo=UTC).replace( + hour=slice_index // 12, minute=(slice_index % 12) * 5 + ) + + timedelta(seconds=300), + resource_usage=ResourceSlot({"cuda.shares": Decimal("300")}), + occupied_slots=raw_slots, + user_uuid=user_uuid, + project_id=project_id, + ) + for slice_index in range(288) + for _ in range(4) + ] + + result = aggregator.aggregate_kernel_usage_to_buckets(specs) + + user_key = UserUsageBucketKey( + user_uuid=user_uuid, + project_id=project_id, + domain_name="default", + resource_group="default", + resource_group_id=RESOURCE_GROUP_ID, + period_date=date(2026, 7, 17), + ) + delta = result.user_usage_deltas[user_key] + assert delta.resource_usage["cuda.shares"] == Decimal("345600") + + +@dataclass(frozen=True) +class _Workload: + """Identical kernels one user runs inside one project.""" + + user_uuid: UUID + project_id: UUID + shares: Decimal + kernel_count: int + + +@dataclass(frozen=True) +class _UserBucketExpectation: + label: str + user_uuid: UUID + project_id: UUID + resource_usage: Decimal + duration_seconds: int + + +@dataclass(frozen=True) +class _ProjectBucketExpectation: + label: str + project_id: UUID + resource_usage: Decimal + duration_seconds: int + + +class TestMultiTenantTick: + """One observation tick spanning two projects, three users and seven kernels. + + The inflation compounds up the hierarchy, since each level multiplies by + the kernel count *it* aggregates: 2x for a user here, 3x for a project, + 7x for the domain. User A runs in both projects, so user buckets must + stay keyed by (user, project). + """ + + @pytest.fixture + def multi_tenant_specs(self) -> list[KernelUsageRecordCreatorSpec]: + """One 5-minute slice per kernel, all on the same day.""" + workloads = [ + _Workload(_USER_A, _PROJECT_1, Decimal("1"), 2), + _Workload(_USER_B, _PROJECT_1, Decimal("2"), 1), + _Workload(_USER_A, _PROJECT_2, Decimal("1"), 3), + _Workload(_USER_C, _PROJECT_2, Decimal("4"), 1), + ] + return [ + make_spec( + period_start=make_datetime(2026, 7, 17, 10, 0, 0), + period_end=make_datetime(2026, 7, 17, 10, 5, 0), + resource_usage=ResourceSlot({"cuda.shares": workload.shares * 300}), + occupied_slots=ResourceSlot({"cuda.shares": workload.shares}), + user_uuid=workload.user_uuid, + project_id=workload.project_id, + ) + for workload in workloads + for _ in range(workload.kernel_count) + ] + + @pytest.mark.parametrize( + "case", + [ + _UserBucketExpectation( + label="user-a-project-1", # 2 kernels * 1 share * 300s + user_uuid=_USER_A, + project_id=_PROJECT_1, + resource_usage=Decimal("600"), + duration_seconds=600, + ), + _UserBucketExpectation( + label="user-b-project-1", # 1 kernel * 2 shares * 300s + user_uuid=_USER_B, + project_id=_PROJECT_1, + resource_usage=Decimal("600"), + duration_seconds=300, + ), + _UserBucketExpectation( + label="user-a-project-2", # 3 kernels * 1 share * 300s + user_uuid=_USER_A, + project_id=_PROJECT_2, + resource_usage=Decimal("900"), + duration_seconds=900, + ), + _UserBucketExpectation( + label="user-c-project-2", # 1 kernel * 4 shares * 300s + user_uuid=_USER_C, + project_id=_PROJECT_2, + resource_usage=Decimal("1200"), + duration_seconds=300, + ), + ], + ids=lambda case: case.label, + ) + def test_user_buckets_are_keyed_by_user_and_project( + self, + aggregator: FairShareAggregator, + multi_tenant_specs: list[KernelUsageRecordCreatorSpec], + case: _UserBucketExpectation, + ) -> None: + result = aggregator.aggregate_kernel_usage_to_buckets(multi_tenant_specs) + + assert len(result.user_usage_deltas) == 4 + delta = result.user_usage_deltas[ + UserUsageBucketKey( + user_uuid=case.user_uuid, + project_id=case.project_id, + domain_name="default", + resource_group="default", + resource_group_id=RESOURCE_GROUP_ID, + period_date=_TICK_DAY, + ) + ] + assert delta.resource_usage["cuda.shares"] == case.resource_usage + assert delta.duration_seconds == case.duration_seconds + + @pytest.mark.parametrize( + "case", + [ + _ProjectBucketExpectation( + label="project-1", # user A 600 + user B 600 + project_id=_PROJECT_1, + resource_usage=Decimal("1200"), + duration_seconds=900, + ), + _ProjectBucketExpectation( + label="project-2", # user A 900 + user C 1200 + project_id=_PROJECT_2, + resource_usage=Decimal("2100"), + duration_seconds=1200, + ), + ], + ids=lambda case: case.label, + ) + def test_project_buckets_sum_their_users( + self, + aggregator: FairShareAggregator, + multi_tenant_specs: list[KernelUsageRecordCreatorSpec], + case: _ProjectBucketExpectation, + ) -> None: + result = aggregator.aggregate_kernel_usage_to_buckets(multi_tenant_specs) + + assert len(result.project_usage_deltas) == 2 + delta = result.project_usage_deltas[ + ProjectUsageBucketKey( + project_id=case.project_id, + domain_name="default", + resource_group="default", + resource_group_id=RESOURCE_GROUP_ID, + period_date=_TICK_DAY, + ) + ] + assert delta.resource_usage["cuda.shares"] == case.resource_usage + assert delta.duration_seconds == case.duration_seconds + + def test_domain_bucket_sums_every_project( + self, + aggregator: FairShareAggregator, + multi_tenant_specs: list[KernelUsageRecordCreatorSpec], + ) -> None: + """Project 1 (1200) + project 2 (2100).""" + result = aggregator.aggregate_kernel_usage_to_buckets(multi_tenant_specs) + + assert len(result.domain_usage_deltas) == 1 + delta = result.domain_usage_deltas[ + DomainUsageBucketKey( + domain_name="default", + resource_group="default", + resource_group_id=RESOURCE_GROUP_ID, + period_date=_TICK_DAY, + ) + ] + assert delta.resource_usage["cuda.shares"] == Decimal("3300") + assert delta.duration_seconds == 2100 From 8f0ed08407a380725001173a2e0ef838be81142d Mon Sep 17 00:00:00 2001 From: Gyubong Date: Mon, 20 Jul 2026 13:58:42 +0900 Subject: [PATCH 2/2] refactor(BA-6927): rename usage_bucket_entries.amount to resource_usage The column holds resource-seconds, not a resource amount. Under the old name `amount * duration_seconds` reads like the quantity a caller wants, when both columns are already sums and their product is a cross product -- the same shape as the bug fixed in the parent commit. resource_usage is what kernel_usage_records and the three parent bucket tables already call this quantity, so the entries table now speaks the same vocabulary rather than introducing a second name for one concept. Mechanical rename across the ORM model, the two repository read paths and the tests, plus the alembic column rename. No behaviour change. Co-Authored-By: Claude Opus 4.8 (1M context) --- changes/12968.misc.md | 1 + ...5a1f26_rename_usage_bucket_entry_amount.py | 45 +++++++++++++++++++ .../models/resource_usage_history/row.py | 11 ++--- .../fair_share/db_source/db_source.py | 12 ++--- .../db_source/db_source.py | 10 ++--- .../test_resource_usage_history_repository.py | 6 +-- .../test_usage_bucket_entries.py | 12 ++--- .../retention/test_retention_repository.py | 2 +- 8 files changed, 73 insertions(+), 26 deletions(-) create mode 100644 changes/12968.misc.md create mode 100644 src/ai/backend/manager/models/alembic/versions/e7d3b95a1f26_rename_usage_bucket_entry_amount.py diff --git a/changes/12968.misc.md b/changes/12968.misc.md new file mode 100644 index 000000000000..e58c5be72cc4 --- /dev/null +++ b/changes/12968.misc.md @@ -0,0 +1 @@ +Rename the internal `usage_bucket_entries.amount` column to `resource_usage`, matching the name the parent bucket tables already use for the same quantity. diff --git a/src/ai/backend/manager/models/alembic/versions/e7d3b95a1f26_rename_usage_bucket_entry_amount.py b/src/ai/backend/manager/models/alembic/versions/e7d3b95a1f26_rename_usage_bucket_entry_amount.py new file mode 100644 index 000000000000..f32c016c1004 --- /dev/null +++ b/src/ai/backend/manager/models/alembic/versions/e7d3b95a1f26_rename_usage_bucket_entry_amount.py @@ -0,0 +1,45 @@ +"""rename usage_bucket_entries.amount to resource_usage + +The column holds resource-seconds, not a resource amount. Under the old name +``amount * duration_seconds`` reads like the quantity a caller wants, when both +columns are already sums and their product is a cross product. + +``resource_usage`` is the name the three parent bucket tables and +``kernel_usage_records`` already use for this quantity, so the table now speaks +the same vocabulary as the rest of the schema. + +Revision ID: e7d3b95a1f26 +Revises: c4a91d7e05b2 +Create Date: 2026-07-20 00:00:00.000000 + +""" + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "e7d3b95a1f26" +down_revision = "c4a91d7e05b2" +# Part of: NEXT_RELEASE_VERSION +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.alter_column( + "usage_bucket_entries", + "amount", + new_column_name="resource_usage", + existing_type=sa.Numeric(precision=32, scale=6), + existing_nullable=False, + ) + + +def downgrade() -> None: + op.alter_column( + "usage_bucket_entries", + "resource_usage", + new_column_name="amount", + existing_type=sa.Numeric(precision=32, scale=6), + existing_nullable=False, + ) diff --git a/src/ai/backend/manager/models/resource_usage_history/row.py b/src/ai/backend/manager/models/resource_usage_history/row.py index 32ca808ba8a8..d77e3406fa24 100644 --- a/src/ai/backend/manager/models/resource_usage_history/row.py +++ b/src/ai/backend/manager/models/resource_usage_history/row.py @@ -442,9 +442,10 @@ class UserUsageBucketRow(LifecycleTimestampsMixin, Base): # type: ignore[misc] class UsageBucketEntryRow(Base): # type: ignore[misc] """Per-slot normalized entry for usage bucket aggregation (Phase 3). - ``amount`` holds ``sum(amount_k * duration_k)`` resource-seconds and is what - the read paths sum. ``duration_seconds`` is for reporting only; both columns - are sums, so they are never multiplied together. + ``resource_usage`` holds ``sum(amount_k * duration_k)`` resource-seconds and + is what the read paths sum, matching the column of the same name on the three + parent bucket tables. ``duration_seconds`` is for reporting only; both + columns are sums, so they are never multiplied together. One entry per (bucket_id, slot_name). ``bucket_type`` is a discriminator indicating which parent table (domain/project/user_usage_buckets) owns @@ -456,8 +457,8 @@ class UsageBucketEntryRow(Base): # type: ignore[misc] bucket_id: Mapped[uuid.UUID] = mapped_column("bucket_id", GUID(), nullable=False) bucket_type: Mapped[str] = mapped_column("bucket_type", sa.String(length=16), nullable=False) slot_name: Mapped[str] = mapped_column("slot_name", sa.String(length=64), nullable=False) - amount: Mapped[Decimal] = mapped_column( - "amount", sa.Numeric(precision=32, scale=6), nullable=False + resource_usage: Mapped[Decimal] = mapped_column( + "resource_usage", sa.Numeric(precision=32, scale=6), nullable=False ) duration_seconds: Mapped[int] = mapped_column("duration_seconds", sa.Integer(), nullable=False) capacity: Mapped[Decimal] = mapped_column( diff --git a/src/ai/backend/manager/repositories/fair_share/db_source/db_source.py b/src/ai/backend/manager/repositories/fair_share/db_source/db_source.py index 5f2cebb08821..5914de10ad4b 100644 --- a/src/ai/backend/manager/repositories/fair_share/db_source/db_source.py +++ b/src/ai/backend/manager/repositories/fair_share/db_source/db_source.py @@ -1585,7 +1585,7 @@ async def _fetch_raw_usage_buckets( UserUsageBucketRow.project_id, UserUsageBucketRow.period_start, ube.c.slot_name, - ube.c.amount, + ube.c.resource_usage, ) .select_from( sa.join( @@ -1612,7 +1612,7 @@ async def _fetch_raw_usage_buckets( ProjectUsageBucketRow.project_id, ProjectUsageBucketRow.period_start, ube.c.slot_name, - ube.c.amount, + ube.c.resource_usage, ) .select_from( sa.join( @@ -1639,7 +1639,7 @@ async def _fetch_raw_usage_buckets( DomainUsageBucketRow.domain_name, DomainUsageBucketRow.period_start, ube.c.slot_name, - ube.c.amount, + ube.c.resource_usage, ) .select_from( sa.join( @@ -1668,7 +1668,7 @@ async def _fetch_raw_usage_buckets( user_buckets[key] = {} if row.period_start not in user_buckets[key]: user_buckets[key][row.period_start] = ResourceSlot() - user_buckets[key][row.period_start][row.slot_name] = Decimal(str(row.amount)) + user_buckets[key][row.period_start][row.slot_name] = Decimal(str(row.resource_usage)) project_buckets: dict[uuid.UUID, dict[date, ResourceSlot]] = {} for row in project_rows: @@ -1677,7 +1677,7 @@ async def _fetch_raw_usage_buckets( if row.period_start not in project_buckets[row.project_id]: project_buckets[row.project_id][row.period_start] = ResourceSlot() project_buckets[row.project_id][row.period_start][row.slot_name] = Decimal( - str(row.amount) + str(row.resource_usage) ) domain_buckets: dict[str, dict[date, ResourceSlot]] = {} @@ -1687,7 +1687,7 @@ async def _fetch_raw_usage_buckets( if row.period_start not in domain_buckets[row.domain_name]: domain_buckets[row.domain_name][row.period_start] = ResourceSlot() domain_buckets[row.domain_name][row.period_start][row.slot_name] = Decimal( - str(row.amount) + str(row.resource_usage) ) return RawUsageBucketsByLevel( diff --git a/src/ai/backend/manager/repositories/resource_usage_history/db_source/db_source.py b/src/ai/backend/manager/repositories/resource_usage_history/db_source/db_source.py index f5a23d0e3380..e209af7b801c 100644 --- a/src/ai/backend/manager/repositories/resource_usage_history/db_source/db_source.py +++ b/src/ai/backend/manager/repositories/resource_usage_history/db_source/db_source.py @@ -398,7 +398,7 @@ async def _fetch_aggregated_usage_by_user( UserUsageBucketRow.user_uuid, UserUsageBucketRow.project_id, ube.c.slot_name, - sa.func.sum(ube.c.amount).label("total_resource_usage"), + sa.func.sum(ube.c.resource_usage).label("total_resource_usage"), ) .select_from( sa.join( @@ -448,7 +448,7 @@ async def get_aggregated_usage_by_project( sa.select( ProjectUsageBucketRow.project_id, ube.c.slot_name, - sa.func.sum(ube.c.amount).label("total_resource_usage"), + sa.func.sum(ube.c.resource_usage).label("total_resource_usage"), ) .select_from( sa.join( @@ -496,7 +496,7 @@ async def get_aggregated_usage_by_domain( sa.select( DomainUsageBucketRow.domain_name, ube.c.slot_name, - sa.func.sum(ube.c.amount).label("total_resource_usage"), + sa.func.sum(ube.c.resource_usage).label("total_resource_usage"), ) .select_from( sa.join( @@ -854,14 +854,14 @@ async def _upsert_bucket_entries( bucket_id=bucket_id, bucket_type=bucket_type, slot_name=slot_name, - amount=value, + resource_usage=value, duration_seconds=bucket_delta.duration_seconds, capacity=0, ) .on_conflict_do_update( constraint="pk_usage_bucket_entries", set_={ - "amount": entry_table.c.amount + value, + "resource_usage": entry_table.c.resource_usage + value, "duration_seconds": ( entry_table.c.duration_seconds + bucket_delta.duration_seconds ), diff --git a/tests/unit/manager/repositories/resource_usage_history/test_resource_usage_history_repository.py b/tests/unit/manager/repositories/resource_usage_history/test_resource_usage_history_repository.py index cee87b32bf86..69b3ba8ec2a6 100644 --- a/tests/unit/manager/repositories/resource_usage_history/test_resource_usage_history_repository.py +++ b/tests/unit/manager/repositories/resource_usage_history/test_resource_usage_history_repository.py @@ -660,7 +660,7 @@ async def test_get_aggregated_usage_by_user( bucket_id=result.id, bucket_type="user", slot_name="cpu", - amount=Decimal("3600"), + resource_usage=Decimal("3600"), duration_seconds=300, capacity=Decimal("0"), ) @@ -732,7 +732,7 @@ async def test_get_aggregated_usage_by_project( bucket_id=result.id, bucket_type="project", slot_name="cpu", - amount=Decimal("7200"), + resource_usage=Decimal("7200"), duration_seconds=300, capacity=Decimal("0"), ) @@ -783,7 +783,7 @@ async def test_get_aggregated_usage_by_domain( bucket_id=result.id, bucket_type="domain", slot_name="cpu", - amount=Decimal("86400"), + resource_usage=Decimal("86400"), duration_seconds=300, capacity=Decimal("0"), ) diff --git a/tests/unit/manager/repositories/resource_usage_history/test_usage_bucket_entries.py b/tests/unit/manager/repositories/resource_usage_history/test_usage_bucket_entries.py index be0bea122bd1..0f414b9a8f03 100644 --- a/tests/unit/manager/repositories/resource_usage_history/test_usage_bucket_entries.py +++ b/tests/unit/manager/repositories/resource_usage_history/test_usage_bucket_entries.py @@ -163,8 +163,8 @@ async def test_increment_domain_buckets_creates_entries( slot_map = {e.slot_name: e for e in entry_rows} assert "cpu" in slot_map assert "mem" in slot_map - assert slot_map["cpu"].amount == Decimal("600") - assert slot_map["mem"].amount == Decimal("1228800000") + assert slot_map["cpu"].resource_usage == Decimal("600") + assert slot_map["mem"].resource_usage == Decimal("1228800000") assert slot_map["cpu"].duration_seconds == 300 assert slot_map["mem"].duration_seconds == 300 @@ -232,7 +232,7 @@ async def test_increment_accumulates_entries( assert len(entry_rows) == 1 assert entry_rows[0].slot_name == "cpu" - assert entry_rows[0].amount == Decimal("1500") + assert entry_rows[0].resource_usage == Decimal("1500") assert entry_rows[0].duration_seconds == 600 stored_resource_group_id = await db_sess.scalar( @@ -291,8 +291,8 @@ async def test_increment_user_buckets_creates_entries( assert len(entry_rows) == 2 slot_map = {e.slot_name: e for e in entry_rows} - assert slot_map["cpu"].amount == Decimal("900") - assert slot_map["cuda.device"].amount == Decimal("600") + assert slot_map["cpu"].resource_usage == Decimal("900") + assert slot_map["cuda.device"].resource_usage == Decimal("600") assert slot_map["cpu"].duration_seconds == 300 async def test_aggregated_usage_reads_from_entries( @@ -332,7 +332,7 @@ async def test_aggregated_usage_reads_from_entries( ) await db_source.increment_usage_buckets(result) - # Query aggregated usage — SUM(amount) across buckets + # Query aggregated usage — SUM(resource_usage) across buckets aggregated = await db_source.get_aggregated_usage_by_domain( resource_group="default", lookback_start=date(2024, 1, 14), diff --git a/tests/unit/manager/repositories/retention/test_retention_repository.py b/tests/unit/manager/repositories/retention/test_retention_repository.py index e4d3bb7889ea..eb029f2664ae 100644 --- a/tests/unit/manager/repositories/retention/test_retention_repository.py +++ b/tests/unit/manager/repositories/retention/test_retention_repository.py @@ -1332,7 +1332,7 @@ async def _add_entry( bucket_id=bucket_id, bucket_type=bucket_type, slot_name=slot, - amount=Decimal("1"), + resource_usage=Decimal("1"), duration_seconds=1, capacity=Decimal("1"), )