refactor: Single-source active/dead set definitions of ContainerStatus (#11213)

Author: achimnol

changes/11213.fix.md

@@ -0,0 +1 @@
+Single-source active/dead flag set definitions of `ContainerStatus` to prevent potential mismatch in future code edits

src/ai/backend/agent/agent.py

@@ -276,17 +276,8 @@
 
 _sentinel = Sentinel.TOKEN
 
-ACTIVE_STATUS_SET = frozenset([
-    ContainerStatus.RUNNING,
-    ContainerStatus.RESTARTING,
-    ContainerStatus.PAUSED,
-])
-
-DEAD_STATUS_SET = frozenset([
-    ContainerStatus.EXITED,
-    ContainerStatus.DEAD,
-    ContainerStatus.REMOVING,
-])
+ACTIVE_STATUS_SET = ContainerStatus.active_set()
+DEAD_STATUS_SET = ContainerStatus.dead_set()
 
 COMMIT_STATUS_EXPIRE: Final[int] = 13
 EVENT_DISPATCHER_CONSUMER_GROUP: Final = "agent"

src/ai/backend/common/types.py

@@ -8,7 +8,6 @@
 import math
 import numbers
 import textwrap
-import uuid
 from abc import ABC, ABCMeta, abstractmethod
 from collections import UserDict, UserString, defaultdict, namedtuple
 from collections.abc import AsyncIterator, Iterable, Mapping, Sequence
@@ -27,7 +26,6 @@
     NewType,
     NotRequired,
     Self,
-    TypeAlias,
     TypedDict,
     TypeVar,
     cast,
@@ -474,6 +472,7 @@ def _validate_slot_name(v: Any) -> SlotName:
 
 
 class ContainerStatus(enum.StrEnum):
+    CREATED = "created"
     RUNNING = "running"
     RESTARTING = "restarting"
     PAUSED = "paused"
@@ -780,7 +779,7 @@ class VFolderMountOptions:
 class VFolderMountRequest:
     """A single vfolder mount request combining reference, destination path, and options."""
 
-    ref: str | uuid.UUID  # vfolder name (with optional /subpath) or UUID
+    ref: str | UUID  # vfolder name (with optional /subpath) or UUID
     dst_path: str | None = None  # custom mount destination path
     options: VFolderMountOptions = attrs.Factory(VFolderMountOptions)
 
@@ -1394,7 +1393,7 @@ def as_trafaret(cls) -> t.Trafaret:
         raise NotImplementedError
 
 
-VolumeID: TypeAlias = uuid.UUID
+VolumeID = NewType("VolumeID", UUID)
 
 
 @attrs.define(slots=True, frozen=True)

src/ai/backend/storage/services/service.py

@@ -126,11 +126,11 @@ async def get_volumes(self) -> list[VolumeMeta]:
         volumes = self._volume_pool.list_volumes()
         return [
             VolumeMeta(
-                volume_id=uuid.UUID(volume_id),
+                volume_id=VolumeID(uuid.UUID(volume_id)),
                 backend=info.backend,
                 path=info.path,
                 fsprefix=info.fsprefix,
-                capabilities=await self._get_capabilities(uuid.UUID(volume_id)),
+                capabilities=await self._get_capabilities(VolumeID(uuid.UUID(volume_id))),
             )
             for volume_id, info in volumes.items()
         ]

src/ai/backend/storage/volumes/pool.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import logging
+import uuid
 from collections.abc import AsyncIterator, Mapping
 from contextlib import asynccontextmanager as actxmgr
 from pathlib import Path
@@ -82,7 +83,7 @@ async def create(
         volumes_by_name: dict[str, AbstractVolume] = {}
         for raw_volume_id, config in local_config.volume.items():
             try:
-                volume_id = VolumeID(raw_volume_id)
+                volume_id = VolumeID(uuid.UUID(raw_volume_id))
             except (ValueError, TypeError):
                 volumes_by_name[raw_volume_id] = await cls._init_volume(
                     config,

tests/unit/storage-proxy/vfolder/test_handler.py

@@ -10,7 +10,7 @@
 from aiohttp import web
 from aiohttp.web import StreamResponse
 
-from ai.backend.common.types import QuotaConfig, QuotaScopeID, QuotaScopeType, VFolderID
+from ai.backend.common.types import QuotaConfig, QuotaScopeID, QuotaScopeType, VFolderID, VolumeID
 from ai.backend.storage.api.vfolder.handler import VFolderHandler, VFolderServiceProtocol
 from ai.backend.storage.volumes.types import (
     QuotaScopeKey,
@@ -27,7 +27,7 @@
 
 class MockVFolderService(VFolderServiceProtocol):
     @override
-    async def get_volume(self, volume_id: uuid.UUID) -> VolumeMeta:
+    async def get_volume(self, volume_id: VolumeID) -> VolumeMeta:
         return VolumeMeta(
             volume_id=volume_id,
             backend="vfs",
@@ -39,8 +39,8 @@ async def get_volume(self, volume_id: uuid.UUID) -> VolumeMeta:
     @override
     async def get_volumes(self) -> list[VolumeMeta]:
         volumes = {
-            UUID1: {"backend": "vfs", "path": "/mnt/volume1", "fsprefix": "vfs-test-1"},
-            UUID2: {"backend": "nfs", "path": "/mnt/volume2", "fsprefix": "nfs-test-2"},
+            VolumeID(UUID1): {"backend": "vfs", "path": "/mnt/volume1", "fsprefix": "vfs-test-1"},
+            VolumeID(UUID2): {"backend": "nfs", "path": "/mnt/volume2", "fsprefix": "nfs-test-2"},
         }
         return [
             VolumeMeta(

tests/unit/storage-proxy/vfolder/test_pool.py

@@ -1,3 +1,4 @@
+import uuid
 from pathlib import Path, PurePath
 from unittest.mock import AsyncMock, MagicMock
 
@@ -53,7 +54,7 @@ def mock_storage_plugin_ctx() -> AsyncMock:
 
 async def test_get_volume(mock_volume: AsyncMock, mock_storage_plugin_ctx: AsyncMock) -> None:
     # Create a VolumePool with mocked volumes
-    volume_id = VolumeID("550e8400-e29b-41d4-a716-446655440000")
+    volume_id = VolumeID(uuid.UUID("550e8400-e29b-41d4-a716-446655440000"))
     volumes = {volume_id: mock_volume}
     volumes_by_name = {"test_volume": mock_volume}
 
@@ -72,7 +73,7 @@ async def test_get_volume(mock_volume: AsyncMock, mock_storage_plugin_ctx: Async
         assert volume is mock_volume
 
     # Test get_volume with invalid volume ID
-    invalid_id = VolumeID("00000000-0000-0000-0000-000000000000")
+    invalid_id = VolumeID(uuid.UUID("00000000-0000-0000-0000-000000000000"))
     with pytest.raises(InvalidVolumeError):
         async with pool.get_volume(invalid_id) as volume:
             pass

tests/unit/storage-proxy/vfolder/test_service.py

@@ -5,7 +5,7 @@
 import pytest
 from aiohttp import web
 
-from ai.backend.common.types import QuotaConfig, QuotaScopeID, QuotaScopeType, VFolderID
+from ai.backend.common.types import QuotaConfig, QuotaScopeID, QuotaScopeType, VFolderID, VolumeID
 from ai.backend.storage.errors import VFolderNotFoundError
 from ai.backend.storage.services.service import VolumeService
 from ai.backend.storage.services.service import log as service_log
@@ -20,6 +20,9 @@
 UUID = uuid.UUID("12345678-1234-5678-1234-567812345678")
 UUID1 = uuid.UUID("12345678-1234-5678-1234-567812345679")
 UUID2 = uuid.UUID("12345678-1234-5678-1234-567812345680")
+VOLUME_ID = VolumeID(UUID)
+VOLUME_ID1 = VolumeID(UUID1)
+VOLUME_ID2 = VolumeID(UUID2)
 
 
 @pytest.fixture
@@ -46,7 +49,7 @@ async def test_get_volume(
 
     mock_volume_pool.get_volume_info.return_value = mock_volume_info
 
-    volume_id = UUID
+    volume_id = VOLUME_ID
     result = await mock_service.get_volume(volume_id)
 
     mock_log.assert_called_once_with(service_log, "get_volume", volume_id)
@@ -108,7 +111,7 @@ async def test_create_quota_scope(
 
     quota_scope_id = QuotaScopeID(scope_type=QuotaScopeType.USER, scope_id=UUID)
     quota_scope_key = QuotaScopeKey(
-        volume_id=uuid.UUID("12345678-1234-5678-1234-567812345678"),
+        volume_id=VOLUME_ID,
         quota_scope_id=quota_scope_id,
     )
     options = QuotaConfig(limit_bytes=1024 * 1024 * 1024)
@@ -132,7 +135,7 @@ async def test_get_quota_scope(
     mock_volume_pool.get_volume.return_value.__aenter__.return_value = mock_volume
 
     quota_scope_id = QuotaScopeID(scope_type=QuotaScopeType.USER, scope_id=UUID)
-    quota_scope_key = QuotaScopeKey(volume_id=UUID, quota_scope_id=quota_scope_id)
+    quota_scope_key = QuotaScopeKey(volume_id=VOLUME_ID, quota_scope_id=quota_scope_id)
 
     result = await mock_service.get_quota_scope(quota_scope_key)
 
@@ -155,7 +158,7 @@ async def test_update_quota_scope(
     mock_volume_pool.get_volume.return_value.__aenter__.return_value = mock_volume
 
     quota_scope_id = QuotaScopeID(scope_type=QuotaScopeType.USER, scope_id=UUID)
-    quota_scope_key = QuotaScopeKey(volume_id=UUID, quota_scope_id=quota_scope_id)
+    quota_scope_key = QuotaScopeKey(volume_id=VOLUME_ID, quota_scope_id=quota_scope_id)
     options = QuotaConfig(limit_bytes=2000)
 
     await mock_service.update_quota_scope(quota_scope_key, options)
@@ -180,7 +183,7 @@ async def test_delete_quota_scope(
     mock_volume_pool.get_volume.return_value.__aenter__.return_value = mock_volume
 
     quota_scope_id = QuotaScopeID(scope_type=QuotaScopeType.USER, scope_id=UUID)
-    quota_scope_key = QuotaScopeKey(volume_id=UUID, quota_scope_id=quota_scope_id)
+    quota_scope_key = QuotaScopeKey(volume_id=VOLUME_ID, quota_scope_id=quota_scope_id)
 
     await mock_service.delete_quota_scope(quota_scope_key)
 
@@ -203,7 +206,7 @@ async def test_create_vfolder(
         quota_scope_id=QuotaScopeID(scope_type=QuotaScopeType.USER, scope_id=UUID),
         folder_id=UUID,
     )
-    vfolder_key = VFolderKey(volume_id=UUID, vfolder_id=vfolder_id)
+    vfolder_key = VFolderKey(volume_id=VOLUME_ID, vfolder_id=vfolder_id)
 
     await mock_service.create_vfolder(vfolder_key)
 
@@ -228,7 +231,7 @@ async def test_clone_vfolder(
         quota_scope_id=QuotaScopeID(scope_type=QuotaScopeType.USER, scope_id=UUID1),
         folder_id=UUID2,
     )
-    vfolder_key = VFolderKey(volume_id=UUID, vfolder_id=src_vfolder_id)
+    vfolder_key = VFolderKey(volume_id=VOLUME_ID, vfolder_id=src_vfolder_id)
 
     await mock_service.clone_vfolder(vfolder_key, dst_vfolder_id)
 
@@ -260,7 +263,7 @@ async def test_get_vfolder_info(
         quota_scope_id=QuotaScopeID(scope_type=QuotaScopeType.USER, scope_id=UUID),
         folder_id=UUID,
     )
-    vfolder_key = VFolderKey(volume_id=UUID, vfolder_id=vfolder_id)
+    vfolder_key = VFolderKey(volume_id=VOLUME_ID, vfolder_id=vfolder_id)
     subpath = "test_subpath"
 
     result = await mock_service.get_vfolder_info(vfolder_key, subpath)
@@ -291,7 +294,7 @@ async def test_delete_vfolder(
         quota_scope_id=QuotaScopeID(scope_type=QuotaScopeType.USER, scope_id=UUID),
         folder_id=UUID,
     )
-    vfolder_key = VFolderKey(volume_id=UUID, vfolder_id=vfolder_id)
+    vfolder_key = VFolderKey(volume_id=VOLUME_ID, vfolder_id=vfolder_id)
 
     with pytest.raises(web.HTTPGone, match="VFolder not found"):
         await mock_service.delete_vfolder(vfolder_key)