flyte 0.2.0b1__py3-none-any.whl → 2.0.0b46__py3-none-any.whl

flyte/_internal/imagebuild/image_builder.py

@@ -3,6 +3,7 @@ from __future__ import annotations
 import asyncio
 import json
 import typing
+from importlib.metadata import entry_points
 from typing import ClassVar, Dict, Optional, Tuple
 
 from async_lru import alru_cache
@@ -10,18 +11,26 @@ from pydantic import BaseModel
 from typing_extensions import Protocol
 
 from flyte._image import Architecture, Image
+from flyte._initialize import _get_init_config
 from flyte._logging import logger
 
 
 class ImageBuilder(Protocol):
     async def build_image(self, image: Image, dry_run: bool) -> str: ...
 
+    def get_checkers(self) -> Optional[typing.List[typing.Type[ImageChecker]]]:
+        """
+        Returns ImageCheckers that can be used to check if the image exists in the registry.
+        If None, then use the default checkers.
+        """
+        return None
+
 
 class ImageChecker(Protocol):
     @classmethod
     async def image_exists(
         cls, repository: str, tag: str, arch: Tuple[Architecture, ...] = ("linux/amd64",)
-    ) -> bool: ...
+    ) -> Optional[str]: ...
 
 
 class DockerAPIImageChecker(ImageChecker):
@@ -32,13 +41,12 @@ class DockerAPIImageChecker(ImageChecker):
     """
 
     @classmethod
-    async def image_exists(cls, repository: str, tag: str, arch: Tuple[Architecture, ...] = ("linux/amd64",)) -> bool:
+    async def image_exists(
+        cls, repository: str, tag: str, arch: Tuple[Architecture, ...] = ("linux/amd64",)
+    ) -> Optional[str]:
         import httpx
 
-        if "/" in repository:
-            if not repository.startswith("library/"):
-                raise ValueError("This checker only works with Docker Hub")
-        else:
+        if "/" not in repository:
             repository = f"library/{repository}"
 
         auth_url = "https://auth.docker.io/token"
@@ -50,8 +58,10 @@ class DockerAPIImageChecker(ImageChecker):
             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"]
 
+            # ghcr.io/union-oss/flyte:latest
             manifest_url = f"https://registry-1.docker.io/v2/{repository}/manifests/{tag}"
             headers = {
                 "Authorization": f"Bearer {token}",
@@ -60,26 +70,30 @@ class DockerAPIImageChecker(ImageChecker):
                     "application/vnd.docker.distribution.manifest.list.v2+json"
                 ),
             }
-            manifest_response = await client.get(manifest_url, headers=headers)
 
+            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}")
+                logger.warning(f"Image not found: {repository}:{tag} (HTTP {manifest_response.status_code})")
+                return None
+
             manifest_list = manifest_response.json()["manifests"]
-            architectures = [f"{x['platform']['os']}/{x['platform']['architecture']}" for x in manifest_list]
+            architectures = [f"{m['platform']['os']}/{m['platform']['architecture']}" for m in manifest_list]
 
-            if set(architectures) >= set(arch):
-                logger.debug(f"Image {repository}:{tag} found for architecture(s) {arch}, has {architectures}")
-                return True
+            if set(arch).issubset(set(architectures)):
+                logger.debug(f"Image {repository}:{tag} found with arch {architectures}")
+                return f"{repository}:{tag}"
             else:
-                logger.debug(f"Image {repository}:{tag} not found for architecture(s) {arch}, only has {architectures}")
-                return False
+                logger.debug(f"Image {repository}:{tag} has {architectures}, but missing {arch}")
+                return None
 
 
 class LocalDockerCommandImageChecker(ImageChecker):
     command_name: ClassVar[str] = "docker"
 
     @classmethod
-    async def image_exists(cls, repository: str, tag: str, arch: Tuple[Architecture, ...] = ("linux/amd64",)) -> bool:
+    async def image_exists(
+        cls, repository: str, tag: str, arch: Tuple[Architecture, ...] = ("linux/amd64",)
+    ) -> Optional[str]:
         # Check if the image exists locally by running the docker inspect command
         process = await asyncio.create_subprocess_exec(
             cls.command_name,
@@ -90,9 +104,9 @@ class LocalDockerCommandImageChecker(ImageChecker):
             stderr=asyncio.subprocess.PIPE,
         )
         stdout, stderr = await process.communicate()
-        if stderr and "manifest unknown" in stderr.decode():
+        if (stderr and "manifest unknown") or "no such manifest" in stderr.decode():
             logger.debug(f"Image {repository}:{tag} not found using the docker command.")
-            return False
+            return None
 
         if process.returncode != 0:
             raise RuntimeError(f"Failed to run docker image inspect {repository}:{tag}")
@@ -104,11 +118,11 @@ class LocalDockerCommandImageChecker(ImageChecker):
         architectures = [f"{x['platform']['os']}/{x['platform']['architecture']}" for x in manifest_list]
         if set(architectures) >= set(arch):
             logger.debug(f"Image {repository}:{tag} found for architecture(s) {arch}, has {architectures}")
-            return True
+            return f"{repository}:{tag}"
 
         # 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
+        return None
 
 
 class LocalPodmanCommandImageChecker(LocalDockerCommandImageChecker):
@@ -120,55 +134,54 @@ 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
+    ImageBuilderType = typing.Literal["local", "remote"]
 
     @staticmethod
     @alru_cache
-    async def image_exists(image: Image) -> bool:
+    async def image_exists(image: Image) -> Optional[str]:
         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}"
+            return image.uri
         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
+            return image.uri
 
-        # 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:
+        builder = None
+        cfg = _get_init_config()
+        if cfg and cfg.image_builder:
+            builder = cfg.image_builder
+        image_builder = ImageBuildEngine._get_builder(builder)
+        image_checker = image_builder.get_checkers()
+        if image_checker is None:
+            logger.info(f"No image checkers found for builder `{image_builder}`, assuming it exists")
+            return image.uri
+        for checker in image_checker:
             try:
-                exists = await checker.image_exists(repository, tag, tuple(image.platform))
-                logger.debug(f"Image {image} {exists=} in registry")
-                return exists
+                repository = image.registry + "/" + image.name if image.registry else image.name
+                image_uri = await checker.image_exists(repository, tag, tuple(image.platform))
+                if image_uri:
+                    logger.debug(f"Image {image_uri} in registry")
+                return image_uri
             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
+        return image.uri
 
     @classmethod
     @alru_cache
     async def build(
-        cls, image: Image, builder: Optional[str] = None, dry_run: bool = False, force: bool = False
+        cls,
+        image: Image,
+        builder: ImageBuildEngine.ImageBuilderType | None = None,
+        dry_run: bool = False,
+        force: bool = False,
     ) -> str:
         """
         Build the image. Images to be tagged with latest will always be built. Otherwise, this engine will check the
@@ -181,30 +194,59 @@ class ImageBuildEngine:
         :return:
         """
         # Always trigger a build if this is a dry run since builder shouldn't really do anything, or a force.
+        image_uri = (await cls.image_exists(image)) or image.uri
         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...")
+            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
+            # If a builder is not specified, use the first registered builder
+            cfg = _get_init_config()
+            if cfg and cfg.image_builder:
+                builder = builder or cfg.image_builder
             img_builder = ImageBuildEngine._get_builder(builder)
+            logger.debug(f"Using `{img_builder}` image builder to build image.")
 
             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
+            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
+    def _get_builder(cls, builder: ImageBuildEngine.ImageBuilderType | None = "local") -> ImageBuilder:
+        if builder is None:
+            builder = "local"
+        if builder == "remote":
+            from flyte._internal.imagebuild.remote_builder import RemoteImageBuilder
+
+            return RemoteImageBuilder()
+        elif builder == "local":
+            from flyte._internal.imagebuild.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]
+        else:
+            return cls._load_custom_type_transformers(builder)
+
+    @classmethod
+    def _load_custom_type_transformers(cls, name: str) -> ImageBuilder:
+        plugins = entry_points(group="flyte.plugins.image_builders")
+        for ep in plugins:
+            if ep.name != name:
+                continue
+            try:
+                logger.info(f"Loading image builder: {ep.name}")
+                builder = ep.load()
+                if callable(builder):
+                    return builder()
+                return builder
+            except Exception as e:
+                raise RuntimeError(f"Failed to load image builder {ep.name} with error: {e}")
+        raise ValueError(
+            f"Unknown image builder type: {name}. Available builders:"
+            f" {[ep.name for ep in plugins] + ['local', 'remote']}"
+        )
 
 
 class ImageCache(BaseModel):
@@ -239,3 +281,17 @@ class ImageCache(BaseModel):
         val = cls.model_validate_json(json_str)
         val.serialized_form = s
         return val
+
+    def repr(self) -> typing.List[typing.List[Tuple[str, str]]]:
+        """
+        Returns a detailed representation of the deployed environments.
+        """
+        tuples = []
+        for k, v in self.image_lookup.items():
+            tuples.append(
+                [
+                    ("Name", k),
+                    ("image", v),
+                ]
+            )
+        return tuples