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

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

@@ -1,393 +0,0 @@
-import collections
-import dataclasses
-import json
-import logging
-import os
-import pathlib
-import shutil
-import subprocess
-import tempfile
-import threading
-import time
-from collections import deque
-from typing import TYPE_CHECKING, Any, Dict, Final, List, Optional, Tuple, 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__)
-
-
-NEURON_MONITOR_DEFAULT_CONFIG: Final[dict] = {
-    "period": "1s",
-    "neuron_runtimes": [
-        {
-            "tag_filter": ".*",
-            "metrics": [
-                {"type": "neuroncore_counters"},
-                {"type": "memory_used"},
-                {"type": "neuron_runtime_vcpu_usage"},
-                # {"type": "execution_stats"},
-            ],
-        }
-    ],
-    "system_metrics": [
-        {"type": "vcpu_usage"},
-        {"type": "memory_info"},
-        {"type": "neuron_hw_counters"},
-    ],
-}
-
-# todo: once a python sdk is released with the Neuron utils, rewrite this
-NEURON_LS_COMMAND: Final[Tuple[str, str]] = (
-    shutil.which("neuron-ls") or "/opt/aws/neuron/bin/neuron-ls",
-    "-j",
-)
-NEURON_MONITOR_PATH: Final[str] = (
-    shutil.which("neuron-monitor") or "/opt/aws/neuron/bin/neuron-monitor"
-)
-
-
-@dataclasses.dataclass
-class _NeuronCoreMemoryUsage:
-    constants: int
-    model_code: int
-    model_shared_scratchpad: int
-    runtime_memory: int
-    tensors: int
-
-
-@dataclasses.dataclass
-class _HostMemoryUsage:
-    application_memory: int
-    constants: int
-    dma_buffers: int
-    tensors: int
-
-
-@dataclasses.dataclass
-class _Stats:
-    neuroncore_utilization: Dict[int, float]  # per neuron core utilization
-    host_total_memory_usage: int  # total memory usage in bytes
-    neuron_device_total_memory_usage: int  # total memory usage
-    host_memory_usage: _HostMemoryUsage  # host memory usage breakdown
-    neuroncore_memory_usage: Dict[
-        int, _NeuronCoreMemoryUsage
-    ]  # per core memory usage breakdown
-
-
-class NeuronCoreStats:
-    """AWS Trainium stats."""
-
-    name: str = "trn.{key}"
-    samples: "Deque[_Stats]"
-
-    def write_neuron_monitor_config(self) -> None:
-        """Write neuron monitor config file."""
-        # mkdir if not exists
-        pathlib.Path(self.neuron_monitor_config_path).parent.mkdir(
-            parents=True, exist_ok=True
-        )
-        # write default config
-        with open(self.neuron_monitor_config_path, "w") as f:
-            json.dump(NEURON_MONITOR_DEFAULT_CONFIG, f, indent=4)
-
-    def neuron_monitor(self) -> None:
-        """Run neuron-monitor in a separate process to collect raw data."""
-        self.write_neuron_monitor_config()
-
-        try:
-            command = [
-                NEURON_MONITOR_PATH,
-                "-c",
-                self.neuron_monitor_config_path,
-            ]
-            with subprocess.Popen(
-                command,
-                stdout=subprocess.PIPE,
-                stderr=None,
-            ) as process:
-                while not self.shutdown_event.is_set():
-                    if process.stdout is None:
-                        self.shutdown_event.wait(0.1)
-                        continue
-
-                    raw_data = process.stdout.readline()
-                    if raw_data:
-                        self.raw_samples.append(raw_data)
-                process.kill()
-                process.wait()
-        except Exception:
-            logger.exception("neuron-monitor failed")
-
-    def __init__(
-        self,
-        pid: int,
-        neuron_monitor_config_path: Optional[str],
-    ) -> None:
-        self.pid = pid
-        # neuron-monitor requires a config file (json)
-        # we provide an option to supply a custom config file path
-        # in case the default temp file path is not writable
-        self.neuron_monitor_config_path = (
-            neuron_monitor_config_path or tempfile.NamedTemporaryFile(delete=False).name
-        )
-        self.raw_samples: Deque[bytes] = deque(maxlen=10)
-        self.samples: Deque[_Stats] = deque()
-        self.shutdown_event = threading.Event()
-
-        self.neuron_monitor_thread: Optional[threading.Thread] = None
-
-    def setup(self) -> None:
-        """Start the neuron-monitor thread for collecting raw data."""
-        if self.neuron_monitor_thread is not None:
-            return
-
-        logger.debug("Starting neuron-monitor thread")
-        self.shutdown_event.clear()
-        self.neuron_monitor_thread = threading.Thread(
-            name="NeuronCoreMntr",
-            target=self.neuron_monitor,
-            daemon=True,
-        )
-        self.neuron_monitor_thread.start()
-
-    def teardown(self) -> None:
-        """Stop the neuron-monitor thread."""
-        logger.debug("Stopping neuron-monitor thread")
-        try:
-            self.shutdown_event.set()
-            assert self.neuron_monitor_thread is not None
-            self.neuron_monitor_thread.join()
-        except Exception:
-            logger.exception("neuron-monitor thread failed to stop")
-        finally:
-            self.neuron_monitor_thread = None
-
-    def _is_matching_entry(self, entry: dict) -> bool:
-        """Check if the entry should be saved.
-
-        Checks if the pid in the entry matches the pid of the process.
-        If not (as in the case of multi-process training with torchrun),
-        checks if the LOCAL_RANK environment variable is set.
-
-        todo: add matching by neuron_runtime_tag
-        """
-        return (int(entry["pid"]) == int(self.pid)) or "LOCAL_RANK" in os.environ
-
-    def sample(self) -> None:
-        try:
-            raw_stats = json.loads(self.raw_samples[-1])
-            neuron_runtime_data = [
-                entry["report"]
-                for entry in raw_stats["neuron_runtime_data"]
-                if self._is_matching_entry(entry)
-            ][0]  # there should be only one entry with the pid
-
-            neuroncores_in_use = neuron_runtime_data["neuroncore_counters"][
-                "neuroncores_in_use"
-            ]
-            # per-core utilization stats:
-            neuroncore_utilization = {
-                int(k): v["neuroncore_utilization"]
-                for k, v in neuroncores_in_use.items()
-            }
-            # memory usage
-            neuron_runtime_used_bytes = neuron_runtime_data["memory_used"][
-                "neuron_runtime_used_bytes"
-            ]
-            # memory usage totals
-            host_total_memory_usage = neuron_runtime_used_bytes["host"]
-            neuron_device_total_memory_usage = neuron_runtime_used_bytes[
-                "neuron_device"
-            ]
-            # memory usage breakdown
-            usage_breakdown = neuron_runtime_used_bytes["usage_breakdown"]
-            host_memory_usage = _HostMemoryUsage(**usage_breakdown["host"])
-            neuroncore_memory_usage = {
-                int(k): _NeuronCoreMemoryUsage(**v)
-                for k, v in usage_breakdown["neuroncore_memory_usage"].items()
-            }
-
-            # When the training script is executed with torchrun,
-            # we only want to keep the relevant LOCAL_RANK stats
-            local_rank = int(os.environ.get("LOCAL_RANK", -1337))
-            if local_rank >= 0:
-                neuroncore_utilization = {
-                    local_rank: neuroncore_utilization[local_rank]
-                }
-                neuroncore_memory_usage = {
-                    local_rank: neuroncore_memory_usage[local_rank]
-                }
-
-            stats: _Stats = _Stats(
-                neuroncore_utilization=neuroncore_utilization,
-                host_total_memory_usage=host_total_memory_usage,
-                neuron_device_total_memory_usage=neuron_device_total_memory_usage,
-                host_memory_usage=host_memory_usage,
-                neuroncore_memory_usage=neuroncore_memory_usage,
-            )
-            self.samples.append(stats)
-
-        except Exception as e:  # noqa
-            pass
-
-    def clear(self) -> None:
-        self.samples.clear()
-
-    @staticmethod
-    def flatten_stats(sample: _Stats) -> dict:
-        """Flatten _Stats object into a flat dict of numbers."""
-        flattened = {}
-
-        def helper(key: str, value: Any) -> None:
-            if isinstance(value, (int, float)):
-                ret = {f"{key}": value}
-                flattened.update(ret)
-                return
-            elif isinstance(value, dict):
-                for kk, vv in value.items():
-                    if isinstance(kk, int):
-                        # top-level keys are neuron core ids,
-                        # so we swap the order to comply with the
-                        # frontend expectations
-                        helper(f"{kk}.{key}", vv)
-                    else:
-                        helper(f"{key}.{kk}", vv)
-                return
-            elif isinstance(value, list):
-                for i, val in enumerate(value):
-                    helper(f"{i}.{key}", val)
-
-        for kkk, vvv in dataclasses.asdict(sample).items():
-            helper(kkk, vvv)
-
-        return flattened
-
-    def aggregate(self) -> dict:
-        if not self.samples:
-            return {}
-
-        stats = {}
-
-        # Stats could be: numbers or dataclass objects or lists of such.
-        # In the latter case that means per-core stats.
-        # The dataclass objects are flat containers of numbers.
-
-        # flatten samples and merge the corresponding values into lists
-        merged_samples: Dict[str, List[Union[int, float]]] = collections.defaultdict(
-            list
-        )
-        for flattened_sample in (self.flatten_stats(sample) for sample in self.samples):
-            for k, v in flattened_sample.items():
-                merged_samples[k].append(v)
-
-        # aggregate the lists
-        for k, v in merged_samples.items():
-            stats[self.name.format(key=k)] = aggregate_mean(v)
-
-        return stats
-
-
-@asset_registry.register
-class Trainium:
-    def __init__(
-        self,
-        interface: "Interface",
-        settings: "SettingsStatic",
-        shutdown_event: threading.Event,
-    ) -> None:
-        self.name = self.__class__.__name__.lower()
-        self.metrics: List[Metric] = [
-            NeuronCoreStats(
-                settings.x_stats_pid,
-                settings.x_stats_neuron_monitor_config_path,
-            ),
-        ]
-        self.metrics_monitor = MetricsMonitor(
-            self.name,
-            self.metrics,
-            interface,
-            settings,
-            shutdown_event,
-        )
-        telemetry_record = telemetry.TelemetryRecord()
-        telemetry_record.env.trainium = True
-        interface._publish_telemetry(telemetry_record)
-
-    @classmethod
-    def is_available(cls) -> bool:
-        # todo: check if neuron-ls is available and if yes, what it reports. see:
-        # https://awsdocs-neuron.readthedocs-hosted.com/en/latest/tools/neuron-sys-tools/neuron-ls.html
-        if not pathlib.Path(NEURON_LS_COMMAND[0]).exists():
-            return False
-        # need to be extra careful as neuron tools could be pre-installed
-        # on some systems that do not have the hardware
-        try:
-            # redirect stderr to null to avoid printing errors to the console
-            # todo: alternative: check /dev/neuron0 ? sysfs support coming soon in neuron tools
-            output = subprocess.check_output(
-                NEURON_LS_COMMAND,
-                universal_newlines=True,
-                stderr=subprocess.DEVNULL,
-            ).strip()
-            if len(json.loads(output)) > 0:
-                return True
-        except (OSError, ValueError, TypeError, subprocess.CalledProcessError):
-            pass
-
-        return False
-
-    def start(self) -> None:
-        self.metrics_monitor.start()
-
-    def finish(self) -> None:
-        self.metrics_monitor.finish()
-
-    def probe(self) -> dict:
-        try:
-            self.metrics[0].check_neuron_monitor_config()  # type: ignore
-            neuron_hardware_info: dict = {}
-            command = [
-                NEURON_MONITOR_PATH,
-                "-c",
-                self.metrics[0].neuron_monitor_config_path,  # type: ignore
-            ]
-            with subprocess.Popen(
-                command,
-                stdout=subprocess.PIPE,
-                stderr=None,
-            ) as process:
-                while True:
-                    if process.stdout is None:
-                        time.sleep(0.1)
-                        continue
-
-                    raw_data = process.stdout.readline()
-                    if raw_data:
-                        parsed_data = json.loads(raw_data)
-                        neuron_hardware_info = parsed_data.get(
-                            "neuron_hardware_info", {}
-                        )
-                        neuron_hardware_info.pop("error", None)
-                        break
-
-            try:
-                process.kill()
-                process.wait()
-            except:  # noqa
-                pass
-
-            return {self.name: neuron_hardware_info}
-        except Exception:
-            logger.exception("neuron-monitor failed")
-            return {}

wandb/sdk/internal/system/env_probe_helpers.py

@@ -1,13 +0,0 @@
-import logging
-
-from sentry_sdk.integrations.aws_lambda import get_lambda_bootstrap  # type: ignore
-
-logger = logging.getLogger(__name__)
-
-
-def is_aws_lambda() -> bool:
-    """Check if we are running in a lambda environment."""
-    lambda_bootstrap = get_lambda_bootstrap()
-    if not lambda_bootstrap or not hasattr(lambda_bootstrap, "handle_event_request"):
-        return False
-    return True

wandb/sdk/internal/system/system_info.py

@@ -1,248 +0,0 @@
-# Information about the system and the environment
-import datetime
-import glob
-import json
-import logging
-import os
-import subprocess
-import sys
-from shutil import copyfile
-from typing import Any, Dict, List, Optional
-from urllib.parse import unquote
-
-from wandb.sdk.internal.settings_static import SettingsStatic
-from wandb.sdk.lib import filesystem
-from wandb.sdk.lib.filenames import CONDA_ENVIRONMENTS_FNAME, DIFF_FNAME, METADATA_FNAME
-from wandb.sdk.lib.gitlib import GitRepo
-
-from .assets.interfaces import Interface
-
-logger = logging.getLogger(__name__)
-
-
-class SystemInfo:
-    # todo: this is mostly a copy of the legacy Meta class, but it should be refactored
-    def __init__(self, settings: SettingsStatic, interface: Interface) -> None:
-        logger.debug("System info init")
-        self.settings = settings
-
-        self.metadata_file_name = os.path.join(self.settings.files_dir, METADATA_FNAME)
-        self.backend_interface = interface
-        self.git = GitRepo(
-            root=self.settings.git_root,
-            remote=self.settings.git_remote,
-            remote_url=self.settings.git_remote_url,
-            commit=self.settings.git_commit,
-        )
-        # Location under "code" directory in files where program was saved.
-        self.saved_program: Optional[os.PathLike] = None
-        # Locations under files directory where diff patches were saved.
-        self.saved_patches: List[str] = []
-        logger.debug("System info init done")
-
-    def _save_code(self) -> None:
-        logger.debug("Saving code")
-        if not self.settings.program_relpath:
-            logger.warning("unable to save code -- program entry not found")
-            return None
-
-        root: str = self.git.root or os.getcwd()
-        program_relative: str = self.settings.program_relpath
-        filesystem.mkdir_exists_ok(
-            os.path.join(
-                self.settings.files_dir, "code", os.path.dirname(program_relative)
-            )
-        )
-        program_absolute = os.path.join(root, program_relative)
-        if not os.path.exists(program_absolute):
-            logger.warning(f"unable to save code -- can't find {program_absolute}")
-            return None
-        saved_program = os.path.join(self.settings.files_dir, "code", program_relative)
-        self.saved_program = program_relative  # type: ignore
-
-        if not os.path.exists(saved_program):
-            copyfile(program_absolute, saved_program)
-        logger.debug("Saving code done")
-
-    def _save_patches(self) -> None:
-        """Save the current state of this repository to one or more patches.
-
-        Makes one patch against HEAD and another one against the most recent
-        commit that occurs in an upstream branch. This way we can be robust
-        to history editing as long as the user never does "push -f" to break
-        history on an upstream branch.
-
-        Writes the first patch to <files_dir>/<DIFF_FNAME> and the second to
-        <files_dir>/upstream_diff_<commit_id>.patch.
-
-        """
-        if not self.git.enabled:
-            return None
-
-        logger.debug("Saving git patches")
-        try:
-            root = self.git.root
-            diff_args = ["git", "diff"]
-            if self.git.has_submodule_diff:
-                diff_args.append("--submodule=diff")
-
-            if self.git.dirty:
-                patch_path = os.path.join(self.settings.files_dir, DIFF_FNAME)
-                with open(patch_path, "wb") as patch:
-                    # we diff against HEAD to ensure we get changes in the index
-                    subprocess.check_call(
-                        diff_args + ["HEAD"], stdout=patch, cwd=root, timeout=5
-                    )
-                    self.saved_patches.append(
-                        os.path.relpath(patch_path, start=self.settings.files_dir)
-                    )
-
-            upstream_commit = self.git.get_upstream_fork_point()
-            if upstream_commit and upstream_commit != self.git.repo.head.commit:  # type: ignore
-                sha = upstream_commit.hexsha
-                upstream_patch_path = os.path.join(
-                    self.settings.files_dir, f"upstream_diff_{sha}.patch"
-                )
-                with open(upstream_patch_path, "wb") as upstream_patch:
-                    subprocess.check_call(
-                        diff_args + [sha], stdout=upstream_patch, cwd=root, timeout=5
-                    )
-                    self.saved_patches.append(
-                        os.path.relpath(
-                            upstream_patch_path, start=self.settings.files_dir
-                        )
-                    )
-        # TODO: A customer saw `ValueError: Reference at 'refs/remotes/origin/foo'
-        #  does not exist` so we now catch ValueError. Catching this error feels
-        #  too generic.
-        except (
-            ValueError,
-            subprocess.CalledProcessError,
-            subprocess.TimeoutExpired,
-        ):
-            logger.exception("Error generating diff.")
-        logger.debug("Saving git patches done")
-
-    def _probe_git(self, data: Dict[str, Any]) -> Dict[str, Any]:
-        if self.settings.disable_git or self.settings.x_disable_machine_info:
-            return data
-
-        # in case of manually passing the git repo info, `enabled` would be False,
-        # but we still want to save the git repo info
-        if not self.git.enabled and self.git.auto:
-            return data
-
-        logger.debug("Probing git")
-
-        data["git"] = {
-            "remote": self.git.remote_url,
-            "commit": self.git.last_commit,
-        }
-        data["email"] = self.git.email
-        data["root"] = self.git.root or data.get("root") or os.getcwd()
-        logger.debug("Probing git done")
-
-        return data
-
-    def probe(self) -> Dict[str, Any]:
-        """Probe the system for information about the current environment."""
-        # todo: refactor this quality code 🤮🤮🤮🤮🤮
-        logger.debug("Probing system")
-        data: Dict[str, Any] = dict()
-
-        data["os"] = self.settings._os
-        data["python"] = self.settings._python
-        data["heartbeatAt"] = datetime.datetime.utcnow().isoformat()
-        data["startedAt"] = (
-            datetime.datetime.utcfromtimestamp(self.settings.x_start_time).isoformat()
-            if self.settings.x_start_time
-            else None
-        )
-
-        data["docker"] = self.settings.docker
-
-        data["args"] = tuple(self.settings._args or ())
-        data["state"] = "running"
-
-        if self.settings.program is not None:
-            data["program"] = self.settings.program
-            # Used during artifact-job creation, always points to the relpath
-            # of code execution, even when in a git repo
-            data["codePathLocal"] = self.settings._code_path_local
-        if not (self.settings.disable_code or self.settings.x_disable_machine_info):
-            if self.settings.program_relpath:
-                data["codePath"] = self.settings.program_relpath
-            elif self.settings._jupyter:
-                if self.settings.notebook_name:
-                    data["program"] = self.settings.notebook_name
-                elif self.settings.x_jupyter_path:
-                    if self.settings.x_jupyter_path.startswith("fileId="):
-                        unescaped = unquote(self.settings.x_jupyter_path)
-                        data["colab"] = (
-                            "https://colab.research.google.com/notebook#" + unescaped
-                        )
-                        data["program"] = self.settings.x_jupyter_name
-                    else:
-                        data["program"] = self.settings.x_jupyter_path
-                        data["root"] = self.settings.x_jupyter_root
-            # get the git repo info
-            data = self._probe_git(data)
-
-        if self.settings.anonymous not in ["allow", "must"]:
-            data["host"] = self.settings.host
-            data["username"] = self.settings.username
-            data["executable"] = sys.executable
-        else:
-            data.pop("email", None)
-            data.pop("root", None)
-
-        logger.debug("Probing system done")
-
-        return data
-
-    def _save_conda(self) -> None:
-        current_shell_is_conda = os.path.exists(os.path.join(sys.prefix, "conda-meta"))
-        if not current_shell_is_conda:
-            return None
-
-        logger.debug(
-            "Saving list of conda packages installed into the current environment"
-        )
-        try:
-            with open(
-                os.path.join(self.settings.files_dir, CONDA_ENVIRONMENTS_FNAME), "w"
-            ) as f:
-                subprocess.call(
-                    ["conda", "env", "export"],
-                    stdout=f,
-                    stderr=subprocess.DEVNULL,
-                    timeout=15,  # add timeout since conda env export could take a really long time
-                )
-        except Exception:
-            logger.exception("Error saving conda packages")
-        logger.debug("Saving conda packages done")
-
-    def publish(self, system_info: dict) -> None:
-        # save pip, conda, code patches to disk
-        if self.settings.x_save_requirements:
-            self._save_conda()
-        if self.settings.save_code:
-            self._save_code()
-            self._save_patches()
-
-        # save system_info to disk
-        with open(self.metadata_file_name, "w") as f:
-            s = json.dumps(system_info, indent=4)
-            f.write(s)
-            f.write("\n")
-        base_name = os.path.basename(self.metadata_file_name)
-        files = dict(files=[(base_name, "now")])
-
-        if self.saved_program:
-            saved_program = os.path.join("code", self.saved_program)
-            files["files"].append((glob.escape(saved_program), "now"))
-        for patch in self.saved_patches:
-            files["files"].append((glob.escape(patch), "now"))
-
-        # publish files to the backend
-        self.backend_interface.publish_files(files)  # type: ignore