refactor(BA-7353): describe a held device once, in both units

Review follow-up. `attached_devices` and the occupancy described the same devices from
two sides — what is attached, and how much of it is held — and both were keyed by device
name and then by device. A unit metered along two axes, as a `cuda` device is with
`cuda.device` and `cuda.shares`, appeared twice with its attributes duplicated.

- Merge them into `occupied_devices`, keyed by `DeviceName` and then `DeviceId`, so each
  unit is described once with every slot it supplies under it.
- Carry the device's own measure of that allocation as `processing_units` and
  `memory_size`, mirroring `AbstractComputeDevice`, rather than the free-form `data`
  mapping the plugins fill. The two accelerator plugins in reach — the enterprise CUDA
  one and the mock — write exactly `smp` and `mem` there, which are those two fields
  projected onto the kernel's share; the intrinsic cpu plugin writes `cores`, which is
  the kernel's total core count copied onto every device and which nothing reads.
- Keep the allocation itself keyed by `ResourceSlotName`. It is what `slot_totals` sums
  into the kernel's occupancy, so a device-measured quantity cannot live in it: `smp` is
  not a slot, and a GPU's `mem` is not the host memory slot of the same name.

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

Author: fregataa

src/ai/backend/common/events/event_types/kernel/types.py

@@ -57,70 +57,53 @@ def from_value(cls, value: str | None) -> Self | None:
         return None
 
 
-class SlotOccupancy(BackendAISchema):
-    """How much of one slot each individual device supplies."""
-
-    amounts: Mapping[DeviceId, Decimal]
-
-
-class DeviceOccupancy(BackendAISchema):
-    """The slots one device supplies.
+class OccupiedDevice(BackendAISchema):
+    """
+    One device unit the kernel holds, and how much of it.
 
-    A device supplies more than one slot when it is metered along more than one axis —
-    `cuda` supplies both `cuda.device` and `cuda.shares`.
+    `allocated` is in the scheduler's units — the slots it accounts by, of which a unit
+    supplies more than one when it is metered along more than one axis, as a `cuda`
+    device does with `cuda.device` and `cuda.shares`. `processing_units` and
+    `memory_size` are that same allocation in the device's own units, mirroring
+    `AbstractComputeDevice`; only an accelerator reports them.
     """
 
-    slots: Mapping[ResourceSlotName, SlotOccupancy]
+    model_name: (
+        str | None
+    )  # Kept for the GPU usage stats, which aggregate the device models a kernel ran on.
+    allocated: Mapping[ResourceSlotName, Decimal]
+    processing_units: int | None
+    memory_size: int | None
 
 
-class KernelOccupancy(BackendAISchema):
+class OccupiedDevices(BackendAISchema):
     """
-    The resources the kernel occupies, attributed to the devices supplying them.
+    The devices the kernel occupies.
 
-    `DeviceName` names a device (`cuda`) and `DeviceId` one of its units (`0`), so the
-    two levels keyed by a device are not the same thing.
+    Keyed by device name (`cuda`) and then by unit (`0`): `DeviceName` names a kind of
+    device, `DeviceId` one of its units.
     """
 
-    devices: Mapping[DeviceName, DeviceOccupancy]
+    units: Mapping[DeviceName, Mapping[DeviceId, OccupiedDevice]]
 
     @property
     def slot_totals(self) -> list[ResourceSlotEntry]:
         """
-        The per-device amounts summed per slot — what a caller records as occupancy.
+        The per-unit allocations summed per slot — what a caller records as occupancy.
 
-        A slot supplied by no device is left out rather than reported as zero, which is
-        the distinction a caller storing the result depends on.
-
-        Not a `computed_field`: it would be written into the payload beside the
-        occupancy it is derived from, and a receiver recomputes it anyway.
+        Not a `computed_field`: it would be written into the payload beside the allocations
+        it is derived from, where nothing reads it — a receiver recomputes it — and it
+        can disagree with the value next to it.
         """
-        totals: list[ResourceSlotEntry] = []
-        for device in self.devices.values():
-            for slot_name, slot in device.slots.items():
-                if not slot.amounts:
-                    continue
-                total = sum(slot.amounts.values(), Decimal(0))
-                totals.append(ResourceSlotEntry(resource_type=slot_name, quantity=str(total)))
-        return totals
-
-
-class DeviceCapacity(BackendAISchema):
-    """
-    What a device reports about itself.
-
-    Both are defaulted, unlike every other field here: the compute plugin's own
-    `ComputedDeviceCapacity` declares them `NotRequired`, so a device that measures
-    neither reports neither.
-    """
-
-    mem: int | None = None
-    proc: int | None = None
-
-
-class AttachedDevice(BackendAISchema):
-    device_id: DeviceId
-    model_name: str
-    data: DeviceCapacity
+        totals: dict[ResourceSlotName, Decimal] = {}
+        for units in self.units.values():
+            for device in units.values():
+                for slot_name, amount in device.allocated.items():
+                    totals[slot_name] = totals.get(slot_name, Decimal(0)) + amount
+        return [
+            ResourceSlotEntry(resource_type=slot_name, quantity=str(total))
+            for slot_name, total in totals.items()
+        ]
 
 
 class ServicePortInfo(BackendAISchema):
@@ -144,5 +127,4 @@ class KernelCreationInfo(BackendAISchema):
     repl_in_port: int
     repl_out_port: int
     service_ports: list[ServicePortInfo]
-    attached_devices: Mapping[DeviceName, list[AttachedDevice]]
-    occupancy: KernelOccupancy
+    occupied_devices: OccupiedDevices

tests/unit/common/events/test_kernel_types.py

@@ -1,18 +1,15 @@
 from __future__ import annotations
 
-from collections import defaultdict
+import json
 from decimal import Decimal
 
 import pytest
 
 from ai.backend.common.events.event_types.kernel.types import (
-    AttachedDevice,
-    DeviceCapacity,
-    DeviceOccupancy,
     KernelCreationInfo,
-    KernelOccupancy,
+    OccupiedDevice,
+    OccupiedDevices,
     ServicePortInfo,
-    SlotOccupancy,
 )
 from ai.backend.common.exception import BackendAISchemaValidationFailed
 from ai.backend.common.identifier.resource_slot import ResourceSlotName
@@ -25,35 +22,53 @@
 )
 
 
-def _info_with(**overrides: object) -> KernelCreationInfo:
-    fields: dict[str, object] = {
-        "container_id": ContainerId("c0ffee"),
-        "kernel_host": "127.0.0.1",
-        "repl_in_port": 2000,
-        "repl_out_port": 2001,
-        "service_ports": [],
-        "attached_devices": {},
-        "occupancy": KernelOccupancy(devices={}),
-        **overrides,
-    }
-    return KernelCreationInfo(**fields)  # type: ignore[arg-type]
-
-
-def _occupancy(**per_slot: dict[str, Decimal]) -> KernelOccupancy:
-    """Build an occupancy, taking the device name from the slot name as the agent does."""
-    by_device: defaultdict[DeviceName, dict[ResourceSlotName, SlotOccupancy]] = defaultdict(dict)
-    for slot, amounts in per_slot.items():
-        by_device[DeviceName(slot.partition(".")[0])][ResourceSlotName(slot)] = SlotOccupancy(
-            amounts={DeviceId(device_id): amount for device_id, amount in amounts.items()}
-        )
-    return KernelOccupancy(
-        devices={name: DeviceOccupancy(slots=slots) for name, slots in by_device.items()}
+def _device(
+    allocated: dict[str, Decimal],
+    *,
+    model_name: str | None = None,
+    processing_units: int | None = None,
+    memory_size: int | None = None,
+) -> OccupiedDevice:
+    return OccupiedDevice(
+        model_name=model_name,
+        allocated={ResourceSlotName(slot): amount for slot, amount in allocated.items()},
+        processing_units=processing_units,
+        memory_size=memory_size,
+    )
+
+
+def _info_with(
+    occupied_devices: OccupiedDevices,
+    service_ports: list[ServicePortInfo] | None = None,
+) -> KernelCreationInfo:
+    return KernelCreationInfo(
+        container_id=ContainerId("c0ffee"),
+        kernel_host="127.0.0.1",
+        repl_in_port=2000,
+        repl_out_port=2001,
+        service_ports=service_ports if service_ports is not None else [],
+        occupied_devices=occupied_devices,
     )
 
 
 @pytest.fixture
 def creation_info() -> KernelCreationInfo:
+    """A kernel holding one cpu core, 1 GiB, and half of one GPU."""
     return _info_with(
+        OccupiedDevices(
+            units={
+                DeviceName("cpu"): {DeviceId("0"): _device({"cpu": Decimal("1")})},
+                DeviceName("mem"): {DeviceId("root"): _device({"mem": Decimal("1073741824")})},
+                DeviceName("cuda"): {
+                    DeviceId("0"): _device(
+                        {"cuda.device": Decimal("1"), "cuda.shares": Decimal("0.5")},
+                        model_name="A100",
+                        processing_units=54,
+                        memory_size=21474836480,
+                    )
+                },
+            }
+        ),
         service_ports=[
             ServicePortInfo(
                 name="jupyter",
@@ -63,16 +78,6 @@ def creation_info() -> KernelCreationInfo:
                 is_inference=False,
             )
         ],
-        attached_devices={
-            DeviceName("cuda"): [
-                AttachedDevice(
-                    device_id=DeviceId("0"),
-                    model_name="A100",
-                    data=DeviceCapacity(mem=1024, proc=8),
-                )
-            ]
-        },
-        occupancy=_occupancy(cpu={"0": Decimal("2")}, mem={"root": Decimal("4294967296")}),
     )
 
 
@@ -84,44 +89,60 @@ def test_roundtrip_keeps_every_typed_leaf(self, creation_info: KernelCreationInf
 
         assert restored == creation_info
         assert restored.service_ports[0].protocol is ServicePortProtocols.HTTP
-        assert restored.attached_devices[DeviceName("cuda")][0].data.mem == 1024
-        assert restored.occupancy.devices[DeviceName("mem")].slots[
-            ResourceSlotName("mem")
-        ].amounts == {DeviceId("root"): Decimal("4294967296")}
+        cuda = restored.occupied_devices.units[DeviceName("cuda")][DeviceId("0")]
+        assert cuda.model_name == "A100"
+        assert (cuda.processing_units, cuda.memory_size) == (54, 21474836480)
+        assert cuda.allocated[ResourceSlotName("cuda.shares")] == Decimal("0.5")
+
+    def test_one_unit_is_described_once(self, creation_info: KernelCreationInfo) -> None:
+        """The unit metered along two axes appears once, with both amounts under it —
+        which is what merging the attached devices into the occupancy buys."""
+        cuda = creation_info.occupied_devices.units[DeviceName("cuda")]
+
+        assert list(cuda) == [DeviceId("0")]
+        assert set(cuda[DeviceId("0")].allocated) == {"cuda.device", "cuda.shares"}
+
+    def test_an_intrinsic_device_reports_neither_unit(
+        self, creation_info: KernelCreationInfo
+    ) -> None:
+        """Only an accelerator measures itself; the cpu and mem plugins report nothing."""
+        cpu = creation_info.occupied_devices.units[DeviceName("cpu")][DeviceId("0")]
+
+        assert (cpu.model_name, cpu.processing_units, cpu.memory_size) == (None, None, None)
 
     def test_derived_totals_are_not_written_to_the_wire(
         self, creation_info: KernelCreationInfo
     ) -> None:
-        """The occupancy is the payload; the per-slot sum is derived and stays off it."""
+        """The allocations are the payload; the per-slot sum is derived and stays off it."""
         assert "slot_totals" not in creation_info.model_dump_json()
 
 
 class TestSlotTotals:
     """`slot_totals` is what a caller records as the kernel's occupancy."""
 
-    def test_one_device_supplying_several_slots(self) -> None:
-        """`cuda` is metered along two axes at once, by the same two units."""
-        occupancy = KernelOccupancy(
-            devices={
-                DeviceName("cuda"): DeviceOccupancy(
-                    slots={
-                        ResourceSlotName("cuda.shares"): SlotOccupancy(
-                            amounts={
-                                DeviceId("0"): Decimal("0.5"),
-                                DeviceId("1"): Decimal("0.25"),
-                            }
-                        ),
-                        ResourceSlotName("cuda.device"): SlotOccupancy(
-                            amounts={DeviceId("0"): Decimal("1"), DeviceId("1"): Decimal("1")}
-                        ),
-                    }
-                )
+    def test_amounts_are_summed_across_units(self) -> None:
+        occupied = OccupiedDevices(
+            units={
+                DeviceName("cuda"): {
+                    DeviceId("0"): _device({"cuda.shares": Decimal("0.5")}),
+                    DeviceId("1"): _device({"cuda.shares": Decimal("0.25")}),
+                }
             }
         )
 
-        totals = {e.resource_type: e.quantity for e in occupancy.slot_totals}
+        totals = {e.resource_type: e.quantity for e in occupied.slot_totals}
+
+        assert totals == {"cuda.shares": "0.75"}
 
-        assert totals == {"cuda.shares": "0.75", "cuda.device": "2"}
+    def test_a_unit_reports_each_of_its_slots(self, creation_info: KernelCreationInfo) -> None:
+        totals = {e.resource_type: e.quantity for e in creation_info.occupied_devices.slot_totals}
+
+        assert totals == {
+            "cpu": "1",
+            "mem": "1073741824",
+            "cuda.device": "1",
+            "cuda.shares": "0.5",
+        }
 
     @pytest.mark.parametrize(
         ("slot", "amount"),
@@ -133,34 +154,35 @@ def test_one_device_supplying_several_slots(self) -> None:
         ids=["exact_bytes", "off_by_one_byte", "fractional"],
     )
     def test_amounts_survive_the_wire_exactly(self, slot: str, amount: Decimal) -> None:
-        info = _info_with(occupancy=_occupancy(**{slot: {"0": amount}}))
+        info = _info_with(
+            OccupiedDevices(units={DeviceName(slot): {DeviceId("0"): _device({slot: amount})}})
+        )
 
         restored = KernelCreationInfo.model_validate_json(info.model_dump_json())
 
-        assert restored.occupancy.slot_totals == [
+        assert restored.occupied_devices.slot_totals == [
             ResourceSlotEntry(resource_type=ResourceSlotName(slot), quantity=str(amount))
         ]
 
     @pytest.mark.parametrize("amount", ["Infinity", "-Infinity", "NaN"], ids=str)
     def test_non_finite_amount_is_rejected(self, amount: str) -> None:
-        """A device supplies a finite share of what it has; an unbounded amount is a
+        """A unit supplies a finite share of what it has; an unbounded amount is a
         limit, which is not what this carries."""
-        payload = (
-            '{"devices":{"cpu":{"slots":{"cpu":{"amounts":{"0":"%s"}}}}}}' % amount  # noqa: UP031
-        )
+        payload = json.dumps({
+            "units": {
+                "cpu": {
+                    "0": {
+                        "model_name": None,
+                        "allocated": {"cpu": amount},
+                        "processing_units": None,
+                        "memory_size": None,
+                    }
+                }
+            }
+        })
 
         with pytest.raises(BackendAISchemaValidationFailed):
-            KernelOccupancy.model_validate_json(payload)
-
-    def test_slot_supplied_by_no_device_is_omitted(self) -> None:
-        """Omitted, not zero — the caller stores the result as occupancy, where a slot
-        present at zero and a slot absent are not the same statement."""
-        occupancy = KernelOccupancy(
-            devices={
-                DeviceName("cuda"): DeviceOccupancy(
-                    slots={ResourceSlotName("cuda.device"): SlotOccupancy(amounts={})}
-                )
-            }
-        )
+            OccupiedDevices.model_validate_json(payload)
 
-        assert occupancy.slot_totals == []
+    def test_a_kernel_holding_nothing_totals_nothing(self) -> None:
+        assert OccupiedDevices(units={}).slot_totals == []