flyte 2.0.0b20__py3-none-any.whl → 2.0.0b22__py3-none-any.whl

flyte/_bin/runtime.py

@@ -21,8 +21,8 @@ import click
 
 ACTION_NAME = "ACTION_NAME"
 RUN_NAME = "RUN_NAME"
-PROJECT_NAME = "FLYTE_INTERNAL_TASK_PROJECT"
-DOMAIN_NAME = "FLYTE_INTERNAL_TASK_DOMAIN"
+PROJECT_NAME = "FLYTE_INTERNAL_EXECUTION_PROJECT"
+DOMAIN_NAME = "FLYTE_INTERNAL_EXECUTION_DOMAIN"
 ORG_NAME = "_U_ORG_NAME"
 ENDPOINT_OVERRIDE = "_U_EP_OVERRIDE"
 RUN_OUTPUT_BASE_DIR = "_U_RUN_BASE"
@@ -137,7 +137,7 @@ def main(
         logger.debug(f"Using controller endpoint: {ep} with kwargs: {controller_kwargs}")
 
     bundle = CodeBundle(tgz=tgz, pkl=pkl, destination=dest, computed_version=version)
-    init(org=org, project=project, domain=domain, **controller_kwargs)
+    init(org=org, project=project, domain=domain, image_builder="remote", **controller_kwargs)
     # Controller is created with the same kwargs as init, so that it can be used to run tasks
     controller = create_controller(ct="remote", **controller_kwargs)
 

flyte/_code_bundle/_ignore.py

@@ -83,8 +83,15 @@ class StandardIgnore(Ignore):
         self.patterns = patterns if patterns else STANDARD_IGNORE_PATTERNS
 
     def _is_ignored(self, path: pathlib.Path) -> bool:
+        # Convert to relative path for pattern matching
+        try:
+            rel_path = path.relative_to(self.root)
+        except ValueError:
+            # If path is not under root, don't ignore it
+            return False
+
         for pattern in self.patterns:
-            if fnmatch(str(path), pattern):
+            if fnmatch(str(rel_path), pattern):
                 return True
         return False
 
@@ -105,9 +112,10 @@ class IgnoreGroup(Ignore):
 
     def list_ignored(self) -> List[str]:
         ignored = []
-        for dir, _, files in self.root.walk():
+        for dir, _, files in os.walk(self.root):
+            dir_path = Path(dir)
             for file in files:
-                abs_path = dir / file
+                abs_path = dir_path / file
                 if self.is_ignored(abs_path):
                     ignored.append(str(abs_path.relative_to(self.root)))
         return ignored

flyte/_code_bundle/_packaging.py

@@ -14,10 +14,9 @@ import typing
 from typing import List, Optional, Tuple, Union
 
 import click
-from rich import print as rich_print
 from rich.tree import Tree
 
-from flyte._logging import logger
+from flyte._logging import _get_console, logger
 
 from ._ignore import Ignore, IgnoreGroup
 from ._utils import CopyFiles, _filehash_update, _pathhash_update, ls_files, tar_strip_file_attributes
@@ -27,10 +26,10 @@ FAST_FILEENDING = ".tar.gz"
 
 
 def print_ls_tree(source: os.PathLike, ls: typing.List[str]):
-    click.secho("Files to be copied for fast registration...", fg="bright_blue")
+    logger.info("Files to be copied for fast registration...")
 
     tree_root = Tree(
-        f":open_file_folder: [link file://{source}]{source} (detected source root)",
+        f"File structure:\n:open_file_folder: {source}",
         guide_style="bold bright_blue",
     )
     trees = {pathlib.Path(source): tree_root}
@@ -49,7 +48,12 @@ def print_ls_tree(source: os.PathLike, ls: typing.List[str]):
                 else:
                     current = trees[current_path]
         trees[fpp.parent].add(f"{fpp.name}", guide_style="bold bright_blue")
-    rich_print(tree_root)
+
+    console = _get_console()
+    with console.capture() as capture:
+        console.print(tree_root, overflow="ignore", no_wrap=True, crop=False)
+    logger.info(f"Root directory: [link=file://{source}]{source}[/link]")
+    logger.info(capture.get(), extra={"console": console})
 
 
 def _compress_tarball(source: pathlib.Path, output: pathlib.Path) -> None:

flyte/_code_bundle/_utils.py

@@ -156,7 +156,7 @@ def list_all_files(source_path: pathlib.Path, deref_symlinks, ignore_group: Opti
 
     # This is needed to prevent infinite recursion when walking with followlinks
     visited_inodes = set()
-    for root, dirnames, files in source_path.walk(top_down=True, follow_symlinks=deref_symlinks):
+    for root, dirnames, files in os.walk(source_path, topdown=True, followlinks=deref_symlinks):
         dirnames[:] = [d for d in dirnames if d not in EXCLUDE_DIRS]
         if deref_symlinks:
             inode = os.stat(root).st_ino
@@ -167,7 +167,7 @@ def list_all_files(source_path: pathlib.Path, deref_symlinks, ignore_group: Opti
         ff = []
         files.sort()
         for fname in files:
-            abspath = (root / fname).absolute()
+            abspath = (pathlib.Path(root) / fname).absolute()
             # Only consider files that exist (e.g. disregard symlinks that point to non-existent files)
             if not os.path.exists(abspath):
                 logger.info(f"Skipping non-existent file {abspath}")

flyte/_deploy.py

@@ -1,10 +1,13 @@
 from __future__ import annotations
 
 import asyncio
+import hashlib
+import sys
 import typing
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
 
+import cloudpickle
 import rich.repr
 
 import flyte.errors
@@ -160,9 +163,11 @@ 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] = image_uri
+            py_version = "{}.{}".format(*env.image.python_version)
+            image_identifier_map[env.image.identifier] = {py_version: image_uri}
         elif env.image == "auto":
-            image_identifier_map["auto"] = image_uri
+            py_version = "{}.{}".format(sys.version_info.major, sys.version_info.minor)
+            image_identifier_map["auto"] = {py_version: image_uri}
 
     return ImageCache(image_lookup=image_identifier_map)
 
@@ -175,15 +180,18 @@ async def apply(deployment_plan: DeploymentPlan, copy_style: CopyFiles, dryrun:
 
     image_cache = await _build_images(deployment_plan)
 
-    version = deployment_plan.version
-    if copy_style == "none" and not version:
+    if copy_style == "none" and not deployment_plan.version:
         raise flyte.errors.DeploymentError("Version must be set when copy_style is none")
     else:
         code_bundle = await build_code_bundle(from_dir=cfg.root_dir, dryrun=dryrun, copy_style=copy_style)
-        version = version or code_bundle.computed_version
-        # TODO we should update the version to include the image cache digest and code bundle digest. This is
-        # to ensure that changes in image dependencies, cause an update to the deployment version.
-        # TODO Also hash the environment and tasks to ensure that changes in the environment or tasks
+        if deployment_plan.version:
+            version = deployment_plan.version
+        else:
+            h = hashlib.md5()
+            h.update(cloudpickle.dumps(deployment_plan.envs))
+            h.update(code_bundle.computed_version.encode("utf-8"))
+            h.update(cloudpickle.dumps(image_cache))
+            version = h.hexdigest()
 
     sc = SerializationContext(
         project=cfg.project,

flyte/_image.py

@@ -203,6 +203,7 @@ class UVProject(PipOption, Layer):
 
         super().update_hash(hasher)
         filehash_update(self.uvlock, hasher)
+        filehash_update(self.pyproject, hasher)
 
 
 @rich.repr.auto
@@ -435,7 +436,10 @@ class Image:
         # Only get the non-None values in the Image to ensure the hash is consistent
         # across different SDK versions.
         # Layers can specify a _compute_identifier optionally, but the default will just stringify
-        image_dict = asdict(self, dict_factory=lambda x: {k: v for (k, v) in x if v is not None and k != "_layers"})
+        image_dict = asdict(
+            self,
+            dict_factory=lambda x: {k: v for (k, v) in x if v is not None and k not in ("_layers", "python_version")},
+        )
         layers_str_repr = "".join([layer.identifier() for layer in self._layers])
         image_dict["layers"] = layers_str_repr
         spec_bytes = image_dict.__str__().encode("utf-8")

flyte/_initialize.py

@@ -33,6 +33,7 @@ class CommonInit:
     project: str | None = None
     domain: str | None = None
     batch_size: int = 1000
+    source_config_path: Optional[Path] = None  # Only used for documentation
 
 
 @dataclass(init=True, kw_only=True, repr=True, eq=True, frozen=True)
@@ -110,6 +111,12 @@ async def _initialize_client(
     )
 
 
+def _initialize_logger(log_level: int | None = None):
+    initialize_logger(enable_rich=True)
+    if log_level:
+        initialize_logger(log_level=log_level, enable_rich=True)
+
+
 @syncify
 async def init(
     org: str | None = None,
@@ -134,6 +141,7 @@ async def init(
     storage: Storage | None = None,
     batch_size: int = 1000,
     image_builder: ImageBuildEngine.ImageBuilderType = "local",
+    source_config_path: Optional[Path] = None,
 ) -> None:
     """
     Initialize the Flyte system with the given configuration. This method should be called before any other Flyte
@@ -169,17 +177,12 @@ async def init(
     :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.
-
+    :param source_config_path: Optional path to the source configuration file (This is only used for documentation)
     :return: None
     """
-    from flyte._tools import ipython_check
     from flyte._utils import get_cwd_editable_install, org_from_endpoint, sanitize_endpoint
 
-    interactive_mode = ipython_check()
-
-    initialize_logger(enable_rich=interactive_mode)
-    if log_level:
-        initialize_logger(log_level=log_level, enable_rich=interactive_mode)
+    _initialize_logger(log_level=log_level)
 
     global _init_config  # noqa: PLW0603
 
@@ -223,6 +226,7 @@ async def init(
             org=org or org_from_endpoint(endpoint),
             batch_size=batch_size,
             image_builder=image_builder,
+            source_config_path=source_config_path,
         )
 
 
@@ -231,6 +235,7 @@ async def init_from_config(
     path_or_config: str | Path | Config | None = None,
     root_dir: Path | None = None,
     log_level: int | None = None,
+    storage: Storage | None = None,
 ) -> None:
     """
     Initialize the Flyte system using a configuration file or Config object. This method should be called before any
@@ -243,11 +248,15 @@ async def init_from_config(
         if not available, the current working directory.
     :param log_level: Optional logging level for the framework logger,
         default is set using the default initialization policies
+    :param storage: Optional blob store (S3, GCS, Azure) configuration if needed to access (i.e. using Minio)
     :return: None
     """
+    from rich.highlighter import ReprHighlighter
+
     import flyte.config as config
 
     cfg: config.Config
+    cfg_path: Optional[Path] = None
     if path_or_config is None:
         # If no path is provided, use the default config file
         cfg = config.auto()
@@ -266,7 +275,9 @@ async def init_from_config(
     else:
         cfg = path_or_config
 
-    logger.debug(f"Flyte config initialized as {cfg}")
+    _initialize_logger(log_level=log_level)
+
+    logger.info(f"Flyte config initialized as {cfg}", extra={"highlighter": ReprHighlighter()})
     await init.aio(
         org=cfg.task.org,
         project=cfg.task.project,
@@ -283,6 +294,8 @@ async def init_from_config(
         root_dir=root_dir,
         log_level=log_level,
         image_builder=cfg.image.builder,
+        storage=storage,
+        source_config_path=cfg_path,
     )
 
 

flyte/_interface.py

@@ -1,7 +1,13 @@
 from __future__ import annotations
 
 import inspect
-from typing import Dict, Generator, Tuple, Type, TypeVar, Union, cast, get_args, get_type_hints
+import typing
+from enum import Enum
+from typing import Dict, Generator, Literal, Tuple, Type, TypeVar, Union, cast, get_args, get_origin, get_type_hints
+
+from flyte._logging import logger
+
+LITERAL_ENUM = "LiteralEnum"
 
 
 def default_output_name(index: int = 0) -> str:
@@ -69,7 +75,15 @@ def extract_return_annotation(return_annotation: Union[Type, Tuple, None]) -> Di
         if len(return_annotation.__args__) == 1:  # type: ignore
             raise TypeError("Tuples should be used to indicate multiple return values, found only one return variable.")
         ra = get_args(return_annotation)
-        return dict(zip(list(output_name_generator(len(ra))), ra))
+        annotations = {}
+        for i, r in enumerate(ra):
+            if r is Ellipsis:
+                raise TypeError("Variable length tuples are not supported as return types.")
+            if get_origin(r) is Literal:
+                annotations[default_output_name(i)] = literal_to_enum(cast(Type, r))
+            else:
+                annotations[default_output_name(i)] = r
+        return annotations
 
     elif isinstance(return_annotation, tuple):
         if len(return_annotation) == 1:
@@ -79,4 +93,25 @@ def extract_return_annotation(return_annotation: Union[Type, Tuple, None]) -> Di
     else:
         # Handle all other single return types
         # Task returns unnamed native tuple
+        if get_origin(return_annotation) is Literal:
+            return {default_output_name(): literal_to_enum(cast(Type, return_annotation))}
         return {default_output_name(): cast(Type, return_annotation)}
+
+
+def literal_to_enum(literal_type: Type) -> Type[Enum | typing.Any]:
+    """Convert a Literal[...] into Union[str, Enum]."""
+
+    if get_origin(literal_type) is not Literal:
+        raise TypeError(f"{literal_type} is not a Literal")
+
+    values = get_args(literal_type)
+    if not all(isinstance(v, str) for v in values):
+        logger.warning(f"Literal type {literal_type} contains non-string values, using Any instead of Enum")
+        return typing.Any
+    # Deduplicate & keep order
+    enum_dict = {str(v).upper(): v for v in values}
+
+    # Dynamically create an Enum
+    literal_enum = Enum(LITERAL_ENUM, enum_dict)  # type: ignore
+
+    return literal_enum  # type: ignore

flyte/_internal/imagebuild/docker_builder.py

@@ -6,7 +6,7 @@ import tempfile
 import typing
 from pathlib import Path
 from string import Template
-from typing import ClassVar, Optional, Protocol, cast
+from typing import ClassVar, List, Optional, Protocol, cast
 
 import aiofiles
 import click
@@ -245,6 +245,9 @@ class UVProjectHandler:
         else:
             # Copy the entire project.
             pyproject_dst = copy_files_to_context(layer.pyproject.parent, context_path)
+            if layer.uvlock:
+                # Sometimes the uv.lock file is in a different folder, if it's specified, let's copy it there explicitly
+                shutil.copy(layer.uvlock, pyproject_dst)
             delta = UV_LOCK_INSTALL_TEMPLATE.substitute(
                 PYPROJECT_PATH=pyproject_dst.relative_to(context_path),
                 PIP_INSTALL_ARGS=" ".join(layer.get_pip_install_args()),
@@ -263,7 +266,37 @@ class DockerIgnoreHandler:
 
 class CopyConfigHandler:
     @staticmethod
-    async def handle(layer: CopyConfig, context_path: Path, dockerfile: str) -> str:
+    def list_dockerignore(root_path: Optional[Path], docker_ignore_file_path: Optional[Path]) -> List[str]:
+        patterns: List[str] = []
+        dockerignore_path: Optional[Path] = None
+        if root_path:
+            dockerignore_path = root_path / ".dockerignore"
+        # DockerIgnore layer should be first priority
+        if docker_ignore_file_path:
+            dockerignore_path = docker_ignore_file_path
+
+        # 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
+
+    @staticmethod
+    async def handle(
+        layer: CopyConfig, context_path: Path, dockerfile: str, docker_ignore_file_path: Optional[Path]
+    ) -> str:
         # Copy the source config file or directory to the context path
         if layer.src.is_absolute() or ".." in str(layer.src):
             dst_path = context_path / str(layer.src.absolute()).replace("/", "./_flyte_abs_context/", 1)
@@ -272,18 +305,26 @@ class CopyConfigHandler:
 
         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, dst_path)
         elif layer.src.is_dir():
             # Copy the entire directory
-            shutil.copytree(abs_path, dst_path, dirs_exist_ok=True)
+            from flyte._initialize import _get_init_config
+
+            init_config = _get_init_config()
+            root_path = init_config.root_dir if init_config else None
+            docker_ignore_patterns = CopyConfigHandler.list_dockerignore(root_path, docker_ignore_file_path)
+            shutil.copytree(
+                abs_path, dst_path, dirs_exist_ok=True, ignore=shutil.ignore_patterns(*docker_ignore_patterns)
+            )
         else:
-            raise ValueError(f"Source path is neither file nor directory: {layer.src}")
+            logger.error(f"Source path not exists: {layer.src}")
+            return dockerfile
 
         # Add a copy command to the dockerfile
         dockerfile += f"\nCOPY {dst_path.relative_to(context_path)} {layer.dst}\n"
-
         return dockerfile
 
 
@@ -349,7 +390,9 @@ def _get_secret_mounts_layer(secrets: typing.Tuple[str | Secret, ...] | None) ->
     return secret_mounts_layer
 
 
-async def _process_layer(layer: Layer, context_path: Path, dockerfile: str) -> str:
+async def _process_layer(
+    layer: Layer, context_path: Path, dockerfile: str, docker_ignore_file_path: Optional[Path]
+) -> str:
     match layer:
         case PythonWheels():
             # Handle Python wheels
@@ -385,7 +428,7 @@ async def _process_layer(layer: Layer, context_path: Path, dockerfile: str) -> s
 
         case CopyConfig():
             # Handle local files and folders
-            dockerfile = await CopyConfigHandler.handle(layer, context_path, dockerfile)
+            dockerfile = await CopyConfigHandler.handle(layer, context_path, dockerfile, docker_ignore_file_path)
 
         case Commands():
             # Handle commands
@@ -512,6 +555,15 @@ class DockerImageBuilder(ImageBuilder):
         else:
             logger.info("Buildx builder already exists.")
 
+    def get_docker_ignore(self, image: Image) -> Optional[Path]:
+        """Get the .dockerignore file path from the image layers."""
+        # Look for DockerIgnore layer in the image layers
+        for layer in image._layers:
+            if isinstance(layer, DockerIgnore) and layer.path.strip():
+                return Path(layer.path)
+
+        return None
+
     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
@@ -541,8 +593,9 @@ class DockerImageBuilder(ImageBuilder):
                 PYTHON_VERSION=f"{image.python_version[0]}.{image.python_version[1]}",
             )
 
+            docker_ignore_file_path = self.get_docker_ignore(image)
             for layer in image._layers:
-                dockerfile = await _process_layer(layer, tmp_path, dockerfile)
+                dockerfile = await _process_layer(layer, tmp_path, dockerfile, docker_ignore_file_path)
 
             dockerfile += DOCKER_FILE_BASE_FOOTER.substitute(F_IMG_ID=image.uri)
 

flyte/_internal/imagebuild/image_builder.py

@@ -235,7 +235,7 @@ class ImageBuildEngine:
 
 
 class ImageCache(BaseModel):
-    image_lookup: Dict[str, str]
+    image_lookup: Dict[str, Dict[str, str]]
     serialized_form: str | None = None
 
     @property
@@ -273,10 +273,11 @@ class ImageCache(BaseModel):
         """
         tuples = []
         for k, v in self.image_lookup.items():
-            tuples.append(
-                [
-                    ("Name", k),
-                    ("image", v),
-                ]
-            )
+            for py_version, image_uri in v.items():
+                tuples.append(
+                    [
+                        ("Name", f"{k} (py{py_version})"),
+                        ("image", image_uri),
+                    ]
+                )
         return tuples

flyte/_internal/imagebuild/remote_builder.py

@@ -10,7 +10,6 @@ from typing import TYPE_CHECKING, Optional, Tuple, cast
 from uuid import uuid4
 
 import aiofiles
-import click
 
 import flyte
 import flyte.errors
@@ -42,7 +41,6 @@ if TYPE_CHECKING:
     from flyte._protos.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"
 
@@ -85,10 +83,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
 
 
@@ -110,37 +108,25 @@ class RemoteImageBuilder(ImageBuilder):
             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]")
         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
             ),
         )
-        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, 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 [cyan]{run.url}[/cyan]")
 
         outputs = await run_details.outputs()
         return _get_fully_qualified_image_name(outputs)
@@ -183,11 +169,8 @@ async def _validate_configuration(image: Image) -> Tuple[str, Optional[str]]:
         context_size = tar_path.stat().st_size
         if context_size > 5 * 1024 * 1024:
             logger.warning(
-                click.style(
-                    f"Context size is {context_size / (1024 * 1024):.2f} MB, which is larger than 5 MB. "
-                    "Upload and build speed will be impacted.",
-                    fg="yellow",
-                )
+                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(context_dst)
     else:

flyte/_internal/runtime/task_serde.py

@@ -4,6 +4,7 @@ It includes a Resolver interface for loading tasks, and functions to load classe
 """
 
 import copy
+import sys
 import typing
 from datetime import timedelta
 from typing import Optional, cast
@@ -217,16 +218,25 @@ def _get_urun_container(
         img_uri = task_template.image.uri
     elif serialize_context.image_cache and image_id not in serialize_context.image_cache.image_lookup:
         img_uri = task_template.image.uri
-        from flyte._version import __version__
 
         logger.warning(
-            f"Image {task_template.image} not found in the image cache: {serialize_context.image_cache.image_lookup}.\n"
-            f"This typically occurs when the Flyte SDK version (`{__version__}`) used in the task environment "
-            f"differs from the version used to compile or deploy it.\n"
-            f"Ensure both environments use the same Flyte SDK version to avoid inconsistencies in image resolution."
+            f"Image {task_template.image} not found in the image cache: {serialize_context.image_cache.image_lookup}."
         )
     else:
-        img_uri = serialize_context.image_cache.image_lookup[image_id]
+        python_version_str = "{}.{}".format(sys.version_info.major, sys.version_info.minor)
+        version_lookup = serialize_context.image_cache.image_lookup[image_id]
+        if python_version_str in version_lookup:
+            img_uri = version_lookup[python_version_str]
+        elif version_lookup:
+            # Fallback: try to get any available version
+            fallback_py_version, img_uri = next(iter(version_lookup.items()))
+            logger.warning(
+                f"Image {task_template.image} for python version {python_version_str} "
+                f"not found in the image cache: {serialize_context.image_cache.image_lookup}.\n"
+                f"Fall back using image {img_uri} for python version {fallback_py_version} ."
+            )
+        else:
+            img_uri = task_template.image.uri
 
     return tasks_pb2.Container(
         image=img_uri,