wandb 0.18.2__py3-none-musllinux_1_2_x86_64.whl

wandb/sdk/launch/builder/build.py

@@ -0,0 +1,297 @@
+import hashlib
+import json
+import logging
+import os
+import pathlib
+import shlex
+from typing import Any, Dict, List, Tuple
+
+from dockerpycreds.utils import find_executable  # type: ignore
+
+import wandb
+import wandb.env
+from wandb import docker
+from wandb.apis.internal import Api
+from wandb.sdk.launch.loader import (
+    builder_from_config,
+    environment_from_config,
+    registry_from_config,
+)
+from wandb.util import get_module
+
+from .._project_spec import EntryPoint, LaunchProject
+from ..errors import ExecutionError, LaunchError
+from ..utils import LOG_PREFIX, event_loop_thread_exec
+from .templates.dockerfile import (
+    ACCELERATOR_SETUP_TEMPLATE,
+    ENTRYPOINT_TEMPLATE,
+    PIP_TEMPLATE,
+    PYTHON_SETUP_TEMPLATE,
+    USER_CREATE_TEMPLATE,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+_WANDB_DOCKERFILE_NAME = "Dockerfile.wandb"
+
+
+async def validate_docker_installation() -> None:
+    """Verify if Docker is installed on host machine."""
+    find_exec = event_loop_thread_exec(find_executable)
+    if not await find_exec("docker"):
+        raise ExecutionError(
+            "Could not find Docker executable. "
+            "Ensure Docker is installed as per the instructions "
+            "at https://docs.docker.com/install/overview/."
+        )
+
+
+def join(split_command: List[str]) -> str:
+    """Return a shell-escaped string from *split_command*.
+
+    Also remove quotes from double quoted strings. Ex:
+    "'local container queue'" --> "local container queue"
+    """
+    return " ".join(shlex.quote(arg.replace("'", "")) for arg in split_command)
+
+
+async def build_image_from_project(
+    launch_project: LaunchProject,
+    api: Api,
+    launch_config: Dict[str, Any],
+) -> str:
+    """Construct a docker image from a project and returns the URI of the image.
+
+    Arguments:
+        launch_project: The project to build an image from.
+        api: The API object to use for fetching the project.
+        launch_config: The launch config to use for building the image.
+
+    Returns:
+        The URI of the built image.
+    """
+    assert launch_project.uri, "To build an image on queue a URI must be set."
+    launch_config = launch_config or {}
+    env_config = launch_config.get("environment", {})
+    if not isinstance(env_config, dict):
+        wrong_type = type(env_config).__name__
+        raise LaunchError(
+            f"Invalid environment config: {env_config} of type {wrong_type} "
+            "loaded from launch config. Expected dict."
+        )
+    environment = environment_from_config(env_config)
+
+    registry_config = launch_config.get("registry", {})
+    if not isinstance(registry_config, dict):
+        wrong_type = type(registry_config).__name__
+        raise LaunchError(
+            f"Invalid registry config: {registry_config} of type {wrong_type}"
+            " loaded from launch config. Expected dict."
+        )
+    registry = registry_from_config(registry_config, environment)
+
+    builder_config = launch_config.get("builder", {})
+    if not isinstance(builder_config, dict):
+        wrong_type = type(builder_config).__name__
+        raise LaunchError(
+            f"Invalid builder config: {builder_config} of type {wrong_type} "
+            "loaded from launch config. Expected dict."
+        )
+    builder = builder_from_config(builder_config, environment, registry)
+
+    if not builder:
+        raise LaunchError("Unable to build image. No builder found.")
+
+    launch_project.fetch_and_validate_project()
+
+    entry_point = (
+        launch_project.get_job_entry_point() or launch_project.override_entrypoint
+    )
+    assert entry_point is not None
+    wandb.termlog(f"{LOG_PREFIX}Building docker image from uri source")
+    image_uri = await builder.build_image(launch_project, entry_point)
+    if not image_uri:
+        raise LaunchError("Error building image uri")
+    else:
+        return image_uri
+
+
+def image_tag_from_dockerfile_and_source(
+    launch_project: LaunchProject, dockerfile_contents: str
+) -> str:
+    """Hashes the source and dockerfile contents into a unique tag."""
+    image_source_string = launch_project.get_image_source_string()
+    unique_id_string = image_source_string + dockerfile_contents
+    image_tag = hashlib.sha256(unique_id_string.encode("utf-8")).hexdigest()[:8]
+    return image_tag
+
+
+def get_docker_user(launch_project: LaunchProject, runner_type: str) -> Tuple[str, int]:
+    import getpass
+
+    username = getpass.getuser()
+
+    if runner_type == "sagemaker" and not launch_project.docker_image:
+        # unless user has provided their own image, sagemaker must run as root but keep the name for workdir etc
+        return username, 0
+
+    userid = launch_project.docker_user_id or os.geteuid()
+    return username, userid
+
+
+def get_base_setup(
+    launch_project: LaunchProject, py_version: str, py_major: str
+) -> str:
+    """Fill in the Dockerfile templates for stage 2 of build.
+
+    CPU version is built on python, Accelerator version is built on user provided.
+    """
+    minor = int(py_version.split(".")[1])
+    if minor < 12:
+        python_base_image = f"python:{py_version}-buster"
+    else:
+        python_base_image = f"python:{py_version}-bookworm"
+    if launch_project.accelerator_base_image:
+        _logger.info(
+            f"Using accelerator base image: {launch_project.accelerator_base_image}"
+        )
+        python_packages = [
+            f"python{py_version}",
+            f"libpython{py_version}",
+            "python3-pip",
+            "python3-setuptools",
+        ]
+        base_setup = ACCELERATOR_SETUP_TEMPLATE.format(
+            accelerator_base_image=launch_project.accelerator_base_image,
+            python_packages=" \\\n".join(python_packages),
+            py_version=py_version,
+        )
+    else:
+        python_packages = [
+            "python3-dev",
+            "gcc",
+        ]  # gcc required for python < 3.7 for some reason
+        base_setup = PYTHON_SETUP_TEMPLATE.format(py_base_image=python_base_image)
+    return base_setup
+
+
+# Move this into the build context manager.
+def get_requirements_section(
+    launch_project: LaunchProject, build_context_dir: str, builder_type: str
+) -> str:
+    if builder_type == "docker":
+        buildx_installed = docker.is_buildx_installed()
+        if not buildx_installed:
+            wandb.termwarn(
+                "Docker BuildX is not installed, for faster builds upgrade docker: https://github.com/docker/buildx#installing"
+            )
+            prefix = "RUN WANDB_DISABLE_CACHE=true"
+    elif builder_type == "kaniko":
+        prefix = "RUN WANDB_DISABLE_CACHE=true"
+        buildx_installed = False
+
+    if buildx_installed:
+        prefix = "RUN --mount=type=cache,mode=0777,target=/root/.cache/pip"
+
+    requirements_files = []
+    deps_install_line = None
+
+    base_path = pathlib.Path(build_context_dir)
+    # If there is a requirements.txt at root of build context, use that.
+    if (base_path / "src" / "requirements.txt").exists():
+        requirements_files += ["src/requirements.txt"]
+        deps_install_line = "pip install uv && uv pip install -r requirements.txt"
+        with open(base_path / "src" / "requirements.txt") as f:
+            requirements = f.readlines()
+        if not any(["wandb" in r for r in requirements]):
+            wandb.termwarn(f"{LOG_PREFIX}wandb is not present in requirements.txt.")
+        return PIP_TEMPLATE.format(
+            buildx_optional_prefix=prefix,
+            requirements_files=" ".join(requirements_files),
+            pip_install=deps_install_line,
+        )
+
+    # Elif there is pyproject.toml at build context, convert the dependencies
+    # section to a requirements.txt and use that.
+    elif (base_path / "src" / "pyproject.toml").exists():
+        tomli = get_module("tomli")
+        if tomli is None:
+            wandb.termwarn(
+                "pyproject.toml found but tomli could not be loaded. To "
+                "install dependencies from pyproject.toml please run "
+                "`pip install tomli` and try again."
+            )
+        else:
+            # First try to read deps from standard pyproject format.
+            with open(base_path / "src" / "pyproject.toml", "rb") as f:
+                contents = tomli.load(f)
+            project_deps = [
+                str(d) for d in contents.get("project", {}).get("dependencies", [])
+            ]
+            if project_deps:
+                if not any(["wandb" in d for d in project_deps]):
+                    wandb.termwarn(
+                        f"{LOG_PREFIX}wandb is not present as a dependency in pyproject.toml."
+                    )
+                with open(base_path / "src" / "requirements.txt", "w") as f:
+                    f.write("\n".join(project_deps))
+                requirements_files += ["src/requirements.txt"]
+                deps_install_line = (
+                    "pip install uv && uv pip install -r requirements.txt"
+                )
+                return PIP_TEMPLATE.format(
+                    buildx_optional_prefix=prefix,
+                    requirements_files=" ".join(requirements_files),
+                    pip_install=deps_install_line,
+                )
+
+    # Else use frozen requirements from wandb run.
+    if (
+        not deps_install_line
+        and (base_path / "src" / "requirements.frozen.txt").exists()
+    ):
+        requirements_files += [
+            "src/requirements.frozen.txt",
+            "_wandb_bootstrap.py",
+        ]
+        deps_install_line = (
+            launch_project.parse_existing_requirements() + "python _wandb_bootstrap.py"
+        )
+
+        if not deps_install_line:
+            raise LaunchError(f"No dependency sources found for {launch_project}")
+
+        with open(base_path / "src" / "requirements.frozen.txt") as f:
+            requirements = f.readlines()
+        if not any(["wandb" in r for r in requirements]):
+            wandb.termwarn(
+                f"{LOG_PREFIX}wandb is not present in requirements.frozen.txt."
+            )
+
+        return PIP_TEMPLATE.format(
+            buildx_optional_prefix=prefix,
+            requirements_files=" ".join(requirements_files),
+            pip_install=deps_install_line,
+        )
+
+    else:
+        # this means no deps file was found
+        requirements_line = "RUN mkdir -p env/"  # Docker fails otherwise
+        wandb.termwarn("No requirements file found. No packages will be installed.")
+        return requirements_line
+
+
+def get_user_setup(username: str, userid: int, runner_type: str) -> str:
+    if runner_type == "sagemaker":
+        # sagemaker must run as root
+        return "USER root"
+    user_create = USER_CREATE_TEMPLATE.format(uid=userid, user=username)
+    user_create += f"\nUSER {username}"
+    return user_create
+
+
+def get_entrypoint_setup(
+    entry_point: EntryPoint,
+) -> str:
+    return ENTRYPOINT_TEMPLATE.format(entrypoint=json.dumps(entry_point.command))

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