Initial public ModelForge release
This commit is contained in:
Jens
@@ -0,0 +1,339 @@
|
||||
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)
|
||||
Reference in New Issue
Block a user