From a39c606b9426efee9f0ff6d045d3e273677431f6 Mon Sep 17 00:00:00 2001 From: gspark-etri Date: Mon, 19 Oct 2026 18:28:59 +0900 Subject: [PATCH] rm-core 0.1.2: airm-placement document normative (PROTOCOL 5.4) The airm-placement ConfigMap was defined by usage: rm-controller writes it, hetero-llm-d documents and ships its empty default, lmd-top reads it. 5.4 states the format they already agree on (no field changes, protocol stays 0.1): location, fields, roles via workload.kind, card refs by accelerator id, the vendor.model label, the generation rule and forward compatibility. - rm_core.placement: PlacementDocument/Decision model, stdlib validate(doc) -> list[str], parse, build (generation bumps only on change), configmap, is_newer (reader rule) - k8s/airm-placement.schema.json, tested against the code's enumerations and examples - airm.placements conformance also validates the document made from a policy's decisions Co-Authored-By: Claude Opus 5.5 --- PROTOCOL.md | 81 +++++++++- README.ko.md | 6 +- README.md | 6 +- k8s/airm-placement.schema.json | 64 ++++++++ pyproject.toml | 2 +- src/rm_core/__init__.py | 2 +- src/rm_core/conformance.py | 13 +- src/rm_core/placement.py | 280 +++++++++++++++++++++++++++++++++ tests/test_placement.py | 218 +++++++++++++++++++++++++ 9 files changed, 665 insertions(+), 7 deletions(-) create mode 100644 k8s/airm-placement.schema.json create mode 100644 src/rm_core/placement.py create mode 100644 tests/test_placement.py diff --git a/PROTOCOL.md b/PROTOCOL.md index 081fca4..ce6c1f5 100644 --- a/PROTOCOL.md +++ b/PROTOCOL.md @@ -99,7 +99,8 @@ limited to `rm_core.workload.SLO_KEYS`. - SHOULD name the KV transport to use between prefill and decode in `kv_transport`. The placement controller runs this policy over `AcceleratorProfile`s and live metrics, and writes -its decisions where the llm-d EPP scorer and serving pods read them (see ai-rm docs). +its decisions to the placement document (§5.4), where the llm-d EPP scorer, serving pods and +lmd-top read them. ### 5.3 Workload runners @@ -108,6 +109,75 @@ callable that accepts a `Workload` plus keyword configuration and starts the wor configuration keys MUST be documented in the component's README. The runner MUST select its KV transport by name (`kv_transport=`), never by hard-coded import. (Formalized in v0.2.) +### 5.4 Placement document (`airm-placement`, document version 0.1) + +The hand-off between the placement controller (writer) and its readers is one JSON document. +`rm_core.placement` implements it (`validate(doc)`, `PlacementDocument`, `build`, `parse`, +`configmap`); `k8s/airm-placement.schema.json` is the JSON Schema for the structure. The document +version is the `protocol` field and follows the protocol version (§1). + +**Location.** ConfigMap `airm-placement` in the namespace of the llm-d InferencePools it describes +(the llm-d namespace, e.g. `llm-serving`), key `decisions.json`, value a UTF-8 JSON object. Writers +MUST NOT put decisions under other keys. The ConfigMap SHOULD carry the annotation +`airm.ai-computing.io/protocol: ""` and the label `app.kubernetes.io/part-of: ai-rm`. + +```json +{ + "protocol": "0.1", + "generation": 3, + "updatedAt": "2026-10-19T00:00:00Z", + "placement": "topology", + "decisions": [ + {"pool": "llama8b-pd-pool", + "workload": {"kind": "prefill", "model": "meta-llama/Llama-3.1-8B-Instruct", + "slo": {"tpot_p90_ms": 100}, + "labels": {"airm.ai-computing.io/accelerator": "furiosa.RNGD"}}, + "acceleratorIds": ["/furiosa:0"], "kvTransport": "shm", "reason": "highest peak"}, + {"pool": "llama8b-pd-pool", + "workload": {"kind": "decode", "model": "meta-llama/Llama-3.1-8B-Instruct", + "slo": {"tpot_p90_ms": 100}, + "labels": {"airm.ai-computing.io/accelerator": "furiosa.RNGD"}}, + "acceleratorIds": ["/furiosa:1"], "kvTransport": "shm", "reason": "most memory"} + ] +} +``` + +| Field | Rule | +|---|---| +| `protocol` | REQUIRED. `"MAJOR.MINOR"`, the protocol version the writer implements | +| `generation` | REQUIRED. Integer ≥ 0; `0` only for the empty document | +| `updatedAt` | REQUIRED. RFC 3339 UTC (`YYYY-MM-DDTHH:MM:SSZ`) time of the last change; `null` only when `generation` is 0 | +| `placement` | Name of the `airm.placements` policy that produced the decisions. REQUIRED when `generation` > 0 | +| `decisions` | REQUIRED. List, empty in the empty document | +| `decisions[].pool` | REQUIRED. InferencePool name (= the EPP's `--pool-name`) | +| `decisions[].workload` | REQUIRED. `Workload` in camelCase: `kind` and `model` REQUIRED; `slo` (keys from `SLO_KEYS`), `labels` (strings), `precision`, `batchHint` OPTIONAL | +| `decisions[].workload.kind` | The role. A prefill/decode pool has one `prefill` and one `decode` decision; an aggregated pool (any role) has one `serve` decision | +| `decisions[].workload.labels["airm.ai-computing.io/accelerator"]` | OPTIONAL. Card type the role was restricted to, `.` as in the node labels (§3.1). When set, every card in `acceleratorIds` MUST be of that vendor | +| `decisions[].acceleratorIds` | REQUIRED. Card references: `Accelerator.id`s `/:` (= `AcceleratorProfile.spec.id`), no duplicates. The node of a card is the part before `/` | +| `decisions[].kvTransport` | `airm.kv_transports` name for the prefill→decode hand-off; `""` or absent if none | +| `decisions[].reason` | Human-readable. REQUIRED (non-empty) when `acceleratorIds` is empty, which means nothing fits | + +The empty document, valid before the first placement, is +`{"protocol": "0.1", "generation": 0, "updatedAt": null, "decisions": []}`. + +Writers: +- MUST write only documents that `rm_core.placement.validate` accepts. +- MUST write at most one decision per (`pool`, `workload.kind`). +- MUST increase `generation` by at least one whenever `placement` or `decisions` change, and MUST + NOT decrease it. A writer SHOULD NOT rewrite a document whose `placement` and `decisions` are + unchanged; if it does, it MUST keep the same `generation`. +- MUST NOT put addresses or credentials in the document (§7): cards are referenced by id only and + resolved through `AcceleratorProfile`s. + +Readers: +- MUST ignore fields they do not know, at every level, and MUST accept a newer minor `protocol` + with the same major version. +- MUST NOT apply a document whose `generation` is not greater than the last one they applied + (`rm_core.placement.is_newer`); equal generations have equal content. +- MUST behave as if the ConfigMap did not exist when it is absent, when `decisions` is empty, or + when it has no decision for their pool. +- SHOULD treat a document that fails validation as absent and report it, not crash. + ## 6. Metrics - Components MUST report the quantities listed in `rm_core.metrics.METRICS` under exactly those @@ -130,12 +200,19 @@ keys. Site values come from the environment (`cluster.local.env`), never from co ## 8. Conformance `python -m rm_core.conformance [group [name ...]]` loads every registered implementation (or the -named ones) and runs the checks for its group. A component is conformant when its adapters pass, +named ones) and runs the checks for its group. For `airm.placements` the checks include that the +policy's decisions make a valid placement document (§5.4). A component is conformant when its adapters pass, its tests pass, and `check_sensitive.py` reports clean. Each component's CI SHOULD run the checks against its own entry points. ## Change log +- 0.1.2 (2026-10-19, rm-core patch; protocol stays 0.1): §5.4 makes the `airm-placement` document + normative as already written by rm-controller and hetero-llm-d and read by lmd-top (no field + changed). `rm_core.placement` (typed model, stdlib validator, `build` with the generation rule), + `k8s/airm-placement.schema.json`, and the `airm.placements` conformance check now also validates + the document made from a policy's decisions. + - 0.1.1 (2026-10-19, rm-core patch; protocol stays 0.1): `AcceleratorProfile` keeps `extra` (CRD field added) and the provider check now verifies the profile round-trip; `NoHardware` moved to `rm_core.errors` so `python -m rm_core.conformance` reports SKIP correctly; §6 migration rule. diff --git a/README.ko.md b/README.ko.md index 2306ff4..2d656d4 100644 --- a/README.ko.md +++ b/README.ko.md @@ -193,6 +193,7 @@ src/rm_core/ ├── accelerator.py # L1: Accelerator, AcceleratorProvider; 벤더·메모리·정밀도·인터커넥트 열거형 ├── memory.py # L2: MemoryTier, KVLayout, TransferPlan/Handle/Result, KVTransport ├── workload.py # L3–L4: Workload, PlacementDecision, Placement; 워크로드 종류, SLO 키 +├── placement.py # L4: airm-placement 문서(decisions.json): 모델, validate(), build() ├── metrics.py # METRICS, TAG_KEYS, Measurement, prometheus_name ├── k8s.py # AcceleratorProfile 객체와 노드 레이블; from_profile() ├── registry.py # entry-point 그룹: available(), load(), groups() @@ -200,6 +201,7 @@ src/rm_core/ ├── conformance.py # python -m rm_core.conformance └── errors.py # ContractError, TransportError, NoHardware k8s/acceleratorprofile.crd.yaml # AcceleratorProfile CRD (airm.ai-computing.io/v1alpha1) +k8s/airm-placement.schema.json # airm-placement 문서의 JSON Schema (PROTOCOL §5.4) ``` 노드 데몬은 장치마다 클러스터 범위 `AcceleratorProfile`을 게시하고, 요약 노드 레이블 @@ -220,6 +222,8 @@ CI는 Python 3.9–3.12에서 테스트를 돌린 뒤 `python -m rm_core.conform - [x] 프로토콜 v0.1: 계약, `AcceleratorProfile` CRD, 참조 구현, 적합성 검사 - [x] 0.1.1: `AcceleratorProfile`이 `extra` 유지, `python -m` 실행 시 `NoHardware`를 SKIP으로 보고, 메트릭 이전 규칙 +- [x] 0.1.2 (미릴리스): `airm-placement` 문서 규범화(§5.4), `rm_core.placement` 검증기, JSON Schema, + 배치 적합성 검사가 문서까지 검증 - [ ] 프로토콜 v0.2: 전송 계획에 `MemoryRegion`(owner, fd/offset, size)을 넣어, 장치 간 전송이 실제 버퍼를 묶고 하드웨어 fixture로 적합성 검사를 통과할 수 있게 함 - [ ] 프로토콜 v0.2: 워크로드 runner 인터페이스(`airm.workloads`) 정식화 @@ -230,7 +234,7 @@ CI는 Python 3.9–3.12에서 테스트를 돌린 뒤 `python -m rm_core.conform | | | |---|---| -| **명세** | [PROTOCOL.md](PROTOCOL.md) · [AcceleratorProfile CRD](k8s/acceleratorprofile.crd.yaml) | +| **명세** | [PROTOCOL.md](PROTOCOL.md) · [AcceleratorProfile CRD](k8s/acceleratorprofile.crd.yaml) · [배치 스키마](k8s/airm-placement.schema.json) | | **ai-rm** | [Umbrella](https://github.com/ai-computing/ai-rm) · [인터페이스](https://github.com/ai-computing/ai-rm/blob/master/docs/interfaces.md) · [통합 안내](https://github.com/ai-computing/ai-rm/blob/master/docs/integration-guide.md) | ## 📄 라이선스 diff --git a/README.md b/README.md index 0b9462e..6d1ca26 100644 --- a/README.md +++ b/README.md @@ -195,6 +195,7 @@ src/rm_core/ ├── accelerator.py # L1: Accelerator, AcceleratorProvider; vendor, memory, precision, interconnect enums ├── memory.py # L2: MemoryTier, KVLayout, TransferPlan/Handle/Result, KVTransport ├── workload.py # L3–L4: Workload, PlacementDecision, Placement; workload kinds, SLO keys +├── placement.py # L4: airm-placement document (decisions.json): model, validate(), build() ├── metrics.py # METRICS, TAG_KEYS, Measurement, prometheus_name ├── k8s.py # AcceleratorProfile objects and node labels; from_profile() ├── registry.py # entry-point groups: available(), load(), groups() @@ -202,6 +203,7 @@ src/rm_core/ ├── conformance.py # python -m rm_core.conformance └── errors.py # ContractError, TransportError, NoHardware k8s/acceleratorprofile.crd.yaml # AcceleratorProfile CRD (airm.ai-computing.io/v1alpha1) +k8s/airm-placement.schema.json # JSON Schema of the airm-placement document (PROTOCOL §5.4) ``` A node daemon publishes each device as a cluster-scoped `AcceleratorProfile` plus summary node labels @@ -222,6 +224,8 @@ CI runs the tests on Python 3.9–3.12, then `python -m rm_core.conformance` and - [x] Protocol v0.1: contracts, `AcceleratorProfile` CRD, reference implementations, conformance checks - [x] 0.1.1: `AcceleratorProfile` keeps `extra`, `NoHardware` reported as SKIP under `python -m`, metric migration rule +- [x] 0.1.2 (unreleased): `airm-placement` document normative (§5.4), `rm_core.placement` validator, JSON Schema, + placement conformance validates the document - [ ] Protocol v0.2: `MemoryRegion` (owner, fd/offset, size) in transfer plans, so device-to-device transports can bind live buffers and pass conformance with a hardware fixture - [ ] Protocol v0.2: formalized workload-runner interface (`airm.workloads`) @@ -232,7 +236,7 @@ Change log: [PROTOCOL.md § Change log](PROTOCOL.md#change-log). | | | |---|---| -| **Specification** | [PROTOCOL.md](PROTOCOL.md) · [AcceleratorProfile CRD](k8s/acceleratorprofile.crd.yaml) | +| **Specification** | [PROTOCOL.md](PROTOCOL.md) · [AcceleratorProfile CRD](k8s/acceleratorprofile.crd.yaml) · [placement schema](k8s/airm-placement.schema.json) | | **ai-rm** | [Umbrella](https://github.com/ai-computing/ai-rm) · [Interfaces](https://github.com/ai-computing/ai-rm/blob/master/docs/interfaces.md) · [Integration guide](https://github.com/ai-computing/ai-rm/blob/master/docs/integration-guide.md) | ## 📄 License diff --git a/k8s/airm-placement.schema.json b/k8s/airm-placement.schema.json new file mode 100644 index 0000000..6769666 --- /dev/null +++ b/k8s/airm-placement.schema.json @@ -0,0 +1,64 @@ +{ + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "https://github.com/ai-computing/rm-core/blob/main/k8s/airm-placement.schema.json", + "title": "airm-placement decisions.json (ai-rm protocol 0.1, PROTOCOL.md section 5.4)", + "description": "Value of key decisions.json in ConfigMap airm-placement. Unknown properties are allowed and ignored by readers. Rules a schema cannot express (generation > 0 when there are decisions, updatedAt and placement set when generation > 0, one decision per pool and workload.kind, card vendor matching the airm.ai-computing.io/accelerator label) are checked by rm_core.placement.validate.", + "type": "object", + "required": ["protocol", "generation", "updatedAt", "decisions"], + "properties": { + "protocol": {"type": "string", "pattern": "^0\\.\\d+$"}, + "generation": {"type": "integer", "minimum": 0}, + "updatedAt": { + "oneOf": [ + {"type": "null"}, + {"type": "string", "pattern": "^\\d{4}-\\d{2}-\\d{2}T\\d{2}:\\d{2}:\\d{2}(\\.\\d+)?Z$"} + ] + }, + "placement": {"type": "string", "minLength": 1}, + "decisions": {"type": "array", "items": {"$ref": "#/$defs/decision"}} + }, + "$defs": { + "decision": { + "type": "object", + "required": ["pool", "workload", "acceleratorIds"], + "properties": { + "pool": {"type": "string", "minLength": 1}, + "workload": {"$ref": "#/$defs/workload"}, + "acceleratorIds": { + "type": "array", + "uniqueItems": true, + "items": {"type": "string", "pattern": "^[a-z0-9][a-z0-9._-]*/[a-z0-9]+:\\d+$"} + }, + "kvTransport": {"type": "string"}, + "reason": {"type": "string"} + }, + "if": {"properties": {"acceleratorIds": {"maxItems": 0}}}, + "then": {"required": ["reason"], "properties": {"reason": {"minLength": 1}}} + }, + "workload": { + "type": "object", + "required": ["kind", "model"], + "properties": { + "kind": {"enum": ["prefill", "decode", "serve", "train", "embed", "search"]}, + "model": {"type": "string", "minLength": 1}, + "slo": { + "type": "object", + "propertyNames": {"enum": ["ttft_p90_ms", "ttft_p99_ms", "tpot_p90_ms", "tpot_p99_ms", "e2e_p90_ms", "tok_s_min"]}, + "additionalProperties": {"type": "number"} + }, + "labels": { + "type": "object", + "additionalProperties": {"type": "string"}, + "properties": { + "airm.ai-computing.io/accelerator": { + "type": "string", + "pattern": "^(nvidia|furiosa|rebellions|cpu|amd|intel|other)\\.[A-Za-z0-9][A-Za-z0-9._-]*$" + } + } + }, + "precision": {"enum": ["fp32", "tf32", "bf16", "fp16", "fp8", "int8", "int4"]}, + "batchHint": {"oneOf": [{"type": "null"}, {"type": "integer", "minimum": 1}]} + } + } + } +} diff --git a/pyproject.toml b/pyproject.toml index e222d28..1d1ed23 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "rm-core" -version = "0.1.1" +version = "0.1.2" description = "Contracts that connect ai-rm components: accelerators, memory tiers, KV transport, placement, metrics" readme = "README.md" requires-python = ">=3.9" diff --git a/src/rm_core/__init__.py b/src/rm_core/__init__.py index c4a5d35..916e9db 100644 --- a/src/rm_core/__init__.py +++ b/src/rm_core/__init__.py @@ -10,7 +10,7 @@ from .metrics import METRICS, Measurement, prometheus_name from .workload import Placement, PlacementDecision, Workload -__version__ = "0.1.1" +__version__ = "0.1.2" PROTOCOL_VERSION = "0.1" __all__ = [ diff --git a/src/rm_core/conformance.py b/src/rm_core/conformance.py index 5198431..1bd8fc3 100644 --- a/src/rm_core/conformance.py +++ b/src/rm_core/conformance.py @@ -9,9 +9,11 @@ """ from __future__ import annotations +import json import sys from typing import Callable, Dict, List, Optional, Tuple +from . import placement as placement_doc from . import registry from .accelerator import Accelerator, AcceleratorProvider from .errors import ContractError, TransportError @@ -80,14 +82,23 @@ def check_placement(placement, pool: Optional[List[Accelerator]] = None) -> None assert isinstance(placement, Placement), "missing name/place()" pool = pool or SimAcceleratorProvider("conf-node").discover() ids = {a.id for a in pool} + entries = [] for kind in ("prefill", "decode"): w = Workload(kind=kind, model="conformance-model", slo={"tpot_p90_ms": 100}) d1, d2 = placement.place(w, pool), placement.place(w, pool) assert isinstance(d1, PlacementDecision), "place() must return a PlacementDecision" assert set(d1.accelerator_ids) <= ids, f"placed outside the pool: {d1.accelerator_ids}" assert d1 == d2, "place() must be deterministic for the same inputs" - empty = placement.place(Workload(kind="decode", model="m"), []) + entries.append(placement_doc.Decision.from_placement("conf-pool", w, d1)) + w = Workload(kind="decode", model="m") + empty = placement.place(w, []) assert empty.accelerator_ids == [] and empty.reason, "empty pool -> empty decision with a reason" + entries.append(placement_doc.Decision.from_placement("conf-empty-pool", w, empty)) + # §5.4: the decisions must make a valid airm-placement document. + doc = placement_doc.PlacementDocument(1, "2026-01-01T00:00:00Z", placement.name, entries) + errors = placement_doc.validate(json.loads(doc.to_json())) + assert not errors, f"decisions do not make a valid airm-placement document: {errors}" + assert placement_doc.PlacementDocument.from_dict(doc.to_dict()) == doc, "document round-trip changed it" def check_measurements(measurements: List[Measurement]) -> None: diff --git a/src/rm_core/placement.py b/src/rm_core/placement.py new file mode 100644 index 0000000..3a4978c --- /dev/null +++ b/src/rm_core/placement.py @@ -0,0 +1,280 @@ +"""L4: the placement document — ConfigMap `airm-placement`, key `decisions.json` (PROTOCOL.md §5.4). + +The placement controller writes it; the llm-d EPP scorer, serving pods and lmd-top read it. + + {"protocol": "0.1", "generation": 3, "updatedAt": "2026-10-19T00:00:00Z", "placement": "topology", + "decisions": [{"pool": "llama8b-pd-pool", + "workload": {"kind": "decode", "model": "...", "slo": {...}, "labels": {...}}, + "acceleratorIds": ["/furiosa:1"], "kvTransport": "shm", "reason": "..."}]} + +`validate(doc)` lists every violation (empty list = valid) and ignores unknown fields. +`PlacementDocument` / `Decision` are the typed model; `build()` bumps `generation` only when the +placement or the decisions change. JSON Schema: k8s/airm-placement.schema.json. +""" +from __future__ import annotations + +import json +import re +from dataclasses import dataclass, field +from datetime import datetime, timezone +from typing import Dict, List, Optional, Sequence, Tuple + +from . import PROTOCOL_VERSION +from .accelerator import _ID, PRECISIONS, VENDORS +from .errors import ContractError +from .workload import SLO_KEYS, WORKLOAD_KINDS, PlacementDecision, Workload + +CONFIGMAP = "airm-placement" +KEY = "decisions.json" +# Optional workload label: the card type a role was restricted to, `.` as in the +# node labels `airm.ai-computing.io/.=` (rm_core.k8s.node_labels). +ACCELERATOR_LABEL = "airm.ai-computing.io/accelerator" +PROTOCOL_ANNOTATION = "airm.ai-computing.io/protocol" + +_PROTOCOL = re.compile(r"^\d+\.\d+$") +_TIME = re.compile(r"^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(\.\d+)?Z$") +_VENDOR_MODEL = re.compile(r"^([a-z0-9]+)\.([A-Za-z0-9][A-Za-z0-9._-]*)$") + + +def _is_int(v) -> bool: + return isinstance(v, int) and not isinstance(v, bool) + + +def _is_num(v) -> bool: + return isinstance(v, (int, float)) and not isinstance(v, bool) + + +def _check_workload(w, at: str, errors: List[str]) -> None: + if not isinstance(w, dict): + errors.append(f"{at}: must be an object") + return + if w.get("kind") not in WORKLOAD_KINDS: + errors.append(f"{at}.kind: {w.get('kind')!r} not in {WORKLOAD_KINDS}") + if not isinstance(w.get("model"), str) or not w.get("model"): + errors.append(f"{at}.model: must be a non-empty string") + slo = w.get("slo", {}) + if not isinstance(slo, dict): + errors.append(f"{at}.slo: must be an object") + else: + for k, v in slo.items(): + if k not in SLO_KEYS: + errors.append(f"{at}.slo: unknown key {k!r}; allowed {SLO_KEYS}") + elif not _is_num(v): + errors.append(f"{at}.slo.{k}: must be a number") + labels = w.get("labels", {}) + if not isinstance(labels, dict) or not all(isinstance(k, str) and isinstance(v, str) for k, v in labels.items()): + errors.append(f"{at}.labels: must be an object of strings") + elif ACCELERATOR_LABEL in labels: + m = _VENDOR_MODEL.match(labels[ACCELERATOR_LABEL]) + if not m or m.group(1) not in VENDORS: + errors.append(f"{at}.labels[{ACCELERATOR_LABEL}]: {labels[ACCELERATOR_LABEL]!r} must be " + f"'.' with vendor in {VENDORS}") + if "precision" in w and w["precision"] not in PRECISIONS: + errors.append(f"{at}.precision: {w['precision']!r} not in {PRECISIONS}") + if w.get("batchHint") is not None and not (_is_int(w["batchHint"]) and w["batchHint"] > 0): + errors.append(f"{at}.batchHint: must be a positive integer or null") + + +def _check_decision(d, at: str, errors: List[str]) -> None: + if not isinstance(d, dict): + errors.append(f"{at}: must be an object") + return + if not isinstance(d.get("pool"), str) or not d.get("pool"): + errors.append(f"{at}.pool: must be a non-empty string") + _check_workload(d.get("workload"), f"{at}.workload", errors) + ids = d.get("acceleratorIds") + if not isinstance(ids, list): + errors.append(f"{at}.acceleratorIds: must be a list") + ids = [] + for i, a in enumerate(ids): + if not isinstance(a, str) or not _ID.match(a): + errors.append(f"{at}.acceleratorIds[{i}]: {a!r} must look like '/:'") + if len(ids) != len({repr(a) for a in ids}): + errors.append(f"{at}.acceleratorIds: duplicate ids {ids}") + for k in ("kvTransport", "reason"): + if k in d and not isinstance(d[k], str): + errors.append(f"{at}.{k}: must be a string") + if isinstance(d.get("acceleratorIds"), list) and not ids and not d.get("reason"): + errors.append(f"{at}: empty acceleratorIds needs a reason") + # The vendor in each id must agree with the card type the role was restricted to. + labels = (d.get("workload") or {}).get("labels") if isinstance(d.get("workload"), dict) else None + m = _VENDOR_MODEL.match(labels.get(ACCELERATOR_LABEL, "")) if isinstance(labels, dict) else None + if m: + for a in ids: + if isinstance(a, str) and _ID.match(a) and a.split("/", 1)[1].split(":")[0] != m.group(1): + errors.append(f"{at}.acceleratorIds: {a!r} is not a {m.group(1)} card ({ACCELERATOR_LABEL})") + + +def validate(doc) -> List[str]: + """Every violation of PROTOCOL §5.4 in `doc` (a parsed `decisions.json`); [] when valid. + + Unknown fields are ignored at every level (forward compatibility). + """ + errors: List[str] = [] + if not isinstance(doc, dict): + return ["document: must be a JSON object"] + p = doc.get("protocol") + if not isinstance(p, str) or not _PROTOCOL.match(p): + errors.append(f"protocol: {p!r} must be a 'MAJOR.MINOR' string") + elif p.split(".")[0] != PROTOCOL_VERSION.split(".")[0]: + errors.append(f"protocol: major version {p!r} is not {PROTOCOL_VERSION}'s") + gen = doc.get("generation") + if not _is_int(gen) or gen < 0: + errors.append(f"generation: {gen!r} must be an integer >= 0") + gen = None + if "updatedAt" not in doc: + errors.append("updatedAt: required (null before the first write)") + else: + t = doc["updatedAt"] + if t is None: + if gen: + errors.append("updatedAt: must be set when generation > 0") + elif not isinstance(t, str) or not _TIME.match(t): + errors.append(f"updatedAt: {t!r} must be RFC 3339 UTC ('YYYY-MM-DDTHH:MM:SSZ') or null") + pol = doc.get("placement") + if pol is not None and (not isinstance(pol, str) or not pol): + errors.append("placement: must be a non-empty string") + elif pol is None and gen: + errors.append("placement: required when generation > 0 (the policy that wrote the decisions)") + decisions = doc.get("decisions") + if not isinstance(decisions, list): + errors.append("decisions: must be a list") + return errors + if decisions and gen == 0: + errors.append("generation: must be > 0 when there are decisions") + seen = set() + for i, d in enumerate(decisions): + _check_decision(d, f"decisions[{i}]", errors) + if isinstance(d, dict) and isinstance(d.get("workload"), dict): + key = (d.get("pool"), d["workload"].get("kind")) + if key in seen: + errors.append(f"decisions[{i}]: second decision for pool {key[0]!r} kind {key[1]!r}") + seen.add(key) + return errors + + +def _workload_to_dict(w: Workload) -> dict: + out = {"kind": w.kind, "model": w.model, "slo": dict(w.slo), "labels": dict(w.labels)} + if w.precision != "bf16": + out["precision"] = w.precision + if w.batch_hint is not None: + out["batchHint"] = w.batch_hint + return out + + +def _workload_from_dict(d: dict) -> Workload: + return Workload(kind=d["kind"], model=d["model"], slo=dict(d.get("slo") or {}), + precision=d.get("precision") or "bf16", batch_hint=d.get("batchHint"), + labels=dict(d.get("labels") or {})) + + +@dataclass(frozen=True) +class Decision: + """One entry of `decisions`: the accelerators one role (`workload.kind`) of one pool got.""" + + pool: str + workload: Workload + accelerator_ids: List[str] = field(default_factory=list) + kv_transport: str = "" + reason: str = "" + + @classmethod + def from_placement(cls, pool: str, workload: Workload, d: PlacementDecision) -> "Decision": + return cls(pool, workload, list(d.accelerator_ids), d.kv_transport, d.reason) + + @property + def decision(self) -> PlacementDecision: + return PlacementDecision(list(self.accelerator_ids), self.kv_transport, self.reason) + + @property + def nodes(self) -> List[str]: + """Nodes of the cards, in order of first appearance (the part of each id before '/').""" + return list(dict.fromkeys(a.split("/", 1)[0] for a in self.accelerator_ids)) + + def to_dict(self) -> dict: + return {"pool": self.pool, "workload": _workload_to_dict(self.workload), + "acceleratorIds": list(self.accelerator_ids), "kvTransport": self.kv_transport, + "reason": self.reason} + + @classmethod + def from_dict(cls, d: dict) -> "Decision": + return cls(d["pool"], _workload_from_dict(d["workload"]), list(d.get("acceleratorIds") or []), + d.get("kvTransport") or "", d.get("reason") or "") + + +@dataclass(frozen=True) +class PlacementDocument: + """The value of `airm-placement` / `decisions.json`. The default is the empty v0.1 document.""" + + generation: int = 0 + updated_at: Optional[str] = None + placement: Optional[str] = None + decisions: List[Decision] = field(default_factory=list) + protocol: str = PROTOCOL_VERSION + + def validate(self) -> "PlacementDocument": + errors = validate(self.to_dict()) + if errors: + raise ContractError("; ".join(errors)) + return self + + def for_pool(self, pool: str) -> List[Decision]: + return [d for d in self.decisions if d.pool == pool] + + def to_dict(self) -> dict: + out = {"protocol": self.protocol, "generation": self.generation, "updatedAt": self.updated_at} + if self.placement is not None: + out["placement"] = self.placement + out["decisions"] = [d.to_dict() for d in self.decisions] + return out + + def to_json(self) -> str: + return json.dumps(self.to_dict(), indent=2) + "\n" + + @classmethod + def from_dict(cls, doc: dict) -> "PlacementDocument": + """Raises ContractError listing every violation; unknown fields are dropped.""" + errors = validate(doc) + if errors: + raise ContractError("; ".join(errors)) + return cls(doc["generation"], doc["updatedAt"], doc.get("placement"), + [Decision.from_dict(d) for d in doc["decisions"]], doc["protocol"]) + + +def parse(cm: Optional[dict]) -> PlacementDocument: + """The document inside a ConfigMap object; the empty document when the ConfigMap or key is absent.""" + raw = ((cm or {}).get("data") or {}).get(KEY) + if not raw: + return PlacementDocument() + try: + doc = json.loads(raw) + except ValueError as e: + raise ContractError(f"{CONFIGMAP}/{KEY}: not JSON ({e})") from None + return PlacementDocument.from_dict(doc) + + +def build(decisions: Sequence[Decision], previous: PlacementDocument, placement: str, + now: Optional[datetime] = None) -> Tuple[PlacementDocument, bool]: + """The next document and whether it changed. `generation` goes up only on a change of + `placement` or `decisions`, so an unchanged rewrite wakes no reader.""" + decisions = list(decisions) + if previous.decisions == decisions and previous.placement == placement: + return previous, False + now = (now or datetime.now(timezone.utc)).astimezone(timezone.utc) + return PlacementDocument(previous.generation + 1, now.strftime("%Y-%m-%dT%H:%M:%SZ"), placement, + decisions).validate(), True + + +def configmap(doc: PlacementDocument, namespace: str, labels: Optional[Dict[str, str]] = None) -> dict: + """The ConfigMap object to apply in the llm-d namespace (`namespace`).""" + return {"apiVersion": "v1", "kind": "ConfigMap", + "metadata": {"name": CONFIGMAP, "namespace": namespace, + "labels": {"app.kubernetes.io/part-of": "ai-rm", "airm.ai-computing.io/managed": "true", + **(labels or {})}, + "annotations": {PROTOCOL_ANNOTATION: doc.protocol}}, + "data": {KEY: doc.validate().to_json()}} + + +def is_newer(doc: PlacementDocument, applied_generation: Optional[int]) -> bool: + """Reader rule: apply a document only if its generation is above the last one applied.""" + return applied_generation is None or doc.generation > applied_generation diff --git a/tests/test_placement.py b/tests/test_placement.py new file mode 100644 index 0000000..52bdcbb --- /dev/null +++ b/tests/test_placement.py @@ -0,0 +1,218 @@ +import copy +import json +import os +import unittest +from datetime import datetime, timezone + +from rm_core import ContractError, PlacementDecision, Workload +from rm_core import accelerator, conformance, placement, workload +from rm_core.placement import Decision, PlacementDocument + +SCHEMA = os.path.join(os.path.dirname(__file__), "..", "k8s", "airm-placement.schema.json") + +# What the existing writers and readers already exchange (protocol 0.1): +EMPTY = {"protocol": "0.1", "generation": 0, "updatedAt": None, "decisions": []} # hetero-llm-d default +HETERO_LLM_D_EXAMPLE = { # hetero-llm-d deploy/airm/README.md + "protocol": "0.1", "generation": 12, "updatedAt": "2026-10-19T00:00:00Z", "placement": "greedy", + "decisions": [{"pool": "llama8b-pd-pool", + "workload": {"kind": "decode", "model": "meta-llama/Llama-3.1-8B-Instruct", + "slo": {"tpot_p90_ms": 40}, "labels": {}}, + "acceleratorIds": ["node-a/furiosa:0", "node-a/furiosa:1"], "kvTransport": "p2p", + "reason": "lowest tpot_p90 within SLO"}]} +RM_CONTROLLER_OUTPUT = { # ai-rm rm-controller (document.py), hetero P2P example spec + "protocol": "0.1", "generation": 1, "updatedAt": "2026-10-19T01:02:03Z", "placement": "topology", + "decisions": [ + {"pool": "koni-hetero-pd-pool", + "workload": {"kind": "prefill", "model": "KISTI-KONI/KONI-Llama3.1-8B-Instruct-20241024", "slo": {}, + "labels": {"airm.ai-computing.io/accelerator": "furiosa.RNGD"}}, + "acceleratorIds": ["node-a/furiosa:0"], "kvTransport": "p2p", "reason": "highest peak"}, + {"pool": "koni-hetero-pd-pool", + "workload": {"kind": "decode", "model": "KISTI-KONI/KONI-Llama3.1-8B-Instruct-20241024", "slo": {}, + "labels": {"airm.ai-computing.io/accelerator": "rebellions.ATOM"}}, + "acceleratorIds": ["node-a/rebellions:0", "node-a/rebellions:1"], "kvTransport": "p2p", + "reason": "most memory"}]} +LMD_TOP_FIXTURE = { # lmd-top src/airm.rs test + "protocol": "0.1", "generation": 1, "placement": "manual", "updatedAt": "2026-10-19T00:00:00Z", + "decisions": [ + {"pool": "llama8b-pd-pool", "workload": {"kind": "prefill", "model": "m"}, + "acceleratorIds": ["n/furiosa:0"], "kvTransport": "shm", "reason": "pinned"}, + {"pool": "llama8b-pd-pool", "workload": {"kind": "decode", "model": "m"}, + "acceleratorIds": ["n/furiosa:1"], "kvTransport": "shm", "reason": "next to prefill"}]} +VALID = [EMPTY, HETERO_LLM_D_EXAMPLE, RM_CONTROLLER_OUTPUT, LMD_TOP_FIXTURE] + + +def _change(doc, path, value): + """A copy of doc with the value at path (keys/indexes) replaced; value=KeyError deletes it.""" + out = copy.deepcopy(doc) + target = out + for k in path[:-1]: + target = target[k] + if value is KeyError: + del target[path[-1]] + else: + target[path[-1]] = value + return out + + +INVALID = [ + ("not an object", []), + ("no protocol", _change(EMPTY, ["protocol"], KeyError)), + ("protocol not MAJOR.MINOR", _change(EMPTY, ["protocol"], "v0.1")), + ("other major", _change(EMPTY, ["protocol"], "1.0")), + ("negative generation", _change(EMPTY, ["generation"], -1)), + ("bool generation", _change(EMPTY, ["generation"], True)), + ("no updatedAt", _change(EMPTY, ["updatedAt"], KeyError)), + ("local time", _change(HETERO_LLM_D_EXAMPLE, ["updatedAt"], "2026-10-19T09:00:00+09:00")), + ("null updatedAt after a write", _change(HETERO_LLM_D_EXAMPLE, ["updatedAt"], None)), + ("no policy name after a write", _change(HETERO_LLM_D_EXAMPLE, ["placement"], KeyError)), + ("decisions at generation 0", _change(HETERO_LLM_D_EXAMPLE, ["generation"], 0)), + ("decisions not a list", _change(EMPTY, ["decisions"], {})), + ("no pool", _change(LMD_TOP_FIXTURE, ["decisions", 0, "pool"], "")), + ("unknown kind", _change(LMD_TOP_FIXTURE, ["decisions", 0, "workload", "kind"], "any")), + ("no model", _change(LMD_TOP_FIXTURE, ["decisions", 0, "workload", "model"], KeyError)), + ("unknown SLO", _change(HETERO_LLM_D_EXAMPLE, ["decisions", 0, "workload", "slo"], {"latency": 1})), + ("bad accelerator id", _change(LMD_TOP_FIXTURE, ["decisions", 0, "acceleratorIds"], ["furiosa:0"])), + ("duplicate card", _change(LMD_TOP_FIXTURE, ["decisions", 0, "acceleratorIds"], ["n/furiosa:0"] * 2)), + ("empty without reason", _change(_change(LMD_TOP_FIXTURE, ["decisions", 0, "acceleratorIds"], []), + ["decisions", 0, "reason"], "")), + ("bad label", _change(RM_CONTROLLER_OUTPUT, ["decisions", 0, "workload", "labels", + "airm.ai-computing.io/accelerator"], "RNGD")), + ("card of another vendor", _change(RM_CONTROLLER_OUTPUT, ["decisions", 1, "acceleratorIds"], + ["node-a/furiosa:1"])), + ("same pool and kind twice", _change(LMD_TOP_FIXTURE, ["decisions", 1, "workload", "kind"], "prefill")), + ("kvTransport not a string", _change(LMD_TOP_FIXTURE, ["decisions", 0, "kvTransport"], None)), +] + + +class ValidateTest(unittest.TestCase): + def test_existing_documents_are_valid(self): + for doc in VALID: + with self.subTest(doc=doc.get("placement")): + self.assertEqual(placement.validate(doc), []) + + def test_violations_are_reported(self): + for name, doc in INVALID: + with self.subTest(name), self.assertRaises(ContractError): + self.assertTrue(placement.validate(doc)) + PlacementDocument.from_dict(doc) + + def test_unknown_fields_are_ignored(self): + doc = copy.deepcopy(LMD_TOP_FIXTURE) + doc["future"] = {"x": 1} + doc["decisions"][0]["cost"] = 3 + doc["decisions"][0]["workload"]["shape"] = "chat" + self.assertEqual(placement.validate(doc), []) + self.assertEqual(PlacementDocument.from_dict(doc), PlacementDocument.from_dict(LMD_TOP_FIXTURE)) + + def test_newer_minor_is_accepted(self): + self.assertEqual(placement.validate(_change(EMPTY, ["protocol"], "0.9")), []) + + +class ModelTest(unittest.TestCase): + def test_round_trip(self): + for doc in VALID: + with self.subTest(doc=doc.get("placement")): + model = PlacementDocument.from_dict(doc) + self.assertEqual(PlacementDocument.from_dict(json.loads(model.to_json())), model) + + def test_rm_controller_output_is_reproduced_exactly(self): + # rm-controller can switch to rm_core.placement without bumping generation on the live document. + self.assertEqual(PlacementDocument.from_dict(RM_CONTROLLER_OUTPUT).to_dict(), RM_CONTROLLER_OUTPUT) + + def test_decision_accessors(self): + d = PlacementDocument.from_dict(RM_CONTROLLER_OUTPUT).for_pool("koni-hetero-pd-pool")[1] + self.assertEqual(d.workload.kind, "decode") + self.assertEqual(d.nodes, ["node-a"]) + self.assertEqual(d.decision, PlacementDecision(["node-a/rebellions:0", "node-a/rebellions:1"], "p2p", + "most memory")) + + def test_parse_configmap(self): + self.assertEqual(placement.parse(None), PlacementDocument()) + self.assertEqual(placement.parse({"data": {}}), PlacementDocument()) + cm = placement.configmap(PlacementDocument.from_dict(LMD_TOP_FIXTURE), "llm-serving") + self.assertEqual(cm["metadata"]["name"], "airm-placement") + self.assertEqual(cm["metadata"]["annotations"]["airm.ai-computing.io/protocol"], "0.1") + self.assertEqual(placement.parse(cm), PlacementDocument.from_dict(LMD_TOP_FIXTURE)) + with self.assertRaises(ContractError): + placement.parse({"data": {"decisions.json": "{not json"}}) + + +class BuildTest(unittest.TestCase): + T = datetime(2026, 10, 19, 1, 2, 3, tzinfo=timezone.utc) + + def entries(self, ids): + w = Workload(kind="decode", model="m") + return [Decision.from_placement("p", w, PlacementDecision(ids, "shm", "most memory"))] + + def test_generation_moves_only_on_change(self): + doc0 = PlacementDocument() + doc1, changed = placement.build(self.entries(["n/furiosa:0"]), doc0, "topology", now=self.T) + self.assertTrue(changed) + self.assertEqual((doc1.generation, doc1.updated_at), (1, "2026-10-19T01:02:03Z")) + same, changed = placement.build(self.entries(["n/furiosa:0"]), doc1, "topology") + self.assertFalse(changed) + self.assertIs(same, doc1) + doc2, changed = placement.build(self.entries(["n/furiosa:1"]), doc1, "topology") + self.assertTrue(changed and doc2.generation == 2) + doc3, changed = placement.build(self.entries(["n/furiosa:1"]), doc2, "optimal") + self.assertTrue(changed and doc3.generation == 3) + + def test_reader_rule(self): + doc = PlacementDocument.from_dict(HETERO_LLM_D_EXAMPLE) + self.assertTrue(placement.is_newer(doc, None)) + self.assertTrue(placement.is_newer(doc, 11)) + self.assertFalse(placement.is_newer(doc, 12)) + + def test_invalid_decisions_are_not_written(self): + w = Workload(kind="decode", model="m") + bad = [Decision.from_placement("p", w, PlacementDecision([], "shm", ""))] + with self.assertRaises(ContractError): + placement.build(bad, PlacementDocument(), "topology") + + +class SchemaTest(unittest.TestCase): + def setUp(self): + with open(SCHEMA, encoding="utf-8") as f: + self.schema = json.load(f) + + def test_enumerations_match_the_code(self): + wl = self.schema["$defs"]["workload"]["properties"] + self.assertEqual(tuple(wl["kind"]["enum"]), workload.WORKLOAD_KINDS) + self.assertEqual(tuple(wl["slo"]["propertyNames"]["enum"]), workload.SLO_KEYS) + self.assertEqual(tuple(wl["precision"]["enum"]), accelerator.PRECISIONS) + label = wl["labels"]["properties"][placement.ACCELERATOR_LABEL]["pattern"] + self.assertTrue(label.startswith("^(" + "|".join(accelerator.VENDORS) + ")\\.")) + ids = self.schema["$defs"]["decision"]["properties"]["acceleratorIds"]["items"]["pattern"] + self.assertEqual(ids, accelerator._ID.pattern) + + def test_schema_agrees_on_the_examples(self): + try: + import jsonschema + except ImportError: + self.skipTest("jsonschema not installed (optional; the stdlib validator is normative in code)") + v = jsonschema.Draft202012Validator(self.schema) + for doc in VALID: + self.assertEqual([e.message for e in v.iter_errors(doc)], []) + # Structural violations the schema can express must fail it too. + for name, doc in INVALID: + if name in ("decisions at generation 0", "null updatedAt after a write", "no policy name after a write", + "card of another vendor", "same pool and kind twice"): + continue # cross-field rules: rm_core.placement.validate only + with self.subTest(name): + self.assertTrue(list(v.iter_errors(doc)), name) + + +class ConformanceTest(unittest.TestCase): + def test_bad_document_fails_placement_check(self): + class NoReason: + name = "no-reason" + + def place(self, workload, pool, metrics=None): + return PlacementDecision([], "", "empty pool" if not pool else "") + + with self.assertRaisesRegex(AssertionError, "airm-placement document"): + conformance.check_placement(NoReason()) + + +if __name__ == "__main__": + unittest.main()