wandb 0.15.9__py3-none-any.whl → 0.15.11__py3-none-any.whl

wandb/sdk/internal/job_builder.py

@@ -87,6 +87,7 @@ class JobBuilder:
     _aliases: List[str]
     _job_seq_id: Optional[str]
     _job_version_alias: Optional[str]
+    _is_notebook_run: bool
 
     def __init__(self, settings: SettingsStatic):
         self._settings = settings
@@ -103,6 +104,7 @@ class JobBuilder:
         self._source_type: Optional[
             Literal["repo", "artifact", "image"]
         ] = settings.job_source  # type: ignore[assignment]
+        self._is_notebook_run = self._get_is_notebook_run()
 
     def set_config(self, config: Dict[str, Any]) -> None:
         self._config = config
@@ -153,7 +155,7 @@ class JobBuilder:
         commit = git_info.get("commit")
         assert remote is not None
         assert commit is not None
-        if self._is_notebook_run():
+        if self._is_notebook_run:
             if not os.path.exists(
                 os.path.join(os.getcwd(), os.path.basename(program_relpath))
             ):
@@ -194,7 +196,7 @@ class JobBuilder:
                 os.path.basename(sys.executable),
                 full_program_path,
             ],
-            "notebook": self._is_notebook_run(),
+            "notebook": self._is_notebook_run,
         }
 
         if self._settings.job_name:
@@ -219,18 +221,16 @@ class JobBuilder:
     ) -> Tuple[Optional[ArtifactSourceDict], Optional[str]]:
         assert isinstance(self._logged_code_artifact, dict)
         # TODO: should we just always exit early if the path doesn't exist?
-        if self._is_notebook_run() and not self._is_colab_run():
+        if self._is_notebook_run and not self._is_colab_run():
             full_program_relpath = os.path.relpath(program_relpath, os.getcwd())
             # if the resolved path doesn't exist, then we shouldn't make a job because it will fail
             if not os.path.exists(full_program_relpath):
                 # when users call log code in a notebook the code artifact starts
-                # at the directory the notebook is in instead of the jupyter
-                # core
-                if os.path.exists(os.path.basename(program_relpath)):
-                    full_program_relpath = os.path.basename(program_relpath)
-                else:
+                # at the directory the notebook is in instead of the jupyter core
+                if not os.path.exists(os.path.basename(program_relpath)):
                     _logger.info("target path does not exist, exiting")
                     return None, None
+                full_program_relpath = os.path.basename(program_relpath)
         else:
             full_program_relpath = program_relpath
         entrypoint = [
@@ -240,7 +240,7 @@ class JobBuilder:
         # TODO: update executable to a method that supports pex
         source: ArtifactSourceDict = {
             "entrypoint": entrypoint,
-            "notebook": self._is_notebook_run(),
+            "notebook": self._is_notebook_run,
             "artifact": f"wandb-artifact://_id/{self._logged_code_artifact['id']}",
         }
 
@@ -271,7 +271,7 @@ class JobBuilder:
         }
         return source, name
 
-    def _is_notebook_run(self) -> bool:
+    def _get_is_notebook_run(self) -> bool:
         return hasattr(self._settings, "_jupyter") and bool(self._settings._jupyter)
 
     def _is_colab_run(self) -> bool:
@@ -288,15 +288,10 @@ class JobBuilder:
             return None
 
         runtime: Optional[str] = metadata.get("python")
-        program_relpath: Optional[str] = metadata.get("codePath")
         # can't build a job without a python version
         if runtime is None:
             return None
 
-        if self._is_notebook_run():
-            _logger.info("run is notebook based run")
-            program_relpath = metadata.get("program")
-
         input_types = TypeRegistry.type_of(self._config).to_json()
         output_types = TypeRegistry.type_of(self._summary).to_json()
 
@@ -315,10 +310,14 @@ class JobBuilder:
             source_type = source_info.get("source_type")
         else:
             # configure job from environment
-            source_type = self._get_source_type(metadata, program_relpath)
+            source_type = self._get_source_type(metadata)
             if not source_type:
                 return None
 
+            program_relpath = self._get_program_relpath(source_type, metadata)
+            if source_type != "image" and not program_relpath:
+                return None
+
             source: Union[
                 Optional[GitSourceDict],
                 Optional[ArtifactSourceDict],
@@ -326,19 +325,15 @@ class JobBuilder:
             ] = None
 
             # make source dict
-            if source_type == "repo" and self._has_git_job_ingredients(
-                metadata, program_relpath
-            ):
-                assert program_relpath is not None
+            if source_type == "repo":
+                assert program_relpath
                 source, name = self._build_repo_job_source(
                     metadata,
                     program_relpath,
                     metadata.get("root"),
                 )
-            elif source_type == "artifact" and self._has_artifact_job_ingredients(
-                program_relpath
-            ):
-                assert program_relpath is not None
+            elif source_type == "artifact":
+                assert program_relpath
                 source, name = self._build_artifact_job_source(program_relpath)
             elif source_type == "image" and self._has_image_job_ingredients(metadata):
                 source, name = self._build_image_job_source(metadata)
@@ -390,17 +385,15 @@ class JobBuilder:
 
         return artifact
 
-    def _get_source_type(
-        self, metadata: Dict[str, Any], relpath: Optional[str]
-    ) -> Optional[str]:
+    def _get_source_type(self, metadata: Dict[str, Any]) -> Optional[str]:
         if self._source_type:
             return self._source_type
 
-        if self._has_git_job_ingredients(metadata, relpath):
+        if self._has_git_job_ingredients(metadata):
             _logger.info("is repo sourced job")
             return "repo"
 
-        if self._has_artifact_job_ingredients(relpath):
+        if self._has_artifact_job_ingredients():
             _logger.info("is artifact sourced job")
             return "artifact"
 
@@ -411,6 +404,21 @@ class JobBuilder:
         _logger.info("no source found")
         return None
 
+    def _get_program_relpath(
+        self, source_type: str, metadata: Dict[str, Any]
+    ) -> Optional[str]:
+        if self._is_notebook_run:
+            _logger.info("run is notebook based run")
+            return metadata.get("program")
+
+        if source_type == "artifact" or self._settings.job_source == "artifact":
+            # if the job is set to be an artifact, use relpath guaranteed
+            # to be correct. 'codePath' uses the root path when in git repo
+            # fallback to codePath if strictly local relpath not present
+            return metadata.get("codePathLocal") or metadata.get("codePath")
+
+        return metadata.get("codePath")
+
     def _handle_metadata_file(
         self,
     ) -> Optional[Dict]:
@@ -421,18 +429,14 @@ class JobBuilder:
 
         return None
 
-    def _has_git_job_ingredients(
-        self, metadata: Dict[str, Any], program_relpath: Optional[str]
-    ) -> bool:
+    def _has_git_job_ingredients(self, metadata: Dict[str, Any]) -> bool:
         git_info: Dict[str, str] = metadata.get("git", {})
-        if program_relpath is None:
-            return False
-        if self._is_notebook_run() and metadata.get("root") is None:
+        if self._is_notebook_run and metadata.get("root") is None:
             return False
         return git_info.get("remote") is not None and git_info.get("commit") is not None
 
-    def _has_artifact_job_ingredients(self, program_relpath: Optional[str]) -> bool:
-        return self._logged_code_artifact is not None and program_relpath is not None
+    def _has_artifact_job_ingredients(self) -> bool:
+        return self._logged_code_artifact is not None
 
     def _has_image_job_ingredients(self, metadata: Dict[str, Any]) -> bool:
         return metadata.get("docker") is not None

wandb/sdk/internal/sender.py

@@ -33,7 +33,7 @@ from wandb.errors import CommError, UsageError
 from wandb.errors.util import ProtobufErrorHandler
 from wandb.filesync.dir_watcher import DirWatcher
 from wandb.proto import wandb_internal_pb2
-from wandb.sdk.artifacts import artifact_saver
+from wandb.sdk.artifacts.artifact_saver import ArtifactSaver
 from wandb.sdk.interface import interface
 from wandb.sdk.interface.interface_queue import InterfaceQueue
 from wandb.sdk.internal import (
@@ -268,6 +268,7 @@ class SendManager:
         self._cached_summary: Dict[str, Any] = dict()
         self._config_metric_index_dict: Dict[str, int] = {}
         self._config_metric_dict: Dict[str, wandb_internal_pb2.MetricRecord] = {}
+        self._consolidated_summary: Dict[str, Any] = dict()
 
         self._cached_server_info = dict()
         self._cached_viewer = dict()
@@ -1149,7 +1150,9 @@ class SendManager:
         summary_dict.pop("_wandb", None)
         if self._metadata_summary:
             summary_dict["_wandb"] = self._metadata_summary
-        json_summary = json.dumps(summary_dict)
+        # merge with consolidated summary
+        self._consolidated_summary.update(summary_dict)
+        json_summary = json.dumps(self._consolidated_summary)
         if self._fs:
             self._fs.push(filenames.SUMMARY_FNAME, json_summary)
         # TODO(jhr): we should only write this at the end of the script
@@ -1446,28 +1449,6 @@ class SendManager:
 
         self._respond_result(result)
 
-    def send_request_artifact_send(self, record: "Record") -> None:
-        # TODO: combine and eventually remove send_request_log_artifact()
-
-        # for now, we are using req/resp uuid for transaction id
-        # in the future this should be part of the message to handle idempotency
-        xid = record.uuid
-
-        done_msg = wandb_internal_pb2.ArtifactDoneRequest(xid=xid)
-        artifact = record.request.artifact_send.artifact
-        try:
-            res = self._send_artifact(artifact)
-            assert res, "Unable to send artifact"
-            done_msg.artifact_id = res["id"]
-            logger.info(f"logged artifact {artifact.name} - {res}")
-        except Exception as e:
-            done_msg.error_message = 'error logging artifact "{}/{}": {}'.format(
-                artifact.type, artifact.name, e
-            )
-
-        logger.info("send artifact done")
-        self._interface._publish_artifact_done(done_msg)
-
     def send_artifact(self, record: "Record") -> None:
         artifact = record.artifact
         try:
@@ -1486,7 +1467,7 @@ class SendManager:
         from pkg_resources import parse_version
 
         assert self._pusher
-        saver = artifact_saver.ArtifactSaver(
+        saver = ArtifactSaver(
             api=self._api,
             digest=artifact.digest,
             manifest_json=_manifest_json_from_proto(artifact.manifest),
@@ -1512,6 +1493,7 @@ class SendManager:
             client_id=artifact.client_id,
             sequence_client_id=artifact.sequence_client_id,
             metadata=metadata,
+            ttl_duration_seconds=artifact.ttl_duration_seconds or None,
             description=artifact.description or None,
             aliases=artifact.aliases,
             use_after_commit=artifact.use_after_commit,

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

@@ -1,12 +1,14 @@
 import threading
 from collections import deque
-from typing import TYPE_CHECKING, List
+from typing import TYPE_CHECKING, List, Optional
 
 try:
     import psutil
 except ImportError:
     psutil = None
 
+from wandb.errors.term import termwarn
+
 from .aggregators import aggregate_mean
 from .asset_registry import asset_registry
 from .interfaces import Interface, Metric, MetricsMonitor
@@ -17,18 +19,132 @@ if TYPE_CHECKING:
     from wandb.sdk.internal.settings_static import SettingsStatic
 
 
-class DiskUsage:
+class DiskUsagePercent:
     """Total system disk usage in percent."""
 
-    # name = "disk_usage"
-    name = "disk"
+    name = "disk.{path}.usagePercent"
+    samples: "Deque[List[float]]"
+
+    def __init__(self, paths: List[str]) -> None:
+        self.samples = deque([])
+        # check if we have access to the disk paths:
+        self.paths: List[str] = []
+        for path in paths:
+            try:
+                psutil.disk_usage(path)
+                self.paths.append(path)
+            except Exception as e:  # noqa
+                termwarn(f"Could not access disk path {path}: {e}", repeat=False)
+
+    def sample(self) -> None:
+        # self.samples.append(psutil.disk_usage("/").percent)
+        disk_usage: List[float] = []
+        for path in self.paths:
+            disk_usage.append(psutil.disk_usage(path).percent)
+        if disk_usage:
+            self.samples.append(disk_usage)
+
+    def clear(self) -> None:
+        self.samples.clear()
+
+    def aggregate(self) -> dict:
+        if not self.samples:
+            return {}
+        disk_metrics = {}
+        for i, _path in enumerate(self.paths):
+            aggregate_i = aggregate_mean([sample[i] for sample in self.samples])
+            # ugly hack to please the frontend:
+            _path = _path.replace("/", "\\")
+            disk_metrics[self.name.format(path=_path)] = aggregate_i
+
+        return disk_metrics
+
+
+class DiskUsage:
+    """Total system disk usage in GB."""
+
+    name = "disk.{path}.usageGB"
+    samples: "Deque[List[float]]"
+
+    def __init__(self, paths: List[str]) -> None:
+        self.samples = deque([])
+        # check if we have access to the disk paths:
+        self.paths: List[str] = []
+        for path in paths:
+            try:
+                psutil.disk_usage(path)
+                self.paths.append(path)
+            except Exception as e:  # noqa
+                termwarn(f"Could not access disk path {path}: {e}", repeat=False)
+
+    def sample(self) -> None:
+        disk_usage: List[float] = []
+        for path in self.paths:
+            disk_usage.append(psutil.disk_usage(path).used / 1024 / 1024 / 1024)
+        if disk_usage:
+            self.samples.append(disk_usage)
+
+    def clear(self) -> None:
+        self.samples.clear()
+
+    def aggregate(self) -> dict:
+        if not self.samples:
+            return {}
+        disk_metrics = {}
+        for i, _path in enumerate(self.paths):
+            aggregate_i = aggregate_mean([sample[i] for sample in self.samples])
+            # ugly hack to please the frontend:
+            _path = _path.replace("/", "\\")
+            disk_metrics[self.name.format(path=_path)] = aggregate_i
+
+        return disk_metrics
+
+
+class DiskIn:
+    """Total system disk read in MB."""
+
+    name = "disk.in"
+    samples: "Deque[float]"
+
+    def __init__(self) -> None:
+        self.samples = deque([])
+        self.read_init: Optional[int] = None
+
+    def sample(self) -> None:
+        if self.read_init is None:
+            # initialize the read_init value on first sample
+            self.read_init = psutil.disk_io_counters().read_bytes
+        self.samples.append(
+            (psutil.disk_io_counters().read_bytes - self.read_init) / 1024 / 1024
+        )
+
+    def clear(self) -> None:
+        self.samples.clear()
+
+    def aggregate(self) -> dict:
+        if not self.samples:
+            return {}
+        aggregate = aggregate_mean(self.samples)
+        return {self.name: aggregate}
+
+
+class DiskOut:
+    """Total system disk write in MB."""
+
+    name = "disk.out"
     samples: "Deque[float]"
 
     def __init__(self) -> None:
         self.samples = deque([])
+        self.write_init: Optional[int] = None
 
     def sample(self) -> None:
-        self.samples.append(psutil.disk_usage("/").percent)
+        if self.write_init is None:
+            # init on first sample
+            self.write_init = psutil.disk_io_counters().write_bytes
+        self.samples.append(
+            (psutil.disk_io_counters().write_bytes - self.write_init) / 1024 / 1024
+        )
 
     def clear(self) -> None:
         self.samples.clear()
@@ -49,7 +165,13 @@ class Disk:
         shutdown_event: threading.Event,
     ) -> None:
         self.name = self.__class__.__name__.lower()
-        self.metrics: List[Metric] = [DiskUsage()]
+        self.settings = settings
+        self.metrics: List[Metric] = [
+            DiskUsagePercent(list(settings._stats_disk_paths or ["/"])),
+            DiskUsage(list(settings._stats_disk_paths or ["/"])),
+            DiskIn(),
+            DiskOut(),
+        ]
         self.metrics_monitor = MetricsMonitor(
             self.name,
             self.metrics,
@@ -64,11 +186,22 @@ class Disk:
         return psutil is not None
 
     def probe(self) -> dict:
-        # total disk space:
-        total = psutil.disk_usage("/").total / 1024 / 1024 / 1024
-        # total disk space used:
-        used = psutil.disk_usage("/").used / 1024 / 1024 / 1024
-        return {self.name: {"total": total, "used": used}}
+        disk_paths = list(self.settings._stats_disk_paths or ["/"])
+        disk_metrics = {}
+        for disk_path in disk_paths:
+            try:
+                # total disk space in GB:
+                total = psutil.disk_usage(disk_path).total / 1024 / 1024 / 1024
+                # total disk space used in GB:
+                used = psutil.disk_usage(disk_path).used / 1024 / 1024 / 1024
+                disk_metrics[disk_path] = {
+                    "total": total,
+                    "used": used,
+                }
+            except Exception as e:  # noqa
+                termwarn(f"Could not access disk path {disk_path}: {e}", repeat=False)
+
+        return {self.name: disk_metrics}
 
     def start(self) -> None:
         self.metrics_monitor.start()

wandb/sdk/internal/system/system_info.py

@@ -19,6 +19,7 @@ from wandb.sdk.lib.filenames import (
     REQUIREMENTS_FNAME,
 )
 from wandb.sdk.lib.gitlib import GitRepo
+from wandb.sdk.wandb_settings import _get_program_relpath
 
 from .assets.interfaces import Interface
 
@@ -87,7 +88,7 @@ class SystemInfo:
 
     def _save_code(self) -> None:
         logger.debug("Saving code")
-        if self.settings.program_relpath is None:
+        if not self.settings.program_relpath:
             logger.warning("unable to save code -- program entry not found")
             return None
 
@@ -210,8 +211,11 @@ class SystemInfo:
 
         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"] = _get_program_relpath(self.settings.program)
         if not self.settings.disable_code:
-            if self.settings.program_relpath is not None:
+            if self.settings.program_relpath:
                 data["codePath"] = self.settings.program_relpath
             elif self.settings._jupyter:
                 if self.settings.notebook_name:

wandb/sdk/launch/__init__.py

@@ -0,0 +1,5 @@
+from ._launch import launch
+from ._launch_add import launch_add
+from .agent.agent import LaunchAgent
+
+__all__ = ["LaunchAgent", "launch", "launch_add"]

wandb/sdk/launch/{launch.py → _launch.py}

@@ -124,24 +124,28 @@ def create_and_run_agent(
     agent.loop()
 
 
-def _run(
-    uri: Optional[str],
-    job: Optional[str],
-    name: Optional[str],
-    project: Optional[str],
-    entity: Optional[str],
-    docker_image: Optional[str],
-    entry_point: Optional[List[str]],
-    version: Optional[str],
-    resource: str,
-    resource_args: Optional[Dict[str, Any]],
-    launch_config: Optional[Dict[str, Any]],
-    synchronous: Optional[bool],
+def _launch(
     api: Api,
-    run_id: Optional[str],
-    repository: Optional[str],
+    uri: Optional[str] = None,
+    job: Optional[str] = None,
+    name: Optional[str] = None,
+    project: Optional[str] = None,
+    entity: Optional[str] = None,
+    docker_image: Optional[str] = None,
+    entry_point: Optional[List[str]] = None,
+    version: Optional[str] = None,
+    resource: Optional[str] = None,
+    resource_args: Optional[Dict[str, Any]] = None,
+    launch_config: Optional[Dict[str, Any]] = None,
+    synchronous: Optional[bool] = None,
+    run_id: Optional[str] = None,
+    repository: Optional[str] = None,
 ) -> AbstractRun:
     """Helper that delegates to the project-running method corresponding to the passed-in backend."""
+    if launch_config is None:
+        launch_config = {}
+    if resource is None:
+        resource = "local-container"
     launch_spec = construct_launch_spec(
         uri,
         job,
@@ -193,9 +197,8 @@ def _run(
         )
 
 
-def run(
+def launch(
     api: Api,
-    uri: Optional[str] = None,
     job: Optional[str] = None,
     entry_point: Optional[List[str]] = None,
     version: Optional[str] = None,
@@ -210,41 +213,43 @@ def run(
     run_id: Optional[str] = None,
     repository: Optional[str] = None,
 ) -> AbstractRun:
-    """Run a W&B launch experiment. The project can be wandb uri or a Git URI.
+    """Launch a W&B launch experiment.
 
     Arguments:
-    uri: URI of experiment to run. A wandb run uri or a Git repository URI.
-    job: string reference to a wandb.Job eg: wandb/test/my-job:latest
-    api: An instance of a wandb Api from wandb.apis.internal.
-    entry_point: Entry point to run within the project. Defaults to using the entry point used
-        in the original run for wandb URIs, or main.py for git repository URIs.
-    version: For Git-based projects, either a commit hash or a branch name.
-    name: Name run under which to launch the run.
-    resource: Execution backend for the run.
-    resource_args: Resource related arguments for launching runs onto a remote backend.
-        Will be stored on the constructed launch config under ``resource_args``.
-    project: Target project to send launched run to
-    entity: Target entity to send launched run to
-    config: A dictionary containing the configuration for the run. May also contain
-    resource specific arguments under the key "resource_args".
-    synchronous: Whether to block while waiting for a run to complete. Defaults to True.
-        Note that if ``synchronous`` is False and ``backend`` is "local-container", this
-        method will return, but the current process will block when exiting until
-        the local run completes. If the current process is interrupted, any
-        asynchronous runs launched via this method will be terminated. If
-        ``synchronous`` is True and the run fails, the current process will
-        error out as well.
-    run_id: ID for the run (To ultimately replace the :name: field)
-    repository: string name of repository path for remote registry
+        job: string reference to a wandb.Job eg: wandb/test/my-job:latest
+        api: An instance of a wandb Api from wandb.apis.internal.
+        entry_point: Entry point to run within the project. Defaults to using the entry point used
+            in the original run for wandb URIs, or main.py for git repository URIs.
+        version: For Git-based projects, either a commit hash or a branch name.
+        name: Name run under which to launch the run.
+        resource: Execution backend for the run.
+        resource_args: Resource related arguments for launching runs onto a remote backend.
+            Will be stored on the constructed launch config under ``resource_args``.
+        project: Target project to send launched run to
+        entity: Target entity to send launched run to
+        config: A dictionary containing the configuration for the run. May also contain
+        resource specific arguments under the key "resource_args".
+        synchronous: Whether to block while waiting for a run to complete. Defaults to True.
+            Note that if ``synchronous`` is False and ``backend`` is "local-container", this
+            method will return, but the current process will block when exiting until
+            the local run completes. If the current process is interrupted, any
+            asynchronous runs launched via this method will be terminated. If
+            ``synchronous`` is True and the run fails, the current process will
+            error out as well.
+        run_id: ID for the run (To ultimately replace the :name: field)
+        repository: string name of repository path for remote registry
 
     Example:
-        import wandb
-        project_uri = "https://github.com/wandb/examples"
-        params = {"alpha": 0.5, "l1_ratio": 0.01}
+        ```python
+        from wandb.sdk.launch import launch
+
+        job = "wandb/jobs/Hello World:latest"
+        params = {"epochs": 5}
         # Run W&B project and create a reproducible docker environment
         # on a local host
         api = wandb.apis.internal.Api()
-        wandb.launch(project_uri, api, parameters=params)
+        launch(api, job, parameters=params)
+        ```
 
 
     Returns:
@@ -255,15 +260,9 @@ def run(
         `wandb.exceptions.ExecutionError` If a run launched in blocking mode
         is unsuccessful.
     """
-    if config is None:
-        config = {}
-
-    # default to local container for runs without a queue
-    if resource is None:
-        resource = "local-container"
-
-    submitted_run_obj = _run(
-        uri=uri,
+    submitted_run_obj = _launch(
+        # TODO: fully deprecate URI path
+        uri=None,
         job=job,
         name=name,
         project=project,