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

wandb/sdk/launch/builder/context_manager.py

@@ -0,0 +1,235 @@
+import logging
+import os
+import shutil
+import tempfile
+from typing import Tuple
+
+from wandb.sdk.launch._project_spec import LaunchProject
+from wandb.sdk.launch.builder.build import image_tag_from_dockerfile_and_source
+from wandb.sdk.launch.errors import LaunchError
+from wandb.sdk.launch.utils import get_current_python_version
+
+from .build import (
+    _WANDB_DOCKERFILE_NAME,
+    get_base_setup,
+    get_docker_user,
+    get_entrypoint_setup,
+    get_requirements_section,
+    get_user_setup,
+)
+from .templates.dockerfile import DOCKERFILE_TEMPLATE
+
+_logger = logging.getLogger(__name__)
+
+
+class BuildContextManager:
+    """Creates a build context for a container image from job source code.
+
+    The dockerfile and build context may be specified by the job itself. If not,
+    the behavior for creating the build context is as follows:
+
+    - If a Dockerfile.wandb is found adjacent to the entrypoint, the directory
+        containing the entrypoint is used as the build context and Dockerfile.wandb
+        is used as the Dockerfile.
+
+    - If `override_dockerfile` is set on the LaunchProject, the directory
+        containing the Dockerfile is used as the build context and the Dockerfile
+        is used as the Dockerfile. `override_dockerfile` can be set in a launch
+        spec via the `-D` flag to `wandb launch` or in the `overrides` section
+        of the launch drawer.
+
+    - If no dockerfile is set, a Dockerfile is generated from the job's
+        requirements and entrypoint.
+    """
+
+    def __init__(self, launch_project: LaunchProject):
+        """Initialize a BuildContextManager.
+
+        Arguments:
+            launch_project: The launch project.
+        """
+        self._launch_project = launch_project
+        assert self._launch_project.project_dir is not None
+        self._directory = tempfile.mkdtemp()
+
+    def _generate_dockerfile(self, builder_type: str) -> str:
+        """Generate a Dockerfile for the container image.
+
+        Arguments:
+            builder_type: The type of builder to use. One of "docker" or "kaniko".
+
+        Returns:
+            The contents of the Dockerfile.
+        """
+        launch_project = self._launch_project
+        entry_point = (
+            launch_project.override_entrypoint or launch_project.get_job_entry_point()
+        )
+
+        # get python versions truncated to major.minor to ensure image availability
+        if launch_project.python_version:
+            spl = launch_project.python_version.split(".")[:2]
+            py_version, py_major = (".".join(spl), spl[0])
+        else:
+            py_version, py_major = get_current_python_version()
+
+        python_build_image = (
+            f"python:{py_version}"  # use full python image for package installation
+        )
+        requirements_section = get_requirements_section(
+            launch_project, self._directory, builder_type
+        )
+        # ----- stage 2: base -----
+        python_base_setup = get_base_setup(launch_project, py_version, py_major)
+
+        # set up user info
+        username, userid = get_docker_user(launch_project, launch_project.resource)
+        user_setup = get_user_setup(username, userid, launch_project.resource)
+        workdir = f"/home/{username}"
+
+        assert entry_point is not None
+        entrypoint_section = get_entrypoint_setup(entry_point)
+
+        dockerfile_contents = DOCKERFILE_TEMPLATE.format(
+            py_build_image=python_build_image,
+            requirements_section=requirements_section,
+            base_setup=python_base_setup,
+            uid=userid,
+            user_setup=user_setup,
+            workdir=workdir,
+            entrypoint_section=entrypoint_section,
+        )
+        return dockerfile_contents
+
+    def create_build_context(self, builder_type: str) -> Tuple[str, str]:
+        """Create the build context for the container image.
+
+        Returns:
+            A pair of str: the path to the build context locally and the image
+            tag computed from the Dockerfile.
+        """
+        entrypoint = (
+            self._launch_project.get_job_entry_point()
+            or self._launch_project.override_entrypoint
+        )
+        assert entrypoint is not None
+        assert entrypoint.name is not None
+        assert self._launch_project.project_dir is not None
+
+        # we use that as the build context.
+        build_context_root_dir = self._launch_project.project_dir
+        job_build_context = self._launch_project.job_build_context
+        if job_build_context:
+            full_path = os.path.join(build_context_root_dir, job_build_context)
+            if not os.path.exists(full_path):
+                raise LaunchError(f"Build context does not exist at {full_path}")
+            build_context_root_dir = full_path
+
+        # This is the case where the user specifies a Dockerfile to use.
+        # We use the directory containing the Dockerfile as the build context.
+        override_dockerfile = self._launch_project.override_dockerfile
+        if override_dockerfile:
+            full_path = os.path.join(
+                build_context_root_dir,
+                override_dockerfile,
+            )
+            if not os.path.exists(full_path):
+                raise LaunchError(f"Dockerfile does not exist at {full_path}")
+            shutil.copytree(
+                build_context_root_dir,
+                self._directory,
+                symlinks=True,
+                dirs_exist_ok=True,
+                ignore=shutil.ignore_patterns("fsmonitor--daemon.ipc"),
+            )
+            shutil.copy(
+                full_path,
+                os.path.join(self._directory, _WANDB_DOCKERFILE_NAME),
+            )
+            return self._directory, image_tag_from_dockerfile_and_source(
+                self._launch_project, open(full_path).read()
+            )
+
+        # If the job specifies a Dockerfile, we use that as the Dockerfile.
+        job_dockerfile = self._launch_project.job_dockerfile
+        if job_dockerfile:
+            dockerfile_path = os.path.join(build_context_root_dir, job_dockerfile)
+            if not os.path.exists(dockerfile_path):
+                raise LaunchError(f"Dockerfile does not exist at {dockerfile_path}")
+            shutil.copytree(
+                build_context_root_dir,
+                self._directory,
+                symlinks=True,
+                dirs_exist_ok=True,
+                ignore=shutil.ignore_patterns("fsmonitor--daemon.ipc"),
+            )
+            shutil.copy(
+                dockerfile_path,
+                os.path.join(self._directory, _WANDB_DOCKERFILE_NAME),
+            )
+            return self._directory, image_tag_from_dockerfile_and_source(
+                self._launch_project, open(dockerfile_path).read()
+            )
+
+        # This is the case where we find Dockerfile.wandb adjacent to the
+        # entrypoint. We use the entrypoint directory as the build context.
+        entrypoint_dir = os.path.dirname(entrypoint.name)
+        if entrypoint_dir:
+            path = os.path.join(
+                build_context_root_dir,
+                entrypoint_dir,
+                _WANDB_DOCKERFILE_NAME,
+            )
+        else:
+            path = os.path.join(build_context_root_dir, _WANDB_DOCKERFILE_NAME)
+        if os.path.exists(
+            path
+        ):  # We found a Dockerfile.wandb adjacent to the entrypoint.
+            shutil.copytree(
+                os.path.dirname(path),
+                self._directory,
+                symlinks=True,
+                dirs_exist_ok=True,
+                ignore=shutil.ignore_patterns("fsmonitor--daemon.ipc"),
+            )
+            # TODO: remove this once we make things more explicit for users
+            if entrypoint_dir:
+                new_path = os.path.basename(entrypoint.name)
+                entrypoint = self._launch_project.get_job_entry_point()
+                if entrypoint is not None:
+                    entrypoint.update_entrypoint_path(new_path)
+            with open(path) as f:
+                docker_file_contents = f.read()
+            return self._directory, image_tag_from_dockerfile_and_source(
+                self._launch_project, docker_file_contents
+            )
+
+        # This is the case where we use our own Dockerfile template. We move
+        # the user code into a src directory in the build context.
+        dst_path = os.path.join(self._directory, "src")
+        assert self._launch_project.project_dir is not None
+        shutil.copytree(
+            src=self._launch_project.project_dir,
+            dst=dst_path,
+            symlinks=True,
+            ignore=shutil.ignore_patterns("fsmonitor--daemon.ipc"),
+        )
+        shutil.copy(
+            os.path.join(os.path.dirname(__file__), "templates", "_wandb_bootstrap.py"),
+            os.path.join(self._directory),
+        )
+        if self._launch_project.python_version:
+            runtime_path = os.path.join(dst_path, "runtime.txt")
+            with open(runtime_path, "w") as fp:
+                fp.write(f"python-{self._launch_project.python_version}")
+
+        # TODO: we likely don't need to pass the whole git repo into the container
+        # with open(os.path.join(directory, ".dockerignore"), "w") as f:
+        #    f.write("**/.git")
+        with open(os.path.join(self._directory, _WANDB_DOCKERFILE_NAME), "w") as handle:
+            docker_file_contents = self._generate_dockerfile(builder_type=builder_type)
+            handle.write(docker_file_contents)
+        image_tag = image_tag_from_dockerfile_and_source(
+            self._launch_project, docker_file_contents
+        )
+        return self._directory, image_tag

wandb/sdk/launch/builder/docker_builder.py

@@ -1,4 +1,5 @@
 """Implementation of the docker builder."""
+
 import logging
 import os
 from typing import Any, Dict, Optional
@@ -6,8 +7,7 @@ from typing import Any, Dict, Optional
 import wandb
 import wandb.docker as docker
 from wandb.sdk.launch.agent.job_status_tracker import JobAndRunStatusTracker
-from wandb.sdk.launch.builder.abstract import AbstractBuilder
-from wandb.sdk.launch.builder.build import registry_from_uri
+from wandb.sdk.launch.builder.abstract import AbstractBuilder, registry_from_uri
 from wandb.sdk.launch.environment.abstract import AbstractEnvironment
 from wandb.sdk.launch.registry.abstract import AbstractRegistry
 
@@ -20,13 +20,8 @@ from ..utils import (
     event_loop_thread_exec,
     warn_failed_packages_from_build_logs,
 )
-from .build import (
-    _WANDB_DOCKERFILE_NAME,
-    _create_docker_build_ctx,
-    generate_dockerfile,
-    image_tag_from_dockerfile_and_source,
-    validate_docker_installation,
-)
+from .build import _WANDB_DOCKERFILE_NAME, validate_docker_installation
+from .context_manager import BuildContextManager
 
 _logger = logging.getLogger(__name__)
 
@@ -40,7 +35,6 @@ class DockerBuilder(AbstractBuilder):
     """
 
     builder_type = "docker"
-    base_image = "python:3.8"
     target_platform = "linux/amd64"
 
     def __init__(
@@ -123,17 +117,11 @@ class DockerBuilder(AbstractBuilder):
         await self.verify()
         await self.login()
 
-        dockerfile_str = generate_dockerfile(
-            launch_project=launch_project,
-            entry_point=entrypoint,
-            runner_type=launch_project.resource,
-            builder_type="docker",
-            dockerfile=launch_project.override_dockerfile,
-        )
-
-        image_tag = image_tag_from_dockerfile_and_source(launch_project, dockerfile_str)
-
+        build_context_manager = BuildContextManager(launch_project=launch_project)
+        build_ctx_path, image_tag = build_context_manager.create_build_context("docker")
+        dockerfile = os.path.join(build_ctx_path, _WANDB_DOCKERFILE_NAME)
         repository = None if not self.registry else await self.registry.get_repo_uri()
+
         # if repo is set, use the repo name as the image name
         if repository:
             image_uri = f"{repository}:{image_tag}"
@@ -151,9 +139,6 @@ class DockerBuilder(AbstractBuilder):
         _logger.info(
             f"image {image_uri} does not already exist in repository, building."
         )
-
-        build_ctx_path = _create_docker_build_ctx(launch_project, dockerfile_str)
-        dockerfile = os.path.join(build_ctx_path, _WANDB_DOCKERFILE_NAME)
         try:
             output = await event_loop_thread_exec(docker.build)(
                 tags=[image_uri],