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

wandb/sdk/launch/create_job.py

@@ -10,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, get_entrypoint_file
+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
 
@@ -34,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.
 
@@ -46,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.
@@ -85,6 +93,8 @@ def create_job(
         runtime,
         entrypoint,
         git_hash,
+        build_context,
+        dockerfile,
     )
 
     return artifact_job
@@ -102,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}...")
 
@@ -172,7 +184,11 @@ 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(
+        api.api,
+        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")
@@ -195,7 +211,7 @@ def _create_job(
         project_name=project,
         run_name=run.id,  # type: ignore # run will be deleted after creation
         description=description,
-        metadata=metadata,
+        metadata={"_partial": True},
         is_user_created=True,
         aliases=[{"artifactCollectionName": name, "alias": a} for a in aliases],
     )
@@ -229,8 +245,9 @@ def _make_metadata_for_partial_job(
     entrypoint: Optional[str],
 ) -> Tuple[Optional[Dict[str, Any]], Optional[List[str]]]:
     """Create metadata for partial jobs, return metadata and requirements."""
-    metadata = {"_partial": "v0"}
+    metadata = {}
     if job_type == "git":
+        assert entrypoint is not None
         repo_metadata = _create_repo_metadata(
             path=path,
             tempdir=tempdir.name,
@@ -245,12 +262,7 @@ def _make_metadata_for_partial_job(
         return metadata, None
 
     if job_type == "code":
-        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
         )
@@ -276,10 +288,18 @@ def _make_metadata_for_partial_job(
     return None, None
 
 
+def _maybe_warn_python_no_executable(entrypoint: str):
+    entrypoint_list = entrypoint.split(" ")
+    if len(entrypoint_list) == 1 and entrypoint_list[0].endswith(".py"):
+        wandb.termwarn(
+            f"Entrypoint {entrypoint} is a python file without an executable, you may want to use `python {entrypoint}` as the entrypoint instead."
+        )
+
+
 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]]:
@@ -287,6 +307,9 @@ def _create_repo_metadata(
     if entrypoint and ".." in entrypoint:
         wandb.termerror("Entrypoint cannot contain backward path traversal")
         return None
+
+    _maybe_warn_python_no_executable(entrypoint)
+
     if not _is_git_uri(path):
         wandb.termerror("Path must be a git URI")
         return None
@@ -315,32 +338,16 @@ def _create_repo_metadata(
             with open(os.path.join(local_dir, ".python-version")) as f:
                 python_version = f.read().strip().splitlines()[0]
         else:
-            _, python_version = get_current_python_version()
+            python_version, _ = get_current_python_version()
 
     python_version = _clean_python_version(python_version)
 
-    # check if entrypoint is valid
-    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
-
-    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 = {
         "git": {
             "commit": commit,
             "remote": ref.url,
         },
-        "codePathLocal": entrypoint_file,  # not in git context, optionally also set local
-        "codePath": entrypoint_file,
-        "entrypoint": entrypoint_list,
+        "entrypoint": entrypoint.split(" "),
         "python": python_version,  # used to build container
         "notebook": False,  # partial jobs from notebooks not supported
     }
@@ -354,13 +361,11 @@ def _create_artifact_metadata(
     if not os.path.isdir(path):
         wandb.termerror("Path must be a valid file or directory")
         return {}, []
+
+    _maybe_warn_python_no_executable(entrypoint)
+
     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 = []
@@ -369,6 +374,9 @@ def _create_artifact_metadata(
         with open(depspath) as f:
             requirements = f.read().splitlines()
 
+    if not any(["wandb" in r for r in requirements]):
+        wandb.termwarn("wandb is not present in requirements.txt.")
+
     if runtime:
         python_version = _clean_python_version(runtime)
     else:
@@ -399,6 +407,7 @@ def _configure_job_builder_for_partial(tmpdir: str, job_source: str) -> JobBuild
         settings=settings,  # type: ignore
         verbose=True,
     )
+    job_builder._partial = True
     # never allow notebook runs
     job_builder._is_notebook_run = False
     # set run inputs and outputs to empty dicts
@@ -421,15 +430,12 @@ def _make_code_artifact(
 
     Returns the name of the eventual job.
     """
-    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
-
+    # 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,

wandb/sdk/launch/environment/aws_environment.py

@@ -7,7 +7,7 @@ from typing import Dict, Optional
 from wandb.sdk.launch.errors import LaunchError
 from wandb.util import get_module
 
-from ..utils import S3_URI_RE, event_loop_thread_exec
+from ..utils import ARN_PARTITION_RE, S3_URI_RE, event_loop_thread_exec
 from .abstract import AbstractEnvironment
 
 boto3 = get_module(
@@ -49,6 +49,7 @@ class AwsEnvironment(AbstractEnvironment):
         self._secret_key = secret_key
         self._session_token = session_token
         self._account = None
+        self._partition = None
 
     @classmethod
     def from_default(cls, region: Optional[str] = None) -> "AwsEnvironment":
@@ -122,6 +123,30 @@ class AwsEnvironment(AbstractEnvironment):
     def region(self, region: str) -> None:
         self._region = region
 
+    async def get_partition(self) -> str:
+        """Set the partition for the AWS environment."""
+        try:
+            session = await self.get_session()
+            client = await event_loop_thread_exec(session.client)("sts")
+            get_caller_identity = event_loop_thread_exec(client.get_caller_identity)
+            identity = await get_caller_identity()
+            arn = identity.get("Arn")
+            if not arn:
+                raise LaunchError(
+                    "Could not set partition for AWS environment. ARN not found."
+                )
+            matched_partition = ARN_PARTITION_RE.match(arn)
+            if not matched_partition:
+                raise LaunchError(
+                    f"Could not set partition for AWS environment. ARN {arn} is not valid."
+                )
+            partition = matched_partition.group(1)
+            return partition
+        except botocore.exceptions.ClientError as e:
+            raise LaunchError(
+                f"Could not set partition for AWS environment. {e}"
+            ) from e
+
     async def verify(self) -> None:
         """Verify that the AWS environment is configured correctly.
 

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,224 @@
+"""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 = "\\"
+LAUNCH_MANAGED_CONFIGS_DIR = "_wandb_configs"
+
+
+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, LAUNCH_MANAGED_CONFIGS_DIR)
+            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)
+    dest_dir = os.path.dirname(dest)
+    if not os.path.exists(dest_dir):
+        os.makedirs(dest_dir)
+    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