Restore Rust installation step from main

theap06 · theap06 · commit ccc912980b47 · 2026-03-20T14:37:35.000-07:00
diff --git a/.github/workflows/gcm_python.yml b/.github/workflows/gcm_python.yml
@@ -64,6 +64,11 @@ jobs:
           path: ~/.cache/venv-ci
           key: ${{ env.pythonLocation }}-${{ hashFiles('dev-requirements.txt') }}
 
+      - name: Install Rust
+        run: |
+          curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y
+          cp -r $HOME/.cargo/bin/* $HOME/.cache/venv-ci/bin/
+
       - name: Install build dependencies
         run: |
           sudo apt update --yes
diff --git a/gcm/accelerator/__init__.py b/gcm/accelerator/__init__.py
@@ -1,33 +1,26 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
-from gcm.monitoring.accelerator.backend import (
+from gcm.accelerator.backend import (
     AcceleratorBackend,
     BackendName,
     DeviceHandle,
     ProbeResult,
 )
-from gcm.monitoring.accelerator.errors import (
+from gcm.accelerator.errors import (
     AcceleratorError,
     BackendUnavailableError,
     UnsupportedOperationError,
 )
-from gcm.monitoring.accelerator.manager import AcceleratorManager
-from gcm.monitoring.accelerator.metrics import (
-    Capability,
-    CapabilitySet,
-    MetricRequest,
-    MetricSet,
-)
-from gcm.monitoring.accelerator.registry import default_backend_factories
+from gcm.accelerator.manager import AcceleratorManager
+from gcm.accelerator.metrics import MetricRequest, MetricSet
+from gcm.accelerator.registry import default_backend_factories
 
 __all__ = [
     "AcceleratorBackend",
     "AcceleratorError",
     "AcceleratorManager",
     "BackendName",
     "BackendUnavailableError",
-    "Capability",
-    "CapabilitySet",
     "DeviceHandle",
     "MetricRequest",
     "MetricSet",
diff --git a/gcm/accelerator/backend.py b/gcm/accelerator/backend.py
@@ -5,7 +5,7 @@
 from enum import Enum
 from typing import Callable, List, Protocol
 
-from gcm.monitoring.accelerator.metrics import CapabilitySet, MetricRequest, MetricSet
+from gcm.accelerator.metrics import MetricRequest, MetricSet
 
 
 class BackendName(str, Enum):
@@ -39,8 +39,6 @@ def probe(self) -> ProbeResult: ...
 
     def enumerate_devices(self) -> List[DeviceHandle]: ...
 
-    def capabilities(self, device: DeviceHandle) -> CapabilitySet: ...
-
     def read_metrics(
         self, device: DeviceHandle, request: MetricRequest
     ) -> MetricSet: ...
diff --git a/gcm/accelerator/backends/__init__.py b/gcm/accelerator/backends/__init__.py
diff --git a/gcm/accelerator/backends/nvml.py b/gcm/accelerator/backends/nvml.py
@@ -2,25 +2,19 @@
 # All rights reserved.
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
-from typing import Callable, Optional, TypeVar
+from typing import Any, Callable, Optional, TypeVar
 
-from gcm.monitoring.accelerator.backend import BackendName, DeviceHandle, ProbeResult
-from gcm.monitoring.accelerator.errors import (
-    BackendUnavailableError,
-    UnsupportedOperationError,
-)
-from gcm.monitoring.accelerator.metrics import (
-    Capability,
-    CapabilitySet,
-    MetricRequest,
-    MetricSet,
-)
-from gcm.monitoring.accelerator.probe import find_and_load_library
+from gcm.accelerator.backend import BackendName, DeviceHandle, ProbeResult
+from gcm.accelerator.errors import BackendUnavailableError, UnsupportedOperationError
+from gcm.accelerator.metrics import MetricRequest, MetricSet
+from gcm.accelerator.probe import find_and_load_library
 from gcm.monitoring.device_telemetry_client import (
     DeviceTelemetryClient,
     DeviceTelemetryException,
 )
+from gcm.monitoring.utils.error import safe_call
 from gcm.schemas.gpu.application_clock import ApplicationClockInfo
+
 from gcm.schemas.gpu.memory import GPUMemory
 from gcm.schemas.gpu.utilization import GPUUtilization
 
@@ -50,6 +44,7 @@ class NVMLBackend:
     _client: Optional[DeviceTelemetryClient] = field(
         default=None, init=False, repr=False
     )
+    _handles: dict[str, Any] = field(default_factory=dict, init=False, repr=False)
 
     def name(self) -> BackendName:
         return BackendName.NVML
@@ -83,7 +78,15 @@ def enumerate_devices(self) -> list[DeviceHandle]:
             devices: list[DeviceHandle] = []
             for index in range(device_count):
                 model: Optional[str] = None
-                handle = client.get_device_by_index(index)
+
+                # Check cache first or fetch handle
+                dev_id = str(index)
+                if dev_id in self._handles:
+                    handle = self._handles[dev_id]
+                else:
+                    handle = client.get_device_by_index(index)
+                    self._handles[dev_id] = handle
+
                 model_getter = getattr(handle, "get_name", None)
                 if callable(model_getter):
                     maybe_model = self._safe_call(model_getter)
@@ -92,7 +95,7 @@ def enumerate_devices(self) -> list[DeviceHandle]:
                 devices.append(
                     DeviceHandle(
                         backend=self.name(),
-                        id=str(index),
+                        id=dev_id,
                         vendor="nvidia",
                         model=model,
                     )
@@ -101,33 +104,21 @@ def enumerate_devices(self) -> list[DeviceHandle]:
         except DeviceTelemetryException as e:
             raise UnsupportedOperationError("NVML enumerate_devices failed") from e
 
-    def capabilities(self, _device: DeviceHandle) -> CapabilitySet:
-        return CapabilitySet(
-            values={
-                Capability.UTILIZATION,
-                Capability.MEMORY,
-                Capability.POWER,
-                Capability.THERMALS,
-                Capability.CLOCKS,
-                Capability.ECC,
-                Capability.PROCESSES,
-            }
-        )
-
     @staticmethod
     def _safe_call(func: Callable[[], _T]) -> _T | None:
-        try:
-            return func()
-        except DeviceTelemetryException:
-            return None
+        return safe_call(func, DeviceTelemetryException, logger_name=__name__)
 
     def read_metrics(self, device: DeviceHandle, _request: MetricRequest) -> MetricSet:
         # TODO: Wire MetricRequest.include_process_info once process telemetry
         # is available through HAL MetricSet.
         client = self._ensure_client()
         try:
-            index = int(device.id)
-            handle = client.get_device_by_index(index)
+            if device.id in self._handles:
+                handle = self._handles[device.id]
+            else:
+                index = int(device.id)
+                handle = client.get_device_by_index(index)
+                self._handles[device.id] = handle
         except (ValueError, DeviceTelemetryException) as e:
             raise UnsupportedOperationError(
                 f"invalid NVML device id: {device.id}"
diff --git a/gcm/accelerator/errors.py b/gcm/accelerator/errors.py
@@ -2,7 +2,7 @@
 # All rights reserved.
 from dataclasses import dataclass
 
-from gcm.monitoring.accelerator.backend import BackendName
+from gcm.accelerator.backend import BackendName
 
 
 class AcceleratorError(Exception):
diff --git a/gcm/accelerator/manager.py b/gcm/accelerator/manager.py
@@ -1,14 +1,14 @@
 # Copyright (c) Meta Platforms, Inc. and affiliates.
 # All rights reserved.
-from gcm.monitoring.accelerator.backend import (
+from gcm.accelerator.backend import (
     AcceleratorBackend,
     BackendFactory,
     BackendName,
     DeviceHandle,
     ProbeResult,
 )
-from gcm.monitoring.accelerator.errors import BackendOperationError
-from gcm.monitoring.accelerator.metrics import MetricRequest, MetricSet
+from gcm.accelerator.errors import BackendOperationError
+from gcm.accelerator.metrics import MetricRequest, MetricSet
 
 
 class AcceleratorManager:
diff --git a/gcm/accelerator/metrics.py b/gcm/accelerator/metrics.py
@@ -2,26 +2,6 @@
 # All rights reserved.
 from dataclasses import dataclass, field
 from datetime import datetime, timezone
-from enum import Enum
-
-
-class Capability(str, Enum):
-    UTILIZATION = "utilization"
-    MEMORY = "memory"
-    POWER = "power"
-    THERMALS = "thermals"
-    CLOCKS = "clocks"
-    ECC = "ecc"
-    TOPOLOGY = "topology"
-    PROCESSES = "processes"
-
-
-@dataclass(frozen=True)
-class CapabilitySet:
-    values: set[Capability]
-
-    def supports(self, capability: Capability) -> bool:
-        return capability in self.values
 
 
 @dataclass(frozen=True)
diff --git a/gcm/accelerator/probe.py b/gcm/accelerator/probe.py
diff --git a/gcm/accelerator/registry.py b/gcm/accelerator/registry.py
@@ -0,0 +1,10 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+from gcm.accelerator.backend import BackendFactory, BackendName
+from gcm.accelerator.backends.nvml import NVMLBackend
+
+
+def default_backend_factories() -> dict[BackendName, BackendFactory]:
+    return {
+        BackendName.NVML: lambda: NVMLBackend(),
+    }
diff --git a/gcm/health_checks/checks/check_nvidia_smi.py b/gcm/health_checks/checks/check_nvidia_smi.py
@@ -37,7 +37,6 @@
     DeviceTelemetryClient,
     DeviceTelemetryException,
 )
-from gcm.monitoring.device_telemetry_nvml import NVMLDeviceTelemetryClient
 from gcm.monitoring.features.gen.generated_features_healthchecksfeatures import (
     FeatureValueHealthChecksFeatures,
 )
@@ -61,6 +60,9 @@ 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()
 
 
diff --git a/gcm/monitoring/accelerator/registry.py b/gcm/monitoring/accelerator/registry.py
diff --git a/gcm/monitoring/cli/nvml_monitor.py b/gcm/monitoring/cli/nvml_monitor.py
@@ -25,10 +25,10 @@
 )
 
 import click
+from gcm.accelerator.backend import BackendName
+from gcm.accelerator.manager import AcceleratorManager
+from gcm.accelerator.registry import default_backend_factories
 from gcm.exporters import registry
-from gcm.monitoring.accelerator.backend import BackendName
-from gcm.monitoring.accelerator.manager import AcceleratorManager
-from gcm.monitoring.accelerator.registry import default_backend_factories
 from gcm.monitoring.accumulate import Accumulator
 from gcm.monitoring.click import (
     click_default_cmd,
@@ -53,7 +53,6 @@
     DeviceTelemetryException,
     GPUDevice,
 )
-from gcm.monitoring.device_telemetry_nvml import NVMLDeviceTelemetryClient
 from gcm.monitoring.sink.protocol import DataType, SinkAdditionalParams, SinkImpl
 from gcm.monitoring.sink.utils import Factory, HasRegistry
 from gcm.monitoring.utils import error
@@ -279,6 +278,9 @@ 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()
 
     def read_env(self, process_id: int) -> Env:
diff --git a/gcm/monitoring/utils/error.py b/gcm/monitoring/utils/error.py
@@ -15,6 +15,26 @@
 _P = ParamSpec("_P")
 
 
+def safe_call(
+    func: Callable[[], _T],
+    *catch: type[BaseException],
+    logger_name: Optional[str] = None,
+) -> Optional[_T]:
+    """Call *func* and return None if it raises a matching exception.
+
+    If no exception types are passed, catches all ``Exception`` subclasses.
+    Failures are logged at WARNING level.
+    """
+    catch_types: tuple[type[BaseException], ...] = catch or (Exception,)
+    try:
+        return func()
+    except catch_types:
+        logging.getLogger(logger_name or __name__).warning(
+            "safe_call: %s failed", func, exc_info=True
+        )
+        return None
+
+
 def fmt_exc_for_log() -> str:
     parts = traceback.format_exc(-1).strip().split("\n")
     return "{}: {}".format(parts[-1].strip(), parts[1].strip())
diff --git a/gcm/tests/test_accelerator_hal.py b/gcm/tests/test_accelerator_hal.py
diff --git a/website/docs/GCM_Monitoring/accelerator_hal.md b/website/docs/GCM_Monitoring/accelerator_hal.md
