wandb 0.20.1__py3-none-win32.whl → 0.20.2rc20250616__py3-none-win32.whl

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

@@ -1,205 +0,0 @@
-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:
-                        logger.exception("Failed to sample metric.")
-                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:
-                logger.exception("Failed to serialize metric.")
-        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:
-            logger.exception("Failed to publish metrics.")
-
-    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

@@ -1,177 +0,0 @@
-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",
-            }
-        }

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

@@ -1,166 +0,0 @@
-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.x_stats_pid),
-            ProcessMemoryPercent(settings.x_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,
-            }
-        }

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

@@ -1,125 +0,0 @@
-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 {}