diff --git a/meridian_control/routes.py b/meridian_control/routes.py index d79cceb..1d8e304 100644 --- a/meridian_control/routes.py +++ b/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} diff --git a/meridian_control/service.py b/meridian_control/service.py index 021b113..6a7557e 100644 --- a/meridian_control/service.py +++ b/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 diff --git a/tests/control/test_service.py b/tests/control/test_service.py index a979f37..b0ea4bd 100644 --- a/tests/control/test_service.py +++ b/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"]