diff --git a/changes/12965.fix.md b/changes/12965.fix.md new file mode 100644 index 00000000000..4e1835190a8 --- /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/__init__.py b/src/ai/backend/manager/data/fair_share/__init__.py index d4b5d24a63f..246c8c61b90 100644 --- a/src/ai/backend/manager/data/fair_share/__init__.py +++ b/src/ai/backend/manager/data/fair_share/__init__.py @@ -1,5 +1,4 @@ from .types import ( - BucketDelta, DomainFactorResult, DomainFairShareData, DomainFairShareSearchResult, @@ -59,6 +58,5 @@ "UserUsageBucketKey", "ProjectUsageBucketKey", "DomainUsageBucketKey", - "BucketDelta", "UsageBucketAggregationResult", ) diff --git a/src/ai/backend/manager/data/fair_share/types.py b/src/ai/backend/manager/data/fair_share/types.py index 0afb56717db..e472fe7893c 100644 --- a/src/ai/backend/manager/data/fair_share/types.py +++ b/src/ai/backend/manager/data/fair_share/types.py @@ -419,23 +419,15 @@ class DomainUsageBucketKey: @dataclass -class BucketDelta: - """Separated resource amount and duration for a usage bucket. +class UsageBucketAggregationResult: + """Resource-seconds to add to each bucket, from one observation tick. - 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. + Each value is ``sum(amount_k * duration_k)`` over the slices folded into that + bucket. It must be accumulated as a sum of per-slice products: summing the + amounts and the durations separately and multiplying afterwards gives a cross + product inflated by the number of slices. """ - slots: ResourceSlot = field(default_factory=ResourceSlot) - duration_seconds: int = 0 - - -@dataclass -class UsageBucketAggregationResult: - """Result of aggregating kernel usage into daily buckets.""" - - user_usage_deltas: dict[UserUsageBucketKey, BucketDelta] = field(default_factory=dict) - project_usage_deltas: dict[ProjectUsageBucketKey, BucketDelta] = field(default_factory=dict) - domain_usage_deltas: dict[DomainUsageBucketKey, BucketDelta] = field(default_factory=dict) + user_usage_deltas: dict[UserUsageBucketKey, ResourceSlot] = field(default_factory=dict) + project_usage_deltas: dict[ProjectUsageBucketKey, ResourceSlot] = field(default_factory=dict) + domain_usage_deltas: dict[DomainUsageBucketKey, ResourceSlot] = field(default_factory=dict) 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 00000000000..edfdaddb04b --- /dev/null +++ b/src/ai/backend/manager/models/alembic/versions/c4a91d7e05b2_rebuild_inflated_usage_buckets.py @@ -0,0 +1,304 @@ +"""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`` becomes ``resource_usage``, matching what the +JSONB mirror on ``kernel_usage_records`` and the three bucket tables already call +this quantity, and drops its precision limit. It now holds the product directly +rather than a factor readers had to multiply back out: a domain-level daily mem +bucket runs past any fixed precision, and unconstrained NUMERIC has no ceiling. + +``duration_seconds`` goes with it. It only ever existed to reconstitute that +product, no reader consulted it on its own, and it counted kernel-seconds rather +than wall-clock, so leaving it would leave a column that invites the same +misreading the product form did. + +Revision ID: c4a91d7e05b2 +Revises: 5405ee0d8eed +Create Date: 2026-07-20 00:00:00.000000 + +""" + +import logging +from datetime import date, timedelta + +import sqlalchemy as sa +from alembic import op + +# revision identifiers, used by Alembic. +revision = "c4a91d7e05b2" +down_revision = "5405ee0d8eed" +# Part of: NEXT_RELEASE_VERSION +branch_labels = None +depends_on = None + +log = logging.getLogger("alembic.runtime.migration") + + +def upgrade() -> None: + # Schema change: amount -> resource_usage (now the product, so no fixed + # precision), and duration_seconds is no longer needed. + op.alter_column( + "usage_bucket_entries", + "amount", + new_column_name="resource_usage", + existing_type=sa.Numeric(precision=24, scale=6), + type_=sa.Numeric(), + existing_nullable=False, + ) + op.drop_column("usage_bucket_entries", "duration_seconds") + + # Data change: rebuild the corrupted values from kernel_usage_records. + # Delete every corrupted entry first (all three levels at once), then rebuild + # each level. The three rebuilds share the same shape and differ only in which + # table and key columns they use, so they are written out one by one below. + conn = op.get_bind() + window = _rebuildable_date_range(conn) + if window is None: + # No usage records to rebuild from (fresh install, or everything purged). + return + rebuild_from, rebuild_to = window + _purge_corrupted_usage(conn, rebuild_from, rebuild_to) + _rebuild_user_buckets(conn, rebuild_from, rebuild_to) + _rebuild_project_buckets(conn, rebuild_from, rebuild_to) + _rebuild_domain_buckets(conn, rebuild_from, rebuild_to) + + +def downgrade() -> None: + op.add_column( + "usage_bucket_entries", + sa.Column("duration_seconds", sa.Integer(), nullable=False, server_default="0"), + ) + # The rebuilt values are the correct ones and the inflated originals cannot + # be reconstructed, so only the column definition is reverted. Values that + # exceed the restored precision will fail the cast, which is the honest + # outcome: they do not fit the old column. + op.alter_column( + "usage_bucket_entries", + "resource_usage", + new_column_name="amount", + existing_type=sa.Numeric(), + type_=sa.Numeric(precision=24, scale=6), + existing_nullable=False, + ) + log.warning( + "usage_bucket_entries is left corrupt by this downgrade: the column restored to " + "'amount' now holds resource-seconds products, not raw amounts, and " + "duration_seconds is reset to 0. Re-apply revision %s to rebuild the correct " + "values from kernel_usage_records.", + revision, + ) + + +def _rebuildable_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 _purge_corrupted_usage( + conn: sa.engine.Connection, rebuild_from: date, rebuild_to: date +) -> None: + """Delete every corrupted entry the rebuild will replace, across all levels. + + Each entry belongs to exactly one bucket via ``bucket_id``, so listing the + in-window buckets of the three tables and deleting entries that point at them + clears user, project and domain in a single statement. The rebuilds then + start from a clean slate and only need to insert. + """ + conn.execute( + sa.text( + """ + DELETE FROM usage_bucket_entries + WHERE (bucket_id, bucket_type) IN ( + SELECT id, 'user' FROM user_usage_buckets + WHERE period_start BETWEEN :rebuild_from AND :rebuild_to + UNION ALL + SELECT id, 'project' FROM project_usage_buckets + WHERE period_start BETWEEN :rebuild_from AND :rebuild_to + UNION ALL + SELECT id, 'domain' FROM domain_usage_buckets + WHERE period_start BETWEEN :rebuild_from AND :rebuild_to + ) + """ + ), + {"rebuild_from": rebuild_from, "rebuild_to": rebuild_to}, + ) + + +def _rebuild_user_buckets(conn: sa.engine.Connection, rebuild_from: date, rebuild_to: date) -> None: + """Recompute user bucket entries and their JSONB mirror from kernel records.""" + params = {"rebuild_from": rebuild_from, "rebuild_to": rebuild_to} + + # 1. Insert one entry per (bucket, slot), summing the per-slice resource-seconds + # that kernel_usage_records already stores correctly. capacity is refilled by + # the next observation tick, so 0 here is safe. + conn.execute( + sa.text( + """ + INSERT INTO usage_bucket_entries + (bucket_id, bucket_type, slot_name, resource_usage, capacity) + SELECT user_usage_buckets.id, 'user', slot.key, SUM(slot.value::numeric), 0 + FROM user_usage_buckets + JOIN kernel_usage_records + ON kernel_usage_records.user_uuid = user_usage_buckets.user_uuid + AND kernel_usage_records.project_id = user_usage_buckets.project_id + AND kernel_usage_records.resource_group_id = user_usage_buckets.resource_group_id + AND (kernel_usage_records.period_start AT TIME ZONE 'UTC')::date + = user_usage_buckets.period_start + CROSS JOIN LATERAL jsonb_each_text(kernel_usage_records.resource_usage) AS slot + WHERE user_usage_buckets.period_start BETWEEN :rebuild_from AND :rebuild_to + GROUP BY user_usage_buckets.id, slot.key + """ + ), + params, + ) + # 2. The JSONB mirror is just the slot map of the bucket's entries, or {} when + # the bucket has no kernel records left to rebuild from. + conn.execute( + sa.text( + """ + UPDATE user_usage_buckets + SET resource_usage = COALESCE( + ( + SELECT jsonb_object_agg( + usage_bucket_entries.slot_name, + usage_bucket_entries.resource_usage + ) + FROM usage_bucket_entries + WHERE usage_bucket_entries.bucket_id = user_usage_buckets.id + AND usage_bucket_entries.bucket_type = 'user' + ), + '{}'::jsonb + ) + WHERE user_usage_buckets.period_start BETWEEN :rebuild_from AND :rebuild_to + """ + ), + params, + ) + + +def _rebuild_project_buckets( + conn: sa.engine.Connection, rebuild_from: date, rebuild_to: date +) -> None: + """Recompute project bucket entries and their JSONB mirror from kernel records.""" + params = {"rebuild_from": rebuild_from, "rebuild_to": rebuild_to} + + # 1. Insert one entry per (bucket, slot), summing the per-slice resource-seconds + # that kernel_usage_records already stores correctly. capacity is refilled by + # the next observation tick, so 0 here is safe. + conn.execute( + sa.text( + """ + INSERT INTO usage_bucket_entries + (bucket_id, bucket_type, slot_name, resource_usage, capacity) + SELECT project_usage_buckets.id, 'project', slot.key, SUM(slot.value::numeric), 0 + FROM project_usage_buckets + JOIN kernel_usage_records + ON kernel_usage_records.project_id = project_usage_buckets.project_id + AND kernel_usage_records.resource_group_id = project_usage_buckets.resource_group_id + AND (kernel_usage_records.period_start AT TIME ZONE 'UTC')::date + = project_usage_buckets.period_start + CROSS JOIN LATERAL jsonb_each_text(kernel_usage_records.resource_usage) AS slot + WHERE project_usage_buckets.period_start BETWEEN :rebuild_from AND :rebuild_to + GROUP BY project_usage_buckets.id, slot.key + """ + ), + params, + ) + # 2. The JSONB mirror is just the slot map of the bucket's entries, or {} when + # the bucket has no kernel records left to rebuild from. + conn.execute( + sa.text( + """ + UPDATE project_usage_buckets + SET resource_usage = COALESCE( + ( + SELECT jsonb_object_agg( + usage_bucket_entries.slot_name, + usage_bucket_entries.resource_usage + ) + FROM usage_bucket_entries + WHERE usage_bucket_entries.bucket_id = project_usage_buckets.id + AND usage_bucket_entries.bucket_type = 'project' + ), + '{}'::jsonb + ) + WHERE project_usage_buckets.period_start BETWEEN :rebuild_from AND :rebuild_to + """ + ), + params, + ) + + +def _rebuild_domain_buckets( + conn: sa.engine.Connection, rebuild_from: date, rebuild_to: date +) -> None: + """Recompute domain bucket entries and their JSONB mirror from kernel records.""" + params = {"rebuild_from": rebuild_from, "rebuild_to": rebuild_to} + + # 1. Insert one entry per (bucket, slot), summing the per-slice resource-seconds + # that kernel_usage_records already stores correctly. capacity is refilled by + # the next observation tick, so 0 here is safe. + conn.execute( + sa.text( + """ + INSERT INTO usage_bucket_entries + (bucket_id, bucket_type, slot_name, resource_usage, capacity) + SELECT domain_usage_buckets.id, 'domain', slot.key, SUM(slot.value::numeric), 0 + FROM domain_usage_buckets + JOIN kernel_usage_records + ON kernel_usage_records.domain_name = domain_usage_buckets.domain_name + AND kernel_usage_records.resource_group_id = domain_usage_buckets.resource_group_id + AND (kernel_usage_records.period_start AT TIME ZONE 'UTC')::date + = domain_usage_buckets.period_start + CROSS JOIN LATERAL jsonb_each_text(kernel_usage_records.resource_usage) AS slot + WHERE domain_usage_buckets.period_start BETWEEN :rebuild_from AND :rebuild_to + GROUP BY domain_usage_buckets.id, slot.key + """ + ), + params, + ) + # 2. The JSONB mirror is just the slot map of the bucket's entries, or {} when + # the bucket has no kernel records left to rebuild from. + conn.execute( + sa.text( + """ + UPDATE domain_usage_buckets + SET resource_usage = COALESCE( + ( + SELECT jsonb_object_agg( + usage_bucket_entries.slot_name, + usage_bucket_entries.resource_usage + ) + FROM usage_bucket_entries + WHERE usage_bucket_entries.bucket_id = domain_usage_buckets.id + AND usage_bucket_entries.bucket_type = 'domain' + ), + '{}'::jsonb + ) + WHERE domain_usage_buckets.period_start BETWEEN :rebuild_from AND :rebuild_to + """ + ), + params, + ) 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 2227aaf06dc..d161aceba6f 100644 --- a/src/ai/backend/manager/models/resource_usage_history/row.py +++ b/src/ai/backend/manager/models/resource_usage_history/row.py @@ -443,10 +443,15 @@ 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. + ``resource_usage`` accumulates a kernel's ``occupied_slots`` integrated + over the time it held them. It records an allocation, not a measurement: a + kernel holding a GPU idle for an hour counts the same as one saturating it. + That is what fair share wants -- holding a resource denies it to others -- but + it means "usage" would overstate what this column knows. + + Declared as unconstrained NUMERIC on purpose: a domain-level daily mem bucket + runs to ~1e18 byte-seconds on a large cluster, past any fixed precision worth + writing down, and PostgreSQL's unconstrained numeric has no such ceiling. One entry per (bucket_id, slot_name). ``bucket_type`` is a discriminator indicating which parent table (domain/project/user_usage_buckets) owns @@ -458,10 +463,7 @@ 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=24, scale=6), nullable=False - ) - duration_seconds: Mapped[int] = mapped_column("duration_seconds", sa.Integer(), nullable=False) + resource_usage: Mapped[Decimal] = mapped_column("resource_usage", sa.Numeric(), nullable=False) capacity: Mapped[Decimal] = mapped_column( "capacity", sa.Numeric(precision=24, scale=6), nullable=False ) 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 09b85a39d3f..b76dca598aa 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 @@ -1659,7 +1659,7 @@ async def _fetch_raw_usage_buckets( UserUsageBucketRow.project_id, UserUsageBucketRow.period_start, ube.c.slot_name, - ube.c.amount, + ube.c.resource_usage.label("resource_usage"), ) .select_from( sa.join( @@ -1686,7 +1686,7 @@ async def _fetch_raw_usage_buckets( ProjectUsageBucketRow.project_id, ProjectUsageBucketRow.period_start, ube.c.slot_name, - ube.c.amount, + ube.c.resource_usage.label("resource_usage"), ) .select_from( sa.join( @@ -1713,7 +1713,7 @@ async def _fetch_raw_usage_buckets( DomainUsageBucketRow.domain_name, DomainUsageBucketRow.period_start, ube.c.slot_name, - ube.c.amount, + ube.c.resource_usage.label("resource_usage"), ) .select_from( sa.join( @@ -1742,7 +1742,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: @@ -1751,7 +1751,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]] = {} @@ -1761,7 +1761,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 ac5f361ccfb..23781772609 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 @@ -52,7 +52,6 @@ from sqlalchemy.ext.asyncio import AsyncSession as SASession from ai.backend.manager.data.fair_share import ( - BucketDelta, DomainUsageBucketKey, ProjectUsageBucketKey, UsageBucketAggregationResult, @@ -399,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.resource_usage).label("total_resource_usage"), ) .select_from( sa.join( @@ -430,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( @@ -449,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.resource_usage).label("total_resource_usage"), ) .select_from( sa.join( @@ -478,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( @@ -497,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.resource_usage).label("total_resource_usage"), ) .select_from( sa.join( @@ -526,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 ==================== @@ -574,7 +573,7 @@ async def increment_usage_buckets( async def _increment_user_usage_buckets( self, db_sess: SASession, - deltas: Mapping[UserUsageBucketKey, BucketDelta], + deltas: Mapping[UserUsageBucketKey, ResourceSlot], decay_unit_days: int, ) -> None: """Increment user usage buckets with deltas.""" @@ -585,14 +584,10 @@ async def _increment_user_usage_buckets( keys_list = list(deltas.keys()) existing = await self._fetch_existing_user_buckets(db_sess, keys_list) - for key, bucket_delta in deltas.items(): + for key, delta_usage in deltas.items(): lookup_key = (key.user_uuid, key.project_id, key.resource_group_id, 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 + delta_usage # Upsert with merged usage (JSONB) stmt = ( @@ -627,7 +622,7 @@ async def _increment_user_usage_buckets( db_sess, bucket_id, "user", - bucket_delta, + delta_usage, ) async def _fetch_existing_user_buckets( @@ -672,7 +667,7 @@ async def _fetch_existing_user_buckets( async def _increment_project_usage_buckets( self, db_sess: SASession, - deltas: Mapping[ProjectUsageBucketKey, BucketDelta], + deltas: Mapping[ProjectUsageBucketKey, ResourceSlot], decay_unit_days: int, ) -> None: """Increment project usage buckets with deltas.""" @@ -683,13 +678,10 @@ async def _increment_project_usage_buckets( keys_list = list(deltas.keys()) existing = await self._fetch_existing_project_buckets(db_sess, keys_list) - for key, bucket_delta in deltas.items(): + for key, delta_usage in deltas.items(): lookup_key = (key.project_id, key.resource_group_id, 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 + delta_usage # Upsert with merged usage (JSONB) stmt = ( @@ -718,7 +710,7 @@ async def _increment_project_usage_buckets( db_sess, bucket_id, "project", - bucket_delta, + delta_usage, ) async def _fetch_existing_project_buckets( @@ -755,7 +747,7 @@ async def _fetch_existing_project_buckets( async def _increment_domain_usage_buckets( self, db_sess: SASession, - deltas: Mapping[DomainUsageBucketKey, BucketDelta], + deltas: Mapping[DomainUsageBucketKey, ResourceSlot], decay_unit_days: int, ) -> None: """Increment domain usage buckets with deltas.""" @@ -766,13 +758,10 @@ async def _increment_domain_usage_buckets( keys_list = list(deltas.keys()) existing = await self._fetch_existing_domain_buckets(db_sess, keys_list) - for key, bucket_delta in deltas.items(): + for key, delta_usage in deltas.items(): lookup_key = (key.domain_name, key.resource_group_id, 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 + delta_usage # Upsert with merged usage (JSONB) stmt = ( @@ -800,7 +789,7 @@ async def _increment_domain_usage_buckets( db_sess, bucket_id, "domain", - bucket_delta, + delta_usage, ) async def _fetch_existing_domain_buckets( @@ -841,39 +830,30 @@ async def _upsert_bucket_entries( db_sess: SASession, bucket_id: uuid.UUID, bucket_type: str, - bucket_delta: BucketDelta, + usage: ResourceSlot, ) -> 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. + ``resource_usage`` accumulates server-side. ``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, delta_slot_seconds in usage.items(): stmt = ( pg_insert(entry_table) .values( bucket_id=bucket_id, bucket_type=bucket_type, slot_name=slot_name, - amount=value, - duration_seconds=bucket_delta.duration_seconds, + resource_usage=delta_slot_seconds, capacity=0, ) .on_conflict_do_update( constraint="pk_usage_bucket_entries", set_={ - "amount": entry_table.c.amount + value, - "duration_seconds": ( - entry_table.c.duration_seconds + bucket_delta.duration_seconds - ), + "resource_usage": (entry_table.c.resource_usage + delta_slot_seconds), }, ) ) @@ -924,15 +904,3 @@ async def update_bucket_entry_capacities( .values(capacity=capacity) ) await db_sess.execute(stmt) - - @staticmethod - def _calculate_resource_seconds( - slots: ResourceSlot, - seconds: int, - ) -> ResourceSlot: - """Convert resource slots to resource-seconds for legacy JSONB storage. - - Multiplies each resource value by the number of seconds to get - the total resource-seconds consumed during the period. - """ - return ResourceSlot({key: value * Decimal(str(seconds)) for key, value in slots.items()}) 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 54fc9926c7b..efbeadb97de 100644 --- a/src/ai/backend/manager/sokovan/scheduler/fair_share/aggregator.py +++ b/src/ai/backend/manager/sokovan/scheduler/fair_share/aggregator.py @@ -24,7 +24,6 @@ from ai.backend.common.types import ResourceSlot from ai.backend.logging.utils import BraceStyleAdapter from ai.backend.manager.data.fair_share import ( - BucketDelta, DomainUsageBucketKey, ProjectUsageBucketKey, UsageBucketAggregationResult, @@ -133,9 +132,9 @@ def aggregate_kernel_usage_to_buckets( Returns: UsageBucketAggregationResult with deltas for each bucket """ - user_deltas: dict[UserUsageBucketKey, BucketDelta] = defaultdict(BucketDelta) - project_deltas: dict[ProjectUsageBucketKey, BucketDelta] = defaultdict(BucketDelta) - domain_deltas: dict[DomainUsageBucketKey, BucketDelta] = defaultdict(BucketDelta) + user_deltas: dict[UserUsageBucketKey, ResourceSlot] = defaultdict(ResourceSlot) + project_deltas: dict[ProjectUsageBucketKey, ResourceSlot] = defaultdict(ResourceSlot) + domain_deltas: dict[DomainUsageBucketKey, ResourceSlot] = defaultdict(ResourceSlot) for spec in specs: # Split spec across day boundaries and aggregate @@ -218,16 +217,15 @@ def _add_to_bucket_deltas( period_date: date, raw_slots: ResourceSlot, segment_seconds: int, - user_deltas: dict[UserUsageBucketKey, BucketDelta], - project_deltas: dict[ProjectUsageBucketKey, BucketDelta], - domain_deltas: dict[DomainUsageBucketKey, BucketDelta], + user_deltas: dict[UserUsageBucketKey, ResourceSlot], + project_deltas: dict[ProjectUsageBucketKey, ResourceSlot], + domain_deltas: dict[DomainUsageBucketKey, ResourceSlot], ) -> 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. + Accumulating the amounts and the durations separately and multiplying + afterwards would give a cross product inflated by the slice count. Args: spec: Original spec (for entity identifiers) @@ -238,6 +236,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, @@ -247,11 +247,7 @@ def _add_to_bucket_deltas( resource_group_id=spec.resource_group_id, period_date=period_date, ) - ud = user_deltas[user_key] - user_deltas[user_key] = BucketDelta( - slots=ud.slots + raw_slots, - duration_seconds=ud.duration_seconds + segment_seconds, - ) + user_deltas[user_key] = user_deltas[user_key] + segment_usage # Project bucket key project_key = ProjectUsageBucketKey( @@ -261,11 +257,7 @@ def _add_to_bucket_deltas( resource_group_id=spec.resource_group_id, period_date=period_date, ) - pd = project_deltas[project_key] - project_deltas[project_key] = BucketDelta( - slots=pd.slots + raw_slots, - duration_seconds=pd.duration_seconds + segment_seconds, - ) + project_deltas[project_key] = project_deltas[project_key] + segment_usage # Domain bucket key domain_key = DomainUsageBucketKey( @@ -274,11 +266,7 @@ def _add_to_bucket_deltas( resource_group_id=spec.resource_group_id, period_date=period_date, ) - dd = domain_deltas[domain_key] - domain_deltas[domain_key] = BucketDelta( - slots=dd.slots + raw_slots, - duration_seconds=dd.duration_seconds + segment_seconds, - ) + domain_deltas[domain_key] = domain_deltas[domain_key] + segment_usage def _prepare_kernel_usage_specs( self, 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 4c4987da1da..276f7691548 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 @@ -276,7 +276,10 @@ async def test_create_kernel_usage_record( resource_group_id=test_resource_group_id, period_start=now - timedelta(minutes=5), period_end=now, - resource_usage=ResourceSlot({"cpu": Decimal("300"), "mem": Decimal("1073741824")}), + resource_usage=ResourceSlot({ + "cpu": Decimal("300"), + "mem": Decimal("1073741824"), + }), ) ) @@ -686,8 +689,7 @@ async def test_get_aggregated_usage_by_user( bucket_id=result.id, bucket_type="user", slot_name="cpu", - amount=Decimal("3600"), - duration_seconds=300, + resource_usage=Decimal("3600"), capacity=Decimal("0"), ) ) @@ -759,8 +761,7 @@ async def test_get_aggregated_usage_by_project( bucket_id=result.id, bucket_type="project", slot_name="cpu", - amount=Decimal("7200"), - duration_seconds=300, + resource_usage=Decimal("7200"), capacity=Decimal("0"), ) ) @@ -811,8 +812,7 @@ async def test_get_aggregated_usage_by_domain( bucket_id=result.id, bucket_type="domain", slot_name="cpu", - amount=Decimal("86400"), - duration_seconds=300, + resource_usage=Decimal("86400"), 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 b67b8cf1405..04c047b9593 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. An entry stores +occupied_slots integrated over time, which the read paths sum. """ from __future__ import annotations @@ -17,7 +18,6 @@ from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ResourceSlot from ai.backend.manager.data.fair_share import ( - BucketDelta, DomainUsageBucketKey, UsageBucketAggregationResult, UserUsageBucketKey, @@ -124,8 +124,7 @@ 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 + resource_usage = ResourceSlot({"cpu": Decimal("600"), "mem": Decimal("1228800000")}) period = date(2024, 1, 15) resource_group_id = ResourceGroupID(uuid.uuid4()) @@ -138,13 +137,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), + ): resource_usage, }, ) 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 = ( ( @@ -162,10 +161,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"].duration_seconds == 300 - assert slot_map["mem"].duration_seconds == 300 + assert slot_map["cpu"].resource_usage == Decimal("600") + assert slot_map["mem"].resource_usage == Decimal("1228800000") async def test_increment_accumulates_entries( self, @@ -183,33 +180,27 @@ 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")}), - duration_seconds=300, - ), + key: ResourceSlot({"cpu": Decimal("600")}), }, ) await db_source.increment_usage_buckets(result1) - # Second increment: 3 CPUs for 300 seconds + # Second increment: 3 CPUs for 300 seconds -> 900 CPU-seconds result2 = UsageBucketAggregationResult( user_usage_deltas={}, project_usage_deltas={}, domain_usage_deltas={ - key: BucketDelta( - slots=ResourceSlot({"cpu": Decimal("3")}), - duration_seconds=300, - ), + key: ResourceSlot({"cpu": Decimal("900")}), }, ) await db_source.increment_usage_buckets(result2) - # Verify accumulated: amount = 2 + 3 = 5, duration = 300 + 300 = 600 + # Verify accumulated: 600 + 900 = 1500 CPU-seconds async with db_with_cleanup.begin_readonly_session() as db_sess: entry_rows = ( ( @@ -225,8 +216,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].duration_seconds == 600 + assert entry_rows[0].resource_usage == Decimal("1500") stored_resource_group_id = await db_sess.scalar( sa.select(DomainUsageBucketRow.resource_group_id).where( @@ -246,8 +236,7 @@ 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 + resource_usage = ResourceSlot({"cpu": Decimal("900"), "cuda.device": Decimal("600")}) period = date(2024, 1, 15) resource_group_id = ResourceGroupID(uuid.uuid4()) @@ -260,7 +249,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), + ): resource_usage, }, project_usage_deltas={}, domain_usage_deltas={}, @@ -268,7 +257,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 = ( ( @@ -284,9 +273,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"].duration_seconds == 300 + assert slot_map["cpu"].resource_usage == Decimal("900") + assert slot_map["cuda.device"].resource_usage == Decimal("600") async def test_aggregated_usage_reads_from_entries( self, @@ -299,7 +287,7 @@ async def test_aggregated_usage_reads_from_entries( period2 = date(2024, 1, 16) resource_group_id = ResourceGroupID(uuid.uuid4()) - # 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={}, @@ -309,24 +297,18 @@ async def test_aggregated_usage_reads_from_entries( resource_group="default", resource_group_id=resource_group_id, period_date=period1, - ): BucketDelta( - slots=ResourceSlot({"cpu": Decimal("2")}), - duration_seconds=300, - ), + ): ResourceSlot({"cpu": Decimal("600")}), DomainUsageBucketKey( domain_name=test_domain_name, resource_group="default", resource_group_id=resource_group_id, period_date=period2, - ): BucketDelta( - slots=ResourceSlot({"cpu": Decimal("3")}), - duration_seconds=300, - ), + ): ResourceSlot({"cpu": Decimal("900")}), }, ) 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_id=resource_group_id, lookback_start=date(2024, 1, 14), @@ -334,5 +316,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/repositories/retention/test_retention_repository.py b/tests/unit/manager/repositories/retention/test_retention_repository.py index 14a68f3f348..303bd1c3ab7 100644 --- a/tests/unit/manager/repositories/retention/test_retention_repository.py +++ b/tests/unit/manager/repositories/retention/test_retention_repository.py @@ -1332,8 +1332,7 @@ async def _add_entry( bucket_id=bucket_id, bucket_type=bucket_type, slot_name=slot, - amount=Decimal("1"), - duration_seconds=1, + resource_usage=Decimal("1"), capacity=Decimal("1"), ) ) 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 430d2cb1d76..b987c1b6938 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. +Each bucket delta is the resource-seconds to add, summed per slice. """ 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,7 @@ 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. + Each delta is ``sum(amount * duration)`` over the slices in that bucket. """ def test_single_spec_single_day(self, aggregator: FairShareAggregator) -> None: @@ -200,8 +208,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") - assert delta.duration_seconds == 300 + # 2 CPU for 300s -> 600 CPU-seconds + assert delta["cpu"] == Decimal("600") # Project bucket assert len(result.project_usage_deltas) == 1 @@ -213,7 +221,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]["cpu"] == Decimal("600") # Domain bucket assert len(result.domain_usage_deltas) == 1 @@ -224,7 +232,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]["cpu"] == Decimal("600") def test_multiple_specs_same_user_same_day_aggregated( self, aggregator: FairShareAggregator @@ -255,14 +263,12 @@ 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["cpu"] == Decimal("1200") # Durations accumulate: 300 + 300 = 600 - assert delta.duration_seconds == 600 def test_spec_crossing_midnight_creates_two_buckets( self, aggregator: FairShareAggregator @@ -308,12 +314,10 @@ 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") - 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") - assert result.user_usage_deltas[day2_key].duration_seconds == 180 + # Day 1: 2 CPU for 180s -> 360 CPU-seconds + assert result.user_usage_deltas[day1_key]["cpu"] == Decimal("360") + # Day 2: 2 CPU for 180s -> 360 CPU-seconds + assert result.user_usage_deltas[day2_key]["cpu"] == Decimal("360") class TestBackloggedUsageScenario: @@ -436,17 +440,13 @@ 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.duration_seconds == 1440 + assert d1["cpu"] == Decimal("2880") - # 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.duration_seconds == 780 + assert d2["cpu"] == Decimal("1560") def test_backlogged_usage_with_midnight_crossing_spec( self, aggregator: FairShareAggregator @@ -540,18 +540,13 @@ 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.duration_seconds == 1440 + assert d1["cpu"] == Decimal("2880") - # 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.duration_seconds == 780 + assert d2["cpu"] == Decimal("1560") def test_backlogged_multiple_users(self, aggregator: FairShareAggregator) -> None: """Backlogged usage from multiple users is correctly separated.""" @@ -597,11 +592,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.duration_seconds == 600 # 300 + 300 + assert pd1["cpu"] == Decimal("1200") # 300 + 300 project_day2_key = ProjectUsageBucketKey( project_id=project_id, @@ -610,10 +603,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.duration_seconds == 300 + assert pd2["cpu"] == Decimal("600") class TestEdgeCases: @@ -649,8 +641,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.duration_seconds == 300 + assert delta["cpu"] == Decimal("600") def test_spec_starting_exactly_at_midnight(self, aggregator: FairShareAggregator) -> None: """Spec starting exactly at midnight is in the new day.""" @@ -696,8 +687,288 @@ 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.duration_seconds == 120 + assert delta["cpu"] == Decimal("240") + assert delta["mem"] == Decimal("491520") + assert delta["cuda.shares"] == Decimal("120") + + +@dataclass(frozen=True) +class _ConcurrentKernelsCase: + """One tick observing ``kernel_count`` kernels of the same user.""" + + kernel_count: int + expected_resource_usage: Decimal + + +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"), + ), + _ConcurrentKernelsCase( + kernel_count=2, + expected_resource_usage=Decimal("600"), + ), + _ConcurrentKernelsCase( + kernel_count=4, + expected_resource_usage=Decimal("1200"), + ), + _ConcurrentKernelsCase( + kernel_count=10, + expected_resource_usage=Decimal("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["cuda.shares"] == case.expected_resource_usage + + 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["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["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 + + +@dataclass(frozen=True) +class _ProjectBucketExpectation: + label: str + project_id: UUID + resource_usage: Decimal + + +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"), + ), + _UserBucketExpectation( + label="user-b-project-1", # 1 kernel * 2 shares * 300s + user_uuid=_USER_B, + project_id=_PROJECT_1, + resource_usage=Decimal("600"), + ), + _UserBucketExpectation( + label="user-a-project-2", # 3 kernels * 1 share * 300s + user_uuid=_USER_A, + project_id=_PROJECT_2, + resource_usage=Decimal("900"), + ), + _UserBucketExpectation( + label="user-c-project-2", # 1 kernel * 4 shares * 300s + user_uuid=_USER_C, + project_id=_PROJECT_2, + resource_usage=Decimal("1200"), + ), + ], + 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["cuda.shares"] == case.resource_usage + + @pytest.mark.parametrize( + "case", + [ + _ProjectBucketExpectation( + label="project-1", # user A 600 + user B 600 + project_id=_PROJECT_1, + resource_usage=Decimal("1200"), + ), + _ProjectBucketExpectation( + label="project-2", # user A 900 + user C 1200 + project_id=_PROJECT_2, + resource_usage=Decimal("2100"), + ), + ], + 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["cuda.shares"] == case.resource_usage + + 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["cuda.shares"] == Decimal("3300")