wandb 0.18.2__py3-none-musllinux_1_2_x86_64.whl

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

@@ -0,0 +1,399 @@
+import collections
+import dataclasses
+import json
+import logging
+import os
+import pathlib
+import shutil
+import subprocess
+import sys
+import tempfile
+import threading
+import time
+from collections import deque
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
+
+if sys.version_info >= (3, 8):
+    from typing import Final
+else:
+    from typing_extensions import Final
+
+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 as e:
+            logger.error("neuron-monitor failed: {}".format(e))
+
+    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 as e:
+            logger.error("neuron-monitor thread failed to stop: {}".format(e))
+        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._stats_pid,
+                settings._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 as e:
+            logger.error("neuron-monitor failed: {}".format(e))
+            return {}

wandb/sdk/internal/system/env_probe_helpers.py

@@ -0,0 +1,13 @@
+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

@@ -0,0 +1,249 @@
+# 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(
+                "unable to save code -- can't find {}".format(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,
+        ) as e:
+            logger.error("Error generating diff: {}".format(e))
+        logger.debug("Saving git patches done")
+
+    def _probe_git(self, data: Dict[str, Any]) -> Dict[str, Any]:
+        if self.settings.disable_git:
+            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._start_time
+        ).isoformat()
+
+        data["docker"] = self.settings.docker
+
+        data["cuda"] = self.settings._cuda
+        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:
+            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._jupyter_path:
+                    if self.settings._jupyter_path.startswith("fileId="):
+                        unescaped = unquote(self.settings._jupyter_path)
+                        data["colab"] = (
+                            "https://colab.research.google.com/notebook#" + unescaped
+                        )
+                        data["program"] = self.settings._jupyter_name
+                    else:
+                        data["program"] = self.settings._jupyter_path
+                        data["root"] = self.settings._jupyter_root
+            # get the git repo info
+            data = self._probe_git(data)
+
+        if self.settings.anonymous != "true":
+            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 as e:
+            logger.exception(f"Error saving conda packages: {e}")
+        logger.debug("Saving conda packages done")
+
+    def publish(self, system_info: dict) -> None:
+        # save pip, conda, code patches to disk
+        if self.settings._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