wandb 0.16.4__py3-none-any.whl → 0.16.6__py3-none-any.whl

wandb/sdk/launch/create_job.py

@@ -11,7 +11,7 @@ from wandb.sdk.artifacts.artifact import Artifact
 from wandb.sdk.internal.job_builder import JobBuilder
 from wandb.sdk.launch.builder.build import get_current_python_version
 from wandb.sdk.launch.git_reference import GitReference
-from wandb.sdk.launch.utils import _is_git_uri
+from wandb.sdk.launch.utils import _is_git_uri, get_entrypoint_file
 from wandb.sdk.lib import filesystem
 from wandb.util import make_artifact_name_safe
 
@@ -145,6 +145,7 @@ def _create_job(
 
     job_builder = _configure_job_builder_for_partial(tempdir.name, job_source=job_type)
     if job_type == "code":
+        assert entrypoint is not None
         job_name = _make_code_artifact(
             api=api,
             job_builder=job_builder,
@@ -233,7 +234,6 @@ def _make_metadata_for_partial_job(
         return metadata, None
 
     if job_type == "code":
-        path, entrypoint = _handle_artifact_entrypoint(path, entrypoint)
         if not entrypoint:
             wandb.termerror(
                 "Artifact jobs must have an entrypoint, either included in the path or specified with -E"
@@ -304,15 +304,22 @@ def _create_repo_metadata(
             with open(os.path.join(local_dir, ".python-version")) as f:
                 python_version = f.read().strip().splitlines()[0]
         else:
-            major, minor = get_current_python_version()
-            python_version = f"{major}.{minor}"
+            _, python_version = get_current_python_version()
 
     python_version = _clean_python_version(python_version)
 
     # check if entrypoint is valid
     assert entrypoint is not None
-    if not os.path.exists(os.path.join(local_dir, entrypoint)):
-        wandb.termerror(f"Entrypoint {entrypoint} not found in git repo")
+    entrypoint_list = entrypoint.split(" ")
+    entrypoint_file = get_entrypoint_file(entrypoint_list)
+    if not entrypoint_file:
+        wandb.termerror(
+            f"Entrypoint {entrypoint} is invalid. An entrypoint should include both an executable and a file, for example 'python train.py'"
+        )
+        return None
+
+    if not os.path.exists(os.path.join(local_dir, entrypoint_file)):
+        wandb.termerror(f"Entrypoint file {entrypoint_file} not found in git repo")
         return None
 
     metadata = {
@@ -320,9 +327,9 @@ def _create_repo_metadata(
             "commit": commit,
             "remote": ref.url,
         },
-        "codePathLocal": entrypoint,  # not in git context, optionally also set local
-        "codePath": entrypoint,
-        "entrypoint": [f"python{python_version}", entrypoint],
+        "codePathLocal": entrypoint_file,  # not in git context, optionally also set local
+        "codePath": entrypoint_file,
+        "entrypoint": entrypoint_list,
         "python": python_version,  # used to build container
         "notebook": False,  # partial jobs from notebooks not supported
     }
@@ -332,10 +339,17 @@ def _create_repo_metadata(
 
 def _create_artifact_metadata(
     path: str, entrypoint: str, runtime: Optional[str] = None
-) -> Tuple[Dict[str, Any], List[str]]:
+) -> Tuple[Optional[Dict[str, Any]], Optional[List[str]]]:
     if not os.path.isdir(path):
         wandb.termerror("Path must be a valid file or directory")
         return {}, []
+    entrypoint_list = entrypoint.split(" ")
+    entrypoint_file = get_entrypoint_file(entrypoint_list)
+    if not entrypoint_file:
+        wandb.termerror(
+            f"Entrypoint {entrypoint} is invalid. An entrypoint should include both an executable and a file, for example 'python train.py'"
+        )
+        return None, None
 
     # read local requirements.txt and dump to temp dir for builder
     requirements = []
@@ -347,41 +361,17 @@ def _create_artifact_metadata(
     if runtime:
         python_version = _clean_python_version(runtime)
     else:
-        python_version = ".".join(get_current_python_version())
+        python_version, _ = get_current_python_version()
+        python_version = _clean_python_version(python_version)
 
-    metadata = {"python": python_version, "codePath": entrypoint}
+    metadata = {
+        "python": python_version,
+        "codePath": entrypoint_file,
+        "entrypoint": entrypoint_list,
+    }
     return metadata, requirements
 
 
-def _handle_artifact_entrypoint(
-    path: str, entrypoint: Optional[str] = None
-) -> Tuple[str, Optional[str]]:
-    if os.path.isfile(path):
-        if entrypoint and path.endswith(entrypoint):
-            path = path.replace(entrypoint, "")
-            wandb.termwarn(
-                f"Both entrypoint provided and path contains file. Using provided entrypoint: {entrypoint}, path is now: {path}"
-            )
-        elif entrypoint:
-            wandb.termwarn(
-                f"Ignoring passed in entrypoint as it does not match file path found in 'path'. Path entrypoint: {path.split('/')[-1]}"
-            )
-        entrypoint = path.split("/")[-1]
-        path = "/".join(path.split("/")[:-1])
-    elif not entrypoint:
-        wandb.termerror("Entrypoint not valid")
-        return "", None
-    path = path or "."  # when path is just an entrypoint, use cdw
-
-    if not os.path.exists(os.path.join(path, entrypoint)):
-        wandb.termerror(
-            f"Could not find execution point: {os.path.join(path, entrypoint)}"
-        )
-        return "", None
-
-    return path, entrypoint
-
-
 def _configure_job_builder_for_partial(tmpdir: str, job_source: str) -> JobBuilder:
     """Configure job builder with temp dir and job source."""
     # adjust git source to repo
@@ -396,6 +386,7 @@ def _configure_job_builder_for_partial(tmpdir: str, job_source: str) -> JobBuild
     settings.update({"files_dir": tmpdir, "job_source": job_source})
     job_builder = JobBuilder(
         settings=settings,  # type: ignore
+        verbose=True,
     )
     # never allow notebook runs
     job_builder._is_notebook_run = False
@@ -410,7 +401,7 @@ def _make_code_artifact(
     job_builder: JobBuilder,
     run: "wandb.sdk.wandb_run.Run",
     path: str,
-    entrypoint: Optional[str],
+    entrypoint: str,
     entity: Optional[str],
     project: Optional[str],
     name: Optional[str],
@@ -419,17 +410,22 @@ def _make_code_artifact(
 
     Returns the name of the eventual job.
     """
-    artifact_name = _make_code_artifact_name(os.path.join(path, entrypoint or ""), name)
+    assert entrypoint is not None
+    entrypoint_list = entrypoint.split(" ")
+    entrypoint_file = get_entrypoint_file(entrypoint_list)
+    if not entrypoint_file:
+        wandb.termerror(
+            f"Entrypoint {entrypoint} is invalid. An entrypoint should include both an executable and a file, for example 'python train.py'"
+        )
+        return None
+
+    artifact_name = _make_code_artifact_name(os.path.join(path, entrypoint_file), name)
     code_artifact = wandb.Artifact(
         name=artifact_name,
         type="code",
         description="Code artifact for job",
     )
 
-    # Update path and entrypoint vars to match metadata
-    # TODO(gst): consolidate into one place
-    path, entrypoint = _handle_artifact_entrypoint(path, entrypoint)
-
     try:
         code_artifact.add_dir(path)
     except Exception as e:
@@ -450,7 +446,7 @@ def _make_code_artifact(
         project_name=project,
         run_name=run.id,  # run will be deleted after creation
         description="Code artifact for job",
-        metadata={"codePath": path, "entrypoint": entrypoint},
+        metadata={"codePath": path, "entrypoint": entrypoint_file},
         is_user_created=True,
         aliases=[
             {"artifactCollectionName": artifact_name, "alias": a} for a in ["latest"]

wandb/sdk/launch/runner/kubernetes_monitor.py

@@ -433,6 +433,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,6 +1,7 @@
 """Implementation of KubernetesRunner class for wandb launch."""
 import asyncio
 import base64
+import datetime
 import json
 import logging
 import os
@@ -23,6 +24,7 @@ from wandb.sdk.launch.runner.kubernetes_monitor import (
     CustomResource,
     LaunchKubernetesMonitor,
 )
+from wandb.sdk.lib.retry import ExponentialBackoff, retry_async
 from wandb.util import get_module
 
 from .._project_spec import EntryPoint, LaunchProject
@@ -59,6 +61,7 @@ from kubernetes_asyncio.client.models.v1_secret import (  # type: ignore # noqa:
 from kubernetes_asyncio.client.rest import ApiException  # type: ignore # noqa: E402
 
 TIMEOUT = 5
+API_KEY_SECRET_MAX_RETRIES = 5
 
 _logger = logging.getLogger(__name__)
 
@@ -421,8 +424,23 @@ class KubernetesRunner(AbstractRunner):
                     else:
                         secret_name += f"-{launch_project.run_id}"
 
-                    api_key_secret = await ensure_api_key_secret(
-                        core_api, secret_name, namespace, value
+                    def handle_exception(e):
+                        wandb.termwarn(
+                            f"Exception when ensuring Kubernetes API key secret: {e}. Retrying..."
+                        )
+
+                    api_key_secret = await retry_async(
+                        backoff=ExponentialBackoff(
+                            initial_sleep=datetime.timedelta(seconds=1),
+                            max_sleep=datetime.timedelta(minutes=1),
+                            max_retries=API_KEY_SECRET_MAX_RETRIES,
+                        ),
+                        fn=ensure_api_key_secret,
+                        on_exc=handle_exception,
+                        core_api=core_api,
+                        secret_name=secret_name,
+                        namespace=namespace,
+                        api_key=value,
                     )
                     env.append(
                         {

wandb/sdk/launch/sweeps/scheduler.py

@@ -157,7 +157,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/utils.py

@@ -222,14 +222,14 @@ def get_default_entity(api: Api, launch_config: Optional[Dict[str, Any]]):
 
 
 def strip_resource_args_and_template_vars(launch_spec: Dict[str, Any]) -> None:
-    wandb.termwarn(
-        "Launch spec contains both resource_args and template_variables, "
-        "only one can be set. Using template_variables."
-    )
     if launch_spec.get("resource_args", None) and launch_spec.get(
         "template_variables", None
     ):
-        launch_spec["resource_args"] = None
+        wandb.termwarn(
+            "Launch spec contains both resource_args and template_variables, "
+            "only one can be set. Using template_variables."
+        )
+        launch_spec.pop("resource_args")
 
 
 def construct_launch_spec(
@@ -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/__init__.py

@@ -1,8 +1,5 @@
 from . import lazyloader
 from .disabled import RunDisabled, SummaryDisabled
+from .run_moment import RunMoment
 
-__all__ = (
-    "lazyloader",
-    "RunDisabled",
-    "SummaryDisabled",
-)
+__all__ = ("lazyloader", "RunDisabled", "SummaryDisabled", "RunMoment")

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",
@@ -102,6 +103,7 @@ _Setting = Literal[
     "entity",
     "files_dir",
     "force",
+    "fork_from",
     "git_commit",
     "git_remote",
     "git_remote_url",

wandb/sdk/lib/filesystem.py

@@ -227,7 +227,17 @@ def safe_open(
 
 
 def safe_copy(source_path: StrPath, target_path: StrPath) -> StrPath:
-    """Copy a file, ensuring any changes only apply atomically once finished."""
+    """Copy a file atomically.
+
+    Copying is not usually atomic, and on operating systems that allow multiple
+    writers to the same file, the result can get corrupted. If two writers copy
+    to the same file, the contents can become interleaved.
+
+    We mitigate the issue somewhat by copying to a temporary file first and
+    then renaming. Renaming is atomic: if process 1 renames file A to X and
+    process 2 renames file B to X, then X will either contain the contents
+    of A or the contents of B, not some mixture of both.
+    """
     # TODO (hugh): check that there is enough free space.
     output_path = Path(target_path).resolve()
     output_path.parent.mkdir(parents=True, exist_ok=True)

wandb/sdk/lib/run_moment.py

@@ -0,0 +1,78 @@
+import sys
+from dataclasses import dataclass
+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"]
+
+
+@dataclass
+class RunMoment:
+    """A moment in a run."""
+
+    run: str  # run name
+
+    # currently, the _step value to fork from. in future, this will be optional
+    value: Union[int, float]
+
+    # only step for now, in future this will be relaxed to be any metric
+    metric: _STEP = "_step"
+
+    def __post_init__(self):
+        if self.metric != "_step":
+            raise ValueError(
+                f"Only the metric '_step' is supported, got '{self.metric}'."
+            )
+        if not isinstance(self.value, (int, float)):
+            raise ValueError(
+                f"Only int or float values are supported, got '{self.value}'."
+            )
+        if not isinstance(self.run, str):
+            raise ValueError(f"Only string run names are supported, got '{self.run}'.")
+
+    @classmethod
+    def from_uri(cls, uri: str) -> "RunMoment":
+        parsable = "runmoment://" + uri
+        parse_err = ValueError(
+            f"Could not parse passed run moment string '{uri}', "
+            f"expected format '<run>?<metric>=<numeric_value>'. "
+            f"Currently, only the metric '_step' is supported. "
+            f"Example: 'ans3bsax?_step=123'."
+        )
+
+        try:
+            parsed = parse.urlparse(parsable)
+        except ValueError as e:
+            raise parse_err from e
+
+        if parsed.scheme != "runmoment":
+            raise parse_err
+
+        # extract run, metric, value from parsed
+        if not parsed.netloc:
+            raise parse_err
+
+        run = parsed.netloc
+
+        if parsed.path or parsed.params or parsed.fragment:
+            raise parse_err
+
+        query = parse.parse_qs(parsed.query)
+        if len(query) != 1:
+            raise parse_err
+
+        metric = list(query.keys())[0]
+        if metric != "_step":
+            raise parse_err
+        value: str = query[metric][0]
+        try:
+            num_value = int(value) if value.isdigit() else float(value)
+        except ValueError as e:
+            raise parse_err from e
+
+        return cls(run=run, metric=cast(_STEP, metric), value=num_value)

wandb/sdk/service/streams.py

@@ -5,6 +5,7 @@ StreamRecord: All the external state for the internal thread (queues, etc)
 StreamAction: Lightweight record for stream ops for thread safety
 StreamMux: Container for dictionary of stream threads per runid
 """
+
 import functools
 import multiprocessing
 import queue
@@ -327,7 +328,6 @@ class StreamMux:
             result = internal_messages_handle.wait(timeout=-1)
             assert result
             internal_messages_response = result.response.internal_messages_response
-            job_info_handle = stream.interface.deliver_request_job_info()
 
             # wait for them, it's ok to do this serially but this can be improved
             result = poll_exit_handle.wait(timeout=-1)
@@ -346,17 +346,12 @@ class StreamMux:
             assert result
             final_summary = result.response.get_summary_response
 
-            result = job_info_handle.wait(timeout=-1)
-            assert result
-            job_info = result.response.job_info_response
-
             Run._footer(
                 sampled_history=sampled_history,
                 final_summary=final_summary,
                 poll_exit_response=poll_exit_response,
                 server_info_response=server_info_response,
                 internal_messages_response=internal_messages_response,
-                job_info=job_info,
                 settings=stream._settings,  # type: ignore
                 printer=printer,
             )

wandb/sdk/wandb_init.py

@@ -7,6 +7,7 @@ your evaluation script, and each step would be tracked as a run in W&B.
 For more on using `wandb.init()`, including code snippets, check out our
 [guide and FAQs](https://docs.wandb.ai/guides/track/launch).
 """
+
 import copy
 import json
 import logging
@@ -194,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)
@@ -828,7 +823,7 @@ class _WandbInit:
             and self.settings.launch_config_path
             and os.path.exists(self.settings.launch_config_path)
         ):
-            run._save(self.settings.launch_config_path)
+            run.save(self.settings.launch_config_path)
         # put artifacts in run config here
         # since doing so earlier will cause an error
         # as the run is not upserted
@@ -961,6 +956,7 @@ def init(
     monitor_gym: Optional[bool] = None,
     save_code: Optional[bool] = None,
     id: Optional[str] = None,
+    fork_from: Optional[str] = None,
     settings: Union[Settings, Dict[str, Any], None] = None,
 ) -> Union[Run, RunDisabled, None]:
     r"""Start a new run to track and log to W&B.
@@ -1122,6 +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
+            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. Example: `fork_from="my-run-id?_step=1234"`.
 
     Examples:
     ### Set where the run is logged
@@ -1167,6 +1167,11 @@ def init(
     error_seen = None
     except_exit = None
     run: Optional[Union[Run, RunDisabled]] = None
+
+    # convert fork_from into a version that can be passed to settings
+    if fork_from is not None and resume is not None:
+        raise ValueError("Cannot specify both `fork_from` and `resume`")
+
     try:
         wi = _WandbInit()
         wi.setup(kwargs)

wandb/sdk/wandb_login.py

@@ -22,7 +22,7 @@ from wandb.old.settings import Settings as OldSettings
 from ..apis import InternalApi
 from .internal.internal_api import Api
 from .lib import apikey
-from .wandb_settings import Settings, Source
+from .wandb_settings import Settings
 
 
 def _handle_host_wandb_setting(host: Optional[str], cloud: bool = False) -> None:
@@ -80,11 +80,17 @@ def login(
     _handle_host_wandb_setting(host)
     if wandb.setup()._settings._noop:
         return True
-    kwargs = dict(locals())
-    _verify = kwargs.pop("verify", False)
-    configured = _login(**kwargs)
 
-    if _verify:
+    configured = _login(
+        anonymous=anonymous,
+        key=key,
+        relogin=relogin,
+        host=host,
+        force=force,
+        timeout=timeout,
+    )
+
+    if verify:
         from . import wandb_setup
 
         singleton = wandb_setup._WandbSetup._instance
@@ -115,22 +121,32 @@ class _WandbLogin:
         self._key = None
         self._relogin = None
 
-    def setup(self, kwargs):
-        self.kwargs = kwargs
+    def setup(
+        self,
+        *,
+        anonymous: Optional[Literal["must", "allow", "never"]] = None,
+        key: Optional[str] = None,
+        relogin: Optional[bool] = None,
+        host: Optional[str] = None,
+        force: Optional[bool] = None,
+        timeout: Optional[int] = None,
+    ):
+        self._relogin = relogin
 
         # built up login settings
         login_settings: Settings = wandb.Settings()
-        settings_param = kwargs.pop("_settings", None)
-        # note that this case does not come up anywhere except for the tests
-        if settings_param is not None:
-            if isinstance(settings_param, Settings):
-                login_settings._apply_settings(settings_param)
-            elif isinstance(settings_param, dict):
-                login_settings.update(settings_param, source=Source.LOGIN)
-        _logger = wandb.setup()._get_logger()
-        # Do not save relogin into settings as we just want to relogin once
-        self._relogin = kwargs.pop("relogin", None)
-        login_settings._apply_login(kwargs, _logger=_logger)
+        logger = wandb.setup()._get_logger()
+
+        login_settings._apply_login(
+            {
+                "anonymous": anonymous,
+                "key": key,
+                "host": host,
+                "force": force,
+                "timeout": timeout,
+            },
+            _logger=logger,
+        )
 
         # make sure they are applied globally
         self._wl = wandb.setup(settings=login_settings)
@@ -259,6 +275,7 @@ class _WandbLogin:
 
 
 def _login(
+    *,
     anonymous: Optional[Literal["must", "allow", "never"]] = None,
     key: Optional[str] = None,
     relogin: Optional[bool] = None,
@@ -270,9 +287,6 @@ def _login(
     _disable_warning: Optional[bool] = None,
     _entity: Optional[str] = None,
 ):
-    kwargs = dict(locals())
-    _disable_warning = kwargs.pop("_disable_warning", None)
-
     if wandb.run is not None:
         if not _disable_warning:
             wandb.termwarn("Calling wandb.login() after wandb.init() has no effect.")
@@ -280,20 +294,24 @@ def _login(
 
     wlogin = _WandbLogin()
 
-    _backend = kwargs.pop("_backend", None)
     if _backend:
         wlogin.set_backend(_backend)
 
-    _silent = kwargs.pop("_silent", None)
     if _silent:
         wlogin.set_silent(_silent)
 
-    _entity = kwargs.pop("_entity", None)
     if _entity:
         wlogin.set_entity(_entity)
 
     # configure login object
-    wlogin.setup(kwargs)
+    wlogin.setup(
+        anonymous=anonymous,
+        key=key,
+        relogin=relogin,
+        host=host,
+        force=force,
+        timeout=timeout,
+    )
 
     if wlogin._settings._offline:
         return False
@@ -306,7 +324,6 @@ def _login(
     # perform a login
     logged_in = wlogin.login()
 
-    key = kwargs.get("key")
     if key:
         wlogin.configure_api_key(key)