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

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
@@ -9,9 +10,12 @@ import wandb
 from wandb.apis.internal import Api
 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_current_python_version,
+    get_entrypoint_file,
+)
 from wandb.sdk.lib import filesystem
 from wandb.util import make_artifact_name_safe
 
@@ -19,6 +23,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,
@@ -30,6 +37,8 @@ def create_job(
     runtime: Optional[str] = None,
     entrypoint: Optional[str] = None,
     git_hash: Optional[str] = None,
+    build_context: Optional[str] = None,
+    dockerfile: Optional[str] = None,
 ) -> Optional[Artifact]:
     """Create a job from a path, not as the output of a run.
 
@@ -42,9 +51,12 @@ def create_job(
         description (Optional[str]): Description of the job.
         aliases (Optional[List[str]]): Aliases for the job.
         runtime (Optional[str]): Python runtime of the job, like 3.9.
-        entrypoint (Optional[str]): Entrypoint of the job.
+        entrypoint (Optional[str]): Entrypoint of the job. If build_context is
+            provided, path is relative to build_context.
         git_hash (Optional[str]): Git hash of a specific commit, when using git type jobs.
-
+        build_context (Optional[str]): Path to the build context, when using image type jobs.
+        dockerfile (Optional[str]): Path to the Dockerfile, when using image type jobs.
+            If build_context is provided, path is relative to build_context.
 
     Returns:
         Optional[Artifact]: The artifact created by the job, the action (for printing), and job aliases.
@@ -81,6 +93,8 @@ def create_job(
         runtime,
         entrypoint,
         git_hash,
+        build_context,
+        dockerfile,
     )
 
     return artifact_job
@@ -98,6 +112,8 @@ def _create_job(
     runtime: Optional[str] = None,
     entrypoint: Optional[str] = None,
     git_hash: Optional[str] = None,
+    build_context: Optional[str] = None,
+    dockerfile: Optional[str] = None,
 ) -> Tuple[Optional[Artifact], str, List[str]]:
     wandb.termlog(f"Creating launch job of type: {job_type}...")
 
@@ -107,6 +123,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:
@@ -145,6 +168,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,
@@ -160,7 +184,10 @@ def _create_job(
         name = job_name
 
     # build job artifact, loads wandb-metadata and creates wandb-job.json here
-    artifact = job_builder.build()
+    artifact = job_builder.build(
+        dockerfile=dockerfile,
+        build_context=build_context,
+    )
     if not artifact:
         wandb.termerror("JobBuilder failed to build a job")
         _logger.debug("Failed to build job, check job source and metadata")
@@ -219,6 +246,7 @@ def _make_metadata_for_partial_job(
     """Create metadata for partial jobs, return metadata and requirements."""
     metadata = {"_partial": "v0"}
     if job_type == "git":
+        assert entrypoint is not None
         repo_metadata = _create_repo_metadata(
             path=path,
             tempdir=tempdir.name,
@@ -233,13 +261,7 @@ 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"
-            )
-            return None, None
-
+        assert entrypoint is not None
         artifact_metadata, requirements = _create_artifact_metadata(
             path=path, entrypoint=entrypoint, runtime=runtime
         )
@@ -268,7 +290,7 @@ def _make_metadata_for_partial_job(
 def _create_repo_metadata(
     path: str,
     tempdir: str,
-    entrypoint: Optional[str] = None,
+    entrypoint: str,
     git_hash: Optional[str] = None,
     runtime: Optional[str] = None,
 ) -> Optional[Dict[str, Any]]:
@@ -304,25 +326,16 @@ 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")
-        return None
-
     metadata = {
         "git": {
             "commit": commit,
             "remote": ref.url,
         },
-        "codePathLocal": entrypoint,  # not in git context, optionally also set local
-        "codePath": entrypoint,
-        "entrypoint": [f"python{python_version}", entrypoint],
+        "entrypoint": entrypoint.split(" "),
         "python": python_version,  # used to build container
         "notebook": False,  # partial jobs from notebooks not supported
     }
@@ -332,10 +345,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 +367,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
@@ -411,7 +407,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],
@@ -420,17 +416,19 @@ def _make_code_artifact(
 
     Returns the name of the eventual job.
     """
-    artifact_name = _make_code_artifact_name(os.path.join(path, entrypoint or ""), name)
+    entrypoint_list = entrypoint.split(" ")
+    # We no longer require the entrypoint to end in an existing file. But we
+    # need something to use as the default job artifact name. In the future we
+    # may require the user to provide a job name explicitly when calling
+    # wandb job create.
+    entrypoint_file = entrypoint_list[-1]
+    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:
@@ -441,6 +439,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,
@@ -451,7 +456,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/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/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