fix(BA-2404): Generalize gpu_allocated field with improved SlotName and ResourceSlot typing

docs/manager/rest-reference/openapi.json

@@ -884,20 +884,11 @@
             "type": "string"
           },
           "resources": {
-            "additionalProperties": {
-              "anyOf": [
-                {
-                  "type": "string"
-                },
-                {
-                  "type": "integer"
-                }
-              ]
-            },
+            "additionalProperties": true,
             "examples": [
               {
                 "cpu": 4,
-                "cuda.shares": 2.5,
+                "cuda.shares": "2.5",
                 "mem": "32g"
               }
             ],

src/ai/backend/accelerator/cuda_open/plugin.py

@@ -428,7 +428,7 @@ async def generate_mounts(
 
     def get_metadata(self) -> AcceleratorMetadata:
         return {
-            "slot_name": self.slot_types[0][0],
+            "slot_name": str(self.slot_types[0][0]),
             "human_readable_name": "GPU",
             "description": "CUDA-capable GPU",
             "display_unit": "GPU",

src/ai/backend/agent/agent.py

@@ -947,12 +947,11 @@ async def __ainit__(self) -> None:
         self.timer_tasks.append(aiotools.create_timer(self.update_slots_periodically, 30.0))
 
         # Use ValkeyStatClient batch operations for better performance
-        field_value_map = {}
+        encoded_metadata_map: dict[str, bytes] = {}
         for metadata in metadatas:
-            field_value_map[metadata["slot_name"]] = dump_json_str(metadata).encode()
-
-        if field_value_map:
-            await self.valkey_stat_client.store_computer_metadata(field_value_map)
+            encoded_metadata_map[str(metadata["slot_name"])] = dump_json_str(metadata).encode()
+        if encoded_metadata_map:
+            await self.valkey_stat_client.store_computer_metadata(encoded_metadata_map)
 
         self.affinity_map = AffinityMap.build(all_devices)
 

src/ai/backend/agent/alloc_map.py

@@ -114,8 +114,8 @@ def check_exclusive(self, a: SlotName, b: SlotName) -> bool:
             return True
         for t in self.exclusive_slot_types:
             if "*" in t:
-                a_in_exclusive_set = a_in_exclusive_set or fnmatch.fnmatchcase(a, t)
-                b_in_exclusive_set = b_in_exclusive_set or fnmatch.fnmatchcase(b, t)
+                a_in_exclusive_set = a_in_exclusive_set or fnmatch.fnmatchcase(str(a), str(t))
+                b_in_exclusive_set = b_in_exclusive_set or fnmatch.fnmatchcase(str(b), str(t))
         return a_in_exclusive_set and b_in_exclusive_set
 
     def format_current_allocations(self) -> str:
@@ -713,7 +713,7 @@ def distribute_evenly(
         def allocate_across_devices(
             dev_allocs: list[tuple[DeviceId, Decimal]],
             remaining_alloc: Decimal,
-            slot_name: str,
+            slot_name: SlotName,
         ) -> dict[DeviceId, Decimal]:
             slot_allocation: dict[DeviceId, Decimal] = {}
             n_devices = len(dev_allocs)

src/ai/backend/agent/docker/agent.py

@@ -336,7 +336,7 @@ async def prepare_resource_spec(self) -> Tuple[KernelResourceSpec, Optional[Mapp
             slots = slots.normalize_slots(ignore_unknown=True)
             resource_spec = KernelResourceSpec(
                 allocations={},
-                slots={**slots},  # copy
+                slots=ResourceSlot(slots),  # copy
                 mounts=[],
                 scratch_disk_size=0,  # TODO: implement (#70)
             )

src/ai/backend/agent/dummy/agent.py

@@ -111,7 +111,7 @@ async def prepare_resource_spec(
         slots = slots.normalize_slots(ignore_unknown=True)
         resource_spec = KernelResourceSpec(
             allocations={},
-            slots={**slots},  # copy
+            slots=ResourceSlot(slots),  # copy
             mounts=[],
             scratch_disk_size=0,  # TODO: implement (#70)
         )

src/ai/backend/agent/kubernetes/agent.py

@@ -193,7 +193,7 @@ def _kernel_resource_spec_read():
             slots = slots.normalize_slots(ignore_unknown=True)
             resource_spec = KernelResourceSpec(
                 allocations={},
-                slots={**slots},  # copy
+                slots=ResourceSlot(slots),  # copy
                 mounts=[],
                 scratch_disk_size=0,  # TODO: implement (#70)
             )

src/ai/backend/agent/resources.py

@@ -88,7 +88,7 @@ class KernelResourceSpec:
     while kernel containers are running.
     """
 
-    slots: Mapping[SlotName, str]
+    slots: ResourceSlot
     """Stores the original user-requested resource slots."""
 
     allocations: MutableMapping[DeviceName, Mapping[SlotName, Mapping[DeviceId, Decimal]]]
@@ -116,7 +116,7 @@ def freeze(self) -> None:
 
     def write_to_string(self) -> str:
         mounts_str = ",".join(map(str, self.mounts))
-        slots_str = dump_json_str({k: str(v) for k, v in self.slots.items()})
+        slots_str = dump_json_str(self.slots.to_json())
 
         resource_str = ""
         resource_str += f"SCRATCH_SIZE={BinarySize(self.scratch_disk_size):m}\n"

src/ai/backend/agent/stage/kernel_lifecycle/docker/resource.py

@@ -119,7 +119,7 @@ def _prepare_resource_spec(
         slots = slots.normalize_slots(ignore_unknown=True)
         resource_spec = KernelResourceSpec(
             allocations={},
-            slots={**slots},  # copy
+            slots=ResourceSlot(slots),  # copy
             mounts=[],
             scratch_disk_size=0,  # TODO: implement (#70)
         )

src/ai/backend/common/msgpack.py

@@ -16,7 +16,7 @@
 import temporenc
 
 from .typed_validators import AutoDirectoryPath
-from .types import BinarySize, ResourceSlot
+from .types import BinarySize, ResourceSlot, SlotName
 
 __all__ = ("packb", "unpackb")
 
@@ -31,6 +31,7 @@ class ExtTypes(enum.IntEnum):
     ENUM = 6
     IMAGE_REF = 7
     RESOURCE_SLOT = 8
+    SLOT_NAME = 9
     BACKENDAI_BINARY_SIZE = 16
     AUTO_DIRECTORY_PATH = 17
 
@@ -55,6 +56,8 @@ def _default(obj: object) -> _msgpack.ExtType:
             return _msgpack.ExtType(ExtTypes.PURE_POSIX_PATH, os.fsencode(obj))
         case AutoDirectoryPath():
             return _msgpack.ExtType(ExtTypes.AUTO_DIRECTORY_PATH, os.fsencode(obj))
+        case SlotName():
+            return _msgpack.ExtType(ExtTypes.SLOT_NAME, pickle.dumps(obj, protocol=5))
         case ResourceSlot():
             return _msgpack.ExtType(ExtTypes.RESOURCE_SLOT, pickle.dumps(obj, protocol=5))
         case enum.Enum():
@@ -65,8 +68,7 @@ def _default(obj: object) -> _msgpack.ExtType:
 
 
 class ExtFunc(Protocol):
-    def __call__(self, data: bytes) -> Any:
-        pass
+    def __call__(self, data: bytes) -> Any: ...
 
 
 _DEFAULT_EXT_HOOK: Mapping[ExtTypes, ExtFunc] = {
@@ -78,6 +80,7 @@ def __call__(self, data: bytes) -> Any:
     ExtTypes.AUTO_DIRECTORY_PATH: lambda data: AutoDirectoryPath(os.fsdecode(data)),
     ExtTypes.ENUM: lambda data: pickle.loads(data),
     ExtTypes.RESOURCE_SLOT: lambda data: pickle.loads(data),
+    ExtTypes.SLOT_NAME: lambda data: pickle.loads(data),
     ExtTypes.BACKENDAI_BINARY_SIZE: lambda data: pickle.loads(data),
     ExtTypes.IMAGE_REF: lambda data: pickle.loads(data),
 }