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

flyte/_map.py

@@ -1,4 +1,6 @@
 import asyncio
+import functools
+import logging
 from typing import Any, AsyncGenerator, AsyncIterator, Generic, Iterable, Iterator, List, Union, cast
 
 from flyte.syncify import syncify
@@ -11,7 +13,14 @@ from ._task import P, R, TaskTemplate
 class MapAsyncIterator(Generic[P, R]):
     """AsyncIterator implementation for the map function results"""
 
-    def __init__(self, func: TaskTemplate[P, R], args: tuple, name: str, concurrency: int, return_exceptions: bool):
+    def __init__(
+        self,
+        func: TaskTemplate[P, R] | functools.partial[R],
+        args: tuple,
+        name: str,
+        concurrency: int,
+        return_exceptions: bool,
+    ):
         self.func = func
         self.args = args
         self.name = name
@@ -49,13 +58,16 @@ class MapAsyncIterator(Generic[P, R]):
             return result
         except Exception as e:
             self._exception_count += 1
-            logger.debug(f"Task {self._current_index - 1} failed with exception: {e}")
+            logger.debug(
+                f"Task {self._current_index - 1} failed with exception: {e}, return_exceptions={self.return_exceptions}"
+            )
             if self.return_exceptions:
                 return e
             else:
                 # Cancel remaining tasks
                 for remaining_task in self._tasks[self._current_index + 1 :]:
                     remaining_task.cancel()
+                logger.warning("Exception raising is `ON`, raising exception and cancelling remaining tasks")
                 raise e
 
     async def _initialize(self):
@@ -64,10 +76,26 @@ class MapAsyncIterator(Generic[P, R]):
         tasks = []
         task_count = 0
 
-        for arg_tuple in zip(*self.args):
-            task = asyncio.create_task(self.func.aio(*arg_tuple))
-            tasks.append(task)
-            task_count += 1
+        if isinstance(self.func, functools.partial):
+            # Handle partial functions by merging bound args/kwargs with mapped args
+            base_func = cast(TaskTemplate, self.func.func)
+            bound_args = self.func.args
+            bound_kwargs = self.func.keywords or {}
+
+            for arg_tuple in zip(*self.args):
+                # Merge bound positional args with mapped args
+                merged_args = bound_args + arg_tuple
+                if logger.isEnabledFor(logging.DEBUG):
+                    logger.debug(f"Running {base_func.name} with args: {merged_args} and kwargs: {bound_kwargs}")
+                task = asyncio.create_task(base_func.aio(*merged_args, **bound_kwargs))
+                tasks.append(task)
+                task_count += 1
+        else:
+            # Handle regular TaskTemplate functions
+            for arg_tuple in zip(*self.args):
+                task = asyncio.create_task(self.func.aio(*arg_tuple))
+                tasks.append(task)
+                task_count += 1
 
         if task_count == 0:
             logger.info(f"Group '{self.name}' has no tasks to process")
@@ -107,9 +135,46 @@ class _Mapper(Generic[P, R]):
         """Get the name of the group, defaulting to 'map' if not provided."""
         return f"{task_name}_{group_name or 'map'}"
 
+    @staticmethod
+    def validate_partial(func: functools.partial[R]):
+        """
+        This method validates that the provided partial function is valid for mapping, i.e. only the one argument
+        is left for mapping and the rest are provided as keywords or args.
+
+        :param func: partial function to validate
+        :raises TypeError: if the partial function is not valid for mapping
+        """
+        f = cast(TaskTemplate, func.func)
+        inputs = f.native_interface.inputs
+        params = list(inputs.keys())
+        total_params = len(params)
+        provided_args = len(func.args)
+        provided_kwargs = len(func.keywords or {})
+
+        # Calculate how many parameters are left unspecified
+        unspecified_count = total_params - provided_args - provided_kwargs
+
+        # Exactly one parameter should be left for mapping
+        if unspecified_count != 1:
+            raise TypeError(
+                f"Partial function must leave exactly one parameter unspecified for mapping. "
+                f"Found {unspecified_count} unspecified parameters in {f.name}, "
+                f"params: {inputs.keys()}"
+            )
+
+        # Validate that no parameter is both in args and keywords
+        if func.keywords:
+            param_names = list(inputs.keys())
+            for i, arg_name in enumerate(param_names[: provided_args + 1]):
+                if arg_name in func.keywords:
+                    raise TypeError(
+                        f"Parameter '{arg_name}' is provided both as positional argument and keyword argument "
+                        f"in partial function {f.name}."
+                    )
+
     def __call__(
         self,
-        func: TaskTemplate[P, R],
+        func: TaskTemplate[P, R] | functools.partial[R],
         *args: Iterable[Any],
         group_name: str | None = None,
         concurrency: int = 0,
@@ -128,7 +193,13 @@ class _Mapper(Generic[P, R]):
         if not args:
             return
 
-        name = self._get_name(func.name, group_name)
+        if isinstance(func, functools.partial):
+            f = cast(TaskTemplate, func.func)
+            self.validate_partial(func)
+        else:
+            f = cast(TaskTemplate, func)
+
+        name = self._get_name(f.name, group_name)
         logger.debug(f"Blocking Map for {name}")
         with group(name):
             import flyte
@@ -154,7 +225,7 @@ class _Mapper(Generic[P, R]):
                     *args,
                     name=name,
                     concurrency=concurrency,
-                    return_exceptions=True,
+                    return_exceptions=return_exceptions,
                 ),
             ):
                 logger.debug(f"Mapped {x}, task {i}")
@@ -163,7 +234,7 @@ class _Mapper(Generic[P, R]):
 
     async def aio(
         self,
-        func: TaskTemplate[P, R],
+        func: TaskTemplate[P, R] | functools.partial[R],
         *args: Iterable[Any],
         group_name: str | None = None,
         concurrency: int = 0,
@@ -171,7 +242,14 @@ class _Mapper(Generic[P, R]):
     ) -> AsyncGenerator[Union[R, Exception], None]:
         if not args:
             return
-        name = self._get_name(func.name, group_name)
+
+        if isinstance(func, functools.partial):
+            f = cast(TaskTemplate, func.func)
+            self.validate_partial(func)
+        else:
+            f = cast(TaskTemplate, func)
+
+        name = self._get_name(f.name, group_name)
         with group(name):
             import flyte
 
@@ -199,7 +277,7 @@ class _Mapper(Generic[P, R]):
 
 @syncify
 async def _map(
-    func: TaskTemplate[P, R],
+    func: TaskTemplate[P, R] | functools.partial[R],
     *args: Iterable[Any],
     name: str = "map",
     concurrency: int = 0,

flyte/_task.py

@@ -258,6 +258,9 @@ class TaskTemplate(Generic[P, R]):
             else:
                 raise RuntimeSystemError("BadContext", "Controller is not initialized.")
         else:
+            from flyte._logging import logger
+
+            logger.warning(f"Task {self.name} running aio outside of a task context.")
             # Local execute, just stay out of the way, but because .aio is used, we want to return an awaitable,
             # even for synchronous tasks. This is to support migration.
             return self.forward(*args, **kwargs)

flyte/_version.py

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

flyte/cli/_create.py

@@ -98,7 +98,7 @@ def secret(
     "-o",
     "--output",
     type=click.Path(exists=False, writable=True),
-    default=Path.cwd() / "config.yaml",
+    default=Path.cwd() / ".flyte" / "config.yaml",
     help="Path to the output directory where the configuration will be saved. Defaults to current directory.",
     show_default=True,
 )
@@ -147,6 +147,9 @@ def config(
 
     output_path = Path(output)
 
+    if not output_path.parent.exists():
+        output_path.parent.mkdir(parents=True)
+
     if output_path.exists() and not force:
         force = click.confirm(f"Overwrite [{output_path}]?", default=False)
         if not force:

flyte/cli/_deploy.py

@@ -4,8 +4,7 @@ from pathlib import Path
 from types import ModuleType
 from typing import Any, Dict, List, cast, get_args
 
-import click
-from click import Context
+import rich_click as click
 
 import flyte
 
@@ -87,14 +86,14 @@ class DeployArguments:
         return [common.get_option_from_metadata(f.metadata) for f in fields(cls) if f.metadata]
 
 
-class DeployEnvCommand(click.Command):
+class DeployEnvCommand(click.RichCommand):
     def __init__(self, env_name: str, env: Any, deploy_args: DeployArguments, *args, **kwargs):
         self.env_name = env_name
         self.env = env
         self.deploy_args = deploy_args
         super().__init__(*args, **kwargs)
 
-    def invoke(self, ctx: Context):
+    def invoke(self, ctx: click.Context):
         from rich.console import Console
 
         console = Console()
@@ -125,7 +124,7 @@ class DeployEnvRecursiveCommand(click.Command):
         self.deploy_args = deploy_args
         super().__init__(*args, **kwargs)
 
-    def invoke(self, ctx: Context):
+    def invoke(self, ctx: click.Context):
         from rich.console import Console
 
         from flyte._environment import list_loaded_environments

flyte/cli/_params.py

@@ -283,13 +283,17 @@ class UnionParamType(click.ParamType):
     A composite type that allows for multiple types to be specified. This is used for union types.
     """
 
-    def __init__(self, types: typing.List[click.ParamType]):
+    def __init__(self, types: typing.List[click.ParamType | None]):
         super().__init__()
         self._types = self._sort_precedence(types)
-        self.name = "|".join([t.name for t in self._types])
+        self.name = "|".join([t.name for t in self._types if t is not None])
+        self.optional = False
+        if None in types:
+            self.name = f"Optional[{self.name}]"
+            self.optional = True
 
     @staticmethod
-    def _sort_precedence(tp: typing.List[click.ParamType]) -> typing.List[click.ParamType]:
+    def _sort_precedence(tp: typing.List[click.ParamType | None]) -> typing.List[click.ParamType]:
         unprocessed = []
         str_types = []
         others = []
@@ -311,6 +315,8 @@ class UnionParamType(click.ParamType):
         """
         for p in self._types:
             try:
+                if p is None and value is None:
+                    return None
                 return p.convert(value, param, ctx)
             except Exception as e:
                 logger.debug(f"Ignoring conversion error for type {p} trying other variants in Union. Error: {e}")
@@ -433,7 +439,10 @@ def literal_type_to_click_type(lt: LiteralType, python_type: typing.Type) -> cli
         for i in range(len(lt.union_type.variants)):
             variant = lt.union_type.variants[i]
             variant_python_type = typing.get_args(python_type)[i]
-            cts.append(literal_type_to_click_type(variant, variant_python_type))
+            if variant_python_type is type(None):
+                cts.append(None)
+            else:
+                cts.append(literal_type_to_click_type(variant, variant_python_type))
         return UnionParamType(cts)
 
     if lt.HasField("enum_type"):
@@ -461,6 +470,9 @@ class FlyteLiteralConverter(object):
     def is_bool(self) -> bool:
         return self.click_type == click.BOOL
 
+    def is_optional(self) -> bool:
+        return isinstance(self.click_type, UnionParamType) and self.click_type.optional
+
     def convert(
         self, ctx: click.Context, param: typing.Optional[click.Parameter], value: typing.Any
     ) -> typing.Union[Literal, typing.Any]:
@@ -525,6 +537,8 @@ def to_click_option(
     if literal_converter.is_bool():
         required = False
         is_flag = True
+    if literal_converter.is_optional():
+        required = False
 
     return click.Option(
         param_decls=[f"--{input_name}"],

flyte/cli/_run.py

@@ -112,7 +112,7 @@ class RunArguments:
         return [common.get_option_from_metadata(f.metadata) for f in fields(cls) if f.metadata]
 
 
-class RunTaskCommand(click.Command):
+class RunTaskCommand(click.RichCommand):
     def __init__(self, obj_name: str, obj: Any, run_args: RunArguments, *args, **kwargs):
         self.obj_name = obj_name
         self.obj = cast(TaskTemplate, obj)
@@ -196,7 +196,7 @@ class TaskPerFileGroup(common.ObjectsPerFileGroup):
         )
 
 
-class RunReferenceTaskCommand(click.Command):
+class RunReferenceTaskCommand(click.RichCommand):
     def __init__(self, task_name: str, run_args: RunArguments, version: str | None, *args, **kwargs):
         self.task_name = task_name
         self.run_args = run_args

flyte/config/_config.py

@@ -192,7 +192,7 @@ class Config(object):
         )
 
     @classmethod
-    def auto(cls, config_file: typing.Union[str, ConfigFile, None] = None) -> "Config":
+    def auto(cls, config_file: typing.Union[str, pathlib.Path, ConfigFile, None] = None) -> "Config":
         """
         Automatically constructs the Config Object. The order of precedence is as follows
           1. first try to find any env vars that match the config vars specified in the FLYTE_CONFIG format.
@@ -225,7 +225,7 @@ def set_if_exists(d: dict, k: str, val: typing.Any) -> dict:
     return d
 
 
-def auto(config_file: typing.Union[str, ConfigFile, None] = None) -> Config:
+def auto(config_file: typing.Union[str, pathlib.Path, ConfigFile, None] = None) -> Config:
     """
     Automatically constructs the Config Object. The order of precedence is as follows
       1. If specified, read the config from the provided file path.

flyte/config/_reader.py

@@ -108,7 +108,7 @@ class ConfigFile(object):
         return pathlib.Path(self._location)
 
     @staticmethod
-    def _read_yaml_config(location: str) -> typing.Optional[typing.Dict[str, typing.Any]]:
+    def _read_yaml_config(location: str | pathlib.Path) -> typing.Optional[typing.Dict[str, typing.Any]]:
         with open(location, "r") as fh:
             try:
                 yaml_contents = yaml.safe_load(fh)
@@ -139,16 +139,22 @@ def resolve_config_path() -> pathlib.Path | None:
     """
     Config is read from the following locations in order of precedence:
     1. ./config.yaml if it exists
-    2. `UCTL_CONFIG` environment variable
-    3. `FLYTECTL_CONFIG` environment variable
-    4. ~/.union/config.yaml if it exists
-    5. ~/.flyte/config.yaml if it exists
+    2. ./.flyte/config.yaml if it exists
+    3. `UCTL_CONFIG` environment variable
+    4. `FLYTECTL_CONFIG` environment variable
+    5. ~/.union/config.yaml if it exists
+    6. ~/.flyte/config.yaml if it exists
     """
     current_location_config = Path("config.yaml")
     if current_location_config.exists():
         return current_location_config
     logger.debug("No ./config.yaml found")
 
+    dot_flyte_config = Path(".flyte", "config.yaml")
+    if dot_flyte_config.exists():
+        return dot_flyte_config
+    logger.debug("No ./.flyte/config.yaml found")
+
     uctl_path_from_env = getenv(UCTL_CONFIG_ENV_VAR, None)
     if uctl_path_from_env:
         return pathlib.Path(uctl_path_from_env)
@@ -173,13 +179,13 @@ def resolve_config_path() -> pathlib.Path | None:
 
 
 @lru_cache
-def get_config_file(c: typing.Union[str, ConfigFile, None]) -> ConfigFile | None:
+def get_config_file(c: typing.Union[str, pathlib.Path, ConfigFile, None]) -> ConfigFile | None:
     """
     Checks if the given argument is a file or a configFile and returns a loaded configFile else returns None
     """
-    if isinstance(c, str):
+    if isinstance(c, (str, pathlib.Path)):
         logger.debug(f"Using specified config file at {c}")
-        return ConfigFile(c)
+        return ConfigFile(str(c))
     elif isinstance(c, ConfigFile):
         return c
     config_path = resolve_config_path()

flyte/errors.py

@@ -132,7 +132,9 @@ class CustomError(RuntimeUserError):
         Create a CustomError from an exception. The exception's class name is used as the error code and the exception
         message is used as the error message.
         """
-        return cls(e.__class__.__name__, str(e))
+        new_exc = cls(e.__class__.__name__, str(e))
+        new_exc.__cause__ = e
+        return new_exc
 
 
 class NotInTaskContextError(RuntimeUserError):
@@ -221,3 +223,12 @@ class RunAbortedError(RuntimeUserError):
 
     def __init__(self, message: str):
         super().__init__("RunAbortedError", message, "user")
+
+
+class SlowDownError(RuntimeUserError):
+    """
+    This error is raised when the user tries to access a resource that does not exist or is invalid.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("SlowDownError", message, "user")

flyte/git/__init__.py

@@ -0,0 +1,3 @@
+from ._config import config_from_root
+
+__all__ = ["config_from_root"]

flyte/git/_config.py

@@ -0,0 +1,17 @@
+import pathlib
+import subprocess
+
+import flyte.config
+
+
+def config_from_root(path: pathlib.Path | str = ".flyte/config.yaml") -> flyte.config.Config:
+    """Get the config file from the git root directory.
+
+    By default, the config file is expected to be in `.flyte/config.yaml` in the git root directory.
+    """
+
+    result = subprocess.run(["git", "rev-parse", "--show-toplevel"], check=False, capture_output=True, text=True)
+    if result.returncode != 0:
+        raise RuntimeError(f"Failed to get git root directory: {result.stderr}")
+    root = pathlib.Path(result.stdout.strip())
+    return flyte.config.auto(root / path)

flyte/io/_dataframe/basic_dfs.py

@@ -58,16 +58,16 @@ class PandasToCSVEncodingHandler(DataFrameEncoder):
 
         if not storage.is_remote(uri):
             Path(uri).mkdir(parents=True, exist_ok=True)
-        path = os.path.join(uri, ".csv")
+        csv_file = storage.join(uri, "data.csv")
         df = typing.cast(pd.DataFrame, dataframe.val)
         df.to_csv(
-            path,
+            csv_file,
             index=False,
-            storage_options=get_pandas_storage_options(uri=path),
+            storage_options=get_pandas_storage_options(uri=csv_file),
         )
         structured_dataset_type.format = CSV
         return literals_pb2.StructuredDataset(
-            uri=uri, metadata=literals_pb2.StructuredDatasetMetadata(structured_dataset_type)
+            uri=uri, metadata=literals_pb2.StructuredDatasetMetadata(structured_dataset_type=structured_dataset_type)
         )
 
 
@@ -83,16 +83,25 @@ class CSVToPandasDecodingHandler(DataFrameDecoder):
         uri = proto_value.uri
         columns = None
         kwargs = get_pandas_storage_options(uri=uri)
-        path = os.path.join(uri, ".csv")
+        csv_file = storage.join(uri, "data.csv")
         if current_task_metadata.structured_dataset_type and current_task_metadata.structured_dataset_type.columns:
             columns = [c.name for c in current_task_metadata.structured_dataset_type.columns]
         try:
-            return pd.read_csv(path, usecols=columns, storage_options=kwargs)
+            import io
+
+            # The pattern used here is a bit wonky because of obstore issues with csv, getting early eof error.
+            buf = io.BytesIO()
+            async for chunk in storage.get_stream(csv_file):
+                buf.write(chunk)
+            buf.seek(0)
+            df = pd.read_csv(buf)
+            return df
+
         except Exception as exc:
             if exc.__class__.__name__ == "NoCredentialsError":
                 logger.debug("S3 source detected, attempting anonymous S3 access")
                 kwargs = get_pandas_storage_options(uri=uri, anonymous=True)
-                return pd.read_csv(path, usecols=columns, storage_options=kwargs)
+                return pd.read_csv(csv_file, usecols=columns, storage_options=kwargs)
             else:
                 raise