340 lines
15 KiB
Python
340 lines
15 KiB
Python
from __future__ import annotations
|
|
|
|
import platform
|
|
import socket
|
|
import sys
|
|
import uuid
|
|
from collections.abc import Callable
|
|
from contextlib import suppress
|
|
from datetime import UTC, datetime
|
|
from pathlib import Path
|
|
from typing import Any, TypeVar
|
|
|
|
import psutil # type: ignore[import-untyped]
|
|
|
|
from modelforge_api import __version__
|
|
from modelforge_api.domain.enums import Availability
|
|
from modelforge_api.domain.hardware import (
|
|
AcceleratorInventory,
|
|
AcceleratorTelemetry,
|
|
HostInventory,
|
|
NvidiaCollection,
|
|
ObservedValue,
|
|
StorageObservation,
|
|
)
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
class NodeIdentityProvider:
|
|
"""Stable node identity: explicit override, OS machine ID, then persisted UUID."""
|
|
|
|
def __init__(
|
|
self,
|
|
identity_file: Path,
|
|
explicit_identity: str | None = None,
|
|
force_persisted: bool = False,
|
|
) -> None:
|
|
self.identity_file = identity_file
|
|
self.explicit_identity = explicit_identity
|
|
self.force_persisted = force_persisted
|
|
|
|
def resolve(self) -> tuple[str, str]:
|
|
if self.explicit_identity:
|
|
return self.explicit_identity, "configured"
|
|
if not self.force_persisted:
|
|
system_id = self._system_machine_id()
|
|
if system_id:
|
|
return str(uuid.uuid5(uuid.NAMESPACE_OID, system_id)), "os_machine_id"
|
|
try:
|
|
if self.identity_file.exists():
|
|
return self.identity_file.read_text(encoding="utf-8").strip(), "persisted_uuid"
|
|
value = str(uuid.uuid4())
|
|
self.identity_file.parent.mkdir(parents=True, exist_ok=True)
|
|
self.identity_file.write_text(value, encoding="utf-8")
|
|
return value, "persisted_uuid"
|
|
except OSError as exc:
|
|
raise RuntimeError("unable to establish stable node identity") from exc
|
|
|
|
@staticmethod
|
|
def _system_machine_id() -> str | None:
|
|
# `sys.platform`, not `os.name`: a type checker narrows on the former and not on the latter,
|
|
# so under `os.name` this block was analysed on Linux too — where winreg has no attributes.
|
|
# The images run on Linux, so the platform the release actually ships on was the one the
|
|
# type check could not see. They mean the same thing at runtime.
|
|
if sys.platform == "win32":
|
|
try:
|
|
import winreg
|
|
|
|
with winreg.OpenKey(
|
|
winreg.HKEY_LOCAL_MACHINE,
|
|
r"SOFTWARE\Microsoft\Cryptography",
|
|
) as key:
|
|
return str(winreg.QueryValueEx(key, "MachineGuid")[0])
|
|
except (OSError, ImportError):
|
|
return None
|
|
for path in (Path("/etc/machine-id"), Path("/var/lib/dbus/machine-id")):
|
|
try:
|
|
value = path.read_text(encoding="utf-8").strip()
|
|
if value:
|
|
return value
|
|
except OSError:
|
|
continue
|
|
return None
|
|
|
|
|
|
class SystemHostCollector:
|
|
def __init__(
|
|
self,
|
|
identity_provider: NodeIdentityProvider,
|
|
storage_paths: dict[str, Path],
|
|
) -> None:
|
|
self.identity_provider = identity_provider
|
|
self.storage_paths = storage_paths
|
|
|
|
@staticmethod
|
|
def _known_or_unknown(value: T | None, reason: str) -> ObservedValue[T]:
|
|
return (
|
|
ObservedValue.known(value)
|
|
if value is not None
|
|
else ObservedValue.absent(Availability.UNKNOWN, reason)
|
|
)
|
|
|
|
def collect(self) -> HostInventory:
|
|
identity_key, identity_source = self.identity_provider.resolve()
|
|
memory = psutil.virtual_memory()
|
|
cpu_model = platform.processor().strip() or None
|
|
physical = psutil.cpu_count(logical=False)
|
|
logical = psutil.cpu_count(logical=True)
|
|
storage: list[StorageObservation] = []
|
|
for purpose, path in self.storage_paths.items():
|
|
try:
|
|
usage = psutil.disk_usage(str(path))
|
|
storage.append(
|
|
StorageObservation(
|
|
purpose=purpose,
|
|
path=str(path.resolve()),
|
|
total_bytes=ObservedValue.known(usage.total),
|
|
used_bytes=ObservedValue.known(usage.used),
|
|
free_bytes=ObservedValue.known(usage.free),
|
|
)
|
|
)
|
|
except OSError as exc:
|
|
reason = f"{type(exc).__name__}: path unavailable"
|
|
unavailable = ObservedValue[int].absent(Availability.UNAVAILABLE, reason)
|
|
storage.append(
|
|
StorageObservation(
|
|
purpose=purpose,
|
|
path=str(path),
|
|
total_bytes=unavailable,
|
|
used_bytes=unavailable,
|
|
free_bytes=unavailable,
|
|
)
|
|
)
|
|
release = platform.release()
|
|
environment = "wsl2" if "microsoft" in release.lower() else "native"
|
|
return HostInventory(
|
|
identity_key=identity_key,
|
|
identity_source=identity_source,
|
|
hostname=socket.gethostname(),
|
|
display_name=socket.gethostname(),
|
|
os_name=platform.system() or "unknown",
|
|
os_version=self._known_or_unknown(platform.version() or None, "OS version unavailable"),
|
|
architecture=platform.machine() or "unknown",
|
|
kernel_version=self._known_or_unknown(release or None, "kernel unavailable"),
|
|
cpu_model=self._known_or_unknown(cpu_model, "CPU model unavailable"),
|
|
logical_cpu_count=self._known_or_unknown(logical, "logical CPU count unavailable"),
|
|
physical_core_count=self._known_or_unknown(physical, "physical core count unavailable"),
|
|
total_ram_bytes=ObservedValue.known(memory.total),
|
|
available_ram_bytes=ObservedValue.known(memory.available),
|
|
agent_version=__version__,
|
|
storage=storage,
|
|
metadata={"environment": environment},
|
|
)
|
|
|
|
|
|
class NvidiaNvmlCollector:
|
|
def __init__(self, api: Any) -> None:
|
|
self.api = api
|
|
self.nvml_error = getattr(api, "NVMLError", Exception)
|
|
self.not_supported = getattr(api, "NVMLError_NotSupported", ())
|
|
|
|
def _optional(
|
|
self, call: Callable[[], T], transform: Callable[[T], Any] | None = None
|
|
) -> ObservedValue[Any]:
|
|
try:
|
|
value = call()
|
|
return ObservedValue.known(transform(value) if transform else value)
|
|
except self.not_supported:
|
|
return ObservedValue.absent(Availability.UNSUPPORTED, "NVML metric not supported")
|
|
except self.nvml_error as exc:
|
|
return ObservedValue.absent(Availability.TEMPORARILY_FAILED, type(exc).__name__)
|
|
|
|
@staticmethod
|
|
def _text(value: Any) -> str:
|
|
return value.decode("utf-8", errors="replace") if isinstance(value, bytes) else str(value)
|
|
|
|
@staticmethod
|
|
def _cuda_version(value: int) -> str:
|
|
return f"{value // 1000}.{(value % 1000) // 10}"
|
|
|
|
def _architecture(self, handle: Any) -> ObservedValue[str]:
|
|
def transform(value: Any) -> str:
|
|
for name in dir(self.api):
|
|
if name.startswith("NVML_DEVICE_ARCH_") and getattr(self.api, name) == value:
|
|
return name.removeprefix("NVML_DEVICE_ARCH_").lower()
|
|
return f"nvml_arch_{value}"
|
|
|
|
function = getattr(self.api, "nvmlDeviceGetArchitecture", None)
|
|
if function is None:
|
|
return ObservedValue.absent(Availability.UNSUPPORTED, "architecture API unavailable")
|
|
return self._optional(lambda: function(handle), transform)
|
|
|
|
def collect(self) -> NvidiaCollection:
|
|
initialized = False
|
|
try:
|
|
self.api.nvmlInit()
|
|
initialized = True
|
|
count = self.api.nvmlDeviceGetCount()
|
|
driver = self._optional(self.api.nvmlSystemGetDriverVersion, self._text)
|
|
cuda_function = getattr(self.api, "nvmlSystemGetCudaDriverVersion_v2", None) or getattr(
|
|
self.api, "nvmlSystemGetCudaDriverVersion", None
|
|
)
|
|
cuda_driver = (
|
|
self._optional(cuda_function, self._cuda_version)
|
|
if cuda_function
|
|
else ObservedValue.absent(Availability.UNSUPPORTED, "CUDA driver API unavailable")
|
|
)
|
|
inventory: list[AcceleratorInventory] = []
|
|
telemetry: list[AcceleratorTelemetry] = []
|
|
failures: list[str] = []
|
|
observed_at = datetime.now(UTC)
|
|
for index in range(count):
|
|
try:
|
|
handle = self.api.nvmlDeviceGetHandleByIndex(index)
|
|
device_uuid = self._text(self.api.nvmlDeviceGetUUID(handle))
|
|
name = self._text(self.api.nvmlDeviceGetName(handle))
|
|
memory = self.api.nvmlDeviceGetMemoryInfo(handle)
|
|
pci = self._optional(
|
|
lambda: self.api.nvmlDeviceGetPciInfo(handle),
|
|
lambda value: self._text(value.busId),
|
|
)
|
|
compute = self._optional(
|
|
lambda: self.api.nvmlDeviceGetCudaComputeCapability(handle)
|
|
)
|
|
major = (
|
|
ObservedValue.known(int(compute.value[0]))
|
|
if compute.availability is Availability.KNOWN and compute.value is not None
|
|
else ObservedValue.absent(compute.availability, compute.reason)
|
|
)
|
|
minor = (
|
|
ObservedValue.known(int(compute.value[1]))
|
|
if compute.availability is Availability.KNOWN and compute.value is not None
|
|
else ObservedValue.absent(compute.availability, compute.reason)
|
|
)
|
|
mig_function = getattr(self.api, "nvmlDeviceGetMigMode", None)
|
|
mig = (
|
|
self._optional(lambda: mig_function(handle), lambda value: bool(value[0]))
|
|
if mig_function
|
|
else ObservedValue.absent(Availability.UNSUPPORTED, "MIG API unavailable")
|
|
)
|
|
inventory.append(
|
|
AcceleratorInventory(
|
|
device_index=index,
|
|
device_uuid=device_uuid,
|
|
pci_bus_id=pci,
|
|
name=name,
|
|
architecture=self._architecture(handle),
|
|
compute_capability_major=major,
|
|
compute_capability_minor=minor,
|
|
total_vram_bytes=ObservedValue.known(int(memory.total)),
|
|
driver_version=driver,
|
|
cuda_driver_version=cuda_driver,
|
|
mig_mode_current=mig,
|
|
inventory_at=observed_at,
|
|
)
|
|
)
|
|
utilization = self._optional(
|
|
lambda: self.api.nvmlDeviceGetUtilizationRates(handle)
|
|
)
|
|
gpu_util = (
|
|
ObservedValue.known(int(utilization.value.gpu))
|
|
if utilization.availability is Availability.KNOWN
|
|
and utilization.value is not None
|
|
else ObservedValue.absent(utilization.availability, utilization.reason)
|
|
)
|
|
mem_util = (
|
|
ObservedValue.known(int(utilization.value.memory))
|
|
if utilization.availability is Availability.KNOWN
|
|
and utilization.value is not None
|
|
else ObservedValue.absent(utilization.availability, utilization.reason)
|
|
)
|
|
telemetry.append(
|
|
AcceleratorTelemetry(
|
|
device_uuid=device_uuid,
|
|
observed_at=observed_at,
|
|
used_vram_bytes=ObservedValue.known(int(memory.used)),
|
|
free_vram_bytes=ObservedValue.known(int(memory.free)),
|
|
gpu_utilization_percent=gpu_util,
|
|
memory_utilization_percent=mem_util,
|
|
temperature_c=self._optional(
|
|
lambda: self.api.nvmlDeviceGetTemperature(
|
|
handle, self.api.NVML_TEMPERATURE_GPU
|
|
),
|
|
int,
|
|
),
|
|
power_draw_w=self._optional(
|
|
lambda: self.api.nvmlDeviceGetPowerUsage(handle),
|
|
lambda value: round(value / 1000, 3),
|
|
),
|
|
power_limit_w=self._optional(
|
|
lambda: self.api.nvmlDeviceGetEnforcedPowerLimit(handle),
|
|
lambda value: round(value / 1000, 3),
|
|
),
|
|
graphics_clock_mhz=self._optional(
|
|
lambda: self.api.nvmlDeviceGetClockInfo(
|
|
handle, self.api.NVML_CLOCK_GRAPHICS
|
|
),
|
|
int,
|
|
),
|
|
memory_clock_mhz=self._optional(
|
|
lambda: self.api.nvmlDeviceGetClockInfo(
|
|
handle, self.api.NVML_CLOCK_MEM
|
|
),
|
|
int,
|
|
),
|
|
fan_speed_percent=self._optional(
|
|
lambda: self.api.nvmlDeviceGetFanSpeed(handle), int
|
|
),
|
|
performance_state=self._optional(
|
|
lambda: self.api.nvmlDeviceGetPerformanceState(handle),
|
|
lambda value: f"P{value}",
|
|
),
|
|
)
|
|
)
|
|
except self.nvml_error as exc:
|
|
failures.append(f"device {index}: {type(exc).__name__}")
|
|
availability = Availability.KNOWN if not failures else Availability.TEMPORARILY_FAILED
|
|
return NvidiaCollection(
|
|
availability=availability,
|
|
reason="; ".join(failures)
|
|
or ("no NVIDIA devices detected" if count == 0 else None),
|
|
inventory=inventory,
|
|
telemetry=telemetry,
|
|
observed_at=observed_at,
|
|
)
|
|
except self.nvml_error as exc:
|
|
return NvidiaCollection(
|
|
availability=Availability.UNAVAILABLE, reason=type(exc).__name__
|
|
)
|
|
finally:
|
|
if initialized:
|
|
with suppress(self.nvml_error):
|
|
self.api.nvmlShutdown()
|
|
|
|
|
|
def build_nvml_collector() -> NvidiaNvmlCollector:
|
|
import pynvml # type: ignore[import-untyped]
|
|
|
|
return NvidiaNvmlCollector(pynvml)
|