wandb 0.16.6__py3-none-any.whl → 0.17.0rc1__py3-none-any.whl

wandb/sdk/launch/agent/config.py

@@ -80,17 +80,7 @@ class RegistryConfig(BaseModel):
     @validator("uri")  # type: ignore
     @classmethod
     def validate_uri(cls, uri: str) -> str:
-        for regex in [
-            GCP_ARTIFACT_REGISTRY_URI_REGEX,
-            AZURE_CONTAINER_REGISTRY_URI_REGEX,
-            ELASTIC_CONTAINER_REGISTRY_URI_REGEX,
-        ]:
-            if regex.match(uri):
-                return uri
-        raise ValueError(
-            "Invalid uri. URI must be a repository URI for an "
-            "ECR, ACR, or GCP Artifact Registry."
-        )
+        return validate_registry_uri(uri)
 
 
 class EnvironmentConfig(BaseModel):
@@ -186,6 +176,14 @@ class BuilderConfig(BaseModel):
         """Right now there are no required fields for docker builds."""
         return values
 
+    @validator("destination")  # type: ignore
+    @classmethod
+    def validate_destination(cls, destination: Optional[str]) -> Optional[str]:
+        """Validate that the destination is a valid container registry URI."""
+        if destination is None:
+            return None
+        return validate_registry_uri(destination)
+
 
 class AgentConfig(BaseModel):
     """Configuration for the Launch agent."""
@@ -236,3 +234,66 @@ class AgentConfig(BaseModel):
 
     class Config:
         extra = "forbid"
+
+
+def validate_registry_uri(uri: str) -> str:
+    """Validate that the registry URI is a valid container registry URI.
+
+    The URI should resolve to an image name in a container registry. The recognized
+    formats are for ECR, ACR, and GCP Artifact Registry. If the URI does not match
+    any of these formats, a warning is printed indicating the registry type is not
+    recognized and the agent can't guarantee that images can be pushed.
+
+    If the format is recognized but does not resolve to an image name, an
+    error is raised. For example, if the URI is an ECR URI but does not include
+    an image name or includes a tag as well as an image name, an error is raised.
+    """
+    tag_msg = (
+        "Destination for built images may not include a tag, but the URI provided "
+        "includes the suffix '{tag}'. Please remove the tag and try again. The agent "
+        "will automatically tag each image with a unique hash of the source code."
+    )
+    if uri.startswith("https://"):
+        uri = uri[8:]
+
+    match = GCP_ARTIFACT_REGISTRY_URI_REGEX.match(uri)
+    if match:
+        if match.group("tag"):
+            raise ValueError(tag_msg.format(tag=match.group("tag")))
+        if not match.group("image_name"):
+            raise ValueError(
+                "An image name must be specified in the URI for a GCP Artifact Registry. "
+                "Please provide a uri with the format "
+                "'https://<region>-docker.pkg.dev/<project>/<repository>/<image>'."
+            )
+        return uri
+
+    match = AZURE_CONTAINER_REGISTRY_URI_REGEX.match(uri)
+    if match:
+        if match.group("tag"):
+            raise ValueError(tag_msg.format(tag=match.group("tag")))
+        if not match.group("repository"):
+            raise ValueError(
+                "A repository name must be specified in the URI for an "
+                "Azure Container Registry. Please provide a uri with the format "
+                "'https://<registry-name>.azurecr.io/<repository>'."
+            )
+        return uri
+
+    match = ELASTIC_CONTAINER_REGISTRY_URI_REGEX.match(uri)
+    if match:
+        if match.group("tag"):
+            raise ValueError(tag_msg.format(tag=match.group("tag")))
+        if not match.group("repository"):
+            raise ValueError(
+                "A repository name must be specified in the URI for an "
+                "Elastic Container Registry. Please provide a uri with the format "
+                "'https://<account-id>.dkr.ecr.<region>.amazonaws.com/<repository>'."
+            )
+        return uri
+
+    wandb.termwarn(
+        f"Unable to recognize registry type in URI {uri}. You are responsible "
+        "for ensuring the agent can push images to this registry."
+    )
+    return uri

wandb/sdk/launch/builder/abstract.py

@@ -1,4 +1,5 @@
 """Abstract plugin class defining the interface needed to build container images for W&B Launch."""
+
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Dict, Optional
 

wandb/sdk/launch/builder/build.py

@@ -237,7 +237,11 @@ def get_base_setup(
 
     CPU version is built on python, Accelerator version is built on user provided.
     """
-    python_base_image = f"python:{py_version}-buster"
+    minor = int(py_version.split(".")[1])
+    if minor < 12:
+        python_base_image = f"python:{py_version}-buster"
+    else:
+        python_base_image = f"python:{py_version}-bookworm"
     if launch_project.accelerator_base_image:
         _logger.info(
             f"Using accelerator base image: {launch_project.accelerator_base_image}"
@@ -311,6 +315,11 @@ def get_env_vars_dict(
     _inject_wandb_config_env_vars(
         launch_project.override_config, env_vars, max_env_length
     )
+
+    _inject_file_overrides_env_vars(
+        launch_project.override_files, env_vars, max_env_length
+    )
+
     artifacts = {}
     # if we're spinning up a launch process from a job
     # we should tell the run to use that artifact
@@ -677,3 +686,21 @@ def _inject_wandb_config_env_vars(
     ]
     config_chunks_dict = {f"WANDB_CONFIG_{i}": chunk for i, chunk in enumerate(chunks)}
     env_dict.update(config_chunks_dict)
+
+
+def _inject_file_overrides_env_vars(
+    overrides: Dict[str, Any], env_dict: Dict[str, Any], maximum_env_length: int
+) -> None:
+    str_overrides = json.dumps(overrides)
+    if len(str_overrides) <= maximum_env_length:
+        env_dict["WANDB_LAUNCH_FILE_OVERRIDES"] = str_overrides
+        return
+
+    chunks = [
+        str_overrides[i : i + maximum_env_length]
+        for i in range(0, len(str_overrides), maximum_env_length)
+    ]
+    overrides_chunks_dict = {
+        f"WANDB_LAUNCH_FILE_OVERRIDES_{i}": chunk for i, chunk in enumerate(chunks)
+    }
+    env_dict.update(overrides_chunks_dict)

wandb/sdk/launch/builder/docker_builder.py

@@ -1,4 +1,5 @@
 """Implementation of the docker builder."""
+
 import logging
 import os
 from typing import Any, Dict, Optional

wandb/sdk/launch/builder/noop.py

@@ -1,4 +1,5 @@
 """NoOp builder implementation."""
+
 from typing import Any, Dict, Optional
 
 from wandb.sdk.launch.builder.abstract import AbstractBuilder

wandb/sdk/launch/create_job.py

@@ -1,6 +1,7 @@
 import json
 import logging
 import os
+import re
 import sys
 import tempfile
 from typing import Any, Dict, List, Optional, Tuple
@@ -19,6 +20,9 @@ logging.basicConfig(stream=sys.stdout, level=logging.INFO)
 _logger = logging.getLogger("wandb")
 
 
+CODE_ARTIFACT_EXCLUDE_PATHS = ["wandb", ".git"]
+
+
 def create_job(
     path: str,
     job_type: str,
@@ -107,6 +111,13 @@ def _create_job(
         )
         return None, "", []
 
+    if runtime is not None:
+        if not re.match(r"^3\.\d+$", runtime):
+            wandb.termerror(
+                f"Runtime (-r, --runtime) must be a minor version of Python 3, "
+                f"e.g. 3.9 or 3.10, received {runtime}"
+            )
+            return None, "", []
     aliases = aliases or []
     tempdir = tempfile.TemporaryDirectory()
     try:
@@ -436,6 +447,13 @@ def _make_code_artifact(
         wandb.termerror(f"Error adding to code artifact: {e}")
         return None
 
+    # Remove paths we don't want to include, if present
+    for item in CODE_ARTIFACT_EXCLUDE_PATHS:
+        try:
+            code_artifact.remove(item)
+        except FileNotFoundError:
+            pass
+
     res, _ = api.create_artifact(
         artifact_type_name="code",
         artifact_collection_name=artifact_name,

wandb/sdk/launch/environment/abstract.py

@@ -1,4 +1,5 @@
 """Abstract base class for environments."""
+
 from abc import ABC, abstractmethod
 
 

wandb/sdk/launch/environment/gcp_environment.py

@@ -1,4 +1,5 @@
 """Implementation of the GCP environment for wandb launch."""
+
 import logging
 import os
 import subprocess

wandb/sdk/launch/environment/local_environment.py

@@ -1,4 +1,5 @@
 """Dummy local environment implementation. This is the default environment."""
+
 from typing import Any, Dict, Union
 
 from wandb.sdk.launch.errors import LaunchError

wandb/sdk/launch/loader.py

@@ -1,4 +1,5 @@
 """Utilities for the agent."""
+
 from typing import Any, Dict, Optional
 
 import wandb

wandb/sdk/launch/registry/abstract.py

@@ -1,4 +1,5 @@
 """Abstract base class for registries."""
+
 from abc import ABC, abstractmethod
 from typing import Tuple
 

wandb/sdk/launch/registry/azure_container_registry.py

@@ -1,4 +1,5 @@
 """Implementation of AzureContainerRegistry class."""
+
 import re
 from typing import TYPE_CHECKING, Optional, Tuple
 

wandb/sdk/launch/registry/elastic_container_registry.py

@@ -1,4 +1,5 @@
 """Implementation of Elastic Container Registry class for wandb launch."""
+
 import base64
 import logging
 from typing import Dict, Optional, Tuple

wandb/sdk/launch/registry/google_artifact_registry.py

@@ -1,4 +1,5 @@
 """Implementation of Google Artifact Registry for wandb launch."""
+
 import logging
 from typing import Optional, Tuple
 

wandb/sdk/launch/registry/local_registry.py

@@ -1,4 +1,5 @@
 """Local registry implementation."""
+
 import logging
 from typing import Tuple
 

wandb/sdk/launch/runner/abstract.py

@@ -3,6 +3,7 @@
 This class defines the interface that the W&B launch runner uses to manage the lifecycle
 of runs launched in different environments (e.g. runs launched locally or in a cluster).
 """
+
 import logging
 import os
 import subprocess

wandb/sdk/launch/runner/kubernetes_monitor.py

@@ -1,4 +1,5 @@
 """Monitors kubernetes resources managed by the launch agent."""
+
 import asyncio
 import logging
 import sys

wandb/sdk/launch/runner/kubernetes_runner.py

@@ -1,4 +1,5 @@
 """Implementation of KubernetesRunner class for wandb launch."""
+
 import asyncio
 import base64
 import datetime
@@ -539,9 +540,9 @@ class KubernetesRunner(AbstractRunner):
                     WANDB_K8S_LABEL_MONITOR,
                     LaunchAgent.name(),
                 )
-                resource_args["metadata"]["labels"][
-                    WANDB_K8S_LABEL_AGENT
-                ] = LaunchAgent.name()
+                resource_args["metadata"]["labels"][WANDB_K8S_LABEL_AGENT] = (
+                    LaunchAgent.name()
+                )
 
             overrides = {}
             if launch_project.override_args:

wandb/sdk/launch/runner/sagemaker_runner.py

@@ -1,4 +1,5 @@
 """Implementation of the SageMakerRunner class."""
+
 import asyncio
 import logging
 from typing import Any, Dict, List, Optional, cast
@@ -324,16 +325,16 @@ def build_sagemaker_args(
     sagemaker_args["TrainingJobName"] = training_job_name
     entry_cmd = entry_point.command if entry_point else []
 
-    sagemaker_args[
-        "AlgorithmSpecification"
-    ] = merge_image_uri_with_algorithm_specification(
-        given_sagemaker_args.get(
-            "AlgorithmSpecification",
-            given_sagemaker_args.get("algorithm_specification"),
-        ),
-        image_uri,
-        entry_cmd,
-        args,
+    sagemaker_args["AlgorithmSpecification"] = (
+        merge_image_uri_with_algorithm_specification(
+            given_sagemaker_args.get(
+                "AlgorithmSpecification",
+                given_sagemaker_args.get("algorithm_specification"),
+            ),
+            image_uri,
+            entry_cmd,
+            args,
+        )
     )
 
     sagemaker_args["RoleArn"] = role_arn

wandb/sdk/launch/sweeps/scheduler.py

@@ -1,4 +1,5 @@
 """Abstract Scheduler class."""
+
 import asyncio
 import base64
 import copy

wandb/sdk/launch/sweeps/scheduler_sweep.py

@@ -1,4 +1,5 @@
 """Scheduler for classic wandb Sweeps."""
+
 import logging
 from pprint import pformat as pf
 from typing import Any, Dict, List, Optional

wandb/sdk/launch/sweeps/utils.py

@@ -296,7 +296,7 @@ def check_job_exists(public_api: "PublicApi", job: Optional[str]) -> bool:
 
 
 def get_previous_args(
-    run_spec: Dict[str, Any]
+    run_spec: Dict[str, Any],
 ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
     """Parse through previous scheduler run_spec.
 

wandb/sdk/launch/utils.py

@@ -57,15 +57,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,
 )
 

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/lazyloader.py

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

wandb/sdk/lib/proto_util.py

@@ -29,7 +29,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/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/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
 
@@ -161,28 +162,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()
 
 

wandb/sdk/wandb_manager.py

@@ -5,7 +5,7 @@ Create a manager channel.
 
 import atexit
 import os
-from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
+from typing import TYPE_CHECKING, Callable, Optional
 
 import psutil
 
@@ -205,7 +205,7 @@ class _Manager:
         svc_iface = self._get_service_interface()
         svc_iface._svc_inform_start(settings=settings, run_id=run_id)
 
-    def _inform_attach(self, attach_id: str) -> Optional[Dict[str, Any]]:
+    def _inform_attach(self, attach_id: str) -> Optional["wandb_settings_pb2.Settings"]:
         svc_iface = self._get_service_interface()
         try:
             response = svc_iface._svc_inform_attach(attach_id=attach_id)

wandb/sdk/wandb_require.py

@@ -9,9 +9,11 @@ Example:
     wandb.require("incremental-artifacts@beta")
 """
 
+import os
 from typing import Optional, Sequence, Union
 
 import wandb
+from wandb.env import _REQUIRE_CORE
 from wandb.errors import UnsupportedError
 from wandb.sdk import wandb_run
 from wandb.sdk.lib.wburls import wburls
@@ -38,6 +40,9 @@ class _Requires:
     def require_service(self) -> None:
         self._require_service()
 
+    def require_core(self) -> None:
+        os.environ[_REQUIRE_CORE] = "true"
+
     def apply(self) -> None:
         """Call require_* method for supported features."""
         last_message: str = ""