flyte 2.0.0b13__py3-none-any.whl → 2.0.0b30__py3-none-any.whl

flyte/_internal/imagebuild/remote_builder.py

@@ -1,5 +1,7 @@
+import gzip
 import os
 import shutil
+import tarfile
 import tempfile
 import typing
 from datetime import datetime, timezone
@@ -7,19 +9,23 @@ from pathlib import Path
 from typing import TYPE_CHECKING, Optional, Tuple, cast
 from uuid import uuid4
 
-import click
+import aiofiles
 
 import flyte
 import flyte.errors
 from flyte import Image, remote
+from flyte._code_bundle._utils import tar_strip_file_attributes
 from flyte._image import (
+    _BASE_REGISTRY,
     AptPackages,
     Architecture,
     Commands,
     CopyConfig,
     DockerIgnore,
     Env,
+    PipOption,
     PipPackages,
+    PoetryProject,
     PythonWheels,
     Requirements,
     UVProject,
@@ -27,15 +33,16 @@ from flyte._image import (
     WorkDir,
 )
 from flyte._internal.imagebuild.image_builder import ImageBuilder, ImageChecker
-from flyte._internal.imagebuild.utils import copy_files_to_context
+from flyte._internal.imagebuild.utils import copy_files_to_context, get_and_list_dockerignore
+from flyte._internal.runtime.task_serde import get_security_context
 from flyte._logging import logger
+from flyte._secret import Secret
 from flyte.remote import ActionOutputs, Run
 
 if TYPE_CHECKING:
-    from flyte._protos.imagebuilder import definition_pb2 as image_definition_pb2
+    from flyteidl2.imagebuilder import definition_pb2 as image_definition_pb2
 
 IMAGE_TASK_NAME = "build-image"
-OPTIMIZE_TASK_NAME = "optimize_task"
 IMAGE_TASK_PROJECT = "system"
 IMAGE_TASK_DOMAIN = "production"
 
@@ -63,10 +70,11 @@ class RemoteImageChecker(ImageChecker):
         image_name = f"{repository.split('/')[-1]}:{tag}"
 
         try:
+            from flyteidl2.imagebuilder import definition_pb2 as image_definition__pb2
+            from flyteidl2.imagebuilder import payload_pb2 as image_payload__pb2
+            from flyteidl2.imagebuilder import service_pb2_grpc as image_service_pb2_grpc
+
             from flyte._initialize import _get_init_config
-            from flyte._protos.imagebuilder import definition_pb2 as image_definition__pb2
-            from flyte._protos.imagebuilder import payload_pb2 as image_payload__pb2
-            from flyte._protos.imagebuilder import service_pb2_grpc as image_service_pb2_grpc
 
             cfg = _get_init_config()
             if cfg is None:
@@ -78,10 +86,10 @@ class RemoteImageChecker(ImageChecker):
                     raise ValueError("remote client should not be None")
                 cls._images_client = image_service_pb2_grpc.ImageServiceStub(cfg.client._channel)
             resp = await cls._images_client.GetImage(req)
-            logger.warning(click.style(f"Image {resp.image.fqin} found. Skip building.", fg="blue"))
+            logger.warning(f"[blue]Image {resp.image.fqin} found. Skip building.[/blue]")
             return resp.image.fqin
         except Exception:
-            logger.warning(click.style(f"Image {image_name} was not found or has expired.", fg="blue"))
+            logger.warning(f"[blue]Image {image_name} was not found or has expired.[/blue]", extra={"highlight": False})
             return None
 
 
@@ -91,49 +99,45 @@ class RemoteImageBuilder(ImageBuilder):
         return [RemoteImageChecker]
 
     async def build_image(self, image: Image, dry_run: bool = False) -> str:
-        from flyte._protos.workflow import run_definition_pb2
+        from flyteidl2.workflow import run_definition_pb2
 
         image_name = f"{image.name}:{image._final_tag}"
         spec, context = await _validate_configuration(image)
 
         start = datetime.now(timezone.utc)
-        entity = remote.Task.get(
+        entity = await remote.Task.get(
             name=IMAGE_TASK_NAME,
             project=IMAGE_TASK_PROJECT,
             domain=IMAGE_TASK_DOMAIN,
             auto_version="latest",
-        )
+        ).override.aio(secrets=_get_build_secrets_from_image(image))
+
+        logger.warning("[bold blue]🐳 Submitting a new build...[/bold blue]")
+        if image.registry and image.registry != _BASE_REGISTRY:
+            target_image = f"{image.registry}/{image_name}"
+        else:
+            # Use the default system registry in the backend.
+            target_image = image_name
+
+        from flyte._initialize import get_init_config
+        cfg = get_init_config()
         run = cast(
             Run,
-            await flyte.with_runcontext(project=IMAGE_TASK_PROJECT, domain=IMAGE_TASK_DOMAIN).run.aio(
-                entity, spec=spec, context=context, target_image=image_name
+            await flyte.with_runcontext(project=cfg.project, domain=cfg.domain).run.aio(
+                entity, spec=spec, context=context, target_image=target_image
             ),
         )
-        logger.warning(click.style("🐳 Submitting a new build...", fg="blue", bold=True))
+        logger.warning(f"⏳ Waiting for build to finish at: [bold cyan link={run.url}]{run.url}[/bold cyan link]")
 
-        logger.warning(click.style("⏳ Waiting for build to finish at: " + click.style(run.url, fg="cyan"), bold=True))
         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 == run_definition_pb2.PHASE_SUCCEEDED:
-            logger.warning(click.style(f"✅ Build completed in {elapsed}!", bold=True, fg="green"))
-            try:
-                entity = remote.Task.get(
-                    name=OPTIMIZE_TASK_NAME,
-                    project=IMAGE_TASK_PROJECT,
-                    domain=IMAGE_TASK_DOMAIN,
-                    auto_version="latest",
-                )
-                await flyte.with_runcontext(project=IMAGE_TASK_PROJECT, domain=IMAGE_TASK_DOMAIN).run.aio(
-                    entity, spec=spec, context=context, target_image=image_name
-                )
-            except Exception as e:
-                # Ignore the error if optimize is not enabled in the backend.
-                logger.warning(f"Failed to run optimize task with error: {e}")
+            logger.warning(f"[bold green]✅ Build completed in {elapsed}![/bold green]")
         else:
-            raise flyte.errors.ImageBuildError(f"❌ Build failed in {elapsed} at {click.style(run.url, fg='cyan')}")
+            raise flyte.errors.ImageBuildError(f"❌ Build failed in {elapsed} at {run.url}")
 
         outputs = await run_details.outputs()
         return _get_fully_qualified_image_name(outputs)
@@ -157,17 +161,29 @@ async def _validate_configuration(image: Image) -> Tuple[str, Optional[str]]:
 
     if any(context_path.iterdir()):
         # If there are files in the context directory, upload it
-        archive = Path(shutil.make_archive(str(tmp_path / "context"), "xztar", context_path))
-        st = archive.stat()
-        if st.st_size > 5 * 1024 * 1024:
-            logger.warning(
-                click.style(
-                    f"Context size is {st.st_size / (1024 * 1024):.2f} MB, which is larger than 5 MB. "
-                    "Upload and build speed will be impacted.",
-                    fg="yellow",
+        tar_path = tmp_path / "context.tar"
+        with tarfile.open(tar_path, "w", dereference=False) as tar:
+            files: typing.List[str] = os.listdir(context_path)
+            for ws_file in files:
+                tar.add(
+                    os.path.join(context_path, ws_file),
+                    recursive=True,
+                    arcname=ws_file,
+                    filter=tar_strip_file_attributes,
                 )
+        context_dst = Path(f"{tar_path!s}.gz")
+        with gzip.GzipFile(filename=context_dst, mode="wb", mtime=0) as gzipped:
+            async with aiofiles.open(tar_path, "rb") as tar_file:
+                content = await tar_file.read()
+                gzipped.write(content)
+
+        context_size = tar_path.stat().st_size
+        if context_size > 5 * 1024 * 1024:
+            logger.warning(
+                f"[yellow]Context size is {context_size / (1024 * 1024):.2f} MB, which is larger than 5 MB. "
+                "Upload and build speed will be impacted.[/yellow]",
             )
-        _, context_url = await remote.upload_file.aio(archive)
+        _, context_url = await remote.upload_file.aio(context_dst)
     else:
         context_url = ""
 
@@ -175,13 +191,36 @@ async def _validate_configuration(image: Image) -> Tuple[str, Optional[str]]:
 
 
 def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2.ImageSpec":
-    from flyte._protos.imagebuilder import definition_pb2 as image_definition_pb2
+    from flyteidl2.imagebuilder import definition_pb2 as image_definition_pb2
+
+    if image.dockerfile is not None:
+        raise flyte.errors.ImageBuildError(
+            "Custom Dockerfile is not supported with remote image builder.You can use local image builder instead."
+        )
 
     layers = []
     for layer in image._layers:
+        secret_mounts = None
+        pip_options = image_definition_pb2.PipOptions()
+
+        if isinstance(layer, PipOption):
+            pip_options = image_definition_pb2.PipOptions(
+                index_url=layer.index_url,
+                extra_index_urls=layer.extra_index_urls,
+                pre=layer.pre,
+                extra_args=layer.extra_args,
+            )
+
+        if hasattr(layer, "secret_mounts"):
+            sc = get_security_context(layer.secret_mounts)
+            secret_mounts = sc.secrets if sc else None
+
         if isinstance(layer, AptPackages):
             apt_layer = image_definition_pb2.Layer(
-                apt_packages=image_definition_pb2.AptPackages(packages=layer.packages)
+                apt_packages=image_definition_pb2.AptPackages(
+                    packages=layer.packages,
+                    secret_mounts=secret_mounts,
+                ),
             )
             layers.append(apt_layer)
         elif isinstance(layer, PythonWheels):
@@ -189,12 +228,8 @@ def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2
             wheel_layer = image_definition_pb2.Layer(
                 python_wheels=image_definition_pb2.PythonWheels(
                     dir=str(dst_path.relative_to(context_path)),
-                    options=image_definition_pb2.PipOptions(
-                        index_url=layer.index_url,
-                        extra_index_urls=layer.extra_index_urls,
-                        pre=layer.pre,
-                        extra_args=layer.extra_args,
-                    ),
+                    options=pip_options,
+                    secret_mounts=secret_mounts,
                 )
             )
             layers.append(wheel_layer)
@@ -204,12 +239,8 @@ def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2
             requirements_layer = image_definition_pb2.Layer(
                 requirements=image_definition_pb2.Requirements(
                     file=str(dst_path.relative_to(context_path)),
-                    options=image_definition_pb2.PipOptions(
-                        index_url=layer.index_url,
-                        extra_index_urls=layer.extra_index_urls,
-                        pre=layer.pre,
-                        extra_args=layer.extra_args,
-                    ),
+                    options=pip_options,
+                    secret_mounts=secret_mounts,
                 )
             )
             layers.append(requirements_layer)
@@ -226,12 +257,8 @@ def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2
             pip_layer = image_definition_pb2.Layer(
                 pip_packages=image_definition_pb2.PipPackages(
                     packages=packages,
-                    options=image_definition_pb2.PipOptions(
-                        index_url=layer.index_url,
-                        extra_index_urls=layer.extra_index_urls,
-                        pre=layer.pre,
-                        extra_args=layer.extra_args,
-                    ),
+                    options=pip_options,
+                    secret_mounts=secret_mounts,
                 )
             )
             layers.append(pip_layer)
@@ -239,18 +266,60 @@ def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2
             for line in layer.pyproject.read_text().splitlines():
                 if "tool.uv.index" in line:
                     raise ValueError("External sources are not supported in pyproject.toml")
-            shutil.copy2(layer.pyproject, context_path / layer.pyproject.name)
+
+            if layer.project_install_mode == "dependencies_only":
+                # Copy pyproject itself
+                pyproject_dst = copy_files_to_context(layer.pyproject, context_path)
+                if pip_options.extra_args:
+                    if "--no-install-project" not in pip_options.extra_args:
+                        pip_options.extra_args += " --no-install-project"
+                else:
+                    pip_options.extra_args = " --no-install-project"
+                if "--no-sources" not in pip_options.extra_args:
+                    pip_options.extra_args += " --no-sources"
+            else:
+                # Copy the entire project
+                docker_ignore_patterns = get_and_list_dockerignore(image)
+                pyproject_dst = copy_files_to_context(layer.pyproject.parent, context_path, docker_ignore_patterns)
 
             uv_layer = image_definition_pb2.Layer(
                 uv_project=image_definition_pb2.UVProject(
-                    pyproject=str(layer.pyproject.name),
-                    uvlock=str(layer.uvlock.name),
+                    pyproject=str(pyproject_dst.relative_to(context_path)),
+                    uvlock=str(copy_files_to_context(layer.uvlock, context_path).relative_to(context_path)),
+                    options=pip_options,
+                    secret_mounts=secret_mounts,
                 )
             )
             layers.append(uv_layer)
+        elif isinstance(layer, PoetryProject):
+            for line in layer.pyproject.read_text().splitlines():
+                if "tool.poetry.source" in line:
+                    raise ValueError("External sources are not supported in pyproject.toml")
+            extra_args = layer.extra_args or ""
+            if layer.project_install_mode == "dependencies_only":
+                # Copy pyproject itself
+                if "--no-root" not in extra_args:
+                    extra_args += " --no-root"
+                pyproject_dst = copy_files_to_context(layer.pyproject, context_path)
+            else:
+                # Copy the entire project
+                pyproject_dst = copy_files_to_context(layer.pyproject.parent, context_path)
+
+            poetry_layer = image_definition_pb2.Layer(
+                poetry_project=image_definition_pb2.PoetryProject(
+                    pyproject=str(pyproject_dst.relative_to(context_path)),
+                    poetry_lock=str(copy_files_to_context(layer.poetry_lock, context_path).relative_to(context_path)),
+                    extra_args=extra_args,
+                    secret_mounts=secret_mounts,
+                )
+            )
+            layers.append(poetry_layer)
         elif isinstance(layer, Commands):
             commands_layer = image_definition_pb2.Layer(
-                commands=image_definition_pb2.Commands(cmd=list(layer.commands))
+                commands=image_definition_pb2.Commands(
+                    cmd=list(layer.commands),
+                    secret_mounts=secret_mounts,
+                )
             )
             layers.append(commands_layer)
         elif isinstance(layer, DockerIgnore):
@@ -287,3 +356,25 @@ def _get_layers_proto(image: Image, context_path: Path) -> "image_definition_pb2
 
 def _get_fully_qualified_image_name(outputs: ActionOutputs) -> str:
     return outputs.pb2.literals[0].value.scalar.primitive.string_value
+
+
+def _get_build_secrets_from_image(image: Image) -> Optional[typing.List[Secret]]:
+    secrets = []
+    DEFAULT_SECRET_DIR = Path("/etc/flyte/secrets")
+    for layer in image._layers:
+        if isinstance(layer, (PipOption, Commands, AptPackages)) and layer.secret_mounts is not None:
+            for secret_mount in layer.secret_mounts:
+                # Mount all the image secrets to a default directory that will be passed to the BuildKit server.
+                if isinstance(secret_mount, Secret):
+                    secrets.append(Secret(key=secret_mount.key, group=secret_mount.group, mount=DEFAULT_SECRET_DIR))
+                elif isinstance(secret_mount, str):
+                    secrets.append(Secret(key=secret_mount, mount=DEFAULT_SECRET_DIR))
+                else:
+                    raise ValueError(f"Unsupported secret_mount type: {type(secret_mount)}")
+
+    image_registry_secret = image._image_registry_secret
+    if image_registry_secret:
+        secrets.append(
+            Secret(key=image_registry_secret.key, group=image_registry_secret.group, mount=DEFAULT_SECRET_DIR)
+        )
+    return secrets

flyte/_internal/imagebuild/utils.py

@@ -1,8 +1,12 @@
 import shutil
 from pathlib import Path
+from typing import List, Optional
 
+from flyte._image import DockerIgnore, Image
+from flyte._logging import logger
 
-def copy_files_to_context(src: Path, context_path: Path) -> Path:
+
+def copy_files_to_context(src: Path, context_path: Path, ignore_patterns: list[str] = []) -> Path:
     """
     This helper function ensures that absolute paths that users specify are converted correctly to a path in the
     context directory. Doing this prevents collisions while ensuring files are available in the context.
@@ -23,7 +27,52 @@ def copy_files_to_context(src: Path, context_path: Path) -> Path:
         dst_path = context_path / src
     dst_path.parent.mkdir(parents=True, exist_ok=True)
     if src.is_dir():
-        shutil.copytree(src, dst_path, dirs_exist_ok=True)
+        default_ignore_patterns = [".idea", ".venv"]
+        ignore_patterns = list(set(ignore_patterns + default_ignore_patterns))
+        shutil.copytree(src, dst_path, dirs_exist_ok=True, ignore=shutil.ignore_patterns(*ignore_patterns))
     else:
         shutil.copy(src, dst_path)
     return dst_path
+
+
+def get_and_list_dockerignore(image: Image) -> List[str]:
+    """
+    Get and parse dockerignore patterns from .dockerignore file.
+
+    This function first looks for a DockerIgnore layer in the image's layers. If found, it uses
+    the path specified in that layer. If no DockerIgnore layer is found, it falls back to looking
+    for a .dockerignore file in the root_path directory.
+
+    :param image: The Image object
+    """
+    from flyte._initialize import _get_init_config
+
+    # Look for DockerIgnore layer in the image layers
+    dockerignore_path: Optional[Path] = None
+    patterns: List[str] = []
+
+    for layer in image._layers:
+        if isinstance(layer, DockerIgnore) and layer.path.strip():
+            dockerignore_path = Path(layer.path)
+    # If DockerIgnore layer not specified, set dockerignore_path under root_path
+    init_config = _get_init_config()
+    root_path = init_config.root_dir if init_config else None
+    if not dockerignore_path and root_path:
+        dockerignore_path = Path(root_path) / ".dockerignore"
+    # Return empty list if no .dockerignore file found
+    if not dockerignore_path or not dockerignore_path.exists() or not dockerignore_path.is_file():
+        logger.info(f".dockerignore file not found at path: {dockerignore_path}")
+        return patterns
+
+    try:
+        with open(dockerignore_path, "r", encoding="utf-8") as f:
+            for line in f:
+                stripped_line = line.strip()
+                # Skip empty lines, whitespace-only lines, and comments
+                if not stripped_line or stripped_line.startswith("#"):
+                    continue
+                patterns.append(stripped_line)
+    except Exception as e:
+        logger.error(f"Failed to read .dockerignore file at {dockerignore_path}: {e}")
+        return []
+    return patterns

flyte/_internal/resolvers/_task_module.py

@@ -1,13 +1,11 @@
-import inspect
-import os
 import pathlib
-import sys
 from typing import Tuple
 
+from flyte._module import extract_obj_module
 from flyte._task import AsyncFunctionTaskTemplate, TaskTemplate
 
 
-def extract_task_module(task: TaskTemplate, /, source_dir: pathlib.Path | None = None) -> Tuple[str, str]:
+def extract_task_module(task: TaskTemplate, /, source_dir: pathlib.Path) -> Tuple[str, str]:
     """
     Extract the task module from the task template.
 
@@ -15,40 +13,9 @@ def extract_task_module(task: TaskTemplate, /, source_dir: pathlib.Path | None =
     :param source_dir: The source directory to use for relative paths.
     :return: A tuple containing the entity name, module
     """
-    entity_name = task.name
     if isinstance(task, AsyncFunctionTaskTemplate):
-        entity_module = inspect.getmodule(task.func)
-        if entity_module is None:
-            raise ValueError(f"Task {entity_name} has no module.")
-
-        fp = entity_module.__file__
-        if fp is None:
-            raise ValueError(f"Task {entity_name} has no module.")
-
-        file_path = pathlib.Path(fp)
-        # Get the relative path to the current directory
-        # Will raise ValueError if the file is not in the source directory
-        relative_path = file_path.relative_to(str(source_dir))
-
-        if relative_path == pathlib.Path("."):
-            entity_module_name = entity_module.__name__
-        else:
-            # Replace file separators with dots and remove the '.py' extension
-            dotted_path = os.path.splitext(str(relative_path))[0].replace(os.sep, ".")
-            entity_module_name = dotted_path
-
         entity_name = task.func.__name__
+        entity_module_name = extract_obj_module(task.func, source_dir)
+        return entity_name, entity_module_name
     else:
-        raise NotImplementedError(f"Task module {entity_name} not implemented.")
-
-    if entity_module_name == "__main__":
-        """
-        This case is for the case in which the task is run from the main module.
-        """
-        fp = sys.modules["__main__"].__file__
-        if fp is None:
-            raise ValueError(f"Task {entity_name} has no module.")
-        main_path = pathlib.Path(fp)
-        entity_module_name = main_path.stem
-
-    return entity_name, entity_module_name
+        raise NotImplementedError(f"Task module {task.name} not implemented.")

flyte/_internal/resolvers/default.py

@@ -1,6 +1,6 @@
 import importlib
 from pathlib import Path
-from typing import List, Optional
+from typing import List
 
 from flyte._internal.resolvers._task_module import extract_task_module
 from flyte._internal.resolvers.common import Resolver
@@ -23,6 +23,6 @@ class DefaultTaskResolver(Resolver):
         task_def = getattr(task_module, task_name)
         return task_def
 
-    def loader_args(self, task: TaskTemplate, root_dir: Optional[Path] = None) -> List[str]:  # type:ignore
+    def loader_args(self, task: TaskTemplate, root_dir: Path) -> List[str]:  # type:ignore
         t, m = extract_task_module(task, root_dir)
         return ["mod", m, "instance", t]