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

wandb/sdk/launch/inputs/manage.py

@@ -0,0 +1,95 @@
+"""Functions for declaring overridable configuration for launch jobs."""
+
+from typing import List, Optional
+
+
+def manage_config_file(
+    path: str,
+    include: Optional[List[str]] = None,
+    exclude: Optional[List[str]] = None,
+):
+    r"""Declare an overridable configuration file for a launch job.
+
+    If a new job version is created from the active run, the configuration file
+    will be added to the job's inputs. If the job is launched and overrides
+    have been provided for the configuration file, this function will detect
+    the overrides from the environment and update the configuration file on disk.
+    Note that these overrides will only be applied in ephemeral containers.
+    `include` and `exclude` are lists of dot separated paths with the config.
+    The paths are used to filter subtrees of the configuration file out of the
+    job's inputs.
+
+    For example, given the following configuration file:
+        ```yaml
+        model:
+            name: resnet
+            layers: 18
+        training:
+            epochs: 10
+            batch_size: 32
+        ```
+
+    Passing `include=['model']` will only include the `model` subtree in the
+    job's inputs. Passing `exclude=['model.layers']` will exclude the `layers`
+    key from the `model` subtree. Note that `exclude` takes precedence over
+    `include`.
+
+    `.` is used as a separator for nested keys. If a key contains a `.`, it
+    should be escaped with a backslash, e.g. `include=[r'model\.layers']`. Note
+    the use of `r` to denote a raw string when using escape chars.
+
+    Args:
+        path (str): The path to the configuration file. This path must be
+            relative and must not contain backwards traversal, i.e. `..`.
+        include (List[str]): A list of keys to include in the configuration file.
+        exclude (List[str]): A list of keys to exclude from the configuration file.
+
+    Raises:
+        LaunchError: If the path is not valid, or if there is no active run.
+    """
+    from .internal import handle_config_file_input
+
+    return handle_config_file_input(path, include, exclude)
+
+
+def manage_wandb_config(
+    include: Optional[List[str]] = None,
+    exclude: Optional[List[str]] = None,
+):
+    r"""Declare wandb.config as an overridable configuration for a launch job.
+
+    If a new job version is created from the active run, the run config
+    (wandb.config) will become an overridable input of the job. If the job is
+    launched and overrides have been provided for the run config, the overrides
+    will be applied to the run config when `wandb.init` is called.
+    `include` and `exclude` are lists of dot separated paths with the config.
+    The paths are used to filter subtrees of the configuration file out of the
+    job's inputs.
+
+    For example, given the following run config contents:
+        ```yaml
+        model:
+            name: resnet
+            layers: 18
+        training:
+            epochs: 10
+            batch_size: 32
+        ```
+    Passing `include=['model']` will only include the `model` subtree in the
+    job's inputs. Passing `exclude=['model.layers']` will exclude the `layers`
+    key from the `model` subtree. Note that `exclude` takes precedence over
+    `include`.
+    `.` is used as a separator for nested keys. If a key contains a `.`, it
+    should be escaped with a backslash, e.g. `include=[r'model\.layers']`. Note
+    the use of `r` to denote a raw string when using escape chars.
+
+    Args:
+        include (List[str]): A list of subtrees to include in the configuration.
+        exclude (List[str]): A list of subtrees to exclude from the configuration.
+
+    Raises:
+        LaunchError: If there is no active run.
+    """
+    from .internal import handle_run_config_input
+
+    handle_run_config_input(include, exclude)

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
 
@@ -210,7 +211,7 @@ class GoogleArtifactRegistry(AbstractRegistry):
             for image in await list_images(request={"parent": parent}):
                 if tag in image.tags:
                     return True
-        except google.api_core.exceptions.NotFound as e:
+        except google.api_core.exceptions.NotFound as e:  # type: ignore[attr-defined]
             raise LaunchError(
                 f"The Google Artifact Registry repository {self.repository} "
                 f"does not exist. Please create it or modify your registry configuration."

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
@@ -28,7 +29,6 @@ from wandb.sdk.lib.retry import ExponentialBackoff, retry_async
 from wandb.util import get_module
 
 from .._project_spec import EntryPoint, LaunchProject
-from ..builder.build import get_env_vars_dict
 from ..errors import LaunchError
 from ..utils import (
     LOG_PREFIX,
@@ -373,8 +373,7 @@ class KubernetesRunner(AbstractRunner):
                 }
 
         entry_point = (
-            launch_project.override_entrypoint
-            or launch_project.get_single_entry_point()
+            launch_project.override_entrypoint or launch_project.get_job_entry_point()
         )
         if launch_project.docker_image:
             # dont specify run id if user provided image, could have multiple runs
@@ -400,8 +399,8 @@ class KubernetesRunner(AbstractRunner):
             launch_project.override_entrypoint is not None,
         )
 
-        env_vars = get_env_vars_dict(
-            launch_project, self._api, MAX_ENV_LENGTHS[self.__class__.__name__]
+        env_vars = launch_project.get_env_vars_dict(
+            self._api, MAX_ENV_LENGTHS[self.__class__.__name__]
         )
         api_key_secret = None
         for cont in containers:
@@ -510,8 +509,8 @@ class KubernetesRunner(AbstractRunner):
         api_version = resource_args.get("apiVersion", "batch/v1")
 
         if api_version not in ["batch/v1", "batch/v1beta1"]:
-            env_vars = get_env_vars_dict(
-                launch_project, self._api, MAX_ENV_LENGTHS[self.__class__.__name__]
+            env_vars = launch_project.get_env_vars_dict(
+                self._api, MAX_ENV_LENGTHS[self.__class__.__name__]
             )
             # Crawl the resource args and add our env vars to the containers.
             add_wandb_env(resource_args, env_vars)
@@ -539,9 +538,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/local_container.py

@@ -12,7 +12,6 @@ from wandb.sdk.launch.environment.abstract import AbstractEnvironment
 from wandb.sdk.launch.registry.abstract import AbstractRegistry
 
 from .._project_spec import LaunchProject
-from ..builder.build import get_env_vars_dict
 from ..errors import LaunchError
 from ..utils import (
     LOG_PREFIX,
@@ -133,8 +132,8 @@ class LocalContainerRunner(AbstractRunner):
         docker_args = self._populate_docker_args(launch_project, image_uri)
         synchronous: bool = self.backend_config[PROJECT_SYNCHRONOUS]
 
-        env_vars = get_env_vars_dict(
-            launch_project, self._api, MAX_ENV_LENGTHS[self.__class__.__name__]
+        env_vars = launch_project.get_env_vars_dict(
+            self._api, MAX_ENV_LENGTHS[self.__class__.__name__]
         )
 
         # When running against local port, need to swap to local docker host

wandb/sdk/launch/runner/local_process.py

@@ -4,16 +4,12 @@ from typing import Any, List, Optional
 
 import wandb
 
-from .._project_spec import LaunchProject, get_entry_point_command
-from ..builder.build import get_env_vars_dict
+from .._project_spec import LaunchProject
 from ..errors import LaunchError
 from ..utils import (
     LOG_PREFIX,
     MAX_ENV_LENGTHS,
     PROJECT_SYNCHRONOUS,
-    _is_wandb_uri,
-    download_wandb_python_deps,
-    parse_wandb_uri,
     sanitize_wandb_api_key,
     validate_wandb_python_deps,
 )
@@ -47,8 +43,7 @@ class LocalProcessRunner(AbstractRunner):
 
         synchronous: bool = self.backend_config[PROJECT_SYNCHRONOUS]
         entry_point = (
-            launch_project.override_entrypoint
-            or launch_project.get_single_entry_point()
+            launch_project.override_entrypoint or launch_project.get_job_entry_point()
         )
 
         cmd: List[Any] = []
@@ -56,23 +51,7 @@ class LocalProcessRunner(AbstractRunner):
         if launch_project.project_dir is None:
             raise LaunchError("Launch LocalProcessRunner received empty project dir")
 
-        # Check to make sure local python dependencies match run's requirement.txt
-        if launch_project.uri and _is_wandb_uri(launch_project.uri):
-            source_entity, source_project, run_name = parse_wandb_uri(
-                launch_project.uri
-            )
-            run_requirements_file = download_wandb_python_deps(
-                source_entity,
-                source_project,
-                run_name,
-                self._api,
-                launch_project.project_dir,
-            )
-            validate_wandb_python_deps(
-                run_requirements_file,
-                launch_project.project_dir,
-            )
-        elif launch_project.job:
+        if launch_project.job:
             assert launch_project._job_artifact is not None
             try:
                 validate_wandb_python_deps(
@@ -81,14 +60,14 @@ class LocalProcessRunner(AbstractRunner):
                 )
             except Exception:
                 wandb.termwarn("Unable to validate python dependencies")
-        env_vars = get_env_vars_dict(
-            launch_project, self._api, MAX_ENV_LENGTHS[self.__class__.__name__]
+        env_vars = launch_project.get_env_vars_dict(
+            self._api, MAX_ENV_LENGTHS[self.__class__.__name__]
         )
         for env_key, env_value in env_vars.items():
             cmd += [f"{shlex.quote(env_key)}={shlex.quote(env_value)}"]
-
-        entry_cmd = get_entry_point_command(entry_point, launch_project.override_args)
-        cmd += entry_cmd
+        if entry_point is not None:
+            cmd += entry_point.command
+        cmd += launch_project.override_args
 
         command_str = " ".join(cmd).strip()
         _msg = f"{LOG_PREFIX}Launching run as a local-process with command {sanitize_wandb_api_key(command_str)}"

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
@@ -11,8 +12,7 @@ from wandb.apis.internal import Api
 from wandb.sdk.launch.environment.aws_environment import AwsEnvironment
 from wandb.sdk.launch.errors import LaunchError
 
-from .._project_spec import EntryPoint, LaunchProject, get_entry_point_command
-from ..builder.build import get_env_vars_dict
+from .._project_spec import EntryPoint, LaunchProject
 from ..registry.abstract import AbstractRegistry
 from ..utils import (
     LOG_PREFIX,
@@ -67,6 +67,7 @@ class SagemakerSubmittedRun(AbstractRun):
                 logGroupName="/aws/sagemaker/TrainingJobs",
                 logStreamName=log_name,
             )
+            assert "events" in res
             return "\n".join(
                 [f'{event["timestamp"]}:{event["message"]}' for event in res["events"]]
             )
@@ -220,12 +221,12 @@ class SageMakerRunner(AbstractRunner):
         launch_project.fill_macros(image_uri)
         _logger.info("Connecting to sagemaker client")
         entry_point = (
-            launch_project.override_entrypoint
-            or launch_project.get_single_entry_point()
-        )
-        command_args = get_entry_point_command(
-            entry_point, launch_project.override_args
+            launch_project.override_entrypoint or launch_project.get_job_entry_point()
         )
+        command_args = []
+        if entry_point is not None:
+            command_args += entry_point.command
+        command_args += launch_project.override_args
         if command_args:
             command_str = " ".join(command_args)
             wandb.termlog(
@@ -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
@@ -348,18 +349,18 @@ def build_sagemaker_args(
 
     if sagemaker_args.get("ResourceConfig") is None:
         raise LaunchError(
-            "Sagemaker launcher requires a ResourceConfig Sagemaker resource argument"
+            "Sagemaker launcher requires a ResourceConfig resource argument"
         )
 
     if sagemaker_args.get("StoppingCondition") is None:
         raise LaunchError(
-            "Sagemaker launcher requires a StoppingCondition Sagemaker resource argument"
+            "Sagemaker launcher requires a StoppingCondition resource argument"
         )
 
     given_env = given_sagemaker_args.get(
         "Environment", sagemaker_args.get("environment", {})
     )
-    calced_env = get_env_vars_dict(launch_project, api, max_env_length)
+    calced_env = launch_project.get_env_vars_dict(api, max_env_length)
     total_env = {**calced_env, **given_env}
     sagemaker_args["Environment"] = total_env
 

wandb/sdk/launch/runner/vertex_runner.py

@@ -8,8 +8,7 @@ if False:
 from wandb.apis.internal import Api
 from wandb.util import get_module
 
-from .._project_spec import LaunchProject, get_entry_point_command
-from ..builder.build import get_env_vars_dict
+from .._project_spec import LaunchProject
 from ..environment.gcp_environment import GcpEnvironment
 from ..errors import LaunchError
 from ..registry.abstract import AbstractRegistry
@@ -113,14 +112,16 @@ class VertexRunner(AbstractRunner):
         synchronous: bool = self.backend_config[PROJECT_SYNCHRONOUS]
 
         entry_point = (
-            launch_project.override_entrypoint
-            or launch_project.get_single_entry_point()
+            launch_project.override_entrypoint or launch_project.get_job_entry_point()
         )
 
         # TODO: Set entrypoint in each container
-        entry_cmd = get_entry_point_command(entry_point, launch_project.override_args)
-        env_vars = get_env_vars_dict(
-            launch_project=launch_project,
+        entry_cmd = []
+        if entry_point is not None:
+            entry_cmd += entry_point.command
+        entry_cmd += launch_project.override_args
+
+        env_vars = launch_project.get_env_vars_dict(
             api=self._api,
             max_env_length=MAX_ENV_LENGTHS[self.__class__.__name__],
         )

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
@@ -405,7 +408,7 @@ class Scheduler(ABC):
         return count
 
     def _try_load_executable(self) -> bool:
-        """Check existance of valid executable for a run.
+        """Check existence of valid executable for a run.
 
         logs and returns False when job is unreachable
         """
@@ -420,7 +423,7 @@ class Scheduler(ABC):
                 return False
             return True
         elif self._kwargs.get("image_uri"):
-            # TODO(gst): check docker existance? Use registry in launch config?
+            # TODO(gst): check docker existence? Use registry in launch config?
             return True
         else:
             return False
@@ -608,7 +611,7 @@ class Scheduler(ABC):
                     f"Failed to get runstate for run ({run_id}). Error: {traceback.format_exc()}"
                 )
                 run_state = RunState.FAILED
-            else:  # first time we get unknwon state
+            else:  # first time we get unknown state
                 run_state = RunState.UNKNOWN
         except (AttributeError, ValueError):
             wandb.termwarn(

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
@@ -58,7 +59,7 @@ class SweepScheduler(Scheduler):
         return None
 
     def _get_sweep_commands(self, worker_id: int) -> List[Dict[str, Any]]:
-        """Helper to recieve sweep command from backend."""
+        """Helper to receive sweep command from backend."""
         # AgentHeartbeat wants a Dict of runs which are running or queued
         _run_states: Dict[str, bool] = {}
         for run_id, run in self._yield_runs():

wandb/sdk/launch/sweeps/utils.py

@@ -217,7 +217,7 @@ def create_sweep_command_args(command: Dict) -> Dict[str, Any]:
     flags: List[str] = []
     # (2) flags without hyphens (e.g. foo=bar)
     flags_no_hyphens: List[str] = []
-    # (3) flags with false booleans ommited  (e.g. --foo)
+    # (3) flags with false booleans omitted  (e.g. --foo)
     flags_no_booleans: List[str] = []
     # (4) flags as a dictionary (used for constructing a json)
     flags_dict: Dict[str, Any] = {}
@@ -257,7 +257,7 @@ def make_launch_sweep_entrypoint(
     """Use args dict from create_sweep_command_args to construct entrypoint.
 
     If replace is True, remove macros from entrypoint, fill them in with args
-    and then return the args in seperate return value.
+    and then return the args in separate return value.
     """
     if not command:
         return None, None
@@ -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.