flyte 0.2.0b32__py3-none-any.whl → 0.2.0b34__py3-none-any.whl

flyte/_code_bundle/_utils.py

@@ -212,11 +212,11 @@ def list_imported_modules_as_files(source_path: str, modules: List[ModuleType])
     from flyte._utils.lazy_module import is_imported
 
     files = []
-    union_root = os.path.dirname(flyte.__file__)
+    flyte_root = os.path.dirname(flyte.__file__)
 
     # These directories contain installed packages or modules from the Python standard library.
     # If a module is from these directories, then they are not user files.
-    invalid_directories = [union_root, sys.prefix, sys.base_prefix, site.getusersitepackages(), *site.getsitepackages()]
+    invalid_directories = [flyte_root, sys.prefix, sys.base_prefix, site.getusersitepackages(), *site.getsitepackages()]
 
     for mod in modules:
         # Be careful not to import a module with the .__file__ call if not yet imported.

flyte/_image.py

@@ -545,8 +545,6 @@ class Image:
         if not script.suffix == ".py":
             raise ValueError(f"UV script {script} must have a .py extension")
         header = parse_uv_script_file(script)
-        if registry is None:
-            raise ValueError("registry must be specified")
 
         # todo: arch
         img = cls.from_debian_base(registry=registry, name=name, python_version=python_version, platform=platform)

flyte/_internal/imagebuild/remote_builder.py

@@ -10,6 +10,7 @@ from uuid import uuid4
 import click
 
 import flyte
+import flyte.errors
 from flyte import Image, remote
 from flyte._image import (
     AptPackages,
@@ -53,7 +54,7 @@ class RemoteImageChecker(ImageChecker):
             )
         except Exception as e:
             msg = "remote image builder is not enabled. Please contact Union support to enable it."
-            raise click.ClickException(msg) from e
+            raise flyte.errors.ImageBuildError(msg) from e
 
         image_name = f"{repository.split('/')[-1]}:{tag}"
 
@@ -115,7 +116,7 @@ class RemoteImageBuilder(ImageBuilder):
         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"))
         else:
-            raise click.ClickException(f"❌ Build failed in {elapsed} at {click.style(run.url, fg='cyan')}")
+            raise flyte.errors.ImageBuildError(f"❌ Build failed in {elapsed} at {click.style(run.url, fg='cyan')}")
 
         outputs = await run_details.outputs()
         return _get_fully_qualified_image_name(outputs)

flyte/_task.py

@@ -18,6 +18,7 @@ from typing import (
     TypeAlias,
     TypeVar,
     Union,
+    cast,
 )
 
 from flyte._pod import PodTemplate
@@ -90,7 +91,7 @@ class TaskTemplate(Generic[P, R]):
     cache: CacheRequest = "auto"
     interruptable: bool = False
     retries: Union[int, RetryStrategy] = 0
-    reusable: Union[ReusePolicy, Literal["auto"], None] = None
+    reusable: Union[ReusePolicy, None] = None
     docs: Optional[Documentation] = None
     env: Optional[Dict[str, str]] = None
     secrets: Optional[SecretRequest] = None
@@ -307,13 +308,11 @@ class TaskTemplate(Generic[P, R]):
     def override(
         self,
         *,
-        local: Optional[bool] = None,
-        ref: Optional[bool] = None,
         resources: Optional[Resources] = None,
         cache: CacheRequest = "auto",
         retries: Union[int, RetryStrategy] = 0,
         timeout: Optional[TimeoutType] = None,
-        reusable: Union[ReusePolicy, Literal["auto"], None] = None,
+        reusable: Union[ReusePolicy, Literal["off"], None] = None,
         env: Optional[Dict[str, str]] = None,
         secrets: Optional[SecretRequest] = None,
         **kwargs: Any,
@@ -322,15 +321,37 @@ class TaskTemplate(Generic[P, R]):
         Override various parameters of the task template. This allows for dynamic configuration of the task
         when it is called, such as changing the image, resources, cache policy, etc.
         """
-        resources = resources or self.resources
         cache = cache or self.cache
         retries = retries or self.retries
         timeout = timeout or self.timeout
         reusable = reusable or self.reusable
+        if reusable == "off":
+            reusable = None
+
+        if reusable is not None:
+            if resources is not None:
+                raise ValueError(
+                    "Cannot override resources when reusable is set."
+                    " Reusable tasks will use the parent env's resources. You can disable reusability and"
+                    " override resources if needed. (set reusable='off')"
+                )
+            if env is not None:
+                raise ValueError(
+                    "Cannot override env when reusable is set."
+                    " Reusable tasks will use the parent env's env. You can disable reusability and "
+                    "override env if needed. (set reusable='off')"
+                )
+            if secrets is not None:
+                raise ValueError(
+                    "Cannot override secrets when reusable is set."
+                    " Reusable tasks will use the parent env's secrets. You can disable reusability and "
+                    "override secrets if needed. (set reusable='off')"
+                )
+
+        resources = resources or self.resources
         env = env or self.env
         secrets = secrets or self.secrets
-        local = local or self.local
-        ref = ref or self.ref
+
         for k, v in kwargs.items():
             if k == "name":
                 raise ValueError("Name cannot be overridden")
@@ -340,13 +361,14 @@ class TaskTemplate(Generic[P, R]):
                 raise ValueError("Docs cannot be overridden")
             if k == "interface":
                 raise ValueError("Interface cannot be overridden")
+
         return replace(
             self,
             resources=resources,
             cache=cache,
             retries=retries,
             timeout=timeout,
-            reusable=reusable,
+            reusable=cast(Optional[ReusePolicy], reusable),
             env=env,
             secrets=secrets,
         )

flyte/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '0.2.0b32'
-__version_tuple__ = version_tuple = (0, 2, 0, 'b32')
+__version__ = version = '0.2.0b34'
+__version_tuple__ = version_tuple = (0, 2, 0, 'b34')

flyte/cli/_common.py

@@ -62,6 +62,28 @@ def _common_options() -> List[click.Option]:
 
 
 # This is global state for the CLI, it is manipulated by the main command
+_client_secret_options = ["client-secret", "client_secret", "clientsecret", "app-credential", "app_credential"]
+_device_flow_options = ["headless", "device-flow", "device_flow"]
+_pkce_options = ["pkce"]
+_external_command_options = ["external-command", "external_command", "externalcommand", "command", "custom"]
+ALL_AUTH_OPTIONS = _client_secret_options + _device_flow_options + _pkce_options + _external_command_options
+
+
+def sanitize_auth_type(auth_type: str | None) -> str:
+    """
+    Convert the auth type to the mode that is used by the Flyte backend.
+    """
+    if auth_type is None:
+        return "pkce"
+    if auth_type.lower() in _pkce_options:
+        return "Pkce"
+    if auth_type.lower() in _device_flow_options:
+        return "DeviceFlow"
+    if auth_type.lower() in _client_secret_options:
+        return "ClientSecret"
+    if auth_type.lower() in _external_command_options:
+        return "ExternalCommand"
+    raise ValueError(f"Unknown auth type: {auth_type}. Supported types are: {ALL_AUTH_OPTIONS}.")
 
 
 @rich.repr.auto
@@ -78,6 +100,7 @@ class CLIConfig:
     insecure: bool = False
     org: str | None = None
     simple: bool = False
+    auth_type: str | None = None
 
     def replace(self, **kwargs) -> CLIConfig:
         """
@@ -99,6 +122,8 @@ class CLIConfig:
             kwargs["endpoint"] = self.endpoint
         if self.insecure is not None:
             kwargs["insecure"] = self.insecure
+        if self.auth_type:
+            kwargs["auth_mode"] = sanitize_auth_type(self.auth_type)
         platform_cfg = self.config.platform.replace(**kwargs)
 
         updated_config = self.config.with_params(platform_cfg, task_cfg)

flyte/cli/_create.py

@@ -95,6 +95,14 @@ def secret(
     help="Image builder to use for building images. Defaults to 'local'.",
     show_default=True,
 )
+@click.option(
+    "--auth-type",
+    type=click.Choice(common.ALL_AUTH_OPTIONS, case_sensitive=False),
+    default=None,
+    help="Authentication type to use for the Flyte backend. Defaults to 'pkce'.",
+    show_default=True,
+    required=False,
+)
 def config(
     output: str,
     endpoint: str | None = None,
@@ -104,6 +112,7 @@ def config(
     domain: str | None = None,
     force: bool = False,
     image_builder: str | None = None,
+    auth_type: str | None = None,
 ):
     """
     Creates a configuration file for Flyte CLI.
@@ -128,6 +137,8 @@ def config(
         admin["endpoint"] = endpoint
     if insecure:
         admin["insecure"] = insecure
+    if auth_type:
+        admin["authType"] = common.sanitize_auth_type(auth_type)
 
     if not org and endpoint:
         org = org_from_endpoint(endpoint)

flyte/cli/_params.py

@@ -405,7 +405,7 @@ def literal_type_to_click_type(lt: LiteralType, python_type: typing.Type) -> cli
             return ct
         if lt.simple in SIMPLE_TYPE_CONVERTER:
             return SIMPLE_TYPE_CONVERTER[lt.simple]
-        raise NotImplementedError(f"Type {lt.simple} is not supported in pyflyte run")
+        raise NotImplementedError(f"Type {lt.simple} is not supported in `flyte run`")
 
     if lt.HasField("structured_dataset_type"):
         return StructuredDatasetParamType()

flyte/cli/main.py

@@ -2,6 +2,7 @@ import rich_click as click
 
 from flyte._logging import initialize_logger, logger
 
+from . import _common as common
 from ._abort import abort
 from ._build import build
 from ._common import CLIConfig
@@ -74,6 +75,14 @@ def _verbosity_to_loglevel(verbosity: int) -> int | None:
     default=None,
     show_default=True,
 )
+@click.option(
+    "--auth-type",
+    type=click.Choice(common.ALL_AUTH_OPTIONS, case_sensitive=False),
+    default=None,
+    help="Authentication type to use for the Flyte backend. Defaults to 'pkce'.",
+    show_default=True,
+    required=False,
+)
 @click.option(
     "-v",
     "--verbose",
@@ -113,6 +122,7 @@ def main(
     org: str | None,
     config_file: str | None,
     simple: bool = False,
+    auth_type: str | None = None,
 ):
     """
     The Flyte CLI is the command line interface for working with the Flyte SDK and backend.
@@ -167,6 +177,7 @@ def main(
         config=cfg,
         ctx=ctx,
         simple=simple,
+        auth_type=auth_type,
     )
 
 

flyte/errors.py

@@ -170,3 +170,12 @@ class DeploymentError(RuntimeUserError):
 
     def __init__(self, message: str):
         super().__init__("DeploymentError", message, "user")
+
+
+class ImageBuildError(RuntimeUserError):
+    """
+    This error is raised when the image build fails.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("ImageBuildError", message, "user")

flyte/io/__init__.py

@@ -7,21 +7,21 @@ of large datasets in Union.
 """
 
 __all__ = [
+    "DataFrame",
+    "DataFrameDecoder",
+    "DataFrameEncoder",
+    "DataFrameTransformerEngine",
     "Dir",
     "File",
-    "StructuredDataset",
-    "StructuredDatasetDecoder",
-    "StructuredDatasetEncoder",
-    "StructuredDatasetTransformerEngine",
-    "lazy_import_structured_dataset_handler",
+    "lazy_import_dataframe_handler",
 ]
 
+from ._dataframe import (
+    DataFrame,
+    DataFrameDecoder,
+    DataFrameEncoder,
+    DataFrameTransformerEngine,
+    lazy_import_dataframe_handler,
+)
 from ._dir import Dir
 from ._file import File
-from ._structured_dataset import (
-    StructuredDataset,
-    StructuredDatasetDecoder,
-    StructuredDatasetEncoder,
-    StructuredDatasetTransformerEngine,
-    lazy_import_structured_dataset_handler,
-)

flyte/io/{_structured_dataset → _dataframe}/__init__.py

@@ -1,15 +1,15 @@
 """
-Flytekit StructuredDataset
+Flytekit DataFrame
 ==========================================================
-.. currentmodule:: flytekit.types.structured
+.. currentmodule:: flyte.io._dataframe
 
 .. autosummary::
    :template: custom.rst
    :toctree: generated/
 
-    StructuredDataset
-    StructuredDatasetDecoder
-    StructuredDatasetEncoder
+    DataFrame
+    DataFrameDecoder
+    DataFrameEncoder
 """
 
 import functools
@@ -17,12 +17,12 @@ import functools
 from flyte._logging import logger
 from flyte._utils.lazy_module import is_imported
 
-from .structured_dataset import (
+from .dataframe import (
+    DataFrame,
+    DataFrameDecoder,
+    DataFrameEncoder,
+    DataFrameTransformerEngine,
     DuplicateHandlerError,
-    StructuredDataset,
-    StructuredDatasetDecoder,
-    StructuredDatasetEncoder,
-    StructuredDatasetTransformerEngine,
 )
 
 
@@ -30,8 +30,8 @@ from .structured_dataset import (
 def register_csv_handlers():
     from .basic_dfs import CSVToPandasDecodingHandler, PandasToCSVEncodingHandler
 
-    StructuredDatasetTransformerEngine.register(PandasToCSVEncodingHandler(), default_format_for_type=True)
-    StructuredDatasetTransformerEngine.register(CSVToPandasDecodingHandler(), default_format_for_type=True)
+    DataFrameTransformerEngine.register(PandasToCSVEncodingHandler(), default_format_for_type=True)
+    DataFrameTransformerEngine.register(CSVToPandasDecodingHandler(), default_format_for_type=True)
 
 
 @functools.lru_cache(maxsize=None)
@@ -42,9 +42,9 @@ def register_pandas_handlers():
 
     from .basic_dfs import PandasToParquetEncodingHandler, ParquetToPandasDecodingHandler
 
-    StructuredDatasetTransformerEngine.register(PandasToParquetEncodingHandler(), default_format_for_type=True)
-    StructuredDatasetTransformerEngine.register(ParquetToPandasDecodingHandler(), default_format_for_type=True)
-    StructuredDatasetTransformerEngine.register_renderer(pd.DataFrame, TopFrameRenderer())
+    DataFrameTransformerEngine.register(PandasToParquetEncodingHandler(), default_format_for_type=True)
+    DataFrameTransformerEngine.register(ParquetToPandasDecodingHandler(), default_format_for_type=True)
+    DataFrameTransformerEngine.register_renderer(pd.DataFrame, TopFrameRenderer())
 
 
 @functools.lru_cache(maxsize=None)
@@ -55,9 +55,9 @@ def register_arrow_handlers():
 
     from .basic_dfs import ArrowToParquetEncodingHandler, ParquetToArrowDecodingHandler
 
-    StructuredDatasetTransformerEngine.register(ArrowToParquetEncodingHandler(), default_format_for_type=True)
-    StructuredDatasetTransformerEngine.register(ParquetToArrowDecodingHandler(), default_format_for_type=True)
-    StructuredDatasetTransformerEngine.register_renderer(pa.Table, ArrowRenderer())
+    DataFrameTransformerEngine.register(ArrowToParquetEncodingHandler(), default_format_for_type=True)
+    DataFrameTransformerEngine.register(ParquetToArrowDecodingHandler(), default_format_for_type=True)
+    DataFrameTransformerEngine.register_renderer(pa.Table, ArrowRenderer())
 
 
 @functools.lru_cache(maxsize=None)
@@ -70,10 +70,10 @@ def register_bigquery_handlers():
             PandasToBQEncodingHandlers,
         )
 
-        StructuredDatasetTransformerEngine.register(PandasToBQEncodingHandlers())
-        StructuredDatasetTransformerEngine.register(BQToPandasDecodingHandler())
-        StructuredDatasetTransformerEngine.register(ArrowToBQEncodingHandlers())
-        StructuredDatasetTransformerEngine.register(BQToArrowDecodingHandler())
+        DataFrameTransformerEngine.register(PandasToBQEncodingHandlers())
+        DataFrameTransformerEngine.register(BQToPandasDecodingHandler())
+        DataFrameTransformerEngine.register(ArrowToBQEncodingHandlers())
+        DataFrameTransformerEngine.register(BQToArrowDecodingHandler())
     except ImportError:
         logger.info(
             "We won't register bigquery handler for structured dataset because "
@@ -86,8 +86,8 @@ def register_snowflake_handlers():
     try:
         from .snowflake import PandasToSnowflakeEncodingHandlers, SnowflakeToPandasDecodingHandler
 
-        StructuredDatasetTransformerEngine.register(SnowflakeToPandasDecodingHandler())
-        StructuredDatasetTransformerEngine.register(PandasToSnowflakeEncodingHandlers())
+        DataFrameTransformerEngine.register(SnowflakeToPandasDecodingHandler())
+        DataFrameTransformerEngine.register(PandasToSnowflakeEncodingHandlers())
 
     except ImportError:
         logger.info(
@@ -96,7 +96,7 @@ def register_snowflake_handlers():
         )
 
 
-def lazy_import_structured_dataset_handler():
+def lazy_import_dataframe_handler():
     if is_imported("pandas"):
         try:
             register_pandas_handlers()
@@ -121,9 +121,9 @@ def lazy_import_structured_dataset_handler():
 
 
 __all__ = [
-    "StructuredDataset",
-    "StructuredDatasetDecoder",
-    "StructuredDatasetEncoder",
-    "StructuredDatasetTransformerEngine",
-    "lazy_import_structured_dataset_handler",
+    "DataFrame",
+    "DataFrameDecoder",
+    "DataFrameEncoder",
+    "DataFrameTransformerEngine",
+    "lazy_import_dataframe_handler",
 ]

flyte/io/{_structured_dataset → _dataframe}/basic_dfs.py

@@ -9,12 +9,12 @@ from fsspec.core import split_protocol, strip_protocol
 import flyte.storage as storage
 from flyte._logging import logger
 from flyte._utils import lazy_module
-from flyte.io._structured_dataset.structured_dataset import (
+from flyte.io._dataframe.dataframe import (
     CSV,
     PARQUET,
-    StructuredDataset,
-    StructuredDatasetDecoder,
-    StructuredDatasetEncoder,
+    DataFrame,
+    DataFrameDecoder,
+    DataFrameEncoder,
 )
 
 if typing.TYPE_CHECKING:
@@ -39,27 +39,27 @@ def get_pandas_storage_options(uri: str, anonymous: bool = False) -> typing.Opti
     return None
 
 
-class PandasToCSVEncodingHandler(StructuredDatasetEncoder):
+class PandasToCSVEncodingHandler(DataFrameEncoder):
     def __init__(self):
         super().__init__(pd.DataFrame, None, CSV)
 
     async def encode(
         self,
-        structured_dataset: StructuredDataset,
+        dataframe: DataFrame,
         structured_dataset_type: types_pb2.StructuredDatasetType,
     ) -> literals_pb2.StructuredDataset:
-        if not structured_dataset.uri:
+        if not dataframe.uri:
             from flyte._context import internal_ctx
 
             ctx = internal_ctx()
             uri = ctx.raw_data.get_random_remote_path()
         else:
-            uri = typing.cast(str, structured_dataset.uri)
+            uri = typing.cast(str, dataframe.uri)
 
         if not storage.is_remote(uri):
             Path(uri).mkdir(parents=True, exist_ok=True)
         path = os.path.join(uri, ".csv")
-        df = typing.cast(pd.DataFrame, structured_dataset.dataframe)
+        df = typing.cast(pd.DataFrame, dataframe.val)
         df.to_csv(
             path,
             index=False,
@@ -71,7 +71,7 @@ class PandasToCSVEncodingHandler(StructuredDatasetEncoder):
         )
 
 
-class CSVToPandasDecodingHandler(StructuredDatasetDecoder):
+class CSVToPandasDecodingHandler(DataFrameDecoder):
     def __init__(self):
         super().__init__(pd.DataFrame, None, CSV)
 
@@ -97,27 +97,27 @@ class CSVToPandasDecodingHandler(StructuredDatasetDecoder):
                 raise
 
 
-class PandasToParquetEncodingHandler(StructuredDatasetEncoder):
+class PandasToParquetEncodingHandler(DataFrameEncoder):
     def __init__(self):
         super().__init__(pd.DataFrame, None, PARQUET)
 
     async def encode(
         self,
-        structured_dataset: StructuredDataset,
+        dataframe: DataFrame,
         structured_dataset_type: types_pb2.StructuredDatasetType,
     ) -> literals_pb2.StructuredDataset:
-        if not structured_dataset.uri:
+        if not dataframe.uri:
             from flyte._context import internal_ctx
 
             ctx = internal_ctx()
             uri = str(ctx.raw_data.get_random_remote_path())
         else:
-            uri = typing.cast(str, structured_dataset.uri)
+            uri = typing.cast(str, dataframe.uri)
 
         if not storage.is_remote(uri):
             Path(uri).mkdir(parents=True, exist_ok=True)
         path = os.path.join(uri, f"{0:05}")
-        df = typing.cast(pd.DataFrame, structured_dataset.dataframe)
+        df = typing.cast(pd.DataFrame, dataframe.val)
         df.to_parquet(
             path,
             coerce_timestamps="us",
@@ -130,7 +130,7 @@ class PandasToParquetEncodingHandler(StructuredDatasetEncoder):
         )
 
 
-class ParquetToPandasDecodingHandler(StructuredDatasetDecoder):
+class ParquetToPandasDecodingHandler(DataFrameDecoder):
     def __init__(self):
         super().__init__(pd.DataFrame, None, PARQUET)
 
@@ -155,36 +155,34 @@ class ParquetToPandasDecodingHandler(StructuredDatasetDecoder):
                 raise
 
 
-class ArrowToParquetEncodingHandler(StructuredDatasetEncoder):
+class ArrowToParquetEncodingHandler(DataFrameEncoder):
     def __init__(self):
         super().__init__(pa.Table, None, PARQUET)
 
     async def encode(
         self,
-        structured_dataset: StructuredDataset,
-        structured_dataset_type: types_pb2.StructuredDatasetType,
+        dataframe: DataFrame,
+        dataframe_type: types_pb2.StructuredDatasetType,
     ) -> literals_pb2.StructuredDataset:
         import pyarrow.parquet as pq
 
-        if not structured_dataset.uri:
+        if not dataframe.uri:
             from flyte._context import internal_ctx
 
             ctx = internal_ctx()
             uri = ctx.raw_data.get_random_remote_path()
         else:
-            uri = typing.cast(str, structured_dataset.uri)
+            uri = typing.cast(str, dataframe.uri)
 
         if not storage.is_remote(uri):
             Path(uri).mkdir(parents=True, exist_ok=True)
         path = os.path.join(uri, f"{0:05}")
         filesystem = storage.get_underlying_filesystem(path=path)
-        pq.write_table(structured_dataset.dataframe, strip_protocol(path), filesystem=filesystem)
-        return literals_pb2.StructuredDataset(
-            uri=uri, metadata=literals_pb2.StructuredDatasetMetadata(structured_dataset_type)
-        )
+        pq.write_table(dataframe.val, strip_protocol(path), filesystem=filesystem)
+        return literals_pb2.StructuredDataset(uri=uri, metadata=literals_pb2.StructuredDatasetMetadata(dataframe_type))
 
 
-class ParquetToArrowDecodingHandler(StructuredDatasetDecoder):
+class ParquetToArrowDecodingHandler(DataFrameDecoder):
     def __init__(self):
         super().__init__(pa.Table, None, PARQUET)
 
@@ -211,5 +209,6 @@ class ParquetToArrowDecodingHandler(StructuredDatasetDecoder):
                 fs = storage.get_underlying_filesystem(path=uri, anonymous=True)
                 if fs is not None:
                     return pq.read_table(path, filesystem=fs, columns=columns)
+                return None
             else:
                 raise