P5: placement breadth — artifact locality, load, NVLink affinity

meridian_control/routes.py

@@ -132,7 +132,11 @@ def projection(request: Request):
 
 @router.post("/admin/place")
 def place(request: Request, body: dict = Body(...)):
-    placement = _svc(request).select_placement(int(body.get("required_vram_bytes", 0)))
+    placement = _svc(request).select_placement(
+        int(body.get("required_vram_bytes", 0)),
+        count=int(body.get("count", 1)),
+        artifact_digest=body.get("artifact_digest"),
+    )
     return {"placement": placement}
 
 

meridian_control/service.py

@@ -365,11 +365,17 @@ def post_observation(self, node_id: str, observation: dict) -> dict:
             return {"received": True}
 
     # --- placement (Phase 5, control side) -----------------------------
-    def select_placement(self, required_vram_bytes: int) -> Optional[dict]:
-        """Pick a routable node + device with enough reported allocatable VRAM
-        (DESIGN.md 24 P3). Consumes the capacity nodes report in observations;
-        prefers the device with the most headroom (spread). None if nothing fits."""
-        best: Optional[tuple[int, str, str]] = None  # (allocatable, node_id, device_id)
+    def select_placement(
+        self, required_vram_bytes: int, count: int = 1, artifact_digest: Optional[str] = None
+    ) -> Optional[dict]:
+        """Pick a routable node + device(s) for a new engine (DESIGN.md 24 P3).
+        Consumes the capacity nodes report: allocatable VRAM, held artifacts,
+        running-engine load, and NVLink groups. Preference order: nodes that
+        already hold the artifact (locality) > lower load > more headroom. For
+        count>1, the devices must share an NVLink group when groups are declared.
+        Returns None if nothing fits."""
+        best_key: Optional[tuple] = None
+        best: Optional[tuple[str, list[str], int]] = None  # (node_id, device_ids, headroom)
         with self._sf() as s:
             for node in s.scalars(select(Node)):
                 lease_valid = node.lease_expires_at is not None and self._now() < _aware(node.lease_expires_at)
@@ -378,14 +384,23 @@ def select_placement(self, required_vram_bytes: int) -> Optional[dict]:
                 obs = s.get(ObservationRow, node.node_id)
                 if obs is None:
                     continue
-                allocatable = obs.observation.get("capacity", {}).get("allocatable_vram_bytes", {})
-                for device_id, free in allocatable.items():
-                    free_i = int(free)
-                    if free_i >= required_vram_bytes and (best is None or free_i > best[0]):
-                        best = (free_i, node.node_id, device_id)
+                cap = obs.observation.get("capacity", {})
+                alloc = {d: int(v) for d, v in cap.get("allocatable_vram_bytes", {}).items()}
+                devices = _pick_devices(alloc, cap.get("nvlink_groups", {}), required_vram_bytes, count)
+                if devices is None:
+                    continue
+                headroom = sum(alloc[d] for d in devices)
+                artifact_hit = 1 if (artifact_digest and artifact_digest in cap.get("held_artifacts", [])) else 0
+                load = int(cap.get("running_engines", 0))
+                key = (artifact_hit, -load, headroom)  # maximize
+                if best_key is None or key > best_key:
+                    best_key, best = key, (node.node_id, devices, headroom)
         if best is None:
             return None
-        return {"node_id": best[1], "device_id": best[2], "allocatable_vram_bytes": best[0]}
+        return {
+            "node_id": best[0], "device_ids": best[1], "device_id": best[1][0],
+            "allocatable_vram_bytes": best[2],
+        }
 
     # --- read models (operator / gateway projection) -------------------
     def serving_projection(self) -> list[dict]:
@@ -410,3 +425,27 @@ def serving_projection(self) -> list[dict]:
 
 def _aware(d: datetime) -> datetime:
     return d if d.tzinfo is not None else d.replace(tzinfo=timezone.utc)
+
+
+def _pick_devices(
+    alloc: dict[str, int], groups: dict, required: int, count: int
+) -> Optional[list[str]]:
+    """Choose `count` devices each with >= required allocatable VRAM. For count>1,
+    when NVLink groups are declared the devices must share one group (affinity)."""
+    fitting = [d for d, free in alloc.items() if free >= required]
+    if not fitting:
+        return None
+    ranked = sorted(fitting, key=lambda d: alloc[d], reverse=True)  # most headroom first
+    if count == 1:
+        return [ranked[0]]
+    if groups:
+        by_group: dict[str, list[str]] = {}
+        for d in ranked:
+            g = groups.get(d)
+            if g is not None:
+                by_group.setdefault(g, []).append(d)
+        for devs in by_group.values():
+            if len(devs) >= count:
+                return devs[:count]
+        return None  # no single NVLink group has enough fitting devices
+    return ranked[:count] if len(ranked) >= count else None

tests/control/test_service.py

@@ -166,15 +166,18 @@ def test_rotate_certificate_reissues_for_same_key(tmp_path, clock, node_key):
     svc._ca.verify_cert(rotated["certificate"])
 
 
-def _ready_node_with_capacity(svc, node_key, allocatable):
+def _ready_node(svc, node_key, capacity):
     node_id = _enroll_auto(svc, node_key)["node_id"]
     svc.establish_session(node_id, {"agent_session_id": "s-" + node_id})
     svc.heartbeat(node_id, {"agent_session_id": "s-" + node_id, "fencing_epoch": 1, "sequence": 1})
-    svc.post_observation(node_id, {"sequence": 1, "engines": [],
-                                   "capacity": {"allocatable_vram_bytes": allocatable}})
+    svc.post_observation(node_id, {"sequence": 1, "engines": [], "capacity": capacity})
     return node_id
 
 
+def _ready_node_with_capacity(svc, node_key, allocatable):
+    return _ready_node(svc, node_key, {"allocatable_vram_bytes": allocatable})
+
+
 def test_select_placement_prefers_most_headroom(tmp_path, clock, node_key):
     g = 1024**3
     svc = make_service(tmp_path, clock)
@@ -197,6 +200,38 @@ def test_select_placement_skips_expired_lease(tmp_path, clock, node_key):
     assert svc.select_placement(10 * g) is None
 
 
+def test_placement_prefers_artifact_locality(tmp_path, clock, node_key):
+    g = 1024**3
+    svc = make_service(tmp_path, clock)
+    holder = _ready_node(svc, node_key, {"allocatable_vram_bytes": {"GPU-0": 20 * g},
+                                         "held_artifacts": ["sha256:aa"]})
+    _ready_node(svc, node_key, {"allocatable_vram_bytes": {"GPU-0": 40 * g}})  # more headroom, no artifact
+    pick = svc.select_placement(10 * g, artifact_digest="sha256:aa")
+    assert pick["node_id"] == holder  # locality beats headroom
+
+
+def test_placement_load_tiebreak(tmp_path, clock, node_key):
+    g = 1024**3
+    svc = make_service(tmp_path, clock)
+    _ready_node(svc, node_key, {"allocatable_vram_bytes": {"GPU-0": 40 * g}, "running_engines": 3})
+    idle = _ready_node(svc, node_key, {"allocatable_vram_bytes": {"GPU-0": 40 * g}, "running_engines": 0})
+    assert svc.select_placement(10 * g)["node_id"] == idle  # equal headroom -> lower load
+
+
+def test_placement_multi_gpu_requires_nvlink_group(tmp_path, clock, node_key):
+    g = 1024**3
+    svc = make_service(tmp_path, clock)
+    affine = _ready_node(svc, node_key, {
+        "allocatable_vram_bytes": {"GPU-0": 10 * g, "GPU-1": 10 * g},
+        "nvlink_groups": {"GPU-0": "nv0", "GPU-1": "nv0"}})
+    _ready_node(svc, node_key, {  # two devices, but split across groups
+        "allocatable_vram_bytes": {"GPU-0": 10 * g, "GPU-1": 10 * g},
+        "nvlink_groups": {"GPU-0": "a", "GPU-1": "b"}})
+    pick = svc.select_placement(5 * g, count=2)
+    assert pick["node_id"] == affine
+    assert len(pick["device_ids"]) == 2
+
+
 def test_rotate_certificate_rejects_revoked_node(tmp_path, clock, node_key):
     svc = make_service(tmp_path, clock)
     node_id = _enroll_auto(svc, node_key)["node_id"]