wandb 0.19.1rc1__py3-none-musllinux_1_2_aarch64.whl

wandb/sdk/internal/system/assets/gpu_amd.py

@@ -0,0 +1,233 @@
+import json
+import logging
+import shutil
+import subprocess
+import threading
+from collections import deque
+from typing import TYPE_CHECKING, Any, Dict, Final, List, Literal, Union
+
+from wandb.sdk.lib import telemetry
+
+from .aggregators import aggregate_mean
+from .asset_registry import asset_registry
+from .interfaces import Interface, Metric, MetricsMonitor
+
+if TYPE_CHECKING:
+    from typing import Deque
+
+    from wandb.sdk.internal.settings_static import SettingsStatic
+
+
+logger = logging.getLogger(__name__)
+ROCM_SMI_CMD: Final[str] = shutil.which("rocm-smi") or "/usr/bin/rocm-smi"
+
+
+_StatsKeys = Literal[
+    "gpu",
+    "memoryAllocated",
+    "temp",
+    "powerWatts",
+    "powerPercent",
+]
+_Stats = Dict[_StatsKeys, float]
+
+
+_InfoDict = Dict[str, Union[int, List[Dict[str, Any]]]]
+
+
+def get_rocm_smi_stats() -> Dict[str, Any]:
+    command = [str(ROCM_SMI_CMD), "-a", "--json"]
+    output = subprocess.check_output(command, universal_newlines=True).strip()
+    if "No AMD GPUs specified" in output:
+        return {}
+    return json.loads(output.split("\n")[0])  # type: ignore
+
+
+def parse_stats(stats: Dict[str, str]) -> _Stats:
+    """Parse stats from rocm-smi output."""
+    parsed_stats: _Stats = {}
+
+    try:
+        parsed_stats["gpu"] = float(stats.get("GPU use (%)"))  # type: ignore
+    except (TypeError, ValueError):
+        logger.warning("Could not parse GPU usage as float")
+    try:
+        parsed_stats["memoryAllocated"] = float(stats.get("GPU memory use (%)"))  # type: ignore
+    except (TypeError, ValueError):
+        logger.warning("Could not parse GPU memory allocation as float")
+    try:
+        parsed_stats["temp"] = float(stats.get("Temperature (Sensor memory) (C)"))  # type: ignore
+    except (TypeError, ValueError):
+        logger.warning("Could not parse GPU temperature as float")
+    try:
+        parsed_stats["powerWatts"] = float(
+            stats.get("Average Graphics Package Power (W)")  # type: ignore
+        )
+    except (TypeError, ValueError):
+        logger.warning("Could not parse GPU power as float")
+    try:
+        parsed_stats["powerPercent"] = (
+            float(stats.get("Average Graphics Package Power (W)"))  # type: ignore
+            / float(stats.get("Max Graphics Package Power (W)"))  # type: ignore
+            * 100
+        )
+    except (TypeError, ValueError):
+        logger.warning("Could not parse GPU average/max power as float")
+
+    return parsed_stats
+
+
+class GPUAMDStats:
+    """Stats for AMD GPU devices."""
+
+    name = "gpu.{gpu_id}.{key}"
+    samples: "Deque[List[_Stats]]"
+
+    def __init__(self) -> None:
+        self.samples = deque()
+
+    def sample(self) -> None:
+        try:
+            raw_stats = get_rocm_smi_stats()
+            cards = []
+
+            card_keys = [
+                key for key in sorted(raw_stats.keys()) if key.startswith("card")
+            ]
+
+            for card_key in card_keys:
+                card_stats = raw_stats[card_key]
+                stats = parse_stats(card_stats)
+                if stats:
+                    cards.append(stats)
+
+            if cards:
+                self.samples.append(cards)
+
+        except (OSError, ValueError, TypeError, subprocess.CalledProcessError) as e:
+            logger.exception(f"GPU stats error: {e}")
+
+    def clear(self) -> None:
+        self.samples.clear()
+
+    def aggregate(self) -> dict:
+        if not self.samples:
+            return {}
+        stats = {}
+        device_count = len(self.samples[0])
+
+        for i in range(device_count):
+            samples = [sample[i] for sample in self.samples]
+
+            for key in samples[0].keys():
+                samples_key = [s[key] for s in samples]
+                aggregate = aggregate_mean(samples_key)
+                stats[self.name.format(gpu_id=i, key=key)] = aggregate
+
+        return stats
+
+
+@asset_registry.register
+class GPUAMD:
+    """GPUAMD is a class for monitoring AMD GPU devices.
+
+    Uses AMD's rocm_smi tool to get GPU stats.
+    For the list of supported environments and devices, see
+    https://github.com/RadeonOpenCompute/ROCm/blob/develop/docs/deploy/
+    """
+
+    def __init__(
+        self,
+        interface: "Interface",
+        settings: "SettingsStatic",
+        shutdown_event: threading.Event,
+    ) -> None:
+        self.name = self.__class__.__name__.lower()
+        self.metrics: List[Metric] = [
+            GPUAMDStats(),
+        ]
+        self.metrics_monitor = MetricsMonitor(
+            self.name,
+            self.metrics,
+            interface,
+            settings,
+            shutdown_event,
+        )
+        telemetry_record = telemetry.TelemetryRecord()
+        telemetry_record.env.amd_gpu = True
+        interface._publish_telemetry(telemetry_record)
+
+    @classmethod
+    def is_available(cls) -> bool:
+        rocm_smi_available = shutil.which(ROCM_SMI_CMD) is not None
+        if not rocm_smi_available:
+            # If rocm-smi is not available, we can't monitor AMD GPUs
+            return False
+
+        is_driver_initialized = False
+
+        try:
+            # inspired by https://github.com/ROCm/rocm_smi_lib/blob/5d2cd0c2715ae45b8f9cfe1e777c6c2cd52fb601/python_smi_tools/rocm_smi.py#L71C1-L81C17
+            with open("/sys/module/amdgpu/initstate") as file:
+                file_content = file.read()
+                if "live" in file_content:
+                    is_driver_initialized = True
+        except FileNotFoundError:
+            pass
+
+        can_read_rocm_smi = False
+        try:
+            # try to read stats from rocm-smi and parse them
+            raw_stats = get_rocm_smi_stats()
+            card_keys = [
+                key for key in sorted(raw_stats.keys()) if key.startswith("card")
+            ]
+
+            for card_key in card_keys:
+                card_stats = raw_stats[card_key]
+                parse_stats(card_stats)
+
+            can_read_rocm_smi = True
+        except Exception:
+            pass
+
+        return is_driver_initialized and can_read_rocm_smi
+
+    def start(self) -> None:
+        self.metrics_monitor.start()
+
+    def finish(self) -> None:
+        self.metrics_monitor.finish()
+
+    def probe(self) -> dict:
+        info: _InfoDict = {}
+        try:
+            stats = get_rocm_smi_stats()
+
+            info["gpu_count"] = len(
+                [key for key in stats.keys() if key.startswith("card")]
+            )
+            key_mapping = {
+                "id": "GPU ID",
+                "unique_id": "Unique ID",
+                "vbios_version": "VBIOS version",
+                "performance_level": "Performance Level",
+                "gpu_overdrive": "GPU OverDrive value (%)",
+                "gpu_memory_overdrive": "GPU Memory OverDrive value (%)",
+                "max_power": "Max Graphics Package Power (W)",
+                "series": "Card series",
+                "model": "Card model",
+                "vendor": "Card vendor",
+                "sku": "Card SKU",
+                "sclk_range": "Valid sclk range",
+                "mclk_range": "Valid mclk range",
+            }
+
+            info["gpu_devices"] = [
+                {k: stats[key][v] for k, v in key_mapping.items() if stats[key].get(v)}
+                for key in stats.keys()
+                if key.startswith("card")
+            ]
+        except Exception as e:
+            logger.exception(f"GPUAMD probe error: {e}")
+        return info

wandb/sdk/internal/system/assets/interfaces.py

@@ -0,0 +1,205 @@
+import datetime
+import logging
+import threading
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    List,
+    Optional,
+    Protocol,
+    TypeVar,
+    runtime_checkable,
+)
+
+if TYPE_CHECKING:
+    from typing import Deque
+
+    from wandb.proto.wandb_telemetry_pb2 import TelemetryRecord
+    from wandb.sdk.interface.interface import FilesDict
+    from wandb.sdk.internal.settings_static import SettingsStatic
+
+import psutil
+
+TimeStamp = TypeVar("TimeStamp", bound=datetime.datetime)
+
+
+logger = logging.getLogger(__name__)
+
+
+class Metric(Protocol):
+    """Base protocol for individual metrics."""
+
+    name: str
+    # samples: Sequence[Tuple[TimeStamp, Sample]]
+    samples: "Deque[Any]"
+
+    def sample(self) -> None:
+        """Sample the metric."""
+        ...  # pragma: no cover
+
+    def clear(self) -> None:
+        """Clear the samples."""
+        ...  # pragma: no cover
+
+    def aggregate(self) -> dict:
+        """Aggregate the samples."""
+        ...  # pragma: no cover
+
+
+@runtime_checkable
+class SetupTeardown(Protocol):
+    """Protocol for classes that require setup and teardown."""
+
+    def setup(self) -> None:
+        """Extra setup required for the metric beyond __init__."""
+        ...  # pragma: no cover
+
+    def teardown(self) -> None:
+        """Extra teardown required for the metric."""
+        ...  # pragma: no cover
+
+
+@runtime_checkable
+class Asset(Protocol):
+    """Base protocol encapsulate everything relating to an "Asset".
+
+    An asset can be CPU, GPU, TPU, Network, I/O etc.
+    """
+
+    name: str
+    metrics: List[Metric]
+    metrics_monitor: "MetricsMonitor"
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None: ...  # pragma: no cover
+
+    @classmethod
+    def is_available(cls) -> bool:
+        """Check if the resource is available."""
+        ...  # pragma: no cover
+
+    def start(self) -> None:
+        """Start monitoring the resource."""
+        ...  # pragma: no cover
+
+    def finish(self) -> None:
+        """Finish monitoring the resource."""
+        ...  # pragma: no cover
+
+    def probe(self) -> dict:
+        """Get static information about the resource."""
+        ...  # pragma: no cover
+
+
+class Interface(Protocol):
+    def publish_stats(self, stats: dict) -> None: ...  # pragma: no cover
+
+    def _publish_telemetry(
+        self, telemetry: "TelemetryRecord"
+    ) -> None: ...  # pragma: no cover
+
+    def publish_files(self, files_dict: "FilesDict") -> None: ...  # pragma: no cover
+
+
+class MetricsMonitor:
+    """Takes care of collecting, sampling, serializing, and publishing a set of metrics."""
+
+    def __init__(
+        self,
+        asset_name: str,
+        metrics: List[Metric],
+        interface: Interface,
+        settings: "SettingsStatic",
+        shutdown_event: threading.Event,
+    ) -> None:
+        self.metrics = metrics
+        self.asset_name = asset_name
+        self._interface = interface
+        self._process: Optional[threading.Thread] = None
+        self._shutdown_event: threading.Event = shutdown_event
+
+        self.sampling_interval: float = float(
+            max(
+                0.1,
+                settings.x_stats_sampling_interval,
+            )
+        )  # seconds
+        self.samples_to_aggregate = 1
+
+    def monitor(self) -> None:
+        """Poll the Asset metrics."""
+        while not self._shutdown_event.is_set():
+            for _ in range(self.samples_to_aggregate):
+                for metric in self.metrics:
+                    try:
+                        metric.sample()
+                    except psutil.NoSuchProcess:
+                        logger.info(f"Process {metric.name} has exited.")
+                        self._shutdown_event.set()
+                        break
+                    except Exception as e:
+                        logger.error(f"Failed to sample metric: {e}")
+                self._shutdown_event.wait(self.sampling_interval)
+                if self._shutdown_event.is_set():
+                    break
+            self.publish()
+
+    def aggregate(self) -> dict:
+        """Return a dict of metrics."""
+        aggregated_metrics = {}
+        for metric in self.metrics:
+            try:
+                serialized_metric = metric.aggregate()
+                aggregated_metrics.update(serialized_metric)
+                # aggregated_metrics = wandb.util.merge_dicts(
+                #     aggregated_metrics, metric.serialize()
+                # )
+            except Exception as e:
+                logger.error(f"Failed to serialize metric: {e}")
+        return aggregated_metrics
+
+    def publish(self) -> None:
+        """Publish the Asset metrics."""
+        try:
+            aggregated_metrics = self.aggregate()
+            if aggregated_metrics:
+                self._interface.publish_stats(aggregated_metrics)
+            for metric in self.metrics:
+                metric.clear()
+        except Exception as e:
+            logger.error(f"Failed to publish metrics: {e}")
+
+    def start(self) -> None:
+        if (self._process is not None) or self._shutdown_event.is_set():
+            return None
+
+        thread_name = f"{self.asset_name[:15]}"  # thread names are limited to 15 chars
+        try:
+            for metric in self.metrics:
+                if isinstance(metric, SetupTeardown):
+                    metric.setup()
+            self._process = threading.Thread(
+                target=self.monitor,
+                daemon=True,
+                name=thread_name,
+            )
+            self._process.start()
+            logger.info(f"Started {thread_name} monitoring")
+        except Exception as e:
+            logger.warning(f"Failed to start {thread_name} monitoring: {e}")
+            self._process = None
+
+    def finish(self) -> None:
+        if self._process is None:
+            return None
+
+        thread_name = f"{self.asset_name[:15]}"
+        try:
+            self._process.join()
+            logger.info(f"Joined {thread_name} monitor")
+            for metric in self.metrics:
+                if isinstance(metric, SetupTeardown):
+                    metric.teardown()
+        except Exception as e:
+            logger.warning(f"Failed to finish {thread_name} monitoring: {e}")
+        finally:
+            self._process = None

wandb/sdk/internal/system/assets/ipu.py

@@ -0,0 +1,177 @@
+import threading
+from collections import deque
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union
+
+try:
+    import gcipuinfo  # type: ignore
+except ImportError:
+    gcipuinfo = None
+
+import wandb
+
+from .aggregators import aggregate_mean
+from .asset_registry import asset_registry
+from .interfaces import Interface, Metric, MetricsMonitor
+
+if TYPE_CHECKING:
+    from typing import Deque
+
+    from wandb.sdk.internal.settings_static import SettingsStatic
+
+
+class IPUStats:
+    """Stats for Graphcore IPU devices."""
+
+    name = "ipu.{}.{}"
+    samples: "Deque[dict]"
+
+    # The metrics that change over time.
+    # Only these are returned on each invocation
+    # to avoid sending a load of unnecessary data.
+    variable_metric_keys = {
+        "average board temp",
+        "average die temp",
+        "clock",
+        "ipu power",
+        "ipu utilisation",
+        "ipu utilisation (session)",
+    }
+
+    def __init__(self, pid: int, gc_ipu_info: Optional[Any] = None) -> None:
+        self.samples: Deque[dict] = deque()
+
+        if gc_ipu_info is None:
+            if not gcipuinfo:
+                raise ImportError(
+                    "Monitoring IPU stats requires gcipuinfo to be installed"
+                )
+
+            self._gc_ipu_info = gcipuinfo.gcipuinfo()
+        else:
+            self._gc_ipu_info = gc_ipu_info
+        self._gc_ipu_info.setUpdateMode(True)
+
+        self._pid = pid
+        self._devices_called: Set[str] = set()
+
+    @staticmethod
+    def parse_metric(key: str, value: str) -> Optional[Tuple[str, Union[int, float]]]:
+        metric_suffixes = {
+            "temp": "C",
+            "clock": "MHz",
+            "power": "W",
+            "utilisation": "%",
+            "utilisation (session)": "%",
+            "speed": "GT/s",
+        }
+
+        for metric, suffix in metric_suffixes.items():
+            if key.endswith(metric) and value.endswith(suffix):
+                value = value[: -len(suffix)]
+                key = f"{key} ({suffix})"
+
+        try:
+            float_value = float(value)
+            num_value = int(float_value) if float_value.is_integer() else float_value
+        except ValueError:
+            return None
+
+        return key, num_value
+
+    def sample(self) -> None:
+        try:
+            stats = {}
+            devices = self._gc_ipu_info.getDevices()
+
+            for device in devices:
+                device_metrics: Dict[str, str] = dict(device)
+
+                pid = device_metrics.get("user process id")
+                if pid is None or int(pid) != self._pid:
+                    continue
+
+                device_id = device_metrics.get("id")
+                initial_call = device_id not in self._devices_called
+                if device_id is not None:
+                    self._devices_called.add(device_id)
+
+                for key, value in device_metrics.items():
+                    log_metric = initial_call or key in self.variable_metric_keys
+                    if not log_metric:
+                        continue
+                    parsed = self.parse_metric(key, value)
+                    if parsed is None:
+                        continue
+                    parsed_key, parsed_value = parsed
+                    stats[self.name.format(device_id, parsed_key)] = parsed_value
+
+            self.samples.append(stats)
+
+        except Exception as e:
+            wandb.termwarn(f"IPU stats error {e}", repeat=False)
+
+    def clear(self) -> None:
+        self.samples.clear()
+
+    def aggregate(self) -> dict:
+        if not self.samples:
+            return {}
+        stats = {}
+        for key in self.samples[0].keys():
+            samples = [s[key] for s in self.samples if key in s]
+            aggregate = aggregate_mean(samples)
+            stats[key] = aggregate
+        return stats
+
+
+@asset_registry.register
+class IPU:
+    def __init__(
+        self,
+        interface: "Interface",
+        settings: "SettingsStatic",
+        shutdown_event: threading.Event,
+    ) -> None:
+        self.name = self.__class__.__name__.lower()
+        self.metrics: List[Metric] = [
+            IPUStats(settings.x_stats_pid),
+        ]
+        self.metrics_monitor = MetricsMonitor(
+            self.name,
+            self.metrics,
+            interface,
+            settings,
+            shutdown_event,
+        )
+
+    @classmethod
+    def is_available(cls) -> bool:
+        return gcipuinfo is not None
+
+    def start(self) -> None:
+        self.metrics_monitor.start()
+
+    def finish(self) -> None:
+        self.metrics_monitor.finish()
+
+    def probe(self) -> dict:
+        device_data = self.metrics[0]._gc_ipu_info.getDevices()  # type: ignore
+        device_count = len(device_data)
+        devices = []
+        for i, device in enumerate(device_data):
+            device_metrics: Dict[str, str] = dict(device)
+            devices.append(
+                {
+                    "id": device_metrics.get("id") or i,
+                    "board ipu index": device_metrics.get("board ipu index"),
+                    "board type": device_metrics.get("board type") or "unknown",
+                }
+            )
+
+        return {
+            self.name: {
+                "device_count": device_count,
+                "devices": devices,
+                "vendor": "Graphcore",
+            }
+        }