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

flyte/io/_hashing_io.py

@@ -0,0 +1,342 @@
+from __future__ import annotations
+
+import hashlib
+import inspect
+from typing import Any, Iterable, Optional, Protocol, Union, runtime_checkable
+
+
+@runtime_checkable
+class HashMethod(Protocol):
+    def update(self, data: memoryview, /) -> None: ...
+    def result(self) -> str: ...
+
+    # Optional convenience; not required by the writers.
+    def reset(self) -> None: ...
+
+
+class PrecomputedValue(HashMethod):
+    def __init__(self, value: str):
+        self._value = value
+
+    def update(self, data: memoryview, /) -> None: ...
+
+    def result(self) -> str:
+        return self._value
+
+
+class HashlibAccumulator(HashMethod):
+    """
+    Wrap a hashlib-like object to the Accumulator protocol.
+    h = hashlib.new("sha256")
+    acc = HashlibAccumulator(h)
+    """
+
+    def __init__(self, h):
+        self._h = h
+
+    def update(self, data: memoryview, /) -> None:
+        self._h.update(data)
+
+    def result(self) -> Any:
+        return self._h.hexdigest()
+
+    @classmethod
+    def from_hash_name(cls, name: str) -> HashlibAccumulator:
+        """
+        Create an accumulator from a hashlib algorithm name.
+        """
+        h = hashlib.new(name)
+        return cls(h)
+
+
+class HashingWriter:
+    """
+    Sync writer that updates a user-supplied accumulator on every write.
+
+    Hashing covers the exact bytes you pass in. If you write str, it is encoded
+    using the underlying file's .encoding if available, else UTF-8 (for hashing only).
+    """
+
+    def __init__(
+        self,
+        fh,
+        accumulator: HashMethod,
+        *,
+        encoding: Optional[str] = None,
+        errors: str = "strict",
+    ):
+        self._fh = fh
+        self._acc: HashMethod = accumulator
+        self._encoding = encoding or getattr(fh, "encoding", None)
+        self._errors = errors
+
+    def result(self) -> Any:
+        return self._acc.result()
+
+    def _to_bytes_mv(self, data) -> memoryview:
+        if isinstance(data, str):
+            b = data.encode(self._encoding or "utf-8", self._errors)
+            return memoryview(b)
+        if isinstance(data, (bytes, bytearray, memoryview)):
+            return memoryview(data)
+        # Accept any buffer-protocol object (e.g., numpy arrays)
+        return memoryview(data)
+
+    def write(self, data):
+        mv = self._to_bytes_mv(data)
+        self._acc.update(mv)
+        return self._fh.write(data)
+
+    def writelines(self, lines: Iterable[Union[str, bytes, bytearray, memoryview]]):
+        for line in lines:
+            self.write(line)
+
+    def flush(self):
+        return self._fh.flush()
+
+    def close(self):
+        return self._fh.close()
+
+    def __getattr__(self, name):
+        return getattr(self._fh, name)
+
+
+class AsyncHashingWriter:
+    """
+    Async version of HashingWriter with the same behavior.
+    """
+
+    def __init__(
+        self,
+        fh,
+        accumulator: HashMethod,
+        *,
+        encoding: Optional[str] = None,
+        errors: str = "strict",
+    ):
+        self._fh = fh
+        self._acc = accumulator
+        self._encoding = encoding or getattr(fh, "encoding", None)
+        self._errors = errors
+
+    def result(self) -> Any:
+        return self._acc.result()
+
+    def _to_bytes_mv(self, data) -> memoryview:
+        if isinstance(data, str):
+            b = data.encode(self._encoding or "utf-8", self._errors)
+            return memoryview(b)
+        if isinstance(data, (bytes, bytearray, memoryview)):
+            return memoryview(data)
+        return memoryview(data)
+
+    async def write(self, data):
+        mv = self._to_bytes_mv(data)
+        self._acc.update(mv)
+        return await self._fh.write(data)
+
+    async def writelines(self, lines: Iterable[Union[str, bytes, bytearray, memoryview]]):
+        for line in lines:
+            await self.write(line)
+
+    async def flush(self):
+        fn = getattr(self._fh, "flush", None)
+        if fn is None:
+            return None
+        res = fn()
+        if inspect.isawaitable(res):
+            return await res
+        return res
+
+    async def close(self):
+        fn = getattr(self._fh, "close", None)
+        if fn is None:
+            return None
+        res = fn()
+        if inspect.isawaitable(res):
+            return await res
+        return res
+
+    def __getattr__(self, name):
+        return getattr(self._fh, name)
+
+
+class HashingReader:
+    """
+    Sync reader that updates a user-supplied accumulator on every read operation.
+
+    If the underlying handle returns str (text mode), we encode it for hashing only,
+    using the handle's .encoding if present, else the explicit 'encoding' arg, else UTF-8.
+    """
+
+    def __init__(
+        self,
+        fh,
+        accumulator: HashMethod,
+        *,
+        encoding: Optional[str] = None,
+        errors: str = "strict",
+    ):
+        self._fh = fh
+        self._acc = accumulator
+        self._encoding = encoding or getattr(fh, "encoding", None)
+        self._errors = errors
+
+    def result(self) -> str:
+        return self._acc.result()
+
+    def _to_bytes_mv(self, data) -> Optional[memoryview]:
+        if data is None:
+            return None
+        if isinstance(data, str):
+            return memoryview(data.encode(self._encoding or "utf-8", self._errors))
+        if isinstance(data, (bytes, bytearray, memoryview)):
+            return memoryview(data)
+        # Accept any buffer-protocol object (rare for read paths, but safe)
+        return memoryview(data)
+
+    def read(self, size: int = -1):
+        data = self._fh.read(size)
+        mv = self._to_bytes_mv(data)
+        if mv is not None and len(mv) > 0:
+            self._acc.update(mv)
+        return data
+
+    def readline(self, size: int = -1):
+        line = self._fh.readline(size)
+        mv = self._to_bytes_mv(line)
+        if mv is not None and len(mv) > 0:
+            self._acc.update(mv)
+        return line
+
+    def readlines(self, hint: int = -1):
+        lines = self._fh.readlines(hint)
+        # Update in order to reflect exact concatenation
+        for line in lines:
+            mv = self._to_bytes_mv(line)
+            if mv is not None and len(mv) > 0:
+                self._acc.update(mv)
+        return lines
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        # Delegate to the underlying iterator to preserve semantics (including buffering),
+        # but intercept the produced line to hash it.
+        line = next(self._fh)
+        mv = self._to_bytes_mv(line)
+        if mv is not None and len(mv) > 0:
+            self._acc.update(mv)
+        return line
+
+    # ---- passthrough ----
+    def __getattr__(self, name):
+        # Avoid leaking private lookups to the underlying object if we're missing something internal
+        if name.startswith("_"):
+            raise AttributeError(name)
+        return getattr(self._fh, name)
+
+
+class AsyncHashingReader:
+    """
+    Async reader that updates a user-supplied accumulator on every read operation.
+
+    Works with aiofiles/fsspec async handles. `flush`/`close` may be awaitable or sync.
+    """
+
+    def __init__(
+        self,
+        fh,
+        accumulator: HashMethod,
+        *,
+        encoding: Optional[str] = None,
+        errors: str = "strict",
+    ):
+        self._fh = fh
+        self._acc = accumulator
+        self._encoding = encoding or getattr(fh, "encoding", None)
+        self._errors = errors
+
+    def result(self) -> str:
+        return self._acc.result()
+
+    def _to_bytes_mv(self, data) -> Optional[memoryview]:
+        if data is None:
+            return None
+        if isinstance(data, str):
+            return memoryview(data.encode(self._encoding or "utf-8", self._errors))
+        if isinstance(data, (bytes, bytearray, memoryview)):
+            return memoryview(data)
+        return memoryview(data)
+
+    async def read(self, size: int = -1):
+        data = await self._fh.read(size)
+        mv = self._to_bytes_mv(data)
+        if mv is not None and len(mv) > 0:
+            self._acc.update(mv)
+        return data
+
+    async def readline(self, size: int = -1):
+        line = await self._fh.readline(size)
+        mv = self._to_bytes_mv(line)
+        if mv is not None and len(mv) > 0:
+            self._acc.update(mv)
+        return line
+
+    async def readlines(self, hint: int = -1):
+        # Some async filehandles implement readlines(); if not, fall back to manual loop.
+        if hasattr(self._fh, "readlines"):
+            lines = await self._fh.readlines(hint)
+            for line in lines:
+                mv = self._to_bytes_mv(line)
+                if mv is not None and len(mv) > 0:
+                    self._acc.update(mv)
+            return lines
+        # Fallback: read all via iteration
+        lines = []
+        async for line in self:
+            lines.append(line)
+        return lines
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        # Prefer the underlying async iterator if present
+        anext_fn = getattr(self._fh, "__anext__", None)
+        if anext_fn is not None:
+            try:
+                line = await anext_fn()
+            except StopAsyncIteration:
+                raise
+            mv = self._to_bytes_mv(line)
+            if mv is not None and len(mv) > 0:
+                self._acc.update(mv)
+            return line
+
+        # Fallback to readline-based iteration
+        line = await self.readline()
+        if line == "" or line == b"":
+            raise StopAsyncIteration
+        return line
+
+    async def flush(self):
+        fn = getattr(self._fh, "flush", None)
+        if fn is None:
+            return None
+        res = fn()
+        return await res if inspect.isawaitable(res) else res
+
+    async def close(self):
+        fn = getattr(self._fh, "close", None)
+        if fn is None:
+            return None
+        res = fn()
+        return await res if inspect.isawaitable(res) else res
+
+    # ---- passthrough ----
+    def __getattr__(self, name):
+        if name.startswith("_"):
+            raise AttributeError(name)
+        return getattr(self._fh, name)

flyte/models.py

@@ -3,17 +3,18 @@ from __future__ import annotations
 import inspect
 import os
 import pathlib
+import typing
 from dataclasses import dataclass, field, replace
 from typing import TYPE_CHECKING, Any, Callable, ClassVar, Dict, List, Literal, Optional, Tuple, Type
 
 import rich.repr
 
 from flyte._docstring import Docstring
-from flyte._interface import extract_return_annotation
+from flyte._interface import extract_return_annotation, literal_to_enum
 from flyte._logging import logger
 
 if TYPE_CHECKING:
-    from flyteidl.core import literals_pb2
+    from flyteidl2.core import literals_pb2
 
     from flyte._internal.imagebuild.image_builder import ImageCache
     from flyte.report import Report
@@ -76,6 +77,37 @@ class ActionID:
         return self.new_sub_action(new_name)
 
 
+@rich.repr.auto
+@dataclass
+class PathRewrite:
+    """
+    Configuration for rewriting paths during input loading.
+    """
+
+    # If set, rewrites any path starting with this prefix to the new prefix.
+    old_prefix: str
+    new_prefix: str
+
+    def __post_init__(self):
+        if not self.old_prefix or not self.new_prefix:
+            raise ValueError("Both old_prefix and new_prefix must be non-empty strings.")
+        if self.old_prefix == self.new_prefix:
+            raise ValueError("old_prefix and new_prefix must be different.")
+
+    @classmethod
+    def from_str(cls, pattern: str) -> PathRewrite:
+        """
+        Create a PathRewrite from a string pattern of the form `old_prefix->new_prefix`.
+        """
+        parts = pattern.split("->")
+        if len(parts) != 2:
+            raise ValueError(f"Invalid path rewrite pattern: {pattern}. Expected format 'old_prefix->new_prefix'.")
+        return cls(old_prefix=parts[0], new_prefix=parts[1])
+
+    def __repr__(self) -> str:
+        return f"{self.old_prefix}->{self.new_prefix}"
+
+
 @rich.repr.auto
 @dataclass(frozen=True, kw_only=True)
 class RawDataPath:
@@ -85,6 +117,7 @@ class RawDataPath:
     """
 
     path: str
+    path_rewrite: Optional[PathRewrite] = None
 
     @classmethod
     def from_local_folder(cls, local_folder: str | pathlib.Path | None = None) -> RawDataPath:
@@ -111,7 +144,7 @@ class RawDataPath:
 
     def get_random_remote_path(self, file_name: Optional[str] = None) -> str:
         """
-        Returns a random path for uploading a file/directory to.
+        Returns a random path for uploading a file/directory to. This file/folder will not be created, it's just a path.
 
         :param file_name: If given, will be joined after a randomly generated portion.
         :return:
@@ -127,13 +160,14 @@ class RawDataPath:
 
         protocol = get_protocol(file_prefix)
         if "file" in protocol:
-            local_path = pathlib.Path(file_prefix) / random_string
+            parent_folder = pathlib.Path(file_prefix)
+            parent_folder.mkdir(exist_ok=True, parents=True)
             if file_name:
-                # Only if file name is given do we create the parent, because it may be needed as a folder otherwise
-                local_path = local_path / file_name
-                if not local_path.exists():
-                    local_path.parent.mkdir(exist_ok=True, parents=True)
-                    local_path.touch()
+                random_folder = parent_folder / random_string
+                random_folder.mkdir()
+                local_path = random_folder / file_name
+            else:
+                local_path = parent_folder / random_string
             return str(local_path.absolute())
 
         fs = fsspec.filesystem(protocol)
@@ -161,11 +195,14 @@ class TaskContext:
     :param action: The action ID of the current execution. This is always set, within a run.
     :param version: The version of the executed task. This is set when the task is executed by an action and will be
       set on all sub-actions.
+    :param custom_context: Context metadata for the action. If an action receives context, it'll automatically pass it
+      to any actions it spawns. Context will not be used for cache key computation.
     """
 
     action: ActionID
     version: str
     raw_data_path: RawDataPath
+    input_path: str | None = None
     output_path: str
     run_base_dir: str
     report: Report
@@ -175,6 +212,8 @@ class TaskContext:
     compiled_image_cache: ImageCache | None = None
     data: Dict[str, Any] = field(default_factory=dict)
     mode: Literal["local", "remote", "hybrid"] = "remote"
+    interactive_mode: bool = False
+    custom_context: Dict[str, str] = field(default_factory=dict)
 
     def replace(self, **kwargs) -> TaskContext:
         if "data" in kwargs:
@@ -316,10 +355,18 @@ class NativeInterface:
         """
         Extract the native interface from the given function. This is used to create a native interface for the task.
         """
+        # Get function parameters, defaults, varargs info (POSITIONAL_ONLY, VAR_POSITIONAL, KEYWORD_ONLY, etc.).
         sig = inspect.signature(func)
 
         # Extract parameter details (name, type, default value)
         param_info = {}
+        try:
+            # Get fully evaluated, real Python types for type checking.
+            hints = typing.get_type_hints(func, include_extras=True)
+        except Exception as e:
+            logger.warning(f"Could not get type hints for function {func.__name__}: {e}")
+            raise
+
         for name, param in sig.parameters.items():
             if param.kind in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD):
                 raise ValueError(f"Function {func.__name__} cannot have variable positional or keyword arguments.")
@@ -327,10 +374,14 @@ class NativeInterface:
                 logger.warning(
                     f"Function {func.__name__} has parameter {name} without type annotation. Data will be pickled."
                 )
-            param_info[name] = (param.annotation, param.default)
+            arg_type = hints.get(name, param.annotation)
+            if typing.get_origin(arg_type) is Literal:
+                param_info[name] = (literal_to_enum(arg_type), param.default)
+            else:
+                param_info[name] = (arg_type, param.default)
 
         # Get return type
-        outputs = extract_return_annotation(sig.return_annotation)
+        outputs = extract_return_annotation(hints.get("return", sig.return_annotation))
         return cls(inputs=param_info, outputs=outputs)
 
     def convert_to_kwargs(self, *args, **kwargs) -> Dict[str, Any]:
@@ -340,9 +391,15 @@ class NativeInterface:
         """
         # Convert positional arguments to keyword arguments
         if len(args) > len(self.inputs):
-            raise ValueError(f"Too many positional arguments provided, inputs {self.inputs.keys()}, args {len(args)}")
+            raise ValueError(
+                f"Too many positional arguments provided, expected inputs {self.inputs.keys()}, args {len(args)}"
+            )
         for arg, input_name in zip(args, self.inputs.keys()):
             kwargs[input_name] = arg
+        if len(kwargs) > len(self.inputs):
+            raise ValueError(
+                f"Too many keyword arguments provided, expected inputs {self.inputs.keys()}, args {kwargs.keys()}"
+            )
         return kwargs
 
     def get_input_types(self) -> Dict[str, Type]:
@@ -408,13 +465,15 @@ class SerializationContext:
     code_bundle: Optional[CodeBundle] = None
     input_path: str = "{{.input}}"
     output_path: str = "{{.outputPrefix}}"
-    _entrypoint_path: str = field(default="_bin/runtime.py", init=False)
+    interpreter_path: str = "/opt/venv/bin/python"
     image_cache: ImageCache | None = None
     root_dir: Optional[pathlib.Path] = None
 
-    def get_entrypoint_path(self, interpreter_path: str) -> str:
+    def get_entrypoint_path(self, interpreter_path: Optional[str] = None) -> str:
         """
         Get the entrypoint path for the task. This is used to determine the entrypoint for the task execution.
         :param interpreter_path: The path to the interpreter (python)
         """
-        return os.path.join(os.path.dirname(os.path.dirname(interpreter_path)), self._entrypoint_path)
+        if interpreter_path is None:
+            interpreter_path = self.interpreter_path
+        return os.path.join(os.path.dirname(interpreter_path), "runtime.py")

flyte/remote/__init__.py

@@ -7,12 +7,15 @@ __all__ = [
     "ActionDetails",
     "ActionInputs",
     "ActionOutputs",
+    "Phase",
     "Project",
     "Run",
     "RunDetails",
     "Secret",
     "SecretTypes",
     "Task",
+    "Trigger",
+    "User",
     "create_channel",
     "upload_dir",
     "upload_file",
@@ -22,6 +25,8 @@ from ._action import Action, ActionDetails, ActionInputs, ActionOutputs
 from ._client.auth import create_channel
 from ._data import upload_dir, upload_file
 from ._project import Project
-from ._run import Run, RunDetails
+from ._run import Phase, Run, RunDetails
 from ._secret import Secret, SecretTypes
 from ._task import Task
+from ._trigger import Trigger
+from ._user import User