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

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

@@ -92,9 +92,7 @@ class _Stats:
 
 
 class NeuronCoreStats:
-    """
-    AWS Trainium stats.
-    """
+    """AWS Trainium stats."""
 
     name: str = "trn.{key}"
     samples: "Deque[_Stats]"
@@ -124,7 +122,7 @@ class NeuronCoreStats:
             ) as process:
                 while not self.shutdown_event.is_set():
                     if process.stdout is None:
-                        time.sleep(0.1)
+                        self.shutdown_event.wait(0.1)
                         continue
 
                     raw_data = process.stdout.readline()
@@ -151,6 +149,15 @@ class NeuronCoreStats:
         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,
@@ -158,9 +165,20 @@ class NeuronCoreStats:
         )
         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: %s" % e)
+        finally:
+            self.neuron_monitor_thread = None
+
     def _is_matching_entry(self, entry: dict) -> bool:
-        """
-        For now, only check if the pid in the entry matches the pid of the process.
+        """For now, only check if the pid in the entry matches the pid of the process.
 
         todo: add matching by neuron_runtime_tag
         """
@@ -218,9 +236,7 @@ class NeuronCoreStats:
 
     @staticmethod
     def flatten_stats(sample: _Stats) -> dict:
-        """
-        Flatten _Stats object into a flat dict of numbers.
-        """
+        """Flatten _Stats object into a flat dict of numbers."""
         flattened = {}
 
         def helper(key: str, value: Any) -> None:
@@ -302,6 +318,7 @@ class Trainium:
         # 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,
@@ -319,11 +336,6 @@ class Trainium:
 
     def finish(self) -> None:
         self.metrics_monitor.finish()
-        # stop the raw data acquisition threads
-        for metric in self.metrics:
-            if hasattr(metric, "shutdown_event"):
-                logger.debug("Stopping neuron-monitor thread")
-                metric.shutdown_event.set()
 
     def probe(self) -> dict:
         try:

wandb/sdk/internal/system/system_info.py

@@ -47,7 +47,7 @@ class SystemInfo:
 
     # todo: refactor these _save_* methods
     def _save_pip(self) -> None:
-        """Saves the current working set of pip packages to {REQUIREMENTS_FNAME}"""
+        """Save the current working set of pip packages to {REQUIREMENTS_FNAME}."""
         logger.debug(
             "Saving list of pip packages installed into the current environment"
         )
@@ -220,8 +220,7 @@ class SystemInfo:
                     if self.settings._jupyter_path.startswith("fileId="):
                         unescaped = unquote(self.settings._jupyter_path)
                         data["colab"] = (
-                            "https://colab.research.google.com/notebook#"
-                            + unescaped  # noqa
+                            "https://colab.research.google.com/notebook#" + unescaped
                         )
                         data["program"] = self.settings._jupyter_name
                     else:

wandb/sdk/internal/system/system_monitor.py

@@ -6,6 +6,7 @@ from typing import TYPE_CHECKING, List, Optional, Union
 
 from .assets.asset_registry import asset_registry
 from .assets.interfaces import Asset, Interface
+from .assets.open_metrics import OpenMetrics
 from .system_info import SystemInfo
 
 if TYPE_CHECKING:
@@ -47,6 +48,8 @@ class SystemMonitor:
         self._shutdown_event: mp.synchronize.Event = mp.Event()
         self._process: Optional[Union[mp.Process, threading.Thread]] = None
 
+        self.settings = settings
+
         # settings._stats_join_assets controls whether we should join stats from different assets
         # before publishing them to the backend. If set to False, we will publish stats from each
         # asset separately, using the backend interface. If set to True, we will aggregate stats from
@@ -59,14 +62,16 @@ class SystemMonitor:
         sampling_interval: float = float(
             max(
                 0.1,
-                settings._stats_sample_rate_seconds,
+                self.settings._stats_sample_rate_seconds,
             )
         )  # seconds
         # The number of samples to aggregate (e.g. average or compute max/min etc.)
         # before publishing; defaults to 15; valid range: [1:30]
-        samples_to_aggregate: int = min(30, max(1, settings._stats_samples_to_average))
+        samples_to_aggregate: int = min(
+            30, max(1, self.settings._stats_samples_to_average)
+        )
         self.publishing_interval: float = sampling_interval * samples_to_aggregate
-        self.join_assets: bool = settings._stats_join_assets
+        self.join_assets: bool = self.settings._stats_join_assets
 
         self.backend_interface = interface
         self.asset_interface: Optional[AssetInterface] = (
@@ -74,21 +79,47 @@ class SystemMonitor:
         )
 
         # hardware assets
-        self.assets: List["Asset"] = []
-        for asset_class in asset_registry:
-            self.assets.append(
-                asset_class(
-                    interface=self.asset_interface or self.backend_interface,
-                    settings=settings,
-                    shutdown_event=self._shutdown_event,
-                )
-            )
+        self.assets: List["Asset"] = self._get_assets()
+
+        # OpenMetrics/Prometheus-compatible endpoints
+        self.assets.extend(self._get_open_metrics_assets())
 
         # static system info, both hardware and software
         self.system_info: SystemInfo = SystemInfo(
-            settings=settings, interface=interface
+            settings=self.settings, interface=interface
         )
 
+    def _get_assets(self) -> List["Asset"]:
+        return [
+            asset_class(
+                interface=self.asset_interface or self.backend_interface,
+                settings=self.settings,
+                shutdown_event=self._shutdown_event,
+            )
+            for asset_class in asset_registry
+        ]
+
+    def _get_open_metrics_assets(self) -> List["Asset"]:
+        open_metrics_endpoints = self.settings._stats_open_metrics_endpoints
+        if not open_metrics_endpoints:
+            return []
+
+        assets: List[Asset] = []
+        for name, endpoint in open_metrics_endpoints.items():
+            if not OpenMetrics.is_available(url=endpoint):
+                continue
+            logger.debug(f"Monitoring OpenMetrics endpoint: {endpoint}")
+            open_metrics = OpenMetrics(
+                interface=self.asset_interface or self.backend_interface,
+                settings=self.settings,
+                shutdown_event=self._shutdown_event,
+                name=name,
+                url=endpoint,
+            )
+            assets.append(open_metrics)  # type: ignore
+
+        return assets
+
     def aggregate_and_publish_asset_metrics(self) -> None:
         if self.asset_interface is None:
             return None
@@ -147,13 +178,14 @@ class SystemMonitor:
 
     def start(self) -> None:
         self._shutdown_event.clear()
-        if self._process is None:
-            logger.info("Starting system monitor")
-            # self._process = mp.Process(target=self._start, name="SystemMonitor")
-            self._process = threading.Thread(
-                target=self._start, daemon=True, name="SystemMonitor"
-            )
-            self._process.start()
+        if self._process is not None:
+            return None
+        logger.info("Starting system monitor")
+        # self._process = mp.Process(target=self._start, name="SystemMonitor")
+        self._process = threading.Thread(
+            target=self._start, daemon=True, name="SystemMonitor"
+        )
+        self._process.start()
 
     def finish(self) -> None:
         if self._process is None:

wandb/sdk/internal/tb_watcher.py

@@ -1,6 +1,4 @@
-"""
-tensorboard watcher.
-"""
+"""tensorboard watcher."""
 
 import glob
 import logging
@@ -61,7 +59,7 @@ def _link_and_save_file(
 
 
 def is_tfevents_file_created_by(path: str, hostname: str, start_time: float) -> bool:
-    """Checks if a path is a tfevents file created by hostname.
+    """Check if a path is a tfevents file created by hostname.
 
     tensorboard tfevents filename format:
         https://github.com/tensorflow/tensorboard/blob/f3f26b46981da5bd46a5bb93fcf02d9eb7608bc1/tensorboard/summary/writer/event_file_writer.py#L81
@@ -95,7 +93,7 @@ def is_tfevents_file_created_by(path: str, hostname: str, start_time: float) ->
     # TODO: we should also check the PID (also contained in the tfevents
     #     filename). Can we assume that our parent pid is the user process
     #     that wrote these files?
-    return created_time >= int(start_time)  # noqa: W503
+    return created_time >= int(start_time)
 
 
 class TBWatcher:
@@ -216,7 +214,7 @@ class TBDirWatcher:
         self._thread.start()
 
     def _is_our_tfevents_file(self, path: str) -> bool:
-        """Checks if a path has been modified since launch and contains tfevents"""
+        """Check if a path has been modified since launch and contains tfevents."""
         if not path:
             raise ValueError("Path must be a nonempty string")
         if self._force:
@@ -229,7 +227,7 @@ class TBDirWatcher:
     def _loader(
         self, save: bool = True, namespace: Optional[str] = None
     ) -> "EventFileLoader":
-        """Incredibly hacky class generator to optionally save / prefix tfevent files"""
+        """Incredibly hacky class generator to optionally save / prefix tfevent files."""
         _loader_interface = self._tbwatcher._interface
         _loader_settings = self._tbwatcher._settings
         try:
@@ -285,7 +283,7 @@ class TBDirWatcher:
             raise e
 
     def _thread_body(self) -> None:
-        """Check for new events every second"""
+        """Check for new events every second."""
         shutdown_time: Optional[float] = None
         while True:
             self._process_events()
@@ -318,7 +316,7 @@ class TBDirWatcher:
 
 
 class Event:
-    """An event wrapper to enable priority queueing"""
+    """An event wrapper to enable priority queueing."""
 
     def __init__(self, event: "ProtoEvent", namespace: Optional[str]):
         self.event = event
@@ -332,10 +330,11 @@ class Event:
 
 
 class TBEventConsumer:
-    """Consumes tfevents from a priority queue.  There should always
-    only be one of these per run_manager.  We wait for 10 seconds of queued
-    events to reduce the chance of multiple tfevent files triggering
-    out of order steps.
+    """Consume tfevents from a priority queue.
+
+    There should always only be one of these per run_manager.  We wait for 10 seconds of
+    queued events to reduce the chance of multiple tfevent files triggering out of order
+    steps.
     """
 
     def __init__(

wandb/sdk/launch/_project_spec.py

@@ -1,8 +1,7 @@
+"""Convert launch arguments into a runnable wandb launch script.
+
+Arguments can come from a launch spec or call to wandb launch.
 """
-Internal utility for converting arguments from a launch spec or call to wandb launch
-into a runnable wandb launch script
-"""
-import binascii
 import enum
 import json
 import logging
@@ -15,11 +14,11 @@ import wandb
 import wandb.docker as docker
 from wandb.apis.internal import Api
 from wandb.apis.public import Artifact as PublicArtifact
-from wandb.errors import CommError, LaunchError
+from wandb.errors import CommError
 from wandb.sdk.lib.runid import generate_id
 
 from . import utils
-from .utils import LOG_PREFIX
+from .utils import LOG_PREFIX, LaunchError
 
 _logger = logging.getLogger(__name__)
 
@@ -60,7 +59,6 @@ class LaunchProject:
         overrides: Dict[str, Any],
         resource: str,
         resource_args: Dict[str, Any],
-        cuda: Optional[bool],
         run_id: Optional[str],
     ):
         if uri is not None and utils.is_bare_wandb_uri(uri):
@@ -68,17 +66,24 @@ class LaunchProject:
             _logger.info(f"{LOG_PREFIX}Updating uri with base uri: {uri}")
         self.uri = uri
         self.job = job
-        wandb.termlog(f"{LOG_PREFIX}Launch project got job {job}")
+        if job is not None:
+            wandb.termlog(f"{LOG_PREFIX}Launching job: {job}")
         self._job_artifact: Optional[PublicArtifact] = None
         self.api = api
         self.launch_spec = launch_spec
         self.target_entity = target_entity
         self.target_project = target_project.lower()
         self.name = name  # TODO: replace with run_id
+        # the builder key can be passed in through the resource args
+        # but these resource_args are then passed to the appropriate
+        # runner, so we need to pop the builder key out
+        resource_args_build = resource_args.get(resource, {}).pop("builder", {})
         self.resource = resource
         self.resource_args = resource_args
         self.python_version: Optional[str] = launch_spec.get("python_version")
-        self.cuda_version: Optional[str] = launch_spec.get("cuda_version")
+        self.cuda_base_image: Optional[str] = resource_args_build.get("cuda", {}).get(
+            "base_image"
+        )
         self._base_image: Optional[str] = launch_spec.get("base_image")
         self.docker_image: Optional[str] = docker_config.get(
             "docker_image"
@@ -95,11 +100,8 @@ class LaunchProject:
         self.override_artifacts: Dict[str, Any] = overrides.get("artifacts", {})
         self.override_entrypoint: Optional[EntryPoint] = None
         self.deps_type: Optional[str] = None
-        self.cuda = cuda
         self._runtime: Optional[str] = None
         self.run_id = run_id or generate_id()
-        self._image_tag: str = self._initialize_image_job_tag() or self.run_id
-        wandb.termlog(f"{LOG_PREFIX}Launch project using image tag {self._image_tag}")
         self._entry_points: Dict[
             str, EntryPoint
         ] = {}  # todo: keep multiple entrypoint support?
@@ -139,15 +141,13 @@ class LaunchProject:
                 )
             self.source = LaunchSource.LOCAL
             self.project_dir = self.uri
-        if launch_spec.get("resource_args"):
-            self.resource_args = launch_spec["resource_args"]
 
         self.aux_dir = tempfile.mkdtemp()
         self.clear_parameter_run_config_collisions()
 
     @property
     def base_image(self) -> str:
-        """Returns {PROJECT}_base:{PYTHON_VERSION}"""
+        """Returns {PROJECT}_base:{PYTHON_VERSION}."""
         # TODO: this should likely be source_project when we have it...
 
         # don't make up a separate base image name if user provides a docker image
@@ -174,25 +174,15 @@ class LaunchProject:
             assert self.job is not None
             return wandb.util.make_docker_image_name_safe(self.job.split(":")[0])
 
-    def _initialize_image_job_tag(self) -> Optional[str]:
-        if self.job is not None:
-            job_name, alias = self.job.split(":")
-            # Alias is used to differentiate images between jobs of the same sequence
-            _image_tag = f"{alias}-{job_name}"
-            _logger.debug(f"{LOG_PREFIX}Setting image tag {_image_tag}")
-            return wandb.util.make_docker_image_name_safe(_image_tag)
-        return None
-
-    @property
-    def image_uri(self) -> str:
-        if self.docker_image:
-            return self.docker_image
-        return f"{self.image_name}:{self.image_tag}"
-
-    @property
-    def image_tag(self) -> str:
-
-        return self._image_tag[:IMAGE_TAG_MAX_LENGTH]
+    def build_required(self) -> bool:
+        """Checks the source to see if a build is required."""
+        # since the image tag for images built from jobs
+        # is based on the job version index, which is immutable
+        # we don't need to build the image for a job if that tag
+        # already exists
+        if self.source != LaunchSource.JOB:
+            return True
+        return False
 
     @property
     def docker_image(self) -> Optional[str]:
@@ -225,7 +215,7 @@ class LaunchProject:
         return list(self._entry_points.values())[0]
 
     def add_entry_point(self, command: List[str]) -> "EntryPoint":
-        """Adds an entry point to the project."""
+        """Add an entry point to the project."""
         entry_point = command[-1]
         new_entrypoint = EntryPoint(name=entry_point, command=command)
         self._entry_points[entry_point] = new_entrypoint
@@ -243,10 +233,37 @@ class LaunchProject:
         try:
             job = public_api.job(self.job, path=job_dir)
         except CommError:
-            raise LaunchError(f"Job {self.job} not found")
+            raise LaunchError(
+                f"Job {self.job} not found. Jobs have the format: <entity>/<project>/<name>:<alias>"
+            )
         job.configure_launch_project(self)
         self._job_artifact = job._job_artifact
 
+    def get_image_source_string(self) -> str:
+        """Returns a unique string identifying the source of an image."""
+        if self.source == LaunchSource.LOCAL:
+            # TODO: more correct to get a hash of local uri contents
+            assert isinstance(self.uri, str)
+            return self.uri
+        elif self.source == LaunchSource.JOB:
+            assert self._job_artifact is not None
+            return f"{self._job_artifact.name}:v{self._job_artifact.version}"
+        elif self.source == LaunchSource.GIT:
+            assert isinstance(self.uri, str)
+            ret = self.uri
+            if self.git_version:
+                ret += self.git_version
+            return ret
+        elif self.source == LaunchSource.WANDB:
+            assert isinstance(self.uri, str)
+            return self.uri
+        elif self.source == LaunchSource.DOCKER:
+            assert isinstance(self.docker_image, str)
+            _logger.debug("")
+            return self.docker_image
+        else:
+            raise LaunchError("Unknown source type when determing image source string")
+
     def _fetch_project_local(self, internal_api: Api) -> None:
         """Fetch a project (either wandb run or git repo) into a local directory, returning the path to the local project directory."""
         # these asserts are all guaranteed to pass, but are required by mypy
@@ -263,24 +280,6 @@ class LaunchProject:
             )
             program_name = run_info.get("codePath") or run_info["program"]
 
-            if run_info.get("cudaVersion"):
-                original_cuda_version = ".".join(run_info["cudaVersion"].split(".")[:2])
-
-                if self.cuda is None:
-                    # only set cuda on by default if cuda is None (unspecified), not False (user specifically requested cpu image)
-                    wandb.termlog(
-                        f"{LOG_PREFIX}Original wandb run {source_run_name} was run with cuda version {original_cuda_version}. Enabling cuda builds by default; to build on a CPU-only image, run again with --cuda=False"
-                    )
-                    self.cuda_version = original_cuda_version
-                    self.cuda = True
-                if (
-                    self.cuda
-                    and self.cuda_version
-                    and self.cuda_version != original_cuda_version
-                ):
-                    wandb.termlog(
-                        f"{LOG_PREFIX}Specified cuda version {self.cuda_version} differs from original cuda version {original_cuda_version}. Running with specified version {self.cuda_version}"
-                    )
             self.python_version = run_info.get("python", "3")
             downloaded_code_artifact = utils.check_and_download_code_artifacts(
                 source_entity,
@@ -289,11 +288,7 @@ class LaunchProject:
                 internal_api,
                 self.project_dir,
             )
-            if downloaded_code_artifact:
-                self._image_tag = binascii.hexlify(
-                    downloaded_code_artifact.digest.encode()
-                ).decode()
-            else:
+            if not downloaded_code_artifact:
                 if not run_info["git"]:
                     raise LaunchError(
                         "Reproducing a run requires either an associated git repo or a code artifact logged with `run.log_code()`"
@@ -308,12 +303,8 @@ class LaunchProject:
                 patch = utils.fetch_project_diff(
                     source_entity, source_project, source_run_name, internal_api
                 )
-                tag_string = run_info["git"]["remote"] + run_info["git"]["commit"]
                 if patch:
                     utils.apply_patch(patch, self.project_dir)
-                    tag_string += patch
-
-                self._image_tag = binascii.hexlify(tag_string.encode()).decode()
 
                 # For cases where the entry point wasn't checked into git
                 if not os.path.exists(os.path.join(self.project_dir, program_name)):
@@ -434,7 +425,6 @@ def create_project_from_spec(launch_spec: Dict[str, Any], api: Api) -> LaunchPro
     Returns:
         An initialized `LaunchProject` object
     """
-
     name: Optional[str] = None
     if launch_spec.get("name"):
         name = launch_spec["name"]
@@ -451,7 +441,6 @@ def create_project_from_spec(launch_spec: Dict[str, Any], api: Api) -> LaunchPro
         launch_spec.get("overrides", {}),
         launch_spec.get("resource", None),
         launch_spec.get("resource_args", {}),
-        launch_spec.get("cuda", None),
         launch_spec.get("run_id", None),
     )