|
|
@@ -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,
|
|
|
+ )
|