From f301430c965ba6598bd46baf18ce9a95dbf75b09 Mon Sep 17 00:00:00 2001 From: BoKeum Date: Fri, 3 Jul 2026 13:21:47 +0900 Subject: [PATCH 1/6] feat(BA-6644): add session scope id columns Co-Authored-By: Claude Fable 5 --- .../backend/manager/data/session/creation.py | 4 + src/ai/backend/manager/data/session/draft.py | 6 +- src/ai/backend/manager/data/session/spec.py | 6 +- ...a41cb881bb_add_session_scope_id_columns.py | 74 +++++++++++++++++++ src/ai/backend/manager/models/domain/row.py | 6 +- .../manager/models/scaling_group/row.py | 6 +- src/ai/backend/manager/models/session/row.py | 26 ++++++- src/ai/backend/manager/registry.py | 9 ++- .../deployment/db_source/db_source.py | 34 ++++++++- .../repositories/scheduler/creators.py | 2 + .../scheduler/db_source/db_source.py | 19 ++++- .../repositories/scheduler/repository.py | 9 +++ .../model_serving/services/model_serving.py | 9 ++- .../manager/services/session/service.py | 9 ++- .../deployment/deployment_draft_builder.py | 2 + 15 files changed, 208 insertions(+), 13 deletions(-) create mode 100644 src/ai/backend/manager/models/alembic/versions/ada41cb881bb_add_session_scope_id_columns.py diff --git a/src/ai/backend/manager/data/session/creation.py b/src/ai/backend/manager/data/session/creation.py index 28a905ac3d1..9b70c9afb11 100644 --- a/src/ai/backend/manager/data/session/creation.py +++ b/src/ai/backend/manager/data/session/creation.py @@ -11,6 +11,8 @@ from uuid import UUID from ai.backend.common.docker import ImageRef +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import AccessKey @@ -62,7 +64,9 @@ class DeploymentContext: created_user: UserContext session_owner: UserContext container_user: ContainerUserContext + domain_id: DomainID group_id: UUID + resource_group_id: ResourceGroupID resource_policy: dict[str, Any] image: ImageContext resolved_presets: ResolvedPresetValues | None = None diff --git a/src/ai/backend/manager/data/session/draft.py b/src/ai/backend/manager/data/session/draft.py index 479128463bc..561406bc322 100644 --- a/src/ai/backend/manager/data/session/draft.py +++ b/src/ai/backend/manager/data/session/draft.py @@ -37,10 +37,10 @@ import yarl from pydantic import ConfigDict, Field -from ai.backend.common.identifier.domain import DomainName +from ai.backend.common.identifier.domain import DomainID, DomainName from ai.backend.common.identifier.image import ImageID from ai.backend.common.identifier.project import ProjectID -from ai.backend.common.identifier.resource_group import ResourceGroupName +from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName from ai.backend.common.identifier.session import SessionID from ai.backend.common.types import ( AccessKey, @@ -179,8 +179,10 @@ class SessionNetworkDraft(_DraftBaseModel): class SessionScopeDraft(_DraftBaseModel): """Optional-heavy mirror of ``SessionScope``.""" + domain_id: DomainID | None = None domain_name: DomainName | None = None project_id: ProjectID | None = None + resource_group_id: ResourceGroupID | None = None resource_group_name: ResourceGroupName | None = None diff --git a/src/ai/backend/manager/data/session/spec.py b/src/ai/backend/manager/data/session/spec.py index d6a12d730e5..c99e30f06f0 100644 --- a/src/ai/backend/manager/data/session/spec.py +++ b/src/ai/backend/manager/data/session/spec.py @@ -24,9 +24,9 @@ from pydantic import ConfigDict, Field from ai.backend.common.exception import BackendAIError -from ai.backend.common.identifier.domain import DomainName +from ai.backend.common.identifier.domain import DomainID, DomainName from ai.backend.common.identifier.project import ProjectID -from ai.backend.common.identifier.resource_group import ResourceGroupName +from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName from ai.backend.common.identifier.session import SessionID from ai.backend.common.types import ( AccessKey, @@ -69,8 +69,10 @@ class SessionIdentity(_SpecBaseModel): class SessionScope(_SpecBaseModel): """Ownership / placement scope of the session.""" + domain_id: DomainID domain_name: DomainName project_id: ProjectID + resource_group_id: ResourceGroupID resource_group_name: ResourceGroupName diff --git a/src/ai/backend/manager/models/alembic/versions/ada41cb881bb_add_session_scope_id_columns.py b/src/ai/backend/manager/models/alembic/versions/ada41cb881bb_add_session_scope_id_columns.py new file mode 100644 index 00000000000..255bab14ab9 --- /dev/null +++ b/src/ai/backend/manager/models/alembic/versions/ada41cb881bb_add_session_scope_id_columns.py @@ -0,0 +1,74 @@ +"""add session scope id columns + +Revision ID: ada41cb881bb +Revises: b4c5d6e7f8a9 +Create Date: 2026-06-30 + +""" + +import sqlalchemy as sa +from alembic import op + +from ai.backend.manager.models.base import GUID + +# revision identifiers, used by Alembic. +revision = "ada41cb881bb" +down_revision = "b4c5d6e7f8a9" +# Part of: 26.7.0 +branch_labels = None +depends_on = None + + +def upgrade() -> None: + op.add_column("sessions", sa.Column("domain_id", GUID(), nullable=True)) + op.add_column("sessions", sa.Column("resource_group_id", GUID(), nullable=True)) + op.create_index("ix_sessions_domain_id", "sessions", ["domain_id"]) + op.create_index("ix_sessions_resource_group_id", "sessions", ["resource_group_id"]) + + op.execute( + sa.text(""" + UPDATE sessions + SET domain_id = domains.id + FROM domains + WHERE sessions.domain_name = domains.name + AND sessions.domain_id IS NULL + """) + ) + op.execute( + sa.text(""" + UPDATE sessions + SET resource_group_id = scaling_groups.id + FROM scaling_groups + WHERE sessions.scaling_group_name = scaling_groups.name + AND sessions.resource_group_id IS NULL + """) + ) + op.alter_column("sessions", "domain_id", nullable=False) + op.alter_column("sessions", "resource_group_id", nullable=False) + op.create_foreign_key( + op.f("fk_sessions_domain_id_domains"), + "sessions", + "domains", + ["domain_id"], + ["id"], + ) + op.create_foreign_key( + op.f("fk_sessions_resource_group_id_scaling_groups"), + "sessions", + "scaling_groups", + ["resource_group_id"], + ["id"], + ) + + +def downgrade() -> None: + op.drop_constraint( + op.f("fk_sessions_resource_group_id_scaling_groups"), + "sessions", + type_="foreignkey", + ) + op.drop_constraint(op.f("fk_sessions_domain_id_domains"), "sessions", type_="foreignkey") + op.drop_index("ix_sessions_resource_group_id", table_name="sessions") + op.drop_index("ix_sessions_domain_id", table_name="sessions") + op.drop_column("sessions", "resource_group_id") + op.drop_column("sessions", "domain_id") diff --git a/src/ai/backend/manager/models/domain/row.py b/src/ai/backend/manager/models/domain/row.py index ad73dfafafc..e8a188f8326 100644 --- a/src/ai/backend/manager/models/domain/row.py +++ b/src/ai/backend/manager/models/domain/row.py @@ -135,7 +135,11 @@ class DomainRow(Base): # type: ignore[misc] "dotfiles", sa.LargeBinary(length=MAXIMUM_DOTFILE_SIZE), nullable=False, default=b"\x90" ) - sessions: Mapped[list[SessionRow]] = relationship("SessionRow", back_populates="domain") + sessions: Mapped[list[SessionRow]] = relationship( + "SessionRow", + back_populates="domain", + foreign_keys="[SessionRow.domain_name]", + ) users: Mapped[list[UserRow]] = relationship("UserRow", back_populates="domain") groups: Mapped[list[GroupRow]] = relationship("GroupRow", back_populates="domain") sgroup_for_domains_rows: Mapped[list[ScalingGroupForDomainRow]] = relationship( diff --git a/src/ai/backend/manager/models/scaling_group/row.py b/src/ai/backend/manager/models/scaling_group/row.py index 24443026642..6481ffef6f2 100644 --- a/src/ai/backend/manager/models/scaling_group/row.py +++ b/src/ai/backend/manager/models/scaling_group/row.py @@ -330,7 +330,11 @@ class ScalingGroupRow(Base): # type: ignore[misc] default=DefaultSessionOptions, ) - sessions: Mapped[list[SessionRow]] = relationship("SessionRow", back_populates="scaling_group") + sessions: Mapped[list[SessionRow]] = relationship( + "SessionRow", + back_populates="scaling_group", + foreign_keys="[SessionRow.scaling_group_name]", + ) agents: Mapped[list[AgentRow]] = relationship("AgentRow", back_populates="scaling_group_row") sgroup_for_domains_rows: Mapped[list[ScalingGroupForDomainRow]] = relationship( diff --git a/src/ai/backend/manager/models/session/row.py b/src/ai/backend/manager/models/session/row.py index 371acabc8ff..930c535d440 100644 --- a/src/ai/backend/manager/models/session/row.py +++ b/src/ai/backend/manager/models/session/row.py @@ -39,6 +39,8 @@ from ai.backend.common.defs.session import SESSION_PRIORITY_DEFAULT from ai.backend.common.exception import BackendAIError +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ( AccessKey, ClusterMode, @@ -453,11 +455,20 @@ class SessionRow(Base): # type: ignore[misc] kernels: Mapped[list[KernelRow]] = relationship("KernelRow", back_populates="session") # Resource ownership + resource_group_id: Mapped[ResourceGroupID] = mapped_column( + "resource_group_id", + GUID(ResourceGroupID), + sa.ForeignKey("scaling_groups.id"), + index=True, + nullable=False, + ) scaling_group_name: Mapped[str] = mapped_column( "scaling_group_name", sa.ForeignKey("scaling_groups.name"), index=True, nullable=False ) scaling_group: Mapped[ScalingGroupRow] = relationship( - "ScalingGroupRow", back_populates="sessions" + "ScalingGroupRow", + back_populates="sessions", + foreign_keys=[scaling_group_name], ) target_sgroup_names: Mapped[list[str] | None] = mapped_column( "target_sgroup_names", @@ -469,7 +480,18 @@ class SessionRow(Base): # type: ignore[misc] domain_name: Mapped[str] = mapped_column( "domain_name", sa.String(length=64), sa.ForeignKey("domains.name"), nullable=False ) - domain: Mapped[DomainRow] = relationship("DomainRow", back_populates="sessions") + domain_id: Mapped[DomainID] = mapped_column( + "domain_id", + GUID(DomainID), + sa.ForeignKey("domains.id"), + index=True, + nullable=False, + ) + domain: Mapped[DomainRow] = relationship( + "DomainRow", + back_populates="sessions", + foreign_keys=[domain_name], + ) group_id: Mapped[UUID] = mapped_column( "group_id", GUID, sa.ForeignKey("groups.id"), nullable=False ) diff --git a/src/ai/backend/manager/registry.py b/src/ai/backend/manager/registry.py index 124795d8e8f..db6aca16062 100644 --- a/src/ai/backend/manager/registry.py +++ b/src/ai/backend/manager/registry.py @@ -1119,8 +1119,13 @@ def _build_execution_spec( if not groups_by_role: raise InvalidAPIParameters("No kernel groups resolved from the enqueue request.") + domain_name = DomainName(user_scope.domain_name) + domain_id = await self._scheduler_repository.get_domain_id_by_name(domain_name) if scaling_group: resource_group_name = ResourceGroupName(scaling_group) + resource_group_id = await self._scheduler_repository.get_resource_group_id_by_name( + resource_group_name + ) else: resource_group_id = await self._scheduler_repository.pick_default_resource_group( access_key=access_key, @@ -1140,8 +1145,10 @@ def _build_execution_spec( user_uuid=user_scope.user_uuid, ), scope=SessionScopeDraft( - domain_name=DomainName(user_scope.domain_name), + domain_id=domain_id, + domain_name=domain_name, project_id=ProjectID(user_scope.group_id), + resource_group_id=resource_group_id, resource_group_name=resource_group_name, ), classification=SessionClassificationDraft( diff --git a/src/ai/backend/manager/repositories/deployment/db_source/db_source.py b/src/ai/backend/manager/repositories/deployment/db_source/db_source.py index e1d2669ad20..43ba09daf8e 100644 --- a/src/ai/backend/manager/repositories/deployment/db_source/db_source.py +++ b/src/ai/backend/manager/repositories/deployment/db_source/db_source.py @@ -29,10 +29,11 @@ from ai.backend.common.identifier.deployment import DeploymentID from ai.backend.common.identifier.deployment_preset import DeploymentPresetID from ai.backend.common.identifier.deployment_revision import DeploymentRevisionID +from ai.backend.common.identifier.domain import DomainID from ai.backend.common.identifier.image import ImageID from ai.backend.common.identifier.replica import ReplicaID from ai.backend.common.identifier.replica_group import ReplicaGroupID -from ai.backend.common.identifier.resource_group import ResourceGroupName +from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName from ai.backend.common.identifier.runtime_variant import RuntimeVariantID from ai.backend.common.identifier.vfolder import VFolderUUID from ai.backend.common.types import ( @@ -105,6 +106,7 @@ UserNotFoundInDeployment, ) from ai.backend.manager.errors.resource import ( + DomainNotFound, ProjectNotFound, RuntimeVariantNotFound, ScalingGroupNotFound, @@ -124,6 +126,7 @@ from ai.backend.manager.models.deployment_revision_preset.row import ( DeploymentRevisionPresetRow, ) +from ai.backend.manager.models.domain import DomainRow from ai.backend.manager.models.endpoint import ( EndpointAutoScalingRuleRow, EndpointRow, @@ -2099,6 +2102,29 @@ async def _resolve_user_and_active_access_key( raise UserNotFoundInDeployment(f"{user_uuid} not found") raise NoActiveKeypairForDeployment(f"{user_uuid} has no active keypair") + async def _resolve_deployment_scope_ids( + self, + db_sess: SASession, + deployment_info: DeploymentInfo, + ) -> tuple[DomainID, ResourceGroupID]: + domain_id = await db_sess.scalar( + sa.select(DomainRow.id).where(DomainRow.name == deployment_info.metadata.domain) + ) + if domain_id is None: + raise DomainNotFound(deployment_info.metadata.domain) + + resource_group_id = await db_sess.scalar( + sa.select(ScalingGroupRow.id).where( + ScalingGroupRow.name == deployment_info.metadata.resource_group + ) + ) + if resource_group_id is None: + raise ScalingGroupNotFound( + f"Resource group {deployment_info.metadata.resource_group!r} not found" + ) + + return DomainID(domain_id), ResourceGroupID(resource_group_id) + async def fetch_deployment_context( self, deployment_info: DeploymentInfo, @@ -2147,6 +2173,10 @@ async def fetch_deployment_context( group_id = user_info.group_id resource_policy = user_info.resource_policy + domain_id, resource_group_id = await self._resolve_deployment_scope_ids( + db_sess, deployment_info + ) + revision_query = ( sa.select(DeploymentRevisionRow) .where(DeploymentRevisionRow.id == revision_id) @@ -2212,7 +2242,9 @@ async def fetch_deployment_context( main_gid=session_owner_user.container_main_gid, supplementary_gids=session_owner_user.container_gids or [], ), + domain_id=domain_id, group_id=group_id, + resource_group_id=resource_group_id, resource_policy=dict(resource_policy), image=ImageContext( ref=image_row.image_ref, diff --git a/src/ai/backend/manager/repositories/scheduler/creators.py b/src/ai/backend/manager/repositories/scheduler/creators.py index 57c5846e9bd..6d41b4f93e1 100644 --- a/src/ai/backend/manager/repositories/scheduler/creators.py +++ b/src/ai/backend/manager/repositories/scheduler/creators.py @@ -190,7 +190,9 @@ def build_row(self) -> SessionRow: access_key=spec.identity.access_key, user_uuid=spec.identity.user_uuid, group_id=spec.scope.project_id, + domain_id=spec.scope.domain_id, domain_name=str(spec.scope.domain_name), + resource_group_id=spec.scope.resource_group_id, scaling_group_name=str(spec.scope.resource_group_name), session_type=spec.classification.session_type, cluster_mode=cluster_mode_value, diff --git a/src/ai/backend/manager/repositories/scheduler/db_source/db_source.py b/src/ai/backend/manager/repositories/scheduler/db_source/db_source.py index ac7fa8ea8cf..cc7e6eac129 100644 --- a/src/ai/backend/manager/repositories/scheduler/db_source/db_source.py +++ b/src/ai/backend/manager/repositories/scheduler/db_source/db_source.py @@ -29,6 +29,7 @@ RBACElementType, ) from ai.backend.common.docker import ImageRef +from ai.backend.common.identifier.domain import DomainID, DomainName from ai.backend.common.identifier.image import ImageID from ai.backend.common.identifier.project import ProjectID from ai.backend.common.identifier.resource_group import ( @@ -86,7 +87,7 @@ from ai.backend.manager.errors.api import InvalidAPIParameters from ai.backend.manager.errors.common import InternalServerError from ai.backend.manager.errors.image import ImageNotFound -from ai.backend.manager.errors.resource import ScalingGroupNotFound +from ai.backend.manager.errors.resource import DomainNotFound, ScalingGroupNotFound from ai.backend.manager.errors.resource_slot import AgentResourceCapacityExceeded from ai.backend.manager.exceptions import ErrorStatusInfo from ai.backend.manager.models.agent import AgentRow @@ -1762,6 +1763,15 @@ async def pick_default_resource_group( raise InvalidAPIParameters("No accessible scaling group available") return allowed_rgs[0].id + async def get_resource_group_id_by_name(self, name: ResourceGroupName) -> ResourceGroupID: + async with self._begin_readonly_session_read_committed() as db_sess: + resource_group_id = await db_sess.scalar( + sa.select(ScalingGroupRow.id).where(ScalingGroupRow.name == name) + ) + if resource_group_id is None: + raise ScalingGroupNotFound(name) + return ResourceGroupID(resource_group_id) + async def get_resource_group_name_by_id( self, resource_group_id: ResourceGroupID ) -> ResourceGroupName: @@ -1773,6 +1783,13 @@ async def get_resource_group_name_by_id( raise ScalingGroupNotFound(f"Resource group not found (id:{resource_group_id})") return ResourceGroupName(resource_group_name) + async def get_domain_id_by_name(self, name: DomainName) -> DomainID: + async with self._begin_readonly_session_read_committed() as db_sess: + domain_id = await db_sess.scalar(sa.select(DomainRow.id).where(DomainRow.name == name)) + if domain_id is None: + raise DomainNotFound(name) + return DomainID(domain_id) + async def _get_scaling_group_network_info( self, db_sess: SASession, scaling_group_name: str ) -> ScalingGroupNetworkInfo: diff --git a/src/ai/backend/manager/repositories/scheduler/repository.py b/src/ai/backend/manager/repositories/scheduler/repository.py index b4d98c805b9..3ef8350a789 100644 --- a/src/ai/backend/manager/repositories/scheduler/repository.py +++ b/src/ai/backend/manager/repositories/scheduler/repository.py @@ -15,6 +15,7 @@ from ai.backend.common.clients.valkey_client.valkey_stat.client import ValkeyStatClient from ai.backend.common.exception import BackendAIError +from ai.backend.common.identifier.domain import DomainID, DomainName from ai.backend.common.identifier.image import ImageID from ai.backend.common.identifier.project import ProjectID from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName @@ -283,12 +284,20 @@ async def pick_default_resource_group( project_id=project_id, ) + @scheduler_repository_resilience.apply() + async def get_resource_group_id_by_name(self, name: ResourceGroupName) -> ResourceGroupID: + return await self._db_source.get_resource_group_id_by_name(name) + @scheduler_repository_resilience.apply() async def get_resource_group_name_by_id( self, resource_group_id: ResourceGroupID ) -> ResourceGroupName: return await self._db_source.get_resource_group_name_by_id(resource_group_id) + @scheduler_repository_resilience.apply() + async def get_domain_id_by_name(self, name: DomainName) -> DomainID: + return await self._db_source.get_domain_id_by_name(name) + @scheduler_repository_resilience.apply() async def prepare_vfolder_mounts( self, diff --git a/src/ai/backend/manager/services/model_serving/services/model_serving.py b/src/ai/backend/manager/services/model_serving/services/model_serving.py index 0871a2f04c2..dc4a0dbd894 100644 --- a/src/ai/backend/manager/services/model_serving/services/model_serving.py +++ b/src/ai/backend/manager/services/model_serving/services/model_serving.py @@ -452,8 +452,13 @@ async def dry_run(self, action: DryRunModelServiceAction) -> DryRunModelServiceA ), ) + domain_name = DomainName(action.domain_name) + domain_id = await self._scheduler_repository.get_domain_id_by_name(domain_name) if service_prepare_ctx.scaling_group: resource_group_name = ResourceGroupName(service_prepare_ctx.scaling_group) + resource_group_id = await self._scheduler_repository.get_resource_group_id_by_name( + resource_group_name + ) else: resource_group_id = await self._scheduler_repository.pick_default_resource_group( access_key=AccessKey(service_prepare_ctx.owner_access_key), @@ -473,8 +478,10 @@ async def dry_run(self, action: DryRunModelServiceAction) -> DryRunModelServiceA user_uuid=created_user.uuid, ), scope=SessionScopeDraft( - domain_name=DomainName(action.domain_name), + domain_id=domain_id, + domain_name=domain_name, project_id=ProjectID(service_prepare_ctx.group_id), + resource_group_id=resource_group_id, resource_group_name=resource_group_name, ), classification=SessionClassificationDraft( diff --git a/src/ai/backend/manager/services/session/service.py b/src/ai/backend/manager/services/session/service.py index 75c0cba461e..894ed19042e 100644 --- a/src/ai/backend/manager/services/session/service.py +++ b/src/ai/backend/manager/services/session/service.py @@ -1578,11 +1578,15 @@ async def enqueue_session(self, action: EnqueueSessionAction) -> EnqueueSessionA callback_url = yarl.URL(action.callback_url) if action.callback_url else None if action.resource.resource_group_id is not None: + resource_group_id = action.resource.resource_group_id resource_group_name = await self._scheduler_repository.get_resource_group_name_by_id( - action.resource.resource_group_id + resource_group_id ) elif action.resource.resource_group: resource_group_name = ResourceGroupName(action.resource.resource_group) + resource_group_id = await self._scheduler_repository.get_resource_group_id_by_name( + resource_group_name + ) else: resource_group_id = await self._scheduler_repository.pick_default_resource_group( access_key=access_key, @@ -1592,6 +1596,7 @@ async def enqueue_session(self, action: EnqueueSessionAction) -> EnqueueSessionA resource_group_name = await self._scheduler_repository.get_resource_group_name_by_id( resource_group_id ) + domain_id = await self._scheduler_repository.get_domain_id_by_name(DomainName(domain_name)) kernel_groups = await self._resolve_kernel_groups( cluster_size=action.resource.cluster_size, preopen_ports=preopen_ports, @@ -1617,8 +1622,10 @@ async def enqueue_session(self, action: EnqueueSessionAction) -> EnqueueSessionA user_uuid=user_id, ), scope=SessionScopeDraft( + domain_id=domain_id, domain_name=DomainName(domain_name), project_id=ProjectID(action.group_id), + resource_group_id=resource_group_id, resource_group_name=resource_group_name, ), classification=SessionClassificationDraft( diff --git a/src/ai/backend/manager/sokovan/deployment/deployment_draft_builder.py b/src/ai/backend/manager/sokovan/deployment/deployment_draft_builder.py index 8eb15bf0cbf..69c5a8b2de0 100644 --- a/src/ai/backend/manager/sokovan/deployment/deployment_draft_builder.py +++ b/src/ai/backend/manager/sokovan/deployment/deployment_draft_builder.py @@ -102,8 +102,10 @@ def build( user_uuid=context.session_owner.uuid, ), scope=SessionScopeDraft( + domain_id=context.domain_id, domain_name=DomainName(deployment_info.metadata.domain), project_id=ProjectID(context.group_id), + resource_group_id=context.resource_group_id, resource_group_name=ResourceGroupName(deployment_info.metadata.resource_group), ), classification=SessionClassificationDraft( From aa76655f461802c9e2f5f3231dcecb419915d2c5 Mon Sep 17 00:00:00 2001 From: BoKeum Date: Fri, 3 Jul 2026 13:27:09 +0900 Subject: [PATCH 2/6] test(BA-6644): update fixtures and tests for session scope id columns Co-Authored-By: Claude Fable 5 --- src/ai/backend/testutils/fixtures.py | 8 + tests/component/agent_api/conftest.py | 5 +- tests/component/agent_api/test_agent_api.py | 15 +- tests/component/auto_scaling_rule/conftest.py | 6 +- tests/component/conftest.py | 19 +- tests/component/deployment/test_deployment.py | 46 ++-- .../deployment/test_deployment_lifecycle.py | 22 +- .../domain/test_domain_management.py | 63 ++++-- tests/component/fair_share/conftest.py | 6 +- tests/component/fair_share/test_fair_share.py | 124 +++++----- .../test_fair_share_domain_weights.py | 52 ++--- .../test_fair_share_project_weights.py | 46 ++-- .../test_fair_share_user_weights.py | 46 ++-- .../component/group/test_group_management.py | 18 +- tests/component/infra/test_infra.py | 5 +- tests/component/operations/conftest.py | 5 +- .../test_resource_allocation.py | 18 +- .../test_resource_allocation_visibility.py | 17 +- .../resource_preset/test_resource_preset.py | 9 +- .../test_resource_preset_full.py | 7 +- .../scaling_group/test_scaling_group.py | 29 ++- .../test_scaling_group_associations.py | 17 +- .../scaling_group/test_scaling_group_crud.py | 16 +- tests/component/session/conftest.py | 30 ++- tests/component/session/test_session_batch.py | 5 +- .../session/test_session_execution.py | 41 ++-- .../session/test_session_lifecycle.py | 18 +- tests/component/session/test_session_query.py | 28 ++- tests/component/session_v2/conftest.py | 20 +- tests/component/stream/conftest.py | 10 +- tests/component/streaming/conftest.py | 18 +- tests/unit/manager/models/test_session.py | 71 +++--- .../repositories/agent/test_repository.py | 18 ++ .../deployment/test_deployment_repository.py | 22 ++ .../model_card/test_model_card_delete.py | 38 +++- .../resource_preset/test_check_presets.py | 30 +++ .../test_scaling_group_repository.py | 20 ++ .../repositories/scheduler/conftest.py | 20 ++ .../test_cancellation_resource_freeing.py | 26 +++ .../scheduler/test_reserved_allocation.py | 40 ++++ .../scheduler/test_resource_allocation.py | 46 ++++ .../scheduler/test_resource_concurrency.py | 14 ++ .../scheduler/test_resource_deallocation.py | 66 ++++++ .../test_scheduling_history_recording.py | 40 ++++ .../test_session_network_propagation.py | 10 + .../scheduler/test_termination.py | 18 ++ .../scheduler/test_update_with_history.py | 55 ++++- .../session/test_session_repository.py | 40 +++- .../session/test_session_search_in_project.py | 16 ++ .../stream/test_stream_repository.py | 29 ++- .../services/session/test_session_service.py | 27 ++- .../preparers/test_session_spec_preparer.py | 6 +- .../test_enqueue_session_from_draft.py | 6 +- .../validators/test_session_spec_rules.py | 6 +- .../manager/test_reconcile_agent_resources.py | 211 +++++++++++------- 55 files changed, 1164 insertions(+), 480 deletions(-) diff --git a/src/ai/backend/testutils/fixtures.py b/src/ai/backend/testutils/fixtures.py index 192df7c2ae0..2c3bd1b0a7d 100644 --- a/src/ai/backend/testutils/fixtures.py +++ b/src/ai/backend/testutils/fixtures.py @@ -4,10 +4,12 @@ from dataclasses import dataclass from ai.backend.common.identifier.domain import DomainID, DomainName +from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName __all__ = ( "DomainFactory", "DomainFixtureData", + "ScalingGroupFixtureData", ) @@ -17,4 +19,10 @@ class DomainFixtureData: domain_id: DomainID +@dataclass(frozen=True) +class ScalingGroupFixtureData: + scaling_group_name: ResourceGroupName + scaling_group_id: ResourceGroupID + + DomainFactory = Callable[..., Awaitable[DomainFixtureData]] diff --git a/tests/component/agent_api/conftest.py b/tests/component/agent_api/conftest.py index c0302b9565f..d6aebb2475c 100644 --- a/tests/component/agent_api/conftest.py +++ b/tests/component/agent_api/conftest.py @@ -28,6 +28,7 @@ from ai.backend.manager.repositories.scheduler.repository import SchedulerRepository from ai.backend.manager.services.agent.processors import AgentProcessors from ai.backend.manager.services.agent.service import AgentService +from ai.backend.testutils.fixtures import ScalingGroupFixtureData @pytest.fixture() @@ -105,7 +106,7 @@ def server_module_registries( @pytest.fixture() async def agent_fixture( db_engine: SAEngine, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> AsyncIterator[str]: """Insert a test agent row and yield its ID. @@ -119,7 +120,7 @@ async def agent_fixture( id=agent_id, status=AgentStatus.ALIVE, region="local", - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, schedulable=True, available_slots=ResourceSlot({"cpu": "4", "mem": "8589934592"}), occupied_slots=ResourceSlot(), diff --git a/tests/component/agent_api/test_agent_api.py b/tests/component/agent_api/test_agent_api.py index 52b77015b12..8769a4077fd 100644 --- a/tests/component/agent_api/test_agent_api.py +++ b/tests/component/agent_api/test_agent_api.py @@ -26,6 +26,7 @@ OrderDirection, ) from ai.backend.common.dto.manager.query import StringFilter +from ai.backend.testutils.fixtures import ScalingGroupFixtureData class TestSearchAgents: @@ -83,13 +84,14 @@ async def test_admin_searches_agents_with_resource_group_filter( self, admin_registry: BackendAIClientRegistry, agent_fixture: str, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: """Filtering by resource_group returns agents in that scaling group.""" + sgroup_name = scaling_group_fixture.scaling_group_name result = await admin_registry.agent.search_agents( SearchAgentsRequest( filter=AgentFilter( - resource_group=StringFilter(equals=scaling_group_fixture), + resource_group=StringFilter(equals=sgroup_name), ), ), ) @@ -97,7 +99,7 @@ async def test_admin_searches_agents_with_resource_group_filter( agent_ids = [item.id for item in result.items] assert agent_fixture in agent_ids for item in result.items: - assert item.resource_group == scaling_group_fixture + assert item.resource_group == sgroup_name async def test_admin_searches_agents_with_ordering( self, @@ -138,15 +140,16 @@ async def test_compound_filters_status_and_resource_group( self, admin_registry: BackendAIClientRegistry, agent_fixture: str, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: """Compound filters (status + resource_group) return intersection of conditions.""" # Search with both status=ALIVE and resource_group filters + sgroup_name = scaling_group_fixture.scaling_group_name result = await admin_registry.agent.search_agents( SearchAgentsRequest( filter=AgentFilter( status=AgentStatusEnumFilter(equals=AgentStatusEnum.ALIVE), - resource_group=StringFilter(equals=scaling_group_fixture), + resource_group=StringFilter(equals=sgroup_name), ), ), ) @@ -154,7 +157,7 @@ async def test_compound_filters_status_and_resource_group( # All results should match both filters for item in result.items: assert item.status == AgentStatusEnum.ALIVE - assert item.resource_group == scaling_group_fixture + assert item.resource_group == sgroup_name # Our fixture agent should be in the results agent_ids = [item.id for item in result.items] diff --git a/tests/component/auto_scaling_rule/conftest.py b/tests/component/auto_scaling_rule/conftest.py index 14f770e1b7f..9cc303888c6 100644 --- a/tests/component/auto_scaling_rule/conftest.py +++ b/tests/component/auto_scaling_rule/conftest.py @@ -37,7 +37,7 @@ ) from ai.backend.manager.services.deployment.processors import DeploymentProcessors from ai.backend.manager.services.deployment.service import DeploymentService -from ai.backend.testutils.fixtures import DomainFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData @dataclass @@ -125,7 +125,7 @@ async def model_deployment_fixture( db_engine: SAEngine, domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, admin_user_fixture: UserFixtureData, ) -> AsyncIterator[uuid.UUID]: """Insert a minimal EndpointRow (model deployment) and yield its UUID. @@ -176,7 +176,7 @@ async def model_deployment_fixture( session_owner=str(admin_user_fixture.user_uuid), domain=domain_fixture.domain_name, project=str(group_fixture), - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, lifecycle_stage=EndpointLifecycle.CREATED, url=None, ) diff --git a/tests/component/conftest.py b/tests/component/conftest.py index 26ce97f317e..85e408e9815 100644 --- a/tests/component/conftest.py +++ b/tests/component/conftest.py @@ -57,7 +57,7 @@ ) from ai.backend.common.etcd import AsyncEtcd, ConfigScopes from ai.backend.common.events.dispatcher import EventProducer -from ai.backend.common.identifier.resource_group import ResourceGroupName +from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName from ai.backend.common.message_queue.redis_queue.queue import RedisMQArgs, RedisQueue from ai.backend.common.plugin.hook import HookPluginContext from ai.backend.common.plugin.monitor import ErrorPluginContext, StatsPluginContext @@ -152,7 +152,7 @@ postgres_container, redis_container, ) -from ai.backend.testutils.fixtures import DomainFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData from ai.backend.testutils.pants import get_parallel_slot log = logging.getLogger("tests.component.conftest") @@ -676,12 +676,13 @@ async def resource_policy_fixture( async def scaling_group_fixture( db_engine: SAEngine, domain_fixture: DomainFixtureData, -) -> AsyncIterator[ResourceGroupName]: - """Insert a scaling group and its domain association; yield the name.""" +) -> AsyncIterator[ScalingGroupFixtureData]: + """Insert a scaling group and its domain association; yield its identifiers.""" sgroup_name = ResourceGroupName(f"sgroup-{secrets.token_hex(6)}") async with db_engine.begin() as conn: - await conn.execute( - sa.insert(scaling_groups).values( + result = await conn.execute( + sa.insert(scaling_groups) + .values( name=sgroup_name, description=f"Test scaling group {sgroup_name}", is_active=True, @@ -690,14 +691,16 @@ async def scaling_group_fixture( scheduler="fifo", scheduler_opts=ScalingGroupOpts(), ) + .returning(scaling_groups.c.id) ) + sgroup_id = ResourceGroupID(result.scalar_one()) await conn.execute( sa.insert(sgroups_for_domains).values( scaling_group=sgroup_name, domain=domain_fixture.domain_name, ) ) - yield sgroup_name + yield ScalingGroupFixtureData(scaling_group_name=sgroup_name, scaling_group_id=sgroup_id) async with db_engine.begin() as conn: await conn.execute( sgroups_for_domains.delete().where(sgroups_for_domains.c.scaling_group == sgroup_name) @@ -933,7 +936,7 @@ async def regular_user_fixture( async def database_fixture( admin_user_fixture: UserFixtureData, regular_user_fixture: UserFixtureData, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> AsyncIterator[None]: """Backward-compatible aggregate: requests all seed data fixtures.""" yield diff --git a/tests/component/deployment/test_deployment.py b/tests/component/deployment/test_deployment.py index b7d6e5fc06a..9cb6741a7b8 100644 --- a/tests/component/deployment/test_deployment.py +++ b/tests/component/deployment/test_deployment.py @@ -56,7 +56,7 @@ from ai.backend.manager.services.deployment.processors import DeploymentProcessors from ai.backend.manager.services.deployment.service import _map_lifecycle_to_status from ai.backend.manager.services.processors import Processors -from ai.backend.testutils.fixtures import DomainFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData if TYPE_CHECKING: from tests.component.conftest import UserFixtureData @@ -96,7 +96,7 @@ async def test_search_deployments_paginated( admin_registry: BackendAIClientRegistry, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ResourceGroupName, + scaling_group_fixture: ScalingGroupFixtureData, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """Search deployments with pagination returns correct page.""" @@ -108,7 +108,7 @@ async def test_search_deployments_paginated( metadata=DeploymentMetadataInput( project_id=group_fixture, domain_name=domain_fixture.domain_name, - resource_group_name=scaling_group_fixture, + resource_group_name=scaling_group_fixture.scaling_group_name, name=f"test-deployment-{i}-{secrets.token_hex(4)}", ), network_access=NetworkAccessInput(open_to_public=False), @@ -158,7 +158,7 @@ async def test_get_deployment_by_id( admin_registry: BackendAIClientRegistry, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ResourceGroupName, + scaling_group_fixture: ScalingGroupFixtureData, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """Get deployment by ID returns correct deployment details.""" @@ -167,7 +167,7 @@ async def test_get_deployment_by_id( metadata=DeploymentMetadataInput( project_id=group_fixture, domain_name=domain_fixture.domain_name, - resource_group_name=scaling_group_fixture, + resource_group_name=scaling_group_fixture.scaling_group_name, name=f"test-deployment-{secrets.token_hex(4)}", ), network_access=NetworkAccessInput(open_to_public=False), @@ -269,7 +269,7 @@ async def test_search_routes_for_deployment( admin_registry: BackendAIClientRegistry, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ResourceGroupName, + scaling_group_fixture: ScalingGroupFixtureData, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """Search routes for a deployment returns route list.""" @@ -278,7 +278,7 @@ async def test_search_routes_for_deployment( metadata=DeploymentMetadataInput( project_id=group_fixture, domain_name=domain_fixture.domain_name, - resource_group_name=scaling_group_fixture, + resource_group_name=scaling_group_fixture.scaling_group_name, name=f"test-deployment-{secrets.token_hex(4)}", ), network_access=NetworkAccessInput(open_to_public=False), @@ -408,7 +408,7 @@ async def test_my_search_and_filter_returns_intersection( deployment_adapter: DeploymentAdapter, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ResourceGroupName, + scaling_group_fixture: ScalingGroupFixtureData, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """AND clause must narrow results to deployments matching every nested filter.""" @@ -416,7 +416,7 @@ async def test_my_search_and_filter_returns_intersection( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture, + scaling_group_fixture.scaling_group_name, deployment_seed_data, ["alpha", "production"], ) @@ -424,7 +424,7 @@ async def test_my_search_and_filter_returns_intersection( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture, + scaling_group_fixture.scaling_group_name, deployment_seed_data, ["beta", "production"], ) @@ -432,7 +432,7 @@ async def test_my_search_and_filter_returns_intersection( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture, + scaling_group_fixture.scaling_group_name, deployment_seed_data, ["alpha", "beta"], ) @@ -460,7 +460,7 @@ async def test_my_search_or_filter_returns_union( deployment_adapter: DeploymentAdapter, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ResourceGroupName, + scaling_group_fixture: ScalingGroupFixtureData, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """OR clause must widen results to deployments matching any nested filter.""" @@ -468,7 +468,7 @@ async def test_my_search_or_filter_returns_union( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture, + scaling_group_fixture.scaling_group_name, deployment_seed_data, ["alpha"], ) @@ -476,7 +476,7 @@ async def test_my_search_or_filter_returns_union( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture, + scaling_group_fixture.scaling_group_name, deployment_seed_data, ["beta"], ) @@ -484,7 +484,7 @@ async def test_my_search_or_filter_returns_union( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture, + scaling_group_fixture.scaling_group_name, deployment_seed_data, ["gamma"], ) @@ -512,7 +512,7 @@ async def test_my_search_or_filter_groups_multi_field_subfilters( deployment_adapter: DeploymentAdapter, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ResourceGroupName, + scaling_group_fixture: ScalingGroupFixtureData, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """OR with multi-field sub-filters must AND fields within each branch. @@ -527,7 +527,7 @@ async def test_my_search_or_filter_groups_multi_field_subfilters( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture, + scaling_group_fixture.scaling_group_name, deployment_seed_data, ["alpha"], name=f"bar-x-{suffix}", @@ -536,7 +536,7 @@ async def test_my_search_or_filter_groups_multi_field_subfilters( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture, + scaling_group_fixture.scaling_group_name, deployment_seed_data, ["beta"], name=f"foo-x-{suffix}", @@ -545,7 +545,7 @@ async def test_my_search_or_filter_groups_multi_field_subfilters( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture, + scaling_group_fixture.scaling_group_name, deployment_seed_data, ["beta"], name=f"bar-y-{suffix}", @@ -554,7 +554,7 @@ async def test_my_search_or_filter_groups_multi_field_subfilters( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture, + scaling_group_fixture.scaling_group_name, deployment_seed_data, ["alpha"], name=f"foo-y-{suffix}", @@ -589,7 +589,7 @@ async def test_project_search_and_filter_returns_intersection( deployment_adapter: DeploymentAdapter, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ResourceGroupName, + scaling_group_fixture: ScalingGroupFixtureData, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """project_search must also honor AND across nested filters.""" @@ -597,7 +597,7 @@ async def test_project_search_and_filter_returns_intersection( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture, + scaling_group_fixture.scaling_group_name, deployment_seed_data, ["alpha"], ) @@ -605,7 +605,7 @@ async def test_project_search_and_filter_returns_intersection( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture, + scaling_group_fixture.scaling_group_name, deployment_seed_data, ["alpha", "beta"], ) diff --git a/tests/component/deployment/test_deployment_lifecycle.py b/tests/component/deployment/test_deployment_lifecycle.py index 8c1c92798d6..5121363578d 100644 --- a/tests/component/deployment/test_deployment_lifecycle.py +++ b/tests/component/deployment/test_deployment_lifecycle.py @@ -37,7 +37,7 @@ from ai.backend.common.identifier.resource_group import ResourceGroupName from ai.backend.common.identifier.vfolder import VFolderUUID from ai.backend.common.types import ClusterMode -from ai.backend.testutils.fixtures import DomainFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData _RANDOM_DEPLOYMENT_ID = uuid.uuid4() _RANDOM_REVISION_ID = uuid.uuid4() @@ -91,7 +91,7 @@ async def test_create_deployment_success( admin_registry: BackendAIClientRegistry, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ResourceGroupName, + scaling_group_fixture: ScalingGroupFixtureData, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """Creating a deployment with valid config returns deployment with initial status.""" @@ -100,7 +100,7 @@ async def test_create_deployment_success( metadata=DeploymentMetadataInput( project_id=group_fixture, domain_name=domain_fixture.domain_name, - resource_group_name=scaling_group_fixture, + resource_group_name=scaling_group_fixture.scaling_group_name, name=f"test-deployment-{secrets.token_hex(4)}", ), network_access=NetworkAccessInput(open_to_public=False), @@ -158,7 +158,7 @@ async def test_update_deployment_config( admin_registry: BackendAIClientRegistry, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ResourceGroupName, + scaling_group_fixture: ScalingGroupFixtureData, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """Updating deployment config (name, replica_count) succeeds.""" @@ -168,7 +168,7 @@ async def test_update_deployment_config( metadata=DeploymentMetadataInput( project_id=group_fixture, domain_name=domain_fixture.domain_name, - resource_group_name=scaling_group_fixture, + resource_group_name=scaling_group_fixture.scaling_group_name, name=f"test-deployment-{secrets.token_hex(4)}", ), network_access=NetworkAccessInput(open_to_public=False), @@ -228,7 +228,7 @@ async def test_destroy_deployment( admin_registry: BackendAIClientRegistry, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ResourceGroupName, + scaling_group_fixture: ScalingGroupFixtureData, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """Destroying a deployment terminates it successfully.""" @@ -238,7 +238,7 @@ async def test_destroy_deployment( metadata=DeploymentMetadataInput( project_id=group_fixture, domain_name=domain_fixture.domain_name, - resource_group_name=scaling_group_fixture, + resource_group_name=scaling_group_fixture.scaling_group_name, name=f"test-deployment-{secrets.token_hex(4)}", ), network_access=NetworkAccessInput(open_to_public=False), @@ -294,7 +294,7 @@ async def test_activate_deactivate_revision( admin_registry: BackendAIClientRegistry, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ResourceGroupName, + scaling_group_fixture: ScalingGroupFixtureData, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """Adding a revision and searching revisions works correctly.""" @@ -319,7 +319,7 @@ async def test_activate_deactivate_revision( metadata=DeploymentMetadataInput( project_id=group_fixture, domain_name=domain_fixture.domain_name, - resource_group_name=scaling_group_fixture, + resource_group_name=scaling_group_fixture.scaling_group_name, name=f"test-deployment-{secrets.token_hex(4)}", ), network_access=NetworkAccessInput(open_to_public=False), @@ -375,7 +375,7 @@ async def test_change_replica_count( admin_registry: BackendAIClientRegistry, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ResourceGroupName, + scaling_group_fixture: ScalingGroupFixtureData, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """Changing replica_count updates the replica count.""" @@ -384,7 +384,7 @@ async def test_change_replica_count( metadata=DeploymentMetadataInput( project_id=group_fixture, domain_name=domain_fixture.domain_name, - resource_group_name=scaling_group_fixture, + resource_group_name=scaling_group_fixture.scaling_group_name, name=f"test-deployment-{secrets.token_hex(4)}", ), network_access=NetworkAccessInput(open_to_public=False), diff --git a/tests/component/domain/test_domain_management.py b/tests/component/domain/test_domain_management.py index e613cf22bb2..4321d79299e 100644 --- a/tests/component/domain/test_domain_management.py +++ b/tests/component/domain/test_domain_management.py @@ -26,6 +26,7 @@ from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm from ai.backend.manager.data.kernel.types import KernelStatus from ai.backend.manager.data.user.types import UserStatus +from ai.backend.manager.models.domain import domains from ai.backend.manager.models.group import groups from ai.backend.manager.models.hasher.types import PasswordInfo from ai.backend.manager.models.kernel import kernels @@ -138,27 +139,38 @@ async def test_purge_domain_with_active_kernels_raises_conflict( sa.select(groups.c.id).where(groups.c.domain_name == domain_name).limit(1) ) group_id = group_result.scalar_one() + domain_id = ( + await conn.execute(sa.select(domains.c.id).where(domains.c.name == domain_name)) + ).scalar_one() sgroup_name = f"test-sg-{secrets.token_hex(4)}" async with db_engine.begin() as conn: - await conn.execute( - sa.insert(scaling_groups).values( - name=sgroup_name, - description="Test scaling group", - is_active=True, - driver="static", - driver_opts={}, - scheduler="fifo", - scheduler_opts=ScalingGroupOpts(), + sgroup_id = ( + await conn.execute( + sa.insert(scaling_groups) + .values( + name=sgroup_name, + description="Test scaling group", + is_active=True, + driver="static", + driver_opts={}, + scheduler="fifo", + scheduler_opts=ScalingGroupOpts(), + ) + .returning(scaling_groups.c.id) ) - ) + ).scalar_one() await conn.execute( sa.insert(SessionRow.__table__).values( id=session_id, domain_name=domain_name, + # domain_id/resource_group_id only satisfy NOT NULL FKs; + # the purge check reads kernels.domain_name. + domain_id=domain_id, group_id=group_id, user_uuid=user_uuid, scaling_group_name=sgroup_name, + resource_group_id=sgroup_id, occupying_slots=ResourceSlot(), requested_slots=ResourceSlot(), priority=10, @@ -230,27 +242,38 @@ async def test_purge_validation_order_kernels_checked_first( sa.select(groups.c.id).where(groups.c.domain_name == domain_name).limit(1) ) group_id = group_result.scalar_one() + domain_id = ( + await conn.execute(sa.select(domains.c.id).where(domains.c.name == domain_name)) + ).scalar_one() sgroup_name = f"test-sg-{secrets.token_hex(4)}" async with db_engine.begin() as conn: - await conn.execute( - sa.insert(scaling_groups).values( - name=sgroup_name, - description="Test scaling group", - is_active=True, - driver="static", - driver_opts={}, - scheduler="fifo", - scheduler_opts=ScalingGroupOpts(), + sgroup_id = ( + await conn.execute( + sa.insert(scaling_groups) + .values( + name=sgroup_name, + description="Test scaling group", + is_active=True, + driver="static", + driver_opts={}, + scheduler="fifo", + scheduler_opts=ScalingGroupOpts(), + ) + .returning(scaling_groups.c.id) ) - ) + ).scalar_one() await conn.execute( sa.insert(SessionRow.__table__).values( id=session_id, domain_name=domain_name, + # domain_id/resource_group_id only satisfy NOT NULL FKs; + # the purge check reads kernels.domain_name. + domain_id=domain_id, group_id=group_id, user_uuid=user_uuid, scaling_group_name=sgroup_name, + resource_group_id=sgroup_id, occupying_slots=ResourceSlot(), requested_slots=ResourceSlot(), priority=10, diff --git a/tests/component/fair_share/conftest.py b/tests/component/fair_share/conftest.py index ee9e22fd550..2073c98539e 100644 --- a/tests/component/fair_share/conftest.py +++ b/tests/component/fair_share/conftest.py @@ -28,7 +28,7 @@ from ai.backend.manager.services.resource_usage.service import ResourceUsageService from ai.backend.manager.services.scaling_group.processors import ScalingGroupProcessors from ai.backend.manager.services.scaling_group.service import ScalingGroupService -from ai.backend.testutils.fixtures import DomainFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData @pytest.fixture() @@ -83,7 +83,7 @@ async def group_fixture( db_engine: SAEngine, domain_fixture: DomainFixtureData, resource_policy_fixture: str, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> AsyncIterator[uuid.UUID]: """Insert a test group with scaling-group association for fair-share tests.""" group_id = uuid.uuid4() @@ -101,7 +101,7 @@ async def group_fixture( ) await conn.execute( sa.insert(sgroups_for_groups).values( - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, group=group_id, ) ) diff --git a/tests/component/fair_share/test_fair_share.py b/tests/component/fair_share/test_fair_share.py index 3bf33bb066a..ef6b44baf2b 100644 --- a/tests/component/fair_share/test_fair_share.py +++ b/tests/component/fair_share/test_fair_share.py @@ -44,7 +44,7 @@ UpsertUserFairShareWeightResponse, UserWeightEntryInput, ) -from ai.backend.testutils.fixtures import DomainFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData # ---- Domain Fair Share (Global-scoped) ---- @@ -53,11 +53,11 @@ class TestGetDomainFairShare: async def test_admin_get_domain_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, ) -> None: result = await admin_registry.fair_share.get_domain_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, ) assert isinstance(result, GetDomainFairShareResponse) @@ -65,12 +65,12 @@ async def test_admin_get_domain_fair_share( async def test_user_get_domain_fair_share_forbidden( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, ) -> None: with pytest.raises(PermissionDeniedError): await user_registry.fair_share.get_domain_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, ) @@ -103,11 +103,11 @@ class TestGetProjectFairShare: async def test_admin_get_project_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, ) -> None: result = await admin_registry.fair_share.get_project_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, project_id=group_fixture, ) assert isinstance(result, GetProjectFairShareResponse) @@ -115,12 +115,12 @@ async def test_admin_get_project_fair_share( async def test_user_get_project_fair_share_forbidden( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, ) -> None: with pytest.raises(PermissionDeniedError): await user_registry.fair_share.get_project_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, project_id=group_fixture, ) @@ -144,12 +144,12 @@ class TestGetUserFairShare: async def test_admin_get_user_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, admin_user_fixture: Any, ) -> None: result = await admin_registry.fair_share.get_user_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, ) @@ -158,13 +158,13 @@ async def test_admin_get_user_fair_share( async def test_user_get_user_fair_share_forbidden( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, regular_user_fixture: Any, ) -> None: with pytest.raises(PermissionDeniedError): await user_registry.fair_share.get_user_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, project_id=group_fixture, user_uuid=regular_user_fixture.user_uuid, ) @@ -228,11 +228,11 @@ class TestRGGetDomainFairShare: async def test_admin_rg_get_domain_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, ) -> None: result = await admin_registry.fair_share.rg_get_domain_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, ) assert isinstance(result, GetDomainFairShareResponse) @@ -240,11 +240,11 @@ async def test_admin_rg_get_domain_fair_share( async def test_user_rg_get_domain_fair_share( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, ) -> None: result = await user_registry.fair_share.rg_get_domain_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, ) assert isinstance(result, GetDomainFairShareResponse) @@ -254,10 +254,10 @@ class TestRGSearchDomainFairShares: async def test_admin_rg_search_domain_fair_shares( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: result = await admin_registry.fair_share.rg_search_domain_fair_shares( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, request=SearchDomainFairSharesRequest(), ) assert isinstance(result, SearchDomainFairSharesResponse) @@ -271,12 +271,12 @@ class TestRGGetProjectFairShare: async def test_admin_rg_get_project_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, ) -> None: result = await admin_registry.fair_share.rg_get_project_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, project_id=group_fixture, ) @@ -287,11 +287,11 @@ class TestRGSearchProjectFairShares: async def test_admin_rg_search_project_fair_shares( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, ) -> None: result = await admin_registry.fair_share.rg_search_project_fair_shares( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, request=SearchProjectFairSharesRequest(), ) @@ -306,13 +306,13 @@ class TestRGGetUserFairShare: async def test_admin_rg_get_user_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: Any, ) -> None: result = await admin_registry.fair_share.rg_get_user_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, @@ -324,12 +324,12 @@ class TestRGSearchUserFairShares: async def test_admin_rg_search_user_fair_shares( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, ) -> None: result = await admin_registry.fair_share.rg_search_user_fair_shares( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, project_id=group_fixture, request=SearchUserFairSharesRequest(), @@ -345,10 +345,10 @@ class TestRGSearchDomainUsageBuckets: async def test_admin_rg_search_domain_usage_buckets( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: result = await admin_registry.fair_share.rg_search_domain_usage_buckets( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, request=SearchDomainUsageBucketsRequest(), ) assert isinstance(result, SearchDomainUsageBucketsResponse) @@ -359,11 +359,11 @@ class TestRGSearchProjectUsageBuckets: async def test_admin_rg_search_project_usage_buckets( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, ) -> None: result = await admin_registry.fair_share.rg_search_project_usage_buckets( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, request=SearchProjectUsageBucketsRequest(), ) @@ -375,12 +375,12 @@ class TestRGSearchUserUsageBuckets: async def test_admin_rg_search_user_usage_buckets( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, ) -> None: result = await admin_registry.fair_share.rg_search_user_usage_buckets( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, project_id=group_fixture, request=SearchUserUsageBucketsRequest(), @@ -396,27 +396,27 @@ class TestUpsertDomainFairShareWeight: async def test_admin_upsert_domain_weight( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, ) -> None: result = await admin_registry.fair_share.upsert_domain_fair_share_weight( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, request=UpsertDomainFairShareWeightRequest(weight=Decimal("1.5")), ) assert isinstance(result, UpsertDomainFairShareWeightResponse) assert result.item.domain_name == domain_fixture.domain_name - assert result.item.resource_group == scaling_group_fixture + assert result.item.resource_group == scaling_group_fixture.scaling_group_name async def test_user_upsert_domain_weight_forbidden( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, ) -> None: with pytest.raises(PermissionDeniedError): await user_registry.fair_share.upsert_domain_fair_share_weight( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, request=UpsertDomainFairShareWeightRequest(weight=Decimal("1.0")), ) @@ -426,12 +426,12 @@ class TestUpsertProjectFairShareWeight: async def test_admin_upsert_project_weight( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, ) -> None: result = await admin_registry.fair_share.upsert_project_fair_share_weight( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, project_id=group_fixture, request=UpsertProjectFairShareWeightRequest( domain_name=domain_fixture.domain_name, @@ -440,20 +440,20 @@ async def test_admin_upsert_project_weight( ) assert isinstance(result, UpsertProjectFairShareWeightResponse) assert result.item.project_id == group_fixture - assert result.item.resource_group == scaling_group_fixture + assert result.item.resource_group == scaling_group_fixture.scaling_group_name class TestUpsertUserFairShareWeight: async def test_admin_upsert_user_weight( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, admin_user_fixture: Any, ) -> None: result = await admin_registry.fair_share.upsert_user_fair_share_weight( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, request=UpsertUserFairShareWeightRequest( @@ -463,7 +463,7 @@ async def test_admin_upsert_user_weight( ) assert isinstance(result, UpsertUserFairShareWeightResponse) assert result.item.user_uuid == admin_user_fixture.user_uuid - assert result.item.resource_group == scaling_group_fixture + assert result.item.resource_group == scaling_group_fixture.scaling_group_name # ---- Bulk Upsert Weights (Global-scoped) ---- @@ -473,12 +473,12 @@ class TestBulkUpsertDomainFairShareWeight: async def test_admin_bulk_upsert_domain_weight( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, ) -> None: result = await admin_registry.fair_share.bulk_upsert_domain_fair_share_weight( BulkUpsertDomainFairShareWeightRequest( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, inputs=[ DomainWeightEntryInput( domain_name=domain_fixture.domain_name, @@ -495,13 +495,13 @@ class TestBulkUpsertProjectFairShareWeight: async def test_admin_bulk_upsert_project_weight( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, ) -> None: result = await admin_registry.fair_share.bulk_upsert_project_fair_share_weight( BulkUpsertProjectFairShareWeightRequest( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, inputs=[ ProjectWeightEntryInput( project_id=group_fixture, @@ -519,14 +519,14 @@ class TestBulkUpsertUserFairShareWeight: async def test_admin_bulk_upsert_user_weight( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, admin_user_fixture: Any, ) -> None: result = await admin_registry.fair_share.bulk_upsert_user_fair_share_weight( BulkUpsertUserFairShareWeightRequest( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, inputs=[ UserWeightEntryInput( user_uuid=admin_user_fixture.user_uuid, @@ -548,22 +548,22 @@ class TestGetResourceGroupFairShareSpec: async def test_admin_get_rg_spec( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: result = await admin_registry.fair_share.get_resource_group_fair_share_spec( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, ) assert isinstance(result, GetResourceGroupFairShareSpecResponse) - assert result.resource_group == scaling_group_fixture + assert result.resource_group == scaling_group_fixture.scaling_group_name async def test_user_get_rg_spec_forbidden( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: with pytest.raises(PermissionDeniedError): await user_registry.fair_share.get_resource_group_fair_share_spec( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, ) @@ -581,33 +581,33 @@ class TestUpdateResourceGroupFairShareSpec: async def test_admin_update_rg_spec( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: result = await admin_registry.fair_share.update_resource_group_fair_share_spec( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, request=UpdateResourceGroupFairShareSpecRequest( half_life_days=14, ), ) assert isinstance(result, UpdateResourceGroupFairShareSpecResponse) - assert result.resource_group == scaling_group_fixture + assert result.resource_group == scaling_group_fixture.scaling_group_name assert result.fair_share_spec.half_life_days == 14 async def test_admin_partial_update_preserves_other_fields( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: """Partial update (merge logic) changes only the specified field.""" current = await admin_registry.fair_share.get_resource_group_fair_share_spec( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, ) original_half_life = current.fair_share_spec.half_life_days original_lookback = current.fair_share_spec.lookback_days new_half_life = original_half_life + 1 result = await admin_registry.fair_share.update_resource_group_fair_share_spec( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, request=UpdateResourceGroupFairShareSpecRequest( half_life_days=new_half_life, ), @@ -618,7 +618,7 @@ async def test_admin_partial_update_preserves_other_fields( # Restore await admin_registry.fair_share.update_resource_group_fair_share_spec( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, request=UpdateResourceGroupFairShareSpecRequest( half_life_days=original_half_life, ), diff --git a/tests/component/fair_share/test_fair_share_domain_weights.py b/tests/component/fair_share/test_fair_share_domain_weights.py index fcb88b4dc36..646f5001e7a 100644 --- a/tests/component/fair_share/test_fair_share_domain_weights.py +++ b/tests/component/fair_share/test_fair_share_domain_weights.py @@ -16,7 +16,7 @@ UpsertDomainFairShareWeightRequest, UpsertDomainFairShareWeightResponse, ) -from ai.backend.testutils.fixtures import DomainFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData class TestDomainFairShareWeights: @@ -25,18 +25,18 @@ class TestDomainFairShareWeights: async def test_get_domain_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, ) -> None: """Get domain fair share → returns weight data.""" result = await admin_registry.fair_share.get_domain_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, ) assert isinstance(result, GetDomainFairShareResponse) assert result.item is not None assert result.item.domain_name == domain_fixture.domain_name - assert result.item.resource_group == scaling_group_fixture + assert result.item.resource_group == scaling_group_fixture.scaling_group_name async def test_search_domain_fair_shares( self, @@ -52,24 +52,24 @@ async def test_search_domain_fair_shares( async def test_upsert_domain_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, ) -> None: """Upsert domain fair share → weight created/updated.""" weight = Decimal("2.5") result = await admin_registry.fair_share.upsert_domain_fair_share_weight( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, request=UpsertDomainFairShareWeightRequest(weight=weight), ) assert isinstance(result, UpsertDomainFairShareWeightResponse) assert result.item.domain_name == domain_fixture.domain_name - assert result.item.resource_group == scaling_group_fixture + assert result.item.resource_group == scaling_group_fixture.scaling_group_name assert result.item.spec.weight == weight # Verify the weight persists get_result = await admin_registry.fair_share.get_domain_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, ) assert get_result.item is not None @@ -82,13 +82,13 @@ class TestBulkUpsertDomainWeights: async def test_bulk_upsert_success( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, ) -> None: """Bulk upsert success → all weights updated.""" result = await admin_registry.fair_share.bulk_upsert_domain_fair_share_weight( BulkUpsertDomainFairShareWeightRequest( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, inputs=[ DomainWeightEntryInput( domain_name=domain_fixture.domain_name, @@ -103,12 +103,12 @@ async def test_bulk_upsert_success( async def test_bulk_upsert_empty_input( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: """Bulk upsert empty input → empty result (no error).""" result = await admin_registry.fair_share.bulk_upsert_domain_fair_share_weight( BulkUpsertDomainFairShareWeightRequest( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, inputs=[], ), ) @@ -118,12 +118,12 @@ async def test_bulk_upsert_empty_input( async def test_bulk_upsert_overwrite( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, ) -> None: """Bulk upsert overwrites existing weight.""" await admin_registry.fair_share.upsert_domain_fair_share_weight( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, request=UpsertDomainFairShareWeightRequest(weight=Decimal("10.0")), ) @@ -131,7 +131,7 @@ async def test_bulk_upsert_overwrite( new_weight = Decimal("3.0") result = await admin_registry.fair_share.bulk_upsert_domain_fair_share_weight( BulkUpsertDomainFairShareWeightRequest( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, inputs=[ DomainWeightEntryInput( domain_name=domain_fixture.domain_name, @@ -144,7 +144,7 @@ async def test_bulk_upsert_overwrite( assert result.upserted_count == 1 get_result = await admin_registry.fair_share.get_domain_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, ) assert get_result.item is not None @@ -157,13 +157,13 @@ class TestDomainScopeAccessControl: async def test_global_scope_regular_user_denied( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, ) -> None: """Global-scoped domain access as regular user → 403 (denied).""" with pytest.raises(PermissionDeniedError): await user_registry.fair_share.get_domain_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, ) @@ -174,7 +174,7 @@ async def test_global_scope_regular_user_denied( with pytest.raises(PermissionDeniedError): await user_registry.fair_share.upsert_domain_fair_share_weight( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, request=UpsertDomainFairShareWeightRequest(weight=Decimal("1.0")), ) @@ -182,18 +182,18 @@ async def test_global_scope_regular_user_denied( async def test_rg_scope_regular_user_allowed( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, ) -> None: """RG-scoped domain access as regular user → 200 (allowed).""" get_result = await user_registry.fair_share.rg_get_domain_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, ) assert isinstance(get_result, GetDomainFairShareResponse) search_result = await user_registry.fair_share.rg_search_domain_fair_shares( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, request=SearchDomainFairSharesRequest(), ) assert isinstance(search_result, SearchDomainFairSharesResponse) @@ -201,12 +201,12 @@ async def test_rg_scope_regular_user_allowed( async def test_global_scope_admin_allowed( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, ) -> None: """Admin global scope access → 200 (allowed).""" result = await admin_registry.fair_share.get_domain_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, ) assert isinstance(result, GetDomainFairShareResponse) @@ -218,12 +218,12 @@ class TestDomainDefaultValueFallback: async def test_get_domain_without_fair_share_default_value( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, ) -> None: """Get existing domain with no fair-share row → default value returned.""" result = await admin_registry.fair_share.get_domain_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, ) assert isinstance(result, GetDomainFairShareResponse) diff --git a/tests/component/fair_share/test_fair_share_project_weights.py b/tests/component/fair_share/test_fair_share_project_weights.py index bb5cc24d105..ee94daf9577 100644 --- a/tests/component/fair_share/test_fair_share_project_weights.py +++ b/tests/component/fair_share/test_fair_share_project_weights.py @@ -17,7 +17,7 @@ UpsertProjectFairShareWeightRequest, UpsertProjectFairShareWeightResponse, ) -from ai.backend.testutils.fixtures import DomainFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData class TestProjectFairShareWeights: @@ -26,18 +26,18 @@ class TestProjectFairShareWeights: async def test_get_project_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, ) -> None: """Get project fair share → returns weight data.""" result = await admin_registry.fair_share.get_project_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, project_id=group_fixture, ) assert isinstance(result, GetProjectFairShareResponse) assert result.item is not None assert result.item.project_id == group_fixture - assert result.item.resource_group == scaling_group_fixture + assert result.item.resource_group == scaling_group_fixture.scaling_group_name async def test_search_project_fair_shares( self, @@ -53,14 +53,14 @@ async def test_search_project_fair_shares( async def test_upsert_project_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, ) -> None: """Upsert project fair share → weight created/updated.""" weight = Decimal("3.5") result = await admin_registry.fair_share.upsert_project_fair_share_weight( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, project_id=group_fixture, request=UpsertProjectFairShareWeightRequest( domain_name=domain_fixture.domain_name, @@ -69,12 +69,12 @@ async def test_upsert_project_fair_share( ) assert isinstance(result, UpsertProjectFairShareWeightResponse) assert result.item.project_id == group_fixture - assert result.item.resource_group == scaling_group_fixture + assert result.item.resource_group == scaling_group_fixture.scaling_group_name assert result.item.spec.weight == weight # Verify the weight persists get_result = await admin_registry.fair_share.get_project_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, project_id=group_fixture, ) assert get_result.item is not None @@ -87,14 +87,14 @@ class TestBulkUpsertProjectWeights: async def test_bulk_upsert_success( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, ) -> None: """Bulk upsert success → all weights updated.""" result = await admin_registry.fair_share.bulk_upsert_project_fair_share_weight( BulkUpsertProjectFairShareWeightRequest( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, inputs=[ ProjectWeightEntryInput( project_id=group_fixture, @@ -110,12 +110,12 @@ async def test_bulk_upsert_success( async def test_bulk_upsert_empty_input( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: """Bulk upsert empty input → empty result (no error).""" result = await admin_registry.fair_share.bulk_upsert_project_fair_share_weight( BulkUpsertProjectFairShareWeightRequest( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, inputs=[], ), ) @@ -125,13 +125,13 @@ async def test_bulk_upsert_empty_input( async def test_bulk_upsert_overwrite( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, ) -> None: """Bulk upsert overwrites existing weight.""" await admin_registry.fair_share.upsert_project_fair_share_weight( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, project_id=group_fixture, request=UpsertProjectFairShareWeightRequest( domain_name=domain_fixture.domain_name, @@ -142,7 +142,7 @@ async def test_bulk_upsert_overwrite( new_weight = Decimal("3.0") result = await admin_registry.fair_share.bulk_upsert_project_fair_share_weight( BulkUpsertProjectFairShareWeightRequest( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, inputs=[ ProjectWeightEntryInput( project_id=group_fixture, @@ -156,7 +156,7 @@ async def test_bulk_upsert_overwrite( assert result.upserted_count == 1 get_result = await admin_registry.fair_share.get_project_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, project_id=group_fixture, ) assert get_result.item is not None @@ -169,13 +169,13 @@ class TestProjectScopeAccessControl: async def test_global_scope_regular_user_denied( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, ) -> None: """Global-scoped project access as regular user → 403 (denied).""" with pytest.raises(PermissionDeniedError): await user_registry.fair_share.get_project_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, project_id=group_fixture, ) @@ -187,20 +187,20 @@ async def test_global_scope_regular_user_denied( async def test_rg_scope_regular_user_allowed( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, ) -> None: """RG-scoped project access as regular user → 200 (allowed).""" get_result = await user_registry.fair_share.rg_get_project_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, project_id=group_fixture, ) assert isinstance(get_result, GetProjectFairShareResponse) search_result = await user_registry.fair_share.rg_search_project_fair_shares( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, request=SearchProjectFairSharesRequest(), ) @@ -213,12 +213,12 @@ class TestProjectDefaultValueFallback: async def test_get_project_without_fair_share_default_value( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, ) -> None: """Get existing project with no fair-share row → default value returned.""" result = await admin_registry.fair_share.get_project_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, project_id=group_fixture, ) assert isinstance(result, GetProjectFairShareResponse) diff --git a/tests/component/fair_share/test_fair_share_user_weights.py b/tests/component/fair_share/test_fair_share_user_weights.py index 74525fba730..f2ac07eb394 100644 --- a/tests/component/fair_share/test_fair_share_user_weights.py +++ b/tests/component/fair_share/test_fair_share_user_weights.py @@ -18,7 +18,7 @@ UpsertUserFairShareWeightResponse, UserWeightEntryInput, ) -from ai.backend.testutils.fixtures import DomainFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData class TestUserFairShareWeights: @@ -27,20 +27,20 @@ class TestUserFairShareWeights: async def test_get_user_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, admin_user_fixture: Any, ) -> None: """Get user fair share → returns weight data.""" result = await admin_registry.fair_share.get_user_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, ) assert isinstance(result, GetUserFairShareResponse) assert result.item is not None assert result.item.user_uuid == admin_user_fixture.user_uuid - assert result.item.resource_group == scaling_group_fixture + assert result.item.resource_group == scaling_group_fixture.scaling_group_name async def test_search_user_fair_shares( self, @@ -56,7 +56,7 @@ async def test_search_user_fair_shares( async def test_upsert_user_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, admin_user_fixture: Any, @@ -64,7 +64,7 @@ async def test_upsert_user_fair_share( """Upsert user fair share → weight created/updated.""" weight = Decimal("4.5") result = await admin_registry.fair_share.upsert_user_fair_share_weight( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, request=UpsertUserFairShareWeightRequest( @@ -74,12 +74,12 @@ async def test_upsert_user_fair_share( ) assert isinstance(result, UpsertUserFairShareWeightResponse) assert result.item.user_uuid == admin_user_fixture.user_uuid - assert result.item.resource_group == scaling_group_fixture + assert result.item.resource_group == scaling_group_fixture.scaling_group_name assert result.item.spec.weight == weight # Verify the weight persists get_result = await admin_registry.fair_share.get_user_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, ) @@ -93,7 +93,7 @@ class TestBulkUpsertUserWeights: async def test_bulk_upsert_success( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, admin_user_fixture: Any, @@ -101,7 +101,7 @@ async def test_bulk_upsert_success( """Bulk upsert success → all weights updated.""" result = await admin_registry.fair_share.bulk_upsert_user_fair_share_weight( BulkUpsertUserFairShareWeightRequest( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, inputs=[ UserWeightEntryInput( user_uuid=admin_user_fixture.user_uuid, @@ -118,12 +118,12 @@ async def test_bulk_upsert_success( async def test_bulk_upsert_empty_input( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: """Bulk upsert empty input → empty result (no error).""" result = await admin_registry.fair_share.bulk_upsert_user_fair_share_weight( BulkUpsertUserFairShareWeightRequest( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, inputs=[], ), ) @@ -133,14 +133,14 @@ async def test_bulk_upsert_empty_input( async def test_bulk_upsert_overwrite( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, admin_user_fixture: Any, ) -> None: """Bulk upsert overwrites existing weight.""" await admin_registry.fair_share.upsert_user_fair_share_weight( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, request=UpsertUserFairShareWeightRequest( @@ -152,7 +152,7 @@ async def test_bulk_upsert_overwrite( new_weight = Decimal("3.0") result = await admin_registry.fair_share.bulk_upsert_user_fair_share_weight( BulkUpsertUserFairShareWeightRequest( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, inputs=[ UserWeightEntryInput( user_uuid=admin_user_fixture.user_uuid, @@ -167,7 +167,7 @@ async def test_bulk_upsert_overwrite( assert result.upserted_count == 1 get_result = await admin_registry.fair_share.get_user_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, ) @@ -181,14 +181,14 @@ class TestUserScopeAccessControl: async def test_global_scope_regular_user_denied( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, admin_user_fixture: Any, ) -> None: """Global-scoped user access as regular user → 403 (denied).""" with pytest.raises(PermissionDeniedError): await user_registry.fair_share.get_user_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, ) @@ -201,14 +201,14 @@ async def test_global_scope_regular_user_denied( async def test_rg_scope_regular_user_allowed( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: Any, ) -> None: """RG-scoped user access as regular user → 200 (allowed).""" get_result = await user_registry.fair_share.rg_get_user_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, @@ -216,7 +216,7 @@ async def test_rg_scope_regular_user_allowed( assert isinstance(get_result, GetUserFairShareResponse) search_result = await user_registry.fair_share.rg_search_user_fair_shares( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, domain_name=domain_fixture.domain_name, project_id=group_fixture, request=SearchUserFairSharesRequest(), @@ -230,13 +230,13 @@ class TestUserDefaultValueFallback: async def test_get_user_without_fair_share_default_value( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, admin_user_fixture: Any, ) -> None: """Get existing user with no fair-share row → default value returned.""" result = await admin_registry.fair_share.get_user_fair_share( - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, ) diff --git a/tests/component/group/test_group_management.py b/tests/component/group/test_group_management.py index bab39a94bab..65d78663f8c 100644 --- a/tests/component/group/test_group_management.py +++ b/tests/component/group/test_group_management.py @@ -49,7 +49,7 @@ from ai.backend.manager.models.kernel import KernelRow, KernelStatus from ai.backend.manager.models.session import SessionRow, SessionStatus from ai.backend.manager.models.vfolder import VFolderRow -from ai.backend.testutils.fixtures import DomainFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData @pytest.fixture() @@ -82,7 +82,7 @@ async def test_group_for_deletion( async def group_with_vfolder_mounted( db_engine: SAEngine, domain_fixture: DomainFixtureData, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, resource_policy_fixture: str, regular_user_fixture: Any, ) -> AsyncIterator[uuid.UUID]: @@ -131,7 +131,7 @@ async def group_with_vfolder_mounted( domain_name=domain_fixture.domain_name, group_id=group_id, user_uuid=user_uuid, - scaling_group_name=scaling_group_fixture, + scaling_group_name=scaling_group_fixture.scaling_group_name, status=SessionStatus.RUNNING, occupying_slots=ResourceSlot(), requested_slots=ResourceSlot(), @@ -153,7 +153,7 @@ async def group_with_vfolder_mounted( group_id=group_id, user_uuid=user_uuid, domain_name=domain_fixture.domain_name, - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, status=KernelStatus.RUNNING, image="python:3.9", occupied_slots=ResourceSlot({}), @@ -186,7 +186,7 @@ async def group_with_vfolder_mounted( async def group_with_active_kernel( db_engine: SAEngine, domain_fixture: DomainFixtureData, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, resource_policy_fixture: str, regular_user_fixture: Any, ) -> AsyncIterator[uuid.UUID]: @@ -222,7 +222,7 @@ async def group_with_active_kernel( domain_name=domain_fixture.domain_name, group_id=group_id, user_uuid=user_uuid, - scaling_group_name=scaling_group_fixture, + scaling_group_name=scaling_group_fixture.scaling_group_name, status=SessionStatus.RUNNING, occupying_slots=ResourceSlot(), requested_slots=ResourceSlot(), @@ -244,7 +244,7 @@ async def group_with_active_kernel( group_id=group_id, user_uuid=user_uuid, domain_name=domain_fixture.domain_name, - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, status=KernelStatus.RUNNING, image="python:3.9", occupied_slots=ResourceSlot({}), @@ -275,7 +275,7 @@ async def group_with_active_endpoint( domain_fixture: DomainFixtureData, resource_policy_fixture: str, regular_user_fixture: Any, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> AsyncIterator[uuid.UUID]: """Create group with active endpoint.""" group_id = uuid.uuid4() @@ -335,7 +335,7 @@ async def group_with_active_endpoint( domain=domain_fixture.domain_name, created_user=user_uuid, session_owner=user_uuid, - resource_group=scaling_group_fixture, + resource_group=scaling_group_fixture.scaling_group_name, image=image_id, resource_slots=ResourceSlot({"cpu": "1", "mem": "1073741824"}), lifecycle_stage=EndpointLifecycle.CREATED, diff --git a/tests/component/infra/test_infra.py b/tests/component/infra/test_infra.py index 7546cbb7fca..4e76726c6b5 100644 --- a/tests/component/infra/test_infra.py +++ b/tests/component/infra/test_infra.py @@ -29,6 +29,7 @@ UsagePerPeriodRequest, UsagePerPeriodResponse, ) +from ai.backend.testutils.fixtures import ScalingGroupFixtureData class TestEtcdConfigRead: @@ -229,11 +230,11 @@ async def test_admin_checks_presets( async def test_list_presets_with_scaling_group_filter( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: """Filtering presets by scaling group.""" await admin_registry.infra.list_presets( - ListPresetsRequest(scaling_group=scaling_group_fixture) + ListPresetsRequest(scaling_group=scaling_group_fixture.scaling_group_name) ) diff --git a/tests/component/operations/conftest.py b/tests/component/operations/conftest.py index fe38cda72ca..557addb9512 100644 --- a/tests/component/operations/conftest.py +++ b/tests/component/operations/conftest.py @@ -30,6 +30,7 @@ from ai.backend.manager.services.error_log.service import ErrorLogService from ai.backend.manager.services.manager_admin.processors import ManagerAdminProcessors from ai.backend.manager.services.manager_admin.service import ManagerAdminService +from ai.backend.testutils.fixtures import ScalingGroupFixtureData @pytest.fixture() @@ -96,7 +97,7 @@ def server_module_registries( @pytest.fixture() async def agent_fixture( db_engine: SAEngine, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> AsyncIterator[str]: """Insert a test agent record and yield its ID.""" agent_id = f"i-test-agent-{secrets.token_hex(4)}" @@ -105,7 +106,7 @@ async def agent_fixture( sa.insert(agents).values( id=agent_id, region="local", - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, available_slots=ResourceSlot(), occupied_slots=ResourceSlot(), addr="127.0.0.1:6001", diff --git a/tests/component/resource_allocation/test_resource_allocation.py b/tests/component/resource_allocation/test_resource_allocation.py index 8c23590a8a6..6c6d63579e7 100644 --- a/tests/component/resource_allocation/test_resource_allocation.py +++ b/tests/component/resource_allocation/test_resource_allocation.py @@ -19,7 +19,7 @@ ProjectResourceAllocationPayload, ResourceGroupResourceAllocationPayload, ) -from ai.backend.testutils.fixtures import DomainFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData class TestKeypairUsage: @@ -102,11 +102,11 @@ class TestResourceGroupUsage: async def test_resource_group_usage_returns_structure( self, admin_v2_registry: V2ClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: """Query resource group usage returns capacity, used, free, max_per_node fields.""" result = await admin_v2_registry.resource_allocation.resource_group_usage( - scaling_group_fixture, + scaling_group_fixture.scaling_group_name, ) assert isinstance(result, ResourceGroupResourceAllocationPayload) assert result.resource_group is not None @@ -134,13 +134,13 @@ async def test_effective_returns_assignable_and_breakdown( self, admin_v2_registry: V2ClientRegistry, group_fixture: uuid.UUID, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: """Effective allocation returns assignable list and breakdown structure.""" result = await admin_v2_registry.resource_allocation.effective( EffectiveResourceAllocationInput( project_id=group_fixture, - resource_group_name=scaling_group_fixture, + resource_group_name=scaling_group_fixture.scaling_group_name, ), ) assert isinstance(result, EffectiveResourceAllocationPayload) @@ -152,7 +152,7 @@ async def test_effective_returns_assignable_and_breakdown( async def test_effective_with_invalid_project_returns_error( self, admin_v2_registry: V2ClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: """Effective allocation with non-existent project returns an error.""" fake_id = uuid.uuid4() @@ -160,7 +160,7 @@ async def test_effective_with_invalid_project_returns_error( await admin_v2_registry.resource_allocation.effective( EffectiveResourceAllocationInput( project_id=fake_id, - resource_group_name=scaling_group_fixture, + resource_group_name=scaling_group_fixture.scaling_group_name, ), ) assert exc_info.value.status in (404, 500) @@ -173,13 +173,13 @@ async def test_preset_availability_returns_list( self, admin_v2_registry: V2ClientRegistry, group_fixture: uuid.UUID, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: """Check preset availability returns a list of presets with availability status.""" result = await admin_v2_registry.resource_allocation.check_availability( CheckPresetAvailabilityInput( project_id=group_fixture, - resource_group_name=scaling_group_fixture, + resource_group_name=scaling_group_fixture.scaling_group_name, ), ) assert isinstance(result.presets, list) diff --git a/tests/component/resource_allocation/test_resource_allocation_visibility.py b/tests/component/resource_allocation/test_resource_allocation_visibility.py index 84983db1794..145e0d9deee 100644 --- a/tests/component/resource_allocation/test_resource_allocation_visibility.py +++ b/tests/component/resource_allocation/test_resource_allocation_visibility.py @@ -51,6 +51,7 @@ from ai.backend.manager.services.resource_allocation.service import ( ResourceAllocationService, ) +from ai.backend.testutils.fixtures import ScalingGroupFixtureData if TYPE_CHECKING: from tests.component.conftest import ServerInfo, UserFixtureData @@ -204,13 +205,13 @@ async def test_regular_user_effective_hides_resource_group( self, user_v2_registry: V2ClientRegistry, group_fixture: uuid.UUID, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: """Regular user with hide_agents=True should get resource_group=null.""" result = await user_v2_registry.resource_allocation.effective( EffectiveResourceAllocationInput( project_id=group_fixture, - resource_group_name=scaling_group_fixture, + resource_group_name=scaling_group_fixture.scaling_group_name, ), ) assert isinstance(result, EffectiveResourceAllocationPayload) @@ -223,13 +224,13 @@ async def test_admin_effective_sees_resource_group_despite_hide_agents( self, admin_v2_registry: V2ClientRegistry, group_fixture: uuid.UUID, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: """Admin should see resource_group even when hide_agents=True.""" result = await admin_v2_registry.resource_allocation.effective( EffectiveResourceAllocationInput( project_id=group_fixture, - resource_group_name=scaling_group_fixture, + resource_group_name=scaling_group_fixture.scaling_group_name, ), ) assert isinstance(result, EffectiveResourceAllocationPayload) @@ -334,13 +335,13 @@ async def test_regular_user_effective_hides_project( self, user_v2_registry: V2ClientRegistry, group_fixture: uuid.UUID, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: """Regular user with group_resource_visibility=False should get project=null.""" result = await user_v2_registry.resource_allocation.effective( EffectiveResourceAllocationInput( project_id=group_fixture, - resource_group_name=scaling_group_fixture, + resource_group_name=scaling_group_fixture.scaling_group_name, ), ) assert isinstance(result, EffectiveResourceAllocationPayload) @@ -353,13 +354,13 @@ async def test_admin_effective_sees_project_despite_grv_off( self, admin_v2_registry: V2ClientRegistry, group_fixture: uuid.UUID, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: """Admin should see project even when group_resource_visibility=False.""" result = await admin_v2_registry.resource_allocation.effective( EffectiveResourceAllocationInput( project_id=group_fixture, - resource_group_name=scaling_group_fixture, + resource_group_name=scaling_group_fixture.scaling_group_name, ), ) assert isinstance(result, EffectiveResourceAllocationPayload) diff --git a/tests/component/resource_preset/test_resource_preset.py b/tests/component/resource_preset/test_resource_preset.py index 61b33807754..2891e0acc7f 100644 --- a/tests/component/resource_preset/test_resource_preset.py +++ b/tests/component/resource_preset/test_resource_preset.py @@ -16,6 +16,7 @@ ListPresetsRequest, ListPresetsResponse, ) +from ai.backend.testutils.fixtures import ScalingGroupFixtureData PresetFixtureData = dict[str, Any] PresetFactory = Callable[..., Coroutine[Any, Any, PresetFixtureData]] @@ -53,12 +54,12 @@ async def test_admin_lists_presets_empty( async def test_admin_lists_presets_with_scaling_group_filter( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, target_preset: PresetFixtureData, ) -> None: """Filtering by scaling group still returns system presets.""" result = await admin_registry.infra.list_presets( - ListPresetsRequest(scaling_group=scaling_group_fixture) + ListPresetsRequest(scaling_group=scaling_group_fixture.scaling_group_name) ) assert isinstance(result, ListPresetsResponse) assert isinstance(result.presets, list) @@ -93,13 +94,13 @@ async def test_admin_checks_presets_with_scaling_group( admin_registry: BackendAIClientRegistry, target_preset: PresetFixtureData, group_name_fixture: str, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: """Admin checks presets filtered by scaling group.""" result = await admin_registry.infra.check_presets( CheckPresetsRequest( group=group_name_fixture, - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, ) ) assert isinstance(result, CheckPresetsResponse) diff --git a/tests/component/resource_preset/test_resource_preset_full.py b/tests/component/resource_preset/test_resource_preset_full.py index 717b95f34b4..cf5c6a9508f 100644 --- a/tests/component/resource_preset/test_resource_preset_full.py +++ b/tests/component/resource_preset/test_resource_preset_full.py @@ -48,6 +48,7 @@ ) from ai.backend.manager.services.resource_preset.processors import ResourcePresetProcessors from ai.backend.manager.types import OptionalState +from ai.backend.testutils.fixtures import ScalingGroupFixtureData PresetFixtureData = dict[str, Any] @@ -454,7 +455,7 @@ async def test_s6_check_presets_with_scaling_group_filter( resource_preset_processors: ResourcePresetProcessors, admin_registry: BackendAIClientRegistry, group_name_fixture: str, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, database_fixture: None, ) -> None: """S-6: check_presets with scaling_group filter returns only that SG.""" @@ -466,14 +467,14 @@ async def test_s6_check_presets_with_scaling_group_filter( result = await admin_registry.infra.check_presets( CheckPresetsRequest( group=group_name_fixture, - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, ) ) assert isinstance(result, CheckPresetsResponse) assert isinstance(result.presets, list) # When filtered, scaling_groups should contain only the specified SG if result.scaling_groups: - assert scaling_group_fixture in result.scaling_groups + assert scaling_group_fixture.scaling_group_name in result.scaling_groups assert len(result.scaling_groups) == 1 async def test_s7_check_presets_no_sessions_zero_usage( diff --git a/tests/component/scaling_group/test_scaling_group.py b/tests/component/scaling_group/test_scaling_group.py index 39a9920a7b9..5c5afb4fc4c 100644 --- a/tests/component/scaling_group/test_scaling_group.py +++ b/tests/component/scaling_group/test_scaling_group.py @@ -7,13 +7,14 @@ from ai.backend.client.v2.exceptions import NotFoundError from ai.backend.client.v2.registry import BackendAIClientRegistry from ai.backend.common.dto.manager.scaling_group import ListScalingGroupsResponse +from ai.backend.testutils.fixtures import ScalingGroupFixtureData class TestScalingGroupList: async def test_admin_lists_scaling_groups( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, ) -> None: """Admin can list scaling groups; the fixture sgroup is visible via domain association.""" @@ -22,12 +23,12 @@ async def test_admin_lists_scaling_groups( ) assert isinstance(result, ListScalingGroupsResponse) names = [sg.name for sg in result.scaling_groups] - assert scaling_group_fixture in names + assert scaling_group_fixture.scaling_group_name in names async def test_list_scaling_groups_with_group_filter( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, ) -> None: """Filtering by group UUID returns the domain-associated scaling group.""" @@ -35,12 +36,14 @@ async def test_list_scaling_groups_with_group_filter( group=str(group_fixture), ) assert isinstance(result, ListScalingGroupsResponse) - assert any(sg.name == scaling_group_fixture for sg in result.scaling_groups) + assert any( + sg.name == scaling_group_fixture.scaling_group_name for sg in result.scaling_groups + ) async def test_regular_user_lists_public_scaling_groups( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, ) -> None: """Regular users can list scaling groups; only public ones are returned. @@ -56,12 +59,12 @@ async def test_regular_user_lists_public_scaling_groups( assert isinstance(result, ListScalingGroupsResponse) # The fixture sgroup defaults to is_public=True, so it is visible. names = [sg.name for sg in result.scaling_groups] - assert scaling_group_fixture in names + assert scaling_group_fixture.scaling_group_name in names async def test_list_scaling_groups_with_uuid_string_parameter( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, ) -> None: """BA-5411: Passing group UUID as a string correctly resolves the group.""" @@ -70,14 +73,16 @@ async def test_list_scaling_groups_with_uuid_string_parameter( group=str(group_fixture), ) assert isinstance(result, ListScalingGroupsResponse) - assert any(sg.name == scaling_group_fixture for sg in result.scaling_groups) + assert any( + sg.name == scaling_group_fixture.scaling_group_name for sg in result.scaling_groups + ) class TestScalingGroupWsproxyVersion: async def test_admin_gets_wsproxy_version( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: """Admin requests wsproxy version for a scaling group without wsproxy_addr configured. @@ -87,7 +92,7 @@ async def test_admin_gets_wsproxy_version( """ with pytest.raises(NotFoundError): await admin_registry.scaling_group.get_wsproxy_version( - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, ) async def test_get_wsproxy_version_nonexistent( @@ -103,7 +108,7 @@ async def test_get_wsproxy_version_nonexistent( async def test_regular_user_gets_wsproxy_version_not_found( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: """Regular users are not blocked from querying wsproxy version. @@ -113,5 +118,5 @@ async def test_regular_user_gets_wsproxy_version_not_found( """ with pytest.raises(NotFoundError): await user_registry.scaling_group.get_wsproxy_version( - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, ) diff --git a/tests/component/scaling_group/test_scaling_group_associations.py b/tests/component/scaling_group/test_scaling_group_associations.py index dc0ed6260a0..2908a54f10b 100644 --- a/tests/component/scaling_group/test_scaling_group_associations.py +++ b/tests/component/scaling_group/test_scaling_group_associations.py @@ -28,7 +28,7 @@ scaling_groups, sgroups_for_domains, ) -from ai.backend.testutils.fixtures import DomainFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData @pytest.fixture() @@ -78,7 +78,7 @@ async def test_s_visibility_public_sgroup_visible_to_user( self, admin_registry: BackendAIClientRegistry, user_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, ) -> None: """Public scaling group is visible to both admin and regular user.""" @@ -86,13 +86,18 @@ async def test_s_visibility_public_sgroup_visible_to_user( group=str(group_fixture), ) assert isinstance(admin_result, ListScalingGroupsResponse) - assert any(sg.name == scaling_group_fixture for sg in admin_result.scaling_groups) + assert any( + sg.name == scaling_group_fixture.scaling_group_name + for sg in admin_result.scaling_groups + ) user_result = await user_registry.scaling_group.list_scaling_groups( group=str(group_fixture), ) assert isinstance(user_result, ListScalingGroupsResponse) - assert any(sg.name == scaling_group_fixture for sg in user_result.scaling_groups) + assert any( + sg.name == scaling_group_fixture.scaling_group_name for sg in user_result.scaling_groups + ) async def test_s_visibility_private_sgroup_hidden_from_user( self, @@ -140,7 +145,7 @@ class TestScalingGroupPermissions: async def test_regular_user_can_list_scaling_groups( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, ) -> None: """F-AUTH-LIST: Regular user can list public scaling groups (auth_required, not superadmin). @@ -155,4 +160,4 @@ async def test_regular_user_can_list_scaling_groups( assert isinstance(result, ListScalingGroupsResponse) names = [sg.name for sg in result.scaling_groups] # The fixture sgroup is public — regular user should see it - assert scaling_group_fixture in names + assert scaling_group_fixture.scaling_group_name in names diff --git a/tests/component/scaling_group/test_scaling_group_crud.py b/tests/component/scaling_group/test_scaling_group_crud.py index 03c7466b4d2..b92f4434f10 100644 --- a/tests/component/scaling_group/test_scaling_group_crud.py +++ b/tests/component/scaling_group/test_scaling_group_crud.py @@ -17,7 +17,7 @@ sgroups_for_domains, sgroups_for_groups, ) -from ai.backend.testutils.fixtures import DomainFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData @pytest.fixture() @@ -104,7 +104,7 @@ class TestScalingGroupCRUDLifecycle: async def test_multiple_scaling_groups_listed( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, extra_scaling_group_fixture: str, group_fixture: uuid.UUID, ) -> None: @@ -114,13 +114,13 @@ async def test_multiple_scaling_groups_listed( ) assert isinstance(result, ListScalingGroupsResponse) names = [sg.name for sg in result.scaling_groups] - assert scaling_group_fixture in names + assert scaling_group_fixture.scaling_group_name in names assert extra_scaling_group_fixture in names async def test_scaling_group_domain_association( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, group_fixture: uuid.UUID, ) -> None: """A scaling group associated with the test domain is visible via group filter.""" @@ -128,12 +128,14 @@ async def test_scaling_group_domain_association( group=str(group_fixture), ) assert isinstance(result, ListScalingGroupsResponse) - assert any(sg.name == scaling_group_fixture for sg in result.scaling_groups) + assert any( + sg.name == scaling_group_fixture.scaling_group_name for sg in result.scaling_groups + ) async def test_regular_user_sees_public_scaling_groups( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, private_scaling_group_fixture: str, group_fixture: uuid.UUID, ) -> None: @@ -148,7 +150,7 @@ async def test_regular_user_sees_public_scaling_groups( ) assert isinstance(result, ListScalingGroupsResponse) names = [sg.name for sg in result.scaling_groups] - assert scaling_group_fixture in names + assert scaling_group_fixture.scaling_group_name in names assert private_scaling_group_fixture not in names async def test_admin_sees_private_scaling_groups( diff --git a/tests/component/session/conftest.py b/tests/component/session/conftest.py index ef657f9903c..876dfe3b2f5 100644 --- a/tests/component/session/conftest.py +++ b/tests/component/session/conftest.py @@ -39,7 +39,7 @@ from ai.backend.manager.services.session.processors import SessionProcessors from ai.backend.manager.services.session.service import SessionService, SessionServiceArgs from ai.backend.manager.services.vfolder.processors.vfolder import VFolderProcessors -from ai.backend.testutils.fixtures import DomainFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData @dataclass @@ -164,7 +164,7 @@ def server_module_registries( @pytest.fixture() async def agent_fixture( db_engine: SAEngine, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> AsyncIterator[AgentId]: """Insert a test agent record and yield its ID.""" agent_id = AgentId(f"i-test-agent-{secrets.token_hex(4)}") @@ -173,7 +173,7 @@ async def agent_fixture( sa.insert(AgentRow).values( id=agent_id, region="local", - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, available_slots=ResourceSlot(), occupied_slots=ResourceSlot(), addr="127.0.0.1:6001", @@ -192,7 +192,7 @@ async def session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, agent_fixture: AgentId, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING session + kernel directly in the database. @@ -223,10 +223,12 @@ async def session_seed( cluster_size=1, cluster_mode="single-node", domain_name=domain_fixture.domain_name, + domain_id=domain_fixture.domain_id, group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture, + scaling_group_name=scaling_group_fixture.scaling_group_name, + resource_group_id=scaling_group_fixture.scaling_group_id, status=SessionStatus.RUNNING, status_info="", status_history=status_history, @@ -251,7 +253,7 @@ async def session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, agent=agent_fixture, status=KernelStatus.RUNNING, status_info="", @@ -287,7 +289,7 @@ async def terminated_session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> AsyncIterator[SessionSeedData]: """Seed a TERMINATED session with container_log in the kernel. @@ -316,10 +318,12 @@ async def terminated_session_seed( cluster_size=1, cluster_mode="single-node", domain_name=domain_fixture.domain_name, + domain_id=domain_fixture.domain_id, group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture, + scaling_group_name=scaling_group_fixture.scaling_group_name, + resource_group_id=scaling_group_fixture.scaling_group_id, status=SessionStatus.TERMINATED, status_info="user-requested", status_history=status_history, @@ -345,7 +349,7 @@ async def terminated_session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, status=KernelStatus.TERMINATED, status_info="user-requested", occupied_slots=ResourceSlot(), @@ -382,7 +386,7 @@ async def user_session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, regular_user_fixture: UserFixtureData, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING session owned by the regular user.""" unique = secrets.token_hex(4) @@ -406,10 +410,12 @@ async def user_session_seed( cluster_size=1, cluster_mode="single-node", domain_name=domain_fixture.domain_name, + domain_id=domain_fixture.domain_id, group_id=group_fixture, user_uuid=regular_user_fixture.user_uuid, access_key=regular_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture, + scaling_group_name=scaling_group_fixture.scaling_group_name, + resource_group_id=scaling_group_fixture.scaling_group_id, status=SessionStatus.RUNNING, status_info="", status_history=status_history, @@ -434,7 +440,7 @@ async def user_session_seed( group_id=group_fixture, user_uuid=regular_user_fixture.user_uuid, access_key=regular_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, status=KernelStatus.RUNNING, status_info="", occupied_slots=ResourceSlot(), diff --git a/tests/component/session/test_session_batch.py b/tests/component/session/test_session_batch.py index efa53db34dc..3e6ea75aab6 100644 --- a/tests/component/session/test_session_batch.py +++ b/tests/component/session/test_session_batch.py @@ -21,6 +21,7 @@ ) from ai.backend.manager.models.kernel import kernels from ai.backend.manager.models.scaling_group import scaling_groups +from ai.backend.testutils.fixtures import ScalingGroupFixtureData from .conftest import SessionSeedData @@ -131,7 +132,7 @@ async def test_start_service_returns_token_and_wsproxy( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, agent_registry: AsyncMock, appproxy_client_pool: AsyncMock, ) -> None: @@ -141,7 +142,7 @@ async def test_start_service_returns_token_and_wsproxy( async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture) + .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) .values(wsproxy_addr=wsproxy_addr) ) diff --git a/tests/component/session/test_session_execution.py b/tests/component/session/test_session_execution.py index 81731a21caa..c198c09a08e 100644 --- a/tests/component/session/test_session_execution.py +++ b/tests/component/session/test_session_execution.py @@ -22,6 +22,7 @@ ) from ai.backend.manager.models.kernel import kernels from ai.backend.manager.models.scaling_group import scaling_groups +from ai.backend.testutils.fixtures import ScalingGroupFixtureData from .conftest import SessionSeedData @@ -332,7 +333,7 @@ async def test_start_service_returns_token_and_wsproxy( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, agent_registry: AsyncMock, appproxy_client_pool: AsyncMock, ) -> None: @@ -340,7 +341,7 @@ async def test_start_service_returns_token_and_wsproxy( async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture) + .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) .values(wsproxy_addr=_WSPROXY_ADDR) ) async with db_engine.begin() as conn: @@ -382,7 +383,7 @@ async def test_start_service_with_specific_port( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, agent_registry: AsyncMock, appproxy_client_pool: AsyncMock, ) -> None: @@ -400,7 +401,7 @@ async def test_start_service_with_specific_port( async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture) + .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) .values(wsproxy_addr=_WSPROXY_ADDR) ) async with db_engine.begin() as conn: @@ -440,7 +441,7 @@ async def test_start_service_with_arguments_and_envs( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, agent_registry: AsyncMock, appproxy_client_pool: AsyncMock, ) -> None: @@ -448,7 +449,7 @@ async def test_start_service_with_arguments_and_envs( async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture) + .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) .values(wsproxy_addr=_WSPROXY_ADDR) ) async with db_engine.begin() as conn: @@ -496,7 +497,7 @@ async def test_start_service_kernel_host_from_agent_addr( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, agent_registry: AsyncMock, appproxy_client_pool: AsyncMock, ) -> None: @@ -504,7 +505,7 @@ async def test_start_service_kernel_host_from_agent_addr( async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture) + .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) .values(wsproxy_addr=_WSPROXY_ADDR) ) async with db_engine.begin() as conn: @@ -556,7 +557,7 @@ async def test_start_service_wsproxy_advertise_address_fallback( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, agent_registry: AsyncMock, appproxy_client_pool: AsyncMock, ) -> None: @@ -564,7 +565,7 @@ async def test_start_service_wsproxy_advertise_address_fallback( async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture) + .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) .values(wsproxy_addr=_WSPROXY_ADDR) ) async with db_engine.begin() as conn: @@ -614,14 +615,14 @@ async def test_start_nonexistent_service( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, appproxy_client_pool: AsyncMock, ) -> None: """F-BIZ-1: Service name not in service_ports → NotFoundError (AppNotFound).""" async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture) + .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) .values(wsproxy_addr=_WSPROXY_ADDR) ) async with db_engine.begin() as conn: @@ -652,7 +653,7 @@ async def test_start_inference_app_rejected( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, appproxy_client_pool: AsyncMock, ) -> None: """F-BIZ-2: Inference app → InvalidRequestError.""" @@ -669,7 +670,7 @@ async def test_start_inference_app_rejected( async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture) + .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) .values(wsproxy_addr=_WSPROXY_ADDR) ) async with db_engine.begin() as conn: @@ -700,14 +701,14 @@ async def test_start_with_invalid_port( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, appproxy_client_pool: AsyncMock, ) -> None: """F-BIZ-3: Port not in container_ports → InvalidRequestError.""" async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture) + .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) .values(wsproxy_addr=_WSPROXY_ADDR) ) async with db_engine.begin() as conn: @@ -738,14 +739,14 @@ async def test_start_with_no_wsproxy_addr( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> None: """F-BIZ-5: Scaling group has no wsproxy_addr → ServerError (ServiceUnavailable).""" # Ensure wsproxy_addr is NULL on the scaling group async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture) + .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) .values(wsproxy_addr=None) ) @@ -760,7 +761,7 @@ async def test_start_agent_returns_failed_status( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, agent_registry: AsyncMock, appproxy_client_pool: AsyncMock, ) -> None: @@ -768,7 +769,7 @@ async def test_start_agent_returns_failed_status( async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture) + .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) .values(wsproxy_addr=_WSPROXY_ADDR) ) async with db_engine.begin() as conn: diff --git a/tests/component/session/test_session_lifecycle.py b/tests/component/session/test_session_lifecycle.py index 0b9b65bd218..15e49e29e38 100644 --- a/tests/component/session/test_session_lifecycle.py +++ b/tests/component/session/test_session_lifecycle.py @@ -25,7 +25,7 @@ from ai.backend.manager.data.session.types import SessionStatus from ai.backend.manager.models.kernel import kernels from ai.backend.manager.models.session import SessionRow -from ai.backend.testutils.fixtures import DomainFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData from .conftest import SessionSeedData, UserFixtureData @@ -73,7 +73,7 @@ async def degraded_session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING_DEGRADED session with two kernels (one RUNNING, one ERROR).""" unique = secrets.token_hex(4) @@ -96,7 +96,7 @@ async def degraded_session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, now=now, ) @@ -110,10 +110,12 @@ async def degraded_session_seed( cluster_size=2, cluster_mode="single-node", domain_name=domain_fixture.domain_name, + domain_id=domain_fixture.domain_id, group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture, + scaling_group_name=scaling_group_fixture.scaling_group_name, + resource_group_id=scaling_group_fixture.scaling_group_id, status=SessionStatus.RUNNING_DEGRADED, status_info="", status_history=status_history, @@ -165,7 +167,7 @@ async def full_lifecycle_session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING session with a full lifecycle status_history (PENDING → SCHEDULED → PREPARING → PULLING → CREATING → RUNNING). @@ -194,7 +196,7 @@ async def full_lifecycle_session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, now=now, ) @@ -208,10 +210,12 @@ async def full_lifecycle_session_seed( cluster_size=1, cluster_mode="single-node", domain_name=domain_fixture.domain_name, + domain_id=domain_fixture.domain_id, group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture, + scaling_group_name=scaling_group_fixture.scaling_group_name, + resource_group_id=scaling_group_fixture.scaling_group_id, status=SessionStatus.RUNNING, status_info="", status_history=status_history, diff --git a/tests/component/session/test_session_query.py b/tests/component/session/test_session_query.py index 53a2dd4b94f..4dfb847b603 100644 --- a/tests/component/session/test_session_query.py +++ b/tests/component/session/test_session_query.py @@ -54,7 +54,7 @@ from ai.backend.manager.services.auth.processors import AuthProcessors from ai.backend.manager.services.session.processors import SessionProcessors from ai.backend.manager.services.vfolder.processors.vfolder import VFolderProcessors -from ai.backend.testutils.fixtures import DomainFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData from .conftest import SessionSeedData, UserFixtureData @@ -97,7 +97,7 @@ async def system_session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING SYSTEM session with an agent row for direct access info tests. @@ -133,7 +133,7 @@ async def system_session_seed( id=agent_id, status=AgentStatus.ALIVE, region="local", - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, available_slots=ResourceSlot(), occupied_slots=ResourceSlot(), addr="tcp://127.0.0.1:6001", @@ -152,10 +152,12 @@ async def system_session_seed( cluster_size=1, cluster_mode="single-node", domain_name=domain_fixture.domain_name, + domain_id=domain_fixture.domain_id, group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture, + scaling_group_name=scaling_group_fixture.scaling_group_name, + resource_group_id=scaling_group_fixture.scaling_group_id, status=SessionStatus.RUNNING, status_info="", status_history=status_history, @@ -180,7 +182,7 @@ async def system_session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, status=KernelStatus.RUNNING, status_info="", occupied_slots=ResourceSlot(), @@ -219,7 +221,7 @@ async def system_session_no_agent_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING SYSTEM session without an agent row. @@ -247,10 +249,12 @@ async def system_session_no_agent_seed( cluster_size=1, cluster_mode="single-node", domain_name=domain_fixture.domain_name, + domain_id=domain_fixture.domain_id, group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture, + scaling_group_name=scaling_group_fixture.scaling_group_name, + resource_group_id=scaling_group_fixture.scaling_group_id, status=SessionStatus.RUNNING, status_info="", status_history=status_history, @@ -275,7 +279,7 @@ async def system_session_no_agent_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, status=KernelStatus.RUNNING, status_info="", occupied_slots=ResourceSlot(), @@ -310,7 +314,7 @@ async def second_session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, session_seed: SessionSeedData, ) -> AsyncIterator[SessionSeedData]: """Seed a second INTERACTIVE session alongside session_seed. @@ -340,10 +344,12 @@ async def second_session_seed( cluster_size=1, cluster_mode="single-node", domain_name=domain_fixture.domain_name, + domain_id=domain_fixture.domain_id, group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture, + scaling_group_name=scaling_group_fixture.scaling_group_name, + resource_group_id=scaling_group_fixture.scaling_group_id, status=SessionStatus.RUNNING, status_info="", status_history=status_history, @@ -368,7 +374,7 @@ async def second_session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, status=KernelStatus.RUNNING, status_info="", occupied_slots=ResourceSlot(), diff --git a/tests/component/session_v2/conftest.py b/tests/component/session_v2/conftest.py index 6e646077501..0c301096ca1 100644 --- a/tests/component/session_v2/conftest.py +++ b/tests/component/session_v2/conftest.py @@ -27,6 +27,8 @@ RoleStatus, ScopeType, ) +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.plugin.monitor import ErrorPluginContext from ai.backend.common.types import ResourceSlot, SessionId, SessionTypes from ai.backend.manager.actions.validators import ActionValidators @@ -58,7 +60,7 @@ from ai.backend.manager.services.processors import Processors from ai.backend.manager.services.session.processors import SessionProcessors from ai.backend.manager.services.session.service import SessionService, SessionServiceArgs -from ai.backend.testutils.fixtures import DomainFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData if TYPE_CHECKING: from sqlalchemy.ext.asyncio.engine import AsyncEngine as SAEngine @@ -291,10 +293,12 @@ async def _seed_session( db_engine: SAEngine, *, domain_name: str, + domain_id: DomainID, group_id: uuid.UUID, user_uuid: uuid.UUID, access_key: str, scaling_group: str, + resource_group_id: ResourceGroupID, status: SessionStatus = SessionStatus.RUNNING, ) -> SessionSeedData: """Insert a session + kernel row with RBAC scope association. @@ -323,10 +327,12 @@ async def _seed_session( cluster_size=1, cluster_mode="single-node", domain_name=domain_name, + domain_id=domain_id, group_id=group_id, user_uuid=user_uuid, access_key=access_key, scaling_group_name=scaling_group, + resource_group_id=resource_group_id, status=status, status_info="", status_history=status_history, @@ -415,16 +421,18 @@ async def admin_session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING session owned by the admin user.""" seed = await _seed_session( db_engine, domain_name=domain_fixture.domain_name, + domain_id=domain_fixture.domain_id, group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, + resource_group_id=scaling_group_fixture.scaling_group_id, ) yield seed await _cleanup_session(db_engine, seed.session_id) @@ -436,7 +444,7 @@ async def user_session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, regular_user_fixture: UserFixtureData, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, user_system_role: uuid.UUID, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING session owned by the regular user. @@ -446,10 +454,12 @@ async def user_session_seed( seed = await _seed_session( db_engine, domain_name=domain_fixture.domain_name, + domain_id=domain_fixture.domain_id, group_id=group_fixture, user_uuid=regular_user_fixture.user_uuid, access_key=regular_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, + resource_group_id=scaling_group_fixture.scaling_group_id, ) yield seed await _cleanup_session(db_engine, seed.session_id) diff --git a/tests/component/stream/conftest.py b/tests/component/stream/conftest.py index 9b130e9025e..580493c4ce9 100644 --- a/tests/component/stream/conftest.py +++ b/tests/component/stream/conftest.py @@ -34,7 +34,7 @@ from ai.backend.manager.services.session.service import SessionService, SessionServiceArgs from ai.backend.manager.services.stream.processors import StreamProcessors from ai.backend.manager.services.stream.service import StreamService -from ai.backend.testutils.fixtures import DomainFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData _STREAMING_SERVER_SUBAPP_MODULES = (_auth_api,) @@ -138,7 +138,7 @@ async def session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING session + kernel with service_ports in the database.""" unique = secrets.token_hex(4) @@ -162,10 +162,12 @@ async def session_seed( cluster_size=1, cluster_mode="single-node", domain_name=domain_fixture.domain_name, + domain_id=domain_fixture.domain_id, group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture, + scaling_group_name=scaling_group_fixture.scaling_group_name, + resource_group_id=scaling_group_fixture.scaling_group_id, status=SessionStatus.RUNNING, status_info="", status_history=status_history, @@ -190,7 +192,7 @@ async def session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, status=KernelStatus.RUNNING, status_info="", occupied_slots=ResourceSlot(), diff --git a/tests/component/streaming/conftest.py b/tests/component/streaming/conftest.py index de46f0a3e37..065f1577206 100644 --- a/tests/component/streaming/conftest.py +++ b/tests/component/streaming/conftest.py @@ -31,7 +31,7 @@ from ai.backend.manager.repositories.stream.repository import StreamRepository from ai.backend.manager.services.stream.processors import StreamProcessors from ai.backend.manager.services.stream.service import StreamService -from ai.backend.testutils.fixtures import DomainFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData _STREAMING_SERVER_SUBAPP_MODULES = (_auth_api,) @@ -105,7 +105,7 @@ async def session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING session + kernel with service_ports in the database. @@ -133,10 +133,12 @@ async def session_seed( cluster_size=1, cluster_mode="single-node", domain_name=domain_fixture.domain_name, + domain_id=domain_fixture.domain_id, group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture, + scaling_group_name=scaling_group_fixture.scaling_group_name, + resource_group_id=scaling_group_fixture.scaling_group_id, status=SessionStatus.RUNNING, status_info="", status_history=status_history, @@ -161,7 +163,7 @@ async def session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, status=KernelStatus.RUNNING, status_info="", occupied_slots=ResourceSlot(), @@ -209,7 +211,7 @@ async def session_seed_no_ports( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: str, + scaling_group_fixture: ScalingGroupFixtureData, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING session + kernel with ``service_ports=None``. @@ -237,10 +239,12 @@ async def session_seed_no_ports( cluster_size=1, cluster_mode="single-node", domain_name=domain_fixture.domain_name, + domain_id=domain_fixture.domain_id, group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture, + scaling_group_name=scaling_group_fixture.scaling_group_name, + resource_group_id=scaling_group_fixture.scaling_group_id, status=SessionStatus.RUNNING, status_info="", status_history=status_history, @@ -265,7 +269,7 @@ async def session_seed_no_ports( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture, + scaling_group=scaling_group_fixture.scaling_group_name, status=KernelStatus.RUNNING, status_info="", occupied_slots=ResourceSlot(), diff --git a/tests/unit/manager/models/test_session.py b/tests/unit/manager/models/test_session.py index 423463fd1a4..25e28af4319 100644 --- a/tests/unit/manager/models/test_session.py +++ b/tests/unit/manager/models/test_session.py @@ -7,6 +7,8 @@ import pytest from sqlalchemy.exc import IntegrityError +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import BinarySize, ResourceSlot from ai.backend.manager.data.session.types import SessionStatus from ai.backend.manager.models.agent import AgentRow @@ -31,7 +33,7 @@ from ai.backend.manager.models.resource_preset import ResourcePresetRow from ai.backend.manager.models.routing import RoutingRow from ai.backend.manager.models.runtime_variant import RuntimeVariantRow -from ai.backend.manager.models.scaling_group import ScalingGroupOpts, ScalingGroupRow +from ai.backend.manager.models.scaling_group import ScalingGroupRow from ai.backend.manager.models.session import SessionRow from ai.backend.manager.models.user import UserRow from ai.backend.manager.models.utils import ExtendedAsyncSAEngine @@ -50,7 +52,10 @@ class SessionData: name: str user_uuid: uuid.UUID group_id: uuid.UUID + domain_id: DomainID domain_name: str + resource_group_id: ResourceGroupID + scaling_group_name: str status: SessionStatus @@ -104,7 +109,7 @@ async def domain( self, database_with_tables: ExtendedAsyncSAEngine ) -> AsyncGenerator[DomainRow, None]: """Create test domain.""" - domain = DomainRow(name=f"test-{uuid.uuid4()}") + domain = DomainRow(id=DomainID(uuid.uuid4()), name=f"test-{uuid.uuid4()}") async with database_with_tables.begin_session() as db_sess: db_sess.add(domain) @@ -112,6 +117,24 @@ async def domain( yield domain + @pytest.fixture + async def scaling_group( + self, database_with_tables: ExtendedAsyncSAEngine + ) -> AsyncGenerator[ScalingGroupRow, None]: + """Create test scaling group.""" + sgroup = ScalingGroupRow( + id=ResourceGroupID(uuid.uuid4()), + name=f"test-sg-{uuid.uuid4().hex[:8]}", + driver="static", + scheduler="fifo", + ) + + async with database_with_tables.begin_session() as db_sess: + db_sess.add(sgroup) + await db_sess.flush() + + yield sgroup + @pytest.fixture async def user_policy( self, database_with_tables: ExtendedAsyncSAEngine, domain: DomainRow @@ -211,26 +234,6 @@ async def group( yield group - @pytest.fixture - async def scaling_group( - self, - database_with_tables: ExtendedAsyncSAEngine, - ) -> AsyncGenerator[str, None]: - """Create test scaling group and yield its name.""" - scaling_group = ScalingGroupRow( - name="test-sg", - driver="static", - driver_opts={}, - scheduler="fifo", - scheduler_opts=ScalingGroupOpts(), - ) - - async with database_with_tables.begin_session() as db_sess: - db_sess.add(scaling_group) - await db_sess.flush() - - yield scaling_group.name - @pytest.fixture async def prepared_first_session( self, @@ -238,7 +241,7 @@ async def prepared_first_session( user_one: UserData, group: GroupRow, domain: DomainRow, - scaling_group: str, + scaling_group: ScalingGroupRow, test_config: TestConfig, ) -> AsyncGenerator[SessionData, None]: """Create first session with status from parametrize (indirect fixture)""" @@ -248,8 +251,10 @@ async def prepared_first_session( name=f"test-{uuid.uuid4()!s}", user_uuid=user_a.uuid, group_id=group.id, + domain_id=domain.id, domain_name=domain.name, - scaling_group_name=scaling_group, + resource_group_id=scaling_group.id, + scaling_group_name=scaling_group.name, status=status, occupying_slots=ResourceSlot(), requested_slots=ResourceSlot(), @@ -265,7 +270,10 @@ async def prepared_first_session( name=session.name, user_uuid=user_a.uuid, group_id=group.id, + domain_id=domain.id, domain_name=domain.name, + resource_group_id=scaling_group.id, + scaling_group_name=scaling_group.name, status=status, ) @@ -287,15 +295,16 @@ async def test_duplicate_session_name_with_terminal_status_in_same_user( test_config: TestConfig, database_with_tables: ExtendedAsyncSAEngine, prepared_first_session: SessionData, - scaling_group: str, ) -> None: async with database_with_tables.begin_session() as db_sess: duplicate_session = SessionRow( name=prepared_first_session.name, user_uuid=prepared_first_session.user_uuid, group_id=prepared_first_session.group_id, + domain_id=prepared_first_session.domain_id, domain_name=prepared_first_session.domain_name, - scaling_group_name=scaling_group, + resource_group_id=prepared_first_session.resource_group_id, + scaling_group_name=prepared_first_session.scaling_group_name, status=test_config.second_session_status, occupying_slots=ResourceSlot(), requested_slots=ResourceSlot(), @@ -320,14 +329,15 @@ async def test_duplicate_session_name_with_non_terminal_status_in_same_user( test_config: TestConfig, database_with_tables: ExtendedAsyncSAEngine, prepared_first_session: SessionData, - scaling_group: str, ) -> None: duplicate_session = SessionRow( name=prepared_first_session.name, user_uuid=prepared_first_session.user_uuid, group_id=prepared_first_session.group_id, + domain_id=prepared_first_session.domain_id, domain_name=prepared_first_session.domain_name, - scaling_group_name=scaling_group, + resource_group_id=prepared_first_session.resource_group_id, + scaling_group_name=prepared_first_session.scaling_group_name, status=test_config.second_session_status, occupying_slots=ResourceSlot(), requested_slots=ResourceSlot(), @@ -361,7 +371,6 @@ async def test_duplicate_session_name_different_user( database_with_tables: ExtendedAsyncSAEngine, prepared_first_session: SessionData, user_two: UserData, - scaling_group: str, test_config: TestConfig, ) -> None: async with database_with_tables.begin_session() as db_sess: @@ -369,8 +378,10 @@ async def test_duplicate_session_name_different_user( name=prepared_first_session.name, user_uuid=user_two.uuid, # Different user group_id=prepared_first_session.group_id, + domain_id=prepared_first_session.domain_id, domain_name=prepared_first_session.domain_name, - scaling_group_name=scaling_group, + resource_group_id=prepared_first_session.resource_group_id, + scaling_group_name=prepared_first_session.scaling_group_name, status=test_config.second_session_status, occupying_slots=ResourceSlot(), requested_slots=ResourceSlot(), diff --git a/tests/unit/manager/repositories/agent/test_repository.py b/tests/unit/manager/repositories/agent/test_repository.py index 3912197be7d..2a889b39350 100644 --- a/tests/unit/manager/repositories/agent/test_repository.py +++ b/tests/unit/manager/repositories/agent/test_repository.py @@ -23,6 +23,8 @@ from ai.backend.common.clients.valkey_client.valkey_stat.client import ValkeyStatClient from ai.backend.common.data.agent.types import AgentInfo from ai.backend.common.exception import AgentNotFound +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ( AgentId, ClusterMode, @@ -648,15 +650,25 @@ async def db_with_tables( db_sess.add(ResourceSlotTypeRow(slot_name="cpu", slot_type="count", rank=0)) yield database_connection + @pytest.fixture + def test_domain_id(self) -> DomainID: + return DomainID(uuid4()) + + @pytest.fixture + def test_scaling_group_id(self) -> ResourceGroupID: + return ResourceGroupID(uuid4()) + @pytest.fixture async def test_domain( self, db_with_tables: ExtendedAsyncSAEngine, + test_domain_id: DomainID, ) -> AsyncGenerator[str, None]: """Create default domain""" domain_name = str(uuid4()) async with db_with_tables.begin_session() as db_sess: domain = DomainRow( + id=test_domain_id, name=domain_name, description="Test domain", is_active=True, @@ -707,11 +719,13 @@ async def test_group( async def scaling_group( self, db_with_tables: ExtendedAsyncSAEngine, + test_scaling_group_id: ResourceGroupID, ) -> AsyncGenerator[str, None]: """Create default scaling group""" name = str(uuid4()) async with db_with_tables.begin_session() as db_sess: scaling_group = ScalingGroupRow( + id=test_scaling_group_id, name=name, description="Test scaling group", is_active=True, @@ -731,7 +745,9 @@ async def agent_with_kernels( request: pytest.FixtureRequest, db_with_tables: ExtendedAsyncSAEngine, test_group: tuple[str, str], + test_domain_id: DomainID, scaling_group: str, + test_scaling_group_id: ResourceGroupID, ) -> AsyncGenerator[KernelFilteringTestCase, None]: """Create ONE agent with kernels based on the test case from indirect parametrization""" test_case: KernelFilteringTestCase = request.param @@ -786,8 +802,10 @@ async def agent_with_kernels( id=session_id, name=session_name, session_type=SessionTypes.INTERACTIVE, + domain_id=test_domain_id, domain_name=domain_name, group_id=UUID(group_id_str), + resource_group_id=test_scaling_group_id, scaling_group_name=scaling_group, status=SessionStatus.RUNNING, status_info="test", diff --git a/tests/unit/manager/repositories/deployment/test_deployment_repository.py b/tests/unit/manager/repositories/deployment/test_deployment_repository.py index 310b735b8df..3c25ffe867e 100644 --- a/tests/unit/manager/repositories/deployment/test_deployment_repository.py +++ b/tests/unit/manager/repositories/deployment/test_deployment_repository.py @@ -21,9 +21,11 @@ from ai.backend.common.dto.manager.v2.deployment.types import IntOrPercent from ai.backend.common.identifier.deployment import DeploymentID from ai.backend.common.identifier.deployment_revision import DeploymentRevisionID +from ai.backend.common.identifier.domain import DomainID from ai.backend.common.identifier.image import ImageID from ai.backend.common.identifier.replica import ReplicaID from ai.backend.common.identifier.replica_group import ReplicaGroupID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.identifier.runtime_variant import RuntimeVariantID from ai.backend.common.identifier.vfolder import VFolderUUID from ai.backend.common.types import ( @@ -213,16 +215,26 @@ async def db_with_cleanup( ) yield database_connection + @pytest.fixture + def test_domain_id(self) -> DomainID: + return DomainID(uuid.uuid4()) + + @pytest.fixture + def test_scaling_group_id(self) -> ResourceGroupID: + return ResourceGroupID(uuid.uuid4()) + @pytest.fixture async def test_domain_name( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, ) -> str: """Create test domain and return domain name.""" domain_name = f"test-domain-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: domain = DomainRow( + id=test_domain_id, name=domain_name, description="Test domain for deployment", is_active=True, @@ -239,12 +251,14 @@ async def test_domain_name( async def test_scaling_group_name( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_scaling_group_id: ResourceGroupID, ) -> str: """Create test scaling group and return name.""" sgroup_name = f"test-sgroup-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: sgroup = ScalingGroupRow( + id=test_scaling_group_id, name=sgroup_name, description="Test scaling group", is_active=True, @@ -443,7 +457,9 @@ async def test_session_id( db_with_cleanup: ExtendedAsyncSAEngine, test_access_key: AccessKey, test_scaling_group_name: str, + test_scaling_group_id: ResourceGroupID, test_domain_name: str, + test_domain_id: DomainID, test_user_uuid: uuid.UUID, test_group_id: uuid.UUID, ) -> SessionId: @@ -455,10 +471,12 @@ async def test_session_id( id=session_id, name=f"test-session-{uuid.uuid4().hex[:8]}", session_type=SessionTypes.INTERACTIVE, + domain_id=test_domain_id, domain_name=test_domain_name, group_id=test_group_id, user_uuid=test_user_uuid, access_key=test_access_key, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, status=SessionStatus.RUNNING, cluster_mode=ClusterMode.SINGLE_NODE, @@ -839,7 +857,9 @@ async def test_fetch_multiple_routes( deployment_repository: DeploymentRepository, db_with_cleanup: ExtendedAsyncSAEngine, test_domain_name: str, + test_domain_id: DomainID, test_scaling_group_name: str, + test_scaling_group_id: ResourceGroupID, test_access_key: AccessKey, test_agent_id: AgentId, test_user_uuid: uuid.UUID, @@ -940,10 +960,12 @@ async def test_fetch_multiple_routes( id=session_id, name=f"session-{i}", session_type=SessionTypes.INTERACTIVE, + domain_id=test_domain_id, domain_name=test_domain_name, group_id=test_group_id, user_uuid=test_user_uuid, access_key=test_access_key, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, status=SessionStatus.RUNNING, cluster_mode=ClusterMode.SINGLE_NODE, diff --git a/tests/unit/manager/repositories/model_card/test_model_card_delete.py b/tests/unit/manager/repositories/model_card/test_model_card_delete.py index a566db031f6..583b23d95e0 100644 --- a/tests/unit/manager/repositories/model_card/test_model_card_delete.py +++ b/tests/unit/manager/repositories/model_card/test_model_card_delete.py @@ -12,6 +12,8 @@ import sqlalchemy as sa from ai.backend.common.dto.manager.v2.model_card.request import DeleteModelCardOptions +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.identifier.vfolder import VFolderUUID from ai.backend.common.types import ( MountPermission, @@ -137,13 +139,23 @@ async def db_with_cleanup( ): yield database_connection + @pytest.fixture + def test_domain_id(self) -> DomainID: + return DomainID(uuid.uuid4()) + + @pytest.fixture + def test_scaling_group_id(self) -> ResourceGroupID: + return ResourceGroupID(uuid.uuid4()) + @pytest.fixture async def test_domain( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, ) -> DomainRow: async with db_with_cleanup.begin_session() as db_sess: domain = DomainRow( + id=test_domain_id, name=f"test-domain-{uuid.uuid4().hex[:8]}", description="Test domain", is_active=True, @@ -155,6 +167,25 @@ async def test_domain( await db_sess.flush() return domain + @pytest.fixture + async def test_scaling_group( + self, + db_with_cleanup: ExtendedAsyncSAEngine, + test_scaling_group_id: ResourceGroupID, + ) -> ScalingGroupRow: + async with db_with_cleanup.begin_session() as db_sess: + sgroup = ScalingGroupRow( + id=test_scaling_group_id, + name=f"test-sgroup-{uuid.uuid4().hex[:8]}", + driver="static", + driver_opts={}, + scheduler="fifo", + scheduler_opts=ScalingGroupOpts(), + ) + db_sess.add(sgroup) + await db_sess.flush() + return sgroup + @pytest.fixture async def test_user_resource_policy( self, @@ -273,6 +304,9 @@ async def mounted_vfolder( self, db_with_cleanup: ExtendedAsyncSAEngine, test_domain: DomainRow, + test_domain_id: DomainID, + test_scaling_group: ScalingGroupRow, + test_scaling_group_id: ResourceGroupID, test_user: UserRow, test_group: GroupRow, ) -> VFolderRow: @@ -302,9 +336,11 @@ async def mounted_vfolder( await db_sess.flush() mount_holder = SessionRow( id=uuid.uuid4(), + domain_id=test_domain_id, domain_name=test_domain.name, + resource_group_id=test_scaling_group_id, + scaling_group_name=test_scaling_group.name, group_id=test_group.id, - scaling_group_name="test-sg", user_uuid=test_user.uuid, occupying_slots=ResourceSlot(), requested_slots=ResourceSlot(), diff --git a/tests/unit/manager/repositories/resource_preset/test_check_presets.py b/tests/unit/manager/repositories/resource_preset/test_check_presets.py index ab1621a139a..5f8fdd133ee 100644 --- a/tests/unit/manager/repositories/resource_preset/test_check_presets.py +++ b/tests/unit/manager/repositories/resource_preset/test_check_presets.py @@ -20,6 +20,8 @@ from ai.backend.common.clients.valkey_client.valkey_stat.client import ValkeyStatClient from ai.backend.common.data.user.types import UserRole +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ( AccessKey, AgentId, @@ -152,16 +154,26 @@ async def db_with_cleanup( ) yield database_connection + @pytest.fixture + def test_domain_id(self) -> DomainID: + return DomainID(uuid.uuid4()) + + @pytest.fixture + def test_scaling_group_id(self) -> ResourceGroupID: + return ResourceGroupID(uuid.uuid4()) + @pytest.fixture async def test_domain_name( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, ) -> AsyncGenerator[str, None]: """Create test domain and return domain name""" domain_name = f"test-domain-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: domain = DomainRow( + id=test_domain_id, name=domain_name, total_resource_slots=ResourceSlot({ "cpu": Decimal("1000"), @@ -182,12 +194,14 @@ async def test_scaling_group_name( self, db_with_cleanup: ExtendedAsyncSAEngine, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, ) -> AsyncGenerator[str, None]: """Create test scaling group and return scaling group name""" sg_name = f"test-sgroup-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: sg = ScalingGroupRow( + id=test_scaling_group_id, name=sg_name, driver="test-driver", scheduler="fifo", @@ -586,7 +600,9 @@ async def test_running_kernels_count_towards_occupied_slots( self, repository: ResourcePresetRepository, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_resource_policy_name: str, test_group_id: uuid.UUID, @@ -611,10 +627,12 @@ async def test_running_kernels_count_towards_occupied_slots( status=SessionStatus.RUNNING, status_data={}, created_at=datetime.now(tzutc()), + domain_id=test_domain_id, domain_name=test_domain_name, group_id=test_group_id, user_uuid=test_user_uuid, access_key=test_keypair_access_key, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, result=SessionResult.UNDEFINED, agent_ids=[], @@ -720,7 +738,9 @@ async def test_terminating_kernels_count_towards_occupied_slots( self, repository: ResourcePresetRepository, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_resource_policy_name: str, test_group_id: uuid.UUID, @@ -744,10 +764,12 @@ async def test_terminating_kernels_count_towards_occupied_slots( status=SessionStatus.TERMINATING, status_data={}, created_at=datetime.now(tzutc()), + domain_id=test_domain_id, domain_name=test_domain_name, group_id=test_group_id, user_uuid=test_user_uuid, access_key=test_keypair_access_key, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, result=SessionResult.UNDEFINED, agent_ids=[], @@ -850,7 +872,9 @@ async def test_pending_kernels_do_not_count_towards_occupied_slots( self, repository: ResourcePresetRepository, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_resource_policy_name: str, test_group_id: uuid.UUID, @@ -874,10 +898,12 @@ async def test_pending_kernels_do_not_count_towards_occupied_slots( status=SessionStatus.PENDING, status_data={}, created_at=datetime.now(tzutc()), + domain_id=test_domain_id, domain_name=test_domain_name, group_id=test_group_id, user_uuid=test_user_uuid, access_key=test_keypair_access_key, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, result=SessionResult.UNDEFINED, agent_ids=[], @@ -972,7 +998,9 @@ async def test_ignores_cached_occupied_slots_in_agent_row( self, repository: ResourcePresetRepository, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_resource_policy_name: str, test_group_id: uuid.UUID, @@ -1033,10 +1061,12 @@ async def test_ignores_cached_occupied_slots_in_agent_row( status=SessionStatus.RUNNING, status_data={}, created_at=datetime.now(tzutc()), + domain_id=test_domain_id, domain_name=test_domain_name, group_id=test_group_id, user_uuid=test_user_uuid, access_key=test_keypair_access_key, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, result=SessionResult.UNDEFINED, agent_ids=[], diff --git a/tests/unit/manager/repositories/scaling_group/test_scaling_group_repository.py b/tests/unit/manager/repositories/scaling_group/test_scaling_group_repository.py index 1c53ac29884..de00882534d 100644 --- a/tests/unit/manager/repositories/scaling_group/test_scaling_group_repository.py +++ b/tests/unit/manager/repositories/scaling_group/test_scaling_group_repository.py @@ -420,14 +420,21 @@ async def sample_scaling_group_with_sessions_and_routes( db_sess.add(sgroup) await db_sess.flush() # Flush to ensure scaling group exists before creating references + domain_id = await db_sess.scalar( + sa.select(DomainRow.id).where(DomainRow.name == test_domain) + ) + assert domain_id is not None + # Create 2 sessions with this scaling group for i in range(2): session_id = SessionId(uuid.uuid4()) session = SessionRow( id=session_id, domain_name=test_domain, + domain_id=domain_id, group_id=test_group_id, user_uuid=test_user_uuid, + resource_group_id=sgroup.id, scaling_group_name=sgroup_name, cluster_size=1, vfolder_mounts={}, @@ -1274,12 +1281,25 @@ async def sample_session( test_user_uuid, test_domain, test_group_id = test_user_domain_group session_id = SessionId(uuid.uuid4()) async with db_with_cleanup.begin_session() as db_sess: + domain_id = await db_sess.scalar( + sa.select(DomainRow.id).where(DomainRow.name == test_domain) + ) + resource_group_id = await db_sess.scalar( + sa.select(ScalingGroupRow.id).where( + ScalingGroupRow.name == sample_scaling_group_for_hierarchy + ) + ) + assert domain_id is not None + assert resource_group_id is not None + db_sess.add( SessionRow( id=session_id, domain_name=test_domain, + domain_id=domain_id, group_id=test_group_id, user_uuid=test_user_uuid, + resource_group_id=resource_group_id, scaling_group_name=sample_scaling_group_for_hierarchy, cluster_size=1, vfolder_mounts={}, diff --git a/tests/unit/manager/repositories/scheduler/conftest.py b/tests/unit/manager/repositories/scheduler/conftest.py index c8f566f2f2a..ecc443ed2b5 100644 --- a/tests/unit/manager/repositories/scheduler/conftest.py +++ b/tests/unit/manager/repositories/scheduler/conftest.py @@ -19,6 +19,8 @@ from dateutil.tz import tzutc from ai.backend.common.data.user.types import UserRole +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ( AccessKey, AgentId, @@ -106,14 +108,26 @@ async def db_with_cleanup( yield database_connection +@pytest.fixture +def test_domain_id() -> DomainID: + return DomainID(uuid.uuid4()) + + +@pytest.fixture +def test_scaling_group_id() -> ResourceGroupID: + return ResourceGroupID(uuid.uuid4()) + + @pytest.fixture async def test_domain_name( db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, ) -> AsyncGenerator[str, None]: domain_name = f"test-domain-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: db_sess.add( DomainRow( + id=test_domain_id, name=domain_name, total_resource_slots=ResourceSlot({ "cpu": Decimal("1000"), @@ -128,11 +142,13 @@ async def test_domain_name( @pytest.fixture async def test_scaling_group_name( db_with_cleanup: ExtendedAsyncSAEngine, + test_scaling_group_id: ResourceGroupID, ) -> AsyncGenerator[str, None]: sg_name = f"test-sgroup-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: db_sess.add( ScalingGroupRow( + id=test_scaling_group_id, name=sg_name, driver="static", scheduler="fifo", @@ -371,7 +387,9 @@ async def seed_agent_resources( async def create_pending_session_with_kernels( db: ExtendedAsyncSAEngine, *, + domain_id: DomainID, domain_name: str, + resource_group_id: ResourceGroupID, scaling_group_name: str, group_id: uuid.UUID, user_uuid: uuid.UUID, @@ -397,10 +415,12 @@ async def create_pending_session_with_kernels( id=session_id, name=f"test-session-{uuid.uuid4().hex[:8]}", session_type=SessionTypes.INTERACTIVE, + domain_id=domain_id, domain_name=domain_name, group_id=group_id, user_uuid=user_uuid, access_key=access_key, + resource_group_id=resource_group_id, scaling_group_name=scaling_group_name, status=SessionStatus.PENDING, status_info="test", diff --git a/tests/unit/manager/repositories/scheduler/test_cancellation_resource_freeing.py b/tests/unit/manager/repositories/scheduler/test_cancellation_resource_freeing.py index 99c1dbed02a..b1cdfcc564f 100644 --- a/tests/unit/manager/repositories/scheduler/test_cancellation_resource_freeing.py +++ b/tests/unit/manager/repositories/scheduler/test_cancellation_resource_freeing.py @@ -13,6 +13,8 @@ from dateutil.tz import tzutc from ai.backend.common.data.user.types import UserRole +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ( AccessKey, AgentId, @@ -93,15 +95,25 @@ async def db_with_cleanup( async with with_tables(database_connection, _BASE_TABLES): yield database_connection + @pytest.fixture + def test_domain_id(self) -> DomainID: + return DomainID(uuid.uuid4()) + + @pytest.fixture + def test_scaling_group_id(self) -> ResourceGroupID: + return ResourceGroupID(uuid.uuid4()) + @pytest.fixture async def test_domain_name( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, ) -> AsyncGenerator[str, None]: domain_name = f"test-domain-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: db_sess.add( DomainRow( + id=test_domain_id, name=domain_name, total_resource_slots=ResourceSlot({ "cpu": Decimal("1000"), @@ -116,11 +128,13 @@ async def test_domain_name( async def test_scaling_group_name( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_scaling_group_id: ResourceGroupID, ) -> AsyncGenerator[str, None]: sg_name = f"test-sgroup-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: db_sess.add( ScalingGroupRow( + id=test_scaling_group_id, name=sg_name, driver="static", scheduler="fifo", @@ -311,7 +325,9 @@ async def _insert_session( *, kernel_status: KernelStatus, domain_name: str, + domain_id: DomainID, scaling_group_name: str, + resource_group_id: ResourceGroupID, group_id: uuid.UUID, user_uuid: uuid.UUID, access_key: AccessKey, @@ -335,8 +351,10 @@ async def _insert_session( name=f"test-session-{uuid.uuid4().hex[:8]}", session_type=SessionTypes.INTERACTIVE, domain_name=domain_name, + domain_id=domain_id, group_id=group_id, scaling_group_name=scaling_group_name, + resource_group_id=resource_group_id, status=session_status, status_info="test", cluster_mode=ClusterMode.SINGLE_NODE, @@ -403,7 +421,9 @@ async def pending_session( self, db_with_cleanup: ExtendedAsyncSAEngine, test_domain_name: str, + test_domain_id: DomainID, test_scaling_group_name: str, + test_scaling_group_id: ResourceGroupID, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, test_access_key: AccessKey, @@ -413,7 +433,9 @@ async def pending_session( db_with_cleanup, kernel_status=KernelStatus.PENDING, domain_name=test_domain_name, + domain_id=test_domain_id, scaling_group_name=test_scaling_group_name, + resource_group_id=test_scaling_group_id, group_id=test_group_id, user_uuid=test_user_uuid, access_key=test_access_key, @@ -425,7 +447,9 @@ async def pulling_session_on_agent( self, db_with_cleanup: ExtendedAsyncSAEngine, test_domain_name: str, + test_domain_id: DomainID, test_scaling_group_name: str, + test_scaling_group_id: ResourceGroupID, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, test_access_key: AccessKey, @@ -436,7 +460,9 @@ async def pulling_session_on_agent( db_with_cleanup, kernel_status=KernelStatus.PULLING, domain_name=test_domain_name, + domain_id=test_domain_id, scaling_group_name=test_scaling_group_name, + resource_group_id=test_scaling_group_id, group_id=test_group_id, user_uuid=test_user_uuid, access_key=test_access_key, diff --git a/tests/unit/manager/repositories/scheduler/test_reserved_allocation.py b/tests/unit/manager/repositories/scheduler/test_reserved_allocation.py index 4527a9285e3..87c421a1cc0 100644 --- a/tests/unit/manager/repositories/scheduler/test_reserved_allocation.py +++ b/tests/unit/manager/repositories/scheduler/test_reserved_allocation.py @@ -16,6 +16,8 @@ import sqlalchemy as sa +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import AccessKey, KernelId, SessionId from ai.backend.manager.data.kernel.types import KernelStatus from ai.backend.manager.data.session.types import SessionStatus @@ -40,7 +42,9 @@ class TestAllocateSessionsReservation: async def test_a1_single_kernel_reserves_and_schedules( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -57,7 +61,9 @@ async def test_a1_single_kernel_reserves_and_schedules( ) session_id, kernel_ids = await create_pending_session_with_kernels( db_with_cleanup, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, @@ -97,7 +103,9 @@ async def test_a1_single_kernel_reserves_and_schedules( async def test_a2_multi_kernel_reserves_sum( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -114,7 +122,9 @@ async def test_a2_multi_kernel_reserves_sum( ) session_id, kernel_ids = await create_pending_session_with_kernels( db_with_cleanup, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, @@ -152,7 +162,9 @@ async def test_a2_multi_kernel_reserves_sum( async def test_a3_idempotent_double_allocate( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -169,7 +181,9 @@ async def test_a3_idempotent_double_allocate( ) session_id, kernel_ids = await create_pending_session_with_kernels( db_with_cleanup, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, @@ -196,7 +210,9 @@ async def test_a3_idempotent_double_allocate( async def test_a4_capacity_gate_rolls_back( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -215,7 +231,9 @@ async def test_a4_capacity_gate_rolls_back( ) session_id, kernel_ids = await create_pending_session_with_kernels( db_with_cleanup, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, @@ -247,7 +265,9 @@ async def test_a4_capacity_gate_rolls_back( async def test_a5_batch_atomicity_second_session_exceeds( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -265,7 +285,9 @@ async def test_a5_batch_atomicity_second_session_exceeds( ) session_a, kernels_a = await create_pending_session_with_kernels( db_with_cleanup, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, @@ -274,7 +296,9 @@ async def test_a5_batch_atomicity_second_session_exceeds( ) session_b, kernels_b = await create_pending_session_with_kernels( db_with_cleanup, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, @@ -320,7 +344,9 @@ async def _allocate_one( self, db: ExtendedAsyncSAEngine, *, + domain_id: DomainID, domain_name: str, + resource_group_id: ResourceGroupID, scaling_group_name: str, group_id: uuid.UUID, user_uuid: uuid.UUID, @@ -332,7 +358,9 @@ async def _allocate_one( """Create a PENDING session and allocate it (-> SCHEDULED, reserved up).""" session_id, kernel_ids = await create_pending_session_with_kernels( db, + domain_id=domain_id, domain_name=domain_name, + resource_group_id=resource_group_id, scaling_group_name=scaling_group_name, group_id=group_id, user_uuid=user_uuid, @@ -352,7 +380,9 @@ async def _allocate_one( async def test_c2_cancel_releases_reserved( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -370,7 +400,9 @@ async def test_c2_cancel_releases_reserved( db_source = ScheduleDBSource(db_with_cleanup) _, kernel_id = await self._allocate_one( db_with_cleanup, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, @@ -410,7 +442,9 @@ async def test_c2_cancel_releases_reserved( async def test_c3_bulk_terminate_releases_reserved( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -428,7 +462,9 @@ async def test_c3_bulk_terminate_releases_reserved( db_source = ScheduleDBSource(db_with_cleanup) _, kernel_id = await self._allocate_one( db_with_cleanup, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, @@ -451,7 +487,9 @@ async def test_c3_bulk_terminate_releases_reserved( async def test_full_lifecycle_reserved_to_used_to_zero( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -469,7 +507,9 @@ async def test_full_lifecycle_reserved_to_used_to_zero( db_source = ScheduleDBSource(db_with_cleanup) _, kernel_id = await self._allocate_one( db_with_cleanup, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, diff --git a/tests/unit/manager/repositories/scheduler/test_resource_allocation.py b/tests/unit/manager/repositories/scheduler/test_resource_allocation.py index 68a8437e8c8..f377c440723 100644 --- a/tests/unit/manager/repositories/scheduler/test_resource_allocation.py +++ b/tests/unit/manager/repositories/scheduler/test_resource_allocation.py @@ -16,6 +16,8 @@ from dateutil.tz import tzutc from ai.backend.common.data.user.types import UserRole +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ( AccessKey, ClusterMode, @@ -118,15 +120,25 @@ async def db_with_cleanup( ): yield database_connection + @pytest.fixture + def test_domain_id(self) -> DomainID: + return DomainID(uuid.uuid4()) + + @pytest.fixture + def test_scaling_group_id(self) -> ResourceGroupID: + return ResourceGroupID(uuid.uuid4()) + @pytest.fixture async def test_domain_name( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, ) -> AsyncGenerator[str, None]: domain_name = f"test-domain-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: db_sess.add( DomainRow( + id=test_domain_id, name=domain_name, total_resource_slots=ResourceSlot({ "cpu": Decimal("1000"), @@ -141,11 +153,13 @@ async def test_domain_name( async def test_scaling_group_name( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_scaling_group_id: ResourceGroupID, ) -> AsyncGenerator[str, None]: sg_name = f"test-sgroup-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: db_sess.add( ScalingGroupRow( + id=test_scaling_group_id, name=sg_name, driver="static", scheduler="fifo", @@ -334,7 +348,9 @@ async def _create_kernel_with_pending_allocations( self, db: ExtendedAsyncSAEngine, *, + domain_id: DomainID, domain_name: str, + resource_group_id: ResourceGroupID, scaling_group_name: str, group_id: uuid.UUID, user_uuid: uuid.UUID, @@ -358,8 +374,10 @@ async def _create_kernel_with_pending_allocations( id=session_id, name=f"test-session-{uuid.uuid4().hex[:8]}", session_type=SessionTypes.INTERACTIVE, + domain_id=domain_id, domain_name=domain_name, group_id=group_id, + resource_group_id=resource_group_id, scaling_group_name=scaling_group_name, status=SessionStatus.CREATING, status_info="test", @@ -451,7 +469,9 @@ async def _seed_agent_resources( async def test_sets_used_and_used_at_on_resource_allocations( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -463,7 +483,9 @@ async def test_sets_used_and_used_at_on_resource_allocations( await self._seed_agent_resources(db_with_cleanup, test_agent_id) _, kernel_id = await self._create_kernel_with_pending_allocations( db_with_cleanup, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, @@ -500,7 +522,9 @@ async def test_sets_used_and_used_at_on_resource_allocations( async def test_increments_agent_resources_used( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -514,7 +538,9 @@ async def test_increments_agent_resources_used( ) _, kernel_id = await self._create_kernel_with_pending_allocations( db_with_cleanup, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, @@ -548,7 +574,9 @@ async def test_increments_agent_resources_used( async def test_updates_kernel_status_to_running( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -560,7 +588,9 @@ async def test_updates_kernel_status_to_running( await self._seed_agent_resources(db_with_cleanup, test_agent_id) _, kernel_id = await self._create_kernel_with_pending_allocations( db_with_cleanup, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, @@ -586,7 +616,9 @@ async def test_updates_kernel_status_to_running( async def test_is_idempotent_via_double_call( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -598,7 +630,9 @@ async def test_is_idempotent_via_double_call( await self._seed_agent_resources(db_with_cleanup, test_agent_id) _, kernel_id = await self._create_kernel_with_pending_allocations( db_with_cleanup, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, @@ -635,7 +669,9 @@ async def test_is_idempotent_via_double_call( async def test_running_does_not_reject_over_capacity( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -657,7 +693,9 @@ async def test_running_does_not_reject_over_capacity( ) _, kernel_id = await self._create_kernel_with_pending_allocations( db_with_cleanup, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, @@ -692,7 +730,9 @@ async def test_running_does_not_reject_over_capacity( async def test_returns_false_when_kernel_not_in_valid_status( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -704,7 +744,9 @@ async def test_returns_false_when_kernel_not_in_valid_status( await self._seed_agent_resources(db_with_cleanup, test_agent_id) _, kernel_id = await self._create_kernel_with_pending_allocations( db_with_cleanup, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, @@ -739,7 +781,9 @@ async def test_returns_false_when_kernel_not_in_valid_status( async def test_handles_no_agent_gracefully( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -750,7 +794,9 @@ async def test_handles_no_agent_gracefully( """update_kernel_status_running succeeds without allocating when agent_id is None.""" _, kernel_id = await self._create_kernel_with_pending_allocations( db_with_cleanup, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, diff --git a/tests/unit/manager/repositories/scheduler/test_resource_concurrency.py b/tests/unit/manager/repositories/scheduler/test_resource_concurrency.py index be8edb4faf2..f6748e41622 100644 --- a/tests/unit/manager/repositories/scheduler/test_resource_concurrency.py +++ b/tests/unit/manager/repositories/scheduler/test_resource_concurrency.py @@ -15,6 +15,8 @@ import sqlalchemy as sa +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import AccessKey, KernelId, SessionId from ai.backend.manager.data.kernel.types import KernelStatus from ai.backend.manager.models.kernel import KernelRow @@ -49,7 +51,9 @@ class TestReservedConcurrency: async def test_e1_concurrent_reservation_no_overcommit( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -86,7 +90,9 @@ async def test_e1_concurrent_reservation_no_overcommit( for spec in session_specs: session_id, kernel_ids = await create_pending_session_with_kernels( db_with_cleanup, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, @@ -137,7 +143,9 @@ async def test_e1_concurrent_reservation_no_overcommit( async def test_e2_repeated_lifecycle_no_leak( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -157,7 +165,9 @@ async def test_e2_repeated_lifecycle_no_leak( async def one_cycle() -> None: session_id, kernel_ids = await create_pending_session_with_kernels( db_with_cleanup, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, @@ -198,7 +208,9 @@ async def one_cycle() -> None: async def test_e3_interleaved_allocate_free_consistency( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -218,7 +230,9 @@ async def test_e3_interleaved_allocate_free_consistency( async def allocate_session(*, run: bool) -> tuple[SessionId, KernelId]: session_id, kernel_ids = await create_pending_session_with_kernels( db_with_cleanup, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, diff --git a/tests/unit/manager/repositories/scheduler/test_resource_deallocation.py b/tests/unit/manager/repositories/scheduler/test_resource_deallocation.py index 3542ec9b7f5..a079c632ee5 100644 --- a/tests/unit/manager/repositories/scheduler/test_resource_deallocation.py +++ b/tests/unit/manager/repositories/scheduler/test_resource_deallocation.py @@ -18,6 +18,8 @@ from dateutil.tz import tzutc from ai.backend.common.data.user.types import UserRole +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ( AccessKey, ClusterMode, @@ -100,15 +102,25 @@ async def db_with_cleanup( ): yield database_connection + @pytest.fixture + def test_domain_id(self) -> DomainID: + return DomainID(uuid.uuid4()) + + @pytest.fixture + def test_scaling_group_id(self) -> ResourceGroupID: + return ResourceGroupID(uuid.uuid4()) + @pytest.fixture async def test_domain_name( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, ) -> AsyncGenerator[str, None]: domain_name = f"test-domain-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: db_sess.add( DomainRow( + id=test_domain_id, name=domain_name, total_resource_slots=ResourceSlot({ "cpu": Decimal("1000"), @@ -123,11 +135,13 @@ async def test_domain_name( async def test_scaling_group_name( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_scaling_group_id: ResourceGroupID, ) -> AsyncGenerator[str, None]: sg_name = f"test-sgroup-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: db_sess.add( ScalingGroupRow( + id=test_scaling_group_id, name=sg_name, driver="static", scheduler="fifo", @@ -319,7 +333,9 @@ async def _create_session_with_kernel_and_resources( *, session_status: SessionStatus, kernel_status: KernelStatus, + domain_id: DomainID, domain_name: str, + resource_group_id: ResourceGroupID, scaling_group_name: str, group_id: uuid.UUID, user_uuid: uuid.UUID, @@ -338,8 +354,10 @@ async def _create_session_with_kernel_and_resources( id=session_id, name=f"test-session-{uuid.uuid4().hex[:8]}", session_type=SessionTypes.INTERACTIVE, + domain_id=domain_id, domain_name=domain_name, group_id=group_id, + resource_group_id=resource_group_id, scaling_group_name=scaling_group_name, status=session_status, status_info="test", @@ -447,7 +465,9 @@ async def _create_session_with_kernel_and_resources( async def test_force_terminate_frees_resources( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -462,7 +482,9 @@ async def test_force_terminate_frees_resources( db_with_cleanup, session_status=SessionStatus.RUNNING, kernel_status=KernelStatus.RUNNING, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, @@ -515,7 +537,9 @@ async def test_force_terminate_frees_resources( async def test_force_terminate_with_null_agent_still_sets_free_at( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -529,7 +553,9 @@ async def test_force_terminate_with_null_agent_still_sets_free_at( db_with_cleanup, session_status=SessionStatus.RUNNING, kernel_status=KernelStatus.RUNNING, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, @@ -561,7 +587,9 @@ async def test_force_terminate_with_null_agent_still_sets_free_at( async def test_force_terminate_from_terminating_session_frees_resources( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -576,7 +604,9 @@ async def test_force_terminate_from_terminating_session_frees_resources( db_with_cleanup, session_status=SessionStatus.TERMINATING, kernel_status=KernelStatus.TERMINATING, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, @@ -649,15 +679,25 @@ async def db_with_cleanup( ): yield database_connection + @pytest.fixture + def test_domain_id(self) -> DomainID: + return DomainID(uuid.uuid4()) + + @pytest.fixture + def test_scaling_group_id(self) -> ResourceGroupID: + return ResourceGroupID(uuid.uuid4()) + @pytest.fixture async def test_domain_name( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, ) -> AsyncGenerator[str, None]: domain_name = f"test-domain-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: db_sess.add( DomainRow( + id=test_domain_id, name=domain_name, total_resource_slots=ResourceSlot({ "cpu": Decimal("1000"), @@ -672,11 +712,13 @@ async def test_domain_name( async def test_scaling_group_name( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_scaling_group_id: ResourceGroupID, ) -> AsyncGenerator[str, None]: sg_name = f"test-sgroup-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: db_sess.add( ScalingGroupRow( + id=test_scaling_group_id, name=sg_name, driver="static", scheduler="fifo", @@ -866,7 +908,9 @@ async def _create_kernel_with_resources( db: ExtendedAsyncSAEngine, *, kernel_status: KernelStatus, + domain_id: DomainID, domain_name: str, + resource_group_id: ResourceGroupID, scaling_group_name: str, group_id: uuid.UUID, user_uuid: uuid.UUID, @@ -885,8 +929,10 @@ async def _create_kernel_with_resources( id=session_id, name=f"test-session-{uuid.uuid4().hex[:8]}", session_type=SessionTypes.INTERACTIVE, + domain_id=domain_id, domain_name=domain_name, group_id=group_id, + resource_group_id=resource_group_id, scaling_group_name=scaling_group_name, status=SessionStatus.RUNNING, status_info="test", @@ -990,7 +1036,9 @@ async def _create_kernel_with_resources( async def test_bulk_terminate_frees_resources( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -1004,7 +1052,9 @@ async def test_bulk_terminate_frees_resources( _, kernel_id = await self._create_kernel_with_resources( db_with_cleanup, kernel_status=KernelStatus.RUNNING, + domain_id=test_domain_id, domain_name=test_domain_name, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, group_id=test_group_id, user_uuid=test_user_uuid, @@ -1083,15 +1133,25 @@ async def db_with_cleanup( ): yield database_connection + @pytest.fixture + def test_domain_id(self) -> DomainID: + return DomainID(uuid.uuid4()) + + @pytest.fixture + def test_scaling_group_id(self) -> ResourceGroupID: + return ResourceGroupID(uuid.uuid4()) + @pytest.fixture async def test_domain_name( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, ) -> AsyncGenerator[str, None]: domain_name = f"test-domain-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: db_sess.add( DomainRow( + id=test_domain_id, name=domain_name, total_resource_slots=ResourceSlot({ "cpu": Decimal("1000"), @@ -1106,11 +1166,13 @@ async def test_domain_name( async def test_scaling_group_name( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_scaling_group_id: ResourceGroupID, ) -> AsyncGenerator[str, None]: sg_name = f"test-sgroup-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: db_sess.add( ScalingGroupRow( + id=test_scaling_group_id, name=sg_name, driver="static", scheduler="fifo", @@ -1298,7 +1360,9 @@ async def resource_slot_types( async def test_double_terminate_does_not_go_negative( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, @@ -1319,8 +1383,10 @@ async def test_double_terminate_does_not_go_negative( id=session_id, name=f"test-session-{uuid.uuid4().hex[:8]}", session_type=SessionTypes.INTERACTIVE, + domain_id=test_domain_id, domain_name=test_domain_name, group_id=test_group_id, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, status=SessionStatus.RUNNING, status_info="test", diff --git a/tests/unit/manager/repositories/scheduler/test_scheduling_history_recording.py b/tests/unit/manager/repositories/scheduler/test_scheduling_history_recording.py index 39be01fa8a4..a530fa67424 100644 --- a/tests/unit/manager/repositories/scheduler/test_scheduling_history_recording.py +++ b/tests/unit/manager/repositories/scheduler/test_scheduling_history_recording.py @@ -15,6 +15,8 @@ from dateutil.tz import tzutc from ai.backend.common.data.user.types import UserRole +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ( AccessKey, ClusterMode, @@ -343,16 +345,26 @@ async def db_with_cleanup( ): yield database_connection + @pytest.fixture + def test_domain_id(self) -> DomainID: + return DomainID(uuid.uuid4()) + + @pytest.fixture + def test_scaling_group_id(self) -> ResourceGroupID: + return ResourceGroupID(uuid.uuid4()) + @pytest.fixture async def test_domain_name( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, ) -> AsyncGenerator[str, None]: """Create test domain and return domain name.""" domain_name = f"test-domain-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: domain = DomainRow( + id=test_domain_id, name=domain_name, total_resource_slots=ResourceSlot({ "cpu": Decimal("1000"), @@ -368,12 +380,14 @@ async def test_domain_name( async def test_scaling_group_name( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_scaling_group_id: ResourceGroupID, ) -> AsyncGenerator[str, None]: """Create test scaling group and return scaling group name.""" sg_name = f"test-sgroup-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: sg = ScalingGroupRow( + id=test_scaling_group_id, name=sg_name, driver="static", scheduler="fifo", @@ -569,7 +583,9 @@ async def _create_session_with_kernel( session_status: SessionStatus, kernel_status: KernelStatus, domain_name: str, + domain_id: DomainID, scaling_group_name: str, + resource_group_id: ResourceGroupID, group_id: uuid.UUID, user_uuid: uuid.UUID, access_key: AccessKey, @@ -585,8 +601,10 @@ async def _create_session_with_kernel( name=f"test-session-{uuid.uuid4().hex[:8]}", session_type=SessionTypes.INTERACTIVE, domain_name=domain_name, + domain_id=domain_id, group_id=group_id, scaling_group_name=scaling_group_name, + resource_group_id=resource_group_id, status=session_status, status_info="test", cluster_mode=ClusterMode.SINGLE_NODE, @@ -639,7 +657,9 @@ async def test_mark_sessions_as_terminating_creates_scheduling_history( self, db_with_cleanup: ExtendedAsyncSAEngine, test_domain_name: str, + test_domain_id: DomainID, test_scaling_group_name: str, + test_scaling_group_id: ResourceGroupID, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, test_access_key: AccessKey, @@ -653,7 +673,9 @@ async def test_mark_sessions_as_terminating_creates_scheduling_history( session_status=SessionStatus.RUNNING, kernel_status=KernelStatus.RUNNING, domain_name=test_domain_name, + domain_id=test_domain_id, scaling_group_name=test_scaling_group_name, + resource_group_id=test_scaling_group_id, group_id=test_group_id, user_uuid=test_user_uuid, access_key=test_access_key, @@ -680,7 +702,9 @@ async def test_cancel_pending_creates_scheduling_history( self, db_with_cleanup: ExtendedAsyncSAEngine, test_domain_name: str, + test_domain_id: DomainID, test_scaling_group_name: str, + test_scaling_group_id: ResourceGroupID, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, test_access_key: AccessKey, @@ -694,7 +718,9 @@ async def test_cancel_pending_creates_scheduling_history( session_status=SessionStatus.PENDING, kernel_status=KernelStatus.PENDING, domain_name=test_domain_name, + domain_id=test_domain_id, scaling_group_name=test_scaling_group_name, + resource_group_id=test_scaling_group_id, group_id=test_group_id, user_uuid=test_user_uuid, access_key=test_access_key, @@ -721,7 +747,9 @@ async def test_force_terminate_creates_scheduling_history( self, db_with_cleanup: ExtendedAsyncSAEngine, test_domain_name: str, + test_domain_id: DomainID, test_scaling_group_name: str, + test_scaling_group_id: ResourceGroupID, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, test_access_key: AccessKey, @@ -735,7 +763,9 @@ async def test_force_terminate_creates_scheduling_history( session_status=SessionStatus.RUNNING, kernel_status=KernelStatus.RUNNING, domain_name=test_domain_name, + domain_id=test_domain_id, scaling_group_name=test_scaling_group_name, + resource_group_id=test_scaling_group_id, group_id=test_group_id, user_uuid=test_user_uuid, access_key=test_access_key, @@ -761,7 +791,9 @@ async def test_force_terminate_from_terminating_creates_scheduling_history( self, db_with_cleanup: ExtendedAsyncSAEngine, test_domain_name: str, + test_domain_id: DomainID, test_scaling_group_name: str, + test_scaling_group_id: ResourceGroupID, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, test_access_key: AccessKey, @@ -775,7 +807,9 @@ async def test_force_terminate_from_terminating_creates_scheduling_history( session_status=SessionStatus.TERMINATING, kernel_status=KernelStatus.TERMINATING, domain_name=test_domain_name, + domain_id=test_domain_id, scaling_group_name=test_scaling_group_name, + resource_group_id=test_scaling_group_id, group_id=test_group_id, user_uuid=test_user_uuid, access_key=test_access_key, @@ -801,7 +835,9 @@ async def test_mark_sessions_as_terminating_captures_correct_from_status( self, db_with_cleanup: ExtendedAsyncSAEngine, test_domain_name: str, + test_domain_id: DomainID, test_scaling_group_name: str, + test_scaling_group_id: ResourceGroupID, test_group_id: uuid.UUID, test_user_uuid: uuid.UUID, test_access_key: AccessKey, @@ -816,7 +852,9 @@ async def test_mark_sessions_as_terminating_captures_correct_from_status( session_status=SessionStatus.RUNNING, kernel_status=KernelStatus.RUNNING, domain_name=test_domain_name, + domain_id=test_domain_id, scaling_group_name=test_scaling_group_name, + resource_group_id=test_scaling_group_id, group_id=test_group_id, user_uuid=test_user_uuid, access_key=test_access_key, @@ -827,7 +865,9 @@ async def test_mark_sessions_as_terminating_captures_correct_from_status( session_status=SessionStatus.SCHEDULED, kernel_status=KernelStatus.SCHEDULED, domain_name=test_domain_name, + domain_id=test_domain_id, scaling_group_name=test_scaling_group_name, + resource_group_id=test_scaling_group_id, group_id=test_group_id, user_uuid=test_user_uuid, access_key=test_access_key, diff --git a/tests/unit/manager/repositories/scheduler/test_session_network_propagation.py b/tests/unit/manager/repositories/scheduler/test_session_network_propagation.py index be4c3888c71..accff40a844 100644 --- a/tests/unit/manager/repositories/scheduler/test_session_network_propagation.py +++ b/tests/unit/manager/repositories/scheduler/test_session_network_propagation.py @@ -10,6 +10,8 @@ from dateutil.tz import tzutc from ai.backend.common.data.user.types import UserRole +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ( AccessKey, AgentId, @@ -80,7 +82,9 @@ async def seeded_environment( db_with_cleanup: ExtendedAsyncSAEngine, ) -> AsyncGenerator[dict[str, Any], None]: """Seed FK-required rows (domain, sgroup, policies, user, keypair, group, agent).""" + domain_id = DomainID(uuid.uuid4()) domain_name = f"test-domain-{uuid.uuid4().hex[:8]}" + sg_id = ResourceGroupID(uuid.uuid4()) sg_name = f"test-sgroup-{uuid.uuid4().hex[:8]}" project_policy_name = f"test-project-policy-{uuid.uuid4().hex[:8]}" user_policy_name = f"test-user-policy-{uuid.uuid4().hex[:8]}" @@ -93,6 +97,7 @@ async def seeded_environment( async with db_with_cleanup.begin_session() as db_sess: db_sess.add( DomainRow( + id=domain_id, name=domain_name, total_resource_slots=ResourceSlot({ "cpu": Decimal("1000"), @@ -102,6 +107,7 @@ async def seeded_environment( ) db_sess.add( ScalingGroupRow( + id=sg_id, name=sg_name, driver="static", scheduler="fifo", @@ -203,7 +209,9 @@ async def seeded_environment( await db_sess.flush() yield { + "domain_id": domain_id, "domain_name": domain_name, + "sg_id": sg_id, "sg_name": sg_name, "user_uuid": user_uuid, "access_key": access_key, @@ -228,10 +236,12 @@ async def _add_multi_node_persistent_session( id=session_id, name=f"test-session-{uuid.uuid4().hex[:8]}", session_type=SessionTypes.INTERACTIVE, + domain_id=env["domain_id"], domain_name=env["domain_name"], group_id=env["group_id"], user_uuid=env["user_uuid"], access_key=env["access_key"], + resource_group_id=env["sg_id"], scaling_group_name=env["sg_name"], status=SessionStatus.PREPARED, status_info="prepared", diff --git a/tests/unit/manager/repositories/scheduler/test_termination.py b/tests/unit/manager/repositories/scheduler/test_termination.py index 5d6a33b3a78..b17a08f7c51 100644 --- a/tests/unit/manager/repositories/scheduler/test_termination.py +++ b/tests/unit/manager/repositories/scheduler/test_termination.py @@ -13,6 +13,8 @@ from dateutil.tz import tzutc from ai.backend.common.data.user.types import UserRole +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ( AccessKey, AgentId, @@ -104,16 +106,26 @@ async def db_with_cleanup( ): yield database_connection + @pytest.fixture + def test_domain_id(self) -> DomainID: + return DomainID(uuid.uuid4()) + + @pytest.fixture + def test_scaling_group_id(self) -> ResourceGroupID: + return ResourceGroupID(uuid.uuid4()) + @pytest.fixture async def test_domain_name( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, ) -> AsyncGenerator[str, None]: """Create test domain and return domain name""" domain_name = f"test-domain-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: domain = DomainRow( + id=test_domain_id, name=domain_name, total_resource_slots=ResourceSlot({ "cpu": Decimal("1000"), @@ -129,12 +141,14 @@ async def test_domain_name( async def test_scaling_group_name( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_scaling_group_id: ResourceGroupID, ) -> AsyncGenerator[str, None]: """Create test scaling group and return scaling group name""" sg_name = f"test-sgroup-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: sg = ScalingGroupRow( + id=test_scaling_group_id, name=sg_name, driver="static", scheduler="fifo", @@ -327,7 +341,9 @@ async def test_agent_id( async def test_session_id( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, test_group_id: uuid.UUID, ) -> AsyncGenerator[SessionId, None]: @@ -339,8 +355,10 @@ async def test_session_id( id=session_id, name=f"test-session-{uuid.uuid4().hex[:8]}", session_type=SessionTypes.INTERACTIVE, + domain_id=test_domain_id, domain_name=test_domain_name, group_id=test_group_id, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, status=SessionStatus.TERMINATING, status_info="test-termination", diff --git a/tests/unit/manager/repositories/scheduler/test_update_with_history.py b/tests/unit/manager/repositories/scheduler/test_update_with_history.py index bfd55945e07..cb6fb631762 100644 --- a/tests/unit/manager/repositories/scheduler/test_update_with_history.py +++ b/tests/unit/manager/repositories/scheduler/test_update_with_history.py @@ -5,7 +5,7 @@ import uuid from collections.abc import AsyncGenerator -from datetime import datetime +from datetime import datetime, timedelta from decimal import Decimal import pytest @@ -13,6 +13,8 @@ from dateutil.tz import tzutc from ai.backend.common.data.user.types import UserRole +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ( AccessKey, ClusterMode, @@ -78,16 +80,26 @@ async def db_with_cleanup( ): yield database_connection + @pytest.fixture + def test_domain_id(self) -> DomainID: + return DomainID(uuid.uuid4()) + + @pytest.fixture + def test_scaling_group_id(self) -> ResourceGroupID: + return ResourceGroupID(uuid.uuid4()) + @pytest.fixture async def test_domain_name( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, ) -> AsyncGenerator[str, None]: """Create test domain and return domain name.""" domain_name = f"test-domain-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: domain = DomainRow( + id=test_domain_id, name=domain_name, total_resource_slots=ResourceSlot({ "cpu": Decimal("1000"), @@ -103,22 +115,29 @@ async def test_domain_name( async def test_scaling_group_name( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_scaling_group_id: ResourceGroupID, ) -> AsyncGenerator[str, None]: - """Create test scaling group and return its name.""" - scaling_group_name = f"test-sg-{uuid.uuid4().hex[:8]}" + """Create test scaling group and return scaling group name.""" + sg_name = f"test-sgroup-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: - scaling_group = ScalingGroupRow( - name=scaling_group_name, + sg = ScalingGroupRow( + id=test_scaling_group_id, + name=sg_name, driver="static", - driver_opts={}, scheduler="fifo", - scheduler_opts=ScalingGroupOpts(), + scheduler_opts=ScalingGroupOpts( + allowed_session_types=[], + pending_timeout=timedelta(hours=1), + config={}, + ), + driver_opts={}, + is_active=True, ) - db_sess.add(scaling_group) + db_sess.add(sg) await db_sess.flush() - yield scaling_group_name + yield sg_name @pytest.fixture async def test_resource_policy_name( @@ -270,9 +289,11 @@ async def test_group_id( async def test_session_id( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, - test_group_id: uuid.UUID, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, + test_group_id: uuid.UUID, ) -> AsyncGenerator[SessionId, None]: """Create test session in PREPARING status and return session ID.""" session_id = SessionId(uuid.uuid4()) @@ -283,8 +304,10 @@ async def test_session_id( creation_id=f"creation-{uuid.uuid4().hex[:8]}", name=f"test-session-{uuid.uuid4().hex[:8]}", session_type=SessionTypes.INTERACTIVE, + domain_id=test_domain_id, domain_name=test_domain_name, group_id=test_group_id, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, status=SessionStatus.PREPARING, status_info="preparing", @@ -434,9 +457,11 @@ async def test_update_with_history_failure( async def test_update_with_history_multiple_sessions( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, - test_group_id: uuid.UUID, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, + test_group_id: uuid.UUID, ) -> None: """Test update_with_history handles multiple sessions.""" # Create multiple sessions @@ -451,8 +476,10 @@ async def test_update_with_history_multiple_sessions( creation_id=f"creation-{uuid.uuid4().hex[:8]}", name=f"test-session-{uuid.uuid4().hex[:8]}", session_type=SessionTypes.INTERACTIVE, + domain_id=test_domain_id, domain_name=test_domain_name, group_id=test_group_id, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, status=SessionStatus.PREPARING, status_info="preparing", @@ -924,9 +951,11 @@ async def test_update_with_history_no_merge_different_to_status( async def test_update_with_history_merge_multiple_sessions_batch( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, test_domain_name: str, - test_group_id: uuid.UUID, + test_scaling_group_id: ResourceGroupID, test_scaling_group_name: str, + test_group_id: uuid.UUID, ) -> None: """Test merge logic works correctly with multiple sessions in batch.""" # Create multiple sessions @@ -940,8 +969,10 @@ async def test_update_with_history_merge_multiple_sessions_batch( creation_id=f"creation-{uuid.uuid4().hex[:8]}", name=f"test-session-{uuid.uuid4().hex[:8]}", session_type=SessionTypes.INTERACTIVE, + domain_id=test_domain_id, domain_name=test_domain_name, group_id=test_group_id, + resource_group_id=test_scaling_group_id, scaling_group_name=test_scaling_group_name, status=SessionStatus.PREPARING, status_info="preparing", diff --git a/tests/unit/manager/repositories/session/test_session_repository.py b/tests/unit/manager/repositories/session/test_session_repository.py index 712230285e3..0000924bb31 100644 --- a/tests/unit/manager/repositories/session/test_session_repository.py +++ b/tests/unit/manager/repositories/session/test_session_repository.py @@ -16,6 +16,8 @@ from dateutil.tz import tzutc from sqlalchemy.orm import selectinload +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.identifier.session import SessionID from ai.backend.common.types import ( AccessKey, @@ -56,7 +58,9 @@ @dataclass class SessionTestData: + domain_id: DomainID domain_name: str + scaling_group_id: ResourceGroupID user_id: uuid.UUID group_id: uuid.UUID session_id: SessionId @@ -64,6 +68,16 @@ class SessionTestData: access_key: AccessKey +@pytest.fixture +def test_domain_id() -> DomainID: + return DomainID(uuid.uuid4()) + + +@pytest.fixture +def test_scaling_group_id() -> ResourceGroupID: + return ResourceGroupID(uuid.uuid4()) + + class TestSessionRepository: """Test cases for SessionRepository using real database""" @@ -99,7 +113,12 @@ def repository(self, db_with_cleanup: ExtendedAsyncSAEngine) -> SessionRepositor return SessionRepository(db_with_cleanup) @pytest.fixture - async def session_with_kernel(self, db_with_cleanup: ExtendedAsyncSAEngine) -> SessionTestData: + async def session_with_kernel( + self, + db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, + test_scaling_group_id: ResourceGroupID, + ) -> SessionTestData: """Create a session with kernel for testing search operations.""" domain_name = "test-domain" user_id = uuid.uuid4() @@ -111,6 +130,7 @@ async def session_with_kernel(self, db_with_cleanup: ExtendedAsyncSAEngine) -> S async with db_with_cleanup.begin_session() as db_sess: # Create domain domain = DomainRow( + id=test_domain_id, name=domain_name, description="Test domain", is_active=True, @@ -123,6 +143,7 @@ async def session_with_kernel(self, db_with_cleanup: ExtendedAsyncSAEngine) -> S # Create scaling group scaling_group = ScalingGroupRow( + id=test_scaling_group_id, name="default", is_active=True, is_public=True, @@ -225,6 +246,7 @@ async def session_with_kernel(self, db_with_cleanup: ExtendedAsyncSAEngine) -> S session_type=SessionTypes.INTERACTIVE, cluster_mode=ClusterMode.SINGLE_NODE, cluster_size=1, + domain_id=test_domain_id, domain_name=domain_name, group_id=group_id, user_uuid=user_id, @@ -246,6 +268,7 @@ async def session_with_kernel(self, db_with_cleanup: ExtendedAsyncSAEngine) -> S environ=None, bootstrap_script=None, use_host_network=False, + resource_group_id=test_scaling_group_id, scaling_group_name="default", ) db_sess.add(session) @@ -303,7 +326,9 @@ async def session_with_kernel(self, db_with_cleanup: ExtendedAsyncSAEngine) -> S await db_sess.commit() return SessionTestData( + domain_id=test_domain_id, domain_name=domain_name, + scaling_group_id=test_scaling_group_id, user_id=user_id, group_id=group_id, session_id=session_id, @@ -425,6 +450,7 @@ async def _insert_session( session_type=SessionTypes.INTERACTIVE, cluster_mode=ClusterMode.SINGLE_NODE, cluster_size=1, + domain_id=base.domain_id, domain_name=base.domain_name, group_id=base.group_id, user_uuid=base.user_id, @@ -446,6 +472,7 @@ async def _insert_session( environ=None, bootstrap_script=None, use_host_network=False, + resource_group_id=base.scaling_group_id, scaling_group_name="default", ) ) @@ -590,7 +617,10 @@ async def db_with_resource_tables( @pytest.fixture async def session_with_allocations( - self, db_with_resource_tables: ExtendedAsyncSAEngine + self, + db_with_resource_tables: ExtendedAsyncSAEngine, + test_domain_id: DomainID, + test_scaling_group_id: ResourceGroupID, ) -> SessionTestData: """Create a session with resource_allocations (but empty JSONB occupying_slots).""" domain_name = "test-domain" @@ -602,6 +632,7 @@ async def session_with_allocations( async with db_with_resource_tables.begin_session() as db_sess: domain = DomainRow( + id=test_domain_id, name=domain_name, description="Test domain", is_active=True, @@ -613,6 +644,7 @@ async def session_with_allocations( db_sess.add(domain) scaling_group = ScalingGroupRow( + id=test_scaling_group_id, name="default", is_active=True, is_public=True, @@ -708,6 +740,7 @@ async def session_with_allocations( session_type=SessionTypes.INTERACTIVE, cluster_mode=ClusterMode.SINGLE_NODE, cluster_size=1, + domain_id=test_domain_id, domain_name=domain_name, group_id=group_id, user_uuid=user_id, @@ -729,6 +762,7 @@ async def session_with_allocations( environ=None, bootstrap_script=None, use_host_network=False, + resource_group_id=test_scaling_group_id, scaling_group_name="default", ) db_sess.add(session) @@ -807,7 +841,9 @@ async def session_with_allocations( await db_sess.commit() return SessionTestData( + domain_id=test_domain_id, domain_name=domain_name, + scaling_group_id=test_scaling_group_id, user_id=user_id, group_id=group_id, session_id=session_id, diff --git a/tests/unit/manager/repositories/session/test_session_search_in_project.py b/tests/unit/manager/repositories/session/test_session_search_in_project.py index 21cefce2c2a..320ae95f290 100644 --- a/tests/unit/manager/repositories/session/test_session_search_in_project.py +++ b/tests/unit/manager/repositories/session/test_session_search_in_project.py @@ -12,6 +12,8 @@ import pytest from dateutil.tz import tzutc +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ( AccessKey, ClusterMode, @@ -50,6 +52,14 @@ class TestSessionSearchInProject: """Tests for SessionRepository.search_in_project()""" + @pytest.fixture + def test_domain_id(self) -> DomainID: + return DomainID(uuid.uuid4()) + + @pytest.fixture + def test_scaling_group_id(self) -> ResourceGroupID: + return ResourceGroupID(uuid.uuid4()) + @pytest.fixture async def db_with_cleanup( self, @@ -88,6 +98,8 @@ def session_repository( async def test_data( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, + test_scaling_group_id: ResourceGroupID, ) -> AsyncGenerator[dict[str, uuid.UUID], None]: """Create two projects with sessions: project_a has 2 sessions, project_b has 1.""" domain_name = "test-domain" @@ -102,6 +114,7 @@ async def test_data( async with db_with_cleanup.begin_session() as db_sess: db_sess.add( DomainRow( + id=test_domain_id, name=domain_name, description="Test domain", is_active=True, @@ -113,6 +126,7 @@ async def test_data( ) db_sess.add( ScalingGroupRow( + id=test_scaling_group_id, name="default", is_active=True, is_public=True, @@ -221,6 +235,7 @@ async def test_data( session_type=SessionTypes.INTERACTIVE, cluster_mode=ClusterMode.SINGLE_NODE, cluster_size=1, + domain_id=test_domain_id, domain_name=domain_name, group_id=group_id, user_uuid=user_id, @@ -242,6 +257,7 @@ async def test_data( environ=None, bootstrap_script=None, use_host_network=False, + resource_group_id=test_scaling_group_id, scaling_group_name="default", ) ) diff --git a/tests/unit/manager/repositories/stream/test_stream_repository.py b/tests/unit/manager/repositories/stream/test_stream_repository.py index 603918a2e65..ea2675cd144 100644 --- a/tests/unit/manager/repositories/stream/test_stream_repository.py +++ b/tests/unit/manager/repositories/stream/test_stream_repository.py @@ -8,6 +8,8 @@ import pytest from dateutil.tz import tzutc +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ( AccessKey, ClusterMode, @@ -59,7 +61,9 @@ def _make_session_row( user_uuid: uuid.UUID, access_key: AccessKey, status: SessionStatus, + domain_id: DomainID, domain_name: str, + resource_group_id: ResourceGroupID, group_id: uuid.UUID, created_at: datetime, ) -> SessionRow: @@ -70,6 +74,7 @@ def _make_session_row( session_type=SessionTypes.INTERACTIVE, cluster_mode=ClusterMode.SINGLE_NODE, cluster_size=1, + domain_id=domain_id, domain_name=domain_name, group_id=group_id, user_uuid=user_uuid, @@ -91,6 +96,7 @@ def _make_session_row( environ=None, bootstrap_script=None, use_host_network=False, + resource_group_id=resource_group_id, scaling_group_name="default", ) @@ -157,6 +163,14 @@ def _make_kernel_row( class TestStreamRepository: """Tests for StreamRepository.get_streaming_session() using a real database.""" + @pytest.fixture + def test_domain_id(self) -> DomainID: + return DomainID(uuid.uuid4()) + + @pytest.fixture + def test_scaling_group_id(self) -> ResourceGroupID: + return ResourceGroupID(uuid.uuid4()) + @pytest.fixture async def db_with_cleanup( self, database_connection: ExtendedAsyncSAEngine @@ -188,7 +202,12 @@ def repository(self, db_with_cleanup: ExtendedAsyncSAEngine) -> StreamRepository return StreamRepository(db_with_cleanup) @pytest.fixture - async def stream_session(self, db_with_cleanup: ExtendedAsyncSAEngine) -> StreamSessionFixture: + async def stream_session( + self, + db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, + test_scaling_group_id: ResourceGroupID, + ) -> StreamSessionFixture: """ Seed a RUNNING session with two kernels (main + sub) owned by a user who has both an active and an inactive keypair, plus an unrelated user @@ -211,6 +230,7 @@ async def stream_session(self, db_with_cleanup: ExtendedAsyncSAEngine) -> Stream async with db_with_cleanup.begin_session() as db_sess: db_sess.add( DomainRow( + id=test_domain_id, name=domain_name, description="Test domain", is_active=True, @@ -222,6 +242,7 @@ async def stream_session(self, db_with_cleanup: ExtendedAsyncSAEngine) -> Stream ) db_sess.add( ScalingGroupRow( + id=test_scaling_group_id, name="default", is_active=True, is_public=True, @@ -359,7 +380,9 @@ async def stream_session(self, db_with_cleanup: ExtendedAsyncSAEngine) -> Stream user_uuid=user_uuid, access_key=inactive_access_key, status=SessionStatus.RUNNING, + domain_id=test_domain_id, domain_name=domain_name, + resource_group_id=test_scaling_group_id, group_id=group_id, created_at=now, ) @@ -372,7 +395,9 @@ async def stream_session(self, db_with_cleanup: ExtendedAsyncSAEngine) -> Stream user_uuid=user_uuid, access_key=active_access_key, status=SessionStatus.TERMINATED, + domain_id=test_domain_id, domain_name=domain_name, + resource_group_id=test_scaling_group_id, group_id=group_id, created_at=now, ) @@ -385,7 +410,9 @@ async def stream_session(self, db_with_cleanup: ExtendedAsyncSAEngine) -> Stream user_uuid=other_user_uuid, access_key=other_user_access_key, status=SessionStatus.RUNNING, + domain_id=test_domain_id, domain_name=domain_name, + resource_group_id=test_scaling_group_id, group_id=group_id, created_at=now, ) diff --git a/tests/unit/manager/services/session/test_session_service.py b/tests/unit/manager/services/session/test_session_service.py index 56bd75806e4..e17ee4020ff 100644 --- a/tests/unit/manager/services/session/test_session_service.py +++ b/tests/unit/manager/services/session/test_session_service.py @@ -16,6 +16,7 @@ from ai.backend.common.dto.agent.response import CodeCompletionResp, CodeCompletionResult from ai.backend.common.exception import InvalidAPIParameters +from ai.backend.common.identifier.domain import DomainID from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName from ai.backend.common.types import ( AccessKey, @@ -1733,6 +1734,11 @@ def requested_resource_group_id(self) -> ResourceGroupID: """Resource group id supplied directly by the caller.""" return ResourceGroupID(uuid4()) + @pytest.fixture + def session_domain_id(self) -> DomainID: + """Domain id resolved from the session's domain name.""" + return DomainID(uuid4()) + @pytest.fixture def enqueue_action_without_rg( self, @@ -1792,6 +1798,7 @@ def configured_session_service( sample_session_data: SessionData, picked_resource_group_id: ResourceGroupID, requested_resource_group_id: ResourceGroupID, + session_domain_id: DomainID, ) -> SessionService: """Wire async mocks every ``enqueue_session`` exercise needs.""" mock_session_repository.resolve_image_by_id = AsyncMock() @@ -1809,6 +1816,10 @@ def configured_session_service( mock_scheduler_repository.get_resource_group_name_by_id = AsyncMock( side_effect=lambda resource_group_id: rg_names[resource_group_id] ) + mock_scheduler_repository.get_resource_group_id_by_name = AsyncMock( + return_value=requested_resource_group_id + ) + mock_scheduler_repository.get_domain_id_by_name = AsyncMock(return_value=session_domain_id) return session_service async def test_auto_picks_default_when_resource_group_omitted( @@ -1818,9 +1829,10 @@ async def test_auto_picks_default_when_resource_group_omitted( mock_scheduling_controller: MagicMock, enqueue_action_without_rg: EnqueueSessionAction, picked_resource_group_id: ResourceGroupID, + session_domain_id: DomainID, ) -> None: """BA-5917: when no resource group is given, the default is auto-picked - by id and the draft carries the name resolved from that id. + by id and the draft carries both the id and the resolved name. """ await configured_session_service.enqueue_session(enqueue_action_without_rg) @@ -1829,7 +1841,9 @@ async def test_auto_picks_default_when_resource_group_omitted( picked_resource_group_id ) draft = mock_scheduling_controller.enqueue_session_from_draft.await_args.args[0] + assert draft.scope.resource_group_id == picked_resource_group_id assert draft.scope.resource_group_name == ResourceGroupName("picked-rg") + assert draft.scope.domain_id == session_domain_id async def test_uses_user_supplied_resource_group( self, @@ -1837,15 +1851,20 @@ async def test_uses_user_supplied_resource_group( mock_scheduler_repository: MagicMock, mock_scheduling_controller: MagicMock, enqueue_action_with_rg: EnqueueSessionAction, + requested_resource_group_id: ResourceGroupID, ) -> None: - """When the caller supplies a scaling group name, it is used as-is: - no lookups and no auto-picking. + """When the caller supplies a scaling group name, the name is used + as-is and only its id is resolved: no auto-picking. """ await configured_session_service.enqueue_session(enqueue_action_with_rg) mock_scheduler_repository.pick_default_resource_group.assert_not_called() mock_scheduler_repository.get_resource_group_name_by_id.assert_not_called() + mock_scheduler_repository.get_resource_group_id_by_name.assert_awaited_once_with( + ResourceGroupName("requested-rg") + ) draft = mock_scheduling_controller.enqueue_session_from_draft.await_args.args[0] + assert draft.scope.resource_group_id == requested_resource_group_id assert draft.scope.resource_group_name == ResourceGroupName("requested-rg") async def test_uses_user_supplied_resource_group_id( @@ -1877,8 +1896,10 @@ async def test_uses_user_supplied_resource_group_id( await configured_session_service.enqueue_session(action) mock_scheduler_repository.pick_default_resource_group.assert_not_called() + mock_scheduler_repository.get_resource_group_id_by_name.assert_not_called() mock_scheduler_repository.get_resource_group_name_by_id.assert_awaited_once_with( requested_resource_group_id ) draft = mock_scheduling_controller.enqueue_session_from_draft.await_args.args[0] + assert draft.scope.resource_group_id == requested_resource_group_id assert draft.scope.resource_group_name == ResourceGroupName("requested-rg") diff --git a/tests/unit/manager/sokovan/scheduling_controller/preparers/test_session_spec_preparer.py b/tests/unit/manager/sokovan/scheduling_controller/preparers/test_session_spec_preparer.py index e5d19e8bfab..9a7e35c61b9 100644 --- a/tests/unit/manager/sokovan/scheduling_controller/preparers/test_session_spec_preparer.py +++ b/tests/unit/manager/sokovan/scheduling_controller/preparers/test_session_spec_preparer.py @@ -22,10 +22,10 @@ import pytest -from ai.backend.common.identifier.domain import DomainName +from ai.backend.common.identifier.domain import DomainID, DomainName from ai.backend.common.identifier.image import ImageID from ai.backend.common.identifier.project import ProjectID -from ai.backend.common.identifier.resource_group import ResourceGroupName +from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName from ai.backend.common.identifier.session import SessionID from ai.backend.common.types import AccessKey, ClusterMode, ResourceSlotEntry, SessionTypes from ai.backend.manager.data.session.draft import ( @@ -120,8 +120,10 @@ def complete_draft(image_id: ImageID, minimal_kernel_group: KernelGroupDraft) -> user_uuid=uuid.uuid4(), ), scope=SessionScopeDraft( + domain_id=DomainID(uuid.uuid4()), domain_name=DomainName("default"), project_id=ProjectID(uuid.uuid4()), + resource_group_id=ResourceGroupID(uuid.uuid4()), resource_group_name=ResourceGroupName("default"), ), classification=SessionClassificationDraft( diff --git a/tests/unit/manager/sokovan/scheduling_controller/test_enqueue_session_from_draft.py b/tests/unit/manager/sokovan/scheduling_controller/test_enqueue_session_from_draft.py index 4369a6e0d3f..6efbb205cb1 100644 --- a/tests/unit/manager/sokovan/scheduling_controller/test_enqueue_session_from_draft.py +++ b/tests/unit/manager/sokovan/scheduling_controller/test_enqueue_session_from_draft.py @@ -30,10 +30,10 @@ from ai.backend.common.contexts.user import with_user from ai.backend.common.data.user.types import UserData, UserRole -from ai.backend.common.identifier.domain import DomainName +from ai.backend.common.identifier.domain import DomainID, DomainName from ai.backend.common.identifier.image import ImageID from ai.backend.common.identifier.project import ProjectID -from ai.backend.common.identifier.resource_group import ResourceGroupName +from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName from ai.backend.common.identifier.session import SessionID from ai.backend.common.identifier.vfolder import VFolderUUID from ai.backend.common.plugin.hook import PASSED, HookResult, HookResults @@ -103,8 +103,10 @@ def draft(image_id: ImageID) -> SessionSpecDraft: access_key=AccessKey("AKIAIOSFODNN7EXAMPLE"), ), scope=SessionScopeDraft( + domain_id=DomainID(uuid.uuid4()), domain_name=DomainName("default"), project_id=ProjectID(uuid.uuid4()), + resource_group_id=ResourceGroupID(uuid.uuid4()), resource_group_name=ResourceGroupName("default"), ), classification=SessionClassificationDraft(session_type=SessionTypes.INTERACTIVE), diff --git a/tests/unit/manager/sokovan/scheduling_controller/validators/test_session_spec_rules.py b/tests/unit/manager/sokovan/scheduling_controller/validators/test_session_spec_rules.py index 787ca10b171..21ad0372137 100644 --- a/tests/unit/manager/sokovan/scheduling_controller/validators/test_session_spec_rules.py +++ b/tests/unit/manager/sokovan/scheduling_controller/validators/test_session_spec_rules.py @@ -10,10 +10,10 @@ import pytest -from ai.backend.common.identifier.domain import DomainName +from ai.backend.common.identifier.domain import DomainID, DomainName from ai.backend.common.identifier.image import ImageID from ai.backend.common.identifier.project import ProjectID -from ai.backend.common.identifier.resource_group import ResourceGroupName +from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName from ai.backend.common.identifier.session import SessionID from ai.backend.common.types import ( AccessKey, @@ -146,8 +146,10 @@ def _spec( user_uuid=uuid.uuid4(), ), scope=SessionScope( + domain_id=DomainID(uuid.uuid4()), domain_name=DomainName("default"), project_id=ProjectID(uuid.uuid4()), + resource_group_id=ResourceGroupID(uuid.uuid4()), resource_group_name=ResourceGroupName("default"), ), classification=SessionClassification(session_type=session_type), diff --git a/tests/unit/manager/test_reconcile_agent_resources.py b/tests/unit/manager/test_reconcile_agent_resources.py index 01c16091612..1c08bdcccc8 100644 --- a/tests/unit/manager/test_reconcile_agent_resources.py +++ b/tests/unit/manager/test_reconcile_agent_resources.py @@ -16,6 +16,8 @@ from dateutil.tz import tzutc from ai.backend.common.auth import PublicKey, SecretKey +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.typed_validators import HostPortPair as HostPortPairModel from ai.backend.common.types import ResourceSlot, SlotName from ai.backend.manager.data.agent.types import AgentStatus @@ -40,6 +42,18 @@ from ai.backend.testutils.db import with_tables +@dataclass(frozen=True) +class SeededInfra: + """Explicit ids/names for the seeded domain, project, scaling group, and agent.""" + + domain_name: str + domain_id: DomainID + project_id: uuid.UUID + sg_name: str + resource_group_id: ResourceGroupID + agent_id: str + + class TestReconcileAgentResources: """Tests for AgentRegistry._reconcile_agent_resources().""" @@ -121,14 +135,16 @@ async def _seed_slot_types(self, db: ExtendedAsyncSAEngine) -> None: db_sess.add(ResourceSlotTypeRow(slot_name="cpu", slot_type="count", rank=0)) db_sess.add(ResourceSlotTypeRow(slot_name="mem", slot_type="bytes", rank=1)) - async def _seed_infrastructure(self, db: ExtendedAsyncSAEngine) -> tuple[str, uuid.UUID, str]: + async def _seed_infrastructure(self, db: ExtendedAsyncSAEngine) -> SeededInfra: """Create domain, project, scaling group, and agent.""" domain_name = "test-domain" + domain_id = DomainID(uuid4()) project_id = uuid4() sg_name = "test-sg" + resource_group_id = ResourceGroupID(uuid4()) agent_id = "i-test-agent" async with db.begin_session() as db_sess: - db_sess.add(DomainRow(name=domain_name)) + db_sess.add(DomainRow(id=domain_id, name=domain_name)) async with db.begin_session() as db_sess: db_sess.add( ProjectResourcePolicyRow( @@ -141,6 +157,7 @@ async def _seed_infrastructure(self, db: ExtendedAsyncSAEngine) -> tuple[str, uu async with db.begin_session() as db_sess: db_sess.add( ScalingGroupRow( + id=resource_group_id, name=sg_name, driver="static", driver_opts={}, @@ -173,7 +190,14 @@ async def _seed_infrastructure(self, db: ExtendedAsyncSAEngine) -> tuple[str, uu compute_plugins={}, ) ) - return domain_name, project_id, agent_id + return SeededInfra( + domain_name=domain_name, + domain_id=domain_id, + project_id=project_id, + sg_name=sg_name, + resource_group_id=resource_group_id, + agent_id=agent_id, + ) async def _seed_agent_resources( self, @@ -197,9 +221,7 @@ async def _create_kernel_with_allocations( self, db: ExtendedAsyncSAEngine, *, - domain_name: str, - project_id: uuid.UUID, - agent_id: str, + infra: SeededInfra, status: KernelStatus, allocations: dict[str, tuple[Decimal, Decimal | None]], free_at: datetime | None = None, @@ -211,10 +233,12 @@ async def _create_kernel_with_allocations( db_sess.add( SessionRow( id=session_id, - domain_name=domain_name, - group_id=project_id, - scaling_group_name="test-sg", + domain_id=infra.domain_id, + domain_name=infra.domain_name, + group_id=infra.project_id, user_uuid=uuid4(), + resource_group_id=infra.resource_group_id, + scaling_group_name=infra.sg_name, occupying_slots=empty_slots, requested_slots=empty_slots, ) @@ -224,8 +248,8 @@ async def _create_kernel_with_allocations( KernelRow( id=kernel_id, session_id=session_id, - domain_name=domain_name, - group_id=project_id, + domain_name=infra.domain_name, + group_id=infra.project_id, user_uuid=uuid4(), status=status, occupied_slots=empty_slots, @@ -234,8 +258,8 @@ async def _create_kernel_with_allocations( repl_out_port=0, stdin_port=0, stdout_port=0, - scaling_group="test-sg", - agent=agent_id, + scaling_group=infra.sg_name, + agent=infra.agent_id, ) ) await db_sess.flush() @@ -275,14 +299,12 @@ async def test_no_drift_no_corrections( """When agent_resources.used matches actual allocations, no corrections are made.""" db = db_with_tables await self._seed_slot_types(db) - domain_name, project_id, agent_id = await self._seed_infrastructure(db) + infra = await self._seed_infrastructure(db) # Actual usage: cpu=2, mem=1024 await self._create_kernel_with_allocations( db, - domain_name=domain_name, - project_id=project_id, - agent_id=agent_id, + infra=infra, status=KernelStatus.RUNNING, allocations={ "cpu": (Decimal("2"), Decimal("2")), @@ -293,7 +315,7 @@ async def test_no_drift_no_corrections( # Set agent_resources.used to match actual await self._seed_agent_resources( db, - agent_id, + infra.agent_id, { "cpu": (Decimal("8"), Decimal("2")), "mem": (Decimal("32768"), Decimal("1024")), @@ -305,7 +327,7 @@ async def test_no_drift_no_corrections( assert "agent_resources drift detected" not in caplog.text - used = await self._get_agent_resource_used(db, agent_id) + used = await self._get_agent_resource_used(db, infra.agent_id) assert used["cpu"] == Decimal("2") assert used["mem"] == Decimal("1024") @@ -318,14 +340,12 @@ async def test_over_count_corrected_down( """When agent_resources.used > actual, it is corrected down.""" db = db_with_tables await self._seed_slot_types(db) - domain_name, project_id, agent_id = await self._seed_infrastructure(db) + infra = await self._seed_infrastructure(db) # Actual usage: cpu=2 await self._create_kernel_with_allocations( db, - domain_name=domain_name, - project_id=project_id, - agent_id=agent_id, + infra=infra, status=KernelStatus.RUNNING, allocations={"cpu": (Decimal("2"), Decimal("2"))}, ) @@ -333,7 +353,7 @@ async def test_over_count_corrected_down( # Tracked as 100 but actual is 2 await self._seed_agent_resources( db, - agent_id, + infra.agent_id, { "cpu": (Decimal("200"), Decimal("100")), }, @@ -343,9 +363,9 @@ async def test_over_count_corrected_down( await registry._reconcile_agent_resources() assert "agent_resources drift detected" in caplog.text - assert agent_id in caplog.text + assert infra.agent_id in caplog.text - used = await self._get_agent_resource_used(db, agent_id) + used = await self._get_agent_resource_used(db, infra.agent_id) assert used["cpu"] == Decimal("2") async def test_under_count_corrected_up( @@ -357,22 +377,18 @@ async def test_under_count_corrected_up( """When agent_resources.used < actual, it is corrected up.""" db = db_with_tables await self._seed_slot_types(db) - domain_name, project_id, agent_id = await self._seed_infrastructure(db) + infra = await self._seed_infrastructure(db) # Actual usage: cpu=6 (two kernels) await self._create_kernel_with_allocations( db, - domain_name=domain_name, - project_id=project_id, - agent_id=agent_id, + infra=infra, status=KernelStatus.RUNNING, allocations={"cpu": (Decimal("4"), Decimal("4"))}, ) await self._create_kernel_with_allocations( db, - domain_name=domain_name, - project_id=project_id, - agent_id=agent_id, + infra=infra, status=KernelStatus.RUNNING, allocations={"cpu": (Decimal("2"), Decimal("2"))}, ) @@ -380,7 +396,7 @@ async def test_under_count_corrected_up( # Tracked as 1 but actual is 6 await self._seed_agent_resources( db, - agent_id, + infra.agent_id, { "cpu": (Decimal("8"), Decimal("1")), }, @@ -391,7 +407,7 @@ async def test_under_count_corrected_up( assert "agent_resources drift detected" in caplog.text - used = await self._get_agent_resource_used(db, agent_id) + used = await self._get_agent_resource_used(db, infra.agent_id) assert used["cpu"] == Decimal("6") async def test_no_allocations_resets_used_to_zero( @@ -403,12 +419,12 @@ async def test_no_allocations_resets_used_to_zero( """When there are no active allocations but used > 0, it is corrected to 0.""" db = db_with_tables await self._seed_slot_types(db) - domain_name, _, agent_id = await self._seed_infrastructure(db) + infra = await self._seed_infrastructure(db) # No allocations exist, but agent_resources says used=50 await self._seed_agent_resources( db, - agent_id, + infra.agent_id, { "cpu": (Decimal("8"), Decimal("50")), }, @@ -419,7 +435,7 @@ async def test_no_allocations_resets_used_to_zero( assert "agent_resources drift detected" in caplog.text - used = await self._get_agent_resource_used(db, agent_id) + used = await self._get_agent_resource_used(db, infra.agent_id) assert used["cpu"] == Decimal("0") @@ -515,18 +531,20 @@ async def registry( async def infra( self, db: ExtendedAsyncSAEngine, - ) -> tuple[str, uuid.UUID, str]: + ) -> SeededInfra: """Seed slot types, domain, project, scaling group, agent.""" domain_name = "test-domain" + domain_id = DomainID(uuid4()) project_id = uuid4() sg_name = "test-sg" + resource_group_id = ResourceGroupID(uuid4()) agent_id = "i-test-agent" async with db.begin_session() as db_sess: db_sess.add(ResourceSlotTypeRow(slot_name="cpu", slot_type="count", rank=0)) db_sess.add(ResourceSlotTypeRow(slot_name="mem", slot_type="bytes", rank=1)) async with db.begin_session() as db_sess: - db_sess.add(DomainRow(name=domain_name)) + db_sess.add(DomainRow(id=domain_id, name=domain_name)) async with db.begin_session() as db_sess: db_sess.add( ProjectResourcePolicyRow( @@ -539,6 +557,7 @@ async def infra( async with db.begin_session() as db_sess: db_sess.add( ScalingGroupRow( + id=resource_group_id, name=sg_name, driver="static", driver_opts={}, @@ -571,18 +590,24 @@ async def infra( compute_plugins={}, ) ) - return domain_name, project_id, agent_id + return SeededInfra( + domain_name=domain_name, + domain_id=domain_id, + project_id=project_id, + sg_name=sg_name, + resource_group_id=resource_group_id, + agent_id=agent_id, + ) @pytest.fixture async def kernel( self, request: pytest.FixtureRequest, db: ExtendedAsyncSAEngine, - infra: tuple[str, uuid.UUID, str], + infra: SeededInfra, ) -> uuid.UUID: """Create a kernel from a KernelSpec passed via indirect parametrize.""" spec: KernelSpec = request.param - domain_name, project_id, agent_id = infra session_id = uuid4() kernel_id = uuid4() empty_slots = ResourceSlot({}) @@ -590,10 +615,12 @@ async def kernel( db_sess.add( SessionRow( id=session_id, - domain_name=domain_name, - group_id=project_id, - scaling_group_name="test-sg", + domain_id=infra.domain_id, + domain_name=infra.domain_name, + group_id=infra.project_id, user_uuid=uuid4(), + resource_group_id=infra.resource_group_id, + scaling_group_name=infra.sg_name, occupying_slots=empty_slots, requested_slots=empty_slots, ) @@ -603,8 +630,8 @@ async def kernel( KernelRow( id=kernel_id, session_id=session_id, - domain_name=domain_name, - group_id=project_id, + domain_name=infra.domain_name, + group_id=infra.project_id, user_uuid=uuid4(), status=spec.status, occupied_slots=empty_slots, @@ -613,8 +640,8 @@ async def kernel( repl_out_port=0, stdin_port=0, stdout_port=0, - scaling_group="test-sg", - agent=agent_id, + scaling_group=infra.sg_name, + agent=infra.agent_id, ) ) await db_sess.flush() @@ -631,7 +658,7 @@ async def kernel( async with db.begin_session() as db_sess: db_sess.add( AgentResourceRow( - agent_id=agent_id, + agent_id=infra.agent_id, slot_name="cpu", capacity=Decimal("8"), used=Decimal("0"), @@ -751,13 +778,12 @@ async def test_already_freed_allocation_not_touched( async def orphan_with_drift( self, db: ExtendedAsyncSAEngine, - infra: tuple[str, uuid.UUID, str], + infra: SeededInfra, ) -> tuple[uuid.UUID, str]: """RUNNING kernel(cpu=2) + CANCELLED orphan(cpu=4) + agent used=6 (stale). Returns (orphan_kernel_id, agent_id). """ - domain_name, project_id, agent_id = infra empty_slots = ResourceSlot({}) # RUNNING kernel: legitimate cpu=2 @@ -766,10 +792,12 @@ async def orphan_with_drift( db_sess.add( SessionRow( id=running_sid, - domain_name=domain_name, - group_id=project_id, - scaling_group_name="test-sg", + domain_id=infra.domain_id, + domain_name=infra.domain_name, + group_id=infra.project_id, user_uuid=uuid4(), + resource_group_id=infra.resource_group_id, + scaling_group_name=infra.sg_name, occupying_slots=empty_slots, requested_slots=empty_slots, ) @@ -779,8 +807,8 @@ async def orphan_with_drift( KernelRow( id=running_kid, session_id=running_sid, - domain_name=domain_name, - group_id=project_id, + domain_name=infra.domain_name, + group_id=infra.project_id, user_uuid=uuid4(), status=KernelStatus.RUNNING, occupied_slots=empty_slots, @@ -789,8 +817,8 @@ async def orphan_with_drift( repl_out_port=0, stdin_port=0, stdout_port=0, - scaling_group="test-sg", - agent=agent_id, + scaling_group=infra.sg_name, + agent=infra.agent_id, ) ) await db_sess.flush() @@ -809,10 +837,12 @@ async def orphan_with_drift( db_sess.add( SessionRow( id=orphan_sid, - domain_name=domain_name, - group_id=project_id, - scaling_group_name="test-sg", + domain_id=infra.domain_id, + domain_name=infra.domain_name, + group_id=infra.project_id, user_uuid=uuid4(), + resource_group_id=infra.resource_group_id, + scaling_group_name=infra.sg_name, occupying_slots=empty_slots, requested_slots=empty_slots, ) @@ -822,8 +852,8 @@ async def orphan_with_drift( KernelRow( id=orphan_kid, session_id=orphan_sid, - domain_name=domain_name, - group_id=project_id, + domain_name=infra.domain_name, + group_id=infra.project_id, user_uuid=uuid4(), status=KernelStatus.CANCELLED, occupied_slots=empty_slots, @@ -832,8 +862,8 @@ async def orphan_with_drift( repl_out_port=0, stdin_port=0, stdout_port=0, - scaling_group="test-sg", - agent=agent_id, + scaling_group=infra.sg_name, + agent=infra.agent_id, ) ) await db_sess.flush() @@ -849,14 +879,14 @@ async def orphan_with_drift( async with db.begin_session() as db_sess: db_sess.add( AgentResourceRow( - agent_id=agent_id, + agent_id=infra.agent_id, slot_name="cpu", capacity=Decimal("8"), used=Decimal("6"), ) ) - return orphan_kid, agent_id + return orphan_kid, infra.agent_id async def test_orphan_cleanup_runs_before_drift_correction( self, @@ -971,16 +1001,18 @@ async def registry( async def infra( self, db: ExtendedAsyncSAEngine, - ) -> tuple[str, uuid.UUID, str]: + ) -> SeededInfra: domain_name = "test-domain" + domain_id = DomainID(uuid4()) project_id = uuid4() sg_name = "test-sg" + resource_group_id = ResourceGroupID(uuid4()) agent_id = "i-test-agent" async with db.begin_session() as db_sess: db_sess.add(ResourceSlotTypeRow(slot_name="cpu", slot_type="count", rank=0)) db_sess.add(ResourceSlotTypeRow(slot_name="mem", slot_type="bytes", rank=1)) async with db.begin_session() as db_sess: - db_sess.add(DomainRow(name=domain_name)) + db_sess.add(DomainRow(id=domain_id, name=domain_name)) async with db.begin_session() as db_sess: db_sess.add( ProjectResourcePolicyRow( @@ -993,6 +1025,7 @@ async def infra( async with db.begin_session() as db_sess: db_sess.add( ScalingGroupRow( + id=resource_group_id, name=sg_name, driver="static", driver_opts={}, @@ -1025,12 +1058,19 @@ async def infra( compute_plugins={}, ) ) - return domain_name, project_id, agent_id + return SeededInfra( + domain_name=domain_name, + domain_id=domain_id, + project_id=project_id, + sg_name=sg_name, + resource_group_id=resource_group_id, + agent_id=agent_id, + ) async def _seed_drift( self, db: ExtendedAsyncSAEngine, - infra: tuple[str, uuid.UUID, str], + infra: SeededInfra, *, session_status: SessionStatus, kernel_status: KernelStatus, @@ -1043,7 +1083,6 @@ async def _seed_drift( Returns (session_id, kernel_id, agent_id). """ - domain_name, project_id, agent_id = infra empty_slots = ResourceSlot({}) session_id = uuid4() kernel_id = uuid4() @@ -1051,10 +1090,12 @@ async def _seed_drift( db_sess.add( SessionRow( id=session_id, - domain_name=domain_name, - group_id=project_id, - scaling_group_name="test-sg", + domain_id=infra.domain_id, + domain_name=infra.domain_name, + group_id=infra.project_id, user_uuid=uuid4(), + resource_group_id=infra.resource_group_id, + scaling_group_name=infra.sg_name, status=session_status, occupying_slots=empty_slots, requested_slots=empty_slots, @@ -1065,8 +1106,8 @@ async def _seed_drift( KernelRow( id=kernel_id, session_id=session_id, - domain_name=domain_name, - group_id=project_id, + domain_name=infra.domain_name, + group_id=infra.project_id, user_uuid=uuid4(), status=kernel_status, occupied_slots=empty_slots, @@ -1075,8 +1116,8 @@ async def _seed_drift( repl_out_port=0, stdin_port=0, stdout_port=0, - scaling_group="test-sg", - agent=agent_id, + scaling_group=infra.sg_name, + agent=infra.agent_id, ) ) await db_sess.flush() @@ -1091,13 +1132,13 @@ async def _seed_drift( async with db.begin_session() as db_sess: db_sess.add( AgentResourceRow( - agent_id=agent_id, + agent_id=infra.agent_id, slot_name="cpu", capacity=Decimal("8"), used=agent_used, ) ) - return session_id, kernel_id, agent_id + return session_id, kernel_id, infra.agent_id @pytest.mark.parametrize( ("session_status", "kernel_status"), @@ -1123,7 +1164,7 @@ async def test_drift_kernel_is_force_cancelled( self, db: ExtendedAsyncSAEngine, registry: AgentRegistry, - infra: tuple[str, uuid.UUID, str], + infra: SeededInfra, session_status: SessionStatus, kernel_status: KernelStatus, ) -> None: @@ -1171,7 +1212,7 @@ async def test_non_terminal_session_kernel_left_untouched( self, db: ExtendedAsyncSAEngine, registry: AgentRegistry, - infra: tuple[str, uuid.UUID, str], + infra: SeededInfra, ) -> None: """Session is still RUNNING — kernel state must not be touched by Step 1.""" _, kernel_id, _ = await self._seed_drift( From 0320b1a341ffc3acf053f5b09476c3a28ee92978 Mon Sep 17 00:00:00 2001 From: BoKeum Date: Fri, 3 Jul 2026 14:12:49 +0900 Subject: [PATCH 3/6] test(BA-6644): add session scope ids to remaining test fixtures Co-Authored-By: Claude Fable 5 --- .../manager/models/resource_slot/conftest.py | 21 ++- .../domain/test_domain_purgers.py | 5 + .../domain/test_domain_repository.py | 9 +- .../group/test_group_db_source.py | 18 +++ .../repositories/group/test_group_purgers.py | 18 ++- .../group/test_group_repository.py | 22 ++++ .../resource_slot/test_db_source.py | 124 ++++++++++++++++-- .../scaling_group/test_resource_info.py | 22 ++++ tests/unit/manager/repositories/test_utils.py | 9 +- 9 files changed, 231 insertions(+), 17 deletions(-) diff --git a/tests/unit/manager/models/resource_slot/conftest.py b/tests/unit/manager/models/resource_slot/conftest.py index 333ed7d312a..1a312bc5d82 100644 --- a/tests/unit/manager/models/resource_slot/conftest.py +++ b/tests/unit/manager/models/resource_slot/conftest.py @@ -6,6 +6,8 @@ import pytest +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import BinarySize, ResourceSlot from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm from ai.backend.manager.models.agent import AgentRow @@ -67,24 +69,37 @@ async def database_with_resource_slot_tables( yield database_connection +@pytest.fixture +def domain_id() -> DomainID: + return DomainID(uuid.uuid4()) + + +@pytest.fixture +def scaling_group_id() -> ResourceGroupID: + return ResourceGroupID(uuid.uuid4()) + + @pytest.fixture async def domain_name( database_with_resource_slot_tables: ExtendedAsyncSAEngine, + domain_id: DomainID, ) -> AsyncGenerator[str, None]: name = "test-domain" async with database_with_resource_slot_tables.begin_session() as db_sess: - db_sess.add(DomainRow(name=name)) + db_sess.add(DomainRow(id=domain_id, name=name)) yield name @pytest.fixture async def scaling_group( database_with_resource_slot_tables: ExtendedAsyncSAEngine, + scaling_group_id: ResourceGroupID, ) -> AsyncGenerator[str, None]: name = "default" async with database_with_resource_slot_tables.begin_session() as db_sess: db_sess.add( ScalingGroupRow( + id=scaling_group_id, name=name, description="Test scaling group", is_active=True, @@ -208,9 +223,11 @@ async def agent_id( @pytest.fixture async def kernel_id( database_with_resource_slot_tables: ExtendedAsyncSAEngine, + domain_id: DomainID, domain_name: str, project_id: uuid.UUID, user_uuid: uuid.UUID, + scaling_group_id: ResourceGroupID, scaling_group: str, agent_id: str, ) -> AsyncGenerator[uuid.UUID, None]: @@ -220,8 +237,10 @@ async def kernel_id( db_sess.add( SessionRow( id=sid, + domain_id=domain_id, domain_name=domain_name, group_id=project_id, + resource_group_id=scaling_group_id, scaling_group_name=scaling_group, user_uuid=user_uuid, occupying_slots=ResourceSlot({"cpu": Decimal("1"), "mem": Decimal("1073741824")}), diff --git a/tests/unit/manager/repositories/domain/test_domain_purgers.py b/tests/unit/manager/repositories/domain/test_domain_purgers.py index 7ee5d2e17c0..52c0c15447b 100644 --- a/tests/unit/manager/repositories/domain/test_domain_purgers.py +++ b/tests/unit/manager/repositories/domain/test_domain_purgers.py @@ -12,6 +12,7 @@ import pytest import sqlalchemy as sa +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ResourceSlot, VFolderHostPermissionMap from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm from ai.backend.manager.data.kernel.types import KernelStatus @@ -182,9 +183,11 @@ async def sample_sessions( """Create test sessions belonging to the domain.""" sessions: list[SessionRow] = [] sgroup_name = f"default-{uuid.uuid4().hex[:8]}" + sgroup_id = ResourceGroupID(uuid.uuid4()) async with db_with_cleanup.begin_session() as session: sgroup = ScalingGroupRow( name=sgroup_name, + id=sgroup_id, description="Test scaling group", is_active=True, driver="static", @@ -201,8 +204,10 @@ async def sample_sessions( cluster_mode="single-node", cluster_size=1, domain_name=sample_domain.domain_name, + domain_id=sample_domain.domain_id, group_id=sample_group.id, scaling_group_name=sgroup_name, + resource_group_id=sgroup_id, user_uuid=sample_user.uuid, occupying_slots=ResourceSlot({}), requested_slots=ResourceSlot({}), diff --git a/tests/unit/manager/repositories/domain/test_domain_repository.py b/tests/unit/manager/repositories/domain/test_domain_repository.py index 4246a9daf1b..a75878fdb67 100644 --- a/tests/unit/manager/repositories/domain/test_domain_repository.py +++ b/tests/unit/manager/repositories/domain/test_domain_repository.py @@ -12,7 +12,8 @@ import sqlalchemy as sa from ai.backend.common.exception import DomainNotFound, InvalidAPIParameters -from ai.backend.common.identifier.domain import DomainName +from ai.backend.common.identifier.domain import DomainID, DomainName +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ( DefaultForUnspecified, ResourceSlot, @@ -333,13 +334,16 @@ async def domain_with_active_kernel( ) -> str: """Create an inactive domain with an active kernel for purge testing.""" domain_name = f"domain-with-kernel-{uuid.uuid4().hex[:8]}" + domain_id = DomainID(uuid.uuid4()) user_uuid = uuid.uuid4() group_id = uuid.uuid4() session_id = uuid.uuid4() sgroup_name = f"default-{uuid.uuid4().hex[:8]}" + sgroup_id = ResourceGroupID(uuid.uuid4()) async with db_with_default_resource_policies.begin_session() as session: domain = DomainRow( + id=domain_id, name=domain_name, description="Test domain with active kernels", is_active=False, @@ -352,6 +356,7 @@ async def domain_with_active_kernel( session.add(domain) sgroup = ScalingGroupRow( + id=sgroup_id, name=sgroup_name, description="Test scaling group", is_active=True, @@ -403,8 +408,10 @@ async def domain_with_active_kernel( creation_id=str(uuid.uuid4()).replace("-", ""), cluster_size=1, domain_name=domain_name, + domain_id=domain_id, group_id=group_id, scaling_group_name=sgroup_name, + resource_group_id=sgroup_id, user_uuid=user_uuid, occupying_slots=ResourceSlot(), requested_slots=ResourceSlot(), diff --git a/tests/unit/manager/repositories/group/test_group_db_source.py b/tests/unit/manager/repositories/group/test_group_db_source.py index 1c9308e9c59..08149275b82 100644 --- a/tests/unit/manager/repositories/group/test_group_db_source.py +++ b/tests/unit/manager/repositories/group/test_group_db_source.py @@ -10,6 +10,8 @@ import sqlalchemy as sa from ai.backend.common.identifier.deployment import DeploymentID +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ResourceSlot, VFolderHostPermissionMap from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm from ai.backend.manager.data.group.types import ProjectType @@ -111,16 +113,22 @@ async def db_with_cleanup( ): yield database_connection + @pytest.fixture + def test_domain_id(self) -> DomainID: + return DomainID(uuid.uuid4()) + @pytest.fixture async def test_domain( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, ) -> str: """Create test domain""" domain_name = f"test-domain-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as session: domain = DomainRow( + id=test_domain_id, name=domain_name, description="Test domain", is_active=True, @@ -290,17 +298,20 @@ async def inactive_endpoint_with_session_and_routing( self, db_with_cleanup: ExtendedAsyncSAEngine, test_domain: str, + test_domain_id: DomainID, test_user: uuid.UUID, test_group: uuid.UUID, ) -> EndpointWithSessionFixtureData: """Create one inactive endpoint with a session and routing entry""" endpoint_id = DeploymentID(uuid.uuid4()) session_id = uuid.uuid4() + sgroup_id = ResourceGroupID(uuid.uuid4()) sgroup_name = f"default-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as session: # Create scaling group sgroup = ScalingGroupRow( + id=sgroup_id, name=sgroup_name, description="Test scaling group", is_active=True, @@ -332,8 +343,10 @@ async def inactive_endpoint_with_session_and_routing( id=session_id, creation_id=f"test-session-{uuid.uuid4().hex[:8]}", domain_name=test_domain, + domain_id=test_domain_id, group_id=test_group, scaling_group_name=sgroup_name, + resource_group_id=sgroup_id, user_uuid=test_user, access_key="test-access-key", cluster_mode="single-node", @@ -413,17 +426,20 @@ async def multiple_endpoints_with_sessions( self, db_with_cleanup: ExtendedAsyncSAEngine, test_domain: str, + test_domain_id: DomainID, test_user: uuid.UUID, test_group: uuid.UUID, ) -> MultipleEndpointsWithSessionsFixtureData: """Create 3 inactive endpoints, each with a session and routing entry""" endpoint_ids = [] session_ids = [] + sgroup_id = ResourceGroupID(uuid.uuid4()) sgroup_name = f"default-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as session: # Create scaling group sgroup = ScalingGroupRow( + id=sgroup_id, name=sgroup_name, description="Test scaling group", is_active=True, @@ -459,8 +475,10 @@ async def multiple_endpoints_with_sessions( id=session_id, creation_id=f"test-session-{i}-{uuid.uuid4().hex[:8]}", domain_name=test_domain, + domain_id=test_domain_id, group_id=test_group, scaling_group_name=sgroup_name, + resource_group_id=sgroup_id, user_uuid=test_user, access_key=f"test-access-key-{i}", cluster_mode="single-node", diff --git a/tests/unit/manager/repositories/group/test_group_purgers.py b/tests/unit/manager/repositories/group/test_group_purgers.py index d07c0753659..1ea180804e5 100644 --- a/tests/unit/manager/repositories/group/test_group_purgers.py +++ b/tests/unit/manager/repositories/group/test_group_purgers.py @@ -12,6 +12,7 @@ import pytest import sqlalchemy as sa +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ResourceSlot from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm from ai.backend.manager.data.kernel.types import KernelStatus @@ -185,12 +186,21 @@ async def sample_group( return group @pytest.fixture - async def sample_scaling_group(self, db_with_cleanup: ExtendedAsyncSAEngine) -> str: + def sample_scaling_group_id(self) -> ResourceGroupID: + return ResourceGroupID(uuid.uuid4()) + + @pytest.fixture + async def sample_scaling_group( + self, + db_with_cleanup: ExtendedAsyncSAEngine, + sample_scaling_group_id: ResourceGroupID, + ) -> str: """Create a test scaling group.""" sgroup_name = f"test-sgroup-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as session: sgroup = ScalingGroupRow( name=sgroup_name, + id=sample_scaling_group_id, driver="static", driver_opts={}, scheduler="fifo", @@ -207,6 +217,7 @@ async def sample_sessions( sample_domain: DomainFixtureData, sample_user: UserRow, sample_scaling_group: str, + sample_scaling_group_id: ResourceGroupID, ) -> list[SessionRow]: """Create test sessions belonging to the group.""" sessions: list[SessionRow] = [] @@ -218,8 +229,10 @@ async def sample_sessions( cluster_mode="single-node", cluster_size=1, domain_name=sample_domain.domain_name, + domain_id=sample_domain.domain_id, group_id=sample_group.id, scaling_group_name=sample_scaling_group, + resource_group_id=sample_scaling_group_id, user_uuid=sample_user.uuid, occupying_slots=ResourceSlot({}), requested_slots=ResourceSlot({}), @@ -310,6 +323,7 @@ async def sample_endpoint_sessions_with_routings( sample_group: GroupRow, sample_user: UserRow, sample_scaling_group: str, + sample_scaling_group_id: ResourceGroupID, ) -> tuple[list[SessionRow], list[RoutingRow]]: """Create test sessions connected to endpoints via routings.""" sessions: list[SessionRow] = [] @@ -323,8 +337,10 @@ async def sample_endpoint_sessions_with_routings( cluster_mode="single-node", cluster_size=1, domain_name=sample_domain.domain_name, + domain_id=sample_domain.domain_id, group_id=sample_group.id, scaling_group_name=sample_scaling_group, + resource_group_id=sample_scaling_group_id, user_uuid=sample_user.uuid, occupying_slots=ResourceSlot({}), requested_slots=ResourceSlot({}), diff --git a/tests/unit/manager/repositories/group/test_group_repository.py b/tests/unit/manager/repositories/group/test_group_repository.py index f885b4c23d7..b7e75439f75 100644 --- a/tests/unit/manager/repositories/group/test_group_repository.py +++ b/tests/unit/manager/repositories/group/test_group_repository.py @@ -13,7 +13,9 @@ from ai.backend.common.data.permission.types import EntityType, ScopeType from ai.backend.common.exception import InvalidAPIParameters +from ai.backend.common.identifier.domain import DomainID from ai.backend.common.identifier.project import ProjectID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ( QuotaScopeID, QuotaScopeType, @@ -280,16 +282,22 @@ async def db_with_cleanup( ): yield database_connection + @pytest.fixture + def test_domain_id(self) -> DomainID: + return DomainID(uuid.uuid4()) + @pytest.fixture async def test_domain( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, ) -> str: """Create test domain""" domain_name = f"test-domain-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as session: domain = DomainRow( + id=test_domain_id, name=domain_name, description="Test domain", is_active=True, @@ -523,16 +531,22 @@ async def group_repository_with_mock_role_manager( repo._db_source = group_db_source_with_mock_role_manager return repo + @pytest.fixture + def test_scaling_group_id(self) -> ResourceGroupID: + return ResourceGroupID(uuid.uuid4()) + @pytest.fixture async def test_scaling_group( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_scaling_group_id: ResourceGroupID, ) -> str: """Create test scaling group""" sgroup_name = f"test-sgroup-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as session: sgroup = ScalingGroupRow( + id=test_scaling_group_id, name=sgroup_name, description="Test scaling group", is_active=True, @@ -551,9 +565,11 @@ async def group_with_active_kernel( self, db_with_cleanup: ExtendedAsyncSAEngine, test_domain: str, + test_domain_id: DomainID, test_user: uuid.UUID, default_project_resource_policy: str, test_scaling_group: str, + test_scaling_group_id: ResourceGroupID, ) -> uuid.UUID: """Create a group with an active kernel""" group_id = uuid.uuid4() @@ -597,8 +613,10 @@ async def group_with_active_kernel( id=session_id, creation_id=f"test-session-{uuid.uuid4().hex[:8]}", domain_name=test_domain, + domain_id=test_domain_id, group_id=group_id, scaling_group_name=test_scaling_group, + resource_group_id=test_scaling_group_id, user_uuid=test_user, access_key="test-access-key", cluster_mode="single-node", @@ -694,9 +712,11 @@ async def group_with_mounted_vfolders( self, db_with_cleanup: ExtendedAsyncSAEngine, test_domain: str, + test_domain_id: DomainID, test_user: uuid.UUID, default_project_resource_policy: str, test_scaling_group: str, + test_scaling_group_id: ResourceGroupID, ) -> uuid.UUID: """Create a group with vfolders mounted to active kernels""" group_id = uuid.uuid4() @@ -758,8 +778,10 @@ async def group_with_mounted_vfolders( id=session_id, creation_id=f"test-session-{uuid.uuid4().hex[:8]}", domain_name=test_domain, + domain_id=test_domain_id, group_id=group_id, scaling_group_name=test_scaling_group, + resource_group_id=test_scaling_group_id, user_uuid=test_user, access_key="test-access-key", cluster_mode="single-node", diff --git a/tests/unit/manager/repositories/resource_slot/test_db_source.py b/tests/unit/manager/repositories/resource_slot/test_db_source.py index 3065b342b7c..8aa2d12e489 100644 --- a/tests/unit/manager/repositories/resource_slot/test_db_source.py +++ b/tests/unit/manager/repositories/resource_slot/test_db_source.py @@ -11,6 +11,8 @@ import pytest from dateutil.tz import tzutc +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ResourceSlot, SlotName from ai.backend.manager.data.agent.types import AgentStatus from ai.backend.manager.data.kernel.types import KernelStatus @@ -248,6 +250,14 @@ async def test_search_resource_allocations_empty( class TestAggregation: """Tests for aggregate_occupied_by_domain and aggregate_occupied_by_project.""" + @pytest.fixture + def domain_id(self) -> DomainID: + return DomainID(uuid4()) + + @pytest.fixture + def scaling_group_id(self) -> ResourceGroupID: + return ResourceGroupID(uuid4()) + @pytest.fixture async def db_with_full_tables( self, @@ -276,14 +286,19 @@ async def _seed_slot_types(self, db: ExtendedAsyncSAEngine) -> None: db_sess.add(ResourceSlotTypeRow(slot_name="cpu", slot_type="count", rank=0)) db_sess.add(ResourceSlotTypeRow(slot_name="mem", slot_type="bytes", rank=1)) - async def _seed_infrastructure(self, db: ExtendedAsyncSAEngine) -> tuple[str, uuid.UUID, str]: + async def _seed_infrastructure( + self, + db: ExtendedAsyncSAEngine, + domain_id: DomainID, + scaling_group_id: ResourceGroupID, + ) -> tuple[str, uuid.UUID, str]: """Create domain, project, scaling group, and agent. Returns (domain_name, project_id, agent_id).""" domain_name = "test-domain" project_id = uuid4() sg_name = "test-sg" agent_id = "i-test-agent" async with db.begin_session() as db_sess: - db_sess.add(DomainRow(name=domain_name)) + db_sess.add(DomainRow(id=domain_id, name=domain_name)) async with db.begin_session() as db_sess: db_sess.add( ProjectResourcePolicyRow( @@ -296,6 +311,7 @@ async def _seed_infrastructure(self, db: ExtendedAsyncSAEngine) -> tuple[str, uu async with db.begin_session() as db_sess: db_sess.add( ScalingGroupRow( + id=scaling_group_id, name=sg_name, driver="static", driver_opts={}, @@ -337,6 +353,8 @@ async def _create_kernel_with_allocations( domain_name: str, project_id: uuid.UUID, agent_id: str, + domain_id: DomainID, + resource_group_id: ResourceGroupID, status: KernelStatus, allocations: dict[str, tuple[Decimal, Decimal | None]], free_at: datetime | None = None, @@ -353,8 +371,10 @@ async def _create_kernel_with_allocations( db_sess.add( SessionRow( id=session_id, + domain_id=domain_id, domain_name=domain_name, group_id=project_id, + resource_group_id=resource_group_id, scaling_group_name="test-sg", user_uuid=uuid4(), occupying_slots=empty_slots, @@ -396,10 +416,14 @@ async def _create_kernel_with_allocations( async def test_aggregate_by_domain_with_active_kernels( self, db_with_full_tables: ExtendedAsyncSAEngine, + domain_id: DomainID, + scaling_group_id: ResourceGroupID, ) -> None: db = db_with_full_tables await self._seed_slot_types(db) - domain_name, project_id, agent_id = await self._seed_infrastructure(db) + domain_name, project_id, agent_id = await self._seed_infrastructure( + db, domain_id, scaling_group_id + ) # Kernel 1: RUNNING with used values await self._create_kernel_with_allocations( @@ -407,6 +431,8 @@ async def test_aggregate_by_domain_with_active_kernels( domain_name=domain_name, project_id=project_id, agent_id=agent_id, + domain_id=domain_id, + resource_group_id=scaling_group_id, status=KernelStatus.RUNNING, allocations={ "cpu": (Decimal("2"), Decimal("2")), @@ -419,6 +445,8 @@ async def test_aggregate_by_domain_with_active_kernels( domain_name=domain_name, project_id=project_id, agent_id=agent_id, + domain_id=domain_id, + resource_group_id=scaling_group_id, status=KernelStatus.SCHEDULED, allocations={ "cpu": (Decimal("4"), None), @@ -439,10 +467,14 @@ async def test_aggregate_by_domain_with_active_kernels( async def test_aggregate_by_domain_excludes_terminated_kernels( self, db_with_full_tables: ExtendedAsyncSAEngine, + domain_id: DomainID, + scaling_group_id: ResourceGroupID, ) -> None: db = db_with_full_tables await self._seed_slot_types(db) - domain_name, project_id, agent_id = await self._seed_infrastructure(db) + domain_name, project_id, agent_id = await self._seed_infrastructure( + db, domain_id, scaling_group_id + ) # TERMINATED kernel — should be excluded await self._create_kernel_with_allocations( @@ -450,6 +482,8 @@ async def test_aggregate_by_domain_excludes_terminated_kernels( domain_name=domain_name, project_id=project_id, agent_id=agent_id, + domain_id=domain_id, + resource_group_id=scaling_group_id, status=KernelStatus.TERMINATED, allocations={"cpu": (Decimal("4"), Decimal("4"))}, ) @@ -463,10 +497,14 @@ async def test_aggregate_by_domain_excludes_terminated_kernels( async def test_aggregate_by_domain_excludes_freed_allocations( self, db_with_full_tables: ExtendedAsyncSAEngine, + domain_id: DomainID, + scaling_group_id: ResourceGroupID, ) -> None: db = db_with_full_tables await self._seed_slot_types(db) - domain_name, project_id, agent_id = await self._seed_infrastructure(db) + domain_name, project_id, agent_id = await self._seed_infrastructure( + db, domain_id, scaling_group_id + ) # RUNNING kernel but allocations already freed (free_at is set) await self._create_kernel_with_allocations( @@ -474,6 +512,8 @@ async def test_aggregate_by_domain_excludes_freed_allocations( domain_name=domain_name, project_id=project_id, agent_id=agent_id, + domain_id=domain_id, + resource_group_id=scaling_group_id, status=KernelStatus.RUNNING, allocations={"cpu": (Decimal("2"), Decimal("2"))}, free_at=datetime.now(tzutc()), @@ -501,16 +541,22 @@ async def test_aggregate_by_domain_empty( async def test_aggregate_by_project_with_active_kernels( self, db_with_full_tables: ExtendedAsyncSAEngine, + domain_id: DomainID, + scaling_group_id: ResourceGroupID, ) -> None: db = db_with_full_tables await self._seed_slot_types(db) - domain_name, project_id, agent_id = await self._seed_infrastructure(db) + domain_name, project_id, agent_id = await self._seed_infrastructure( + db, domain_id, scaling_group_id + ) await self._create_kernel_with_allocations( db, domain_name=domain_name, project_id=project_id, agent_id=agent_id, + domain_id=domain_id, + resource_group_id=scaling_group_id, status=KernelStatus.RUNNING, allocations={ "cpu": (Decimal("1"), Decimal("1")), @@ -522,6 +568,8 @@ async def test_aggregate_by_project_with_active_kernels( domain_name=domain_name, project_id=project_id, agent_id=agent_id, + domain_id=domain_id, + resource_group_id=scaling_group_id, status=KernelStatus.PREPARING, allocations={ "cpu": (Decimal("3"), None), @@ -553,17 +601,23 @@ async def test_aggregate_by_project_empty( async def test_aggregate_slots_sorted_by_rank( self, db_with_full_tables: ExtendedAsyncSAEngine, + domain_id: DomainID, + scaling_group_id: ResourceGroupID, ) -> None: """Verify used_slots are sorted by rank (cpu=0 before mem=1).""" db = db_with_full_tables await self._seed_slot_types(db) - domain_name, project_id, agent_id = await self._seed_infrastructure(db) + domain_name, project_id, agent_id = await self._seed_infrastructure( + db, domain_id, scaling_group_id + ) await self._create_kernel_with_allocations( db, domain_name=domain_name, project_id=project_id, agent_id=agent_id, + domain_id=domain_id, + resource_group_id=scaling_group_id, status=KernelStatus.RUNNING, allocations={ "cpu": (Decimal("1"), Decimal("1")), @@ -582,6 +636,14 @@ async def test_aggregate_slots_sorted_by_rank( class TestComputeActualAgentResourceUsage: """Tests for compute_actual_agent_resource_usage.""" + @pytest.fixture + def domain_id(self) -> DomainID: + return DomainID(uuid4()) + + @pytest.fixture + def scaling_group_id(self) -> ResourceGroupID: + return ResourceGroupID(uuid4()) + @pytest.fixture async def db_with_full_tables( self, @@ -610,14 +672,19 @@ async def _seed_slot_types(self, db: ExtendedAsyncSAEngine) -> None: db_sess.add(ResourceSlotTypeRow(slot_name="cpu", slot_type="count", rank=0)) db_sess.add(ResourceSlotTypeRow(slot_name="mem", slot_type="bytes", rank=1)) - async def _seed_infrastructure(self, db: ExtendedAsyncSAEngine) -> tuple[str, uuid.UUID, str]: + async def _seed_infrastructure( + self, + db: ExtendedAsyncSAEngine, + domain_id: DomainID, + scaling_group_id: ResourceGroupID, + ) -> tuple[str, uuid.UUID, str]: """Create domain, project, scaling group, and agent.""" domain_name = "test-domain" project_id = uuid4() sg_name = "test-sg" agent_id = "i-test-agent" async with db.begin_session() as db_sess: - db_sess.add(DomainRow(name=domain_name)) + db_sess.add(DomainRow(id=domain_id, name=domain_name)) async with db.begin_session() as db_sess: db_sess.add( ProjectResourcePolicyRow( @@ -630,6 +697,7 @@ async def _seed_infrastructure(self, db: ExtendedAsyncSAEngine) -> tuple[str, uu async with db.begin_session() as db_sess: db_sess.add( ScalingGroupRow( + id=scaling_group_id, name=sg_name, driver="static", driver_opts={}, @@ -671,6 +739,8 @@ async def _create_kernel_with_allocations( domain_name: str, project_id: uuid.UUID, agent_id: str, + domain_id: DomainID, + resource_group_id: ResourceGroupID, status: KernelStatus, allocations: dict[str, tuple[Decimal, Decimal | None]], free_at: datetime | None = None, @@ -682,8 +752,10 @@ async def _create_kernel_with_allocations( db_sess.add( SessionRow( id=session_id, + domain_id=domain_id, domain_name=domain_name, group_id=project_id, + resource_group_id=resource_group_id, scaling_group_name="test-sg", user_uuid=uuid4(), occupying_slots=empty_slots, @@ -738,17 +810,23 @@ async def test_no_active_allocations( async def test_active_allocations_summed_by_agent_and_slot( self, db_with_full_tables: ExtendedAsyncSAEngine, + domain_id: DomainID, + scaling_group_id: ResourceGroupID, ) -> None: """Active allocations are summed per (agent_id, slot_name).""" db = db_with_full_tables await self._seed_slot_types(db) - domain_name, project_id, agent_id = await self._seed_infrastructure(db) + domain_name, project_id, agent_id = await self._seed_infrastructure( + db, domain_id, scaling_group_id + ) await self._create_kernel_with_allocations( db, domain_name=domain_name, project_id=project_id, agent_id=agent_id, + domain_id=domain_id, + resource_group_id=scaling_group_id, status=KernelStatus.RUNNING, allocations={ "cpu": (Decimal("2"), Decimal("2")), @@ -760,6 +838,8 @@ async def test_active_allocations_summed_by_agent_and_slot( domain_name=domain_name, project_id=project_id, agent_id=agent_id, + domain_id=domain_id, + resource_group_id=scaling_group_id, status=KernelStatus.RUNNING, allocations={ "cpu": (Decimal("4"), Decimal("4")), @@ -776,17 +856,23 @@ async def test_active_allocations_summed_by_agent_and_slot( async def test_freed_allocations_excluded( self, db_with_full_tables: ExtendedAsyncSAEngine, + domain_id: DomainID, + scaling_group_id: ResourceGroupID, ) -> None: """Allocations with free_at set are excluded.""" db = db_with_full_tables await self._seed_slot_types(db) - domain_name, project_id, agent_id = await self._seed_infrastructure(db) + domain_name, project_id, agent_id = await self._seed_infrastructure( + db, domain_id, scaling_group_id + ) await self._create_kernel_with_allocations( db, domain_name=domain_name, project_id=project_id, agent_id=agent_id, + domain_id=domain_id, + resource_group_id=scaling_group_id, status=KernelStatus.RUNNING, allocations={"cpu": (Decimal("4"), Decimal("4"))}, free_at=datetime.now(tzutc()), @@ -800,17 +886,23 @@ async def test_freed_allocations_excluded( async def test_terminated_kernels_excluded( self, db_with_full_tables: ExtendedAsyncSAEngine, + domain_id: DomainID, + scaling_group_id: ResourceGroupID, ) -> None: """Kernels with TERMINATED status are excluded.""" db = db_with_full_tables await self._seed_slot_types(db) - domain_name, project_id, agent_id = await self._seed_infrastructure(db) + domain_name, project_id, agent_id = await self._seed_infrastructure( + db, domain_id, scaling_group_id + ) await self._create_kernel_with_allocations( db, domain_name=domain_name, project_id=project_id, agent_id=agent_id, + domain_id=domain_id, + resource_group_id=scaling_group_id, status=KernelStatus.TERMINATED, allocations={"cpu": (Decimal("4"), Decimal("4"))}, ) @@ -823,17 +915,23 @@ async def test_terminated_kernels_excluded( async def test_coalesce_uses_requested_when_used_is_none( self, db_with_full_tables: ExtendedAsyncSAEngine, + domain_id: DomainID, + scaling_group_id: ResourceGroupID, ) -> None: """When used is None, COALESCE falls back to requested.""" db = db_with_full_tables await self._seed_slot_types(db) - domain_name, project_id, agent_id = await self._seed_infrastructure(db) + domain_name, project_id, agent_id = await self._seed_infrastructure( + db, domain_id, scaling_group_id + ) await self._create_kernel_with_allocations( db, domain_name=domain_name, project_id=project_id, agent_id=agent_id, + domain_id=domain_id, + resource_group_id=scaling_group_id, status=KernelStatus.SCHEDULED, allocations={"cpu": (Decimal("4"), None)}, ) diff --git a/tests/unit/manager/repositories/scaling_group/test_resource_info.py b/tests/unit/manager/repositories/scaling_group/test_resource_info.py index ec98718478e..029ac23dbce 100644 --- a/tests/unit/manager/repositories/scaling_group/test_resource_info.py +++ b/tests/unit/manager/repositories/scaling_group/test_resource_info.py @@ -9,6 +9,8 @@ import pytest +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ResourceSlot, SlotQuantity from ai.backend.manager.data.agent.types import AgentStatus from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm @@ -152,15 +154,25 @@ async def db_source( db=db_with_cleanup, ) + @pytest.fixture + def base_scaling_group_id(self) -> ResourceGroupID: + return ResourceGroupID(uuid.uuid4()) + + @pytest.fixture + def test_domain_id(self) -> DomainID: + return DomainID(uuid.uuid4()) + @pytest.fixture async def base_scaling_group( self, db_with_cleanup: ExtendedAsyncSAEngine, + base_scaling_group_id: ResourceGroupID, ) -> AsyncGenerator[str, None]: """Create a basic scaling group for testing.""" sgroup_name = f"test-sgroup-{uuid.uuid4().hex[:8]}" async with db_with_cleanup.begin_session() as db_sess: sgroup = ScalingGroupRow( + id=base_scaling_group_id, name=sgroup_name, description="Test scaling group", is_active=True, @@ -380,6 +392,7 @@ async def scaling_group_with_non_schedulable_agents( async def test_user_domain_group( self, db_with_cleanup: ExtendedAsyncSAEngine, + test_domain_id: DomainID, ) -> AsyncGenerator[tuple[uuid.UUID, str, uuid.UUID], None]: """Create test domain, user, and group for kernel tests. @@ -394,6 +407,7 @@ async def test_user_domain_group( async with db_with_cleanup.begin_session() as db_sess: # Create domain domain = DomainRow( + id=test_domain_id, name=test_domain, description="Test domain", is_active=True, @@ -461,6 +475,8 @@ async def scaling_group_with_running_kernels( self, db_with_cleanup: ExtendedAsyncSAEngine, base_scaling_group: str, + base_scaling_group_id: ResourceGroupID, + test_domain_id: DomainID, test_user_domain_group: tuple[uuid.UUID, str, uuid.UUID], ) -> AsyncGenerator[tuple[str, ResourceSlot, list[ResourceSlot]], None]: """Create scaling group with agent and RUNNING/TERMINATING kernels. @@ -509,9 +525,11 @@ async def scaling_group_with_running_kernels( SessionRow( id=session_id, domain_name=domain_name, + domain_id=test_domain_id, group_id=group_id, user_uuid=user_uuid, scaling_group_name=base_scaling_group, + resource_group_id=base_scaling_group_id, cluster_size=2, vfolder_mounts={}, ) @@ -557,6 +575,8 @@ async def scaling_group_with_mixed_kernel_statuses( self, db_with_cleanup: ExtendedAsyncSAEngine, base_scaling_group: str, + base_scaling_group_id: ResourceGroupID, + test_domain_id: DomainID, test_user_domain_group: tuple[uuid.UUID, str, uuid.UUID], ) -> AsyncGenerator[tuple[str, ResourceSlot, ResourceSlot], None]: """Create scaling group with kernels in various statuses. @@ -609,9 +629,11 @@ async def scaling_group_with_mixed_kernel_statuses( SessionRow( id=session_id, domain_name=domain_name, + domain_id=test_domain_id, group_id=group_id, user_uuid=user_uuid, scaling_group_name=base_scaling_group, + resource_group_id=base_scaling_group_id, cluster_size=4, vfolder_mounts={}, ) diff --git a/tests/unit/manager/repositories/test_utils.py b/tests/unit/manager/repositories/test_utils.py index 62055564666..9ad34939690 100644 --- a/tests/unit/manager/repositories/test_utils.py +++ b/tests/unit/manager/repositories/test_utils.py @@ -10,6 +10,8 @@ from dateutil.tz import tzutc from sqlalchemy.engine import Row +from ai.backend.common.identifier.domain import DomainID +from ai.backend.common.identifier.resource_group import ResourceGroupID from ai.backend.common.types import ResourceSlot from ai.backend.manager.data.auth.hash import PasswordHashAlgorithm from ai.backend.manager.defs import DEFAULT_ROLE @@ -112,9 +114,12 @@ async def session_info( session_id = str(uuid.uuid4()).replace("-", "") session_creation_id = str(uuid.uuid4()).replace("-", "") resource_policy_name = str(uuid.uuid4()).replace("-", "") + resource_group_id = ResourceGroupID(uuid.uuid4()) + domain_id = DomainID(uuid.uuid4()) async with db_with_cleanup.begin_session() as db_sess: scaling_group = ScalingGroupRow( + id=resource_group_id, name=sgroup_name, driver="test", scheduler="test", @@ -122,7 +127,7 @@ async def session_info( ) db_sess.add(scaling_group) - domain = DomainRow(name=domain_name, total_resource_slots=ResourceSlot()) + domain = DomainRow(id=domain_id, name=domain_name, total_resource_slots=ResourceSlot()) db_sess.add(domain) user_resource_policy = UserResourcePolicyRow( @@ -165,7 +170,9 @@ async def session_info( id=session_id, creation_id=session_creation_id, cluster_size=1, + domain_id=domain_id, domain_name=domain_name, + resource_group_id=resource_group_id, scaling_group_name=sgroup_name, group_id=group_id, user_uuid=user_uuid, From 88cdf47228f62e8cd52373d812d03b081dca9362 Mon Sep 17 00:00:00 2001 From: BoKeum Date: Fri, 3 Jul 2026 14:29:13 +0900 Subject: [PATCH 4/6] changelog: add news fragment for PR #12452 --- changes/12452.feature.md | 1 + 1 file changed, 1 insertion(+) create mode 100644 changes/12452.feature.md diff --git a/changes/12452.feature.md b/changes/12452.feature.md new file mode 100644 index 00000000000..681b4c0c0f3 --- /dev/null +++ b/changes/12452.feature.md @@ -0,0 +1 @@ +Add domain_id and resource_group_id columns to session rows and populate them at session creation From ca3eb14e3f8c9de3d9ef7a97de5595f29b07ef66 Mon Sep 17 00:00:00 2001 From: BoKeum Date: Fri, 3 Jul 2026 16:15:55 +0900 Subject: [PATCH 5/6] refactor: split scaling group component fixtures --- src/ai/backend/testutils/fixtures.py | 8 -- tests/component/agent_api/conftest.py | 8 +- tests/component/agent_api/test_agent_api.py | 10 +- tests/component/auto_scaling_rule/conftest.py | 7 +- tests/component/conftest.py | 32 +++-- tests/component/deployment/test_deployment.py | 46 +++---- .../deployment/test_deployment_lifecycle.py | 22 +-- tests/component/fair_share/conftest.py | 7 +- tests/component/fair_share/test_fair_share.py | 125 +++++++++--------- .../test_fair_share_domain_weights.py | 53 ++++---- .../test_fair_share_project_weights.py | 47 +++---- .../test_fair_share_user_weights.py | 47 +++---- .../component/group/test_group_management.py | 19 +-- tests/component/infra/test_infra.py | 6 +- tests/component/operations/conftest.py | 6 +- .../test_resource_allocation.py | 19 +-- .../test_resource_allocation_visibility.py | 18 +-- .../resource_preset/test_resource_preset.py | 10 +- .../test_resource_preset_full.py | 8 +- .../scaling_group/test_scaling_group.py | 30 ++--- .../test_scaling_group_associations.py | 18 +-- .../scaling_group/test_scaling_group_crud.py | 19 ++- tests/component/session/conftest.py | 34 ++--- tests/component/session/test_session_batch.py | 6 +- .../session/test_session_execution.py | 42 +++--- .../session/test_session_lifecycle.py | 21 +-- tests/component/session/test_session_query.py | 32 +++-- tests/component/session_v2/conftest.py | 18 +-- tests/component/stream/conftest.py | 12 +- tests/component/streaming/conftest.py | 21 +-- 30 files changed, 385 insertions(+), 366 deletions(-) diff --git a/src/ai/backend/testutils/fixtures.py b/src/ai/backend/testutils/fixtures.py index 2c3bd1b0a7d..192df7c2ae0 100644 --- a/src/ai/backend/testutils/fixtures.py +++ b/src/ai/backend/testutils/fixtures.py @@ -4,12 +4,10 @@ from dataclasses import dataclass from ai.backend.common.identifier.domain import DomainID, DomainName -from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName __all__ = ( "DomainFactory", "DomainFixtureData", - "ScalingGroupFixtureData", ) @@ -19,10 +17,4 @@ class DomainFixtureData: domain_id: DomainID -@dataclass(frozen=True) -class ScalingGroupFixtureData: - scaling_group_name: ResourceGroupName - scaling_group_id: ResourceGroupID - - DomainFactory = Callable[..., Awaitable[DomainFixtureData]] diff --git a/tests/component/agent_api/conftest.py b/tests/component/agent_api/conftest.py index d6aebb2475c..1fbf92a51e7 100644 --- a/tests/component/agent_api/conftest.py +++ b/tests/component/agent_api/conftest.py @@ -11,6 +11,7 @@ from ai.backend.common.etcd import AsyncEtcd, ConfigScopes from ai.backend.common.events.dispatcher import EventProducer +from ai.backend.common.identifier.resource_group import ResourceGroupName from ai.backend.common.plugin.hook import HookPluginContext from ai.backend.common.types import HostPortPair, ResourceSlot from ai.backend.manager.actions.validators import ActionValidators @@ -28,7 +29,6 @@ from ai.backend.manager.repositories.scheduler.repository import SchedulerRepository from ai.backend.manager.services.agent.processors import AgentProcessors from ai.backend.manager.services.agent.service import AgentService -from ai.backend.testutils.fixtures import ScalingGroupFixtureData @pytest.fixture() @@ -106,11 +106,11 @@ def server_module_registries( @pytest.fixture() async def agent_fixture( db_engine: SAEngine, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> AsyncIterator[str]: """Insert a test agent row and yield its ID. - The agent references the scaling_group_fixture via FK. + The agent references the scaling_group_name via FK. Teardown deletes the agent row (cascade deletes agent_resources). """ agent_id = f"i-test-agent-{secrets.token_hex(6)}" @@ -120,7 +120,7 @@ async def agent_fixture( id=agent_id, status=AgentStatus.ALIVE, region="local", - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, schedulable=True, available_slots=ResourceSlot({"cpu": "4", "mem": "8589934592"}), occupied_slots=ResourceSlot(), diff --git a/tests/component/agent_api/test_agent_api.py b/tests/component/agent_api/test_agent_api.py index 8769a4077fd..db7af7791d1 100644 --- a/tests/component/agent_api/test_agent_api.py +++ b/tests/component/agent_api/test_agent_api.py @@ -26,7 +26,7 @@ OrderDirection, ) from ai.backend.common.dto.manager.query import StringFilter -from ai.backend.testutils.fixtures import ScalingGroupFixtureData +from ai.backend.common.identifier.resource_group import ResourceGroupName class TestSearchAgents: @@ -84,10 +84,10 @@ async def test_admin_searches_agents_with_resource_group_filter( self, admin_registry: BackendAIClientRegistry, agent_fixture: str, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: """Filtering by resource_group returns agents in that scaling group.""" - sgroup_name = scaling_group_fixture.scaling_group_name + sgroup_name = scaling_group_name result = await admin_registry.agent.search_agents( SearchAgentsRequest( filter=AgentFilter( @@ -140,11 +140,11 @@ async def test_compound_filters_status_and_resource_group( self, admin_registry: BackendAIClientRegistry, agent_fixture: str, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: """Compound filters (status + resource_group) return intersection of conditions.""" # Search with both status=ALIVE and resource_group filters - sgroup_name = scaling_group_fixture.scaling_group_name + sgroup_name = scaling_group_name result = await admin_registry.agent.search_agents( SearchAgentsRequest( filter=AgentFilter( diff --git a/tests/component/auto_scaling_rule/conftest.py b/tests/component/auto_scaling_rule/conftest.py index 9cc303888c6..176375ba7f8 100644 --- a/tests/component/auto_scaling_rule/conftest.py +++ b/tests/component/auto_scaling_rule/conftest.py @@ -11,6 +11,7 @@ from ai.backend.common.container_registry import ContainerRegistryType from ai.backend.common.data.endpoint.types import EndpointLifecycle +from ai.backend.common.identifier.resource_group import ResourceGroupName from ai.backend.manager.actions.validators import ActionValidators from ai.backend.manager.actions.validators.rbac import RBACValidators from ai.backend.manager.actions.validators.rbac.bulk import BulkActionRBACValidator @@ -37,7 +38,7 @@ ) from ai.backend.manager.services.deployment.processors import DeploymentProcessors from ai.backend.manager.services.deployment.service import DeploymentService -from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData @dataclass @@ -125,7 +126,7 @@ async def model_deployment_fixture( db_engine: SAEngine, domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, admin_user_fixture: UserFixtureData, ) -> AsyncIterator[uuid.UUID]: """Insert a minimal EndpointRow (model deployment) and yield its UUID. @@ -176,7 +177,7 @@ async def model_deployment_fixture( session_owner=str(admin_user_fixture.user_uuid), domain=domain_fixture.domain_name, project=str(group_fixture), - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, lifecycle_stage=EndpointLifecycle.CREATED, url=None, ) diff --git a/tests/component/conftest.py b/tests/component/conftest.py index 85e408e9815..203ed7bdf3a 100644 --- a/tests/component/conftest.py +++ b/tests/component/conftest.py @@ -152,7 +152,7 @@ postgres_container, redis_container, ) -from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData from ai.backend.testutils.pants import get_parallel_slot log = logging.getLogger("tests.component.conftest") @@ -673,16 +673,15 @@ async def resource_policy_fixture( @pytest.fixture() -async def scaling_group_fixture( +async def scaling_group_name( db_engine: SAEngine, domain_fixture: DomainFixtureData, -) -> AsyncIterator[ScalingGroupFixtureData]: - """Insert a scaling group and its domain association; yield its identifiers.""" +) -> AsyncIterator[ResourceGroupName]: + """Insert a scaling group and its domain association; yield its name.""" sgroup_name = ResourceGroupName(f"sgroup-{secrets.token_hex(6)}") async with db_engine.begin() as conn: - result = await conn.execute( - sa.insert(scaling_groups) - .values( + await conn.execute( + sa.insert(scaling_groups).values( name=sgroup_name, description=f"Test scaling group {sgroup_name}", is_active=True, @@ -691,16 +690,14 @@ async def scaling_group_fixture( scheduler="fifo", scheduler_opts=ScalingGroupOpts(), ) - .returning(scaling_groups.c.id) ) - sgroup_id = ResourceGroupID(result.scalar_one()) await conn.execute( sa.insert(sgroups_for_domains).values( scaling_group=sgroup_name, domain=domain_fixture.domain_name, ) ) - yield ScalingGroupFixtureData(scaling_group_name=sgroup_name, scaling_group_id=sgroup_id) + yield sgroup_name async with db_engine.begin() as conn: await conn.execute( sgroups_for_domains.delete().where(sgroups_for_domains.c.scaling_group == sgroup_name) @@ -708,6 +705,19 @@ async def scaling_group_fixture( await conn.execute(scaling_groups.delete().where(scaling_groups.c.name == sgroup_name)) +@pytest.fixture() +async def scaling_group_id( + db_engine: SAEngine, + scaling_group_name: ResourceGroupName, +) -> ResourceGroupID: + """Return the inserted scaling group's ID.""" + async with db_engine.begin() as conn: + result = await conn.execute( + sa.select(scaling_groups.c.id).where(scaling_groups.c.name == scaling_group_name) + ) + return ResourceGroupID(result.scalar_one()) + + @pytest.fixture() async def group_fixture( db_engine: SAEngine, @@ -936,7 +946,7 @@ async def regular_user_fixture( async def database_fixture( admin_user_fixture: UserFixtureData, regular_user_fixture: UserFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> AsyncIterator[None]: """Backward-compatible aggregate: requests all seed data fixtures.""" yield diff --git a/tests/component/deployment/test_deployment.py b/tests/component/deployment/test_deployment.py index 9cb6741a7b8..cd390f410e2 100644 --- a/tests/component/deployment/test_deployment.py +++ b/tests/component/deployment/test_deployment.py @@ -56,7 +56,7 @@ from ai.backend.manager.services.deployment.processors import DeploymentProcessors from ai.backend.manager.services.deployment.service import _map_lifecycle_to_status from ai.backend.manager.services.processors import Processors -from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData if TYPE_CHECKING: from tests.component.conftest import UserFixtureData @@ -96,7 +96,7 @@ async def test_search_deployments_paginated( admin_registry: BackendAIClientRegistry, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """Search deployments with pagination returns correct page.""" @@ -108,7 +108,7 @@ async def test_search_deployments_paginated( metadata=DeploymentMetadataInput( project_id=group_fixture, domain_name=domain_fixture.domain_name, - resource_group_name=scaling_group_fixture.scaling_group_name, + resource_group_name=scaling_group_name, name=f"test-deployment-{i}-{secrets.token_hex(4)}", ), network_access=NetworkAccessInput(open_to_public=False), @@ -158,7 +158,7 @@ async def test_get_deployment_by_id( admin_registry: BackendAIClientRegistry, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """Get deployment by ID returns correct deployment details.""" @@ -167,7 +167,7 @@ async def test_get_deployment_by_id( metadata=DeploymentMetadataInput( project_id=group_fixture, domain_name=domain_fixture.domain_name, - resource_group_name=scaling_group_fixture.scaling_group_name, + resource_group_name=scaling_group_name, name=f"test-deployment-{secrets.token_hex(4)}", ), network_access=NetworkAccessInput(open_to_public=False), @@ -269,7 +269,7 @@ async def test_search_routes_for_deployment( admin_registry: BackendAIClientRegistry, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """Search routes for a deployment returns route list.""" @@ -278,7 +278,7 @@ async def test_search_routes_for_deployment( metadata=DeploymentMetadataInput( project_id=group_fixture, domain_name=domain_fixture.domain_name, - resource_group_name=scaling_group_fixture.scaling_group_name, + resource_group_name=scaling_group_name, name=f"test-deployment-{secrets.token_hex(4)}", ), network_access=NetworkAccessInput(open_to_public=False), @@ -408,7 +408,7 @@ async def test_my_search_and_filter_returns_intersection( deployment_adapter: DeploymentAdapter, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """AND clause must narrow results to deployments matching every nested filter.""" @@ -416,7 +416,7 @@ async def test_my_search_and_filter_returns_intersection( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture.scaling_group_name, + scaling_group_name, deployment_seed_data, ["alpha", "production"], ) @@ -424,7 +424,7 @@ async def test_my_search_and_filter_returns_intersection( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture.scaling_group_name, + scaling_group_name, deployment_seed_data, ["beta", "production"], ) @@ -432,7 +432,7 @@ async def test_my_search_and_filter_returns_intersection( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture.scaling_group_name, + scaling_group_name, deployment_seed_data, ["alpha", "beta"], ) @@ -460,7 +460,7 @@ async def test_my_search_or_filter_returns_union( deployment_adapter: DeploymentAdapter, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """OR clause must widen results to deployments matching any nested filter.""" @@ -468,7 +468,7 @@ async def test_my_search_or_filter_returns_union( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture.scaling_group_name, + scaling_group_name, deployment_seed_data, ["alpha"], ) @@ -476,7 +476,7 @@ async def test_my_search_or_filter_returns_union( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture.scaling_group_name, + scaling_group_name, deployment_seed_data, ["beta"], ) @@ -484,7 +484,7 @@ async def test_my_search_or_filter_returns_union( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture.scaling_group_name, + scaling_group_name, deployment_seed_data, ["gamma"], ) @@ -512,7 +512,7 @@ async def test_my_search_or_filter_groups_multi_field_subfilters( deployment_adapter: DeploymentAdapter, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """OR with multi-field sub-filters must AND fields within each branch. @@ -527,7 +527,7 @@ async def test_my_search_or_filter_groups_multi_field_subfilters( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture.scaling_group_name, + scaling_group_name, deployment_seed_data, ["alpha"], name=f"bar-x-{suffix}", @@ -536,7 +536,7 @@ async def test_my_search_or_filter_groups_multi_field_subfilters( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture.scaling_group_name, + scaling_group_name, deployment_seed_data, ["beta"], name=f"foo-x-{suffix}", @@ -545,7 +545,7 @@ async def test_my_search_or_filter_groups_multi_field_subfilters( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture.scaling_group_name, + scaling_group_name, deployment_seed_data, ["beta"], name=f"bar-y-{suffix}", @@ -554,7 +554,7 @@ async def test_my_search_or_filter_groups_multi_field_subfilters( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture.scaling_group_name, + scaling_group_name, deployment_seed_data, ["alpha"], name=f"foo-y-{suffix}", @@ -589,7 +589,7 @@ async def test_project_search_and_filter_returns_intersection( deployment_adapter: DeploymentAdapter, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """project_search must also honor AND across nested filters.""" @@ -597,7 +597,7 @@ async def test_project_search_and_filter_returns_intersection( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture.scaling_group_name, + scaling_group_name, deployment_seed_data, ["alpha"], ) @@ -605,7 +605,7 @@ async def test_project_search_and_filter_returns_intersection( admin_registry, group_fixture, domain_fixture.domain_name, - scaling_group_fixture.scaling_group_name, + scaling_group_name, deployment_seed_data, ["alpha", "beta"], ) diff --git a/tests/component/deployment/test_deployment_lifecycle.py b/tests/component/deployment/test_deployment_lifecycle.py index 5121363578d..020950bab6d 100644 --- a/tests/component/deployment/test_deployment_lifecycle.py +++ b/tests/component/deployment/test_deployment_lifecycle.py @@ -37,7 +37,7 @@ from ai.backend.common.identifier.resource_group import ResourceGroupName from ai.backend.common.identifier.vfolder import VFolderUUID from ai.backend.common.types import ClusterMode -from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData _RANDOM_DEPLOYMENT_ID = uuid.uuid4() _RANDOM_REVISION_ID = uuid.uuid4() @@ -91,7 +91,7 @@ async def test_create_deployment_success( admin_registry: BackendAIClientRegistry, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """Creating a deployment with valid config returns deployment with initial status.""" @@ -100,7 +100,7 @@ async def test_create_deployment_success( metadata=DeploymentMetadataInput( project_id=group_fixture, domain_name=domain_fixture.domain_name, - resource_group_name=scaling_group_fixture.scaling_group_name, + resource_group_name=scaling_group_name, name=f"test-deployment-{secrets.token_hex(4)}", ), network_access=NetworkAccessInput(open_to_public=False), @@ -158,7 +158,7 @@ async def test_update_deployment_config( admin_registry: BackendAIClientRegistry, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """Updating deployment config (name, replica_count) succeeds.""" @@ -168,7 +168,7 @@ async def test_update_deployment_config( metadata=DeploymentMetadataInput( project_id=group_fixture, domain_name=domain_fixture.domain_name, - resource_group_name=scaling_group_fixture.scaling_group_name, + resource_group_name=scaling_group_name, name=f"test-deployment-{secrets.token_hex(4)}", ), network_access=NetworkAccessInput(open_to_public=False), @@ -228,7 +228,7 @@ async def test_destroy_deployment( admin_registry: BackendAIClientRegistry, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """Destroying a deployment terminates it successfully.""" @@ -238,7 +238,7 @@ async def test_destroy_deployment( metadata=DeploymentMetadataInput( project_id=group_fixture, domain_name=domain_fixture.domain_name, - resource_group_name=scaling_group_fixture.scaling_group_name, + resource_group_name=scaling_group_name, name=f"test-deployment-{secrets.token_hex(4)}", ), network_access=NetworkAccessInput(open_to_public=False), @@ -294,7 +294,7 @@ async def test_activate_deactivate_revision( admin_registry: BackendAIClientRegistry, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """Adding a revision and searching revisions works correctly.""" @@ -319,7 +319,7 @@ async def test_activate_deactivate_revision( metadata=DeploymentMetadataInput( project_id=group_fixture, domain_name=domain_fixture.domain_name, - resource_group_name=scaling_group_fixture.scaling_group_name, + resource_group_name=scaling_group_name, name=f"test-deployment-{secrets.token_hex(4)}", ), network_access=NetworkAccessInput(open_to_public=False), @@ -375,7 +375,7 @@ async def test_change_replica_count( admin_registry: BackendAIClientRegistry, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, deployment_seed_data: tuple[ImageID, VFolderUUID], ) -> None: """Changing replica_count updates the replica count.""" @@ -384,7 +384,7 @@ async def test_change_replica_count( metadata=DeploymentMetadataInput( project_id=group_fixture, domain_name=domain_fixture.domain_name, - resource_group_name=scaling_group_fixture.scaling_group_name, + resource_group_name=scaling_group_name, name=f"test-deployment-{secrets.token_hex(4)}", ), network_access=NetworkAccessInput(open_to_public=False), diff --git a/tests/component/fair_share/conftest.py b/tests/component/fair_share/conftest.py index 2073c98539e..44981f99ca6 100644 --- a/tests/component/fair_share/conftest.py +++ b/tests/component/fair_share/conftest.py @@ -9,6 +9,7 @@ import sqlalchemy as sa from sqlalchemy.ext.asyncio.engine import AsyncEngine as SAEngine +from ai.backend.common.identifier.resource_group import ResourceGroupName from ai.backend.manager.actions.validators import ActionValidators from ai.backend.manager.api.rest.fair_share.handler import FairShareAPIHandler from ai.backend.manager.api.rest.fair_share.registry import register_fair_share_routes @@ -28,7 +29,7 @@ from ai.backend.manager.services.resource_usage.service import ResourceUsageService from ai.backend.manager.services.scaling_group.processors import ScalingGroupProcessors from ai.backend.manager.services.scaling_group.service import ScalingGroupService -from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData @pytest.fixture() @@ -83,7 +84,7 @@ async def group_fixture( db_engine: SAEngine, domain_fixture: DomainFixtureData, resource_policy_fixture: str, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> AsyncIterator[uuid.UUID]: """Insert a test group with scaling-group association for fair-share tests.""" group_id = uuid.uuid4() @@ -101,7 +102,7 @@ async def group_fixture( ) await conn.execute( sa.insert(sgroups_for_groups).values( - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, group=group_id, ) ) diff --git a/tests/component/fair_share/test_fair_share.py b/tests/component/fair_share/test_fair_share.py index ef6b44baf2b..9a664cccc53 100644 --- a/tests/component/fair_share/test_fair_share.py +++ b/tests/component/fair_share/test_fair_share.py @@ -44,7 +44,8 @@ UpsertUserFairShareWeightResponse, UserWeightEntryInput, ) -from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData +from ai.backend.common.identifier.resource_group import ResourceGroupName +from ai.backend.testutils.fixtures import DomainFixtureData # ---- Domain Fair Share (Global-scoped) ---- @@ -53,11 +54,11 @@ class TestGetDomainFairShare: async def test_admin_get_domain_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, ) -> None: result = await admin_registry.fair_share.get_domain_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, ) assert isinstance(result, GetDomainFairShareResponse) @@ -65,12 +66,12 @@ async def test_admin_get_domain_fair_share( async def test_user_get_domain_fair_share_forbidden( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, ) -> None: with pytest.raises(PermissionDeniedError): await user_registry.fair_share.get_domain_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, ) @@ -103,11 +104,11 @@ class TestGetProjectFairShare: async def test_admin_get_project_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, ) -> None: result = await admin_registry.fair_share.get_project_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, project_id=group_fixture, ) assert isinstance(result, GetProjectFairShareResponse) @@ -115,12 +116,12 @@ async def test_admin_get_project_fair_share( async def test_user_get_project_fair_share_forbidden( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, ) -> None: with pytest.raises(PermissionDeniedError): await user_registry.fair_share.get_project_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, project_id=group_fixture, ) @@ -144,12 +145,12 @@ class TestGetUserFairShare: async def test_admin_get_user_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, admin_user_fixture: Any, ) -> None: result = await admin_registry.fair_share.get_user_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, ) @@ -158,13 +159,13 @@ async def test_admin_get_user_fair_share( async def test_user_get_user_fair_share_forbidden( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, regular_user_fixture: Any, ) -> None: with pytest.raises(PermissionDeniedError): await user_registry.fair_share.get_user_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, project_id=group_fixture, user_uuid=regular_user_fixture.user_uuid, ) @@ -228,11 +229,11 @@ class TestRGGetDomainFairShare: async def test_admin_rg_get_domain_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, ) -> None: result = await admin_registry.fair_share.rg_get_domain_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, ) assert isinstance(result, GetDomainFairShareResponse) @@ -240,11 +241,11 @@ async def test_admin_rg_get_domain_fair_share( async def test_user_rg_get_domain_fair_share( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, ) -> None: result = await user_registry.fair_share.rg_get_domain_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, ) assert isinstance(result, GetDomainFairShareResponse) @@ -254,10 +255,10 @@ class TestRGSearchDomainFairShares: async def test_admin_rg_search_domain_fair_shares( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: result = await admin_registry.fair_share.rg_search_domain_fair_shares( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, request=SearchDomainFairSharesRequest(), ) assert isinstance(result, SearchDomainFairSharesResponse) @@ -271,12 +272,12 @@ class TestRGGetProjectFairShare: async def test_admin_rg_get_project_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, ) -> None: result = await admin_registry.fair_share.rg_get_project_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, project_id=group_fixture, ) @@ -287,11 +288,11 @@ class TestRGSearchProjectFairShares: async def test_admin_rg_search_project_fair_shares( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, ) -> None: result = await admin_registry.fair_share.rg_search_project_fair_shares( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, request=SearchProjectFairSharesRequest(), ) @@ -306,13 +307,13 @@ class TestRGGetUserFairShare: async def test_admin_rg_get_user_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: Any, ) -> None: result = await admin_registry.fair_share.rg_get_user_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, @@ -324,12 +325,12 @@ class TestRGSearchUserFairShares: async def test_admin_rg_search_user_fair_shares( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, ) -> None: result = await admin_registry.fair_share.rg_search_user_fair_shares( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, project_id=group_fixture, request=SearchUserFairSharesRequest(), @@ -345,10 +346,10 @@ class TestRGSearchDomainUsageBuckets: async def test_admin_rg_search_domain_usage_buckets( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: result = await admin_registry.fair_share.rg_search_domain_usage_buckets( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, request=SearchDomainUsageBucketsRequest(), ) assert isinstance(result, SearchDomainUsageBucketsResponse) @@ -359,11 +360,11 @@ class TestRGSearchProjectUsageBuckets: async def test_admin_rg_search_project_usage_buckets( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, ) -> None: result = await admin_registry.fair_share.rg_search_project_usage_buckets( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, request=SearchProjectUsageBucketsRequest(), ) @@ -375,12 +376,12 @@ class TestRGSearchUserUsageBuckets: async def test_admin_rg_search_user_usage_buckets( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, ) -> None: result = await admin_registry.fair_share.rg_search_user_usage_buckets( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, project_id=group_fixture, request=SearchUserUsageBucketsRequest(), @@ -396,27 +397,27 @@ class TestUpsertDomainFairShareWeight: async def test_admin_upsert_domain_weight( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, ) -> None: result = await admin_registry.fair_share.upsert_domain_fair_share_weight( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, request=UpsertDomainFairShareWeightRequest(weight=Decimal("1.5")), ) assert isinstance(result, UpsertDomainFairShareWeightResponse) assert result.item.domain_name == domain_fixture.domain_name - assert result.item.resource_group == scaling_group_fixture.scaling_group_name + assert result.item.resource_group == scaling_group_name async def test_user_upsert_domain_weight_forbidden( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, ) -> None: with pytest.raises(PermissionDeniedError): await user_registry.fair_share.upsert_domain_fair_share_weight( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, request=UpsertDomainFairShareWeightRequest(weight=Decimal("1.0")), ) @@ -426,12 +427,12 @@ class TestUpsertProjectFairShareWeight: async def test_admin_upsert_project_weight( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, ) -> None: result = await admin_registry.fair_share.upsert_project_fair_share_weight( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, project_id=group_fixture, request=UpsertProjectFairShareWeightRequest( domain_name=domain_fixture.domain_name, @@ -440,20 +441,20 @@ async def test_admin_upsert_project_weight( ) assert isinstance(result, UpsertProjectFairShareWeightResponse) assert result.item.project_id == group_fixture - assert result.item.resource_group == scaling_group_fixture.scaling_group_name + assert result.item.resource_group == scaling_group_name class TestUpsertUserFairShareWeight: async def test_admin_upsert_user_weight( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, admin_user_fixture: Any, ) -> None: result = await admin_registry.fair_share.upsert_user_fair_share_weight( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, request=UpsertUserFairShareWeightRequest( @@ -463,7 +464,7 @@ async def test_admin_upsert_user_weight( ) assert isinstance(result, UpsertUserFairShareWeightResponse) assert result.item.user_uuid == admin_user_fixture.user_uuid - assert result.item.resource_group == scaling_group_fixture.scaling_group_name + assert result.item.resource_group == scaling_group_name # ---- Bulk Upsert Weights (Global-scoped) ---- @@ -473,12 +474,12 @@ class TestBulkUpsertDomainFairShareWeight: async def test_admin_bulk_upsert_domain_weight( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, ) -> None: result = await admin_registry.fair_share.bulk_upsert_domain_fair_share_weight( BulkUpsertDomainFairShareWeightRequest( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, inputs=[ DomainWeightEntryInput( domain_name=domain_fixture.domain_name, @@ -495,13 +496,13 @@ class TestBulkUpsertProjectFairShareWeight: async def test_admin_bulk_upsert_project_weight( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, ) -> None: result = await admin_registry.fair_share.bulk_upsert_project_fair_share_weight( BulkUpsertProjectFairShareWeightRequest( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, inputs=[ ProjectWeightEntryInput( project_id=group_fixture, @@ -519,14 +520,14 @@ class TestBulkUpsertUserFairShareWeight: async def test_admin_bulk_upsert_user_weight( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, admin_user_fixture: Any, ) -> None: result = await admin_registry.fair_share.bulk_upsert_user_fair_share_weight( BulkUpsertUserFairShareWeightRequest( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, inputs=[ UserWeightEntryInput( user_uuid=admin_user_fixture.user_uuid, @@ -548,22 +549,22 @@ class TestGetResourceGroupFairShareSpec: async def test_admin_get_rg_spec( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: result = await admin_registry.fair_share.get_resource_group_fair_share_spec( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, ) assert isinstance(result, GetResourceGroupFairShareSpecResponse) - assert result.resource_group == scaling_group_fixture.scaling_group_name + assert result.resource_group == scaling_group_name async def test_user_get_rg_spec_forbidden( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: with pytest.raises(PermissionDeniedError): await user_registry.fair_share.get_resource_group_fair_share_spec( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, ) @@ -581,33 +582,33 @@ class TestUpdateResourceGroupFairShareSpec: async def test_admin_update_rg_spec( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: result = await admin_registry.fair_share.update_resource_group_fair_share_spec( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, request=UpdateResourceGroupFairShareSpecRequest( half_life_days=14, ), ) assert isinstance(result, UpdateResourceGroupFairShareSpecResponse) - assert result.resource_group == scaling_group_fixture.scaling_group_name + assert result.resource_group == scaling_group_name assert result.fair_share_spec.half_life_days == 14 async def test_admin_partial_update_preserves_other_fields( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: """Partial update (merge logic) changes only the specified field.""" current = await admin_registry.fair_share.get_resource_group_fair_share_spec( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, ) original_half_life = current.fair_share_spec.half_life_days original_lookback = current.fair_share_spec.lookback_days new_half_life = original_half_life + 1 result = await admin_registry.fair_share.update_resource_group_fair_share_spec( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, request=UpdateResourceGroupFairShareSpecRequest( half_life_days=new_half_life, ), @@ -618,7 +619,7 @@ async def test_admin_partial_update_preserves_other_fields( # Restore await admin_registry.fair_share.update_resource_group_fair_share_spec( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, request=UpdateResourceGroupFairShareSpecRequest( half_life_days=original_half_life, ), diff --git a/tests/component/fair_share/test_fair_share_domain_weights.py b/tests/component/fair_share/test_fair_share_domain_weights.py index 646f5001e7a..3cf2061fe05 100644 --- a/tests/component/fair_share/test_fair_share_domain_weights.py +++ b/tests/component/fair_share/test_fair_share_domain_weights.py @@ -16,7 +16,8 @@ UpsertDomainFairShareWeightRequest, UpsertDomainFairShareWeightResponse, ) -from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData +from ai.backend.common.identifier.resource_group import ResourceGroupName +from ai.backend.testutils.fixtures import DomainFixtureData class TestDomainFairShareWeights: @@ -25,18 +26,18 @@ class TestDomainFairShareWeights: async def test_get_domain_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, ) -> None: """Get domain fair share → returns weight data.""" result = await admin_registry.fair_share.get_domain_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, ) assert isinstance(result, GetDomainFairShareResponse) assert result.item is not None assert result.item.domain_name == domain_fixture.domain_name - assert result.item.resource_group == scaling_group_fixture.scaling_group_name + assert result.item.resource_group == scaling_group_name async def test_search_domain_fair_shares( self, @@ -52,24 +53,24 @@ async def test_search_domain_fair_shares( async def test_upsert_domain_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, ) -> None: """Upsert domain fair share → weight created/updated.""" weight = Decimal("2.5") result = await admin_registry.fair_share.upsert_domain_fair_share_weight( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, request=UpsertDomainFairShareWeightRequest(weight=weight), ) assert isinstance(result, UpsertDomainFairShareWeightResponse) assert result.item.domain_name == domain_fixture.domain_name - assert result.item.resource_group == scaling_group_fixture.scaling_group_name + assert result.item.resource_group == scaling_group_name assert result.item.spec.weight == weight # Verify the weight persists get_result = await admin_registry.fair_share.get_domain_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, ) assert get_result.item is not None @@ -82,13 +83,13 @@ class TestBulkUpsertDomainWeights: async def test_bulk_upsert_success( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, ) -> None: """Bulk upsert success → all weights updated.""" result = await admin_registry.fair_share.bulk_upsert_domain_fair_share_weight( BulkUpsertDomainFairShareWeightRequest( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, inputs=[ DomainWeightEntryInput( domain_name=domain_fixture.domain_name, @@ -103,12 +104,12 @@ async def test_bulk_upsert_success( async def test_bulk_upsert_empty_input( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: """Bulk upsert empty input → empty result (no error).""" result = await admin_registry.fair_share.bulk_upsert_domain_fair_share_weight( BulkUpsertDomainFairShareWeightRequest( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, inputs=[], ), ) @@ -118,12 +119,12 @@ async def test_bulk_upsert_empty_input( async def test_bulk_upsert_overwrite( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, ) -> None: """Bulk upsert overwrites existing weight.""" await admin_registry.fair_share.upsert_domain_fair_share_weight( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, request=UpsertDomainFairShareWeightRequest(weight=Decimal("10.0")), ) @@ -131,7 +132,7 @@ async def test_bulk_upsert_overwrite( new_weight = Decimal("3.0") result = await admin_registry.fair_share.bulk_upsert_domain_fair_share_weight( BulkUpsertDomainFairShareWeightRequest( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, inputs=[ DomainWeightEntryInput( domain_name=domain_fixture.domain_name, @@ -144,7 +145,7 @@ async def test_bulk_upsert_overwrite( assert result.upserted_count == 1 get_result = await admin_registry.fair_share.get_domain_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, ) assert get_result.item is not None @@ -157,13 +158,13 @@ class TestDomainScopeAccessControl: async def test_global_scope_regular_user_denied( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, ) -> None: """Global-scoped domain access as regular user → 403 (denied).""" with pytest.raises(PermissionDeniedError): await user_registry.fair_share.get_domain_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, ) @@ -174,7 +175,7 @@ async def test_global_scope_regular_user_denied( with pytest.raises(PermissionDeniedError): await user_registry.fair_share.upsert_domain_fair_share_weight( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, request=UpsertDomainFairShareWeightRequest(weight=Decimal("1.0")), ) @@ -182,18 +183,18 @@ async def test_global_scope_regular_user_denied( async def test_rg_scope_regular_user_allowed( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, ) -> None: """RG-scoped domain access as regular user → 200 (allowed).""" get_result = await user_registry.fair_share.rg_get_domain_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, ) assert isinstance(get_result, GetDomainFairShareResponse) search_result = await user_registry.fair_share.rg_search_domain_fair_shares( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, request=SearchDomainFairSharesRequest(), ) assert isinstance(search_result, SearchDomainFairSharesResponse) @@ -201,12 +202,12 @@ async def test_rg_scope_regular_user_allowed( async def test_global_scope_admin_allowed( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, ) -> None: """Admin global scope access → 200 (allowed).""" result = await admin_registry.fair_share.get_domain_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, ) assert isinstance(result, GetDomainFairShareResponse) @@ -218,12 +219,12 @@ class TestDomainDefaultValueFallback: async def test_get_domain_without_fair_share_default_value( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, ) -> None: """Get existing domain with no fair-share row → default value returned.""" result = await admin_registry.fair_share.get_domain_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, ) assert isinstance(result, GetDomainFairShareResponse) diff --git a/tests/component/fair_share/test_fair_share_project_weights.py b/tests/component/fair_share/test_fair_share_project_weights.py index ee94daf9577..3217be90a87 100644 --- a/tests/component/fair_share/test_fair_share_project_weights.py +++ b/tests/component/fair_share/test_fair_share_project_weights.py @@ -17,7 +17,8 @@ UpsertProjectFairShareWeightRequest, UpsertProjectFairShareWeightResponse, ) -from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData +from ai.backend.common.identifier.resource_group import ResourceGroupName +from ai.backend.testutils.fixtures import DomainFixtureData class TestProjectFairShareWeights: @@ -26,18 +27,18 @@ class TestProjectFairShareWeights: async def test_get_project_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, ) -> None: """Get project fair share → returns weight data.""" result = await admin_registry.fair_share.get_project_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, project_id=group_fixture, ) assert isinstance(result, GetProjectFairShareResponse) assert result.item is not None assert result.item.project_id == group_fixture - assert result.item.resource_group == scaling_group_fixture.scaling_group_name + assert result.item.resource_group == scaling_group_name async def test_search_project_fair_shares( self, @@ -53,14 +54,14 @@ async def test_search_project_fair_shares( async def test_upsert_project_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, ) -> None: """Upsert project fair share → weight created/updated.""" weight = Decimal("3.5") result = await admin_registry.fair_share.upsert_project_fair_share_weight( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, project_id=group_fixture, request=UpsertProjectFairShareWeightRequest( domain_name=domain_fixture.domain_name, @@ -69,12 +70,12 @@ async def test_upsert_project_fair_share( ) assert isinstance(result, UpsertProjectFairShareWeightResponse) assert result.item.project_id == group_fixture - assert result.item.resource_group == scaling_group_fixture.scaling_group_name + assert result.item.resource_group == scaling_group_name assert result.item.spec.weight == weight # Verify the weight persists get_result = await admin_registry.fair_share.get_project_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, project_id=group_fixture, ) assert get_result.item is not None @@ -87,14 +88,14 @@ class TestBulkUpsertProjectWeights: async def test_bulk_upsert_success( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, ) -> None: """Bulk upsert success → all weights updated.""" result = await admin_registry.fair_share.bulk_upsert_project_fair_share_weight( BulkUpsertProjectFairShareWeightRequest( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, inputs=[ ProjectWeightEntryInput( project_id=group_fixture, @@ -110,12 +111,12 @@ async def test_bulk_upsert_success( async def test_bulk_upsert_empty_input( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: """Bulk upsert empty input → empty result (no error).""" result = await admin_registry.fair_share.bulk_upsert_project_fair_share_weight( BulkUpsertProjectFairShareWeightRequest( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, inputs=[], ), ) @@ -125,13 +126,13 @@ async def test_bulk_upsert_empty_input( async def test_bulk_upsert_overwrite( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, ) -> None: """Bulk upsert overwrites existing weight.""" await admin_registry.fair_share.upsert_project_fair_share_weight( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, project_id=group_fixture, request=UpsertProjectFairShareWeightRequest( domain_name=domain_fixture.domain_name, @@ -142,7 +143,7 @@ async def test_bulk_upsert_overwrite( new_weight = Decimal("3.0") result = await admin_registry.fair_share.bulk_upsert_project_fair_share_weight( BulkUpsertProjectFairShareWeightRequest( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, inputs=[ ProjectWeightEntryInput( project_id=group_fixture, @@ -156,7 +157,7 @@ async def test_bulk_upsert_overwrite( assert result.upserted_count == 1 get_result = await admin_registry.fair_share.get_project_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, project_id=group_fixture, ) assert get_result.item is not None @@ -169,13 +170,13 @@ class TestProjectScopeAccessControl: async def test_global_scope_regular_user_denied( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, ) -> None: """Global-scoped project access as regular user → 403 (denied).""" with pytest.raises(PermissionDeniedError): await user_registry.fair_share.get_project_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, project_id=group_fixture, ) @@ -187,20 +188,20 @@ async def test_global_scope_regular_user_denied( async def test_rg_scope_regular_user_allowed( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, ) -> None: """RG-scoped project access as regular user → 200 (allowed).""" get_result = await user_registry.fair_share.rg_get_project_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, project_id=group_fixture, ) assert isinstance(get_result, GetProjectFairShareResponse) search_result = await user_registry.fair_share.rg_search_project_fair_shares( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, request=SearchProjectFairSharesRequest(), ) @@ -213,12 +214,12 @@ class TestProjectDefaultValueFallback: async def test_get_project_without_fair_share_default_value( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, ) -> None: """Get existing project with no fair-share row → default value returned.""" result = await admin_registry.fair_share.get_project_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, project_id=group_fixture, ) assert isinstance(result, GetProjectFairShareResponse) diff --git a/tests/component/fair_share/test_fair_share_user_weights.py b/tests/component/fair_share/test_fair_share_user_weights.py index f2ac07eb394..412b55685ee 100644 --- a/tests/component/fair_share/test_fair_share_user_weights.py +++ b/tests/component/fair_share/test_fair_share_user_weights.py @@ -18,7 +18,8 @@ UpsertUserFairShareWeightResponse, UserWeightEntryInput, ) -from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData +from ai.backend.common.identifier.resource_group import ResourceGroupName +from ai.backend.testutils.fixtures import DomainFixtureData class TestUserFairShareWeights: @@ -27,20 +28,20 @@ class TestUserFairShareWeights: async def test_get_user_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, admin_user_fixture: Any, ) -> None: """Get user fair share → returns weight data.""" result = await admin_registry.fair_share.get_user_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, ) assert isinstance(result, GetUserFairShareResponse) assert result.item is not None assert result.item.user_uuid == admin_user_fixture.user_uuid - assert result.item.resource_group == scaling_group_fixture.scaling_group_name + assert result.item.resource_group == scaling_group_name async def test_search_user_fair_shares( self, @@ -56,7 +57,7 @@ async def test_search_user_fair_shares( async def test_upsert_user_fair_share( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, admin_user_fixture: Any, @@ -64,7 +65,7 @@ async def test_upsert_user_fair_share( """Upsert user fair share → weight created/updated.""" weight = Decimal("4.5") result = await admin_registry.fair_share.upsert_user_fair_share_weight( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, request=UpsertUserFairShareWeightRequest( @@ -74,12 +75,12 @@ async def test_upsert_user_fair_share( ) assert isinstance(result, UpsertUserFairShareWeightResponse) assert result.item.user_uuid == admin_user_fixture.user_uuid - assert result.item.resource_group == scaling_group_fixture.scaling_group_name + assert result.item.resource_group == scaling_group_name assert result.item.spec.weight == weight # Verify the weight persists get_result = await admin_registry.fair_share.get_user_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, ) @@ -93,7 +94,7 @@ class TestBulkUpsertUserWeights: async def test_bulk_upsert_success( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, admin_user_fixture: Any, @@ -101,7 +102,7 @@ async def test_bulk_upsert_success( """Bulk upsert success → all weights updated.""" result = await admin_registry.fair_share.bulk_upsert_user_fair_share_weight( BulkUpsertUserFairShareWeightRequest( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, inputs=[ UserWeightEntryInput( user_uuid=admin_user_fixture.user_uuid, @@ -118,12 +119,12 @@ async def test_bulk_upsert_success( async def test_bulk_upsert_empty_input( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: """Bulk upsert empty input → empty result (no error).""" result = await admin_registry.fair_share.bulk_upsert_user_fair_share_weight( BulkUpsertUserFairShareWeightRequest( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, inputs=[], ), ) @@ -133,14 +134,14 @@ async def test_bulk_upsert_empty_input( async def test_bulk_upsert_overwrite( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, domain_fixture: DomainFixtureData, admin_user_fixture: Any, ) -> None: """Bulk upsert overwrites existing weight.""" await admin_registry.fair_share.upsert_user_fair_share_weight( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, request=UpsertUserFairShareWeightRequest( @@ -152,7 +153,7 @@ async def test_bulk_upsert_overwrite( new_weight = Decimal("3.0") result = await admin_registry.fair_share.bulk_upsert_user_fair_share_weight( BulkUpsertUserFairShareWeightRequest( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, inputs=[ UserWeightEntryInput( user_uuid=admin_user_fixture.user_uuid, @@ -167,7 +168,7 @@ async def test_bulk_upsert_overwrite( assert result.upserted_count == 1 get_result = await admin_registry.fair_share.get_user_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, ) @@ -181,14 +182,14 @@ class TestUserScopeAccessControl: async def test_global_scope_regular_user_denied( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, admin_user_fixture: Any, ) -> None: """Global-scoped user access as regular user → 403 (denied).""" with pytest.raises(PermissionDeniedError): await user_registry.fair_share.get_user_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, ) @@ -201,14 +202,14 @@ async def test_global_scope_regular_user_denied( async def test_rg_scope_regular_user_allowed( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: Any, ) -> None: """RG-scoped user access as regular user → 200 (allowed).""" get_result = await user_registry.fair_share.rg_get_user_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, @@ -216,7 +217,7 @@ async def test_rg_scope_regular_user_allowed( assert isinstance(get_result, GetUserFairShareResponse) search_result = await user_registry.fair_share.rg_search_user_fair_shares( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, domain_name=domain_fixture.domain_name, project_id=group_fixture, request=SearchUserFairSharesRequest(), @@ -230,13 +231,13 @@ class TestUserDefaultValueFallback: async def test_get_user_without_fair_share_default_value( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, admin_user_fixture: Any, ) -> None: """Get existing user with no fair-share row → default value returned.""" result = await admin_registry.fair_share.get_user_fair_share( - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, project_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, ) diff --git a/tests/component/group/test_group_management.py b/tests/component/group/test_group_management.py index 65d78663f8c..63da095e41c 100644 --- a/tests/component/group/test_group_management.py +++ b/tests/component/group/test_group_management.py @@ -36,6 +36,7 @@ from ai.backend.common.dto.manager.infra.request import UsagePerPeriodRequest from ai.backend.common.dto.manager.infra.response import UsagePerPeriodResponse from ai.backend.common.dto.manager.query import StringFilter +from ai.backend.common.identifier.resource_group import ResourceGroupName from ai.backend.common.types import ( QuotaScopeID, QuotaScopeType, @@ -49,7 +50,7 @@ from ai.backend.manager.models.kernel import KernelRow, KernelStatus from ai.backend.manager.models.session import SessionRow, SessionStatus from ai.backend.manager.models.vfolder import VFolderRow -from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData @pytest.fixture() @@ -82,7 +83,7 @@ async def test_group_for_deletion( async def group_with_vfolder_mounted( db_engine: SAEngine, domain_fixture: DomainFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, resource_policy_fixture: str, regular_user_fixture: Any, ) -> AsyncIterator[uuid.UUID]: @@ -131,7 +132,7 @@ async def group_with_vfolder_mounted( domain_name=domain_fixture.domain_name, group_id=group_id, user_uuid=user_uuid, - scaling_group_name=scaling_group_fixture.scaling_group_name, + scaling_group_name=scaling_group_name, status=SessionStatus.RUNNING, occupying_slots=ResourceSlot(), requested_slots=ResourceSlot(), @@ -153,7 +154,7 @@ async def group_with_vfolder_mounted( group_id=group_id, user_uuid=user_uuid, domain_name=domain_fixture.domain_name, - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, status=KernelStatus.RUNNING, image="python:3.9", occupied_slots=ResourceSlot({}), @@ -186,7 +187,7 @@ async def group_with_vfolder_mounted( async def group_with_active_kernel( db_engine: SAEngine, domain_fixture: DomainFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, resource_policy_fixture: str, regular_user_fixture: Any, ) -> AsyncIterator[uuid.UUID]: @@ -222,7 +223,7 @@ async def group_with_active_kernel( domain_name=domain_fixture.domain_name, group_id=group_id, user_uuid=user_uuid, - scaling_group_name=scaling_group_fixture.scaling_group_name, + scaling_group_name=scaling_group_name, status=SessionStatus.RUNNING, occupying_slots=ResourceSlot(), requested_slots=ResourceSlot(), @@ -244,7 +245,7 @@ async def group_with_active_kernel( group_id=group_id, user_uuid=user_uuid, domain_name=domain_fixture.domain_name, - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, status=KernelStatus.RUNNING, image="python:3.9", occupied_slots=ResourceSlot({}), @@ -275,7 +276,7 @@ async def group_with_active_endpoint( domain_fixture: DomainFixtureData, resource_policy_fixture: str, regular_user_fixture: Any, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> AsyncIterator[uuid.UUID]: """Create group with active endpoint.""" group_id = uuid.uuid4() @@ -335,7 +336,7 @@ async def group_with_active_endpoint( domain=domain_fixture.domain_name, created_user=user_uuid, session_owner=user_uuid, - resource_group=scaling_group_fixture.scaling_group_name, + resource_group=scaling_group_name, image=image_id, resource_slots=ResourceSlot({"cpu": "1", "mem": "1073741824"}), lifecycle_stage=EndpointLifecycle.CREATED, diff --git a/tests/component/infra/test_infra.py b/tests/component/infra/test_infra.py index 4e76726c6b5..c2b54d11d5d 100644 --- a/tests/component/infra/test_infra.py +++ b/tests/component/infra/test_infra.py @@ -29,7 +29,7 @@ UsagePerPeriodRequest, UsagePerPeriodResponse, ) -from ai.backend.testutils.fixtures import ScalingGroupFixtureData +from ai.backend.common.identifier.resource_group import ResourceGroupName class TestEtcdConfigRead: @@ -230,11 +230,11 @@ async def test_admin_checks_presets( async def test_list_presets_with_scaling_group_filter( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: """Filtering presets by scaling group.""" await admin_registry.infra.list_presets( - ListPresetsRequest(scaling_group=scaling_group_fixture.scaling_group_name) + ListPresetsRequest(scaling_group=scaling_group_name) ) diff --git a/tests/component/operations/conftest.py b/tests/component/operations/conftest.py index 557addb9512..50c8df09fb4 100644 --- a/tests/component/operations/conftest.py +++ b/tests/component/operations/conftest.py @@ -10,6 +10,7 @@ from sqlalchemy.ext.asyncio.engine import AsyncEngine as SAEngine from ai.backend.common.etcd import AsyncEtcd, ConfigScopes +from ai.backend.common.identifier.resource_group import ResourceGroupName from ai.backend.common.types import HostPortPair, ResourceSlot from ai.backend.manager.actions.validators import ActionValidators from ai.backend.manager.api.rest.error_log.handler import ErrorLogHandler @@ -30,7 +31,6 @@ from ai.backend.manager.services.error_log.service import ErrorLogService from ai.backend.manager.services.manager_admin.processors import ManagerAdminProcessors from ai.backend.manager.services.manager_admin.service import ManagerAdminService -from ai.backend.testutils.fixtures import ScalingGroupFixtureData @pytest.fixture() @@ -97,7 +97,7 @@ def server_module_registries( @pytest.fixture() async def agent_fixture( db_engine: SAEngine, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> AsyncIterator[str]: """Insert a test agent record and yield its ID.""" agent_id = f"i-test-agent-{secrets.token_hex(4)}" @@ -106,7 +106,7 @@ async def agent_fixture( sa.insert(agents).values( id=agent_id, region="local", - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, available_slots=ResourceSlot(), occupied_slots=ResourceSlot(), addr="127.0.0.1:6001", diff --git a/tests/component/resource_allocation/test_resource_allocation.py b/tests/component/resource_allocation/test_resource_allocation.py index 6c6d63579e7..edf392c83ff 100644 --- a/tests/component/resource_allocation/test_resource_allocation.py +++ b/tests/component/resource_allocation/test_resource_allocation.py @@ -19,7 +19,8 @@ ProjectResourceAllocationPayload, ResourceGroupResourceAllocationPayload, ) -from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData +from ai.backend.common.identifier.resource_group import ResourceGroupName +from ai.backend.testutils.fixtures import DomainFixtureData class TestKeypairUsage: @@ -102,11 +103,11 @@ class TestResourceGroupUsage: async def test_resource_group_usage_returns_structure( self, admin_v2_registry: V2ClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: """Query resource group usage returns capacity, used, free, max_per_node fields.""" result = await admin_v2_registry.resource_allocation.resource_group_usage( - scaling_group_fixture.scaling_group_name, + scaling_group_name, ) assert isinstance(result, ResourceGroupResourceAllocationPayload) assert result.resource_group is not None @@ -134,13 +135,13 @@ async def test_effective_returns_assignable_and_breakdown( self, admin_v2_registry: V2ClientRegistry, group_fixture: uuid.UUID, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: """Effective allocation returns assignable list and breakdown structure.""" result = await admin_v2_registry.resource_allocation.effective( EffectiveResourceAllocationInput( project_id=group_fixture, - resource_group_name=scaling_group_fixture.scaling_group_name, + resource_group_name=scaling_group_name, ), ) assert isinstance(result, EffectiveResourceAllocationPayload) @@ -152,7 +153,7 @@ async def test_effective_returns_assignable_and_breakdown( async def test_effective_with_invalid_project_returns_error( self, admin_v2_registry: V2ClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: """Effective allocation with non-existent project returns an error.""" fake_id = uuid.uuid4() @@ -160,7 +161,7 @@ async def test_effective_with_invalid_project_returns_error( await admin_v2_registry.resource_allocation.effective( EffectiveResourceAllocationInput( project_id=fake_id, - resource_group_name=scaling_group_fixture.scaling_group_name, + resource_group_name=scaling_group_name, ), ) assert exc_info.value.status in (404, 500) @@ -173,13 +174,13 @@ async def test_preset_availability_returns_list( self, admin_v2_registry: V2ClientRegistry, group_fixture: uuid.UUID, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: """Check preset availability returns a list of presets with availability status.""" result = await admin_v2_registry.resource_allocation.check_availability( CheckPresetAvailabilityInput( project_id=group_fixture, - resource_group_name=scaling_group_fixture.scaling_group_name, + resource_group_name=scaling_group_name, ), ) assert isinstance(result.presets, list) diff --git a/tests/component/resource_allocation/test_resource_allocation_visibility.py b/tests/component/resource_allocation/test_resource_allocation_visibility.py index 145e0d9deee..8b459e9c07e 100644 --- a/tests/component/resource_allocation/test_resource_allocation_visibility.py +++ b/tests/component/resource_allocation/test_resource_allocation_visibility.py @@ -24,6 +24,7 @@ from ai.backend.common.dto.manager.v2.resource_allocation.response import ( EffectiveResourceAllocationPayload, ) +from ai.backend.common.identifier.resource_group import ResourceGroupName from ai.backend.manager.actions.validators import ActionValidators from ai.backend.manager.api.adapters.resource_allocation.adapter import ResourceAllocationAdapter from ai.backend.manager.api.rest.routing import RouteRegistry @@ -51,7 +52,6 @@ from ai.backend.manager.services.resource_allocation.service import ( ResourceAllocationService, ) -from ai.backend.testutils.fixtures import ScalingGroupFixtureData if TYPE_CHECKING: from tests.component.conftest import ServerInfo, UserFixtureData @@ -205,13 +205,13 @@ async def test_regular_user_effective_hides_resource_group( self, user_v2_registry: V2ClientRegistry, group_fixture: uuid.UUID, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: """Regular user with hide_agents=True should get resource_group=null.""" result = await user_v2_registry.resource_allocation.effective( EffectiveResourceAllocationInput( project_id=group_fixture, - resource_group_name=scaling_group_fixture.scaling_group_name, + resource_group_name=scaling_group_name, ), ) assert isinstance(result, EffectiveResourceAllocationPayload) @@ -224,13 +224,13 @@ async def test_admin_effective_sees_resource_group_despite_hide_agents( self, admin_v2_registry: V2ClientRegistry, group_fixture: uuid.UUID, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: """Admin should see resource_group even when hide_agents=True.""" result = await admin_v2_registry.resource_allocation.effective( EffectiveResourceAllocationInput( project_id=group_fixture, - resource_group_name=scaling_group_fixture.scaling_group_name, + resource_group_name=scaling_group_name, ), ) assert isinstance(result, EffectiveResourceAllocationPayload) @@ -335,13 +335,13 @@ async def test_regular_user_effective_hides_project( self, user_v2_registry: V2ClientRegistry, group_fixture: uuid.UUID, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: """Regular user with group_resource_visibility=False should get project=null.""" result = await user_v2_registry.resource_allocation.effective( EffectiveResourceAllocationInput( project_id=group_fixture, - resource_group_name=scaling_group_fixture.scaling_group_name, + resource_group_name=scaling_group_name, ), ) assert isinstance(result, EffectiveResourceAllocationPayload) @@ -354,13 +354,13 @@ async def test_admin_effective_sees_project_despite_grv_off( self, admin_v2_registry: V2ClientRegistry, group_fixture: uuid.UUID, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: """Admin should see project even when group_resource_visibility=False.""" result = await admin_v2_registry.resource_allocation.effective( EffectiveResourceAllocationInput( project_id=group_fixture, - resource_group_name=scaling_group_fixture.scaling_group_name, + resource_group_name=scaling_group_name, ), ) assert isinstance(result, EffectiveResourceAllocationPayload) diff --git a/tests/component/resource_preset/test_resource_preset.py b/tests/component/resource_preset/test_resource_preset.py index 2891e0acc7f..15e8f2216ee 100644 --- a/tests/component/resource_preset/test_resource_preset.py +++ b/tests/component/resource_preset/test_resource_preset.py @@ -16,7 +16,7 @@ ListPresetsRequest, ListPresetsResponse, ) -from ai.backend.testutils.fixtures import ScalingGroupFixtureData +from ai.backend.common.identifier.resource_group import ResourceGroupName PresetFixtureData = dict[str, Any] PresetFactory = Callable[..., Coroutine[Any, Any, PresetFixtureData]] @@ -54,12 +54,12 @@ async def test_admin_lists_presets_empty( async def test_admin_lists_presets_with_scaling_group_filter( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, target_preset: PresetFixtureData, ) -> None: """Filtering by scaling group still returns system presets.""" result = await admin_registry.infra.list_presets( - ListPresetsRequest(scaling_group=scaling_group_fixture.scaling_group_name) + ListPresetsRequest(scaling_group=scaling_group_name) ) assert isinstance(result, ListPresetsResponse) assert isinstance(result.presets, list) @@ -94,13 +94,13 @@ async def test_admin_checks_presets_with_scaling_group( admin_registry: BackendAIClientRegistry, target_preset: PresetFixtureData, group_name_fixture: str, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: """Admin checks presets filtered by scaling group.""" result = await admin_registry.infra.check_presets( CheckPresetsRequest( group=group_name_fixture, - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, ) ) assert isinstance(result, CheckPresetsResponse) diff --git a/tests/component/resource_preset/test_resource_preset_full.py b/tests/component/resource_preset/test_resource_preset_full.py index cf5c6a9508f..20086de8580 100644 --- a/tests/component/resource_preset/test_resource_preset_full.py +++ b/tests/component/resource_preset/test_resource_preset_full.py @@ -26,6 +26,7 @@ CheckPresetsResponse, ListPresetsResponse, ) +from ai.backend.common.identifier.resource_group import ResourceGroupName from ai.backend.common.types import ResourceSlot from ai.backend.manager.data.resource_preset.types import ResourcePresetData from ai.backend.manager.repositories.base.creator import Creator @@ -48,7 +49,6 @@ ) from ai.backend.manager.services.resource_preset.processors import ResourcePresetProcessors from ai.backend.manager.types import OptionalState -from ai.backend.testutils.fixtures import ScalingGroupFixtureData PresetFixtureData = dict[str, Any] @@ -455,7 +455,7 @@ async def test_s6_check_presets_with_scaling_group_filter( resource_preset_processors: ResourcePresetProcessors, admin_registry: BackendAIClientRegistry, group_name_fixture: str, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, database_fixture: None, ) -> None: """S-6: check_presets with scaling_group filter returns only that SG.""" @@ -467,14 +467,14 @@ async def test_s6_check_presets_with_scaling_group_filter( result = await admin_registry.infra.check_presets( CheckPresetsRequest( group=group_name_fixture, - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, ) ) assert isinstance(result, CheckPresetsResponse) assert isinstance(result.presets, list) # When filtered, scaling_groups should contain only the specified SG if result.scaling_groups: - assert scaling_group_fixture.scaling_group_name in result.scaling_groups + assert scaling_group_name in result.scaling_groups assert len(result.scaling_groups) == 1 async def test_s7_check_presets_no_sessions_zero_usage( diff --git a/tests/component/scaling_group/test_scaling_group.py b/tests/component/scaling_group/test_scaling_group.py index 5c5afb4fc4c..173c0003667 100644 --- a/tests/component/scaling_group/test_scaling_group.py +++ b/tests/component/scaling_group/test_scaling_group.py @@ -7,14 +7,14 @@ from ai.backend.client.v2.exceptions import NotFoundError from ai.backend.client.v2.registry import BackendAIClientRegistry from ai.backend.common.dto.manager.scaling_group import ListScalingGroupsResponse -from ai.backend.testutils.fixtures import ScalingGroupFixtureData +from ai.backend.common.identifier.resource_group import ResourceGroupName class TestScalingGroupList: async def test_admin_lists_scaling_groups( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, ) -> None: """Admin can list scaling groups; the fixture sgroup is visible via domain association.""" @@ -23,12 +23,12 @@ async def test_admin_lists_scaling_groups( ) assert isinstance(result, ListScalingGroupsResponse) names = [sg.name for sg in result.scaling_groups] - assert scaling_group_fixture.scaling_group_name in names + assert scaling_group_name in names async def test_list_scaling_groups_with_group_filter( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, ) -> None: """Filtering by group UUID returns the domain-associated scaling group.""" @@ -36,14 +36,12 @@ async def test_list_scaling_groups_with_group_filter( group=str(group_fixture), ) assert isinstance(result, ListScalingGroupsResponse) - assert any( - sg.name == scaling_group_fixture.scaling_group_name for sg in result.scaling_groups - ) + assert any(sg.name == scaling_group_name for sg in result.scaling_groups) async def test_regular_user_lists_public_scaling_groups( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, ) -> None: """Regular users can list scaling groups; only public ones are returned. @@ -59,12 +57,12 @@ async def test_regular_user_lists_public_scaling_groups( assert isinstance(result, ListScalingGroupsResponse) # The fixture sgroup defaults to is_public=True, so it is visible. names = [sg.name for sg in result.scaling_groups] - assert scaling_group_fixture.scaling_group_name in names + assert scaling_group_name in names async def test_list_scaling_groups_with_uuid_string_parameter( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, ) -> None: """BA-5411: Passing group UUID as a string correctly resolves the group.""" @@ -73,16 +71,14 @@ async def test_list_scaling_groups_with_uuid_string_parameter( group=str(group_fixture), ) assert isinstance(result, ListScalingGroupsResponse) - assert any( - sg.name == scaling_group_fixture.scaling_group_name for sg in result.scaling_groups - ) + assert any(sg.name == scaling_group_name for sg in result.scaling_groups) class TestScalingGroupWsproxyVersion: async def test_admin_gets_wsproxy_version( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: """Admin requests wsproxy version for a scaling group without wsproxy_addr configured. @@ -92,7 +88,7 @@ async def test_admin_gets_wsproxy_version( """ with pytest.raises(NotFoundError): await admin_registry.scaling_group.get_wsproxy_version( - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, ) async def test_get_wsproxy_version_nonexistent( @@ -108,7 +104,7 @@ async def test_get_wsproxy_version_nonexistent( async def test_regular_user_gets_wsproxy_version_not_found( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: """Regular users are not blocked from querying wsproxy version. @@ -118,5 +114,5 @@ async def test_regular_user_gets_wsproxy_version_not_found( """ with pytest.raises(NotFoundError): await user_registry.scaling_group.get_wsproxy_version( - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, ) diff --git a/tests/component/scaling_group/test_scaling_group_associations.py b/tests/component/scaling_group/test_scaling_group_associations.py index 2908a54f10b..c3f9ccac674 100644 --- a/tests/component/scaling_group/test_scaling_group_associations.py +++ b/tests/component/scaling_group/test_scaling_group_associations.py @@ -23,12 +23,13 @@ from ai.backend.client.v2.registry import BackendAIClientRegistry from ai.backend.common.dto.manager.scaling_group import ListScalingGroupsResponse +from ai.backend.common.identifier.resource_group import ResourceGroupName from ai.backend.manager.models.scaling_group import ( ScalingGroupOpts, scaling_groups, sgroups_for_domains, ) -from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData @pytest.fixture() @@ -78,7 +79,7 @@ async def test_s_visibility_public_sgroup_visible_to_user( self, admin_registry: BackendAIClientRegistry, user_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, ) -> None: """Public scaling group is visible to both admin and regular user.""" @@ -86,18 +87,13 @@ async def test_s_visibility_public_sgroup_visible_to_user( group=str(group_fixture), ) assert isinstance(admin_result, ListScalingGroupsResponse) - assert any( - sg.name == scaling_group_fixture.scaling_group_name - for sg in admin_result.scaling_groups - ) + assert any(sg.name == scaling_group_name for sg in admin_result.scaling_groups) user_result = await user_registry.scaling_group.list_scaling_groups( group=str(group_fixture), ) assert isinstance(user_result, ListScalingGroupsResponse) - assert any( - sg.name == scaling_group_fixture.scaling_group_name for sg in user_result.scaling_groups - ) + assert any(sg.name == scaling_group_name for sg in user_result.scaling_groups) async def test_s_visibility_private_sgroup_hidden_from_user( self, @@ -145,7 +141,7 @@ class TestScalingGroupPermissions: async def test_regular_user_can_list_scaling_groups( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, ) -> None: """F-AUTH-LIST: Regular user can list public scaling groups (auth_required, not superadmin). @@ -160,4 +156,4 @@ async def test_regular_user_can_list_scaling_groups( assert isinstance(result, ListScalingGroupsResponse) names = [sg.name for sg in result.scaling_groups] # The fixture sgroup is public — regular user should see it - assert scaling_group_fixture.scaling_group_name in names + assert scaling_group_name in names diff --git a/tests/component/scaling_group/test_scaling_group_crud.py b/tests/component/scaling_group/test_scaling_group_crud.py index b92f4434f10..e93d8cbcc57 100644 --- a/tests/component/scaling_group/test_scaling_group_crud.py +++ b/tests/component/scaling_group/test_scaling_group_crud.py @@ -11,13 +11,14 @@ from ai.backend.client.v2.exceptions import NotFoundError from ai.backend.client.v2.registry import BackendAIClientRegistry from ai.backend.common.dto.manager.scaling_group import ListScalingGroupsResponse +from ai.backend.common.identifier.resource_group import ResourceGroupName from ai.backend.manager.models.scaling_group import ( ScalingGroupOpts, scaling_groups, sgroups_for_domains, sgroups_for_groups, ) -from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData @pytest.fixture() @@ -104,7 +105,7 @@ class TestScalingGroupCRUDLifecycle: async def test_multiple_scaling_groups_listed( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, extra_scaling_group_fixture: str, group_fixture: uuid.UUID, ) -> None: @@ -114,13 +115,13 @@ async def test_multiple_scaling_groups_listed( ) assert isinstance(result, ListScalingGroupsResponse) names = [sg.name for sg in result.scaling_groups] - assert scaling_group_fixture.scaling_group_name in names + assert scaling_group_name in names assert extra_scaling_group_fixture in names async def test_scaling_group_domain_association( self, admin_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, group_fixture: uuid.UUID, ) -> None: """A scaling group associated with the test domain is visible via group filter.""" @@ -128,20 +129,18 @@ async def test_scaling_group_domain_association( group=str(group_fixture), ) assert isinstance(result, ListScalingGroupsResponse) - assert any( - sg.name == scaling_group_fixture.scaling_group_name for sg in result.scaling_groups - ) + assert any(sg.name == scaling_group_name for sg in result.scaling_groups) async def test_regular_user_sees_public_scaling_groups( self, user_registry: BackendAIClientRegistry, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, private_scaling_group_fixture: str, group_fixture: uuid.UUID, ) -> None: """Regular user sees public scaling groups but not private ones. - The fixture scaling_group_fixture defaults to is_public=True, so it + The fixture scaling_group_name defaults to is_public=True, so it should appear. private_scaling_group_fixture is is_public=False, so it should NOT appear for a regular user. """ @@ -150,7 +149,7 @@ async def test_regular_user_sees_public_scaling_groups( ) assert isinstance(result, ListScalingGroupsResponse) names = [sg.name for sg in result.scaling_groups] - assert scaling_group_fixture.scaling_group_name in names + assert scaling_group_name in names assert private_scaling_group_fixture not in names async def test_admin_sees_private_scaling_groups( diff --git a/tests/component/session/conftest.py b/tests/component/session/conftest.py index 876dfe3b2f5..c7bd1f91b6d 100644 --- a/tests/component/session/conftest.py +++ b/tests/component/session/conftest.py @@ -14,6 +14,7 @@ from sqlalchemy.ext.asyncio.engine import AsyncEngine as SAEngine from ai.backend.common.bgtask.bgtask import BackgroundTaskManager +from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName from ai.backend.common.plugin.monitor import ErrorPluginContext from ai.backend.common.types import AgentId, ResourceSlot, SessionId, SessionTypes from ai.backend.manager.actions.validators import ActionValidators @@ -39,7 +40,7 @@ from ai.backend.manager.services.session.processors import SessionProcessors from ai.backend.manager.services.session.service import SessionService, SessionServiceArgs from ai.backend.manager.services.vfolder.processors.vfolder import VFolderProcessors -from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData @dataclass @@ -164,7 +165,7 @@ def server_module_registries( @pytest.fixture() async def agent_fixture( db_engine: SAEngine, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> AsyncIterator[AgentId]: """Insert a test agent record and yield its ID.""" agent_id = AgentId(f"i-test-agent-{secrets.token_hex(4)}") @@ -173,7 +174,7 @@ async def agent_fixture( sa.insert(AgentRow).values( id=agent_id, region="local", - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, available_slots=ResourceSlot(), occupied_slots=ResourceSlot(), addr="127.0.0.1:6001", @@ -192,7 +193,8 @@ async def session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, + scaling_group_id: ResourceGroupID, agent_fixture: AgentId, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING session + kernel directly in the database. @@ -227,8 +229,8 @@ async def session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture.scaling_group_name, - resource_group_id=scaling_group_fixture.scaling_group_id, + scaling_group_name=scaling_group_name, + resource_group_id=scaling_group_id, status=SessionStatus.RUNNING, status_info="", status_history=status_history, @@ -253,7 +255,7 @@ async def session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, agent=agent_fixture, status=KernelStatus.RUNNING, status_info="", @@ -289,7 +291,8 @@ async def terminated_session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, + scaling_group_id: ResourceGroupID, ) -> AsyncIterator[SessionSeedData]: """Seed a TERMINATED session with container_log in the kernel. @@ -322,8 +325,8 @@ async def terminated_session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture.scaling_group_name, - resource_group_id=scaling_group_fixture.scaling_group_id, + scaling_group_name=scaling_group_name, + resource_group_id=scaling_group_id, status=SessionStatus.TERMINATED, status_info="user-requested", status_history=status_history, @@ -349,7 +352,7 @@ async def terminated_session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, status=KernelStatus.TERMINATED, status_info="user-requested", occupied_slots=ResourceSlot(), @@ -386,7 +389,8 @@ async def user_session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, regular_user_fixture: UserFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, + scaling_group_id: ResourceGroupID, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING session owned by the regular user.""" unique = secrets.token_hex(4) @@ -414,8 +418,8 @@ async def user_session_seed( group_id=group_fixture, user_uuid=regular_user_fixture.user_uuid, access_key=regular_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture.scaling_group_name, - resource_group_id=scaling_group_fixture.scaling_group_id, + scaling_group_name=scaling_group_name, + resource_group_id=scaling_group_id, status=SessionStatus.RUNNING, status_info="", status_history=status_history, @@ -440,7 +444,7 @@ async def user_session_seed( group_id=group_fixture, user_uuid=regular_user_fixture.user_uuid, access_key=regular_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, status=KernelStatus.RUNNING, status_info="", occupied_slots=ResourceSlot(), diff --git a/tests/component/session/test_session_batch.py b/tests/component/session/test_session_batch.py index 3e6ea75aab6..3333aa5fa87 100644 --- a/tests/component/session/test_session_batch.py +++ b/tests/component/session/test_session_batch.py @@ -19,9 +19,9 @@ GetContainerLogsResponse, StartServiceResponse, ) +from ai.backend.common.identifier.resource_group import ResourceGroupName from ai.backend.manager.models.kernel import kernels from ai.backend.manager.models.scaling_group import scaling_groups -from ai.backend.testutils.fixtures import ScalingGroupFixtureData from .conftest import SessionSeedData @@ -132,7 +132,7 @@ async def test_start_service_returns_token_and_wsproxy( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, agent_registry: AsyncMock, appproxy_client_pool: AsyncMock, ) -> None: @@ -142,7 +142,7 @@ async def test_start_service_returns_token_and_wsproxy( async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) + .where(scaling_groups.c.name == scaling_group_name) .values(wsproxy_addr=wsproxy_addr) ) diff --git a/tests/component/session/test_session_execution.py b/tests/component/session/test_session_execution.py index c198c09a08e..7fcc2b45f30 100644 --- a/tests/component/session/test_session_execution.py +++ b/tests/component/session/test_session_execution.py @@ -20,9 +20,9 @@ ExecuteResponse, StartServiceResponse, ) +from ai.backend.common.identifier.resource_group import ResourceGroupName from ai.backend.manager.models.kernel import kernels from ai.backend.manager.models.scaling_group import scaling_groups -from ai.backend.testutils.fixtures import ScalingGroupFixtureData from .conftest import SessionSeedData @@ -333,7 +333,7 @@ async def test_start_service_returns_token_and_wsproxy( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, agent_registry: AsyncMock, appproxy_client_pool: AsyncMock, ) -> None: @@ -341,7 +341,7 @@ async def test_start_service_returns_token_and_wsproxy( async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) + .where(scaling_groups.c.name == scaling_group_name) .values(wsproxy_addr=_WSPROXY_ADDR) ) async with db_engine.begin() as conn: @@ -383,7 +383,7 @@ async def test_start_service_with_specific_port( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, agent_registry: AsyncMock, appproxy_client_pool: AsyncMock, ) -> None: @@ -401,7 +401,7 @@ async def test_start_service_with_specific_port( async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) + .where(scaling_groups.c.name == scaling_group_name) .values(wsproxy_addr=_WSPROXY_ADDR) ) async with db_engine.begin() as conn: @@ -441,7 +441,7 @@ async def test_start_service_with_arguments_and_envs( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, agent_registry: AsyncMock, appproxy_client_pool: AsyncMock, ) -> None: @@ -449,7 +449,7 @@ async def test_start_service_with_arguments_and_envs( async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) + .where(scaling_groups.c.name == scaling_group_name) .values(wsproxy_addr=_WSPROXY_ADDR) ) async with db_engine.begin() as conn: @@ -497,7 +497,7 @@ async def test_start_service_kernel_host_from_agent_addr( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, agent_registry: AsyncMock, appproxy_client_pool: AsyncMock, ) -> None: @@ -505,7 +505,7 @@ async def test_start_service_kernel_host_from_agent_addr( async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) + .where(scaling_groups.c.name == scaling_group_name) .values(wsproxy_addr=_WSPROXY_ADDR) ) async with db_engine.begin() as conn: @@ -557,7 +557,7 @@ async def test_start_service_wsproxy_advertise_address_fallback( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, agent_registry: AsyncMock, appproxy_client_pool: AsyncMock, ) -> None: @@ -565,7 +565,7 @@ async def test_start_service_wsproxy_advertise_address_fallback( async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) + .where(scaling_groups.c.name == scaling_group_name) .values(wsproxy_addr=_WSPROXY_ADDR) ) async with db_engine.begin() as conn: @@ -615,14 +615,14 @@ async def test_start_nonexistent_service( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, appproxy_client_pool: AsyncMock, ) -> None: """F-BIZ-1: Service name not in service_ports → NotFoundError (AppNotFound).""" async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) + .where(scaling_groups.c.name == scaling_group_name) .values(wsproxy_addr=_WSPROXY_ADDR) ) async with db_engine.begin() as conn: @@ -653,7 +653,7 @@ async def test_start_inference_app_rejected( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, appproxy_client_pool: AsyncMock, ) -> None: """F-BIZ-2: Inference app → InvalidRequestError.""" @@ -670,7 +670,7 @@ async def test_start_inference_app_rejected( async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) + .where(scaling_groups.c.name == scaling_group_name) .values(wsproxy_addr=_WSPROXY_ADDR) ) async with db_engine.begin() as conn: @@ -701,14 +701,14 @@ async def test_start_with_invalid_port( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, appproxy_client_pool: AsyncMock, ) -> None: """F-BIZ-3: Port not in container_ports → InvalidRequestError.""" async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) + .where(scaling_groups.c.name == scaling_group_name) .values(wsproxy_addr=_WSPROXY_ADDR) ) async with db_engine.begin() as conn: @@ -739,14 +739,14 @@ async def test_start_with_no_wsproxy_addr( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, ) -> None: """F-BIZ-5: Scaling group has no wsproxy_addr → ServerError (ServiceUnavailable).""" # Ensure wsproxy_addr is NULL on the scaling group async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) + .where(scaling_groups.c.name == scaling_group_name) .values(wsproxy_addr=None) ) @@ -761,7 +761,7 @@ async def test_start_agent_returns_failed_status( admin_registry: BackendAIClientRegistry, session_seed: SessionSeedData, db_engine: SAEngine, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, agent_registry: AsyncMock, appproxy_client_pool: AsyncMock, ) -> None: @@ -769,7 +769,7 @@ async def test_start_agent_returns_failed_status( async with db_engine.begin() as conn: await conn.execute( sa.update(scaling_groups) - .where(scaling_groups.c.name == scaling_group_fixture.scaling_group_name) + .where(scaling_groups.c.name == scaling_group_name) .values(wsproxy_addr=_WSPROXY_ADDR) ) async with db_engine.begin() as conn: diff --git a/tests/component/session/test_session_lifecycle.py b/tests/component/session/test_session_lifecycle.py index 15e49e29e38..b75929cf974 100644 --- a/tests/component/session/test_session_lifecycle.py +++ b/tests/component/session/test_session_lifecycle.py @@ -20,12 +20,13 @@ from ai.backend.common.dto.manager.session.response import ( GetStatusHistoryResponse, ) +from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName from ai.backend.common.types import ResourceSlot, SessionId, SessionTypes from ai.backend.manager.data.kernel.types import KernelStatus from ai.backend.manager.data.session.types import SessionStatus from ai.backend.manager.models.kernel import kernels from ai.backend.manager.models.session import SessionRow -from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData from .conftest import SessionSeedData, UserFixtureData @@ -73,7 +74,8 @@ async def degraded_session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, + scaling_group_id: ResourceGroupID, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING_DEGRADED session with two kernels (one RUNNING, one ERROR).""" unique = secrets.token_hex(4) @@ -96,7 +98,7 @@ async def degraded_session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, now=now, ) @@ -114,8 +116,8 @@ async def degraded_session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture.scaling_group_name, - resource_group_id=scaling_group_fixture.scaling_group_id, + scaling_group_name=scaling_group_name, + resource_group_id=scaling_group_id, status=SessionStatus.RUNNING_DEGRADED, status_info="", status_history=status_history, @@ -167,7 +169,8 @@ async def full_lifecycle_session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, + scaling_group_id: ResourceGroupID, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING session with a full lifecycle status_history (PENDING → SCHEDULED → PREPARING → PULLING → CREATING → RUNNING). @@ -196,7 +199,7 @@ async def full_lifecycle_session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, now=now, ) @@ -214,8 +217,8 @@ async def full_lifecycle_session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture.scaling_group_name, - resource_group_id=scaling_group_fixture.scaling_group_id, + scaling_group_name=scaling_group_name, + resource_group_id=scaling_group_id, status=SessionStatus.RUNNING, status_info="", status_history=status_history, diff --git a/tests/component/session/test_session_query.py b/tests/component/session/test_session_query.py index 4dfb847b603..4d790d5cdc3 100644 --- a/tests/component/session/test_session_query.py +++ b/tests/component/session/test_session_query.py @@ -32,6 +32,7 @@ GetStatusHistoryResponse, MatchSessionsResponse, ) +from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName from ai.backend.common.types import ResourceSlot, SessionId, SessionTypes from ai.backend.manager.api.rest.compute_sessions.handler import ComputeSessionsHandler from ai.backend.manager.api.rest.compute_sessions.registry import ( @@ -54,7 +55,7 @@ from ai.backend.manager.services.auth.processors import AuthProcessors from ai.backend.manager.services.session.processors import SessionProcessors from ai.backend.manager.services.vfolder.processors.vfolder import VFolderProcessors -from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData from .conftest import SessionSeedData, UserFixtureData @@ -97,7 +98,8 @@ async def system_session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, + scaling_group_id: ResourceGroupID, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING SYSTEM session with an agent row for direct access info tests. @@ -133,7 +135,7 @@ async def system_session_seed( id=agent_id, status=AgentStatus.ALIVE, region="local", - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, available_slots=ResourceSlot(), occupied_slots=ResourceSlot(), addr="tcp://127.0.0.1:6001", @@ -156,8 +158,8 @@ async def system_session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture.scaling_group_name, - resource_group_id=scaling_group_fixture.scaling_group_id, + scaling_group_name=scaling_group_name, + resource_group_id=scaling_group_id, status=SessionStatus.RUNNING, status_info="", status_history=status_history, @@ -182,7 +184,7 @@ async def system_session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, status=KernelStatus.RUNNING, status_info="", occupied_slots=ResourceSlot(), @@ -221,7 +223,8 @@ async def system_session_no_agent_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, + scaling_group_id: ResourceGroupID, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING SYSTEM session without an agent row. @@ -253,8 +256,8 @@ async def system_session_no_agent_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture.scaling_group_name, - resource_group_id=scaling_group_fixture.scaling_group_id, + scaling_group_name=scaling_group_name, + resource_group_id=scaling_group_id, status=SessionStatus.RUNNING, status_info="", status_history=status_history, @@ -279,7 +282,7 @@ async def system_session_no_agent_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, status=KernelStatus.RUNNING, status_info="", occupied_slots=ResourceSlot(), @@ -314,7 +317,8 @@ async def second_session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, + scaling_group_id: ResourceGroupID, session_seed: SessionSeedData, ) -> AsyncIterator[SessionSeedData]: """Seed a second INTERACTIVE session alongside session_seed. @@ -348,8 +352,8 @@ async def second_session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture.scaling_group_name, - resource_group_id=scaling_group_fixture.scaling_group_id, + scaling_group_name=scaling_group_name, + resource_group_id=scaling_group_id, status=SessionStatus.RUNNING, status_info="", status_history=status_history, @@ -374,7 +378,7 @@ async def second_session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, status=KernelStatus.RUNNING, status_info="", occupied_slots=ResourceSlot(), diff --git a/tests/component/session_v2/conftest.py b/tests/component/session_v2/conftest.py index 0c301096ca1..94166fd67e9 100644 --- a/tests/component/session_v2/conftest.py +++ b/tests/component/session_v2/conftest.py @@ -28,7 +28,7 @@ ScopeType, ) from ai.backend.common.identifier.domain import DomainID -from ai.backend.common.identifier.resource_group import ResourceGroupID +from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName from ai.backend.common.plugin.monitor import ErrorPluginContext from ai.backend.common.types import ResourceSlot, SessionId, SessionTypes from ai.backend.manager.actions.validators import ActionValidators @@ -60,7 +60,7 @@ from ai.backend.manager.services.processors import Processors from ai.backend.manager.services.session.processors import SessionProcessors from ai.backend.manager.services.session.service import SessionService, SessionServiceArgs -from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData if TYPE_CHECKING: from sqlalchemy.ext.asyncio.engine import AsyncEngine as SAEngine @@ -421,7 +421,8 @@ async def admin_session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, + scaling_group_id: ResourceGroupID, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING session owned by the admin user.""" seed = await _seed_session( @@ -431,8 +432,8 @@ async def admin_session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture.scaling_group_name, - resource_group_id=scaling_group_fixture.scaling_group_id, + scaling_group=scaling_group_name, + resource_group_id=scaling_group_id, ) yield seed await _cleanup_session(db_engine, seed.session_id) @@ -444,7 +445,8 @@ async def user_session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, regular_user_fixture: UserFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, + scaling_group_id: ResourceGroupID, user_system_role: uuid.UUID, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING session owned by the regular user. @@ -458,8 +460,8 @@ async def user_session_seed( group_id=group_fixture, user_uuid=regular_user_fixture.user_uuid, access_key=regular_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture.scaling_group_name, - resource_group_id=scaling_group_fixture.scaling_group_id, + scaling_group=scaling_group_name, + resource_group_id=scaling_group_id, ) yield seed await _cleanup_session(db_engine, seed.session_id) diff --git a/tests/component/stream/conftest.py b/tests/component/stream/conftest.py index 580493c4ce9..aebbe4615d3 100644 --- a/tests/component/stream/conftest.py +++ b/tests/component/stream/conftest.py @@ -14,6 +14,7 @@ from sqlalchemy.ext.asyncio.engine import AsyncEngine as SAEngine from ai.backend.common.etcd import AsyncEtcd +from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName from ai.backend.common.plugin.monitor import ErrorPluginContext from ai.backend.common.types import ResourceSlot, SessionId, SessionTypes from ai.backend.manager.actions.processor import ActionProcessor @@ -34,7 +35,7 @@ from ai.backend.manager.services.session.service import SessionService, SessionServiceArgs from ai.backend.manager.services.stream.processors import StreamProcessors from ai.backend.manager.services.stream.service import StreamService -from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData _STREAMING_SERVER_SUBAPP_MODULES = (_auth_api,) @@ -138,7 +139,8 @@ async def session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, + scaling_group_id: ResourceGroupID, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING session + kernel with service_ports in the database.""" unique = secrets.token_hex(4) @@ -166,8 +168,8 @@ async def session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture.scaling_group_name, - resource_group_id=scaling_group_fixture.scaling_group_id, + scaling_group_name=scaling_group_name, + resource_group_id=scaling_group_id, status=SessionStatus.RUNNING, status_info="", status_history=status_history, @@ -192,7 +194,7 @@ async def session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, status=KernelStatus.RUNNING, status_info="", occupied_slots=ResourceSlot(), diff --git a/tests/component/streaming/conftest.py b/tests/component/streaming/conftest.py index 065f1577206..2c65acffac0 100644 --- a/tests/component/streaming/conftest.py +++ b/tests/component/streaming/conftest.py @@ -14,6 +14,7 @@ from sqlalchemy.ext.asyncio.engine import AsyncEngine as SAEngine from ai.backend.common.etcd import AsyncEtcd +from ai.backend.common.identifier.resource_group import ResourceGroupID, ResourceGroupName from ai.backend.common.plugin.monitor import ErrorPluginContext from ai.backend.common.types import ResourceSlot, SessionId, SessionTypes from ai.backend.manager.api.rest.middleware import auth as _auth_api @@ -31,7 +32,7 @@ from ai.backend.manager.repositories.stream.repository import StreamRepository from ai.backend.manager.services.stream.processors import StreamProcessors from ai.backend.manager.services.stream.service import StreamService -from ai.backend.testutils.fixtures import DomainFixtureData, ScalingGroupFixtureData +from ai.backend.testutils.fixtures import DomainFixtureData _STREAMING_SERVER_SUBAPP_MODULES = (_auth_api,) @@ -105,7 +106,8 @@ async def session_seed( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, + scaling_group_id: ResourceGroupID, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING session + kernel with service_ports in the database. @@ -137,8 +139,8 @@ async def session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture.scaling_group_name, - resource_group_id=scaling_group_fixture.scaling_group_id, + scaling_group_name=scaling_group_name, + resource_group_id=scaling_group_id, status=SessionStatus.RUNNING, status_info="", status_history=status_history, @@ -163,7 +165,7 @@ async def session_seed( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, status=KernelStatus.RUNNING, status_info="", occupied_slots=ResourceSlot(), @@ -211,7 +213,8 @@ async def session_seed_no_ports( domain_fixture: DomainFixtureData, group_fixture: uuid.UUID, admin_user_fixture: UserFixtureData, - scaling_group_fixture: ScalingGroupFixtureData, + scaling_group_name: ResourceGroupName, + scaling_group_id: ResourceGroupID, ) -> AsyncIterator[SessionSeedData]: """Seed a RUNNING session + kernel with ``service_ports=None``. @@ -243,8 +246,8 @@ async def session_seed_no_ports( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group_name=scaling_group_fixture.scaling_group_name, - resource_group_id=scaling_group_fixture.scaling_group_id, + scaling_group_name=scaling_group_name, + resource_group_id=scaling_group_id, status=SessionStatus.RUNNING, status_info="", status_history=status_history, @@ -269,7 +272,7 @@ async def session_seed_no_ports( group_id=group_fixture, user_uuid=admin_user_fixture.user_uuid, access_key=admin_user_fixture.keypair.access_key, - scaling_group=scaling_group_fixture.scaling_group_name, + scaling_group=scaling_group_name, status=KernelStatus.RUNNING, status_info="", occupied_slots=ResourceSlot(), From 77f9c1efeb297f5cbded9fc41bb6715a0f215746 Mon Sep 17 00:00:00 2001 From: BoKeum Date: Fri, 3 Jul 2026 16:49:42 +0900 Subject: [PATCH 6/6] fix: linearize alembic migration head --- .../versions/ada41cb881bb_add_session_scope_id_columns.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/ai/backend/manager/models/alembic/versions/ada41cb881bb_add_session_scope_id_columns.py b/src/ai/backend/manager/models/alembic/versions/ada41cb881bb_add_session_scope_id_columns.py index 255bab14ab9..9d427832f0f 100644 --- a/src/ai/backend/manager/models/alembic/versions/ada41cb881bb_add_session_scope_id_columns.py +++ b/src/ai/backend/manager/models/alembic/versions/ada41cb881bb_add_session_scope_id_columns.py @@ -1,7 +1,7 @@ """add session scope id columns Revision ID: ada41cb881bb -Revises: b4c5d6e7f8a9 +Revises: a379b72f1206 Create Date: 2026-06-30 """ @@ -13,7 +13,7 @@ # revision identifiers, used by Alembic. revision = "ada41cb881bb" -down_revision = "b4c5d6e7f8a9" +down_revision = "a379b72f1206" # Part of: 26.7.0 branch_labels = None depends_on = None