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

wandb/sdk/launch/_project_spec.py

@@ -2,18 +2,21 @@
 
 Arguments can come from a launch spec or call to wandb launch.
 """
+
 import enum
+import json
 import logging
 import os
 import tempfile
 from copy import deepcopy
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, cast
 
+from six.moves import shlex_quote
+
 import wandb
-import wandb.docker as docker
 from wandb.apis.internal import Api
 from wandb.errors import CommError
-from wandb.sdk.launch import utils
+from wandb.sdk.launch.utils import get_entrypoint_file
 from wandb.sdk.lib.runid import generate_id
 
 from .errors import LaunchError
@@ -32,15 +35,18 @@ IMAGE_TAG_MAX_LENGTH = 32
 
 
 class LaunchSource(enum.IntEnum):
-    WANDB: int = 1
-    GIT: int = 2
-    LOCAL: int = 3
-    DOCKER: int = 4
-    JOB: int = 5
+    """Enumeration of possible sources for a launch project.
 
+    Attributes:
+        DOCKER: Source is a Docker image. This can happen if a user runs
+            `wandb launch -d <docker-image>`.
+        JOB: Source is a job. This is standard case.
+        SCHEDULER: Source is a wandb sweep scheduler command.
+    """
 
-class EntrypointDefaults(List[str]):
-    PYTHON = ["python", "main.py"]
+    DOCKER: int = 1
+    JOB: int = 2
+    SCHEDULER: int = 3
 
 
 class LaunchProject:
@@ -59,8 +65,16 @@ class LaunchProject:
 
     This class is stateful and certain methods can only be called after
     `LaunchProject.fetch_and_validate_project()` has been called.
+
+    Notes on the entrypoint:
+    - The entrypoint is the command that will be run inside the container.
+    - The LaunchProject stores two entrypoints
+        - The job entrypoint is the entrypoint specified in the job's config.
+        - The override entrypoint is the entrypoint specified in the launch spec.
+    - The override entrypoint takes precedence over the job entrypoint.
     """
 
+    # This init is way to long, and there are too many attributes on this sucker.
     def __init__(
         self,
         uri: Optional[str],
@@ -78,9 +92,6 @@ class LaunchProject:
         run_id: Optional[str],
         sweep_id: Optional[str] = None,
     ):
-        if uri is not None and utils.is_bare_wandb_uri(uri):
-            uri = api.settings("base_url") + uri
-            _logger.info(f"{LOG_PREFIX}Updating uri with base uri: {uri}")
         self.uri = uri
         self.job = job
         if job is not None:
@@ -104,74 +115,57 @@ class LaunchProject:
         self.accelerator_base_image: Optional[str] = resource_args_build.get(
             "accelerator", {}
         ).get("base_image") or resource_args_build.get("cuda", {}).get("base_image")
-        self._base_image: Optional[str] = launch_spec.get("base_image")
         self.docker_image: Optional[str] = docker_config.get(
             "docker_image"
         ) or launch_spec.get("image_uri")
-        uid = RESOURCE_UID_MAP.get(resource, 1000)
-        if self._base_image:
-            uid = docker.get_image_uid(self._base_image)
-            _logger.info(f"{LOG_PREFIX}Retrieved base image uid {uid}")
-        self.docker_user_id: int = docker_config.get("user_id", uid)
-        self.git_version: Optional[str] = git_info.get("version")
-        self.git_repo: Optional[str] = git_info.get("repo")
-        self.overrides = overrides
-        self.override_args: List[str] = overrides.get("args", [])
-        self.override_config: Dict[str, Any] = overrides.get("run_config", {})
-        self.override_artifacts: Dict[str, Any] = overrides.get("artifacts", {})
-        self.override_entrypoint: Optional[EntryPoint] = None
-        self.override_dockerfile: Optional[str] = overrides.get("dockerfile")
+        self.docker_user_id = docker_config.get("user_id", 1000)
+        self._entry_point: Optional[EntryPoint] = (
+            None  # todo: keep multiple entrypoint support?
+        )
+        self.init_overrides(overrides)
+        self.init_source()
+        self.init_git(git_info)
         self.deps_type: Optional[str] = None
         self._runtime: Optional[str] = None
         self.run_id = run_id or generate_id()
         self._queue_name: Optional[str] = None
         self._queue_entity: Optional[str] = None
         self._run_queue_item_id: Optional[str] = None
-        self._entry_point: Optional[
-            EntryPoint
-        ] = None  # todo: keep multiple entrypoint support?
+        self._job_dockerfile: Optional[str] = None
+        self._job_build_context: Optional[str] = None
 
-        override_entrypoint = overrides.get("entry_point")
-        if override_entrypoint:
-            _logger.info("Adding override entry point")
-            self.override_entrypoint = EntryPoint(
-                name=_get_entrypoint_file(override_entrypoint),
-                command=override_entrypoint,
-            )
-
-        if overrides.get("sweep_id") is not None:
-            _logger.info("Adding override sweep id")
-            self.sweep_id = overrides["sweep_id"]
+    def init_source(self) -> None:
         if self.docker_image is not None:
             self.source = LaunchSource.DOCKER
             self.project_dir = None
         elif self.job is not None:
             self.source = LaunchSource.JOB
             self.project_dir = tempfile.mkdtemp()
-        elif self.uri is not None and utils._is_wandb_uri(self.uri):
-            _logger.info(f"URI {self.uri} indicates a wandb uri")
-            self.source = LaunchSource.WANDB
-            self.project_dir = tempfile.mkdtemp()
-        elif self.uri is not None and utils._is_git_uri(self.uri):
-            _logger.info(f"URI {self.uri} indicates a git uri")
-            self.source = LaunchSource.GIT
-            self.project_dir = tempfile.mkdtemp()
-        elif self.uri is not None and "placeholder-" in self.uri:
-            wandb.termlog(
-                f"{LOG_PREFIX}Launch received placeholder URI, replacing with local path."
+        if self.uri and self.uri.startswith("placeholder"):
+            self.source = LaunchSource.SCHEDULER
+            self.project_dir = os.getcwd()
+            self._entry_point = self.override_entrypoint
+
+    def init_git(self, git_info: Dict[str, str]) -> None:
+        self.git_version = git_info.get("version")
+        self.git_repo = git_info.get("repo")
+
+    def init_overrides(self, overrides: Dict[str, Any]) -> None:
+        """Initialize override attributes for a launch project."""
+        self.overrides = overrides
+        self.override_args: List[str] = overrides.get("args", [])
+        self.override_config: Dict[str, Any] = overrides.get("run_config", {})
+        self.override_artifacts: Dict[str, Any] = overrides.get("artifacts", {})
+        self.override_files: Dict[str, Any] = overrides.get("files", {})
+        self.override_entrypoint: Optional[EntryPoint] = None
+        self.override_dockerfile: Optional[str] = overrides.get("dockerfile")
+        override_entrypoint = overrides.get("entry_point")
+        if override_entrypoint:
+            _logger.info("Adding override entry point")
+            self.override_entrypoint = EntryPoint(
+                name=get_entrypoint_file(override_entrypoint),
+                command=override_entrypoint,
             )
-            self.uri = os.getcwd()
-            self.source = LaunchSource.LOCAL
-            self.project_dir = self.uri
-        else:
-            _logger.info(f"URI {self.uri} indicates a local uri")
-            # assume local
-            if self.uri is not None and not os.path.exists(self.uri):
-                raise LaunchError(
-                    "Assumed URI supplied is a local path but path is not valid"
-                )
-            self.source = LaunchSource.LOCAL
-            self.project_dir = self.uri
 
     def __repr__(self) -> str:
         """String representation of LaunchProject."""
@@ -210,6 +204,20 @@ class LaunchProject:
             launch_spec.get("sweep_id", {}),
         )
 
+    @property
+    def job_dockerfile(self) -> Optional[str]:
+        return self._job_dockerfile
+
+    @property
+    def job_build_context(self) -> Optional[str]:
+        return self._job_build_context
+
+    def set_job_dockerfile(self, dockerfile: str) -> None:
+        self._job_dockerfile = dockerfile
+
+    def set_job_build_context(self, build_context: str) -> None:
+        self._job_build_context = build_context
+
     @property
     def image_name(self) -> str:
         if self.docker_image is not None:
@@ -273,7 +281,7 @@ class LaunchProject:
             image (str): The image name to fill in for ${wandb-image}.
 
         Returns:
-            None
+            Dict[str, Any]: The resource args with all macros filled in.
         """
         update_dict = {
             "project_name": self.target_project,
@@ -323,8 +331,8 @@ class LaunchProject:
         self._docker_image = value
         self._ensure_not_docker_image_and_local_process()
 
-    def get_single_entry_point(self) -> Optional["EntryPoint"]:
-        """Returns the first entrypoint for the project, or None if no entry point was provided because a docker image was provided."""
+    def get_job_entry_point(self) -> Optional["EntryPoint"]:
+        """Returns the job entrypoint for the project."""
         # assuming project only has 1 entry point, pull that out
         # tmp fn until we figure out if we want to support multiple entry points or not
         if not self._entry_point:
@@ -335,8 +343,8 @@ class LaunchProject:
             return None
         return self._entry_point
 
-    def set_entry_point(self, command: List[str]) -> "EntryPoint":
-        """Add an entry point to the project."""
+    def set_job_entry_point(self, command: List[str]) -> "EntryPoint":
+        """Set job entrypoint for the project."""
         assert (
             self._entry_point is None
         ), "Cannot set entry point twice. Use LaunchProject.override_entrypoint"
@@ -357,51 +365,23 @@ class LaunchProject:
         """
         if self.source == LaunchSource.DOCKER:
             return
-        if self.source == LaunchSource.LOCAL:
-            if not self._entry_point:
-                wandb.termlog(
-                    f"{LOG_PREFIX}Entry point for repo not specified, defaulting to `python main.py`"
-                )
-                self.set_entry_point(EntrypointDefaults.PYTHON)
         elif self.source == LaunchSource.JOB:
             self._fetch_job()
-        else:
-            self._fetch_project_local(internal_api=self.api)
-
         assert self.project_dir is not None
-        # this prioritizes pip, and we don't support any cases where both are present conda projects when uploaded to
-        # wandb become pip projects via requirements.frozen.txt, wandb doesn't preserve conda envs
-        if os.path.exists(
-            os.path.join(self.project_dir, "requirements.txt")
-        ) or os.path.exists(os.path.join(self.project_dir, "requirements.frozen.txt")):
-            self.deps_type = "pip"
-        elif os.path.exists(os.path.join(self.project_dir, "environment.yml")):
-            self.deps_type = "conda"
 
+    # Let's make sure we document this very clearly.
     def get_image_source_string(self) -> str:
         """Returns a unique string identifying the source of an image."""
-        if self.source == LaunchSource.LOCAL:
-            # TODO: more correct to get a hash of local uri contents
-            assert isinstance(self.uri, str)
-            return self.uri
-        elif self.source == LaunchSource.JOB:
+        if self.source == LaunchSource.JOB:
             assert self._job_artifact is not None
             return f"{self._job_artifact.name}:v{self._job_artifact.version}"
-        elif self.source == LaunchSource.GIT:
-            assert isinstance(self.uri, str)
-            ret = self.uri
-            if self.git_version:
-                ret += self.git_version
-            return ret
-        elif self.source == LaunchSource.WANDB:
-            assert isinstance(self.uri, str)
-            return self.uri
         elif self.source == LaunchSource.DOCKER:
             assert isinstance(self.docker_image, str)
-            _logger.debug("")
             return self.docker_image
         else:
-            raise LaunchError("Unknown source type when determing image source string")
+            raise LaunchError(
+                "Unknown source type when determining image source string"
+            )
 
     def _ensure_not_docker_image_and_local_process(self) -> None:
         """Ensure that docker image is not specified with local-process resource runner.
@@ -429,129 +409,84 @@ class LaunchProject:
             raise LaunchError(
                 f"Error accessing job {self.job}: {msg} on {public_api.settings.get('base_url')}"
             )
-        job.configure_launch_project(self)
+        job.configure_launch_project(self)  # Why is this a method of the job?
         self._job_artifact = job._job_artifact
 
-    def _fetch_project_local(self, internal_api: Api) -> None:
-        """Fetch a project (either wandb run or git repo) into a local directory, returning the path to the local project directory."""
-        # these asserts are all guaranteed to pass, but are required by mypy
-        assert self.source != LaunchSource.LOCAL and self.source != LaunchSource.JOB
-        assert isinstance(self.uri, str)
-        assert self.project_dir is not None
-        _logger.info("Fetching project locally...")
-        if utils._is_wandb_uri(self.uri):
-            source_entity, source_project, source_run_name = utils.parse_wandb_uri(
-                self.uri
-            )
-            run_info = utils.fetch_wandb_project_run_info(
-                source_entity, source_project, source_run_name, internal_api
-            )
-            program_name = run_info.get("codePath") or run_info["program"]
-
-            self.python_version = run_info.get("python", "3")
-            downloaded_code_artifact = utils.check_and_download_code_artifacts(
-                source_entity,
-                source_project,
-                source_run_name,
-                internal_api,
-                self.project_dir,
-            )
-            if not downloaded_code_artifact:
-                if not run_info["git"]:
-                    raise LaunchError(
-                        "Reproducing a run requires either an associated git repo or a code artifact logged with `run.log_code()`"
-                    )
-                branch_name = utils._fetch_git_repo(
-                    self.project_dir,
-                    run_info["git"]["remote"],
-                    run_info["git"]["commit"],
-                )
-                if self.git_version is None:
-                    self.git_version = branch_name
-                patch = utils.fetch_project_diff(
-                    source_entity, source_project, source_run_name, internal_api
-                )
-                if patch:
-                    utils.apply_patch(patch, self.project_dir)
-
-                # For cases where the entry point wasn't checked into git
-                if not os.path.exists(os.path.join(self.project_dir, program_name)):
-                    downloaded_entrypoint = utils.download_entry_point(
-                        source_entity,
-                        source_project,
-                        source_run_name,
-                        internal_api,
-                        program_name,
-                        self.project_dir,
-                    )
-
-                    if not downloaded_entrypoint:
-                        raise LaunchError(
-                            f"Entrypoint file: {program_name} does not exist, "
-                            "and could not be downloaded. Please specify the entrypoint for this run."
-                        )
-
-            if (
-                "_session_history.ipynb" in os.listdir(self.project_dir)
-                or ".ipynb" in program_name
-            ):
-                program_name = utils.convert_jupyter_notebook_to_script(
-                    program_name, self.project_dir
-                )
-
-            # Download any frozen requirements
-            utils.download_wandb_python_deps(
-                source_entity,
-                source_project,
-                source_run_name,
-                internal_api,
-                self.project_dir,
-            )
-
-            if not self._entry_point:
-                _, ext = os.path.splitext(program_name)
-                if ext == ".py":
-                    entry_point = ["python", program_name]
-                elif ext == ".sh":
-                    command = os.environ.get("SHELL", "bash")
-                    entry_point = [command, program_name]
-                else:
-                    raise LaunchError(f"Unsupported entrypoint: {program_name}")
-                self.set_entry_point(entry_point)
-            if not self.override_args:
-                self.override_args = run_info["args"]
-        else:
-            assert utils._GIT_URI_REGEX.match(self.uri), (
-                "Non-wandb URI %s should be a Git URI" % self.uri
-            )
-            if not self._entry_point:
-                wandb.termlog(
-                    f"{LOG_PREFIX}Entry point for repo not specified, defaulting to python main.py"
-                )
-                self.set_entry_point(EntrypointDefaults.PYTHON)
-            branch_name = utils._fetch_git_repo(
-                self.project_dir, self.uri, self.git_version
-            )
-            if self.git_version is None:
-                self.git_version = branch_name
+    def get_env_vars_dict(self, api: Api, max_env_length: int) -> Dict[str, str]:
+        """Generate environment variables for the project.
 
+        Arguments:
+        launch_project: LaunchProject to generate environment variables for.
 
-def _get_entrypoint_file(entrypoint: List[str]) -> Optional[str]:
-    """Get the entrypoint file from the given command.
+        Returns:
+            Dictionary of environment variables.
+        """
+        env_vars = {}
+        env_vars["WANDB_BASE_URL"] = api.settings("base_url")
+        override_api_key = self.launch_spec.get("_wandb_api_key")
+        env_vars["WANDB_API_KEY"] = override_api_key or api.api_key
+        if self.target_project:
+            env_vars["WANDB_PROJECT"] = self.target_project
+        env_vars["WANDB_ENTITY"] = self.target_entity
+        env_vars["WANDB_LAUNCH"] = "True"
+        env_vars["WANDB_RUN_ID"] = self.run_id
+        if self.docker_image:
+            env_vars["WANDB_DOCKER"] = self.docker_image
+        if self.name is not None:
+            env_vars["WANDB_NAME"] = self.name
+        if "author" in self.launch_spec and not override_api_key:
+            env_vars["WANDB_USERNAME"] = self.launch_spec["author"]
+        if self.sweep_id:
+            env_vars["WANDB_SWEEP_ID"] = self.sweep_id
+        if self.launch_spec.get("_resume_count", 0) > 0:
+            env_vars["WANDB_RESUME"] = "allow"
+        if self.queue_name:
+            env_vars[wandb.env.LAUNCH_QUEUE_NAME] = self.queue_name
+        if self.queue_entity:
+            env_vars[wandb.env.LAUNCH_QUEUE_ENTITY] = self.queue_entity
+        if self.run_queue_item_id:
+            env_vars[wandb.env.LAUNCH_TRACE_ID] = self.run_queue_item_id
+
+        _inject_wandb_config_env_vars(self.override_config, env_vars, max_env_length)
+        _inject_file_overrides_env_vars(self.override_files, env_vars, max_env_length)
+
+        artifacts = {}
+        # if we're spinning up a launch process from a job
+        # we should tell the run to use that artifact
+        if self.job:
+            artifacts = {wandb.util.LAUNCH_JOB_ARTIFACT_SLOT_NAME: self.job}
+        env_vars["WANDB_ARTIFACTS"] = json.dumps(
+            {**artifacts, **self.override_artifacts}
+        )
+        return env_vars
 
-    Args:
-        entrypoint (List[str]): List of command and arguments.
+    def parse_existing_requirements(self) -> str:
+        import pkg_resources
 
-    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]
+        requirements_line = ""
+        assert self.project_dir is not None
+        base_requirements = os.path.join(self.project_dir, "requirements.txt")
+        if os.path.exists(base_requirements):
+            include_only = set()
+            with open(base_requirements) as f:
+                iter = pkg_resources.parse_requirements(f)
+                while True:
+                    try:
+                        pkg = next(iter)
+                        if hasattr(pkg, "name"):
+                            name = pkg.name.lower()
+                        else:
+                            name = str(pkg)
+                        include_only.add(shlex_quote(name))
+                    except StopIteration:
+                        break
+                    # Different versions of pkg_resources throw different errors
+                    # just catch them all and ignore packages we can't parse
+                    except Exception as e:
+                        _logger.warn(f"Unable to parse requirements.txt: {e}")
+                        continue
+            requirements_line += "WANDB_ONLY_INCLUDE={} ".format(",".join(include_only))
+        return requirements_line
 
 
 class EntryPoint:
@@ -561,31 +496,43 @@ class EntryPoint:
         self.name = name
         self.command = command
 
-    def compute_command(self, user_parameters: Optional[List[str]]) -> List[str]:
-        """Converts user parameter dictionary to a string."""
-        ret = self.command
-        if user_parameters:
-            return ret + user_parameters
-        return ret
-
     def update_entrypoint_path(self, new_path: str) -> None:
         """Updates the entrypoint path to a new path."""
-        if len(self.command) == 2 and self.command[0] in ["python", "bash"]:
+        if len(self.command) == 2 and (
+            self.command[0].startswith("python") or self.command[0] == "bash"
+        ):
             self.command[1] = new_path
 
 
-def get_entry_point_command(
-    entry_point: Optional["EntryPoint"], parameters: List[str]
-) -> List[str]:
-    """Returns the shell command to execute in order to run the specified entry point.
-
-    Arguments:
-    entry_point: Entry point to run
-    parameters: Parameters (dictionary) for the entry point command
-
-    Returns:
-        List of strings representing the shell command to be executed
-    """
-    if entry_point is None:
-        return []
-    return entry_point.compute_command(parameters)
+def _inject_wandb_config_env_vars(
+    config: Dict[str, Any], env_dict: Dict[str, Any], maximum_env_length: int
+) -> None:
+    str_config = json.dumps(config)
+    if len(str_config) <= maximum_env_length:
+        env_dict["WANDB_CONFIG"] = str_config
+        return
+
+    chunks = [
+        str_config[i : i + maximum_env_length]
+        for i in range(0, len(str_config), maximum_env_length)
+    ]
+    config_chunks_dict = {f"WANDB_CONFIG_{i}": chunk for i, chunk in enumerate(chunks)}
+    env_dict.update(config_chunks_dict)
+
+
+def _inject_file_overrides_env_vars(
+    overrides: Dict[str, Any], env_dict: Dict[str, Any], maximum_env_length: int
+) -> None:
+    str_overrides = json.dumps(overrides)
+    if len(str_overrides) <= maximum_env_length:
+        env_dict["WANDB_LAUNCH_FILE_OVERRIDES"] = str_overrides
+        return
+
+    chunks = [
+        str_overrides[i : i + maximum_env_length]
+        for i in range(0, len(str_overrides), maximum_env_length)
+    ]
+    overrides_chunks_dict = {
+        f"WANDB_LAUNCH_FILE_OVERRIDES_{i}": chunk for i, chunk in enumerate(chunks)
+    }
+    env_dict.update(overrides_chunks_dict)