wandb 0.18.0__py3-none-macosx_10_13_x86_64.whl

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._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",
+            }
+        }

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

@@ -0,0 +1,166 @@
+import threading
+from collections import deque
+from typing import TYPE_CHECKING, List, Optional
+
+try:
+    import psutil
+except ImportError:
+    psutil = None
+
+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 ProcessMemoryRSS:
+    """Memory resident set size (RSS) in MB.
+
+    RSS is the portion of memory occupied by a process that is held in main memory (RAM).
+    """
+
+    # name = "memory_rss"
+    name = "proc.memory.rssMB"
+    samples: "Deque[float]"
+
+    def __init__(self, pid: int) -> None:
+        self.pid = pid
+        self.process: Optional[psutil.Process] = None
+        self.samples = deque([])
+
+    def sample(self) -> None:
+        if self.process is None:
+            self.process = psutil.Process(self.pid)
+
+        self.samples.append(self.process.memory_info().rss / 1024 / 1024)
+
+    def clear(self) -> None:
+        self.samples.clear()
+
+    def aggregate(self) -> dict:
+        if not self.samples:
+            return {}
+        aggregate = aggregate_mean(self.samples)
+        return {self.name: aggregate}
+
+
+class ProcessMemoryPercent:
+    """Process memory usage in percent."""
+
+    # name = "process_memory_percent"
+    name = "proc.memory.percent"
+    samples: "Deque[float]"
+
+    def __init__(self, pid: int) -> None:
+        self.pid = pid
+        self.process: Optional[psutil.Process] = None
+        self.samples = deque([])
+
+    def sample(self) -> None:
+        if self.process is None:
+            self.process = psutil.Process(self.pid)
+
+        self.samples.append(self.process.memory_percent())
+
+    def clear(self) -> None:
+        self.samples.clear()
+
+    def aggregate(self) -> dict:
+        if not self.samples:
+            return {}
+        aggregate = aggregate_mean(self.samples)
+        return {self.name: aggregate}
+
+
+class MemoryPercent:
+    """Total system memory usage in percent."""
+
+    # name = "memory_percent"
+    name = "memory"
+    samples: "Deque[float]"
+
+    def __init__(self) -> None:
+        self.samples = deque([])
+
+    def sample(self) -> None:
+        self.samples.append(psutil.virtual_memory().percent)
+
+    def clear(self) -> None:
+        self.samples.clear()
+
+    def aggregate(self) -> dict:
+        if not self.samples:
+            return {}
+        aggregate = aggregate_mean(self.samples)
+        return {self.name: aggregate}
+
+
+class MemoryAvailable:
+    """Total system memory available in MB."""
+
+    # name = "memory_available"
+    name = "proc.memory.availableMB"
+    samples: "Deque[float]"
+
+    def __init__(self) -> None:
+        self.samples = deque([])
+
+    def sample(self) -> None:
+        self.samples.append(psutil.virtual_memory().available / 1024 / 1024)
+
+    def clear(self) -> None:
+        self.samples.clear()
+
+    def aggregate(self) -> dict:
+        if not self.samples:
+            return {}
+        aggregate = aggregate_mean(self.samples)
+        return {self.name: aggregate}
+
+
+@asset_registry.register
+class Memory:
+    def __init__(
+        self,
+        interface: "Interface",
+        settings: "SettingsStatic",
+        shutdown_event: threading.Event,
+    ) -> None:
+        self.name = self.__class__.__name__.lower()
+        self.metrics: List[Metric] = [
+            MemoryAvailable(),
+            MemoryPercent(),
+            ProcessMemoryRSS(settings._stats_pid),
+            ProcessMemoryPercent(settings._stats_pid),
+        ]
+        self.metrics_monitor = MetricsMonitor(
+            self.name,
+            self.metrics,
+            interface,
+            settings,
+            shutdown_event,
+        )
+
+    def start(self) -> None:
+        self.metrics_monitor.start()
+
+    def finish(self) -> None:
+        self.metrics_monitor.finish()
+
+    @classmethod
+    def is_available(cls) -> bool:
+        """Return a new instance of the CPU metrics."""
+        return psutil is not None
+
+    def probe(self) -> dict:
+        """Return a dict of the hardware information."""
+        # total available memory in gigabytes
+        return {
+            "memory": {
+                "total": psutil.virtual_memory().total / 1024 / 1024 / 1024,
+            }
+        }

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

@@ -0,0 +1,125 @@
+import threading
+from collections import deque
+from typing import TYPE_CHECKING, List
+
+try:
+    import psutil
+except ImportError:
+    psutil = None
+
+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 NetworkSent:
+    """Network bytes sent."""
+
+    name = "network.sent"
+    samples: "Deque[float]"
+
+    def __init__(self) -> None:
+        self.samples = deque([])
+        self.sent_init = psutil.net_io_counters().bytes_sent
+
+    def sample(self) -> None:
+        self.samples.append(psutil.net_io_counters().bytes_sent - self.sent_init)
+
+    def clear(self) -> None:
+        self.samples.clear()
+
+    def aggregate(self) -> dict:
+        if not self.samples:
+            return {}
+        aggregate = aggregate_mean(self.samples)
+        # todo: this is an adapter for the legacy metrics system
+        # return {"network": {self.name: aggregate}}
+        return {self.name: aggregate}
+
+
+class NetworkRecv:
+    """Network bytes received."""
+
+    name = "network.recv"
+    samples: "Deque[float]"
+
+    def __init__(self) -> None:
+        self.samples = deque([])
+        self.recv_init = psutil.net_io_counters().bytes_recv
+
+    def sample(self) -> None:
+        self.samples.append(psutil.net_io_counters().bytes_recv - self.recv_init)
+
+    def clear(self) -> None:
+        self.samples.clear()
+
+    def aggregate(self) -> dict:
+        if not self.samples:
+            return {}
+        aggregate = aggregate_mean(self.samples)
+        # todo: this is an adapter for the legacy metrics system
+        # return {"network": {self.name: aggregate}}
+
+        return {self.name: aggregate}
+
+
+@asset_registry.register
+class Network:
+    def __init__(
+        self,
+        interface: "Interface",
+        settings: "SettingsStatic",
+        shutdown_event: threading.Event,
+    ) -> None:
+        self.name = self.__class__.__name__.lower()
+        self.metrics: List[Metric] = [
+            NetworkSent(),
+            NetworkRecv(),
+        ]
+        self.metrics_monitor = MetricsMonitor(
+            self.name,
+            self.metrics,
+            interface,
+            settings,
+            shutdown_event,
+        )
+
+    def start(self) -> None:
+        self.metrics_monitor.start()
+
+    def finish(self) -> None:
+        self.metrics_monitor.finish()
+
+    @classmethod
+    def is_available(cls) -> bool:
+        """Return a new instance of the CPU metrics."""
+        return psutil is not None
+
+    def probe(self) -> dict:
+        """Return a dict of the hardware information."""
+        # net_if_addrs = psutil.net_if_addrs()
+
+        # return {
+        #     self.name: {
+        #         "interfaces": {
+        #             k: {
+        #                 "addresses": [
+        #                     {
+        #                         "address": v.address,
+        #                         "netmask": v.netmask,
+        #                         "broadcast": v.broadcast,
+        #                         "ptp": v.ptp,
+        #                     }
+        #                     for v in v
+        #                 ]
+        #             }
+        #             for k, v in net_if_addrs.items()
+        #         }
+        #     }
+        # }
+        return {}