refactor(BA-7315): carry the kernel resource spec as domain types

`KernelResourceSpecData` described itself entirely in built-ins, so nothing in the
payload said what a key or a mount actually was.

- Type the fields as what they hold: `list[ResourceSlotEntry]` for slots,
  `dict[DeviceName, dict[ResourceSlotName, dict[DeviceId, str]]]` for allocations,
  and `list[MountData]` for mounts, which keeps a mount structured instead of
  flattening it to its `str()` form. `KernelCreationInfo` follows with `KernelId`,
  `ContainerId`, `DeviceId` and `DeviceName`.
- `MountData` mirrors the agent's `Mount` and declares no defaults, so a producer
  states every field rather than inheriting one.
- Omit a slot holding no device allocation from `to_resource_slot()` instead of
  recording it as zero. This restores what the manager helper it replaced did — a
  slot absent and a slot present at zero are not the same occupancy.
- Drop `KernelResourceSpec.to_json_serializable_dict()` / `to_json()`, whose only
  caller was a log line, and `AgentRegistry.convert_resource_spec_to_resource_slot()`,
  which had no caller at all. Slot aggregation now has one implementation rather
  than three.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

Author: fregataa

src/ai/backend/agent/agent.py

@@ -2718,7 +2718,7 @@ async def create_kernel(
                     "create_kernel(kernel:{}, session:{}) resource spec prepared: {}",
                     kernel_id,
                     session_id,
-                    resource_spec.to_json(),
+                    resource_spec.to_data().model_dump_json(),
                 )
 
                 # Mount backend-specific intrinsic mounts (e.g., scratch directories)
@@ -3314,7 +3314,7 @@ async def create_kernel(
                     # Finally we are done.
                     creation_info = KernelCreationInfo(
                         id=kernel_id,
-                        container_id=str(kernel_obj["container_id"]),
+                        container_id=ContainerId(str(kernel_obj["container_id"])),
                         kernel_host=str(kernel_obj["kernel_host"]),
                         repl_in_port=kernel_obj["repl_in_port"],
                         repl_out_port=kernel_obj["repl_out_port"],
@@ -3328,7 +3328,7 @@ async def create_kernel(
                         ],
                         resource_spec=resource_spec.to_data(),
                         attached_devices={
-                            str(dev_name): [
+                            dev_name: [
                                 AttachedDeviceData.model_validate(device) for device in devices
                             ]
                             for dev_name, devices in attached_devices.items()

src/ai/backend/agent/resources.py

@@ -40,8 +40,9 @@
     InvalidResourceConfigError,
     ResourceOverAllocatedError,
 )
-from ai.backend.common.data.kernel.types import KernelResourceSpecData
+from ai.backend.common.data.kernel.types import KernelResourceSpecData, MountData
 from ai.backend.common.etcd import AsyncEtcd
+from ai.backend.common.identifier.resource_slot import ResourceSlotName
 from ai.backend.common.json import dump_json_str, load_json
 from ai.backend.common.plugin import AbstractPlugin, BasePluginContext
 from ai.backend.common.types import (
@@ -56,6 +57,7 @@
     MountPermission,
     MountTypes,
     ResourceSlot,
+    ResourceSlotEntry,
     SlotName,
     SlotTypes,
     aobject,
@@ -271,37 +273,36 @@ def _format_alloc(slot_name: SlotName, alloc: Decimal) -> str:
 
     def to_data(self) -> KernelResourceSpecData:
         """
-        Render this spec with JSON-representable leaves, as events and logs carry it.
+        Render this spec as the value events and logs carry it as.
         """
         return KernelResourceSpecData(
-            slots={
-                str(slot_name): self._format_alloc(SlotName(slot_name), alloc)
-                for slot_name, alloc in self.slots.items()
-            },
+            slots=ResourceSlotEntry.from_resource_slot(self.slots),
             allocations={
-                str(dev_name): {
-                    str(slot_name): {
-                        str(dev_id): self._format_alloc(slot_name, alloc)
+                dev_name: {
+                    ResourceSlotName(str(slot_name)): {
+                        dev_id: self._format_alloc(slot_name, alloc)
                         for dev_id, alloc in per_device_alloc.items()
                     }
                     for slot_name, per_device_alloc in dev_alloc.items()
                 }
                 for dev_name, dev_alloc in self.allocations.items()
             },
             scratch_disk_size=self.scratch_disk_size,
-            mounts=[str(mount) for mount in self.mounts],
+            mounts=[
+                MountData(
+                    type=mount.type,
+                    source=mount.source,
+                    target=mount.target,
+                    permission=mount.permission,
+                )
+                for mount in self.mounts
+            ],
             unified_devices=[
-                (str(device_name), str(slot_name))
+                (device_name, ResourceSlotName(str(slot_name)))
                 for device_name, slot_name in self.unified_devices
             ],
         )
 
-    def to_json_serializable_dict(self) -> Mapping[str, Any]:
-        return self.to_data().model_dump(mode="json")
-
-    def to_json(self) -> str:
-        return dump_json_str(self.to_json_serializable_dict())
-
     @classmethod
     def __get_pydantic_core_schema__(
         cls, source_type: Any, handler: GetCoreSchemaHandler

src/ai/backend/common/data/kernel/types.py

@@ -1,14 +1,22 @@
 from __future__ import annotations
 
-import uuid
 from decimal import Decimal
+from pathlib import Path
 
 from pydantic import Field
 
+from ai.backend.common.identifier.resource_slot import ResourceSlotName
 from ai.backend.common.types import (
     BackendAISchema,
     BinarySize,
+    ContainerId,
+    DeviceId,
+    DeviceName,
+    KernelId,
+    MountPermission,
+    MountTypes,
     ResourceSlot,
+    ResourceSlotEntry,
     ServicePortProtocols,
     SlotName,
 )
@@ -18,6 +26,7 @@
     "DeviceCapacityData",
     "KernelCreationInfo",
     "KernelResourceSpecData",
+    "MountData",
     "ServicePortData",
 )
 
@@ -28,7 +37,7 @@ class DeviceCapacityData(BackendAISchema):
 
 
 class AttachedDeviceData(BackendAISchema):
-    device_id: str
+    device_id: DeviceId
     model_name: str
     data: DeviceCapacityData = Field(default_factory=DeviceCapacityData)
 
@@ -41,30 +50,50 @@ class ServicePortData(BackendAISchema):
     is_inference: bool = False
 
 
+class MountData(BackendAISchema):
+    """
+    A mount as the agent realized it on the container.
+
+    Mirrors the agent's `Mount`, minus the container-runtime `opts` it never carries on
+    a resource spec.
+    """
+
+    type: MountTypes
+    source: Path | None
+    target: Path
+    permission: MountPermission
+
+
 class KernelResourceSpecData(BackendAISchema):
     """
-    The agent's resource spec for one kernel, with every leaf JSON-representable.
+    The agent's resource spec for one kernel.
 
-    Allocations and slot amounts are decimal strings rather than numbers: they are
-    read back as `Decimal`, and a float would lose the exact value on the way.
+    Per-device allocation amounts are decimal strings rather than numbers: they are read
+    back as `Decimal`, and a float would lose the exact value on the way.
     """
 
-    slots: dict[str, str] = Field(default_factory=dict)
-    allocations: dict[str, dict[str, dict[str, str]]] = Field(default_factory=dict)
+    slots: list[ResourceSlotEntry] = Field(default_factory=list)
+    allocations: dict[DeviceName, dict[ResourceSlotName, dict[DeviceId, str]]] = Field(
+        default_factory=dict
+    )
     scratch_disk_size: int = 0
-    mounts: list[str] = Field(default_factory=list)
-    unified_devices: list[tuple[str, str]] = Field(default_factory=list)
+    mounts: list[MountData] = Field(default_factory=list)
+    unified_devices: list[tuple[DeviceName, ResourceSlotName]] = Field(default_factory=list)
 
     def to_resource_slot(self) -> ResourceSlot:
         """
         Sum the per-device allocations into the resource slot the manager accounts by.
 
         A byte-valued amount arrives suffixed (`"4g"`), so it is parsed rather than
-        read as a plain decimal.
+        read as a plain decimal. A slot holding no device allocation is left out rather
+        than recorded as zero, which is the distinction a caller reading the result back
+        as occupancy depends on.
         """
         slots = ResourceSlot()
         for alloc_map in self.allocations.values():
             for slot_name, allocation_by_device in alloc_map.items():
+                if not allocation_by_device:
+                    continue
                 total = Decimal(0)
                 for amount in allocation_by_device.values():
                     if amount and BinarySize.suffix_map.get(amount[-1].lower()) is not None:
@@ -80,8 +109,8 @@ class KernelCreationInfo(BackendAISchema):
     What the agent reports about a kernel once its container is up.
     """
 
-    id: uuid.UUID
-    container_id: str
+    id: KernelId
+    container_id: ContainerId
     kernel_host: str
     repl_in_port: int
     repl_out_port: int
@@ -91,4 +120,4 @@ class KernelCreationInfo(BackendAISchema):
     agent_addr: str
     service_ports: list[ServicePortData] = Field(default_factory=list)
     resource_spec: KernelResourceSpecData = Field(default_factory=KernelResourceSpecData)
-    attached_devices: dict[str, list[AttachedDeviceData]] = Field(default_factory=dict)
+    attached_devices: dict[DeviceName, list[AttachedDeviceData]] = Field(default_factory=dict)

src/ai/backend/manager/registry.py

@@ -12,7 +12,6 @@
     Sequence,
 )
 from datetime import datetime, timedelta
-from decimal import Decimal
 from typing import (
     Any,
     cast,
@@ -80,11 +79,9 @@
     AbuseReport,
     AccessKey,
     AgentId,
-    BinarySize,
     ClusterMode,
     ClusterSSHKeyPair,
     CommitStatus,
-    DeviceId,
     HardwareMetadata,
     ImageAlias,
     ImageRegistry,
@@ -97,7 +94,6 @@
     SessionEnqueueingConfig,
     SessionId,
     SessionTypes,
-    SlotName,
 )
 from ai.backend.common.utils import str_to_timedelta
 from ai.backend.logging import BraceStyleAdapter
@@ -1251,29 +1247,6 @@ async def enqueue_session(
             startup_command=startup_command,
         )
 
-    def convert_resource_spec_to_resource_slot(
-        self,
-        allocations: Mapping[str, Mapping[SlotName, Mapping[DeviceId, str]]],
-    ) -> ResourceSlot:
-        """
-        Convert per-device resource spec allocations (agent-side format)
-        back into a resource slot (manager-side format).
-        """
-        slots = ResourceSlot()
-        for alloc_map in allocations.values():
-            for slot_name, allocation_by_device in alloc_map.items():
-                total_allocs: list[Decimal] = []
-                for allocation in allocation_by_device.values():
-                    if (
-                        isinstance(allocation, (BinarySize, str))
-                        and BinarySize.suffix_map.get(allocation[-1].lower()) is not None
-                    ):
-                        total_allocs.append(Decimal(BinarySize.from_str(allocation)))
-                    else:  # maybe Decimal("Infinity"), etc.
-                        total_allocs.append(Decimal(allocation))
-                slots[slot_name] = str(sum(total_allocs))
-        return slots
-
     async def create_cluster_ssh_keypair(self) -> ClusterSSHKeyPair:
         key = rsa.generate_private_key(
             backend=default_backend(),