flyte 2.0.0b17__py3-none-any.whl → 2.0.0b18__py3-none-any.whl

flyte/_logging.py

@@ -75,39 +75,52 @@ def get_rich_handler(log_level: int) -> Optional[logging.Handler]:
     return handler
 
 
-def get_default_handler(log_level: int) -> logging.Handler:
-    handler = logging.StreamHandler()
-    handler.setLevel(log_level)
-    formatter = logging.Formatter(fmt="[%(name)s] %(message)s")
-    if log_format_from_env() == "json":
-        pass
-        # formatter = jsonlogger.JsonFormatter(fmt="%(asctime)s %(name)s %(levelname)s %(message)s")
-    handler.setFormatter(formatter)
-    return handler
-
-
 def initialize_logger(log_level: int = get_env_log_level(), enable_rich: bool = False):
     """
     Initializes the global loggers to the default configuration.
+    When enable_rich=True, upgrades to Rich handler for local CLI usage.
     """
     global logger  # noqa: PLW0603
-    logger = _create_logger("flyte", log_level, enable_rich)
 
+    # Clear existing handlers to reconfigure
+    root = logging.getLogger()
+    root.handlers.clear()
 
-def _create_logger(name: str, log_level: int = DEFAULT_LOG_LEVEL, enable_rich: bool = False) -> logging.Logger:
-    """
-    Creates a logger with the given name and log level.
-    """
-    logger = logging.getLogger(name)
-    logger.setLevel(log_level)
-    handler = None
-    logger.handlers = []
+    flyte_logger = logging.getLogger("flyte")
+    flyte_logger.handlers.clear()
+
+    # Set up root logger handler
+    root_handler = None
+    if enable_rich:
+        root_handler = get_rich_handler(log_level)
+
+    if root_handler is None:
+        root_handler = logging.StreamHandler()
+
+    # Add context filter to root handler for all logging
+    root_handler.addFilter(ContextFilter())
+    root.addHandler(root_handler)
+
+    # Set up Flyte logger handler
+    flyte_handler = None
     if enable_rich:
-        handler = get_rich_handler(log_level)
-    if handler is None:
-        handler = get_default_handler(log_level)
-    logger.addHandler(handler)
-    return logger
+        flyte_handler = get_rich_handler(log_level)
+
+    if flyte_handler is None:
+        flyte_handler = logging.StreamHandler()
+        flyte_handler.setLevel(log_level)
+        formatter = logging.Formatter(fmt="%(message)s")
+        flyte_handler.setFormatter(formatter)
+
+    # Add both filters to Flyte handler
+    flyte_handler.addFilter(FlyteInternalFilter())
+    flyte_handler.addFilter(ContextFilter())
+
+    flyte_logger.addHandler(flyte_handler)
+    flyte_logger.setLevel(log_level)
+    flyte_logger.propagate = False  # Prevent double logging
+
+    logger = flyte_logger
 
 
 def log(fn=None, *, level=logging.DEBUG, entry=True, exit=True):
@@ -135,4 +148,75 @@ def log(fn=None, *, level=logging.DEBUG, entry=True, exit=True):
     return decorator(fn)
 
 
-logger = _create_logger("flyte", get_env_log_level())
+class ContextFilter(logging.Filter):
+    """
+    A logging filter that adds the current action's run name and name to all log records.
+    Applied globally to capture context for both user and Flyte internal logging.
+    """
+
+    def filter(self, record):
+        from flyte._context import ctx
+
+        c = ctx()
+        if c:
+            action = c.action
+            record.msg = f"[{action.run_name}][{action.name}] {record.msg}"
+        return True
+
+
+class FlyteInternalFilter(logging.Filter):
+    """
+    A logging filter that adds [flyte] prefix to internal Flyte logging only.
+    """
+
+    def filter(self, record):
+        if record.name.startswith("flyte"):
+            record.msg = f"[flyte] {record.msg}"
+        return True
+
+
+def _setup_root_logger():
+    """
+    Configure the root logger to capture all logging with context information.
+    This ensures both user code and Flyte internal logging get the context.
+    """
+    root = logging.getLogger()
+    root.handlers.clear()  # Remove any existing handlers to prevent double logging
+
+    # Create a basic handler for the root logger
+    handler = logging.StreamHandler()
+    # Add context filter to ALL logging
+    handler.addFilter(ContextFilter())
+
+    # Simple formatter since filters handle prefixes
+    root.addHandler(handler)
+
+
+def _create_flyte_logger() -> logging.Logger:
+    """
+    Create the internal Flyte logger with [flyte] prefix.
+    """
+    flyte_logger = logging.getLogger("flyte")
+    flyte_logger.setLevel(get_env_log_level())
+
+    # Add a handler specifically for flyte logging with the prefix filter
+    handler = logging.StreamHandler()
+    handler.setLevel(get_env_log_level())
+    handler.addFilter(FlyteInternalFilter())
+    handler.addFilter(ContextFilter())
+
+    formatter = logging.Formatter(fmt="%(message)s")
+    handler.setFormatter(formatter)
+
+    # Prevent propagation to root to avoid double logging
+    flyte_logger.propagate = False
+    flyte_logger.addHandler(handler)
+
+    return flyte_logger
+
+
+# Initialize root logger for global context
+_setup_root_logger()
+
+# Create the Flyte internal logger
+logger = _create_flyte_logger()

flyte/_pod.py

@@ -3,7 +3,7 @@ from typing import TYPE_CHECKING, Dict, Optional
 
 if TYPE_CHECKING:
     from flyteidl.core.tasks_pb2 import K8sPod
-    from kubernetes.client import ApiClient, V1PodSpec
+    from kubernetes.client import V1PodSpec
 
 
 _PRIMARY_CONTAINER_NAME_FIELD = "primary_container_name"
@@ -21,6 +21,7 @@ class PodTemplate(object):
 
     def to_k8s_pod(self) -> "K8sPod":
         from flyteidl.core.tasks_pb2 import K8sObjectMetadata, K8sPod
+        from kubernetes.client import ApiClient
 
         return K8sPod(
             metadata=K8sObjectMetadata(labels=self.labels, annotations=self.annotations),

flyte/_run.py

@@ -161,7 +161,12 @@ class _Runner:
                 code_bundle = cached_value.code_bundle
                 image_cache = cached_value.image_cache
             else:
-                image_cache = await build_images.aio(cast(Environment, obj.parent_env()))
+                if not self._dry_run:
+                    image_cache = await build_images.aio(cast(Environment, obj.parent_env()))
+                else:
+                    from ._internal.imagebuild.image_builder import ImageCache
+
+                    image_cache = ImageCache(image_lookup={})
 
                 if self._interactive_mode:
                     code_bundle = await build_pkl_bundle(

flyte/_task_environment.py

@@ -18,7 +18,7 @@ from typing import (
 
 import rich.repr
 
-from ._cache import CacheRequest
+from ._cache import Cache, CacheRequest
 from ._doc import Documentation
 from ._environment import Environment
 from ._image import Image
@@ -74,6 +74,10 @@ class TaskEnvironment(Environment):
         super().__post_init__()
         if self.reusable is not None and self.plugin_config is not None:
             raise ValueError("Cannot set plugin_config when environment is reusable.")
+        if self.reusable and not isinstance(self.reusable, ReusePolicy):
+            raise TypeError(f"Expected reusable to be of type ReusePolicy, got {type(self.reusable)}")
+        if self.cache and not isinstance(self.cache, (str, Cache)):
+            raise TypeError(f"Expected cache to be of type str or Cache, got {type(self.cache)}")
 
     def clone_with(
         self,

flyte/_trace.py

@@ -3,6 +3,7 @@ import inspect
 import time
 from typing import Any, AsyncGenerator, AsyncIterator, Awaitable, Callable, TypeGuard, TypeVar, Union, cast
 
+from flyte._logging import logger
 from flyte.models import NativeInterface
 
 T = TypeVar("T")
@@ -33,10 +34,13 @@ def trace(func: Callable[..., T]) -> Callable[..., T]:
             iface = NativeInterface.from_callable(func)
             info, ok = await controller.get_action_outputs(iface, func, *args, **kwargs)
             if ok:
+                logger.info(f"Found existing trace info for {func}, {info}")
                 if info.output:
                     return info.output
                 elif info.error:
                     raise info.error
+            else:
+                logger.debug(f"No existing trace info found for {func}, proceeding to execute.")
             start_time = time.time()
             try:
                 # Cast to Awaitable to satisfy mypy
@@ -44,6 +48,7 @@ def trace(func: Callable[..., T]) -> Callable[..., T]:
                 results = await coroutine_result
                 info.add_outputs(results, start_time=start_time, end_time=time.time())
                 await controller.record_trace(info)
+                logger.debug(f"Finished trace for {func}, {info}")
                 return results
             except Exception as e:
                 # If there is an error, we need to record it

flyte/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '2.0.0b17'
-__version_tuple__ = version_tuple = (2, 0, 0, 'b17')
+__version__ = version = '2.0.0b18'
+__version_tuple__ = version_tuple = (2, 0, 0, 'b18')
 
-__commit_id__ = commit_id = 'gfe0ca1266'
+__commit_id__ = commit_id = 'g930faeaea'

flyte/io/_dir.py

@@ -48,6 +48,7 @@ class Dir(BaseModel, Generic[T], SerializableType):
     path: str
     name: Optional[str] = None
     format: str = ""
+    hash: Optional[str] = None
 
     class Config:
         arbitrary_types_allowed = True
@@ -248,13 +249,20 @@ class Dir(BaseModel, Generic[T], SerializableType):
         raise NotImplementedError("Sync download is not implemented for remote paths")
 
     @classmethod
-    async def from_local(cls, local_path: Union[str, Path], remote_path: Optional[str] = None) -> Dir[T]:
+    async def from_local(
+        cls,
+        local_path: Union[str, Path],
+        remote_path: Optional[str] = None,
+        dir_cache_key: Optional[str] = None,
+    ) -> Dir[T]:
         """
         Asynchronously create a new Dir by uploading a local directory to the configured remote store.
 
         Args:
             local_path: Path to the local directory
             remote_path: Optional path to store the directory remotely. If None, a path will be generated.
+            dir_cache_key: If you have a precomputed hash value you want to use when computing cache keys for
+              discoverable tasks that this File is an input to.
 
         Returns:
             A new Dir instance pointing to the uploaded directory
@@ -262,13 +270,34 @@ class Dir(BaseModel, Generic[T], SerializableType):
         Example:
             ```python
             remote_dir = await Dir[DataFrame].from_local('/tmp/data_dir/', 's3://bucket/data/')
+            # With a known hash value you want to use for cache key calculation
+            remote_dir = await Dir[DataFrame].from_local('/tmp/data_dir/', 's3://bucket/data/', dir_cache_key='abc123')
             ```
         """
         local_path_str = str(local_path)
         dirname = os.path.basename(os.path.normpath(local_path_str))
 
         output_path = await storage.put(from_path=local_path_str, to_path=remote_path, recursive=True)
-        return cls(path=output_path, name=dirname)
+        return cls(path=output_path, name=dirname, hash=dir_cache_key)
+
+    @classmethod
+    def from_existing_remote(cls, remote_path: str, dir_cache_key: Optional[str] = None) -> Dir[T]:
+        """
+        Create a Dir reference from an existing remote directory.
+
+        Args:
+            remote_path: The remote path to the existing directory
+            dir_cache_key: Optional hash value to use for cache key computation. If not specified,
+                            the cache key will be computed based on this object's attributes.
+
+        Example:
+            ```python
+            remote_dir = Dir.from_existing_remote("s3://bucket/data/")
+            # With a known hash
+            remote_dir = Dir.from_existing_remote("s3://bucket/data/", dir_cache_key="abc123")
+            ```
+        """
+        return cls(path=remote_path, hash=dir_cache_key)
 
     @classmethod
     def from_local_sync(cls, local_path: Union[str, Path], remote_path: Optional[str] = None) -> Dir[T]:
@@ -414,7 +443,8 @@ class DirTransformer(TypeTransformer[Dir]):
                     ),
                     uri=python_val.path,
                 )
-            )
+            ),
+            hash=python_val.hash if python_val.hash else None,
         )
 
     async def to_python_value(
@@ -432,7 +462,8 @@ class DirTransformer(TypeTransformer[Dir]):
 
         uri = lv.scalar.blob.uri
         filename = Path(uri).name
-        f: Dir = Dir(path=uri, name=filename, format=lv.scalar.blob.metadata.type.format)
+        hash_value = lv.hash if lv.hash else None
+        f: Dir = Dir(path=uri, name=filename, format=lv.scalar.blob.metadata.type.format, hash=hash_value)
         return f
 
     def guess_python_type(self, literal_type: types_pb2.LiteralType) -> Type[Dir]:

flyte/io/_file.py

@@ -5,6 +5,7 @@ from contextlib import asynccontextmanager, contextmanager
 from pathlib import Path
 from typing import (
     IO,
+    Annotated,
     Any,
     AsyncGenerator,
     Dict,
@@ -21,12 +22,14 @@ from flyteidl.core import literals_pb2, types_pb2
 from fsspec.asyn import AsyncFileSystem
 from fsspec.utils import get_protocol
 from mashumaro.types import SerializableType
-from pydantic import BaseModel, model_validator
+from pydantic import BaseModel, Field, model_validator
+from pydantic.json_schema import SkipJsonSchema
 
 import flyte.storage as storage
 from flyte._context import internal_ctx
 from flyte._initialize import requires_initialization
 from flyte._logging import logger
+from flyte.io._hashing_io import AsyncHashingReader, HashingWriter, HashMethod, PrecomputedValue
 from flyte.types import TypeEngine, TypeTransformer, TypeTransformerFailedError
 
 # Type variable for the file format
@@ -104,6 +107,8 @@ class File(BaseModel, Generic[T], SerializableType):
     path: str
     name: Optional[str] = None
     format: str = ""
+    hash: Optional[str] = None
+    hash_method: Annotated[Optional[HashMethod], Field(default=None, exclude=True), SkipJsonSchema()] = None
 
     class Config:
         arbitrary_types_allowed = True
@@ -139,7 +144,7 @@ class File(BaseModel, Generic[T], SerializableType):
 
     @classmethod
     @requires_initialization
-    def new_remote(cls) -> File[T]:
+    def new_remote(cls, hash_method: Optional[HashMethod | str] = None) -> File[T]:
         """
         Create a new File reference for a remote file that will be written to.
 
@@ -155,11 +160,13 @@ class File(BaseModel, Generic[T], SerializableType):
         ```
         """
         ctx = internal_ctx()
+        known_cache_key = hash_method if isinstance(hash_method, str) else None
+        method = hash_method if isinstance(hash_method, HashMethod) else None
 
-        return cls(path=ctx.raw_data.get_random_remote_path())
+        return cls(path=ctx.raw_data.get_random_remote_path(), hash=known_cache_key, hash_method=method)
 
     @classmethod
-    def from_existing_remote(cls, remote_path: str) -> File[T]:
+    def from_existing_remote(cls, remote_path: str, file_cache_key: Optional[str] = None) -> File[T]:
         """
         Create a File reference from an existing remote file.
 
@@ -172,8 +179,10 @@ class File(BaseModel, Generic[T], SerializableType):
 
         Args:
             remote_path: The remote path to the existing file
+            file_cache_key: Optional hash value to use for discovery purposes. If not specified, the value of this
+              File object will be hashed (basically the path, not the contents).
         """
-        return cls(path=remote_path)
+        return cls(path=remote_path, hash=file_cache_key)
 
     @asynccontextmanager
     async def open(
@@ -184,7 +193,7 @@ class File(BaseModel, Generic[T], SerializableType):
         cache_options: Optional[dict] = None,
         compression: Optional[str] = None,
         **kwargs,
-    ) -> AsyncGenerator[IO[Any]]:
+    ) -> AsyncGenerator[Union[IO[Any], "HashingWriter"], None]:
         """
         Asynchronously open the file and return a file-like object.
 
@@ -245,7 +254,15 @@ class File(BaseModel, Generic[T], SerializableType):
                     file_handle.close()
 
             with fs.open(self.path, mode) as file_handle:
-                yield file_handle
+                if self.hash_method and self.hash is None:
+                    logger.debug(f"Wrapping file handle with hashing writer using {self.hash_method}")
+                    fh = HashingWriter(file_handle, accumulator=self.hash_method)
+                    yield fh
+                    self.hash = fh.result()
+                    fh.close()
+                else:
+                    yield file_handle
+                    file_handle.close()
 
     def exists_sync(self) -> bool:
         """
@@ -351,13 +368,22 @@ class File(BaseModel, Generic[T], SerializableType):
 
     @classmethod
     @requires_initialization
-    async def from_local(cls, local_path: Union[str, Path], remote_destination: Optional[str] = None) -> File[T]:
+    async def from_local(
+        cls,
+        local_path: Union[str, Path],
+        remote_destination: Optional[str] = None,
+        hash_method: Optional[HashMethod | str] = None,
+    ) -> File[T]:
         """
         Create a new File object from a local file that will be uploaded to the configured remote store.
 
         Args:
             local_path: Path to the local file
             remote_destination: Optional path to store the file remotely. If None, a path will be generated.
+            hash_method: Pass this argument either as a set string or a HashMethod to use for
+              determining a task's cache key if this File object is used as an input to said task. If not specified,
+              the cache key will just be computed based on this object's attributes (i.e. path, name, format, etc.).
+              If there is a set value you want to use, please pass an instance of the PrecomputedValue HashMethod.
 
         Returns:
             A new File instance pointing to the uploaded file
@@ -376,20 +402,38 @@ class File(BaseModel, Generic[T], SerializableType):
 
         # If remote_destination was not set by the user, and the configured raw data path is also local,
         # then let's optimize by not uploading.
+        hash_value = hash_method if isinstance(hash_method, str) else None
+        hash_method = hash_method if isinstance(hash_method, HashMethod) else None
         if "file" in protocol:
             if remote_destination is None:
                 path = str(Path(local_path).absolute())
             else:
                 # Otherwise, actually make a copy of the file
-                async with aiofiles.open(remote_path, "rb") as src:
-                    async with aiofiles.open(local_path, "wb") as dst:
-                        await dst.write(await src.read())
+                async with aiofiles.open(local_path, "rb") as src:
+                    async with aiofiles.open(remote_path, "wb") as dst:
+                        if hash_method:
+                            dst_wrapper = HashingWriter(dst, accumulator=hash_method)
+                            await dst_wrapper.write(await src.read())
+                            hash_value = dst_wrapper.result()
+                        else:
+                            await dst.write(await src.read())
                 path = str(Path(remote_path).absolute())
         else:
             # Otherwise upload to remote using async storage layer
-            path = await storage.put(str(local_path), remote_path)
+            if hash_method:
+                # We can skip the wrapper if the hash method is just a precomputed value
+                if not isinstance(hash_method, PrecomputedValue):
+                    async with aiofiles.open(local_path, "rb") as src:
+                        src_wrapper = AsyncHashingReader(src, accumulator=hash_method)
+                        path = await storage.put_stream(src_wrapper, to_path=remote_path)
+                        hash_value = src_wrapper.result()
+                else:
+                    path = await storage.put(str(local_path), remote_path)
+                    hash_value = hash_method.result()
+            else:
+                path = await storage.put(str(local_path), remote_path)
 
-        f = cls(path=path, name=filename)
+        f = cls(path=path, name=filename, hash_method=hash_method, hash=hash_value)
         return f
 
 
@@ -432,7 +476,8 @@ class FileTransformer(TypeTransformer[File]):
                     ),
                     uri=python_val.path,
                 )
-            )
+            ),
+            hash=python_val.hash if python_val.hash else None,
         )
 
     async def to_python_value(
@@ -450,7 +495,8 @@ class FileTransformer(TypeTransformer[File]):
 
         uri = lv.scalar.blob.uri
         filename = Path(uri).name
-        f: File = File(path=uri, name=filename, format=lv.scalar.blob.metadata.type.format)
+        hash_value = lv.hash if lv.hash else None
+        f: File = File(path=uri, name=filename, format=lv.scalar.blob.metadata.type.format, hash=hash_value)
         return f
 
     def guess_python_type(self, literal_type: types_pb2.LiteralType) -> Type[File]: