wandb 0.13.10__py3-none-any.whl → 0.14.0__py3-none-any.whl

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

@@ -18,9 +18,7 @@ if TYPE_CHECKING:
 
 
 class DiskUsage:
-    """
-    Total system disk usage in percent.
-    """
+    """Total system disk usage in percent."""
 
     # name = "disk_usage"
     name = "disk"
@@ -62,7 +60,7 @@ class Disk:
 
     @classmethod
     def is_available(cls) -> bool:
-        """Return a new instance of the CPU metrics"""
+        """Return a new instance of the CPU metrics."""
         return psutil is not None
 
     def probe(self) -> dict:

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

@@ -55,9 +55,7 @@ def gpu_in_use_by_this_process(gpu_handle: "GPUHandle", pid: int) -> bool:
 
 
 class GPUMemoryUtilization:
-    """
-    GPU memory utilization in percent for each GPU.
-    """
+    """GPU memory utilization in percent for each GPU."""
 
     # name = "memory_utilization"
     name = "gpu.{}.memory"
@@ -99,9 +97,7 @@ class GPUMemoryUtilization:
 
 
 class GPUMemoryAllocated:
-    """
-    GPU memory allocated in percent for each GPU.
-    """
+    """GPU memory allocated in percent for each GPU."""
 
     # name = "memory_allocated"
     name = "gpu.{}.memoryAllocated"
@@ -142,9 +138,7 @@ class GPUMemoryAllocated:
 
 
 class GPUUtilization:
-    """
-    GPU utilization in percent for each GPU.
-    """
+    """GPU utilization in percent for each GPU."""
 
     # name = "gpu_utilization"
     name = "gpu.{}.gpu"
@@ -186,9 +180,7 @@ class GPUUtilization:
 
 
 class GPUTemperature:
-    """
-    GPU temperature in Celsius for each GPU.
-    """
+    """GPU temperature in Celsius for each GPU."""
 
     # name = "gpu_temperature"
     name = "gpu.{}.temp"
@@ -233,9 +225,7 @@ class GPUTemperature:
 
 
 class GPUPowerUsageWatts:
-    """
-    GPU power usage in Watts for each GPU.
-    """
+    """GPU power usage in Watts for each GPU."""
 
     name = "gpu.{}.powerWatts"
     # samples: Deque[Tuple[datetime.datetime, float]]
@@ -273,9 +263,7 @@ class GPUPowerUsageWatts:
 
 
 class GPUPowerUsagePercent:
-    """
-    GPU power usage in percent for each GPU.
-    """
+    """GPU power usage in percent for each GPU."""
 
     name = "gpu.{}.powerPercent"
     # samples: Deque[Tuple[datetime.datetime, float]]

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

@@ -34,13 +34,11 @@ class _Stats(TypedDict):
     temp: float
     powerWatts: float  # noqa: N815
     powerPercent: float  # noqa: N815
-    # cpuWaitMs: float  # noqa: N815
+    # cpuWaitMs: float
 
 
 class GPUAppleStats:
-    """
-    Apple GPU stats available on Arm Macs.
-    """
+    """Apple GPU stats available on Arm Macs."""
 
     name = "gpu.0.{}"
     samples: "Deque[_Stats]"

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

@@ -26,29 +26,43 @@ logger = logging.getLogger(__name__)
 
 
 class Metric(Protocol):
-    """
-    Base protocol for individual metrics
-    """
+    """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 to encapsulate everything relating to an "Asset"
-    e.g. CPU, GPU, TPU, Network, I/O etc.
+    """Base protocol encapsulate everything relating to an "Asset".
+
+    An asset can be CPU, GPU, TPU, Network, I/O etc.
     """
 
     name: str
@@ -60,19 +74,19 @@ class Asset(Protocol):
 
     @classmethod
     def is_available(cls) -> bool:
-        """Check if the resource is available"""
+        """Check if the resource is available."""
         ...  # pragma: no cover
 
     def start(self) -> None:
-        """Start monitoring the resource"""
+        """Start monitoring the resource."""
         ...  # pragma: no cover
 
     def finish(self) -> None:
-        """finish monitoring the resource"""
+        """Finish monitoring the resource."""
         ...  # pragma: no cover
 
     def probe(self) -> dict:
-        """Get static information about the resource"""
+        """Get static information about the resource."""
         ...  # pragma: no cover
 
 
@@ -88,9 +102,7 @@ class Interface(Protocol):
 
 
 class MetricsMonitor:
-    """
-    Takes care of collecting, sampling, serializing, and publishing a set of metrics.
-    """
+    """Takes care of collecting, sampling, serializing, and publishing a set of metrics."""
 
     def __init__(
         self,
@@ -119,7 +131,7 @@ class MetricsMonitor:
         )
 
     def monitor(self) -> None:
-        """Poll the Asset metrics"""
+        """Poll the Asset metrics."""
         while not self._shutdown_event.is_set():
             for _ in range(self.samples_to_aggregate):
                 for metric in self.metrics:
@@ -133,7 +145,7 @@ class MetricsMonitor:
             self.publish()
 
     def aggregate(self) -> dict:
-        """Return a dict of metrics"""
+        """Return a dict of metrics."""
         aggregated_metrics = {}
         for metric in self.metrics:
             try:
@@ -147,7 +159,7 @@ class MetricsMonitor:
         return aggregated_metrics
 
     def publish(self) -> None:
-        """Publish the Asset metrics"""
+        """Publish the Asset metrics."""
         try:
             aggregated_metrics = self.aggregate()
             if aggregated_metrics:
@@ -158,21 +170,37 @@ class MetricsMonitor:
             logger.error(f"Failed to publish metrics: {e}")
 
     def start(self) -> None:
-        if self._process is None and not self._shutdown_event.is_set():
+        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=f"{self.asset_name}",
+                name=thread_name,
             )
             self._process.start()
-            logger.info(f"Started {self._process.name}")
+            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 {self._process.name}")
+            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 join {self._process.name}: {e}")
-        self._process = None
+            logger.warning(f"Failed to finish {thread_name} monitoring: {e}")
+        finally:
+            self._process = None

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

@@ -20,9 +20,7 @@ if TYPE_CHECKING:
 
 
 class IPUStats:
-    """
-    Stats for Graphcore IPU devices
-    """
+    """Stats for Graphcore IPU devices."""
 
     name = "ipu.{}.{}"
     samples: "Deque[dict]"

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

@@ -18,8 +18,8 @@ if TYPE_CHECKING:
 
 
 class ProcessMemoryRSS:
-    """
-    Memory resident set size (RSS) in MB.
+    """Memory resident set size (RSS) in MB.
+
     RSS is the portion of memory occupied by a process that is held in main memory (RAM).
     """
 
@@ -49,9 +49,7 @@ class ProcessMemoryRSS:
 
 
 class ProcessMemoryPercent:
-    """
-    Process memory usage in percent.
-    """
+    """Process memory usage in percent."""
 
     # name = "process_memory_percent"
     name = "proc.memory.percent"
@@ -79,9 +77,7 @@ class ProcessMemoryPercent:
 
 
 class MemoryPercent:
-    """
-    Total system memory usage in percent.
-    """
+    """Total system memory usage in percent."""
 
     # name = "memory_percent"
     name = "memory"
@@ -104,9 +100,7 @@ class MemoryPercent:
 
 
 class MemoryAvailable:
-    """
-    Total system memory available in MB.
-    """
+    """Total system memory available in MB."""
 
     # name = "memory_available"
     name = "proc.memory.availableMB"
@@ -159,11 +153,11 @@ class Memory:
 
     @classmethod
     def is_available(cls) -> bool:
-        """Return a new instance of the CPU metrics"""
+        """Return a new instance of the CPU metrics."""
         return psutil is not None
 
     def probe(self) -> dict:
-        """Return a dict of the hardware information"""
+        """Return a dict of the hardware information."""
         # total available memory in gigabytes
         return {
             "memory": {

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

@@ -18,9 +18,7 @@ if TYPE_CHECKING:
 
 
 class NetworkSent:
-    """
-    Network bytes sent.
-    """
+    """Network bytes sent."""
 
     name = "network.sent"
     samples: "Deque[float]"
@@ -45,9 +43,7 @@ class NetworkSent:
 
 
 class NetworkRecv:
-    """
-    Network bytes received.
-    """
+    """Network bytes received."""
 
     name = "network.recv"
     samples: "Deque[float]"
@@ -101,11 +97,11 @@ class Network:
 
     @classmethod
     def is_available(cls) -> bool:
-        """Return a new instance of the CPU metrics"""
+        """Return a new instance of the CPU metrics."""
         return psutil is not None
 
     def probe(self) -> dict:
-        """Return a dict of the hardware information"""
+        """Return a dict of the hardware information."""
         # net_if_addrs = psutil.net_if_addrs()
 
         # return {

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

@@ -0,0 +1,283 @@
+import logging
+import multiprocessing as mp
+import re
+import sys
+from collections import defaultdict, deque
+from functools import lru_cache
+from hashlib import md5
+from types import ModuleType
+from typing import TYPE_CHECKING, Dict, List, Mapping, Tuple, Union
+
+if sys.version_info >= (3, 8):
+    from typing import Final
+else:
+    from typing_extensions import Final
+
+import requests
+import requests.adapters
+import urllib3
+
+import wandb
+from wandb.sdk.lib import telemetry
+
+from .aggregators import aggregate_last, aggregate_mean
+from .interfaces import Interface, Metric, MetricsMonitor
+
+if TYPE_CHECKING:
+    from typing import Deque, Optional
+
+    from wandb.sdk.internal.settings_static import SettingsStatic
+
+
+_PREFIX: Final[str] = "openmetrics"
+
+_REQUEST_RETRY_STRATEGY = urllib3.util.retry.Retry(
+    backoff_factor=1,
+    total=3,
+    status_forcelist=(408, 409, 429, 500, 502, 503, 504),
+)
+_REQUEST_POOL_CONNECTIONS = 4
+_REQUEST_POOL_MAXSIZE = 4
+_REQUEST_TIMEOUT = 3
+
+
+logger = logging.getLogger(__name__)
+
+
+prometheus_client_parser: "Optional[ModuleType]" = None
+try:
+    import prometheus_client.parser  # type: ignore
+
+    prometheus_client_parser = prometheus_client.parser
+except ImportError:
+    pass
+
+
+def _setup_requests_session() -> requests.Session:
+    session = requests.Session()
+    adapter = requests.adapters.HTTPAdapter(
+        max_retries=_REQUEST_RETRY_STRATEGY,
+        pool_connections=_REQUEST_POOL_CONNECTIONS,
+        pool_maxsize=_REQUEST_POOL_MAXSIZE,
+    )
+    session.mount("http://", adapter)
+    session.mount("https://", adapter)
+    return session
+
+
+def _nested_dict_to_tuple(
+    nested_dict: Mapping[str, Mapping[str, str]]
+) -> Tuple[Tuple[str, Tuple[str, str]], ...]:
+    return tuple((k, *v.items()) for k, v in nested_dict.items())  # type: ignore
+
+
+def _tuple_to_nested_dict(
+    nested_tuple: Tuple[Tuple[str, Tuple[str, str]], ...]
+) -> Dict[str, Dict[str, str]]:
+    return {k: dict(v) for k, *v in nested_tuple}
+
+
+@lru_cache(maxsize=128)
+def _should_capture_metric(
+    metric_name: str,
+    metric_labels: Tuple[str, ...],
+    filters: Tuple[Tuple[str, Tuple[str, str]], ...],
+) -> bool:
+    # we use tuples to make the function arguments hashable => usable with lru_cache
+    should_capture = False
+
+    if not filters:
+        return should_capture
+
+    # self.filters keys are regexes, check the name against them
+    # and for the first match, check the labels against the label filters.
+    # assume that if at least one label filter doesn't match, the metric
+    # should not be captured.
+    # it's up to the user to make sure that the filters are not conflicting etc.
+    metric_labels_dict = {t[0]: t[1] for t in metric_labels}
+    filters_dict = _tuple_to_nested_dict(filters)
+    for metric_name_regex, label_filters in filters_dict.items():
+        if not re.match(metric_name_regex, metric_name):
+            continue
+
+        should_capture = True
+
+        for label, label_filter in label_filters.items():
+            if not re.match(label_filter, metric_labels_dict.get(label, "")):
+                should_capture = False
+                break
+        break
+
+    return should_capture
+
+
+class OpenMetricsMetric:
+    """Container for all the COUNTER and GAUGE metrics extracted from an OpenMetrics endpoint."""
+
+    def __init__(
+        self, name: str, url: str, filters: Mapping[str, Mapping[str, str]]
+    ) -> None:
+        self.name = name
+        self.url = url
+        self.filters = filters
+        self.filters_tuple = _nested_dict_to_tuple(filters)
+        self._session: Optional["requests.Session"] = None
+        self.samples: "Deque[dict]" = deque([])
+        # {"<metric name>": {"<labels hash>": <index>}}
+        self.label_map: "Dict[str, Dict[str, int]]" = defaultdict(dict)
+        # {"<labels hash>": <labels>}
+        self.label_hashes: "Dict[str, dict]" = {}
+
+    def setup(self) -> None:
+        if self._session is not None:
+            return
+
+        self._session = _setup_requests_session()
+
+    def teardown(self) -> None:
+        if self._session is None:
+            return
+
+        self._session.close()
+        self._session = None
+
+    def parse_open_metrics_endpoint(self) -> Dict[str, Union[str, int, float]]:
+        assert prometheus_client_parser is not None
+        assert self._session is not None
+
+        response = self._session.get(self.url, timeout=_REQUEST_TIMEOUT)
+        response.raise_for_status()
+
+        text = response.text
+        measurement = {}
+        for family in prometheus_client_parser.text_string_to_metric_families(text):
+            if family.type not in ("counter", "gauge"):
+                # todo: add support for other metric types?
+                # todo: log warning about that?
+                continue
+            for sample in family.samples:
+                name, labels, value = sample.name, sample.labels, sample.value
+
+                if not _should_capture_metric(
+                    name,
+                    tuple(labels.items()),
+                    self.filters_tuple,
+                ):
+                    continue
+
+                # md5 hash of the labels
+                label_hash = md5(str(labels).encode("utf-8")).hexdigest()
+                if label_hash not in self.label_map[name]:
+                    # store the index of the label hash in the label map
+                    self.label_map[name][label_hash] = len(self.label_map[name])
+                    # store the labels themselves
+                    self.label_hashes[label_hash] = labels
+                index = self.label_map[name][label_hash]
+                measurement[f"{name}.{index}"] = value
+
+        return measurement
+
+    def sample(self) -> None:
+        s = self.parse_open_metrics_endpoint()
+        self.samples.append(s)
+
+    def clear(self) -> None:
+        self.samples.clear()
+
+    def aggregate(self) -> dict:
+        if not self.samples:
+            return {}
+
+        prefix = f"{_PREFIX}.{self.name}."
+
+        stats = {}
+        for key in self.samples[0].keys():
+            samples = [s[key] for s in self.samples if key in s]
+            if samples and all(isinstance(s, (int, float)) for s in samples):
+                stats[f"{prefix}{key}"] = aggregate_mean(samples)
+            else:
+                stats[f"{prefix}{key}"] = aggregate_last(samples)
+        return stats
+
+
+class OpenMetrics:
+    # Poll an OpenMetrics endpoint, parse the response and return a dict of metrics
+    # Implements the same Protocol interface as Asset
+
+    def __init__(
+        self,
+        interface: "Interface",
+        settings: "SettingsStatic",
+        shutdown_event: mp.synchronize.Event,
+        name: str,
+        url: str,
+    ) -> None:
+        self.name = name
+        self.url = url
+        self.interface = interface
+        self.settings = settings
+        self.shutdown_event = shutdown_event
+
+        self.metrics: List[Metric] = [
+            OpenMetricsMetric(name, url, settings._stats_open_metrics_filters)
+        ]
+
+        self.metrics_monitor: "MetricsMonitor" = MetricsMonitor(
+            asset_name=self.name,
+            metrics=self.metrics,
+            interface=interface,
+            settings=settings,
+            shutdown_event=shutdown_event,
+        )
+
+        telemetry_record = telemetry.TelemetryRecord()
+        telemetry_record.feature.open_metrics = True
+        interface._publish_telemetry(telemetry_record)
+
+    @classmethod
+    def is_available(cls, url: str) -> bool:
+        _is_available: bool = False
+
+        ret = prometheus_client_parser is not None
+        if not ret:
+            wandb.termwarn(
+                "Monitoring OpenMetrics endpoints requires the `prometheus_client` package. "
+                "To install it, run `pip install prometheus_client`.",
+                repeat=False,
+            )
+            return _is_available
+        # check if the endpoint is available and is a valid OpenMetrics endpoint
+        _session: Optional[requests.Session] = None
+        try:
+            assert prometheus_client_parser is not None
+            _session = _setup_requests_session()
+            response = _session.get(url, timeout=_REQUEST_TIMEOUT)
+            response.raise_for_status()
+
+            # check if the response is a valid OpenMetrics response
+            # text_string_to_metric_families returns a generator
+            if list(
+                prometheus_client_parser.text_string_to_metric_families(response.text)
+            ):
+                _is_available = True
+        except Exception as e:
+            logger.debug(
+                f"OpenMetrics endpoint {url} is not available: {e}", exc_info=True
+            )
+
+        if _session is not None:
+            try:
+                _session.close()
+            except Exception:
+                pass
+        return _is_available
+
+    def start(self) -> None:
+        self.metrics_monitor.start()
+
+    def finish(self) -> None:
+        self.metrics_monitor.finish()
+
+    def probe(self) -> dict:
+        # todo: also return self.label_hashes
+        return {self.name: self.url}

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

@@ -17,9 +17,7 @@ logger = logging.getLogger(__name__)
 
 
 class TPUUtilization:
-    """
-    Google Cloud TPU utilization in percent.
-    """
+    """Google Cloud TPU utilization in percent."""
 
     name = "tpu"
     samples: "Deque[float]"
@@ -130,7 +128,6 @@ class TPU:
 
     @classmethod
     def is_available(cls) -> bool:
-
         if os.environ.get("TPU_NAME", False) is False:
             return False