ModelForge/backend/tests/hardware_fakes.py

from __future__ import annotations

from dataclasses import dataclass
from pathlib import Path

from modelforge_api.domain.enums import Availability
from modelforge_api.domain.hardware import (
    AcceleratorInventory,
    AcceleratorTelemetry,
    HostInventory,
    NvidiaCollection,
    ObservedValue,
    StorageObservation,
)


class FakeHostCollector:
    def __init__(self, identity: str = "node-identity") -> None:
        self.identity = identity

    def collect(self) -> HostInventory:
        return HostInventory(
            identity_key=self.identity,
            identity_source="test",
            hostname="forge-host",
            display_name="Forge Host",
            os_name="TestOS",
            os_version=ObservedValue.known("1"),
            architecture="x86_64",
            kernel_version=ObservedValue.known("1.0"),
            cpu_model=ObservedValue.known("Test CPU"),
            logical_cpu_count=ObservedValue.known(16),
            physical_core_count=ObservedValue.known(8),
            total_ram_bytes=ObservedValue.known(64 * 1024**3),
            available_ram_bytes=ObservedValue.known(32 * 1024**3),
            agent_version="0.1.0",
            storage=[
                StorageObservation(
                    purpose="artifacts",
                    path=str(Path("/artifacts")),
                    total_bytes=ObservedValue.known(1000),
                    used_bytes=ObservedValue.known(400),
                    free_bytes=ObservedValue.known(600),
                )
            ],
        )


def accelerator(device_uuid: str = "GPU-A", name: str = "Fake GPU") -> AcceleratorInventory:
    return AcceleratorInventory(
        device_index=0,
        device_uuid=device_uuid,
        pci_bus_id=ObservedValue.known("0000:01:00.0"),
        name=name,
        architecture=ObservedValue.known("ada"),
        compute_capability_major=ObservedValue.known(8),
        compute_capability_minor=ObservedValue.known(9),
        total_vram_bytes=ObservedValue.known(16 * 1024**3),
        driver_version=ObservedValue.known("600.1"),
        cuda_driver_version=ObservedValue.known("13.0"),
        mig_mode_current=ObservedValue.absent(Availability.UNSUPPORTED),
    )


def telemetry(device_uuid: str = "GPU-A", utilization: int = 25) -> AcceleratorTelemetry:
    return AcceleratorTelemetry(
        device_uuid=device_uuid,
        used_vram_bytes=ObservedValue.known(2 * 1024**3),
        free_vram_bytes=ObservedValue.known(14 * 1024**3),
        gpu_utilization_percent=ObservedValue.known(utilization),
        memory_utilization_percent=ObservedValue.known(10),
        temperature_c=ObservedValue.known(45),
        power_draw_w=ObservedValue.known(80.0),
        power_limit_w=ObservedValue.known(320.0),
        graphics_clock_mhz=ObservedValue.known(2000),
        memory_clock_mhz=ObservedValue.known(10000),
        fan_speed_percent=ObservedValue.absent(Availability.UNSUPPORTED),
        performance_state=ObservedValue.known("P2"),
    )


class FakeAcceleratorCollector:
    def __init__(
        self,
        devices: list[AcceleratorInventory] | None = None,
        availability: Availability = Availability.KNOWN,
        utilization: int = 25,
    ) -> None:
        self.devices = devices or []
        self.availability = availability
        self.utilization = utilization

    def collect(self) -> NvidiaCollection:
        return NvidiaCollection(
            availability=self.availability,
            reason="NVML unavailable"
            if self.availability is not Availability.KNOWN
            else ("no NVIDIA devices detected" if not self.devices else None),
            inventory=self.devices,
            telemetry=[telemetry(item.device_uuid, self.utilization) for item in self.devices],
        )


@dataclass
class Memory:
    total: int = 16 * 1024**3
    used: int = 2 * 1024**3
    free: int = 14 * 1024**3


@dataclass
class Utilization:
    gpu: int = 25
    memory: int = 10


@dataclass
class Pci:
    busId: bytes = b"0000:01:00.0"


class FakeNvml:
    class NVMLError(Exception):
        pass

    class NVMLError_NotSupported(NVMLError):
        pass

    NVML_TEMPERATURE_GPU = 0
    NVML_CLOCK_GRAPHICS = 0
    NVML_CLOCK_MEM = 1
    NVML_DEVICE_ARCH_ADA = 7

    def __init__(
        self,
        count: int = 1,
        init_error: bool = False,
        unsupported_power: bool = False,
        device_error: int | None = None,
    ) -> None:
        self.count = count
        self.init_error = init_error
        self.unsupported_power = unsupported_power
        self.device_error = device_error
        self.shutdown_calls = 0

    def nvmlInit(self):
        if self.init_error:
            raise self.NVMLError("driver")

    def nvmlShutdown(self):
        self.shutdown_calls += 1

    def nvmlDeviceGetCount(self):
        return self.count

    def nvmlSystemGetDriverVersion(self):
        return b"600.1"

    def nvmlSystemGetCudaDriverVersion_v2(self):
        return 13000

    def nvmlDeviceGetHandleByIndex(self, index):
        if self.device_error == index:
            raise self.NVMLError("device")
        return index

    def nvmlDeviceGetUUID(self, handle):
        return f"GPU-{handle}".encode()

    def nvmlDeviceGetName(self, handle):
        return f"Fake GPU {handle}".encode()

    def nvmlDeviceGetMemoryInfo(self, handle):
        return Memory()

    def nvmlDeviceGetPciInfo(self, handle):
        return Pci(busId=f"0000:0{handle + 1}:00.0".encode())

    def nvmlDeviceGetCudaComputeCapability(self, handle):
        return (8, 9)

    def nvmlDeviceGetMigMode(self, handle):
        raise self.NVMLError_NotSupported()

    def nvmlDeviceGetArchitecture(self, handle):
        return self.NVML_DEVICE_ARCH_ADA

    def nvmlDeviceGetUtilizationRates(self, handle):
        return Utilization()

    def nvmlDeviceGetTemperature(self, handle, sensor):
        return 45

    def nvmlDeviceGetPowerUsage(self, handle):
        if self.unsupported_power:
            raise self.NVMLError_NotSupported()
        return 80000

    def nvmlDeviceGetEnforcedPowerLimit(self, handle):
        return 320000

    def nvmlDeviceGetClockInfo(self, handle, clock):
        return 2000

    def nvmlDeviceGetFanSpeed(self, handle):
        raise self.NVMLError_NotSupported()

    def nvmlDeviceGetPerformanceState(self, handle):
        return 2