Skip to content

feat(BA-7353): add the kernel creation payload type and event conversion errors - #13754

Open
fregataa wants to merge 5 commits into
mainfrom
feat/BA-7353-kernel-payload-types
Open

feat(BA-7353): add the kernel creation payload type and event conversion errors#13754
fregataa wants to merge 5 commits into
mainfrom
feat/BA-7353-kernel-payload-types

Conversation

@fregataa

@fregataa fregataa commented Aug 13, 2026

Copy link
Copy Markdown
Member

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.

  • Add KernelCreationInfo to events/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 what update_kernel_status_running() reads — container_id, kernel_host, the two repl ports, service_ports and occupied_devices.
  • Describe a held device once. attached_devices and 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 — a cuda device is, with cuda.device and cuda.shares — appeared twice with its attributes duplicated. OccupiedDevices.units is keyed by DeviceName and then DeviceId, and each OccupiedDevice carries every slot it supplies.
  • Carry the allocation in both units it is expressed in. allocated is in the scheduler's — keyed by ResourceSlotName, and what slot_totals sums into the kernel's occupancy. processing_units and memory_size are the same allocation in the device's own, mirroring AbstractComputeDevice. They cannot share one mapping: smp is not a slot, and a GPU's mem is not the host memory slot of the same name, so folding them together would put both into the resource accounting.
  • 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.

Every field is JSON-representable and round-trips unchanged: amounts stay Decimal, mappings are declared Mapping so 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. 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.

Notes for the reviewer

  • Where the device fields come from. The plugins fill DeviceModelInfo.data freely — its declared ComputedDeviceCapacity has never matched what anyone writes. Both accelerator plugins in reach (the enterprise CUDA one and the mock) write exactly smp and mem, which are AbstractComputeDevice.processing_units / .memory_size 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. Hence two named fields rather than an open mapping, and cores dropped.
  • model_name is kept even though it is not part of the creation contract: the GPU usage stats aggregate the device models a kernel ran on.
  • slot_totals is a plain property, not a computed_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.
  • The middle commits move the payload into a 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_devices as {"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 passed
  • pants test tests/unit/common/events::
  • pants check :: — the new types are unused by production code in this PR, so type checking is the main guard

Resolves BA-7353

@fregataa
fregataa requested a review from a team as a code owner August 13, 2026 02:05
Copilot AI balanced review requested due to automatic review settings August 13, 2026 02:05

Copilot AI left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

@github-actions github-actions Bot added size:XL 500~ LoC comp:agent Related to Agent component comp:common Related to Common component labels Aug 13, 2026
fregataa added a commit that referenced this pull request Aug 13, 2026
@fregataa
fregataa marked this pull request as draft August 13, 2026 02:11
fregataa added a commit that referenced this pull request Aug 13, 2026
@fregataa
fregataa force-pushed the feat/BA-7353-kernel-payload-types branch from 7c4f621 to e36958c Compare August 13, 2026 02:18
@fregataa
fregataa marked this pull request as ready for review August 13, 2026 02:27
Comment on lines +89 to +104
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

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

If possible, could you change it from ResourceSlot to ResourceSlotEntry wherever it appears?

Comment on lines +82 to +84
allocations: dict[DeviceName, dict[ResourceSlotName, dict[DeviceId, AllocationAmount]]] = Field(
default_factory=dict
)

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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?

fregataa added a commit that referenced this pull request Aug 13, 2026
@fregataa
fregataa force-pushed the feat/BA-7353-kernel-payload-types branch from e36958c to e3212ea Compare August 13, 2026 03:22
@github-actions github-actions Bot added size:L 100~500 LoC and removed size:XL 500~ LoC labels Aug 13, 2026
@fregataa
fregataa force-pushed the feat/BA-7353-kernel-payload-types branch from e3212ea to 314912d Compare August 13, 2026 03:28
@fregataa
fregataa requested a review from a team August 13, 2026 03:31

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We don't provide information about sub-packages, though.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

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.

Comment thread src/ai/backend/agent/resources.py Outdated
"""Stores the original user-requested resource slots."""

allocations: MutableMapping[DeviceName, Mapping[SlotName, Mapping[DeviceId, Decimal]]]
allocations: MutableMapping[DeviceName, DeviceAllocation]

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

On the event side, it looks like "ResourceSlotName" should be the key.

fregataa added a commit that referenced this pull request Aug 13, 2026
@fregataa
fregataa force-pushed the feat/BA-7353-kernel-payload-types branch from 314912d to 6ffb749 Compare August 13, 2026 10:25
@fregataa
fregataa requested a review from a team August 13, 2026 10:32
@fregataa
fregataa force-pushed the feat/BA-7353-kernel-payload-types branch from 6ffb749 to 56dc8ed Compare August 13, 2026 10:43
@fregataa
fregataa requested a review from HyeockJinKim August 13, 2026 10:56
@fregataa
fregataa force-pushed the feat/BA-7353-kernel-payload-types branch from 56dc8ed to f9dc4da Compare August 13, 2026 11:12
fregataa and others added 2 commits August 13, 2026 21:42
…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>
fregataa and others added 2 commits August 13, 2026 21:42
…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>
@fregataa
fregataa force-pushed the feat/BA-7353-kernel-payload-types branch from f9dc4da to 5064142 Compare August 13, 2026 12:44
@fregataa fregataa changed the title feat(BA-7353): add JSON-representable kernel payload types and event conversion errors feat(BA-7353): add the kernel creation payload type and event conversion errors Aug 13, 2026
@fregataa
fregataa marked this pull request as draft August 13, 2026 18:52
@fregataa
fregataa marked this pull request as ready for review August 13, 2026 20:01
@fregataa
fregataa force-pushed the feat/BA-7353-kernel-payload-types branch from d2293c9 to 83024c8 Compare August 13, 2026 20:03
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>
@fregataa
fregataa force-pushed the feat/BA-7353-kernel-payload-types branch from 83024c8 to 49ecc44 Compare August 13, 2026 20:06
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

comp:agent Related to Agent component comp:common Related to Common component size:L 100~500 LoC

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants