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

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
@@ -433,6 +434,8 @@ class SafeWatch:
                     del kwargs["resource_version"]
                     self._last_seen_resource_version = None
             except Exception as E:
+                exc_type = type(E).__name__
+                stack_trace = traceback.format_exc()
                 wandb.termerror(
-                    f"Unknown exception in event stream: {E}, attempting to recover"
+                    f"Unknown exception in event stream of type {exc_type}: {E}, attempting to recover. Stack trace: {stack_trace}"
                 )

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
@@ -157,7 +158,9 @@ class Scheduler(ABC):
         self._runs: Dict[str, SweepRun] = {}
         # Threading lock to ensure thread-safe access to the runs dictionary
         self._threading_lock: threading.Lock = threading.Lock()
-        self._polling_sleep = polling_sleep or DEFAULT_POLLING_SLEEP
+        self._polling_sleep = (
+            polling_sleep if polling_sleep is not None else DEFAULT_POLLING_SLEEP
+        )
         self._project_queue = project_queue
         # Optionally run multiple workers in (pseudo-)parallel. Workers do not
         # actually run training workloads, they simply send heartbeat messages

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,
 )
 
@@ -846,3 +846,21 @@ 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]

wandb/sdk/lib/_settings_toposort_generated.py

@@ -22,6 +22,7 @@ _Setting = Literal[
     "_disable_service",
     "_disable_setproctitle",
     "_disable_stats",
+    "_disable_update_check",
     "_disable_viewer",
     "_disable_machine_info",
     "_except_exit",

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

@@ -195,12 +195,6 @@ class _WandbInit:
         # Start with settings from wandb library singleton
         settings: Settings = self._wl.settings.copy()
 
-        # when using launch, we don't want to reuse the same run id from the singleton
-        # since users might launch multiple runs in the same process
-        # TODO(kdg): allow users to control this via launch settings
-        if settings.launch and singleton is not None:
-            settings.update({"run_id": None}, source=Source.INIT)
-
         settings_param = kwargs.pop("settings", None)
         if settings_param is not None and isinstance(settings_param, (Settings, dict)):
             settings.update(settings_param, source=Source.INIT)
@@ -1124,10 +1118,10 @@ def init(
             for saving hyperparameters to compare across runs. The ID cannot
             contain the following special characters: `/\#?%:`.
             See [our guide to resuming runs](https://docs.wandb.com/guides/runs/resuming).
-        fork_from: (str, optional) A string with the format <run_id>?_step=<step> describing
+        fork_from: (str, optional) A string with the format {run_id}?_step={step} describing
             a moment in a previous run to fork a new run from. Creates a new run that picks up
             logging history from the specified run at the specified moment. The target run must
-            be in the current project.
+            be in the current project. Example: `fork_from="my-run-id?_step=1234"`.
 
     Examples:
     ### Set where the run is logged

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 = ""