fabric/agent/systemd_wrap.py

"""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,
        )