wandb 0.16.5__py3-none-any.whl → 0.17.0__py3-none-any.whl

wandb/sdk/launch/utils.py

@@ -1,4 +1,3 @@
-# heavily inspired by https://github.com/mlflow/mlflow/blob/master/mlflow/projects/utils.py
 import asyncio
 import json
 import logging
@@ -16,7 +15,6 @@ import wandb
 import wandb.docker as docker
 from wandb import util
 from wandb.apis.internal import Api
-from wandb.errors import CommError
 from wandb.sdk.launch.errors import LaunchError
 from wandb.sdk.launch.git_reference import GitReference
 from wandb.sdk.launch.wandb_reference import WandbReference
@@ -32,7 +30,6 @@ FAILED_PACKAGES_REGEX = re.compile(
 )
 
 if TYPE_CHECKING:  # pragma: no cover
-    from wandb.sdk.artifacts.artifact import Artifact
     from wandb.sdk.launch.agent.job_status_tracker import JobAndRunStatusTracker
 
 
@@ -57,15 +54,15 @@ API_KEY_REGEX = r"WANDB_API_KEY=\w+(-\w+)?"
 MACRO_REGEX = re.compile(r"\$\{(\w+)\}")
 
 AZURE_CONTAINER_REGISTRY_URI_REGEX = re.compile(
-    r"(?:https://)?([\w]+)\.azurecr\.io/([\w\-]+):?(.*)"
+    r"^(?:https://)?([\w]+)\.azurecr\.io/(?P<repository>[\w\-]+):?(?P<tag>.*)"
 )
 
 ELASTIC_CONTAINER_REGISTRY_URI_REGEX = re.compile(
-    r"^(?P<account>.*)\.dkr\.ecr\.(?P<region>.*)\.amazonaws\.com/(?P<repository>.*)/?$"
+    r"^(?:https://)?(?P<account>[\w-]+)\.dkr\.ecr\.(?P<region>[\w-]+)\.amazonaws\.com/(?P<repository>[\w-]+):?(?P<tag>.*)$"
 )
 
 GCP_ARTIFACT_REGISTRY_URI_REGEX = re.compile(
-    r"^(?P<region>[\w-]+)-docker\.pkg\.dev/(?P<project>[\w-]+)/(?P<repository>[\w-]+)/(?P<image_name>[\w-]+)$",
+    r"^(?:https://)?(?P<region>[\w-]+)-docker\.pkg\.dev/(?P<project>[\w-]+)/(?P<repository>[\w-]+)/?(?P<image_name>[\w-]+)?(?P<tag>:.*)?$",
     re.IGNORECASE,
 )
 
@@ -316,16 +313,13 @@ def construct_launch_spec(
 
 
 def validate_launch_spec_source(launch_spec: Dict[str, Any]) -> None:
-    uri = launch_spec.get("uri")
     job = launch_spec.get("job")
     docker_image = launch_spec.get("docker", {}).get("docker_image")
-
-    if not bool(uri) and not bool(job) and not bool(docker_image):
-        raise LaunchError("Must specify a uri, job or docker image")
-    elif bool(uri) and bool(docker_image):
-        raise LaunchError("Found both uri and docker-image, only one can be set")
-    elif sum(map(bool, [uri, job, docker_image])) > 1:
-        raise LaunchError("Must specify exactly one of uri, job or image")
+    if bool(job) == bool(docker_image):
+        raise LaunchError(
+            "Exactly one of job or docker_image must be specified in the launch "
+            "spec."
+        )
 
 
 def parse_wandb_uri(uri: str) -> Tuple[str, str, str]:
@@ -336,77 +330,6 @@ def parse_wandb_uri(uri: str) -> Tuple[str, str, str]:
     return (ref.entity, ref.project, ref.run_id)
 
 
-def is_bare_wandb_uri(uri: str) -> bool:
-    """Check that a wandb uri is valid.
-
-    URI must be in the format
-    `/<entity>/<project>/runs/<run_name>[other stuff]`
-    or
-    `/<entity>/<project>/artifacts/job/<job_name>[other stuff]`.
-    """
-    _logger.info(f"Checking if uri {uri} is bare...")
-    return uri.startswith("/") and WandbReference.is_uri_job_or_run(uri)
-
-
-def fetch_wandb_project_run_info(
-    entity: str, project: str, run_name: str, api: Api
-) -> Any:
-    _logger.info("Fetching run info...")
-    try:
-        result = api.get_run_info(entity, project, run_name)
-    except CommError:
-        result = None
-    if result is None:
-        raise LaunchError(
-            f"Run info is invalid or doesn't exist for {api.settings('base_url')}/{entity}/{project}/runs/{run_name}"
-        )
-    if result.get("codePath") is None:
-        # TODO: we don't currently expose codePath in the runInfo endpoint, this downloads
-        # it from wandb-metadata.json if we can.
-        metadata = api.download_url(
-            project, "wandb-metadata.json", run=run_name, entity=entity
-        )
-        if metadata is not None:
-            _, response = api.download_file(metadata["url"])
-            data = response.json()
-            result["codePath"] = data.get("codePath")
-            result["cudaVersion"] = data.get("cuda", None)
-
-    return result
-
-
-def download_entry_point(
-    entity: str, project: str, run_name: str, api: Api, entry_point: str, dir: str
-) -> bool:
-    metadata = api.download_url(
-        project, f"code/{entry_point}", run=run_name, entity=entity
-    )
-    if metadata is not None:
-        _, response = api.download_file(metadata["url"])
-        with util.fsync_open(os.path.join(dir, entry_point), "wb") as file:
-            for data in response.iter_content(chunk_size=1024):
-                file.write(data)
-        return True
-    return False
-
-
-def download_wandb_python_deps(
-    entity: str, project: str, run_name: str, api: Api, dir: str
-) -> Optional[str]:
-    reqs = api.download_url(project, "requirements.txt", run=run_name, entity=entity)
-    if reqs is not None:
-        _logger.info("Downloading python dependencies")
-        _, response = api.download_file(reqs["url"])
-
-        with util.fsync_open(
-            os.path.join(dir, "requirements.frozen.txt"), "wb"
-        ) as file:
-            for data in response.iter_content(chunk_size=1024):
-                file.write(data)
-        return "requirements.frozen.txt"
-    return None
-
-
 def get_local_python_deps(
     dir: str, filename: str = "requirements.local.txt"
 ) -> Optional[str]:
@@ -498,19 +421,6 @@ def validate_wandb_python_deps(
     _logger.warning("Unable to validate local python dependencies")
 
 
-def fetch_project_diff(
-    entity: str, project: str, run_name: str, api: Api
-) -> Optional[str]:
-    """Fetches project diff from wandb servers."""
-    _logger.info("Searching for diff.patch")
-    patch = None
-    try:
-        (_, _, patch, _) = api.run_config(project, run_name, entity)
-    except CommError:
-        pass
-    return patch
-
-
 def apply_patch(patch_string: str, dst_dir: str) -> None:
     """Applies a patch file to a directory."""
     _logger.info("Applying diff.patch")
@@ -531,17 +441,6 @@ def apply_patch(patch_string: str, dst_dir: str) -> None:
         raise wandb.Error("Failed to apply diff.patch associated with run.")
 
 
-def _make_refspec_from_version(version: Optional[str]) -> List[str]:
-    """Create a refspec that checks for the existence of origin/main and the version."""
-    if version:
-        return [f"+{version}"]
-
-    return [
-        "+refs/heads/main*:refs/remotes/origin/main*",
-        "+refs/heads/master*:refs/remotes/origin/master*",
-    ]
-
-
 def _fetch_git_repo(dst_dir: str, uri: str, version: Optional[str]) -> Optional[str]:
     """Clones the git repo at ``uri`` into ``dst_dir``.
 
@@ -561,13 +460,6 @@ def _fetch_git_repo(dst_dir: str, uri: str, version: Optional[str]) -> Optional[
     return version
 
 
-def merge_parameters(
-    higher_priority_params: Dict[str, Any], lower_priority_params: Dict[str, Any]
-) -> Dict[str, Any]:
-    """Merge the contents of two dicts, keeping values from higher_priority_params if there are conflicts."""
-    return {**lower_priority_params, **higher_priority_params}
-
-
 def convert_jupyter_notebook_to_script(fname: str, project_dir: str) -> str:
     nbconvert = wandb.util.get_module(
         "nbconvert", "nbformat and nbconvert are required to use launch with notebooks"
@@ -597,25 +489,6 @@ def convert_jupyter_notebook_to_script(fname: str, project_dir: str) -> str:
     return new_name
 
 
-def check_and_download_code_artifacts(
-    entity: str, project: str, run_name: str, internal_api: Api, project_dir: str
-) -> Optional["Artifact"]:
-    _logger.info("Checking for code artifacts")
-    public_api = wandb.PublicApi(
-        overrides={"base_url": internal_api.settings("base_url")}
-    )
-
-    run = public_api.run(f"{entity}/{project}/{run_name}")
-    run_artifacts = run.logged_artifacts()
-
-    for artifact in run_artifacts:
-        if hasattr(artifact, "type") and artifact.type == "code":
-            artifact.download(project_dir)
-            return artifact  # type: ignore
-
-    return None
-
-
 def to_camel_case(maybe_snake_str: str) -> str:
     if "_" not in maybe_snake_str:
         return maybe_snake_str
@@ -623,11 +496,6 @@ def to_camel_case(maybe_snake_str: str) -> str:
     return "".join(x.title() if x else "_" for x in components)
 
 
-def run_shell(args: List[str]) -> Tuple[str, str]:
-    out = subprocess.run(args, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
-    return out.stdout.decode("utf-8").strip(), out.stderr.decode("utf-8").strip()
-
-
 def validate_build_and_registry_configs(
     build_config: Dict[str, Any], registry_config: Dict[str, Any]
 ) -> None:
@@ -846,3 +714,28 @@ def fetch_and_validate_template_variables(
             raise LaunchError(f"Value for {key} must be of type {field_type}.")
         template_variables[key] = val
     return template_variables
+
+
+def get_entrypoint_file(entrypoint: List[str]) -> Optional[str]:
+    """Get the entrypoint file from the given command.
+
+    Args:
+        entrypoint (List[str]): List of command and arguments.
+
+    Returns:
+        Optional[str]: The entrypoint file if found, otherwise None.
+    """
+    if not entrypoint:
+        return None
+    if entrypoint[0].endswith(".py") or entrypoint[0].endswith(".sh"):
+        return entrypoint[0]
+    if len(entrypoint) < 2:
+        return None
+    return entrypoint[1]
+
+
+def get_current_python_version() -> Tuple[str, str]:
+    full_version = sys.version.split()[0].split(".")
+    major = full_version[0]
+    version = ".".join(full_version[:2]) if len(full_version) >= 2 else major + ".0"
+    return version, major

wandb/sdk/lib/_settings_toposort_generated.py

@@ -13,7 +13,6 @@ else:
 _Setting = Literal[
     "_args",
     "_aws_lambda",
-    "_async_upload_concurrency_limit",
     "_cli_only_mode",
     "_code_path_local",
     "_colab",
@@ -22,9 +21,9 @@ _Setting = Literal[
     "_disable_service",
     "_disable_setproctitle",
     "_disable_stats",
+    "_disable_update_check",
     "_disable_viewer",
     "_disable_machine_info",
-    "_except_exit",
     "_executable",
     "_extra_http_headers",
     "_file_stream_retry_max",
@@ -125,7 +124,6 @@ _Setting = Literal[
     "login_timeout",
     "mode",
     "notebook_name",
-    "problem",
     "program",
     "program_abspath",
     "program_relpath",
@@ -178,7 +176,6 @@ _Setting = Literal[
 ]
 
 SETTINGS_TOPOLOGICALLY_SORTED: Final[Tuple[_Setting, ...]] = (
-    "_async_upload_concurrency_limit",
     "_service_wait",
     "_stats_sample_rate_seconds",
     "_stats_samples_to_average",
@@ -188,7 +185,6 @@ SETTINGS_TOPOLOGICALLY_SORTED: Final[Tuple[_Setting, ...]] = (
     "console",
     "job_source",
     "mode",
-    "problem",
     "project",
     "run_id",
     "start_method",

wandb/sdk/lib/fsm.py

@@ -52,43 +52,39 @@ T_FsmContext_contra = TypeVar("T_FsmContext_contra", contravariant=True)
 @runtime_checkable
 class FsmStateCheck(Protocol[T_FsmInputs]):
     @abstractmethod
-    def on_check(self, inputs: T_FsmInputs) -> None:
-        ...  # pragma: no cover
+    def on_check(self, inputs: T_FsmInputs) -> None: ...  # pragma: no cover
 
 
 @runtime_checkable
 class FsmStateOutput(Protocol[T_FsmInputs]):
     @abstractmethod
-    def on_state(self, inputs: T_FsmInputs) -> None:
-        ...  # pragma: no cover
+    def on_state(self, inputs: T_FsmInputs) -> None: ...  # pragma: no cover
 
 
 @runtime_checkable
 class FsmStateEnter(Protocol[T_FsmInputs]):
     @abstractmethod
-    def on_enter(self, inputs: T_FsmInputs) -> None:
-        ...  # pragma: no cover
+    def on_enter(self, inputs: T_FsmInputs) -> None: ...  # pragma: no cover
 
 
 @runtime_checkable
 class FsmStateEnterWithContext(Protocol[T_FsmInputs, T_FsmContext_contra]):
     @abstractmethod
-    def on_enter(self, inputs: T_FsmInputs, context: T_FsmContext_contra) -> None:
-        ...  # pragma: no cover
+    def on_enter(
+        self, inputs: T_FsmInputs, context: T_FsmContext_contra
+    ) -> None: ...  # pragma: no cover
 
 
 @runtime_checkable
 class FsmStateStay(Protocol[T_FsmInputs]):
     @abstractmethod
-    def on_stay(self, inputs: T_FsmInputs) -> None:
-        ...  # pragma: no cover
+    def on_stay(self, inputs: T_FsmInputs) -> None: ...  # pragma: no cover
 
 
 @runtime_checkable
 class FsmStateExit(Protocol[T_FsmInputs, T_FsmContext_cov]):
     @abstractmethod
-    def on_exit(self, inputs: T_FsmInputs) -> T_FsmContext_cov:
-        ...  # pragma: no cover
+    def on_exit(self, inputs: T_FsmInputs) -> T_FsmContext_cov: ...  # pragma: no cover
 
 
 # It would be nice if python provided optional protocol members, but it doesnt as described here:

wandb/sdk/lib/gitlib.py

@@ -14,7 +14,7 @@ try:
         Repo,
     )
 except ImportError:
-    Repo = None
+    Repo = None  # type: ignore
 
 if TYPE_CHECKING:
     from git import Repo
@@ -121,7 +121,7 @@ class GitRepo:
         # TODO: Saw a user getting a Unicode decode error when parsing refs,
         # more details on implementing a real fix in [WB-4064]
         try:
-            if len(self.repo.refs) > 0:
+            if len(self.repo.refs) > 0:  # type: ignore[arg-type]
                 return self.repo.head.commit.hexsha
             else:
                 return self.repo.git.show_ref("--head").split(" ")[0]
@@ -140,7 +140,7 @@ class GitRepo:
         if not self.repo:
             return None
         try:
-            return self.repo.remotes[self.remote_name]
+            return self.repo.remotes[self.remote_name]  # type: ignore[index]
         except IndexError:
             return None
 
@@ -200,7 +200,7 @@ class GitRepo:
                     possible_relatives.append(tracking_branch.commit)
 
             if not possible_relatives:
-                for branch in self.repo.branches:
+                for branch in self.repo.branches:  # type: ignore[attr-defined]
                     tracking_branch = branch.tracking_branch()
                     if tracking_branch is not None:
                         possible_relatives.append(tracking_branch.commit)

wandb/sdk/lib/import_hooks.py

@@ -143,7 +143,7 @@ class _ImportHookChainedLoader:
         # None, so handle None as well. The module may not support attribute
         # assignment, in which case we simply skip it. Note that we also deal
         # with __loader__ not existing at all. This is to future proof things
-        # due to proposal to remove the attribue as described in the GitHub
+        # due to proposal to remove the attribute as described in the GitHub
         # issue at https://github.com/python/cpython/issues/77458. Also prior
         # to Python 3.3, the __loader__ attribute was only set if a custom
         # module loader was used. It isn't clear whether the attribute still

wandb/sdk/lib/lazyloader.py

@@ -1,6 +1,5 @@
 """module lazyloader."""
 
-
 import importlib
 import sys
 import types

wandb/sdk/lib/proto_util.py

@@ -12,7 +12,28 @@ if TYPE_CHECKING:  # pragma: no cover
 
 
 def dict_from_proto_list(obj_list: "RepeatedCompositeFieldContainer") -> Dict[str, Any]:
-    return {item.key: json.loads(item.value_json) for item in obj_list}
+    result: Dict[str, Any] = {}
+
+    for item in obj_list:
+        # Start from the root of the result dict
+        current_level = result
+
+        if len(item.nested_key) > 0:
+            keys = list(item.nested_key)
+        else:
+            keys = [item.key]
+
+        for key in keys[:-1]:
+            if key not in current_level:
+                current_level[key] = {}
+            # Move the reference deeper into the nested dictionary
+            current_level = current_level[key]
+
+        # Set the value at the final key location, parsing JSON from the value_json field
+        final_key = keys[-1]
+        current_level[final_key] = json.loads(item.value_json)
+
+    return result
 
 
 def _result_from_record(record: "pb.Record") -> "pb.Result":
@@ -29,7 +50,7 @@ def _assign_end_offset(record: "pb.Record", end_offset: int) -> None:
 
 
 def proto_encode_to_dict(
-    pb_obj: Union["tpb.TelemetryRecord", "pb.MetricRecord"]
+    pb_obj: Union["tpb.TelemetryRecord", "pb.MetricRecord"],
 ) -> Dict[int, Any]:
     data: Dict[int, Any] = dict()
     fields = pb_obj.ListFields()

wandb/sdk/lib/redirect.py

@@ -224,7 +224,7 @@ class TerminalEmulator:
     def carriage_return(self):
         self.cursor.x = 0
 
-    def cursor_postion(self, line, column):
+    def cursor_position(self, line, column):
         self.cursor.x = min(column, 1) - 1
         self.cursor.y = min(line, 1) - 1
 
@@ -393,25 +393,30 @@ class TerminalEmulator:
                         p = (int(p[0]), 1)
                     else:
                         p = (1, 1)
-                    self.cursor_postion(*p)
+                    self.cursor_position(*p)
         except Exception:
             pass
 
     def _get_line(self, n):
         line = self.buffer[n]
         line_len = self._get_line_len(n)
-        # We have to loop through each character in the line and check if foreground, background and
-        # other attributes (italics, bold, underline, etc) of the ith character are different from those of the
-        # (i-1)th character. If different, the appropriate ascii character for switching the color/attribute
-        # should be appended to the output string before appending the actual character. This loop and subsequent
-        # checks can be expensive, especially because 99% of terminal output use default colors and formatting. Even
-        # in outputs that do contain colors and styles, its unlikely that they will change on a per character basis.
-
-        # So instead we create a character list without any ascii codes (`out`), and a list of all the foregrounds
-        # in the line (`fgs`) on which we call np.diff() and np.where() to find the indices where the foreground change,
-        # and insert the ascii characters in the output list (`out`) on those indices. All of this is the done ony if
-        # there are more than 1 foreground color in the line in the first place (`if len(set(fgs)) > 1 else None`).
-        # Same logic is repeated for background colors and other attributes.
+        # We have to loop through each character in the line and check if foreground,
+        # background and other attributes (italics, bold, underline, etc) of the ith
+        # character are different from those of the (i-1)th character. If different, the
+        # appropriate ascii character for switching the color/attribute should be
+        # appended to the output string before appending the actual character. This loop
+        # and subsequent checks can be expensive, especially because 99% of terminal
+        # output use default colors and formatting. Even in outputs that do contain
+        # colors and styles, its unlikely that they will change on a per character
+        # basis.
+
+        # So instead we create a character list without any ascii codes (`out`), and a
+        # list of all the foregrounds in the line (`fgs`) on which we call np.diff() and
+        # np.where() to find the indices where the foreground change, and insert the
+        # ascii characters in the output list (`out`) on those indices. All of this is
+        # the done only if there are more than 1 foreground color in the line in the
+        # first place (`if len(set(fgs)) > 1 else None`). Same logic is repeated for
+        # background colors and other attributes.
 
         out = [line[i].data for i in range(line_len)]
 

wandb/sdk/lib/retry.py

@@ -248,8 +248,9 @@ class ExponentialBackoff(Backoff):
         if self._timeout_at is not None and NOW_FN() > self._timeout_at:
             raise exc
 
-        result, self._next_sleep = self._next_sleep, min(
-            self._max_sleep, self._next_sleep * (1 + random.random())
+        result, self._next_sleep = (
+            self._next_sleep,
+            min(self._max_sleep, self._next_sleep * (1 + random.random())),
         )
 
         return result

wandb/sdk/lib/run_moment.py

@@ -1,7 +1,13 @@
+import sys
 from dataclasses import dataclass
-from typing import Literal, Union, cast
+from typing import Union, cast
 from urllib import parse
 
+if sys.version_info >= (3, 8):
+    from typing import Literal
+else:
+    from typing_extensions import Literal
+
 _STEP = Literal["_step"]
 
 

wandb/sdk/lib/tracelog.py

@@ -6,7 +6,7 @@ Functions:
     log_message_send    - message sent to socket
     log_message_recv    - message received from socket
     log_message_process - message processed by thread
-    log_message_link    - message linked to another mesage
+    log_message_link    - message linked to another message
     log_message_assert  - message encountered problem
 
 """

wandb/sdk/service/service.py

@@ -2,6 +2,7 @@
 
 Backend server process can be connected to using tcp sockets transport.
 """
+
 import datetime
 import os
 import pathlib
@@ -14,8 +15,8 @@ import time
 from typing import TYPE_CHECKING, Any, Dict, Optional
 
 from wandb import _sentry, termlog
-from wandb.env import error_reporting_enabled
-from wandb.errors import Error
+from wandb.env import core_debug, core_error_reporting_enabled, is_require_core
+from wandb.errors import Error, WandbCoreNotAvailableError
 from wandb.sdk.lib.wburls import wburls
 from wandb.util import get_core_path, get_module
 
@@ -109,7 +110,8 @@ class _Service:
                     f"The wandb service process exited with {proc.returncode}. "
                     "Ensure that `sys.executable` is a valid python interpreter. "
                     "You can override it with the `_executable` setting "
-                    "or with the `WANDB__EXECUTABLE` environment variable.",
+                    "or with the `WANDB__EXECUTABLE` environment variable."
+                    f"\n{context}",
                     context=context,
                 )
             if not os.path.isfile(fname):
@@ -161,28 +163,29 @@ class _Service:
                 exec_cmd_list += ["coverage", "run", "-m"]
 
             service_args = []
-            # NOTE: "wandb-core" is the name of the package that will be distributed
-            #       as the stable version of the wandb core library.
-            #
-            #       Environment variable _WANDB_CORE_PATH is a temporary development feature
-            #       to assist in running the core service from a live development directory.
-            core_path = get_core_path()
-            if core_path:
+
+            if is_require_core():
+                try:
+                    core_path = get_core_path()
+                except WandbCoreNotAvailableError as e:
+                    _sentry.reraise(e)
+
                 service_args.extend([core_path])
-                if not error_reporting_enabled():
+
+                if not core_error_reporting_enabled(default="True"):
                     service_args.append("--no-observability")
-                if os.environ.get("WANDB_CORE_DEBUG", False):
+
+                if core_debug(default="False"):
                     service_args.append("--debug")
+
                 trace_filename = os.environ.get("_WANDB_TRACE")
                 if trace_filename is not None:
                     service_args.extend(["--trace", trace_filename])
 
                 exec_cmd_list = []
-                # TODO: remove this after the wandb-core GA release
-                wandb_core = get_module("wandb_core")
                 termlog(
-                    f"Using wandb-core version {wandb_core.__version__} as the SDK backend. "
-                    f"Please refer to {wburls.get('wandb_core')} for more information.",
+                    "Using wandb-core as the SDK backend."
+                    f" Please refer to {wburls.get('wandb_core')} for more information.",
                     repeat=False,
                 )
             else:

wandb/sdk/verify/verify.py

@@ -1,4 +1,5 @@
 """Utilities for wandb verify."""
+
 import getpass
 import os
 import time
@@ -20,7 +21,7 @@ PROJECT_NAME = "verify"
 GET_RUN_MAX_TIME = 10
 MIN_RETRYS = 3
 CHECKMARK = "\u2705"
-RED_X = "\u274C"
+RED_X = "\u274c"
 ID_PREFIX = runid.generate_id()