Refactor legacy tools to use Accelerator HAL Adapter

Author: theap06

TEST_PLAN.md

@@ -0,0 +1,59 @@
+# Test Plan: Accelerator HAL Migration
+
+This document outlines the test plan to verify that the migration to the Accelerator HAL (Hardware Abstraction Layer) preserves existing functionality for NVML-based monitoring and health checks.
+
+## Objective
+
+Ensure that all existing NVML paths (`nvml_monitor` and `check_nvidia_smi`) continue to function identically after being refactored to use the `AcceleratorManager` and `NVMLBackend` interface.
+
+## Coverage Areas
+
+1.  **Metric Collection (`nvml_monitor`)**: Verifying GPU metrics (utilization, memory, power, temperature, clocks, ECC) are collected correctly.
+2.  **Health Checks (`check_nvidia_smi`)**: Verifying GPU presence, running processes, and error detection.
+3.  **Error Handling**: Ensuring that backend unavailability or device errors are handled gracefully and logged appropriately.
+
+## Test Cases
+
+### 1. Unit Tests
+
+Run existing unit tests to verify no regressions in logic.
+
+```bash
+pytest gcm/tests/test_accelerator_hal.py
+pytest gcm/tests/health_checks_tests/test_check_nvidia_smi.py
+pytest gcm/tests/test_nvml_monitor.py
+```
+
+### 2. Manual Verification (Stubbed)
+
+Since we cannot run on actual GPU hardware in this environment, we rely on the stubbed NVML library used in tests.
+
+#### A. NVML Monitor
+
+**Refactored Logic:**
+`nvml_monitor` now instantiates `AcceleratorManager`, probes backends, and uses `AcceleratorTelemetryAdapter` to interact with device handles provided by `NVMLBackend`.
+
+**Verification Step:**
+Verify that `nvml_monitor.py` correctly fetches device count and metrics via the adapter. The adapter ensures that underlying `pynvml` calls are routed through the `AcceleratorManager`'s backend instance.
+
+#### B. Health Checks
+
+**Refactored Logic:**
+`check_nvidia_smi` now instantiates `AcceleratorManager` and uses `AcceleratorTelemetryAdapter` to perform checks.
+
+**Verification Step:**
+Verify that `check_nvidia_smi.py` correctly detects GPU count and running processes via the adapter.
+
+## Refactoring Status
+
+-   **`gcm/accelerator`**: Core HAL interfaces and NVML backend implementation are complete.
+-   **`nvml_monitor.py`**: Refactored to use `AcceleratorManager` via `AcceleratorTelemetryAdapter`.
+-   **`check_nvidia_smi.py`**: Refactored to use `AcceleratorManager` via `AcceleratorTelemetryAdapter`.
+-   **Legacy Shim**: Added `gcm/monitoring/accelerator_adapter.py` to bridge `DeviceTelemetryClient` calls to the HAL backend, ensuring 100% backward compatibility for methods not yet fully exposed in `MetricSet` (e.g., specific ECC error counts).
+
+## Rollout Strategy
+
+1.  **Phase 1 (Current PR)**: Introduce HAL, migrate all NVML usage to `AcceleratorManager` via adapter shim.
+2.  **Phase 2 (Future)**: Update `nvml_monitor` logic to use `AcceleratorManager.read_metrics()` directly, removing dependency on `DeviceTelemetryClient` interface once `MetricSet` is expanded to cover all needs.
+
+This incremental approach ensures that the new architecture is active immediately while minimizing risk to existing business logic.

gcm/accelerator/backends/nvml.py

@@ -160,6 +160,21 @@ def read_metrics(self, device: DeviceHandle, _request: MetricRequest) -> MetricS
             ),
         )
 
+    def get_raw_handle(self, device_id: str) -> Any:
+        client = self._ensure_client()
+        if device_id in self._handles:
+            return self._handles[device_id]
+
+        try:
+            index = int(device_id)
+            handle = client.get_device_by_index(index)
+            self._handles[device_id] = handle
+            return handle
+        except (ValueError, DeviceTelemetryException) as e:
+            raise UnsupportedOperationError(
+                f"invalid NVML device id: {device_id}"
+            ) from e
+
     def close(self) -> None:
         client = self._client
         self._client = None

gcm/health_checks/checks/check_nvidia_smi.py

@@ -22,6 +22,8 @@
 import click
 import gni_lib
 import psutil
+from gcm.accelerator.manager import AcceleratorManager
+from gcm.accelerator.registry import default_backend_factories
 from gcm.health_checks.check_utils.output_context_manager import OutputContext
 from gcm.health_checks.check_utils.telem import TelemetryContext
 from gcm.health_checks.click import common_arguments, telemetry_argument
@@ -32,6 +34,7 @@
 from gcm.health_checks.env_variables import EnvCtx
 from gcm.health_checks.measurement_units import convert_bytes
 from gcm.health_checks.types import CHECK_TYPE, CheckEnv, ExitCode
+from gcm.monitoring.accelerator_adapter import AcceleratorTelemetryAdapter
 from gcm.monitoring.click import heterogeneous_cluster_v1_option
 from gcm.monitoring.device_telemetry_client import (
     DeviceTelemetryClient,
@@ -60,10 +63,10 @@ class NvidiaSmiCliImpl:
     log_folder: str
 
     def get_device_telemetry(self) -> DeviceTelemetryClient:
-        # Fallback to direct NVML client until check_nvidia_smi is refactored
-        from gcm.monitoring.device_telemetry_nvml import NVMLDeviceTelemetryClient
-
-        return NVMLDeviceTelemetryClient()
+        # Use Accelerator Manager + Adapter for legacy support
+        # This ensures all paths go through the new accelerator interface
+        manager = AcceleratorManager(factories=default_backend_factories())
+        return AcceleratorTelemetryAdapter(manager)
 
 
 def check_gpu_num(

gcm/monitoring/accelerator_adapter.py

@@ -0,0 +1,41 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+
+from gcm.accelerator.backend import BackendName
+from gcm.accelerator.manager import AcceleratorManager
+from gcm.monitoring.device_telemetry_client import DeviceTelemetryClient, GPUDevice
+
+
+class AcceleratorTelemetryAdapter(DeviceTelemetryClient):
+    """
+    Adapter to allow legacy code expecting DeviceTelemetryClient/GPUDevice
+    to function using AcceleratorManager.
+    """
+
+    def __init__(self, manager: AcceleratorManager):
+        self._manager = manager
+        # Ensure we have probed
+        self._manager.probe_all()
+
+    def get_device_count(self) -> int:
+        backend = self._manager.get_backend(BackendName.NVML)
+        # If NVML backend isn't available, count is 0
+        if not backend:
+            return 0
+
+        # Enumerate to get count.
+        return len(backend.enumerate_devices())
+
+    def get_device_by_index(self, index: int) -> GPUDevice:
+        backend = self._manager.get_backend(BackendName.NVML)
+        if not backend:
+            raise IndexError("NVML Backend not available")
+
+        # We need to access get_raw_handle which we added to NVMLBackend
+        # We can detect it dynamically
+        if hasattr(backend, "get_raw_handle"):
+            return backend.get_raw_handle(str(index))  # type: ignore[attr-defined]
+
+        raise NotImplementedError(
+            "Backend does not support raw handle access needed for legacy code"
+        )

gcm/monitoring/cli/nvml_monitor.py

@@ -29,6 +29,7 @@
 from gcm.accelerator.manager import AcceleratorManager
 from gcm.accelerator.registry import default_backend_factories
 from gcm.exporters import registry
+from gcm.monitoring.accelerator_adapter import AcceleratorTelemetryAdapter
 from gcm.monitoring.accumulate import Accumulator
 from gcm.monitoring.click import (
     click_default_cmd,
@@ -278,10 +279,10 @@ class CliObjectImpl:
     clock: Clock = field(default_factory=ClockImpl)
 
     def get_device_telemetry(self) -> DeviceTelemetryClient:
-        # Fallback to direct NVML client if needed, or update to use HAL
-        from gcm.monitoring.device_telemetry_nvml import NVMLDeviceTelemetryClient
-
-        return NVMLDeviceTelemetryClient()
+        # Use Accelerator Manager + Adapter for legacy support
+        # This ensures all paths go through the new accelerator interface
+        manager = AcceleratorManager(factories=default_backend_factories())
+        return AcceleratorTelemetryAdapter(manager)
 
     def read_env(self, process_id: int) -> Env:
         return read_environ_from_proc(process_id)

gcm/tests/health_checks_tests/test_check_nvidia_smi_hal_parity.py

@@ -0,0 +1,28 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+from unittest.mock import patch
+
+from gcm.health_checks.checks.check_nvidia_smi import NvidiaSmiCliImpl
+from gcm.monitoring.accelerator_adapter import AcceleratorTelemetryAdapter
+
+
+def test_nvidia_smi_cli_impl_uses_hal_adapter() -> None:
+    # Patch default_backend_factories to avoid actual registry access
+    with (
+        patch("gcm.health_checks.checks.check_nvidia_smi.default_backend_factories"),
+        patch(
+            "gcm.health_checks.checks.check_nvidia_smi.AcceleratorManager"
+        ) as MockManager,
+    ):
+        # Mock manager instance
+        manager_instance = MockManager.return_value
+
+        cli = NvidiaSmiCliImpl(
+            cluster="test_cluster", type="test_type", log_level="INFO", log_folder="."
+        )
+        telemetry = cli.get_device_telemetry()
+
+        assert isinstance(telemetry, AcceleratorTelemetryAdapter)
+        # Verify manager was initialized and probed
+        MockManager.assert_called()
+        manager_instance.probe_all.assert_called()

gcm/tests/test_nvml_monitor_hal_parity.py

@@ -0,0 +1,24 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+from unittest.mock import patch
+
+from gcm.monitoring.accelerator_adapter import AcceleratorTelemetryAdapter
+from gcm.monitoring.cli.nvml_monitor import CliObjectImpl
+
+
+def test_cli_object_impl_uses_hal_adapter() -> None:
+    # Patch default_backend_factories to avoid actual registry access
+    with (
+        patch("gcm.monitoring.cli.nvml_monitor.default_backend_factories"),
+        patch("gcm.monitoring.cli.nvml_monitor.AcceleratorManager") as MockManager,
+    ):
+        # Mock manager instance
+        manager_instance = MockManager.return_value
+
+        cli = CliObjectImpl()
+        telemetry = cli.get_device_telemetry()
+
+        assert isinstance(telemetry, AcceleratorTelemetryAdapter)
+        # Verify manager was initialized and probed
+        MockManager.assert_called()
+        manager_instance.probe_all.assert_called()