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

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/inputs/files.py

@@ -0,0 +1,148 @@
+import json
+import os
+from typing import Any, Dict
+
+import yaml
+
+from ..errors import LaunchError
+
+FILE_OVERRIDE_ENV_VAR = "WANDB_LAUNCH_FILE_OVERRIDES"
+
+
+class FileOverrides:
+    """Singleton that read file overrides json from environment variables."""
+
+    _instance = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = object.__new__(cls)
+            cls._instance.overrides = {}
+            cls._instance.load()
+        return cls._instance
+
+    def load(self) -> None:
+        """Load overrides from an environment variable."""
+        overrides = os.environ.get(FILE_OVERRIDE_ENV_VAR)
+        if overrides is None:
+            if f"{FILE_OVERRIDE_ENV_VAR}_0" in os.environ:
+                overrides = ""
+                idx = 0
+                while f"{FILE_OVERRIDE_ENV_VAR}_{idx}" in os.environ:
+                    overrides += os.environ[f"{FILE_OVERRIDE_ENV_VAR}_{idx}"]
+                    idx += 1
+        if overrides:
+            try:
+                contents = json.loads(overrides)
+                if not isinstance(contents, dict):
+                    raise LaunchError(f"Invalid JSON in {FILE_OVERRIDE_ENV_VAR}")
+                self.overrides = contents
+            except json.JSONDecodeError:
+                raise LaunchError(f"Invalid JSON in {FILE_OVERRIDE_ENV_VAR}")
+
+
+def config_path_is_valid(path: str) -> None:
+    """Validate a config file path.
+
+    This function checks if a given config file path is valid. A valid path
+    should meet the following criteria:
+
+    - The path must be expressed as a relative path without any upwards path
+      traversal, e.g. `../config.json`.
+    - The file specified by the path must exist.
+    - The file must have a supported extension (`.json`, `.yaml`, or `.yml`).
+
+    Args:
+        path (str): The path to validate.
+
+    Raises:
+        LaunchError: If the path is not valid.
+    """
+    if os.path.isabs(path):
+        raise LaunchError(
+            f"Invalid config path: {path}. Please provide a relative path."
+        )
+    if ".." in path:
+        raise LaunchError(
+            f"Invalid config path: {path}. Please provide a relative path "
+            "without any upward path traversal, e.g. `../config.json`."
+        )
+    path = os.path.normpath(path)
+    if not os.path.exists(path):
+        raise LaunchError(f"Invalid config path: {path}. File does not exist.")
+    if not any(path.endswith(ext) for ext in [".json", ".yaml", ".yml"]):
+        raise LaunchError(
+            f"Invalid config path: {path}. Only JSON and YAML files are supported."
+        )
+
+
+def override_file(path: str) -> None:
+    """Check for file overrides in the environment and apply them if found."""
+    file_overrides = FileOverrides()
+    if path in file_overrides.overrides:
+        overrides = file_overrides.overrides.get(path)
+        if overrides is not None:
+            config = _read_config_file(path)
+            _update_dict(config, overrides)
+            _write_config_file(path, config)
+
+
+def _write_config_file(path: str, config: Any) -> None:
+    """Write a config file to disk.
+
+    Args:
+        path (str): The path to the config file.
+        config (Any): The contents of the config file as a Python object.
+
+    Raises:
+        LaunchError: If the file extension is not supported.
+    """
+    _, ext = os.path.splitext(path)
+    if ext == ".json":
+        with open(path, "w") as f:
+            json.dump(config, f, indent=2)
+    elif ext in [".yaml", ".yml"]:
+        with open(path, "w") as f:
+            yaml.safe_dump(config, f)
+    else:
+        raise LaunchError(f"Unsupported file extension: {ext}")
+
+
+def _read_config_file(path: str) -> Any:
+    """Read a config file from disk.
+
+    Args:
+        path (str): The path to the config file.
+
+    Returns:
+        Any: The contents of the config file as a Python object.
+    """
+    _, ext = os.path.splitext(path)
+    if ext == ".json":
+        with open(
+            path,
+        ) as f:
+            return json.load(f)
+    elif ext in [".yaml", ".yml"]:
+        with open(
+            path,
+        ) as f:
+            return yaml.safe_load(f)
+    else:
+        raise LaunchError(f"Unsupported file extension: {ext}")
+
+
+def _update_dict(target: Dict, source: Dict) -> None:
+    """Update a dictionary with the contents of another dictionary.
+
+    Args:
+        target (Dict): The dictionary to update.
+        source (Dict): The dictionary to update from.
+    """
+    for key, value in source.items():
+        if isinstance(value, dict):
+            if key not in target:
+                target[key] = {}
+            _update_dict(target[key], value)
+        else:
+            target[key] = value

wandb/sdk/launch/inputs/internal.py

@@ -0,0 +1,217 @@
+"""The layer between launch sdk user code and the wandb internal process.
+
+If there is an active run this communication is done through the wandb run's
+backend interface.
+
+If there is no active run, the messages are staged on the StagedLaunchInputs
+singleton and sent when a run is created.
+"""
+
+import os
+import pathlib
+import shutil
+import tempfile
+from typing import List, Optional
+
+import wandb
+import wandb.data_types
+from wandb.sdk.launch.errors import LaunchError
+from wandb.sdk.wandb_run import Run
+
+from .files import config_path_is_valid, override_file
+
+PERIOD = "."
+BACKSLASH = "\\"
+
+
+class ConfigTmpDir:
+    """Singleton for managing temporary directories for configuration files.
+
+    Any configuration files designated as inputs to a launch job are copied to
+    a temporary directory. This singleton manages the temporary directory and
+    provides paths to the configuration files.
+    """
+
+    _instance = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = object.__new__(cls)
+        return cls._instance
+
+    def __init__(self):
+        if not hasattr(self, "_tmp_dir"):
+            self._tmp_dir = tempfile.mkdtemp()
+            self._configs_dir = os.path.join(self._tmp_dir, "configs")
+            os.mkdir(self._configs_dir)
+
+    @property
+    def tmp_dir(self):
+        return pathlib.Path(self._tmp_dir)
+
+    @property
+    def configs_dir(self):
+        return pathlib.Path(self._configs_dir)
+
+
+class JobInputArguments:
+    """Arguments for the publish_job_input of Interface."""
+
+    def __init__(
+        self,
+        include: Optional[List[str]] = None,
+        exclude: Optional[List[str]] = None,
+        file_path: Optional[str] = None,
+        run_config: Optional[bool] = None,
+    ):
+        self.include = include
+        self.exclude = exclude
+        self.file_path = file_path
+        self.run_config = run_config
+
+
+class StagedLaunchInputs:
+    _instance = None
+
+    def __new__(cls):
+        if cls._instance is None:
+            cls._instance = object.__new__(cls)
+        return cls._instance
+
+    def __init__(self) -> None:
+        if not hasattr(self, "_staged_inputs"):
+            self._staged_inputs: List[JobInputArguments] = []
+
+    def add_staged_input(
+        self,
+        input_arguments: JobInputArguments,
+    ):
+        self._staged_inputs.append(input_arguments)
+
+    def apply(self, run: Run):
+        """Apply the staged inputs to the given run."""
+        for input in self._staged_inputs:
+            _publish_job_input(input, run)
+
+
+def _publish_job_input(
+    input: JobInputArguments,
+    run: Run,
+) -> None:
+    """Publish a job input to the backend interface of the given run.
+
+    Arguments:
+        input (JobInputArguments): The arguments for the job input.
+        run (Run): The run to publish the job input to.
+    """
+    assert run._backend is not None
+    assert run._backend.interface is not None
+    assert input.run_config is not None
+
+    interface = run._backend.interface
+    if input.file_path:
+        config_dir = ConfigTmpDir()
+        dest = os.path.join(config_dir.configs_dir, input.file_path)
+        run.save(dest, base_path=config_dir.tmp_dir)
+    interface.publish_job_input(
+        include_paths=[_split_on_unesc_dot(path) for path in input.include]
+        if input.include
+        else [],
+        exclude_paths=[_split_on_unesc_dot(path) for path in input.exclude]
+        if input.exclude
+        else [],
+        run_config=input.run_config,
+        file_path=input.file_path or "",
+    )
+
+
+def handle_config_file_input(
+    path: str,
+    include: Optional[List[str]] = None,
+    exclude: Optional[List[str]] = None,
+):
+    """Declare an overridable configuration file for a launch job.
+
+    The configuration file is copied to a temporary directory and the path to
+    the copy is sent to the backend interface of the active run and used to
+    configure the job builder.
+
+    If there is no active run, the configuration file is staged and sent when a
+    run is created.
+    """
+    config_path_is_valid(path)
+    override_file(path)
+    tmp_dir = ConfigTmpDir()
+    dest = os.path.join(tmp_dir.configs_dir, path)
+    shutil.copy(path, dest)
+    arguments = JobInputArguments(
+        include=include,
+        exclude=exclude,
+        file_path=path,
+        run_config=False,
+    )
+    if wandb.run is not None:
+        _publish_job_input(arguments, wandb.run)
+    else:
+        staged_inputs = StagedLaunchInputs()
+        staged_inputs.add_staged_input(arguments)
+
+
+def handle_run_config_input(
+    include: Optional[List[str]] = None, exclude: Optional[List[str]] = None
+):
+    """Declare wandb.config as an overridable configuration for a launch job.
+
+    The include and exclude paths are sent to the backend interface of the
+    active run and used to configure the job builder.
+
+    If there is no active run, the include and exclude paths are staged and sent
+    when a run is created.
+    """
+    arguments = JobInputArguments(
+        include=include,
+        exclude=exclude,
+        run_config=True,
+        file_path=None,
+    )
+    if wandb.run is not None:
+        _publish_job_input(arguments, wandb.run)
+    else:
+        stage_inputs = StagedLaunchInputs()
+        stage_inputs.add_staged_input(arguments)
+
+
+def _split_on_unesc_dot(path: str) -> List[str]:
+    r"""Split a string on unescaped dots.
+
+    Arguments:
+        path (str): The string to split.
+
+    Raises:
+        ValueError: If the path has a trailing escape character.
+
+    Returns:
+        List[str]: The split string.
+    """
+    parts = []
+    part = ""
+    i = 0
+    while i < len(path):
+        if path[i] == BACKSLASH:
+            if i == len(path) - 1:
+                raise LaunchError(
+                    f"Invalid config path {path}: trailing {BACKSLASH}.",
+                )
+            if path[i + 1] == PERIOD:
+                part += PERIOD
+                i += 2
+        elif path[i] == PERIOD:
+            parts.append(part)
+            part = ""
+            i += 1
+        else:
+            part += path[i]
+            i += 1
+    if part:
+        parts.append(part)
+    return parts

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

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)}"