feat(BA-7353): add the kernel creation payload type and event conversion errors - #13754
feat(BA-7353): add the kernel creation payload type and event conversion errors#13754fregataa wants to merge 5 commits into
Conversation
7c4f621 to
e36958c
Compare
| def to_resource_slot(self) -> ResourceSlot: | ||
| """ | ||
| Sum the per-device allocations into the resource slot the manager accounts by. | ||
|
|
||
| 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 = sum(allocation_by_device.values(), Decimal(0)) | ||
| slots[SlotName(slot_name)] = str(total) | ||
| return slots |
There was a problem hiding this comment.
If possible, could you change it from ResourceSlot to ResourceSlotEntry wherever it appears?
| allocations: dict[DeviceName, dict[ResourceSlotName, dict[DeviceId, AllocationAmount]]] = Field( | ||
| default_factory=dict | ||
| ) |
There was a problem hiding this comment.
If possible, it would be best to manage the Dict by separating the types.
A three-level dictionary just doesn't seem to be very readable.
There was a problem hiding this comment.
I think “schema” might be a better fit than “data,” but even that’s a bit ambiguous.
In my opinion, it would be best to separate only the values that are actually needed so that the values stored in “data” within the Agent are distinct from those passed via the DTO. Could you expand the scope of this task a bit to address this?
e36958c to
e3212ea
Compare
e3212ea to
314912d
Compare
There was a problem hiding this comment.
We don't provide information about sub-packages, though.
There was a problem hiding this comment.
In my opinion, the interchange package doesn't seem necessary; I think the best approach would be to define a separate type in the event section and implement the conversion logic there.
| """Stores the original user-requested resource slots.""" | ||
|
|
||
| allocations: MutableMapping[DeviceName, Mapping[SlotName, Mapping[DeviceId, Decimal]]] | ||
| allocations: MutableMapping[DeviceName, DeviceAllocation] |
There was a problem hiding this comment.
On the event side, it looks like "ResourceSlotName" should be the key.
314912d to
6ffb749
Compare
6ffb749 to
56dc8ed
Compare
56dc8ed to
f9dc4da
Compare
…conversion errors Groundwork for carrying event payloads as JSON (BA-7315). Nothing changes on the wire here — the types and errors land on their own so the cutover can be reviewed without them attached to it. - Add `ai.backend.common.data.kernel.types`: `KernelCreationInfo` and its parts (`ServicePortData`, `AttachedDeviceData`, `DeviceCapacityData`, `MountData`, `KernelResourceSpecData`). Every field is JSON-representable and typed by what it holds — `DeviceName` / `ResourceSlotName` / `DeviceId` keys, `ResourceGroupName`, `ResourceSlotEntry` slots, `Decimal` amounts — rather than by `str` or `Any`. - Give `KernelResourceSpecData` a `to_resource_slot()` that sums the per-device allocations, leaving a slot with no allocation out rather than recording it as zero. - Add `KernelResourceSpec.to_data()` as the single conversion into that form, and drop `to_json_serializable_dict()` / `to_json()`, whose only caller is a log line. The removed pair also mis-rendered `slots`: its loop replaced the whole dict with a string instead of setting one key. - Raise `EventPayloadEncodingError` / `EventPayloadDecodingError` from `AbstractEvent.to_message()` / `from_message()`, so a caller sees one event-layer error rather than a Pydantic one. An event deriving `BackendAISchema` maps `ValidationError` to a `BackendAIError` of its own, so both forms are funnelled in. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…d carry only what the manager reads Review follow-up. `common/data` is where a component keeps its own value objects, so a payload that crosses a component boundary had no documented home — `dto/` is organized by target component, and a payload has two sides. - Add `common/interchange/`, with `AGENTS.md` stating the membership test: one component produces it, another consumes it, so it has to survive serialization. Record the split between `interchange` / `dto` / `schema` / `data` in `common/KNOWLEDGE.md`, and list the package (along with the previously missing `schema/`) in `common/AGENTS.md`. - Carry only what `update_kernel_status_running()` reads, which is its one consumer: `container_id`, `kernel_host`, the two repl ports, `service_ports`, `attached_devices` and `allocations`. Dropped `id` (the event already names the kernel), `agent_addr` (the event names its source), the legacy stdio ports, and the parts of the resource spec — slots, mounts, unified devices, scratch size — that nothing reads. - Return `list[ResourceSlotEntry]` rather than the legacy `ResourceSlot`; a caller that needs the dict form goes through `ResourceSlotEntry.inputs_to_resource_slot()`. - Name the levels of the allocation mapping (`DeviceAllocation`, `PerDeviceAllocation`) instead of nesting three dicts inline, and apply the `DeviceAllocation` alias that `agent/resources.py` already defined but did not use. With the spec no longer travelling whole, `KernelResourceSpecData` / `MountData` / `to_data()` would only serve the agent's own log, so they are gone and `to_json_serializable_dict()` is left as it was. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
…arry it Review follow-up. A payload that only ever travels on a kernel event does not need a package of its own — it belongs with the other kernel event types. - Drop `common/interchange/` and define the payload in `events/event_types/kernel/types.py`, next to `KernelLifecycleEventReason`. Revert `common/AGENTS.md` and `common/KNOWLEDGE.md`, which only described that package, and leave `agent/resources.py` untouched. - Give each level of the occupancy its own model with a mapping field rather than nesting three dicts inline: `KernelOccupancy.devices` (by `DeviceName`) → `DeviceOccupancy.slots` (by `ResourceSlotName`) → `SlotOccupancy.amounts` (by `DeviceId`). A device supplies more than one slot when it is metered along more than one axis, as `cuda` does with `cuda.device` and `cuda.shares`. - Name it after what the manager does with it. It records the value as the kernel's occupancy, so the field is `occupancy` and the per-slot sum is `slot_totals`. - Key slots by `ResourceSlotName` throughout, not the legacy `SlotName`. - Declare the mappings as `Mapping`, so a consumer that tries to edit the payload it received is caught by the type checker. - Drop the field defaults: a producer states the whole payload. The exception is `DeviceCapacity`, whose two fields the compute plugin's own `ComputedDeviceCapacity` declares `NotRequired` — a device that measures neither reports neither. - Let the default `Decimal` constraint reject a non-finite amount. `Infinity` belongs to a resource policy with `DefaultForUnspecified.UNLIMITED`, not to an allocation: a device supplies a finite share of what it has, and `alloc_map` never produces one. `slot_totals` stays a property rather than a `computed_field`: a computed field is written into the payload beside the occupancy it is derived from, where nothing reads it — a receiver recomputes it — and it can disagree with the value next to it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
f9dc4da to
5064142
Compare
d2293c9 to
83024c8
Compare
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. - File the news fragment under `feature`, matching the commit prefix, and describe what the change actually adds. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
83024c8 to
49ecc44
Compare
Summary
Groundwork split out of #13746 so the JSON cutover can be reviewed without the type definitions attached to it. Nothing changes on the wire here — the new types have no caller yet, and the existing
serialize()/deserialize()msgpack path is untouched.KernelCreationInfotoevents/event_types/kernel/types.py, beside the events that carry it. It is the contract with the manager, not a rendering of the agent's resource spec: it carries only whatupdate_kernel_status_running()reads —container_id,kernel_host, the two repl ports,service_portsandoccupied_devices.attached_devicesand the per-slot occupancy used to describe the same devices from two sides, both keyed by device name and then by device, so a unit metered along two axes — acudadevice is, withcuda.deviceandcuda.shares— appeared twice with its attributes duplicated.OccupiedDevices.unitsis keyed byDeviceNameand thenDeviceId, and eachOccupiedDevicecarries every slot it supplies.allocatedis in the scheduler's — keyed byResourceSlotName, and whatslot_totalssums into the kernel's occupancy.processing_unitsandmemory_sizeare the same allocation in the device's own, mirroringAbstractComputeDevice. They cannot share one mapping:smpis not a slot, and a GPU'smemis not the host memory slot of the same name, so folding them together would put both into the resource accounting.EventPayloadEncodingError/EventPayloadDecodingErrorfromAbstractEvent.to_message()/from_message()so a caller sees one event-layer error rather than a Pydantic one. An event derivingBackendAISchemamapsValidationErrorto aBackendAIErrorof its own, so both forms are funnelled in.Every field is JSON-representable and round-trips unchanged: amounts stay
Decimal, mappings are declaredMappingso a consumer editing a payload it received is caught by the type checker, and there are no field defaults — a producer states the whole payload.A non-finite amount is rejected at decode.
Infinitybelongs to a resource policy withDefaultForUnspecified.UNLIMITED, not to an allocation — a device supplies a finite share of what it has, andalloc_mapnever produces one.Notes for the reviewer
DeviceModelInfo.datafreely — its declaredComputedDeviceCapacityhas never matched what anyone writes. Both accelerator plugins in reach (the enterprise CUDA one and the mock) write exactlysmpandmem, which areAbstractComputeDevice.processing_units/.memory_sizeprojected onto the kernel's share; the intrinsic cpu plugin writescores, which is the kernel's total core count copied onto every device and which nothing reads. Hence two named fields rather than an open mapping, andcoresdropped.model_nameis kept even though it is not part of the creation contract: the GPU usage stats aggregate the device models a kernel ran on.slot_totalsis a plain property, not acomputed_field: a computed field is 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.common/interchange/package and then back out again. That package was the answer to "where does an inter-component payload live"; defining it beside its events made it unnecessary. Kept as separate commits so the review discussion stays readable.Follow-up for #13746
The manager writes
kernels.attached_devicesas{"cuda": [{"model_name", "data": {"smp", "mem"}}]}, and the GPU usage stats read it back through that shape. Wiring this payload up has to reshape it, or those stats go to zero.Test plan
tests/unit/common/events/— 41 passedpants test tests/unit/common/events::pants check ::— the new types are unused by production code in this PR, so type checking is the main guardResolves BA-7353