flyte 0.2.0b14__py3-none-any.whl → 0.2.0b16__py3-none-any.whl

flyte/io/_structured_dataset/basic_dfs.py

@@ -27,19 +27,12 @@ else:
 T = TypeVar("T")
 
 
-# pr: add back after storage
-def get_pandas_storage_options(uri: str, data_config=None, anonymous: bool = False) -> typing.Optional[typing.Dict]:
+def get_pandas_storage_options(uri: str, anonymous: bool = False) -> typing.Optional[typing.Dict]:
     from pandas.io.common import is_fsspec_url  # type: ignore
 
     if is_fsspec_url(uri):
         if uri.startswith("s3"):
-            # pr: after storage, replace with real call to get_fsspec_storage_options
-            return {
-                "cache_regions": True,
-                "client_kwargs": {"endpoint_url": "http://localhost:30002"},
-                "key": "minio",
-                "secret": "miniostorage",
-            }
+            return storage.get_configured_fsspec_kwargs("s3", anonymous=anonymous)
         return {}
 
     # Pandas does not allow storage_options for non-fsspec paths e.g. local.
@@ -70,7 +63,7 @@ class PandasToCSVEncodingHandler(StructuredDatasetEncoder):
         df.to_csv(
             path,
             index=False,
-            storage_options=get_pandas_storage_options(uri=path, data_config=None),
+            storage_options=get_pandas_storage_options(uri=path),
         )
         structured_dataset_type.format = CSV
         return literals_pb2.StructuredDataset(
@@ -87,20 +80,21 @@ class CSVToPandasDecodingHandler(StructuredDatasetDecoder):
         proto_value: literals_pb2.StructuredDataset,
         current_task_metadata: literals_pb2.StructuredDatasetMetadata,
     ) -> "pd.DataFrame":
-        from botocore.exceptions import NoCredentialsError
-
         uri = proto_value.uri
         columns = None
-        kwargs = get_pandas_storage_options(uri=uri, data_config=None)
+        kwargs = get_pandas_storage_options(uri=uri)
         path = os.path.join(uri, ".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)
-        except NoCredentialsError:
-            logger.debug("S3 source detected, attempting anonymous S3 access")
-            kwargs = get_pandas_storage_options(uri=uri, data_config=None, anonymous=True)
-            return pd.read_csv(path, usecols=columns, storage_options=kwargs)
+        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)
+            else:
+                raise
 
 
 class PandasToParquetEncodingHandler(StructuredDatasetEncoder):
@@ -128,7 +122,7 @@ class PandasToParquetEncodingHandler(StructuredDatasetEncoder):
             path,
             coerce_timestamps="us",
             allow_truncated_timestamps=False,
-            storage_options=get_pandas_storage_options(uri=path, data_config=None),
+            storage_options=get_pandas_storage_options(uri=path),
         )
         structured_dataset_type.format = PARQUET
         return literals_pb2.StructuredDataset(
@@ -145,19 +139,20 @@ class ParquetToPandasDecodingHandler(StructuredDatasetDecoder):
         flyte_value: literals_pb2.StructuredDataset,
         current_task_metadata: literals_pb2.StructuredDatasetMetadata,
     ) -> "pd.DataFrame":
-        from botocore.exceptions import NoCredentialsError
-
         uri = flyte_value.uri
         columns = None
-        kwargs = get_pandas_storage_options(uri=uri, data_config=None)
+        kwargs = get_pandas_storage_options(uri=uri)
         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_parquet(uri, columns=columns, storage_options=kwargs)
-        except NoCredentialsError:
-            logger.debug("S3 source detected, attempting anonymous S3 access")
-            kwargs = get_pandas_storage_options(uri=uri, data_config=None, anonymous=True)
-            return pd.read_parquet(uri, columns=columns, storage_options=kwargs)
+        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_parquet(uri, columns=columns, storage_options=kwargs)
+            else:
+                raise
 
 
 class ArrowToParquetEncodingHandler(StructuredDatasetEncoder):
@@ -199,7 +194,6 @@ class ParquetToArrowDecodingHandler(StructuredDatasetDecoder):
         current_task_metadata: literals_pb2.StructuredDatasetMetadata,
     ) -> "pa.Table":
         import pyarrow.parquet as pq
-        from botocore.exceptions import NoCredentialsError
 
         uri = proto_value.uri
         if not storage.is_remote(uri):
@@ -211,9 +205,11 @@ class ParquetToArrowDecodingHandler(StructuredDatasetDecoder):
             columns = [c.name for c in current_task_metadata.structured_dataset_type.columns]
         try:
             return pq.read_table(path, columns=columns)
-        except NoCredentialsError as e:
-            logger.debug("S3 source detected, attempting anonymous S3 access")
-            fs = storage.get_underlying_filesystem(path=uri, anonymous=True)
-            if fs is not None:
-                return pq.read_table(path, filesystem=fs, columns=columns)
-            raise e
+        except Exception as exc:
+            if exc.__class__.__name__ == "NoCredentialsError":
+                logger.debug("S3 source detected, attempting anonymous S3 access")
+                fs = storage.get_underlying_filesystem(path=uri, anonymous=True)
+                if fs is not None:
+                    return pq.read_table(path, filesystem=fs, columns=columns)
+            else:
+                raise

flyte/models.py

@@ -5,7 +5,7 @@ import os
 import pathlib
 import tempfile
 from dataclasses import dataclass, field, replace
-from typing import TYPE_CHECKING, Any, Callable, Dict, Literal, Optional, Tuple, Type
+from typing import TYPE_CHECKING, Any, Callable, ClassVar, Dict, Literal, Optional, Tuple, Type
 
 import rich.repr
 
@@ -15,6 +15,8 @@ from flyte._logging import logger
 from flyte._utils.helpers import base36_encode
 
 if TYPE_CHECKING:
+    from flyteidl.core import literals_pb2
+
     from flyte._internal.imagebuild.image_builder import ImageCache
     from flyte.report import Report
 
@@ -231,6 +233,13 @@ class Checkpoints:
     checkpoint_path: str | None
 
 
+class _has_default:
+    """
+    A marker class to indicate that a specific input has a default value or not.
+    This is used to determine if the input is required or not.
+    """
+
+
 @dataclass(frozen=True)
 class NativeInterface:
     """
@@ -240,7 +249,14 @@ class NativeInterface:
 
     inputs: Dict[str, Tuple[Type, Any]]
     outputs: Dict[str, Type]
-    docstring: Optional[Docstring] = field(default=None)
+    docstring: Optional[Docstring] = None
+
+    # This field is used to indicate that the task has a default value for the input, but already in the
+    # remote form.
+    _remote_defaults: Optional[Dict[str, literals_pb2.Literal]] = field(default=None, repr=False)
+
+    has_default: ClassVar[Type[_has_default]] = _has_default  # This can be used to indicate if a specific input
+    # has a default value or not, in the case when the default value is not known. An example would be remote tasks.
 
     def has_outputs(self) -> bool:
         """
@@ -248,12 +264,31 @@ class NativeInterface:
         """
         return self.outputs is not None and len(self.outputs) > 0
 
+    def num_required_inputs(self) -> int:
+        """
+        Get the number of required inputs for the task. This is used to determine how many inputs are required for the
+        task execution.
+        """
+        return sum(1 for t in self.inputs.values() if t[1] is inspect.Parameter.empty)
+
     @classmethod
-    def from_types(cls, inputs: Dict[str, Type], outputs: Dict[str, Type]) -> NativeInterface:
+    def from_types(
+        cls,
+        inputs: Dict[str, Tuple[Type, Type[_has_default] | Type[inspect._empty]]],
+        outputs: Dict[str, Type],
+        default_inputs: Optional[Dict[str, literals_pb2.Literal]] = None,
+    ) -> NativeInterface:
         """
         Create a new NativeInterface from the given types. This is used to create a native interface for the task.
+        :param inputs: A dictionary of input names and their types and a value indicating if they have a default value.
+        :param outputs: A dictionary of output names and their types.
+        :param default_inputs: Optional dictionary of default inputs for remote tasks.
+        :return: A NativeInterface object with the given inputs and outputs.
         """
-        return cls(inputs={k: (v, inspect.Parameter.empty) for k, v in inputs.items()}, outputs=outputs)
+        for k, v in inputs.items():
+            if v[1] is cls.has_default and (default_inputs is None or k not in default_inputs):
+                raise ValueError(f"Input {k} has a default value but no default input provided for remote task.")
+        return cls(inputs=inputs, outputs=outputs, _remote_defaults=default_inputs)
 
     @classmethod
     def from_callable(cls, func: Callable) -> NativeInterface:
@@ -301,7 +336,10 @@ class NativeInterface:
                 tp = tpe[0] if isinstance(tpe[0], str) else tpe[0].__name__
                 i += f"{key}: {tp}"
                 if tpe[1] is not inspect.Parameter.empty:
-                    i += f" = {tpe[1]}"
+                    if tpe[1] is self.has_default:
+                        i += " = ..."
+                    else:
+                        i += f" = {tpe[1]}"
         i += ")"
         if self.outputs:
             initial = True

flyte/remote/_data.py

@@ -16,6 +16,7 @@ from flyteidl.service import dataproxy_pb2
 from google.protobuf import duration_pb2
 
 from flyte._initialize import CommonInit, ensure_client, get_client, get_common_config
+from flyte._logging import make_hyperlink
 from flyte.errors import InitializationError, RuntimeSystemError
 
 _UPLOAD_EXPIRES_IN = timedelta(seconds=60)
@@ -89,7 +90,7 @@ async def _upload_single_file(
             raise RuntimeSystemError(e.code().value, f"Failed to get signed url for {fp}.")
     except Exception as e:
         raise RuntimeSystemError(type(e).__name__, f"Failed to get signed url for {fp}.") from e
-    logger.debug(f"Uploading to signed url {resp.signed_url} for {fp}")
+    logger.debug(f"Uploading to {make_hyperlink('signed url', resp.signed_url)} for {fp}")
     extra_headers = get_extra_headers_for_protocol(resp.native_url)
     extra_headers.update(resp.headers)
     encoded_md5 = b64encode(md5_bytes)

flyte/remote/_run.py

@@ -43,7 +43,10 @@ def _action_time_phase(action: run_definition_pb2.Action | run_definition_pb2.Ac
         yield "run_time", f"{(datetime.now(timezone.utc) - start_time).seconds} secs"
     yield "phase", run_definition_pb2.Phase.Name(action.status.phase)
     if isinstance(action, run_definition_pb2.ActionDetails):
-        yield "error", action.error_info if action.HasField("error_info") else "NA"
+        yield (
+            "error",
+            f"{action.error_info.kind}: {action.error_info.message}" if action.HasField("error_info") else "NA",
+        )
 
 
 def _action_rich_repr(action: run_definition_pb2.Action) -> rich.repr.Result:
@@ -78,7 +81,7 @@ def _action_details_rich_repr(action: run_definition_pb2.ActionDetails) -> rich.
     yield "task_type", action.resolved_task_spec.task_template.type
     yield "task_version", action.resolved_task_spec.task_template.id.version
     yield "attempts", action.attempts
-    yield "error_info", action.error_info if action.HasField("error_info") else "NA"
+    yield "error", f"{action.error_info.kind}: {action.error_info.message}" if action.HasField("error_info") else "NA"
     yield "phase", run_definition_pb2.Phase.Name(action.status.phase)
     yield "group", action.metadata.group
     yield "parent", action.metadata.parent
@@ -639,7 +642,7 @@ class Action:
                             console.print(f"[bold green]Run '{self.name}' completed successfully.[/bold green]")
                         else:
                             console.print(
-                                f"[bold red]Run '{self.name}' exited unsuccessfully in state {ad.phase}"
+                                f"[bold red]Run '{self.name}' exited unsuccessfully in state {ad.phase} "
                                 f"with error: {ad.error_info}[/bold red]"
                             )
                         break
@@ -667,6 +670,8 @@ class Action:
         Rich representation of the Action object.
         """
         yield from _action_rich_repr(self.pb2)
+        if self._details:
+            yield from self._details.__rich_repr__()
 
     def __repr__(self) -> str:
         """

flyte/remote/_task.py

@@ -162,6 +162,20 @@ class TaskDetails:
         """
         return self.pb2.spec.task_template.type
 
+    @property
+    def default_input_args(self) -> Tuple[str, ...]:
+        """
+        The default input arguments of the task.
+        """
+        return tuple(x.name for x in self.pb2.spec.default_inputs)
+
+    @property
+    def required_args(self) -> Tuple[str, ...]:
+        """
+        The required input arguments of the task.
+        """
+        return tuple(x for x, _ in self.interface.inputs.items() if x not in self.default_input_args)
+
     @functools.cached_property
     def interface(self) -> NativeInterface:
         """
@@ -169,7 +183,7 @@ class TaskDetails:
         """
         import flyte.types as types
 
-        return types.guess_interface(self.pb2.spec.task_template.interface)
+        return types.guess_interface(self.pb2.spec.task_template.interface, default_inputs=self.pb2.spec.default_inputs)
 
     @property
     def cache(self) -> flyte.Cache:
@@ -206,6 +220,19 @@ class TaskDetails:
         """
         Forwards the call to the underlying task. The entity will be fetched if not already present
         """
+        # TODO support kwargs, for this we need ordered inputs to be stored in the task spec.
+        if len(args) > 0:
+            raise flyte.errors.ReferenceTaskError(
+                f"Reference task {self.name} does not support positional arguments"
+                f"currently. Please use keyword arguments."
+            )
+        if len(self.required_args) > 0:
+            if len(args) + len(kwargs) < len(self.required_args):
+                raise ValueError(
+                    f"Task {self.name} requires at least {self.required_args} arguments, "
+                    f"but only received args:{args}  kwargs{kwargs}."
+                )
+
         ctx = internal_ctx()
         if ctx.is_task_context():
             # If we are in a task context, that implies we are executing a Run.
@@ -238,6 +265,11 @@ class TaskDetails:
         """
         Rich representation of the task.
         """
+        yield "friendly_name", self.pb2.spec.short_name
+        yield "environment", self.pb2.spec.environment
+        yield "default_inputs_keys", self.default_input_args
+        yield "required_args", self.required_args
+        yield "raw_default_inputs", [str(x) for x in self.pb2.spec.default_inputs]
         yield "project", self.pb2.task_id.project
         yield "domain", self.pb2.task_id.domain
         yield "name", self.name

flyte/storage/__init__.py

@@ -4,6 +4,7 @@ __all__ = [
     "S3",
     "Storage",
     "get",
+    "get_configured_fsspec_kwargs",
     "get_random_local_directory",
     "get_random_local_path",
     "get_stream",
@@ -18,6 +19,7 @@ __all__ = [
 from ._config import ABFS, GCS, S3, Storage
 from ._storage import (
     get,
+    get_configured_fsspec_kwargs,
     get_random_local_directory,
     get_random_local_path,
     get_stream,

flyte/storage/_storage.py

@@ -62,40 +62,51 @@ def get_random_local_directory() -> pathlib.Path:
     return _dir
 
 
-def get_underlying_filesystem(
-    protocol: typing.Optional[str] = None,
-    anonymous: bool = False,
-    path: typing.Optional[str] = None,
-    **kwargs,
-) -> fsspec.AbstractFileSystem:
-    if protocol is None:
-        # If protocol is None, get it from the path
-        protocol = get_protocol(path)
-
-    storage_config = get_storage()
-    if storage_config:
-        kwargs = storage_config.get_fsspec_kwargs(anonymous, **kwargs)
-    elif protocol:
+def get_configured_fsspec_kwargs(
+    protocol: typing.Optional[str] = None, anonymous: bool = False
+) -> typing.Dict[str, typing.Any]:
+    if protocol:
         match protocol:
             case "s3":
                 # If the protocol is s3, we can use the s3 filesystem
                 from flyte.storage import S3
 
-                kwargs = S3.auto().get_fsspec_kwargs(anonymous=anonymous, **kwargs)
+                return S3.auto().get_fsspec_kwargs(anonymous=anonymous)
             case "gs":
                 # If the protocol is gs, we can use the gs filesystem
                 from flyte.storage import GCS
 
-                kwargs = GCS.auto().get_fsspec_kwargs(anonymous=anonymous, **kwargs)
+                return GCS.auto().get_fsspec_kwargs(anonymous=anonymous)
             case "abfs" | "abfss":
                 # If the protocol is abfs or abfss, we can use the abfs filesystem
                 from flyte.storage import ABFS
 
-                kwargs = ABFS.auto().get_fsspec_kwargs(anonymous=anonymous, **kwargs)
+                return ABFS.auto().get_fsspec_kwargs(anonymous=anonymous)
             case _:
-                pass
+                return {}
+
+    # If no protocol, return args from storage config if set
+    storage_config = get_storage()
+    if storage_config:
+        return storage_config.get_fsspec_kwargs(anonymous)
+
+    return {}
+
+
+def get_underlying_filesystem(
+    protocol: typing.Optional[str] = None,
+    anonymous: bool = False,
+    path: typing.Optional[str] = None,
+    **kwargs,
+) -> fsspec.AbstractFileSystem:
+    if protocol is None:
+        # If protocol is None, get it from the path
+        protocol = get_protocol(path)
+
+    configured_kwargs = get_configured_fsspec_kwargs(protocol, anonymous=anonymous)
+    configured_kwargs.update(kwargs)
 
-    return fsspec.filesystem(protocol, **kwargs)
+    return fsspec.filesystem(protocol, **configured_kwargs)
 
 
 def _get_anonymous_filesystem(from_path):

flyte/types/_interface.py

@@ -1,11 +1,15 @@
-from typing import Any, Dict, Type, cast
+import inspect
+from typing import Any, Dict, Iterable, Tuple, Type, cast
 
-from flyteidl.core import interface_pb2
+from flyteidl.core import interface_pb2, literals_pb2
 
+from flyte._protos.workflow import common_pb2
 from flyte.models import NativeInterface
 
 
-def guess_interface(interface: interface_pb2.TypedInterface) -> NativeInterface:
+def guess_interface(
+    interface: interface_pb2.TypedInterface, default_inputs: Iterable[common_pb2.NamedParameter] | None = None
+) -> NativeInterface:
     """
     Returns the interface of the task with guessed types, as types may not be present in current env.
     """
@@ -14,12 +18,23 @@ def guess_interface(interface: interface_pb2.TypedInterface) -> NativeInterface:
     if interface is None:
         return NativeInterface({}, {})
 
-    guessed_inputs: Dict[str, Type[Any]] = {}
+    default_input_literals: Dict[str, literals_pb2.Literal] = {}
+    if default_inputs is not None:
+        for param in default_inputs:
+            if param.parameter.HasField("default"):
+                default_input_literals[param.name] = param.parameter.default
+
+    guessed_inputs: Dict[str, Tuple[Type[Any], Any] | Any] = {}
     if interface.inputs is not None and len(interface.inputs.variables) > 0:
-        guessed_inputs = flyte.types.TypeEngine.guess_python_types(cast(dict, interface.inputs.variables))
+        input_types = flyte.types.TypeEngine.guess_python_types(cast(dict, interface.inputs.variables))
+        for name, t in input_types.items():
+            if name not in default_input_literals:
+                guessed_inputs[name] = (t, inspect.Parameter.empty)
+            else:
+                guessed_inputs[name] = (t, NativeInterface.has_default)
 
     guessed_outputs: Dict[str, Type[Any]] = {}
     if interface.outputs is not None and len(interface.outputs.variables) > 0:
         guessed_outputs = flyte.types.TypeEngine.guess_python_types(cast(dict, interface.outputs.variables))
 
-    return NativeInterface.from_types(guessed_inputs, guessed_outputs)
+    return NativeInterface.from_types(guessed_inputs, guessed_outputs, default_input_literals)

{flyte-0.2.0b14.dist-info → flyte-0.2.0b16.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flyte
-Version: 0.2.0b14
+Version: 0.2.0b16
 Summary: Add your description here
 Author-email: Ketan Umare <kumare3@users.noreply.github.com>
 Requires-Python: >=3.10
@@ -125,7 +125,7 @@ import flyte
 
 env = flyte.TaskEnvironment(
     name="hello_world",
-    image=flyte.Image.auto().with_apt_packages(...).with_pip_packages(...),
+    image=flyte.Image.from_debian_base().with_apt_packages(...).with_pip_packages(...),
 )
 
 ```