Kaynağa Gözat

feat(fabric): hard anti-OOM via declared ModuleManifest.resources.vram_mb + pytest wiring

- ModuleStatus gains required_vram_mb; systemd_wrap.get_status() populates it
  from manifest.resources.vram_mb ("auto"/unset -> 0). Flows manifest ->
  ModuleStatus -> heartbeat -> registry -> control plane.
- /place anti-OOM now rejects on the DECLARED requirement (503 when the chosen
  node's free VRAM < required), falling back to last observed footprint when
  undeclared. Live declared values: ltx=102400, tts=44000, stt=26624 mb, etc.
- tests/test_vram_capacity.py: +2 tests (required-vram manifest plumbing +
  registry carry); 7/7 pass.
- e2e-test.sh: TEST 0 runs `pytest tests/` (pytest installed in fabric venv;
  standalone fallback) + TEST 3b asserts VRAM capacity reporting.

Additive and backward-compatible. Deployed fleet-wide (12 nodes; schemas.py +
systemd_wrap.py to agents, schemas.py + control/main.py to the control plane).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Claude 1 ay önce
ebeveyn
işleme
a55f68f0bc
5 değiştirilmiş dosya ile 258 ekleme ve 13 silme
  1. 134 0
      agent/systemd_wrap.py
  2. 8 6
      control/main.py
  3. 91 0
      e2e-test.sh
  4. 1 0
      shared/schemas.py
  5. 24 7
      tests/test_vram_capacity.py

+ 134 - 0
agent/systemd_wrap.py

@@ -0,0 +1,134 @@
+"""systemd-wrap: control existing systemd services as Fabric modules."""
+from __future__ import annotations
+import logging
+import subprocess
+import time
+from pathlib import Path
+
+import httpx
+
+from fabric.shared.schemas import ModuleManifest, ModuleState, ModuleStatus
+
+log = logging.getLogger("fabric.agent.systemd_wrap")
+
+
+class SystemdWrapRuntime:
+    def __init__(self, manifest: ModuleManifest):
+        self.manifest = manifest
+        self.unit = manifest.systemd.systemd_unit if manifest.systemd else None
+        self.health_url = manifest.systemd.health_url if manifest.systemd else None
+        self.infer_url = manifest.systemd.infer_url if manifest.systemd else ""
+        self._last_infer_at: float | None = None
+        self._health_ok = False
+        self._state = ModuleState.unloaded
+        self._vram_mb = 0
+
+    def _systemctl(self, action: str) -> tuple[int, str]:
+        if not self.unit:
+            return 1, "no systemd unit configured"
+        try:
+            r = subprocess.run(["systemctl", action, self.unit], capture_output=True, text=True, timeout=30)
+            return r.returncode, r.stdout + r.stderr
+        except Exception as e:
+            return 1, str(e)
+
+    def is_active(self) -> bool:
+        code, out = self._systemctl("is-active")
+        return code == 0 and "active" in out
+
+    async def check_health(self) -> bool:
+        if not self.health_url:
+            return self.is_active()
+        try:
+            async with httpx.AsyncClient(timeout=5) as c:
+                r = await c.get(self.health_url)
+                self._health_ok = r.status_code < 500
+                return self._health_ok
+        except Exception:
+            self._health_ok = False
+            return False
+
+    def start(self) -> bool:
+        if self.is_active():
+            return True
+        log.info("Starting service %s", self.unit)
+        code, out = self._systemctl("start")
+        if code != 0:
+            log.error("Failed to start %s: %s", self.unit, out)
+        return code == 0
+
+    def stop(self) -> bool:
+        if not self.is_active():
+            return True
+        log.info("Stopping service %s", self.unit)
+        code, out = self._systemctl("stop")
+        if code != 0:
+            log.error("Failed to stop %s: %s", self.unit, out)
+        return code == 0
+
+    def get_main_pid(self) -> int | None:
+        if not self.unit:
+            return None
+        try:
+            r = subprocess.run(["systemctl", "show", "-p", "MainPID", "--value", self.unit], capture_output=True, text=True, timeout=5)
+            pid = int(r.stdout.strip())
+            return pid if pid > 0 else None
+        except Exception:
+            return None
+
+    def get_cgroup_pids(self) -> list[int]:
+        """Return all PIDs in the unit's cgroup (main + workers)."""
+        if not self.unit:
+            return []
+        try:
+            r = subprocess.run(["systemctl", "show", "-p", "ControlGroup", "--value", self.unit], capture_output=True, text=True, timeout=5)
+            cg = r.stdout.strip()
+            if not cg:
+                return []
+            procs = Path(f"/sys/fs/cgroup{cg}/cgroup.procs")
+            if not procs.exists():
+                return []
+            return [int(p) for p in procs.read_text().split() if p.strip()]
+        except Exception:
+            return []
+
+    def update_state(self, per_pid_vram: dict[int, int] | None = None):
+        if not self.is_active():
+            self._state = ModuleState.unloaded
+            self._vram_mb = 0
+            return
+        if per_pid_vram:
+            cgroup_pids = self.get_cgroup_pids()
+            total = sum(per_pid_vram.get(pid, 0) for pid in cgroup_pids)
+            if total > 0:
+                self._vram_mb = total
+            elif not cgroup_pids:
+                # cgroup unreadable — try MainPID as last resort
+                pid = self.get_main_pid()
+                self._vram_mb = per_pid_vram.get(pid, 0) if pid else 0
+            else:
+                self._vram_mb = 0
+        else:
+            self._vram_mb = 0
+        self._state = ModuleState.loaded
+        self._health_ok = True
+
+    def record_infer(self):
+        self._last_infer_at = time.time()
+
+    def _required_vram_mb(self) -> int:
+        """Declared manifest VRAM requirement as int (0 when "auto"/unset)."""
+        req = self.manifest.resources.vram_mb
+        return req if isinstance(req, int) else 0
+
+    def get_status(self) -> ModuleStatus:
+        return ModuleStatus(
+            name=self.manifest.name, state=self._state, vram_mb=self._vram_mb,
+            required_vram_mb=self._required_vram_mb(),
+            last_infer_at=self._last_infer_at, health_ok=self._health_ok,
+            infer_url=self.infer_url, modality=self.manifest.modality,
+            aliases=self.manifest.aliases, priority_class=self.manifest.scheduling.priority_class,
+            reclaim=self.manifest.lifecycle.reclaim,
+            idle_evict_s=self.manifest.lifecycle.idle_evict_s,
+            warm_default=self.manifest.lifecycle.warm_default,
+        )

+ 8 - 6
control/main.py

@@ -82,18 +82,20 @@ async def place(req: PlaceRequest):
     if result is None:
         raise HTTPException(status_code=404, detail=f"Model '{req.model}' not found in any node")
     if result.state in ("unloaded", "error"):
-        # Anti-OOM guard: refuse to (re)load onto a node that lacks the headroom
-        # the model is known to need. Footprint estimated from the model's last
-        # observed VRAM usage; skipped when unknown (0) or capacity data absent.
+        # Anti-OOM guard: refuse to (re)load when the chosen node lacks the VRAM
+        # the model needs. Prefer the DECLARED manifest requirement; fall back to
+        # the last observed footprint. Skipped when requirement/capacity unknown.
         candidates = registry.find_module(req.model)
-        footprint = max((m.vram_mb for (_n, m) in candidates), default=0)
+        required = max((m.required_vram_mb for (_n, m) in candidates), default=0)
+        if required <= 0:
+            required = max((m.vram_mb for (_n, m) in candidates), default=0)
         chosen = registry.nodes.get(result.node_id)
         free = (chosen.vram_free_mb if chosen else 0) or 0
-        if footprint > 0 and free > 0 and free < footprint:
+        if required > 0 and free > 0 and free < required:
             raise HTTPException(
                 status_code=503,
                 detail=f"Insufficient free VRAM on '{result.node_id}' for '{req.model}': "
-                       f"needs ~{footprint}mb, free {free}mb",
+                       f"needs ~{required}mb, free {free}mb",
             )
         await scheduler.ensure_loaded(result.node_id, result.module_name)
     return result

+ 91 - 0
e2e-test.sh

@@ -0,0 +1,91 @@
+#!/bin/bash
+exec 2>&1
+echo "===== FABRIC EPIC i-zmaudd54 — FINAL E2E TEST SUITE ====="
+echo
+
+echo "=== TEST 0: Unit tests (schemas / agent / registry / vram-capacity) ==="
+if ( cd /srv/fabric && PYTHONPATH=/srv venv/bin/python -m pytest tests/ -q ); then
+  echo "UNIT TESTS: PASS"
+else
+  echo "UNIT TESTS: FAIL (falling back to standalone runner)"
+  ( cd /srv/fabric && PYTHONPATH=/srv venv/bin/python tests/test_vram_capacity.py )
+fi
+
+echo
+echo "=== TEST 1: Control plane health (v0) ==="
+curl -s http://127.0.0.1:7700/health | python3 -m json.tool
+
+echo
+echo "=== TEST 2: Multi-node registry ==="
+curl -s http://127.0.0.1:7700/registry/nodes | python3 -c "import sys,json; d=json.load(sys.stdin); print(f'Total nodes: {len(d[\"nodes\"])}'); [print(f'  {n[\"node_id\"]:12s} {n[\"node_ip\"]:15s} hyp={n[\"hypervisor\"]:5s} mods={len(n[\"modules\"])} alive={n[\"alive\"]}') for n in d['nodes']]"
+
+echo
+echo "=== TEST 3: Module discovery ==="
+curl -s http://127.0.0.1:7700/registry/models | python3 -c "import sys,json; d=json.load(sys.stdin); print(f'Modules: {len(d[\"models\"])} across {d[\"nodes_count\"]} nodes, total VRAM {d[\"total_gpu_vram_mb\"]}MB'); [print(f'  {m[\"name\"]:18s} state={m[\"state\"]:10s} vram={m[\"vram_mb\"]:>6}MB modality={m[\"modality\"]:10s}') for m in d['models']]"
+
+echo
+echo "=== TEST 3b: VRAM capacity reporting (free/total) ==="
+curl -s http://127.0.0.1:7700/registry/models | python3 -c "import sys,json; d=json.load(sys.stdin); print(f'capacity={d.get(\"total_gpu_capacity_mb\")}MB free={d.get(\"total_gpu_free_mb\")}MB used={d.get(\"total_gpu_vram_mb\")}MB'); assert d.get('total_gpu_capacity_mb',0) > 0, 'no capacity reported'; print('  capacity OK')"
+
+echo
+echo "=== TEST 4: Multi-hypervisor (v2.1) ==="
+curl -s http://127.0.0.1:7700/registry/hypervisors | python3 -c "import sys,json; d=json.load(sys.stdin); print(f'Hypervisors: {d[\"total_hypervisors\"]}'); [print(f'  hyp={hyp}: {len(nodes)} nodes') for hyp, nodes in d['hypervisors'].items()]"
+
+echo
+echo "=== TEST 5: Placement & alias resolution ==="
+PASS=0
+FAIL=0
+for m in echo-canary echo ollama qwen acestep music tts comfyui sd ltx lipsync stt whisper trellis 3d dreamstack seedream audioface dirty pipecat voice; do
+  result=$(curl -s -X POST http://127.0.0.1:7700/place -H 'Content-Type: application/json' -d "{\"model\":\"$m\"}")
+  if echo "$result" | grep -q "node_id"; then
+    node=$(echo "$result" | python3 -c "import sys,json; d=json.load(sys.stdin); print(f'{d[\"module_name\"]}@{d[\"node_id\"]}')")
+    echo "  $m -> $node OK"
+    PASS=$((PASS+1))
+  else
+    echo "  $m -> NOT FOUND FAIL"
+    FAIL=$((FAIL+1))
+  fi
+done
+echo "Placement: $PASS passed, $FAIL failed"
+
+echo
+echo "=== TEST 6: Batch fan-out (v1.5) ==="
+JOB=$(curl -s -X POST "http://127.0.0.1:7700/batch?model=echo-canary&priority=normal" -H 'Content-Type: application/json' -d '[{"i":1},{"i":2},{"i":3},{"i":4},{"i":5},{"i":6},{"i":7},{"i":8},{"i":9},{"i":10},{"i":11},{"i":12},{"i":13},{"i":14},{"i":15},{"i":16}]' | python3 -c "import sys,json; print(json.load(sys.stdin)['job_id'])")
+sleep 2
+curl -s "http://127.0.0.1:7700/batch/$JOB" | python3 -m json.tool
+
+echo
+echo "=== TEST 7: Recipes (v2.4) ==="
+curl -s http://127.0.0.1:7700/install/recipes | python3 -m json.tool
+
+echo
+echo "=== TEST 8: Installer jobs (v2.2) ==="
+curl -s http://127.0.0.1:7700/install | python3 -c "import sys,json; d=json.load(sys.stdin); print(f'Total install jobs: {len(d[\"jobs\"])}'); [print(f'  {j[\"job_id\"]} {j[\"name\"]:15s} state={j[\"state\"]:10s} steps={len(j[\"steps\"])}') for j in d['jobs']]"
+
+echo
+echo "=== TEST 9: CUDA-checkpoint availability (v2.3) ==="
+curl -s http://127.0.0.1:7700/checkpoint/availability | python3 -m json.tool
+
+echo
+echo "=== TEST 10: Scheduler utilization (v1) ==="
+curl -s http://127.0.0.1:7700/scheduler/utilization | python3 -m json.tool
+
+echo
+echo "=== TEST 11: Per-PID NVML accounting (red-team mitigation) ==="
+curl -s http://10.0.2.162:7701/gpu-status | python3 -m json.tool
+
+echo
+echo "=== TEST 12: Module lifecycle (load/unload) ==="
+echo "BEFORE:"
+curl -s http://10.0.2.162:7701/modules | python3 -c "import sys,json; m=[x for x in json.load(sys.stdin)['modules'] if x['name']=='echo-canary'][0]; print(f'  echo-canary state={m[\"state\"]}')"
+echo "Unloading..."
+curl -s -X POST http://10.0.2.162:7701/unload/echo-canary | python3 -m json.tool
+sleep 1
+curl -s http://10.0.2.162:7701/modules | python3 -c "import sys,json; m=[x for x in json.load(sys.stdin)['modules'] if x['name']=='echo-canary'][0]; print(f'  AFTER unload: state={m[\"state\"]}')"
+echo "Loading..."
+curl -s -X POST http://10.0.2.162:7701/load/echo-canary | python3 -m json.tool
+sleep 2
+curl -s http://10.0.2.162:7701/modules | python3 -c "import sys,json; m=[x for x in json.load(sys.stdin)['modules'] if x['name']=='echo-canary'][0]; print(f'  AFTER reload: state={m[\"state\"]}')"
+
+echo
+echo "===== ALL TESTS COMPLETE ====="

+ 1 - 0
shared/schemas.py

@@ -107,6 +107,7 @@ class ModuleStatus(BaseModel):
     name: str
     state: ModuleState
     vram_mb: int = 0
+    required_vram_mb: int = 0  # declared manifest requirement (0 = auto/unknown)
     last_infer_at: float | None = None
     health_ok: bool = False
     infer_url: str = ""

+ 24 - 7
tests/test_vram_capacity.py

@@ -1,8 +1,9 @@
-"""Tests for additive GPU VRAM capacity + telemetry reporting.
+"""Tests for additive GPU VRAM capacity + telemetry + capacity-aware placement.
 
 Covers: nvidia-smi parser (capacity/free/telemetry, N/A + failure robustness),
-heartbeat round-trip, registry persistence + aggregation, and the capacity-aware
-placement tie-break. Runs under pytest OR standalone: `python tests/test_vram_capacity.py`.
+heartbeat round-trip, registry persistence + aggregation, capacity-aware placement
+tie-break, and the declared-requirement (required_vram_mb) plumbing for anti-OOM.
+Runs under pytest OR standalone: `python tests/test_vram_capacity.py`.
 """
 import sys
 from unittest import mock
@@ -11,8 +12,10 @@ sys.path.insert(0, "/srv")  # make `import fabric.*` resolve when run directly
 
 from fabric.shared.schemas import (
     Heartbeat, GpuMemInfo, ModuleStatus, ModuleState, PlaceRequest,
+    ModuleManifest, GpuResources,
 )
 from fabric.control.registry import Registry
+from fabric.agent.systemd_wrap import SystemdWrapRuntime
 import fabric.agent.main as agent
 
 
@@ -49,12 +52,10 @@ def test_get_gpu_info_handles_na_and_failure():
     g = info["gpus"][0]
     assert g["mem_free_mb"] == 24576 - 100
     assert g["util_pct"] is None and g["temp_c"] is None and g["power_w"] is None
-    # non-zero return code -> zeros, no crash
     with mock.patch.object(agent.subprocess, "run",
                            lambda *a, **k: mock.Mock(returncode=9, stdout="")):
         empty = agent._get_gpu_info()
     assert empty["gpus"] == [] and empty["vram_total_mb"] == 0
-    # nvidia-smi missing entirely -> zeros, no crash
     with mock.patch.object(agent.subprocess, "run",
                            mock.Mock(side_effect=FileNotFoundError())):
         empty2 = agent._get_gpu_info()
@@ -70,7 +71,7 @@ def test_heartbeat_roundtrip_carries_capacity():
     d = hb.model_dump()
     assert d["vram_total_mb"] == 49140 and d["vram_free_mb"] == 49000
     assert d["gpus"][0]["mem_free_mb"] == 49000
-    hb2 = Heartbeat(**d)  # what the CP does when parsing the request body
+    hb2 = Heartbeat(**d)
     assert hb2.vram_total_mb == 49140
     assert isinstance(hb2.gpus[0], GpuMemInfo)
 
@@ -95,7 +96,23 @@ def test_place_prefers_node_with_more_free_vram():
     reg.handle_heartbeat(Heartbeat(node_id="high", node_ip="10.0.0.2",
                                    modules=[mod], vram_total_mb=49140, vram_free_mb=40000))
     res = reg.place(PlaceRequest(model="m"))
-    assert res is not None and res.node_id == "high"  # equal state -> most-free wins
+    assert res is not None and res.node_id == "high"
+
+
+def test_systemd_wrap_reports_required_vram_from_manifest():
+    man = ModuleManifest(name="big", resources=GpuResources(gpus=1, vram_mb=20480))
+    assert SystemdWrapRuntime(man).get_status().required_vram_mb == 20480
+    # default resources -> vram_mb="auto" -> reported as 0 (unknown)
+    assert SystemdWrapRuntime(ModuleManifest(name="auto")).get_status().required_vram_mb == 0
+
+
+def test_registry_module_carries_required_vram():
+    reg = Registry()
+    mod = ModuleStatus(name="m", state=ModuleState.unloaded, required_vram_mb=20480)
+    reg.handle_heartbeat(Heartbeat(node_id="n", node_ip="1.1.1.1",
+                                   modules=[mod], vram_total_mb=49140, vram_free_mb=5000))
+    cands = reg.find_module("m")
+    assert cands and cands[0][1].required_vram_mb == 20480
 
 
 if __name__ == "__main__":