Add wait in image builder protocol

examples/image/base_image.py

@@ -6,7 +6,7 @@
 image = (
     Image.from_debian_base(install_flyte=False)
     .with_apt_packages("vim", "wget")
-    .with_pip_packages("mypy", pre=True)
+    .with_pip_packages("mypy", "httpx", pre=True)
     .with_env_vars({"hello": "world1"})
     .with_dockerignore(Path(__file__).parent / ".dockerignore")
     .with_local_v2()
@@ -22,6 +22,5 @@ async def t1(data: str = "hello") -> str:
 
 if __name__ == "__main__":
     flyte.init_from_config()
-    run = flyte.run(t1, data="world")
-    print(run.name)
-    print(run.url)
+    uri = flyte.build(image, force=True, wait=False)
+    print(uri)

examples/image/custom_builder/src/my_builder/my_builder.py

@@ -20,6 +20,11 @@ def get_checkers(self) -> Optional[typing.List[typing.Type[ImageChecker]]]:
         """Return the image checker."""
         return [MyImageChecker]
 
-    async def build_image(self, image: Image, dry_run: bool = False) -> str:
+    async def build_image(
+        self,
+        image: Image,
+        dry_run: bool = False,
+        wait: bool = True,
+    ) -> str:
         print("Building image locally...")
         return image.uri

src/flyte/_build.py

@@ -6,10 +6,25 @@
 
 
 @syncify
-async def build(image: Image) -> str:
+async def build(
+    image: Image,
+    dry_run: bool = False,
+    force: bool = False,
+    wait: bool = True,
+) -> str:
     """
     Build an image. The existing async context will be used.
 
+    Args:
+        image: The image(s) to build.
+        dry_run: Tell the builder to not actually build. Different builders will have different behaviors.
+        force: Skip the existence check. Normally if the image already exists we won't build it.
+        wait: Wait for the build to finish. If wait is False, the function will return immediately and the build will
+            run in the background.
+    Returns:
+        The image URI. If wait is False when using the remote image builder, the function will return the build image
+        task URL.
+
     Example:
     ```
     import flyte
@@ -19,8 +34,12 @@ async def build(image: Image) -> str:
     ```
 
     :param image: The image(s) to build.
+    :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.
+    :param wait: Wait for the build to finish. If wait is False, the function will return immediately and the build will
+        run in the background.
     :return: The image URI.
     """
     from flyte._internal.imagebuild.image_builder import ImageBuildEngine
 
-    return await ImageBuildEngine.build(image)
+    return await ImageBuildEngine.build(image, dry_run=dry_run, force=force, wait=wait)

src/flyte/_internal/imagebuild/docker_builder.py

@@ -544,10 +544,10 @@ 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:
+    async def build_image(self, image: Image, dry_run: bool = False, wait: bool = True) -> str:
         if image.dockerfile:
             # If a dockerfile is provided, use it directly
-            return await self._build_from_dockerfile(image, push=True)
+            return await self._build_from_dockerfile(image, push=True, wait=wait)
 
         if len(image._layers) == 0:
             logger.warning("No layers to build, returning the image URI as is.")
@@ -559,7 +559,7 @@ async def build_image(self, image: Image, dry_run: bool = False) -> str:
             dry_run=dry_run,
         )
 
-    async def _build_from_dockerfile(self, image: Image, push: bool) -> str:
+    async def _build_from_dockerfile(self, image: Image, push: bool, wait: bool = True) -> str:
         """
         Build the image from a provided Dockerfile.
         """
@@ -592,7 +592,10 @@ async def _build_from_dockerfile(self, image: Image, push: bool) -> str:
         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(cast(Path, image.dockerfile).cwd()), check=True)
+        if wait:
+            await asyncio.to_thread(subprocess.run, command, cwd=str(cast(Path, image.dockerfile).cwd()), check=True)
+        else:
+            await asyncio.to_thread(subprocess.Popen, command, cwd=str(cast(Path, image.dockerfile).cwd()))
 
         return image.uri
 
@@ -633,7 +636,7 @@ async def _ensure_buildx_builder():
         else:
             logger.info("Buildx builder already exists.")
 
-    async def _build_image(self, image: Image, *, push: bool = True, dry_run: bool = False) -> str:
+    async def _build_image(self, image: Image, *, push: bool = True, dry_run: bool = False, wait: bool = True) -> str:
         """
         if default image (only base image and locked), raise an error, don't have a dockerfile
         if dockerfile, just build
@@ -715,7 +718,10 @@ async def _build_image(self, image: Image, *, push: bool = True, dry_run: bool =
                 click.secho(f"Run command: {concat_command} ", fg="blue")
 
             try:
-                await asyncio.to_thread(subprocess.run, command, check=True)
+                if wait:
+                    await asyncio.to_thread(subprocess.run, command, check=True)
+                else:
+                    await asyncio.to_thread(subprocess.Popen, command)
             except subprocess.CalledProcessError as e:
                 logger.error(f"Failed to build image: {e}")
                 raise RuntimeError(f"Failed to build image: {e}")

src/flyte/_internal/imagebuild/image_builder.py

@@ -16,7 +16,7 @@
 
 
 class ImageBuilder(Protocol):
-    async def build_image(self, image: Image, dry_run: bool) -> str: ...
+    async def build_image(self, image: Image, dry_run: bool, wait: bool = True) -> str: ...
 
     def get_checkers(self) -> Optional[typing.List[typing.Type[ImageChecker]]]:
         """
@@ -182,6 +182,7 @@ async def build(
         builder: ImageBuildEngine.ImageBuilderType | None = None,
         dry_run: bool = False,
         force: bool = False,
+        wait: bool = True,
     ) -> str:
         """
         Build the image. Images to be tagged with latest will always be built. Otherwise, this engine will check the
@@ -191,7 +192,10 @@ async def build(
         :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:
+        :param wait: Wait for the build to finish. If wait is False when using the remote image builder, the function
+            will return the build image task URL.
+        :return: The image URI. If wait is False when using the remote image builder, the function will return the build
+            image task URL.
         """
         # 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
@@ -208,7 +212,7 @@ async def build(
             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)
+            result = await img_builder.build_image(image, dry_run=dry_run, wait=wait)
             return result
         else:
             logger.info(f"Image {image_uri} already exists in registry. Skipping build.")

src/flyte/_internal/imagebuild/remote_builder.py

@@ -104,7 +104,7 @@ def get_checkers(self) -> Optional[typing.List[typing.Type[ImageChecker]]]:
         """Return the image checker."""
         return [RemoteImageChecker]
 
-    async def build_image(self, image: Image, dry_run: bool = False) -> str:
+    async def build_image(self, image: Image, dry_run: bool = False, wait: bool = True) -> str:
         image_name = f"{image.name}:{image._final_tag}"
         spec, context = await _validate_configuration(image)
 
@@ -137,11 +137,15 @@ async def build_image(self, image: Image, dry_run: bool = False) -> str:
                 project=cfg.project, domain=cfg.domain, cache_lookup_scope="project-domain"
             ).run.aio(entity, spec=spec, context=context, target_image=target_image),
         )
-        logger.warning(f"⏳ Waiting for build to finish at: [bold cyan link={run.url}]{run.url}[/bold cyan link]")
 
+        logger.warning(f"▶️ Started build at: [bold cyan link={run.url}]{run.url}[/bold cyan link]")
+        if not wait:
+            # return the run url of the build image task
+            return run.url
+
+        logger.warning("⏳ Waiting for build to finish")
         await run.wait.aio(quiet=True)
         run_details = await run.details.aio()
-
         elapsed = str(datetime.now(timezone.utc) - start).split(".")[0]
 
         if run_details.action_details.raw_phase == phase_pb2.ACTION_PHASE_SUCCEEDED: