wandb 0.17.0rc1__py3-none-any.whl → 0.17.1__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

@@ -7,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
 
@@ -21,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__)
 
@@ -41,7 +35,6 @@ class DockerBuilder(AbstractBuilder):
     """
 
     builder_type = "docker"
-    base_image = "python:3.8"
     target_platform = "linux/amd64"
 
     def __init__(
@@ -124,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}"
@@ -152,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],

wandb/sdk/launch/builder/kaniko_builder.py

@@ -13,8 +13,7 @@ from typing import Any, Dict, Optional
 
 import wandb
 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.environment.azure_environment import AzureEnvironment
 from wandb.sdk.launch.registry.abstract import AbstractRegistry
@@ -32,12 +31,8 @@ from ..utils import (
     get_kube_context_and_api_client,
     warn_failed_packages_from_build_logs,
 )
-from .build import (
-    _WANDB_DOCKERFILE_NAME,
-    _create_docker_build_ctx,
-    generate_dockerfile,
-    image_tag_from_dockerfile_and_source,
-)
+from .build import _WANDB_DOCKERFILE_NAME
+from .context_manager import BuildContextManager
 
 get_module(
     "kubernetes_asyncio",
@@ -261,17 +256,13 @@ class KanikoBuilder(AbstractBuilder):
         job_tracker: Optional[JobAndRunStatusTracker] = None,
     ) -> str:
         await self.verify()
-        # kaniko builder doesn't seem to work with a custom user id, need more investigation
-        dockerfile_str = generate_dockerfile(
-            launch_project=launch_project,
-            entry_point=entrypoint,
-            runner_type=launch_project.resource,
-            builder_type="kaniko",
-            dockerfile=launch_project.override_dockerfile,
-        )
-        image_tag = image_tag_from_dockerfile_and_source(launch_project, dockerfile_str)
+
+        build_contex_manager = BuildContextManager(launch_project=launch_project)
+        context_path, image_tag = build_contex_manager.create_build_context("kaniko")
+        run_id = launch_project.run_id
         repo_uri = await self.registry.get_repo_uri()
         image_uri = repo_uri + ":" + image_tag
+
         if (
             not launch_project.build_required()
             and await self.registry.check_image_exists(image_uri)
@@ -279,14 +270,10 @@ class KanikoBuilder(AbstractBuilder):
             return image_uri
 
         _logger.info(f"Building image {image_uri}...")
-
-        context_path = _create_docker_build_ctx(launch_project, dockerfile_str)
-        run_id = launch_project.run_id
-
         _, api_client = await get_kube_context_and_api_client(
             kubernetes, launch_project.resource_args
         )
-        # TODO: use same client as kuberentes_runner.py
+        # TODO: use same client as kubernetes_runner.py
         batch_v1 = client.BatchV1Api(api_client)
         core_v1 = client.CoreV1Api(api_client)
 
@@ -492,8 +479,8 @@ class KanikoBuilder(AbstractBuilder):
                     }
                 )
             else:
-                raise LaunchError(
-                    f"Registry type {type(self.registry)} not supported by kaniko"
+                wandb.termwarn(
+                    f"{LOG_PREFIX}Automatic credential handling is not supported for registry type {type(self.registry)}. Build job: {self.build_job_name}"
                 )
             volumes.append(
                 {
@@ -522,7 +509,7 @@ class KanikoBuilder(AbstractBuilder):
             volume_mounts.append(
                 {"name": "docker-config", "mountPath": "/kaniko/.docker/"}
             )
-        # Kaniko doesn't want https:// at the begining of the image tag.
+        # Kaniko doesn't want https:// at the beginning of the image tag.
         destination = image_tag
         if destination.startswith("https://"):
             destination = destination.replace("https://", "")

wandb/sdk/launch/builder/templates/dockerfile.py

@@ -0,0 +1,92 @@
+DOCKERFILE_TEMPLATE = """
+# ----- stage 1: build -----
+FROM {py_build_image} as build
+
+# requirements section depends on pip vs conda, and presence of buildx
+ENV PIP_PROGRESS_BAR off
+{requirements_section}
+
+# ----- stage 2: base -----
+{base_setup}
+
+COPY --from=build /env /env
+ENV PATH="/env/bin:$PATH"
+
+ENV SHELL /bin/bash
+
+# some resources (eg sagemaker) must run on root
+{user_setup}
+
+WORKDIR {workdir}
+RUN chown -R {uid} {workdir}
+
+# make artifacts cache dir unrelated to build
+RUN mkdir -p {workdir}/.cache && chown -R {uid} {workdir}/.cache
+
+# copy code/etc
+COPY --chown={uid} src/ {workdir}
+
+ENV PYTHONUNBUFFERED=1
+
+{entrypoint_section}
+"""
+
+# this goes into base_setup in TEMPLATE
+PYTHON_SETUP_TEMPLATE = """
+FROM {py_base_image} as base
+"""
+
+# this goes into base_setup in TEMPLATE
+ACCELERATOR_SETUP_TEMPLATE = """
+FROM {accelerator_base_image} as base
+
+# make non-interactive so build doesn't block on questions
+ENV DEBIAN_FRONTEND=noninteractive
+
+# install python
+RUN apt-get update -qq && apt-get install --no-install-recommends -y \
+    {python_packages} \
+    && apt-get -qq purge && apt-get -qq clean \
+    && rm -rf /var/lib/apt/lists/*
+
+# make sure `python` points at the right version
+RUN update-alternatives --install /usr/bin/python python /usr/bin/python{py_version} 1 \
+    && update-alternatives --install /usr/local/bin/python python /usr/bin/python{py_version} 1
+"""
+
+# this goes into requirements_section in TEMPLATE
+PIP_TEMPLATE = """
+RUN python -m venv /env
+# make sure we install into the env
+ENV PATH="/env/bin:$PATH"
+
+COPY {requirements_files} ./
+{buildx_optional_prefix} {pip_install}
+"""
+
+# this goes into requirements_section in TEMPLATE
+CONDA_TEMPLATE = """
+COPY src/environment.yml .
+{buildx_optional_prefix} conda env create -f environment.yml -n env
+
+# pack the environment so that we can transfer to the base image
+RUN conda install -c conda-forge conda-pack
+RUN conda pack -n env -o /tmp/env.tar && \
+    mkdir /env && cd /env && tar xf /tmp/env.tar && \
+    rm /tmp/env.tar
+RUN /env/bin/conda-unpack
+"""
+
+USER_CREATE_TEMPLATE = """
+RUN useradd \
+    --create-home \
+    --no-log-init \
+    --shell /bin/bash \
+    --gid 0 \
+    --uid {uid} \
+    {user} || echo ""
+"""
+
+ENTRYPOINT_TEMPLATE = """
+ENTRYPOINT {entrypoint}
+"""