flyte 0.2.0b18__py3-none-any.whl → 0.2.0b19__py3-none-any.whl

flyte/_code_bundle/bundle.py

@@ -80,7 +80,7 @@ async def build_pkl_bundle(
             logger.debug("Uploading pickled code bundle to control plane.")
             from flyte.remote import upload_file
 
-            hash_digest, remote_path = await upload_file(dest)
+            hash_digest, remote_path = await upload_file.aio(dest)
             return CodeBundle(pkl=remote_path, computed_version=hash_digest)
 
         elif upload_from_dataplane_base_path:
@@ -145,7 +145,7 @@ async def build_code_bundle(
         bundle_path, tar_size, archive_size = create_bundle(from_dir, pathlib.Path(tmp_dir), files, digest)
         logger.info(f"Code bundle created at {bundle_path}, size: {tar_size} MB, archive size: {archive_size} MB")
         if not dryrun:
-            hash_digest, remote_path = await upload_file(bundle_path)
+            hash_digest, remote_path = await upload_file.aio(bundle_path)
             logger.debug(f"Code bundle uploaded to {remote_path}")
         else:
             remote_path = "na"

flyte/_deploy.py

@@ -133,7 +133,7 @@ async def _build_images(deployment: DeploymentPlan) -> ImageCache:
     image_identifier_map = {}
     for env_name, env in deployment.envs.items():
         if not isinstance(env.image, str):
-            logger.warning(f"Building Image for environment {env_name}, image: {env.image}")
+            logger.debug(f"Building Image for environment {env_name}, image: {env.image}")
             images.append(_build_image_bg(env_name, env.image))
 
         elif env.image == "auto" and "auto" not in image_identifier_map:
@@ -145,7 +145,7 @@ async def _build_images(deployment: DeploymentPlan) -> ImageCache:
         logger.warning(f"Built Image for environment {env_name}, image: {image_uri}")
         env = deployment.envs[env_name]
         if isinstance(env.image, Image):
-            image_identifier_map[env.image.identifier] = env.image.uri
+            image_identifier_map[env.image.identifier] = image_uri
 
     return ImageCache(image_lookup=image_identifier_map)
 
@@ -158,7 +158,6 @@ async def apply(deployment: DeploymentPlan, copy_style: CopyFiles, dryrun: bool
     image_cache = await _build_images(deployment)
 
     version = deployment.version
-    code_bundle = None
     if copy_style == "none" and not version:
         raise flyte.errors.DeploymentError("Version must be set when copy_style is none")
     else:

flyte/_image.py

@@ -53,7 +53,7 @@ class Layer:
 
         :param hasher: The hash object to update with the layer's data.
         """
-        ...
+        print("hash hash")
 
     def validate(self):
         """
@@ -64,24 +64,32 @@ class Layer:
 
 @rich.repr.auto
 @dataclass(kw_only=True, frozen=True, repr=True)
-class PipPackages(Layer):
-    packages: Optional[Tuple[str, ...]] = None
+class PipOption:
     index_url: Optional[str] = None
     extra_index_urls: Optional[Tuple[str] | Tuple[str, ...] | List[str]] = None
     pre: bool = False
     extra_args: Optional[str] = None
 
-    # todo: to be implemented
-    # secret_mounts: Optional[List[Tuple[str, str]]] = None
+    def get_pip_install_args(self) -> List[str]:
+        pip_install_args = []
+        if self.index_url:
+            pip_install_args.append(f"--index-url {self.index_url}")
+
+        if self.extra_index_urls:
+            pip_install_args.extend([f"--extra-index-url {url}" for url in self.extra_index_urls])
+
+        if self.pre:
+            pip_install_args.append("--pre")
+
+        if self.extra_args:
+            pip_install_args.append(self.extra_args)
+        return pip_install_args
 
     def update_hash(self, hasher: hashlib._Hash):
         """
-        Update the hash with the pip packages
+        Update the hash with the PipOption
         """
         hash_input = ""
-        if self.packages:
-            for package in self.packages:
-                hash_input += package
         if self.index_url:
             hash_input += self.index_url
         if self.extra_index_urls:
@@ -95,6 +103,44 @@ class PipPackages(Layer):
         hasher.update(hash_input.encode("utf-8"))
 
 
+@rich.repr.auto
+@dataclass(kw_only=True, frozen=True, repr=True)
+class PipPackages(PipOption, Layer):
+    packages: Optional[Tuple[str, ...]] = None
+
+    # todo: to be implemented
+    # secret_mounts: Optional[List[Tuple[str, str]]] = None
+
+    def update_hash(self, hasher: hashlib._Hash):
+        """
+        Update the hash with the pip packages
+        """
+        super().update_hash(hasher)
+        hash_input = ""
+        if self.packages:
+            for package in self.packages:
+                hash_input += package
+
+        hasher.update(hash_input.encode("utf-8"))
+
+
+@rich.repr.auto
+@dataclass(kw_only=True, frozen=True, repr=True)
+class PythonWheels(PipOption, Layer):
+    wheel_dir: Path
+
+    def update_hash(self, hasher: hashlib._Hash):
+        super().update_hash(hasher)
+        from ._utils import filehash_update
+
+        # Iterate through all the wheel files in the directory and update the hash
+        for wheel_file in self.wheel_dir.glob("*.whl"):
+            if not wheel_file.is_file():
+                # Skip if it's not a file (e.g., directory or symlink)
+                continue
+            filehash_update(wheel_file, hasher)
+
+
 @rich.repr.auto
 @dataclass(kw_only=True, frozen=True, repr=True)
 class Requirements(PipPackages):
@@ -107,6 +153,19 @@ class Requirements(PipPackages):
         filehash_update(self.file, hasher)
 
 
+@rich.repr.auto
+@dataclass(frozen=True, repr=True)
+class UVProject(PipOption, Layer):
+    pyproject: Path
+    uvlock: Path
+
+    def update_hash(self, hasher: hashlib._Hash):
+        from ._utils import filehash_update
+
+        super().update_hash(hasher)
+        filehash_update(self.uvlock, hasher)
+
+
 @rich.repr.auto
 @dataclass(frozen=True, repr=True)
 class AptPackages(Layer):
@@ -138,35 +197,23 @@ class WorkDir(Layer):
 @dataclass(frozen=True, repr=True)
 class CopyConfig(Layer):
     path_type: CopyConfigType
-    context_source: Path
-    image_dest: str = "."
+    src: Path
+    dst: str = "."
 
     def validate(self):
-        if not self.context_source.exists():
-            raise ValueError(f"Source folder {self.context_source.absolute()} does not exist")
-        if not self.context_source.is_dir() and self.path_type == 1:
-            raise ValueError(f"Source folder {self.context_source.absolute()} is not a directory")
-        if not self.context_source.is_file() and self.path_type == 0:
-            raise ValueError(f"Source file {self.context_source.absolute()} is not a file")
+        if not self.src.exists():
+            raise ValueError(f"Source folder {self.src.absolute()} does not exist")
+        if not self.src.is_dir() and self.path_type == 1:
+            raise ValueError(f"Source folder {self.src.absolute()} is not a directory")
+        if not self.src.is_file() and self.path_type == 0:
+            raise ValueError(f"Source file {self.src.absolute()} is not a file")
 
     def update_hash(self, hasher: hashlib._Hash):
         from ._utils import update_hasher_for_source
 
-        update_hasher_for_source(self.context_source, hasher)
-        if self.image_dest:
-            hasher.update(self.image_dest.encode("utf-8"))
-
-
-@rich.repr.auto
-@dataclass(frozen=True, repr=True)
-class UVProject(Layer):
-    pyproject: Path
-    uvlock: Path
-
-    def update_hash(self, hasher: hashlib._Hash):
-        from ._utils import filehash_update
-
-        filehash_update(self.uvlock, hasher)
+        update_hasher_for_source(self.src, hasher)
+        if self.dst:
+            hasher.update(self.dst.encode("utf-8"))
 
 
 @rich.repr.auto
@@ -384,8 +431,6 @@ class Image:
             python_version = _detect_python_version()
 
         base_image = cls._get_default_image_for(python_version=python_version, flyte_version=flyte_version)
-        if name is not None and registry is None:
-            raise ValueError("Both name and registry must be specified to override the default image name.")
 
         if registry and name:
             return base_image.clone(registry=registry, name=name)
@@ -436,6 +481,8 @@ class Image:
 
         :param name: name of the image
         :param registry: registry to use for the image
+        :param python_version: Python version to use for the image, if not specified, will use the current Python
+        version
         :param script: path to the uv script
         :param arch: architecture to use for the image, default is linux/amd64, use tuple for multiple values
         :param python_version: Python version for the image, if not specified, will use the current Python version
@@ -470,13 +517,17 @@ class Image:
         return img
 
     def clone(
-        self, registry: Optional[str] = None, name: Optional[str] = None, addl_layer: Optional[Layer] = None
+        self,
+        registry: Optional[str] = None,
+        name: Optional[str] = None,
+        addl_layer: Optional[Layer] = None,
     ) -> Image:
         """
         Use this method to clone the current image and change the registry and name
 
         :param registry: Registry to use for the image
         :param name: Name of the image
+        :param addl_layer: Additional layer to add to the image. This will be added to the end of the layers.
         :param addl_layer: Additional layer to add to the image. This will be added on top of the existing layers.
 
         :return:
@@ -566,7 +617,7 @@ class Image:
         new_image = self.clone(addl_layer=WorkDir(workdir=workdir))
         return new_image
 
-    def with_requirements(self, file: Path) -> Image:
+    def with_requirements(self, file: str | Path) -> Image:
         """
         Use this method to create a new image with the specified requirements file layered on top of the current image
         Cannot be used in conjunction with conda
@@ -574,10 +625,8 @@ class Image:
         :param file: path to the requirements file, must be a .txt file
         :return:
         """
-        if not file.exists():
-            raise FileNotFoundError(f"Requirements file {file} does not exist")
-        if not file.is_file():
-            raise ValueError(f"Requirements file {file} is not a file")
+        if isinstance(file, str):
+            file = Path(file)
         if file.suffix != ".txt":
             raise ValueError(f"Requirements file {file} must have a .txt extension")
         new_image = self.clone(addl_layer=Requirements(file=file))
@@ -638,30 +687,28 @@ class Image:
         new_image = self.clone(addl_layer=Env.from_dict(env_vars))
         return new_image
 
-    def with_source_folder(self, context_source: Path, image_dest: Optional[str] = None) -> Image:
+    def with_source_folder(self, src: Path, dst: str = ".") -> Image:
         """
         Use this method to create a new image with the specified local directory layered on top of the current image.
         If dest is not specified, it will be copied to the working directory of the image
 
-        :param context_source: root folder of the source code from the build context to be copied
-        :param image_dest: destination folder in the image
+        :param src: root folder of the source code from the build context to be copied
+        :param dst: destination folder in the image
         :return: Image
         """
-        image_dest = image_dest if image_dest else "."
-        new_image = self.clone(addl_layer=CopyConfig(path_type=1, context_source=context_source, image_dest=image_dest))
+        new_image = self.clone(addl_layer=CopyConfig(path_type=1, src=src, dst=dst))
         return new_image
 
-    def with_source_file(self, context_source: Path, image_dest: Optional[str] = None) -> Image:
+    def with_source_file(self, src: Path, dst: str = ".") -> Image:
         """
         Use this method to create a new image with the specified local file layered on top of the current image.
         If dest is not specified, it will be copied to the working directory of the image
 
-        :param context_source: root folder of the source code from the build context to be copied
-        :param image_dest: destination folder in the image
+        :param src: root folder of the source code from the build context to be copied
+        :param dst: destination folder in the image
         :return: Image
         """
-        image_dest = image_dest if image_dest else "."
-        new_image = self.clone(addl_layer=CopyConfig(path_type=0, context_source=context_source, image_dest=image_dest))
+        new_image = self.clone(addl_layer=CopyConfig(path_type=0, src=src, dst=dst))
         return new_image
 
     def with_uv_project(self, pyproject_file: Path) -> Image:
@@ -714,22 +761,11 @@ class Image:
         :return: Image
         """
         dist_folder = Path(__file__).parent.parent.parent / "dist"
-        # Manually declare the CopyConfig instead of using with_source_folder so we can set the hashing
+        # Manually declare the PythonWheel so we can set the hashing
         # used to compute the identifier. Can remove if we ever decide to expose the lambda in with_ commands
-        with_dist = self.clone(
-            addl_layer=CopyConfig(
-                path_type=1, context_source=dist_folder, image_dest=".", _compute_identifier=lambda x: "/dist"
-            )
-        )
+        with_dist = self.clone(addl_layer=PythonWheels(wheel_dir=dist_folder, _compute_identifier=lambda x: "/dist"))
 
-        return with_dist.with_commands(
-            [
-                "--mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv"
-                " --mount=from=uv,source=/uv,target=/usr/bin/uv"
-                " --mount=source=dist,target=/dist,type=bind"
-                " uv pip install --python $VIRTUALENV $(ls /dist/*whl)"
-            ]
-        )
+        return with_dist
 
     def __img_str__(self) -> str:
         """

flyte/_initialize.py

@@ -14,6 +14,7 @@ from ._logging import initialize_logger, logger
 from ._tools import ipython_check
 
 if TYPE_CHECKING:
+    from flyte._internal.imagebuild import ImageBuildEngine
     from flyte.config import Config
     from flyte.remote._client.auth import AuthType, ClientConfig
     from flyte.remote._client.controlplane import ClientSet
@@ -39,6 +40,7 @@ class CommonInit:
 class _InitConfig(CommonInit):
     client: Optional[ClientSet] = None
     storage: Optional[Storage] = None
+    image_builder: "ImageBuildEngine.ImageBuilderType" = "local"
 
     def replace(self, **kwargs) -> _InitConfig:
         return replace(self, **kwargs)
@@ -132,6 +134,7 @@ async def init(
     http_proxy_url: str | None = None,
     storage: Storage | None = None,
     batch_size: int = 1000,
+    image_builder: ImageBuildEngine.ImageBuilderType = "local",
 ) -> None:
     """
     Initialize the Flyte system with the given configuration. This method should be called before any other Flyte
@@ -166,6 +169,7 @@ async def init(
     :param org: Optional organization override for the client. Should be set by auth instead.
     :param batch_size: Optional batch size for operations that use listings, defaults to 1000, so limit larger than
       batch_size will be split into multiple requests.
+    :param image_builder: Optional image builder configuration, if not provided, the default image builder will be used.
 
     :return: None
     """
@@ -210,12 +214,15 @@ async def init(
             storage=storage,
             org=org or org_from_endpoint(endpoint),
             batch_size=batch_size,
+            image_builder=image_builder,
         )
 
 
 @syncify
 async def init_from_config(
-    path_or_config: str | Config | None = None, root_dir: Path | None = None, log_level: int | None = None
+    path_or_config: str | Config | None = None,
+    root_dir: Path | None = None,
+    log_level: int | None = None,
 ) -> None:
     """
     Initialize the Flyte system using a configuration file or Config object. This method should be called before any
@@ -266,6 +273,7 @@ async def init_from_config(
         client_credentials_secret=cfg.platform.client_credentials_secret,
         root_dir=root_dir,
         log_level=log_level,
+        image_builder=cfg.image.builder,
     )
 
 

flyte/_internal/imagebuild/__init__.py

@@ -3,6 +3,10 @@ from typing import List
 
 from flyte._image import Image
 from flyte._internal.imagebuild.docker_builder import DockerImageBuilder
+from flyte._internal.imagebuild.image_builder import ImageBuildEngine
+from flyte._internal.imagebuild.remote_builder import RemoteImageBuilder
+
+__all__ = ["DockerImageBuilder", "ImageBuildEngine", "RemoteImageBuilder"]
 
 
 async def build(images: List[Image]) -> List[str]:

flyte/_internal/imagebuild/docker_builder.py

@@ -3,9 +3,10 @@ import os
 import shutil
 import subprocess
 import tempfile
+import typing
 from pathlib import Path
 from string import Template
-from typing import ClassVar, Protocol, cast
+from typing import ClassVar, Optional, Protocol, cast
 
 import aiofiles
 import click
@@ -18,11 +19,19 @@ from flyte._image import (
     Image,
     Layer,
     PipPackages,
+    PythonWheels,
     Requirements,
     UVProject,
     WorkDir,
     _DockerLines,
 )
+from flyte._internal.imagebuild.image_builder import (
+    DockerAPIImageChecker,
+    ImageBuilder,
+    ImageChecker,
+    LocalDockerCommandImageChecker,
+    LocalPodmanCommandImageChecker,
+)
 from flyte._logging import logger
 
 _F_IMG_ID = "_F_IMG_ID"
@@ -49,6 +58,12 @@ RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=uv \
     uv pip install --prerelease=allow --python $$UV_PYTHON $PIP_INSTALL_ARGS
 """)
 
+UV_WHEEL_INSTALL_COMMAND_TEMPLATE = Template("""\
+RUN --mount=type=cache,sharing=locked,mode=0777,target=/root/.cache/uv,id=wheel \
+    --mount=source=/dist,target=/dist,type=bind \
+    uv pip install --prerelease=allow --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 \
@@ -104,10 +119,15 @@ class PipAndRequirementsHandler:
     @staticmethod
     async def handle(layer: PipPackages, context_path: Path, dockerfile: str) -> str:
         if isinstance(layer, Requirements):
+            if not layer.file.exists():
+                raise FileNotFoundError(f"Requirements file {layer.file} does not exist")
+            if not layer.file.is_file():
+                raise ValueError(f"Requirements file {layer.file} is not a file")
+
             async with aiofiles.open(layer.file) as f:
                 requirements = []
                 async for line in f:
-                    requirement = await line
+                    requirement = line
                     requirements.append(requirement.strip())
         else:
             requirements = list(layer.packages) if layer.packages else []
@@ -116,22 +136,23 @@ class PipAndRequirementsHandler:
             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}")
+        pip_install_args = layer.get_pip_install_args()
+        pip_install_args.extend(["--requirement", "requirements_uv.txt"])
 
-        if layer.extra_index_urls:
-            pip_install_args.extend([f"--extra-index-url {url}" for url in layer.extra_index_urls])
+        delta = UV_PACKAGE_INSTALL_COMMAND_TEMPLATE.substitute(PIP_INSTALL_ARGS=" ".join(pip_install_args))
+        dockerfile += delta
 
-        if layer.pre:
-            pip_install_args.append("--pre")
+        return dockerfile
 
-        if layer.extra_args:
-            pip_install_args.append(layer.extra_args)
 
-        pip_install_args.extend(["--requirement", "requirements_uv.txt"])
+class PythonWheelHandler:
+    @staticmethod
+    async def handle(layer: PythonWheels, context_path: Path, dockerfile: str) -> str:
+        shutil.copytree(layer.wheel_dir, context_path / "dist", dirs_exist_ok=True)
+        pip_install_args = layer.get_pip_install_args()
+        pip_install_args.extend(["/dist/*.whl"])
 
-        delta = UV_PACKAGE_INSTALL_COMMAND_TEMPLATE.substitute(PIP_INSTALL_ARGS=" ".join(pip_install_args))
+        delta = UV_WHEEL_INSTALL_COMMAND_TEMPLATE.substitute(PIP_INSTALL_ARGS=" ".join(pip_install_args))
         dockerfile += delta
 
         return dockerfile
@@ -188,27 +209,31 @@ 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():
+        if layer.src.is_absolute() or ".." in str(layer.src):
+            dst_path = context_path / str(layer.src.absolute()).replace("/", "./_flyte_abs_context/")
+        else:
+            dst_path = context_path / layer.src
+
+        dst_path.parent.mkdir(parents=True, exist_ok=True)
+        abs_path = layer.src.absolute()
+        if layer.src.is_file():
             # Copy the file
-            shutil.copy(abs_path, dest_path)
-        elif layer.context_source.is_dir():
+            shutil.copy(abs_path, dst_path)
+        elif layer.src.is_dir():
             # Copy the entire directory
-            shutil.copytree(abs_path, dest_path, dirs_exist_ok=True)
+            shutil.copytree(abs_path, dst_path, dirs_exist_ok=True)
         else:
-            raise ValueError(f"Source path is neither file nor directory: {layer.context_source}")
+            raise ValueError(f"Source path is neither file nor directory: {layer.src}")
 
         # Add a copy command to the dockerfile
-        dockerfile += f"\nCOPY {abs_path.name} {image_dest_path}\n"
+        dockerfile += f"\nCOPY {dst_path.relative_to(context_path)} {layer.dst}\n"
 
         return dockerfile
 
 
 class CommandsHandler:
     @staticmethod
-    async def handle(layer: Commands, context_path: Path, dockerfile: str) -> str:
+    async def handle(layer: Commands, _: Path, dockerfile: str) -> str:
         # Append raw commands to the dockerfile
         for command in layer.commands:
             dockerfile += f"\nRUN {command}\n"
@@ -227,6 +252,10 @@ class WorkDirHandler:
 
 async def _process_layer(layer: Layer, context_path: Path, dockerfile: str) -> str:
     match layer:
+        case PythonWheels():
+            # Handle Python wheels
+            dockerfile = await PythonWheelHandler.handle(layer, context_path, dockerfile)
+
         case Requirements() | PipPackages():
             # Handle pip packages and requirements
             dockerfile = await PipAndRequirementsHandler.handle(layer, context_path, dockerfile)
@@ -265,12 +294,16 @@ async def _process_layer(layer: Layer, context_path: Path, dockerfile: str) -> s
     return dockerfile
 
 
-class DockerImageBuilder:
+class DockerImageBuilder(ImageBuilder):
     """Image builder using Docker and buildkit."""
 
     builder_type: ClassVar = "docker"
     _builder_name: ClassVar = "flytex"
 
+    def get_checkers(self) -> Optional[typing.List[typing.Type[ImageChecker]]]:
+        # Can get a public token for docker.io but ghcr requires a pat, so harder to get the manifest anonymously
+        return [LocalDockerCommandImageChecker, LocalPodmanCommandImageChecker, DockerAPIImageChecker]
+
     async def build_image(self, image: Image, dry_run: bool = False) -> str:
         if len(image._layers) == 0:
             logger.warning("No layers to build, returning the image URI as is.")