wandb 0.15.4__py3-none-any.whl → 0.15.5__py3-none-any.whl

wandb/sdk/lib/_settings_toposort_generated.py

@@ -23,6 +23,7 @@ _Setting = Literal[
     "_disable_service",
     "_disable_stats",
     "_disable_viewer",
+    "_disable_setproctitle",
     "_except_exit",
     "_executable",
     "_extra_http_headers",
@@ -48,6 +49,7 @@ _Setting = Literal[
     "_platform",
     "_python",
     "_runqueue_item_id",
+    "_require_nexus",
     "_save_requirements",
     "_service_transport",
     "_service_wait",
@@ -67,6 +69,7 @@ _Setting = Literal[
     "allow_val_change",
     "anonymous",
     "api_key",
+    "azure_account_url_to_access_key",
     "base_url",
     "code_dir",
     "config_paths",

wandb/sdk/service/server.py

@@ -129,15 +129,15 @@ class WandbServer:
             return
         _startup_debug.print_message(message)
 
-    def serve(self) -> None:
-        self._setup_tracelog()
-        mux = StreamMux()
-        self._startup_debug_print("before_network")
-        grpc_port = self._start_grpc(mux=mux) if self._serve_grpc else None
-        sock_port = self._start_sock(mux=mux) if self._serve_sock else None
-        self._startup_debug_print("after_network")
-        self._inform_used_ports(grpc_port=grpc_port, sock_port=sock_port)
-        self._startup_debug_print("after_inform")
+    def _setup_proctitle(
+        self, grpc_port: Optional[int], sock_port: Optional[int]
+    ) -> None:
+        # TODO: similar to _setup_tracelog, the internal_process should have
+        # a better way to have access to settings.
+        disable_setproctitle = os.environ.get("WANDB__DISABLE_SETPROCTITLE")
+        if disable_setproctitle:
+            return
+
         setproctitle = wandb.util.get_optional_module("setproctitle")
         if setproctitle:
             service_ver = 2
@@ -148,7 +148,20 @@ class WandbServer:
             # (consider unifying this in the future)
             service_id = f"{service_ver}-{pid}-{transport}-{port}"
             proc_title = f"wandb-service({service_id})"
+            self._startup_debug_print("before_setproctitle")
             setproctitle.setproctitle(proc_title)
+            self._startup_debug_print("after_setproctitle")
+
+    def serve(self) -> None:
+        self._setup_tracelog()
+        mux = StreamMux()
+        self._startup_debug_print("before_network")
+        grpc_port = self._start_grpc(mux=mux) if self._serve_grpc else None
+        sock_port = self._start_sock(mux=mux) if self._serve_sock else None
+        self._startup_debug_print("after_network")
+        self._inform_used_ports(grpc_port=grpc_port, sock_port=sock_port)
+        self._startup_debug_print("after_inform")
+        self._setup_proctitle(grpc_port=grpc_port, sock_port=sock_port)
         self._startup_debug_print("before_loop")
         mux.loop()
         self._stop_servers()

wandb/sdk/service/service.py

@@ -14,6 +14,7 @@ from typing import TYPE_CHECKING, Any, Dict, Optional
 
 from wandb import _sentry
 from wandb.errors import Error
+from wandb.util import get_module
 
 from . import _startup_debug, port_file
 from .service_base import ServiceInterface
@@ -168,9 +169,24 @@ class _Service:
             # Add coverage collection if needed
             if os.environ.get("YEA_RUN_COVERAGE") and os.environ.get("COVERAGE_RCFILE"):
                 exec_cmd_list += ["coverage", "run", "-m"]
-            service_args = [
-                "wandb",
-                "service",
+
+            service_args = []
+            if self._settings._require_nexus:
+                # NOTE: the wandb_core module will be distributed at first as an alpha
+                #       package as "wandb-core-alpha" to avoid polluting the pypi namespace.
+                #       When the package reaches compatibility milestones, it will be released
+                #       as "wandb-core".
+                wandb_nexus = get_module(
+                    "wandb_core",
+                    required="The nexus experiment requires the wandb_core module.",
+                )
+                nexus_path = wandb_nexus.get_nexus_path()
+                service_args.extend([nexus_path])
+                exec_cmd_list = []
+            else:
+                service_args.extend(["wandb", "service"])
+
+            service_args += [
                 "--port-filename",
                 fname,
                 "--pid",
@@ -181,11 +197,14 @@ class _Service:
                 service_args.append("--serve-grpc")
             else:
                 service_args.append("--serve-sock")
-            internal_proc = subprocess.Popen(
-                exec_cmd_list + service_args,
-                env=os.environ,
-                **kwargs,
-            )
+            try:
+                internal_proc = subprocess.Popen(
+                    exec_cmd_list + service_args,
+                    env=os.environ,
+                    **kwargs,
+                )
+            except Exception as e:
+                _sentry.reraise(e)
             self._startup_debug_print("wait_ports")
             try:
                 self._wait_for_ports(fname, proc=internal_proc)

wandb/sdk/verify/verify.py

@@ -12,8 +12,7 @@ from pkg_resources import parse_version
 from wandb_gql import gql
 
 import wandb
-from wandb.sdk.artifacts.local_artifact import Artifact as LocalArtifact
-from wandb.sdk.artifacts.public_artifact import Artifact as PublicArtifact
+from wandb.sdk.artifacts.artifact import Artifact
 from wandb.sdk.lib import runid
 
 from ...apis.internal import Api
@@ -224,7 +223,7 @@ def verify_manifest(
 
 
 def verify_digest(
-    downloaded: "PublicArtifact", computed: "PublicArtifact", fails_list: List[str]
+    downloaded: "Artifact", computed: "Artifact", fails_list: List[str]
 ) -> None:
     if downloaded.digest != computed.digest:
         fails_list.append(
@@ -234,7 +233,7 @@ def verify_digest(
 
 def artifact_with_path_or_paths(
     name: str, verify_dir: Optional[str] = None, singular: bool = False
-) -> "LocalArtifact":
+) -> "Artifact":
     art = wandb.Artifact(type="artsy", name=name)
     # internal file
     with open("verify_int_test.txt", "w") as f:
@@ -262,13 +261,13 @@ def artifact_with_path_or_paths(
 
 
 def log_use_download_artifact(
-    artifact: "LocalArtifact",
+    artifact: "Artifact",
     alias: str,
     name: str,
     download_dir: str,
     failed_test_strings: List[str],
     add_extra_file: bool,
-) -> Tuple[bool, Optional["PublicArtifact"], List[str]]:
+) -> Tuple[bool, Optional["Artifact"], List[str]]:
     with wandb.init(
         id=nice_id("log_artifact"),
         reinit=True,
@@ -357,9 +356,7 @@ def check_artifacts() -> bool:
     verify_digest(download_artifact, computed, failed_test_strings)
 
     computed_manifest = computed.manifest.to_manifest_json()["contents"]
-    downloaded_manifest = download_artifact._load_manifest().to_manifest_json()[
-        "contents"
-    ]
+    downloaded_manifest = download_artifact.manifest.to_manifest_json()["contents"]
     verify_manifest(downloaded_manifest, computed_manifest, failed_test_strings)
 
     print_results(failed_test_strings, False)

wandb/sdk/wandb_config.py

@@ -4,7 +4,6 @@ import logging
 from typing import Optional
 
 import wandb
-from wandb.sdk.artifacts.public_artifact import Artifact as PublicArtifact
 from wandb.util import (
     _is_artifact_representation,
     check_dict_contains_nested_artifact,
@@ -252,7 +251,7 @@ class Config:
         if _is_artifact_representation(val):
             val = self._artifact_callback(key, val)
         # if the user inserts an artifact into the config
-        if not (isinstance(val, wandb.Artifact) or isinstance(val, PublicArtifact)):
+        if not isinstance(val, wandb.Artifact):
             val = json_friendly_val(val)
         if not allow_val_change:
             if key in self._items and val != self._items[key]:
@@ -271,8 +270,7 @@ class Config:
         # best if we don't allow nested artifacts until we can lock nested keys in the config
         if isinstance(v, dict) and check_dict_contains_nested_artifact(v, nested):
             raise ValueError(
-                "Instances of wandb.Artifact and PublicArtifact"
-                " can only be top level keys in wandb.config"
+                "Instances of wandb.Artifact can only be top level keys in wandb.config"
             )
 
 

wandb/sdk/wandb_init.py

@@ -675,6 +675,8 @@ class _WandbInit:
                 tel.feature.flow_control_disabled = True
             if self.settings._flow_control_custom:
                 tel.feature.flow_control_custom = True
+            if self.settings._require_nexus:
+                tel.feature.nexus = True
 
             tel.env.maybe_mp = _maybe_mp_process(backend)
 

wandb/sdk/wandb_require.py

@@ -9,6 +9,7 @@ Example:
     wandb.require("incremental-artifacts@beta")
 """
 
+import os
 from typing import Optional, Sequence, Union
 
 import wandb
@@ -38,6 +39,12 @@ class _Requires:
     def require_service(self) -> None:
         self._require_service()
 
+    def _require_nexus(self) -> None:
+        os.environ["WANDB_REQUIRE_NEXUS"] = "True"
+
+    def require_nexus(self) -> None:
+        self._require_nexus()
+
     def apply(self) -> None:
         """Call require_* method for supported features."""
         last_message: str = ""

wandb/sdk/wandb_run.py

@@ -47,9 +47,7 @@ from wandb.proto.wandb_internal_pb2 import (
     RunRecord,
     ServerInfoResponse,
 )
-from wandb.sdk.artifacts.lazy_artifact import LazyArtifact
-from wandb.sdk.artifacts.local_artifact import Artifact as LocalArtifact
-from wandb.sdk.artifacts.public_artifact import Artifact as PublicArtifact
+from wandb.sdk.artifacts.artifact import Artifact
 from wandb.sdk.internal import job_builder
 from wandb.sdk.lib.import_hooks import (
     register_post_import_hook,
@@ -1051,7 +1049,7 @@ class Run:
         name: Optional[str] = None,
         include_fn: Callable[[str], bool] = _is_py_path,
         exclude_fn: Callable[[str], bool] = filenames.exclude_wandb_fn,
-    ) -> Optional[LocalArtifact]:
+    ) -> Optional[Artifact]:
         """Save the current state of your code to a W&B Artifact.
 
         By default, it walks the current directory and logs all files that end with `.py`.
@@ -1285,15 +1283,15 @@ class Run:
             self._backend.interface.publish_config(key=key, val=val, data=data)
 
     def _config_artifact_callback(
-        self, key: str, val: Union[str, LocalArtifact, dict]
-    ) -> Union[LocalArtifact, PublicArtifact]:
+        self, key: str, val: Union[str, Artifact, dict]
+    ) -> Artifact:
         # artifacts can look like dicts as they are passed into the run config
         # since the run config stores them on the backend as a dict with fields shown
         # in wandb.util.artifact_to_json
         if _is_artifact_version_weave_dict(val):
             assert isinstance(val, dict)
             public_api = self._public_api()
-            artifact = PublicArtifact.from_id(val["id"], public_api.client)
+            artifact = Artifact._from_id(val["id"], public_api.client)
             return self.use_artifact(artifact, use_as=key)
         elif _is_artifact_string(val):
             # this will never fail, but is required to make mypy happy
@@ -1306,12 +1304,11 @@ class Run:
             else:
                 public_api = self._public_api()
             if is_id:
-                artifact = PublicArtifact.from_id(artifact_string, public_api._client)
+                artifact = Artifact._from_id(artifact_string, public_api._client)
             else:
                 artifact = public_api.artifact(name=artifact_string)
             # in the future we'll need to support using artifacts from
-            # different instances of wandb. simplest way to do that is
-            # likely to convert the retrieved PublicArtifact to a LocalArtifact
+            # different instances of wandb.
 
             return self.use_artifact(artifact, use_as=key)
         elif _is_artifact_object(val):
@@ -2272,7 +2269,7 @@ class Run:
         source_dict: "JobSourceDict",
         installed_packages_list: List[str],
         patch_path: Optional[os.PathLike] = None,
-    ) -> "LocalArtifact":
+    ) -> "Artifact":
         job_artifact = job_builder.JobArtifact(name)
         if patch_path and os.path.exists(patch_path):
             job_artifact.add_file(FilePathStr(str(patch_path)), "diff.patch")
@@ -2290,7 +2287,7 @@ class Run:
         installed_packages_list: List[str],
         docker_image_name: Optional[str] = None,
         args: Optional[List[str]] = None,
-    ) -> Optional["LocalArtifact"]:
+    ) -> Optional["Artifact"]:
         docker_image_name = docker_image_name or os.getenv("WANDB_DOCKER")
 
         if not docker_image_name:
@@ -2314,7 +2311,7 @@ class Run:
 
     def _log_job_artifact_with_image(
         self, docker_image_name: str, args: Optional[List[str]] = None
-    ) -> LocalArtifact:
+    ) -> Artifact:
         packages, in_types, out_types = self._make_job_source_reqs()
         job_artifact = self._create_image_job(
             in_types,
@@ -2590,7 +2587,7 @@ class Run:
     @_run_decorator._attach
     def link_artifact(
         self,
-        artifact: Union[PublicArtifact, LocalArtifact],
+        artifact: Artifact,
         target_path: str,
         aliases: Optional[List[str]] = None,
     ) -> None:
@@ -2615,7 +2612,7 @@ class Run:
             aliases = []
 
         if self._backend and self._backend.interface:
-            if isinstance(artifact, LocalArtifact) and not artifact._logged_artifact:
+            if artifact.is_draft() and not artifact._is_draft_save_started():
                 artifact = self._log_artifact(artifact)
             if not self._settings._offline:
                 self._backend.interface.publish_link_artifact(
@@ -2634,11 +2631,11 @@ class Run:
     @_run_decorator._attach
     def use_artifact(
         self,
-        artifact_or_name: Union[str, PublicArtifact, LocalArtifact],
+        artifact_or_name: Union[str, Artifact],
         type: Optional[str] = None,
         aliases: Optional[List[str]] = None,
         use_as: Optional[str] = None,
-    ) -> Union[PublicArtifact, LocalArtifact]:
+    ) -> Artifact:
         """Declare an artifact as an input to a run.
 
         Call `download` or `file` on the returned object to get the contents locally.
@@ -2703,10 +2700,10 @@ class Run:
                 aliases = []
             elif isinstance(aliases, str):
                 aliases = [aliases]
-            if isinstance(artifact_or_name, wandb.Artifact):
+            if isinstance(artifact_or_name, Artifact) and artifact.is_draft():
                 if use_as is not None:
                     wandb.termwarn(
-                        "Indicating use_as is not supported when using an artifact with an instance of `wandb.Artifact`"
+                        "Indicating use_as is not supported when using a draft artifact"
                     )
                 self._log_artifact(
                     artifact,
@@ -2716,13 +2713,13 @@ class Run:
                 )
                 artifact.wait()
                 artifact._use_as = use_as or artifact.name
-            elif isinstance(artifact, PublicArtifact):
+            elif isinstance(artifact, Artifact) and not artifact.is_draft():
                 if (
                     self._launch_artifact_mapping
                     and artifact.name in self._launch_artifact_mapping.keys()
                 ):
                     wandb.termwarn(
-                        "Swapping artifacts is not supported when using an instance of `PublicArtifact`. "
+                        "Swapping artifacts is not supported when using a non-draft artifact. "
                         f"Using {artifact.name}."
                     )
                 artifact._use_as = use_as or artifact.name
@@ -2742,11 +2739,11 @@ class Run:
     @_run_decorator._attach
     def log_artifact(
         self,
-        artifact_or_path: Union[LocalArtifact, StrPath],
+        artifact_or_path: Union[Artifact, StrPath],
         name: Optional[str] = None,
         type: Optional[str] = None,
         aliases: Optional[List[str]] = None,
-    ) -> LocalArtifact:
+    ) -> Artifact:
         """Declare an artifact as an output of a run.
 
         Arguments:
@@ -2779,12 +2776,12 @@ class Run:
     @_run_decorator._attach
     def upsert_artifact(
         self,
-        artifact_or_path: Union[LocalArtifact, str],
+        artifact_or_path: Union[Artifact, str],
         name: Optional[str] = None,
         type: Optional[str] = None,
         aliases: Optional[List[str]] = None,
         distributed_id: Optional[str] = None,
-    ) -> LocalArtifact:
+    ) -> Artifact:
         """Declare (or append to) a non-finalized artifact as output of a run.
 
         Note that you must call run.finish_artifact() to finalize the artifact.
@@ -2833,12 +2830,12 @@ class Run:
     @_run_decorator._attach
     def finish_artifact(
         self,
-        artifact_or_path: Union[LocalArtifact, str],
+        artifact_or_path: Union[Artifact, str],
         name: Optional[str] = None,
         type: Optional[str] = None,
         aliases: Optional[List[str]] = None,
         distributed_id: Optional[str] = None,
-    ) -> LocalArtifact:
+    ) -> Artifact:
         """Finishes a non-finalized artifact as output of a run.
 
         Subsequent "upserts" with the same distributed ID will result in a new version.
@@ -2885,7 +2882,7 @@ class Run:
 
     def _log_artifact(
         self,
-        artifact_or_path: Union[LocalArtifact, StrPath],
+        artifact_or_path: Union[Artifact, StrPath],
         name: Optional[str] = None,
         type: Optional[str] = None,
         aliases: Optional[List[str]] = None,
@@ -2893,7 +2890,7 @@ class Run:
         finalize: bool = True,
         is_user_created: bool = False,
         use_after_commit: bool = False,
-    ) -> LocalArtifact:
+    ) -> Artifact:
         api = internal.Api()
         if api.settings().get("anonymous") == "true":
             wandb.termwarn(
@@ -2922,7 +2919,7 @@ class Run:
                     is_user_created=is_user_created,
                     use_after_commit=use_after_commit,
                 )
-                artifact._logged_artifact = LazyArtifact(self._public_api(), future)
+                artifact._set_save_future(future, self._public_api().client)
             else:
                 self._backend.interface.publish_artifact(
                     self,
@@ -2955,11 +2952,11 @@ class Run:
         if not self._settings._offline:
             try:
                 public_api = self._public_api()
-                expected_type = PublicArtifact.expected_type(
-                    public_api.client,
-                    artifact.name,
+                expected_type = Artifact._expected_type(
                     public_api.settings["entity"],
                     public_api.settings["project"],
+                    artifact.name,
+                    public_api.client,
                 )
             except requests.exceptions.RequestException:
                 # Just return early if there is a network error. This is
@@ -2974,11 +2971,11 @@ class Run:
 
     def _prepare_artifact(
         self,
-        artifact_or_path: Union[LocalArtifact, StrPath],
+        artifact_or_path: Union[Artifact, StrPath],
         name: Optional[str] = None,
         type: Optional[str] = None,
         aliases: Optional[List[str]] = None,
-    ) -> Tuple[LocalArtifact, List[str]]:
+    ) -> Tuple[Artifact, List[str]]:
         if isinstance(artifact_or_path, (str, os.PathLike)):
             name = name or f"run-{self._run_id}-{os.path.basename(artifact_or_path)}"
             artifact = wandb.Artifact(name, type or "unspecified")

wandb/sdk/wandb_settings.py

@@ -457,10 +457,11 @@ class Settings:
     _config_dict: Config
     _console: SettingsConsole
     _cuda: str
-    _disable_meta: bool
-    _disable_service: bool
-    _disable_stats: bool
+    _disable_meta: bool  # Do not collect system metadata
+    _disable_service: bool  # Disable wandb-service, spin up internal process the old way
+    _disable_stats: bool  # Do not collect system metrics
     _disable_viewer: bool  # Prevent early viewer query
+    _disable_setproctitle: bool  # Do not use setproctitle on internal process
     _except_exit: bool
     _executable: str
     _extra_http_headers: Mapping[str, str]
@@ -486,6 +487,7 @@ class Settings:
     _platform: str
     _python: str
     _runqueue_item_id: str
+    _require_nexus: bool
     _save_requirements: bool
     _service_transport: str
     _service_wait: float
@@ -508,6 +510,7 @@ class Settings:
     allow_val_change: bool
     anonymous: str
     api_key: str
+    azure_account_url_to_access_key: Dict[str, str]
     base_url: str  # The base url for the wandb api
     code_dir: str
     config_paths: Sequence[str]
@@ -616,6 +619,7 @@ class Settings:
                 "preprocessor": _str_as_bool,
                 "is_policy": True,
             },
+            _disable_setproctitle={"value": False, "preprocessor": _str_as_bool},
             _disable_stats={"preprocessor": _str_as_bool},
             _disable_viewer={"preprocessor": _str_as_bool},
             _extra_http_headers={"preprocessor": _str_as_json},
@@ -662,6 +666,7 @@ class Settings:
             },
             _sync={"value": False},
             _platform={"value": util.get_platform_name()},
+            _require_nexus={"value": False, "preprocessor": _str_as_bool},
             _save_requirements={"value": True, "preprocessor": _str_as_bool},
             _service_wait={
                 "value": 30,
@@ -700,6 +705,7 @@ class Settings:
             },
             anonymous={"validator": self._validate_anonymous},
             api_key={"validator": self._validate_api_key},
+            azure_account_url_to_access_key={"value": {}},
             base_url={
                 "value": "https://api.wandb.ai",
                 "preprocessor": lambda x: str(x).strip().rstrip("/"),
@@ -1559,6 +1565,7 @@ class Settings:
             "WANDB_TRACELOG": "_tracelog",
             "WANDB_DISABLE_SERVICE": "_disable_service",
             "WANDB_SERVICE_TRANSPORT": "_service_transport",
+            "WANDB_REQUIRE_NEXUS": "_require_nexus",
             "WANDB_DIR": "root_dir",
             "WANDB_NAME": "run_name",
             "WANDB_NOTES": "run_notes",

wandb/testing/relay.py

@@ -4,6 +4,7 @@ import logging
 import socket
 import sys
 import threading
+import traceback
 import urllib.parse
 from collections import defaultdict, deque
 from copy import deepcopy
@@ -294,7 +295,9 @@ class QueryResolver:
             }
             post_processed_data = {
                 "name": name,
-                "dropped": [request_data["dropped"]],
+                "dropped": [request_data["dropped"]]
+                if "dropped" in request_data
+                else [],
                 "files": files,
             }
             return post_processed_data
@@ -642,7 +645,17 @@ class RelayServer:
         }
         self.context.raw_data.append(raw_data)
 
-        snooped_context = self.resolver.resolve(request_data, response_data, **kwargs)
+        try:
+            snooped_context = self.resolver.resolve(
+                request_data,
+                response_data,
+                **kwargs,
+            )
+        except Exception as e:
+            print("Failed to resolve context: ", e)
+            traceback.print_exc()
+            snooped_context = None
+
         if snooped_context is not None:
             self.context.upsert(snooped_context)