wandb 0.19.1rc1__py3-none-musllinux_1_2_aarch64.whl

wandb/sdk/internal/sender_config.py

@@ -0,0 +1,197 @@
+import json
+from typing import Any, Dict, NewType, Optional, Sequence
+
+from wandb.proto import wandb_internal_pb2
+from wandb.sdk.lib import proto_util, telemetry
+
+BackendConfigDict = NewType("BackendConfigDict", Dict[str, Any])
+"""Run config dictionary in the format used by the backend."""
+
+_WANDB_INTERNAL_KEY = "_wandb"
+
+
+class ConfigState:
+    """The configuration of a run."""
+
+    def __init__(self, tree: Optional[Dict[str, Any]] = None) -> None:
+        self._tree: Dict[str, Any] = tree or {}
+        """A tree with string-valued nodes and JSON leaves.
+
+        Leaves are Python objects that are valid JSON values:
+
+        * Primitives like strings and numbers
+        * Dictionaries from strings to JSON objects
+        * Lists of JSON objects
+        """
+
+    def non_internal_config(self) -> Dict[str, Any]:
+        """Returns the config settings minus "_wandb"."""
+        return {k: v for k, v in self._tree.items() if k != _WANDB_INTERNAL_KEY}
+
+    def update_from_proto(
+        self,
+        config_record: wandb_internal_pb2.ConfigRecord,
+    ) -> None:
+        """Applies update and remove commands."""
+        for config_item in config_record.update:
+            self._update_at_path(
+                _key_path(config_item),
+                json.loads(config_item.value_json),
+            )
+
+        for config_item in config_record.remove:
+            self._delete_at_path(_key_path(config_item))
+
+    def merge_resumed_config(self, old_config_tree: Dict[str, Any]) -> None:
+        """Merges the config from a run that's being resumed."""
+        # Add any top-level keys that aren't already set.
+        self._add_unset_keys_from_subtree(old_config_tree, [])
+
+        # Unfortunately, when a user logs visualizations, we store them in the
+        # run's config. When resuming a run, we want to avoid erasing previously
+        # logged visualizations, hence this special handling:
+        self._add_unset_keys_from_subtree(
+            old_config_tree,
+            [_WANDB_INTERNAL_KEY, "visualize"],
+        )
+        self._add_unset_keys_from_subtree(
+            old_config_tree,
+            [_WANDB_INTERNAL_KEY, "viz"],
+        )
+
+    def _add_unset_keys_from_subtree(
+        self,
+        old_config_tree: Dict[str, Any],
+        path: Sequence[str],
+    ) -> None:
+        """Uses the given subtree for keys that aren't already set."""
+        old_subtree = _subtree(old_config_tree, path, create=False)
+        if not old_subtree:
+            return
+
+        new_subtree = _subtree(self._tree, path, create=True)
+        assert new_subtree is not None
+
+        for key, value in old_subtree.items():
+            if key not in new_subtree:
+                new_subtree[key] = value
+
+    def to_backend_dict(
+        self,
+        telemetry_record: telemetry.TelemetryRecord,
+        framework: Optional[str],
+        start_time_millis: int,
+        metric_pbdicts: Sequence[Dict[int, Any]],
+    ) -> BackendConfigDict:
+        """Returns a dictionary representation expected by the backend.
+
+        The backend expects the configuration in a specific format, and the
+        config is also used to store additional metadata about the run.
+
+        Args:
+            telemetry_record: Telemetry information to insert.
+            framework: The detected framework used in the run (e.g. TensorFlow).
+            start_time_millis: The run's start time in Unix milliseconds.
+            metric_pbdicts: List of dict representations of metric protobuffers.
+        """
+        backend_dict = self._tree.copy()
+        wandb_internal = backend_dict.setdefault(_WANDB_INTERNAL_KEY, {})
+
+        ###################################################
+        # Telemetry information
+        ###################################################
+        py_version = telemetry_record.python_version
+        if py_version:
+            wandb_internal["python_version"] = py_version
+
+        cli_version = telemetry_record.cli_version
+        if cli_version:
+            wandb_internal["cli_version"] = cli_version
+
+        if framework:
+            wandb_internal["framework"] = framework
+
+        huggingface_version = telemetry_record.huggingface_version
+        if huggingface_version:
+            wandb_internal["huggingface_version"] = huggingface_version
+
+        wandb_internal["is_jupyter_run"] = telemetry_record.env.jupyter
+        wandb_internal["is_kaggle_kernel"] = telemetry_record.env.kaggle
+        wandb_internal["start_time"] = start_time_millis
+
+        # The full telemetry record.
+        wandb_internal["t"] = proto_util.proto_encode_to_dict(telemetry_record)
+
+        ###################################################
+        # Metrics
+        ###################################################
+        if metric_pbdicts:
+            wandb_internal["m"] = metric_pbdicts
+
+        return BackendConfigDict(
+            {
+                key: {
+                    # Configurations can be stored in a hand-written YAML file,
+                    # and users can add descriptions to their hyperparameters
+                    # there. However, we don't support a way to set descriptions
+                    # via code, so this is always None.
+                    "desc": None,
+                    "value": value,
+                }
+                for key, value in self._tree.items()
+            }
+        )
+
+    def _update_at_path(
+        self,
+        key_path: Sequence[str],
+        value: Any,
+    ) -> None:
+        """Sets the value at the path in the config tree."""
+        subtree = _subtree(self._tree, key_path[:-1], create=True)
+        assert subtree is not None
+
+        subtree[key_path[-1]] = value
+
+    def _delete_at_path(
+        self,
+        key_path: Sequence[str],
+    ) -> None:
+        """Removes the subtree at the path in the config tree."""
+        subtree = _subtree(self._tree, key_path[:-1], create=False)
+        if subtree:
+            del subtree[key_path[-1]]
+
+
+def _key_path(config_item: wandb_internal_pb2.ConfigItem) -> Sequence[str]:
+    """Returns the key path referenced by the config item."""
+    if config_item.nested_key:
+        return config_item.nested_key
+    elif config_item.key:
+        return [config_item.key]
+    else:
+        raise AssertionError(
+            "Invalid ConfigItem: either key or nested_key must be set",
+        )
+
+
+def _subtree(
+    tree: Dict[str, Any],
+    key_path: Sequence[str],
+    *,
+    create: bool = False,
+) -> Optional[Dict[str, Any]]:
+    """Returns a subtree at the given path."""
+    for key in key_path:
+        subtree = tree.get(key)
+
+        if not subtree:
+            if create:
+                subtree = {}
+                tree[key] = subtree
+            else:
+                return None
+
+        tree = subtree
+
+    return tree

wandb/sdk/internal/settings_static.py

@@ -0,0 +1,97 @@
+from typing import Any, Iterable
+
+from wandb.proto import wandb_settings_pb2
+from wandb.sdk.lib import RunMoment
+from wandb.sdk.wandb_settings import Settings
+
+
+class SettingsStatic(Settings):
+    """A readonly object that wraps a protobuf Settings message.
+
+    Implements the mapping protocol, so you can access settings as
+    attributes or items.
+    """
+
+    def __init__(self, proto: wandb_settings_pb2.Settings) -> None:
+        data = self._proto_to_dict(proto)
+        super().__init__(**data)
+
+    def _proto_to_dict(self, proto: wandb_settings_pb2.Settings) -> dict:
+        data = {}
+
+        forks_specified: list[str] = []
+        for key in Settings.model_fields:  # type: ignore [attr-defined]
+            value: Any = None
+
+            field_info = Settings.model_fields[key]
+            annotation = str(field_info.annotation)
+
+            if key == "_stats_open_metrics_filters":
+                # todo: it's an underscored field, refactor into
+                #  something more elegant?
+                # I'm really about this. It's ugly, but it works.
+                # Do not try to repeat this at home.
+                value_type = getattr(proto, key).WhichOneof("value")
+                if value_type == "sequence":
+                    value = list(getattr(proto, key).sequence.value)
+                elif value_type == "mapping":
+                    unpacked_mapping = {}
+                    for outer_key, outer_value in getattr(
+                        proto, key
+                    ).mapping.value.items():
+                        unpacked_inner = {}
+                        for inner_key, inner_value in outer_value.value.items():
+                            unpacked_inner[inner_key] = inner_value
+                        unpacked_mapping[outer_key] = unpacked_inner
+                    value = unpacked_mapping
+            elif key == "fork_from" or key == "resume_from":
+                value = getattr(proto, key)
+                if value.run:
+                    value = RunMoment(
+                        run=value.run, value=value.value, metric=value.metric
+                    )
+                    forks_specified.append(key)
+                else:
+                    value = None
+            else:
+                if proto.HasField(key):  # type: ignore [arg-type]
+                    value = getattr(proto, key).value
+                    # Convert to list if the field is a sequence
+                    if any(t in annotation for t in ("tuple", "Sequence", "list")):
+                        value = list(value)
+                else:
+                    value = None
+
+            if value is not None:
+                data[key] = value
+
+        if len(forks_specified) > 1:
+            raise ValueError(
+                "Only one of fork_from or resume_from can be specified, not both"
+            )
+
+        return data
+
+    def __setattr__(self, name: str, value: object) -> None:
+        raise AttributeError("Error: SettingsStatic is a readonly object")
+
+    def __setitem__(self, key: str, val: object) -> None:
+        raise AttributeError("Error: SettingsStatic is a readonly object")
+
+    def keys(self) -> "Iterable[str]":
+        return self.__dict__.keys()
+
+    def __getitem__(self, key: str) -> Any:
+        return self.__dict__[key]
+
+    def __getattr__(self, name: str) -> Any:
+        try:
+            return self.__dict__[name]
+        except KeyError:
+            raise AttributeError(f"SettingsStatic has no attribute {name}")
+
+    def __str__(self) -> str:
+        return str(self.__dict__)
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.__dict__

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

@@ -0,0 +1,25 @@
+__all__ = (
+    "asset_registry",
+    "CPU",
+    "Disk",
+    "GPU",
+    "GPUAMD",
+    "IPU",
+    "Memory",
+    "Network",
+    "OpenMetrics",
+    "TPU",
+    "Trainium",
+)
+
+from .asset_registry import asset_registry
+from .cpu import CPU
+from .disk import Disk
+from .gpu import GPU
+from .gpu_amd import GPUAMD
+from .ipu import IPU
+from .memory import Memory
+from .network import Network
+from .open_metrics import OpenMetrics
+from .tpu import TPU
+from .trainium import Trainium

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

@@ -0,0 +1,37 @@
+import sys
+from typing import Union
+
+if sys.version_info >= (3, 9):
+    from collections.abc import Sequence
+else:
+    from typing import Sequence
+
+Number = Union[int, float]
+
+
+def aggregate_mean(samples: Sequence[Number], precision: int = 2) -> float:
+    return round(sum(samples) / len(samples), precision)
+
+
+def aggregate_last(samples: Sequence[Number], precision: int = 2) -> Union[float, int]:
+    if isinstance(samples[-1], int):
+        return samples[-1]
+    return round(samples[-1], precision)
+
+
+def aggregate_max(samples: Sequence[Number], precision: int = 2) -> Union[float, int]:
+    if isinstance(samples[-1], int):
+        return max(samples)
+    return round(max(samples), precision)
+
+
+def aggregate_min(samples: Sequence[Number], precision: int = 2) -> Union[float, int]:
+    if isinstance(samples[-1], int):
+        return min(samples)
+    return round(min(samples), precision)
+
+
+def aggregate_sum(samples: Sequence[Number], precision: int = 2) -> Union[float, int]:
+    if isinstance(samples[-1], int):
+        return sum(samples)
+    return round(sum(samples), precision)

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

@@ -0,0 +1,20 @@
+from typing import Iterator, List, Type
+
+from .interfaces import Asset
+
+
+class AssetRegistry:
+    def __init__(self) -> None:
+        self._registry: List[Type[Asset]] = []
+
+    def register(self, asset: Type[Asset]) -> Type[Asset]:
+        self._registry.append(asset)
+        return asset
+
+    def __iter__(self) -> Iterator[Type[Asset]]:
+        for asset in self._registry:
+            if asset.is_available():
+                yield asset
+
+
+asset_registry = AssetRegistry()

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

@@ -0,0 +1,163 @@
+import threading
+from collections import deque
+from typing import TYPE_CHECKING, List, Optional
+
+try:
+    import psutil
+except ImportError:
+    psutil = None
+from .aggregators import aggregate_last, 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
+
+
+# CPU Metrics
+
+
+class ProcessCpuPercent:
+    """CPU usage of the process in percent normalized by the number of CPUs."""
+
+    # name = "process_cpu_percent"
+    name = "cpu"
+
+    def __init__(self, pid: int) -> None:
+        self.pid = pid
+        self.samples: Deque[float] = deque([])
+        self.process: Optional[psutil.Process] = None
+
+    def sample(self) -> None:
+        # todo: this is what we'd eventually want to do
+        # self.samples.append(
+        #     (
+        #         datetime.datetime.utcnow(),
+        #         psutil.Process(self.pid).cpu_percent(),
+        #     )
+        # )
+        if self.process is None:
+            self.process = psutil.Process(self.pid)
+
+        self.samples.append(self.process.cpu_percent() / psutil.cpu_count())
+
+    def clear(self) -> None:
+        self.samples.clear()
+
+    def aggregate(self) -> dict:
+        # todo: create a statistics class with helper methods to compute
+        #      mean, median, min, max, etc.
+        if not self.samples:
+            return {}
+        aggregate = aggregate_mean(self.samples)
+        return {self.name: aggregate}
+
+
+class CpuPercent:
+    """CPU usage of the system in percent per core."""
+
+    name = "cpu.{i}.cpu_percent"
+
+    def __init__(self, interval: Optional[float] = None) -> None:
+        self.samples: Deque[List[float]] = deque([])
+        self.interval = interval
+
+    def sample(self) -> None:
+        self.samples.append(psutil.cpu_percent(interval=self.interval, percpu=True))
+
+    def clear(self) -> None:
+        self.samples.clear()
+
+    def aggregate(self) -> dict:
+        if not self.samples:
+            return {}
+        num_cpu = len(self.samples[0])
+        cpu_metrics = {}
+        for i in range(num_cpu):
+            aggregate_i = aggregate_mean([sample[i] for sample in self.samples])
+            cpu_metrics[self.name.format(i=i)] = aggregate_i
+
+        return cpu_metrics
+
+
+class ProcessCpuThreads:
+    """Number of threads used by the process."""
+
+    name = "proc.cpu.threads"
+
+    def __init__(self, pid: int) -> None:
+        self.samples: Deque[int] = deque([])
+        self.pid = pid
+        self.process: Optional[psutil.Process] = None
+
+    def sample(self) -> None:
+        if self.process is None:
+            self.process = psutil.Process(self.pid)
+
+        self.samples.append(self.process.num_threads())
+
+    def clear(self) -> None:
+        self.samples.clear()
+
+    def aggregate(self) -> dict:
+        if not self.samples:
+            return {}
+        return {self.name: aggregate_last(self.samples)}
+
+
+@asset_registry.register
+class CPU:
+    def __init__(
+        self,
+        interface: "Interface",
+        settings: "SettingsStatic",
+        shutdown_event: threading.Event,
+    ) -> None:
+        self.name: str = self.__class__.__name__.lower()
+        self.metrics: List[Metric] = [
+            ProcessCpuPercent(settings.x_stats_pid),
+            CpuPercent(),
+            ProcessCpuThreads(settings.x_stats_pid),
+        ]
+        self.metrics_monitor: MetricsMonitor = MetricsMonitor(
+            self.name,
+            self.metrics,
+            interface,
+            settings,
+            shutdown_event,
+        )
+
+    @classmethod
+    def is_available(cls) -> bool:
+        return psutil is not None
+
+    def probe(self) -> dict:
+        asset_info = {
+            "cpu_count": psutil.cpu_count(logical=False),
+            "cpu_count_logical": psutil.cpu_count(logical=True),
+        }
+        try:
+            asset_info["cpu_freq"] = {
+                "current": psutil.cpu_freq().current,
+                "min": psutil.cpu_freq().min,
+                "max": psutil.cpu_freq().max,
+            }
+            asset_info["cpu_freq_per_core"] = [
+                {
+                    "current": freq.current,
+                    "min": freq.min,
+                    "max": freq.max,
+                }
+                for freq in psutil.cpu_freq(percpu=True)
+            ]
+        except Exception:
+            pass
+        return asset_info
+
+    def start(self) -> None:
+        self.metrics_monitor.start()
+
+    def finish(self) -> None:
+        self.metrics_monitor.finish()