flyte 0.0.1b0__py3-none-any.whl

union/_internal/imagebuild/docker_builder.py

@@ -0,0 +1,416 @@
+import asyncio
+import shutil
+import subprocess
+import tempfile
+from pathlib import Path
+from string import Template
+from typing import ClassVar, Protocol
+
+import aiofiles
+import click
+
+from union._image import (
+    AptPackages,
+    Commands,
+    CopyConfig,
+    Env,
+    Image,
+    Layer,
+    PipPackages,
+    Requirements,
+    UVProject,
+    WorkDir,
+    _DockerLines,
+)
+from union._logging import logger
+
+_F_IMG_ID = "_F_IMG_ID"
+
+UV_LOCK_INSTALL_TEMPLATE = Template("""\
+WORKDIR /root
+RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv \
+    --mount=type=bind,target=uv.lock,src=uv.lock \
+    --mount=type=bind,target=pyproject.toml,src=pyproject.toml \
+    uv sync $PIP_INSTALL_ARGS
+WORKDIR /
+
+# Update PATH and UV_PYTHON to point to the venv created by uv sync
+ENV PATH="/root/.venv/bin:$$PATH" \
+    VIRTUALENV=/root/.venv \
+    UV_PYTHON=/root/.venv/bin/python
+""")
+
+UV_PACKAGE_INSTALL_COMMAND_TEMPLATE = Template("""\
+RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv \
+    --mount=type=bind,target=requirements_uv.txt,src=requirements_uv.txt \
+    uv pip install --python $$UV_PYTHON $PIP_INSTALL_ARGS
+""")
+
+APT_INSTALL_COMMAND_TEMPLATE = Template("""\
+RUN --mount=type=cache,sharing=locked,mode=0777,target=/var/cache/apt,id=apt \
+    apt-get update && apt-get install -y --no-install-recommends \
+    $APT_PACKAGES
+""")
+
+UV_PYTHON_INSTALL_COMMAND = Template("""\
+RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv \
+    uv pip install $PIP_INSTALL_ARGS
+""")
+
+# uv pip install --python /root/env/bin/python
+# new template
+DOCKER_FILE_UV_BASE_TEMPLATE = Template("""\
+#syntax=docker/dockerfile:1.5
+FROM ghcr.io/astral-sh/uv:0.6.12 as uv
+FROM $BASE_IMAGE
+
+USER root
+
+# Copy in uv so that later commands don't have to mount it in
+COPY --from=uv /uv /usr/bin/uv
+
+# Configure default envs
+ENV UV_COMPILE_BYTECODE=1 \
+    UV_LINK_MODE=copy \
+    VIRTUALENV=/opt/venv \
+    UV_PYTHON=/opt/venv/bin/python \
+    PATH="/opt/venv/bin:$$PATH"
+
+# Create a virtualenv with the user specified python version
+RUN uv venv $$VIRTUALENV --python=$PYTHON_VERSION
+
+# Adds nvidia just in case it exists
+ENV PATH="$$PATH:/usr/local/nvidia/bin:/usr/local/cuda/bin" \
+    LD_LIBRARY_PATH="/usr/local/nvidia/lib64:$$LD_LIBRARY_PATH"
+""")
+
+# This gets added on to the end of the dockerfile
+DOCKER_FILE_BASE_FOOTER = Template("""\
+ENV _F_IMG_ID=$F_IMG_ID
+WORKDIR /root
+SHELL ["/bin/bash", "-c"]
+""")
+
+
+class Handler(Protocol):
+    @staticmethod
+    async def handle(layer: Layer, context_path: Path, dockerfile: str) -> str: ...
+
+
+class PipAndRequirementsHandler:
+    @staticmethod
+    async def handle(layer: PipPackages, context_path: Path, dockerfile: str) -> str:
+        if isinstance(layer, Requirements):
+            async with aiofiles.open(layer.file) as f:
+                requirements = []
+                async for line in f:
+                    requirement = await line
+                    requirements.append(requirement.strip())
+        else:
+            requirements = layer.packages
+        requirements_uv_path = context_path / "requirements_uv.txt"
+        async with aiofiles.open(requirements_uv_path, "w") as f:
+            reqs = "\n".join(requirements)
+            await f.write(reqs)
+
+        pip_install_args = []
+        if layer.index_url:
+            pip_install_args.append(f"--index-url {layer.index_url}")
+
+        if layer.extra_index_urls:
+            pip_install_args.extend([f"--extra-index-url {url}" for url in layer.extra_index_urls])
+
+        if layer.pre:
+            pip_install_args.append("--pre")
+
+        if layer.extra_args:
+            pip_install_args.append(layer.extra_args)
+
+        pip_install_args.extend(["--requirement", "requirements_uv.txt"])
+
+        delta = UV_PACKAGE_INSTALL_COMMAND_TEMPLATE.substitute(PIP_INSTALL_ARGS=" ".join(pip_install_args))
+        dockerfile += delta
+
+        return dockerfile
+
+
+class _DockerLinesHandler:
+    @staticmethod
+    async def handle(layer: _DockerLines, context_path: Path, dockerfile: str) -> str:
+        # Add the lines to the dockerfile
+        for line in layer.lines:
+            dockerfile += f"\n{line}\n"
+
+        return dockerfile
+
+
+class EnvHandler:
+    @staticmethod
+    async def handle(layer: Env, context_path: Path, dockerfile: str) -> str:
+        # Add the env vars to the dockerfile
+        for key, value in layer.env_vars:
+            dockerfile += f"\nENV {key}={value}\n"
+
+        return dockerfile
+
+
+class AptPackagesHandler:
+    @staticmethod
+    async def handle(layer: AptPackages, context_path: Path, dockerfile: str) -> str:
+        layer = layer.packages
+        delta = APT_INSTALL_COMMAND_TEMPLATE.substitute(APT_PACKAGES=" ".join(layer))
+        dockerfile += delta
+
+        return dockerfile
+
+
+class UVProjectHandler:
+    @staticmethod
+    async def handle(layer: UVProject, context_path: Path, dockerfile: str) -> str:
+        # copy the two files
+        shutil.copy(layer.pyproject, context_path)
+        shutil.copy(layer.uvlock, context_path)
+
+        # --locked: Assert that the `uv.lock` will remain unchanged
+        # --no-dev: Omit the development dependency group
+        # --no-install-project: Do not install the current project
+        additional_pip_install_args = ["--locked", "--no-dev", "--no-install-project"]
+        delta = UV_LOCK_INSTALL_TEMPLATE.substitute(PIP_INSTALL_ARGS=" ".join(additional_pip_install_args))
+        dockerfile += delta
+
+        return dockerfile
+
+
+class CopyConfigHandler:
+    @staticmethod
+    async def handle(layer: CopyConfig, context_path: Path, dockerfile: str) -> str:
+        # Copy the source config file or directory to the context path
+        abs_path = layer.context_source.absolute()
+        dest_path = context_path / abs_path.name
+        image_dest_path = layer.image_dest + "/" + abs_path.name
+        if layer.context_source.is_file():
+            # Copy the file
+            shutil.copy(abs_path, dest_path)
+        elif layer.context_source.is_dir():
+            # Copy the entire directory
+            shutil.copytree(abs_path, dest_path)
+        else:
+            raise ValueError(f"Source path is neither file nor directory: {layer.context_source}")
+
+        # Add a copy command to the dockerfile
+        dockerfile += f"\nCOPY {abs_path.name} {image_dest_path}\n"
+
+        return dockerfile
+
+
+class CommandsHandler:
+    @staticmethod
+    async def handle(layer: Commands, context_path: Path, dockerfile: str) -> str:
+        # Append raw commands to the dockerfile
+        for command in layer.commands:
+            dockerfile += f"\nRUN {command}\n"
+
+        return dockerfile
+
+
+class WorkDirHandler:
+    @staticmethod
+    async def handle(layer: WorkDir, context_path: Path, dockerfile: str) -> str:
+        # cd to the workdir
+        dockerfile += f"\nWORKDIR {layer.workdir}\n"
+
+        return dockerfile
+
+
+async def _process_layer(layer: Layer, context_path: Path, dockerfile: str) -> str:
+    match layer:
+        case Requirements() | PipPackages():
+            # Handle pip packages and requirements
+            dockerfile = await PipAndRequirementsHandler.handle(layer, context_path, dockerfile)
+
+        case AptPackages():
+            # Handle apt packages
+            dockerfile = await AptPackagesHandler.handle(layer, context_path, dockerfile)
+
+        case UVProject():
+            # Handle UV project
+            dockerfile = await UVProjectHandler.handle(layer, context_path, dockerfile)
+
+        case CopyConfig():
+            # Handle local files and folders
+            dockerfile = await CopyConfigHandler.handle(layer, context_path, dockerfile)
+
+        case Commands():
+            # Handle commands
+            dockerfile = await CommandsHandler.handle(layer, context_path, dockerfile)
+
+        case WorkDir():
+            # Handle workdir
+            dockerfile = await WorkDirHandler.handle(layer, context_path, dockerfile)
+
+        case Env():
+            # Handle environment variables
+            dockerfile = await EnvHandler.handle(layer, context_path, dockerfile)
+
+        case _DockerLines():
+            # Only for internal use
+            dockerfile = await _DockerLinesHandler.handle(layer, context_path, dockerfile)
+
+        case _:
+            raise NotImplementedError(f"Layer type {type(layer)} not supported")
+
+    return dockerfile
+
+
+class DockerImageBuilder:
+    """Image builder using Docker and buildkit."""
+
+    builder_type: ClassVar = "docker"
+    _builder_name: ClassVar = "flytex"
+
+    async def build_image(self, image: Image, dry_run: bool = False) -> str:
+        if image.is_final:
+            if image._layers:
+                raise ValueError("Image is a default image and should already be built")
+
+        if image.dockerfile:
+            # If a dockerfile is provided, use it directly
+            return await self._build_from_dockerfile(image, push=True)
+
+        return await self._build_image(
+            image,
+            push=True,
+            dry_run=dry_run,
+        )
+
+    async def _build_from_dockerfile(self, image: Image, push: bool) -> str:
+        """
+        Build the image from a provided Dockerfile.
+        """
+        command = [
+            "docker",
+            "build",
+            "--tag",
+            f"{image.uri}",
+            "--platform",
+            ",".join(image.platform),
+            ".",
+        ]
+
+        if image.registry and push:
+            command.append("--push")
+
+        concat_command = " ".join(command)
+        logger.debug(f"Build command: {concat_command}")
+        click.secho(f"Run command: {concat_command} ", fg="blue")
+
+        await asyncio.to_thread(subprocess.run, command, cwd=str(image.dockerfile.cwd()), check=True)
+
+        return image.uri
+
+    @staticmethod
+    async def _ensure_buildx_builder():
+        """Ensure there is a docker buildx builder called flyte"""
+        # Check if buildx is available
+        try:
+            await asyncio.to_thread(
+                subprocess.run, ["docker", "buildx", "version"], check=True, stdout=subprocess.DEVNULL
+            )
+        except subprocess.CalledProcessError:
+            raise RuntimeError("Docker buildx is not available. Make sure BuildKit is installed and enabled.")
+
+        # List builders
+        result = await asyncio.to_thread(
+            subprocess.run, ["docker", "buildx", "ls"], capture_output=True, text=True, check=True
+        )
+        builders = result.stdout
+
+        # Check if there's any usable builder
+        if DockerImageBuilder._builder_name not in builders:
+            # No default builder found, create one
+            logger.info("No buildx builder found, creating one...")
+            await asyncio.to_thread(
+                subprocess.run,
+                [
+                    "docker",
+                    "buildx",
+                    "create",
+                    "--name",
+                    DockerImageBuilder._builder_name,
+                    "--platform",
+                    "linux/amd64,linux/arm64",
+                ],
+                check=True,
+            )
+        else:
+            logger.info("Buildx builder already exists.")
+
+    async def _build_image(self, image: Image, *, push: bool = True, dry_run: bool = False) -> str:
+        """
+        if default image (only base image and locked), raise an error, don't have a dockerfile
+        if dockerfile, just build
+        in the main case, get the default Dockerfile template
+          - start from the base image
+          - use python to create a default venv and export variables
+
+          Then for the layers
+          - for each layer
+            - find the appropriate layer handler
+            - call layer handler with the context dir and the dockerfile
+              - handler can choose to do something (copy files from local) to the context and update the dockerfile
+                contents, returning the new string
+        """
+        # For testing, set `push=False` to just build the image locally and not push to
+        # registry.
+
+        await DockerImageBuilder._ensure_buildx_builder()
+
+        with tempfile.TemporaryDirectory(delete=False) as tmp_dir:
+            logger.warning(f"Temporary directory: {tmp_dir}")
+            tmp_path = Path(tmp_dir)
+
+            dockerfile = DOCKER_FILE_UV_BASE_TEMPLATE.substitute(
+                BASE_IMAGE=image.base_image,
+                PYTHON_VERSION=f"{image.python_version[0]}.{image.python_version[1]}",
+            )
+
+            for layer in image._layers:
+                dockerfile = await _process_layer(layer, tmp_path, dockerfile)
+
+            dockerfile += DOCKER_FILE_BASE_FOOTER.substitute(F_IMG_ID=image.uri)
+
+            dockerfile_path = tmp_path / "Dockerfile"
+            async with aiofiles.open(dockerfile_path, mode="w") as f:
+                await f.write(dockerfile)
+
+            command = [
+                "docker",
+                "buildx",
+                "build",
+                "--builder",
+                DockerImageBuilder._builder_name,
+                "--tag",
+                f"{image.uri}",
+                "--platform",
+                ",".join(image.platform),
+                "--push" if push else "--load",
+            ]
+
+            if image.registry and push:
+                command.append("--push")
+            command.append(tmp_dir)
+
+            concat_command = " ".join(command)
+            logger.debug(f"Build command: {concat_command}")
+            if dry_run:
+                click.secho("Dry run for docker builder...")
+                click.secho(f"Context path: {tmp_path}")
+                click.secho(f"Dockerfile: {dockerfile}")
+                click.secho(f"Command: {concat_command}")
+                return image.uri
+            else:
+                click.secho(f"Run command: {concat_command} ", fg="blue")
+
+            await asyncio.to_thread(subprocess.run, command, check=True)
+
+            return image.uri

union/_internal/imagebuild/image_builder.py

@@ -0,0 +1,243 @@
+from __future__ import annotations
+
+import asyncio
+import json
+import typing
+from typing import ClassVar, Dict, Optional, Tuple
+
+from async_lru import alru_cache
+from pydantic import BaseModel
+from typing_extensions import Protocol
+
+from union._image import Architecture, Image
+from union._logging import logger
+
+
+class ImageBuilder(Protocol):
+    async def build_image(self, image: Image, dry_run: bool) -> Optional[str]: ...
+
+
+class ImageChecker(Protocol):
+    @classmethod
+    async def image_exists(
+        cls, repository: str, tag: str, arch: Tuple[Architecture, ...] = ("linux/amd64",)
+    ) -> bool: ...
+
+
+class DockerAPIImageChecker(ImageChecker):
+    """
+    Unfortunately only works for docker hub as there's no way to get a public token for ghcr.io. See SO:
+    https://stackoverflow.com/questions/57316115/get-manifest-of-a-public-docker-image-hosted-on-docker-hub-using-the-docker-regi
+    The token used here seems to be short-lived (<1 second), so copy pasting doesn't even work.
+    """
+
+    @classmethod
+    async def image_exists(cls, repository: str, tag: str, arch: Tuple[Architecture, ...] = ("linux/amd64",)) -> bool:
+        import httpx
+
+        if "/" in repository:
+            if not repository.startswith("library/"):
+                raise ValueError("This checker only works with Docker Hub")
+        else:
+            repository = f"library/{repository}"
+
+        auth_url = "https://auth.docker.io/token"
+        service = "registry.docker.io"
+        scope = f"repository:{repository}:pull"
+
+        async with httpx.AsyncClient() as client:
+            # Get auth token
+            auth_response = await client.get(auth_url, params={"service": service, "scope": scope})
+            if auth_response.status_code != 200:
+                raise Exception(f"Failed to get auth token: {auth_response.status_code}")
+            token = auth_response.json()["token"]
+
+            manifest_url = f"https://registry-1.docker.io/v2/{repository}/manifests/{tag}"
+            headers = {
+                "Authorization": f"Bearer {token}",
+                "Accept": (
+                    "application/vnd.docker.distribution.manifest.v2+json,"
+                    "application/vnd.docker.distribution.manifest.list.v2+json"
+                ),
+            }
+            manifest_response = await client.get(manifest_url, headers=headers)
+
+            if manifest_response.status_code != 200:
+                raise Exception(f"Failed to get manifest: {manifest_response.status_code}")
+            manifest_list = manifest_response.json()["manifests"]
+            architectures = [f"{x['platform']['os']}/{x['platform']['architecture']}" for x in manifest_list]
+            architectures = set(architectures)
+
+            if architectures >= set(arch):
+                logger.debug(f"Image {repository}:{tag} found for architecture(s) {arch}, has {architectures}")
+                return True
+            else:
+                logger.debug(f"Image {repository}:{tag} not found for architecture(s) {arch}, only has {architectures}")
+                return False
+
+
+class LocalDockerCommandImageChecker(ImageChecker):
+    command_name: ClassVar[str] = "docker"
+
+    @classmethod
+    async def image_exists(cls, repository: str, tag: str, arch: Tuple[Architecture, ...] = ("linux/amd64",)) -> bool:
+        # Check if the image exists locally by running the docker inspect command
+        process = await asyncio.create_subprocess_exec(
+            cls.command_name,
+            "manifest",
+            "inspect",
+            f"{repository}:{tag}",
+            stdout=asyncio.subprocess.PIPE,
+            stderr=asyncio.subprocess.PIPE,
+        )
+        stdout, stderr = await process.communicate()
+        if stderr and "manifest unknown" in stderr.decode():
+            logger.debug(f"Image {repository}:{tag} not found using the docker command.")
+            return False
+
+        if process.returncode != 0:
+            raise RuntimeError(f"Failed to run docker image inspect {repository}:{tag}")
+
+        inspect_data = json.loads(stdout.decode())
+        if "manifests" not in inspect_data:
+            raise RuntimeError(f"Invalid data returned from docker image inspect for {repository}:{tag}")
+        manifest_list = inspect_data["manifests"]
+        architectures = [f"{x['platform']['os']}/{x['platform']['architecture']}" for x in manifest_list]
+        architectures = set(architectures)
+        if architectures >= set(arch):
+            logger.debug(f"Image {repository}:{tag} found for architecture(s) {arch}, has {architectures}")
+            return True
+
+        # Otherwise write a message and return false to trigger build
+        logger.debug(f"Image {repository}:{tag} not found for architecture(s) {arch}, only has {architectures}")
+        return False
+
+
+class LocalPodmanCommandImageChecker(LocalDockerCommandImageChecker):
+    command_name: ClassVar[str] = "podman"
+
+
+class ImageBuildEngine:
+    """
+    ImageBuildEngine contains a list of builders that can be used to build an ImageSpec.
+    """
+
+    _REGISTRY: typing.ClassVar[typing.Dict[str, Tuple[ImageBuilder, int]]] = {}
+    _SEEN_IMAGES: typing.ClassVar[typing.Dict[str, str]] = {
+        # Set default for the auto container. See Image._identifier_override for more info.
+        "auto": Image.auto().uri,
+    }
+
+    @classmethod
+    def register(cls, builder_type: str, image_builder: ImageBuilder, priority: int = 5):
+        cls._REGISTRY[builder_type] = (image_builder, priority)
+
+    @classmethod
+    def get_registry(cls) -> Dict[str, Tuple[ImageBuilder, int]]:
+        return cls._REGISTRY
+
+    @staticmethod
+    @alru_cache
+    async def image_exists(image: Image) -> bool:
+        if image.base_image is not None and not image._layers:
+            logger.debug(f"Image {image} has a base image: {image.base_image} and no layers. Skip existence check.")
+            return True
+        assert image.registry is not None, f"Image registry is not set for {image}"
+        assert image.name is not None, f"Image name is not set for {image}"
+
+        repository = image.registry + "/" + image.name
+        tag = image._final_tag
+
+        if tag == "latest":
+            logger.debug(f"Image {image} has tag 'latest', skip existence check, always build")
+            return True
+
+        # Can get a public token for docker.io but ghcr requires a pat, so harder to get the manifest anonymously.
+        checkers = [LocalDockerCommandImageChecker, LocalPodmanCommandImageChecker, DockerAPIImageChecker]
+        for checker in checkers:
+            try:
+                exists = await checker.image_exists(repository, tag, set(image.platform))
+                logger.debug(f"Image {image} {exists=} in registry")
+                return exists
+            except Exception as e:
+                logger.debug(f"Error checking image existence with {checker.__name__}: {e}")
+                continue
+
+        # If all checkers fail, then assume the image exists. This is current flytekit behavior
+        logger.info(f"All checkers failed to check existence of {image.uri}, assuming it does exists")
+        return True
+
+    @classmethod
+    @alru_cache
+    async def build(
+        cls, image: Image, builder: Optional[str] = None, dry_run: bool = False, force: bool = False
+    ) -> Optional[str]:
+        """
+        Build the image. Images to be tagged with latest will always be built. Otherwise, this engine will check the
+        registry to see if the manifest exists.
+
+        :param image:
+        :param builder:
+        :param dry_run: Tell the builder to not actually build. Different builders will have different behaviors.
+        :param force: Skip the existence check. Normally if the image already exists we won't build it.
+        :return:
+        """
+        # Always trigger a build if this is a dry run since builder shouldn't really do anything, or a force.
+        if force or dry_run or not await cls.image_exists(image):
+            logger.info(f"Image {image.uri} does not exist in registry or force/dry-run, building...")
+
+            # Validate the image before building
+            image.validate()
+
+            # If builder is not specified, use the first registered builder
+            img_builder = ImageBuildEngine._get_builder(builder)
+
+            result = await img_builder.build_image(image, dry_run=dry_run)
+            return result
+        else:
+            logger.info(f"Image {image.uri} already exists in registry. Skipping build.")
+            return image.uri
+
+    @classmethod
+    def _get_builder(cls, builder: Optional[str]) -> ImageBuilder:
+        if not builder:
+            from .docker_builder import DockerImageBuilder
+
+            return DockerImageBuilder()
+        if builder not in cls._REGISTRY:
+            raise AssertionError(f"Image builder {builder} is not registered.")
+        return cls._REGISTRY[builder][0]
+
+
+class ImageCache(BaseModel):
+    image_lookup: Dict[str, str]
+    serialized_form: str | None = None
+
+    @property
+    def to_transport(self) -> str:
+        """
+        :return: returns the serialization context as a base64encoded, gzip compressed, json string
+        """
+        # This is so that downstream tasks continue to have the same image lookup abilities
+        import base64
+        import gzip
+        from io import BytesIO
+
+        if self.serialized_form:
+            return self.serialized_form
+        json_str = self.model_dump_json(exclude={"serialized_form"})
+        buf = BytesIO()
+        with gzip.GzipFile(mode="wb", fileobj=buf, mtime=0) as f:
+            f.write(json_str.encode("utf-8"))
+        return base64.b64encode(buf.getvalue()).decode("utf-8")
+
+    @classmethod
+    def from_transport(cls, s: str) -> ImageCache:
+        import base64
+        import gzip
+
+        compressed_val = base64.b64decode(s.encode("utf-8"))
+        json_str = gzip.decompress(compressed_val).decode("utf-8")
+        val = cls.model_validate_json(json_str)
+        val.serialized_form = s
+        return val

union/_internal/resolvers/_task_module.py

@@ -0,0 +1,31 @@
+import inspect
+import pathlib
+import sys
+from typing import Tuple
+
+from union._task import AsyncFunctionTaskTemplate, TaskTemplate
+
+
+def extract_task_module(task: TaskTemplate) -> Tuple[str, str]:
+    """
+    Extract the task module from the task template.
+
+    :param task: The task template to extract the module from.
+    :return: A tuple containing the entity name, module
+    """
+    # TODO Work with current working directory, to figure out the module name, in case of nested launches
+    entity_name = task.name
+    if isinstance(task, AsyncFunctionTaskTemplate):
+        entity_module = inspect.getmodule(task.func).__name__
+        entity_name = task.func.__name__
+    else:
+        raise NotImplementedError(f"Task module {entity_name} not implemented.")
+
+    if entity_module == "__main__":
+        """
+        This case is for the case in which the task is run from the main module.
+        """
+        main_path = pathlib.Path(sys.modules["__main__"].__file__)
+        entity_module = main_path.stem
+
+    return entity_name, entity_module