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

flyte/_keyring/file.py

@@ -0,0 +1,85 @@
+from base64 import decodebytes, encodebytes
+from configparser import ConfigParser, NoOptionError, NoSectionError
+from pathlib import Path
+from typing import Optional
+
+from keyring.backend import KeyringBackend
+from keyring.errors import PasswordDeleteError
+
+_FLYTE_KEYRING_PATH: Path = Path.home() / ".flyte" / "keyring.cfg"
+
+
+class SimplePlainTextKeyring(KeyringBackend):
+    """Simple plain text keyring"""
+
+    priority = 0.5
+
+    def get_password(self, service: str, username: str) -> Optional[str]:
+        """Get password."""
+        if not self.file_path.exists():
+            return None
+
+        config = ConfigParser(interpolation=None)
+        config.read(self.file_path, encoding="utf-8")
+
+        try:
+            password_base64 = config.get(service, username).encode("utf-8")
+            return decodebytes(password_base64).decode("utf-8")
+        except (NoOptionError, NoSectionError):
+            return None
+
+    def delete_password(self, service: str, username: str) -> None:
+        """Delete password."""
+        if not self.file_path.exists():
+            raise PasswordDeleteError("Config file does not exist")
+
+        config = ConfigParser(interpolation=None)
+        config.read(self.file_path, encoding="utf-8")
+
+        try:
+            if not config.remove_option(service, username):
+                raise PasswordDeleteError("Password not found")
+        except NoSectionError:
+            raise PasswordDeleteError("Password not found")
+
+        with self.file_path.open("w", encoding="utf-8") as config_file:
+            config.write(config_file)
+
+    def set_password(self, service: str, username: str, password: str) -> None:
+        """Set password."""
+        if not username:
+            raise ValueError("Username must be provided")
+
+        file_path = self._ensure_file_path()
+        value = encodebytes(password.encode("utf-8")).decode("utf-8")
+
+        config = ConfigParser(interpolation=None)
+        config.read(file_path, encoding="utf-8")
+
+        if not config.has_section(service):
+            config.add_section(service)
+
+        config.set(service, username, value)
+
+        with file_path.open("w", encoding="utf-8") as config_file:
+            config.write(config_file)
+
+    def _ensure_file_path(self):
+        self.file_path.parent.mkdir(exist_ok=True, parents=True)
+        if not self.file_path.is_file():
+            self.file_path.touch(0o600)
+        return self.file_path
+
+    @property
+    def file_path(self) -> Path:
+        from flyte._initialize import get_common_config
+
+        config_path = get_common_config().source_config_path
+        if config_path and str(config_path.parent) == ".flyte":
+            # if the config is in a .flyte directory, use that as the path
+            return config_path.parent / "keyring.cfg"
+        # otherwise use the default path
+        return _FLYTE_KEYRING_PATH
+
+    def __repr__(self):
+        return f"<{self.__class__.__name__}> at {self.file_path}>"

flyte/_logging.py

@@ -40,6 +40,21 @@ def log_format_from_env() -> str:
     return os.environ.get("LOG_FORMAT", "json")
 
 
+def _get_console():
+    """
+    Get the console.
+    """
+    from rich.console import Console
+
+    try:
+        width = os.get_terminal_size().columns
+    except Exception as e:
+        logger.debug(f"Failed to get terminal size: {e}")
+        width = 160
+
+    return Console(width=width)
+
+
 def get_rich_handler(log_level: int) -> Optional[logging.Handler]:
     """
     Upgrades the global loggers to use Rich logging.
@@ -51,23 +66,19 @@ def get_rich_handler(log_level: int) -> Optional[logging.Handler]:
         return None
 
     import click
-    from rich.console import Console
+    from rich.highlighter import NullHighlighter
     from rich.logging import RichHandler
 
-    try:
-        width = os.get_terminal_size().columns
-    except Exception as e:
-        logger.debug(f"Failed to get terminal size: {e}")
-        width = 160
-
     handler = RichHandler(
         tracebacks_suppress=[click],
         rich_tracebacks=True,
         omit_repeated_times=False,
         show_path=False,
         log_time_format="%H:%M:%S.%f",
-        console=Console(width=width),
+        console=_get_console(),
         level=log_level,
+        highlighter=NullHighlighter(),
+        markup=True,
     )
 
     formatter = logging.Formatter(fmt="%(filename)s:%(lineno)d - %(message)s")

flyte/_task_environment.py

@@ -135,7 +135,7 @@ class TaskEnvironment(Environment):
 
     def task(
         self,
-        _func=None,
+        _func: Callable[P, R] | None = None,
         *,
         short_name: Optional[str] = None,
         cache: CacheRequest | None = None,

flyte/_utils/coro_management.py

@@ -11,7 +11,8 @@ async def run_coros(*coros: typing.Coroutine, return_when: str = asyncio.FIRST_C
     :param return_when:
     :return:
     """
-    tasks: typing.List[asyncio.Task[typing.Never]] = [asyncio.create_task(c) for c in coros]
+    # tasks: typing.List[asyncio.Task[typing.Never]] = [asyncio.create_task(c) for c in coros] # Python 3.11+
+    tasks: typing.List[asyncio.Task] = [asyncio.create_task(c) for c in coros]
     done, pending = await asyncio.wait(tasks, return_when=return_when)
     # TODO we might want to handle asyncio.CancelledError here, for cases when the `action` is cancelled
     # and we want to propagate it to all tasks. Though the backend will handle it anyway,

flyte/_version.py

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

flyte/cli/_common.py

@@ -111,7 +111,7 @@ class CLIConfig:
         """
         return replace(self, **kwargs)
 
-    def init(self, project: str | None = None, domain: str | None = None):
+    def init(self, project: str | None = None, domain: str | None = None, root_dir: str | None = None):
         from flyte.config._config import TaskConfig
 
         task_cfg = TaskConfig(
@@ -131,7 +131,7 @@ class CLIConfig:
 
         updated_config = self.config.with_params(platform_cfg, task_cfg)
 
-        flyte.init_from_config(updated_config, log_level=self.log_level)
+        flyte.init_from_config(updated_config, log_level=self.log_level, root_dir=root_dir)
 
 
 class InvokeBaseMixin:

flyte/cli/_deploy.py

@@ -43,6 +43,16 @@ class DeployArguments:
             )
         },
     )
+    root_dir: str | None = field(
+        default=None,
+        metadata={
+            "click.option": click.Option(
+                ["--root-dir"],
+                type=str,
+                help="Override the root source directory, helpful when working with monorepos.",
+            )
+        },
+    )
     recursive: bool = field(
         default=False,
         metadata={
@@ -99,7 +109,7 @@ class DeployEnvCommand(click.RichCommand):
         console = Console()
         console.print(f"Deploying root - environment: {self.env_name}")
         obj: CLIConfig = ctx.obj
-        obj.init(self.deploy_args.project, self.deploy_args.domain)
+        obj.init(self.deploy_args.project, self.deploy_args.domain, root_dir=self.deploy_args.root_dir)
         with console.status("Deploying...", spinner="dots"):
             deployment = flyte.deploy(
                 self.env,

flyte/cli/_run.py

@@ -23,14 +23,14 @@ RUN_REMOTE_CMD = "deployed-task"
 
 
 @lru_cache()
-def _initialize_config(ctx: click.Context, project: str, domain: str):
+def _initialize_config(ctx: click.Context, project: str, domain: str, root_dir: str | None = None):
     obj: CLIConfig | None = ctx.obj
     if obj is None:
         import flyte.config
 
         obj = CLIConfig(flyte.config.auto(), ctx)
 
-    obj.init(project, domain)
+    obj.init(project, domain, root_dir)
     return obj
 
 
@@ -77,6 +77,16 @@ class RunArguments:
             )
         },
     )
+    root_dir: str | None = field(
+        default=None,
+        metadata={
+            "click.option": click.Option(
+                ["--root-dir"],
+                type=str,
+                help="Override the root source directory, helpful when working with monorepos.",
+            )
+        },
+    )
     name: str | None = field(
         default=None,
         metadata={
@@ -121,7 +131,7 @@ class RunTaskCommand(click.RichCommand):
         super().__init__(obj_name, *args, **kwargs)
 
     def invoke(self, ctx: click.Context):
-        obj: CLIConfig = _initialize_config(ctx, self.run_args.project, self.run_args.domain)
+        obj: CLIConfig = _initialize_config(ctx, self.run_args.project, self.run_args.domain, self.run_args.root_dir)
 
         async def _run():
             import flyte

flyte/config/_config.py

@@ -231,10 +231,12 @@ def auto(config_file: typing.Union[str, pathlib.Path, ConfigFile, None] = None)
       1. If specified, read the config from the provided file path.
       2. If not specified, the config file is searched in the default locations.
             a. ./config.yaml if it exists  (current working directory)
-            b. `UCTL_CONFIG` environment variable
-            c. `FLYTECTL_CONFIG` environment variable
-            d. ~/.union/config.yaml if it exists
-            e. ~/.flyte/config.yaml if it exists
+            b. ./.flyte/config.yaml if it exists (current working directory)
+            c. <git_root>/.flyte/config.yaml if it exists
+            d. `UCTL_CONFIG` environment variable
+            e. `FLYTECTL_CONFIG` environment variable
+            f. ~/.union/config.yaml if it exists
+            g. ~/.flyte/config.yaml if it exists
     3. If any value is not found in the config file, the default value is used.
     4. For any value there are environment variables that match the config variable names, those will override
 

flyte/config/_reader.py

@@ -135,15 +135,25 @@ class ConfigFile(object):
         return self._yaml_config
 
 
+def _config_path_from_git_root() -> pathlib.Path | None:
+    from flyte.git import config_from_root
+
+    try:
+        return config_from_root().source
+    except RuntimeError:
+        return None
+
+
 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. ./.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
+    3. <git_root>/.flyte/config.yaml if it exists
+    4. `UCTL_CONFIG` environment variable
+    5. `FLYTECTL_CONFIG` environment variable
+    6. ~/.union/config.yaml if it exists
+    7. ~/.flyte/config.yaml if it exists
     """
     current_location_config = Path("config.yaml")
     if current_location_config.exists():
@@ -155,6 +165,11 @@ def resolve_config_path() -> pathlib.Path | None:
         return dot_flyte_config
     logger.debug("No ./.flyte/config.yaml found")
 
+    git_root_config = _config_path_from_git_root()
+    if git_root_config:
+        return git_root_config
+    logger.debug("No .flyte/config.yaml found in git repo root")
+
     uctl_path_from_env = getenv(UCTL_CONFIG_ENV_VAR, None)
     if uctl_path_from_env:
         return pathlib.Path(uctl_path_from_env)

flyte/git/_config.py

@@ -14,4 +14,6 @@ def config_from_root(path: pathlib.Path | str = ".flyte/config.yaml") -> flyte.c
     if result.returncode != 0:
         raise RuntimeError(f"Failed to get git root directory: {result.stderr}")
     root = pathlib.Path(result.stdout.strip())
+    if not (root / path).exists():
+        raise RuntimeError(f"Config file {root / path} does not exist")
     return flyte.config.auto(root / path)

flyte/io/_dataframe/dataframe.py

@@ -904,7 +904,8 @@ class DataFrameTransformerEngine(TypeTransformer[DataFrame]):
         #   t1(input_a: DataFrame)  # or
         #   t1(input_a: Annotated[DataFrame, my_cols])
         if issubclass(expected_python_type, DataFrame):
-            fdf = DataFrame(format=metad.structured_dataset_type.format)
+            fdf = DataFrame(format=metad.structured_dataset_type.format, uri=lv.scalar.structured_dataset.uri)
+            fdf._already_uploaded = True
             fdf._literal_sd = lv.scalar.structured_dataset
             fdf._metadata = metad
             return fdf
@@ -1012,7 +1013,7 @@ class DataFrameTransformerEngine(TypeTransformer[DataFrame]):
     def guess_python_type(self, literal_type: types_pb2.LiteralType) -> Type[DataFrame]:
         # todo: technically we should return the dataframe type specified in the constructor, but to do that,
         #   we'd have to store that, which we don't do today. See possibly #1363
-        if literal_type.HasField("dataframe_type"):
+        if literal_type.HasField("structured_dataset_type"):
             return DataFrame
         raise ValueError(f"DataFrameTransformerEngine cannot reverse {literal_type}")
 

flyte/io/_dir.py

@@ -27,21 +27,21 @@ class Dir(BaseModel, Generic[T], SerializableType):
     The generic type T represents the format of the files in the directory.
 
     Example:
-        ```python
-        # Async usage
-        from pandas import DataFrame
-        data_dir = Dir[DataFrame](path="s3://my-bucket/data/")
-
-        # Walk through files
-        async for file in data_dir.walk():
-            async with file.open() as f:
-                content = await f.read()
-
-        # Sync alternative
-        for file in data_dir.walk_sync():
-            with file.open_sync() as f:
-                content = f.read()
-        ```
+    ```python
+    # Async usage
+    from pandas import DataFrame
+    data_dir = Dir[DataFrame](path="s3://my-bucket/data/")
+
+    # Walk through files
+    async for file in data_dir.walk():
+        async with file.open() as f:
+            content = await f.read()
+
+    # Sync alternative
+    for file in data_dir.walk_sync():
+        with file.open_sync() as f:
+            content = f.read()
+    ```
     """
 
     # Represents either a local or remote path.
@@ -94,11 +94,11 @@ class Dir(BaseModel, Generic[T], SerializableType):
             File objects for each file found in the directory
 
         Example:
-            ```python
-            async for file in directory.walk():
-                local_path = await file.download()
-                # Process the file
-            ```
+        ```python
+        async for file in directory.walk():
+            local_path = await file.download()
+            # Process the file
+        ```
         """
         fs = storage.get_underlying_filesystem(path=self.path)
         if recursive is False:
@@ -134,11 +134,11 @@ class Dir(BaseModel, Generic[T], SerializableType):
             File objects for each file found in the directory
 
         Example:
-            ```python
-            for file in directory.walk_sync():
-                local_path = file.download_sync()
-                # Process the file
-            ```
+        ```python
+        for file in directory.walk_sync():
+            local_path = file.download_sync()
+            # Process the file
+        ```
         """
         fs = storage.get_underlying_filesystem(path=self.path)
         for parent, _, files in fs.walk(self.path, maxdepth=max_depth):
@@ -157,11 +157,11 @@ class Dir(BaseModel, Generic[T], SerializableType):
             A list of File objects
 
         Example:
-            ```python
-            files = await directory.list_files()
-            for file in files:
-                # Process the file
-            ```
+        ```python
+        files = await directory.list_files()
+        for file in files:
+            # Process the file
+        ```
         """
         # todo: this should probably also just defer to fsspec.find()
         files = []
@@ -177,11 +177,11 @@ class Dir(BaseModel, Generic[T], SerializableType):
             A list of File objects
 
         Example:
-            ```python
-            files = directory.list_files_sync()
-            for file in files:
-                # Process the file
-            ```
+        ```python
+        files = directory.list_files_sync()
+        for file in files:
+            # Process the file
+        ```
         """
         return list(self.walk_sync(recursive=False))
 
@@ -197,9 +197,9 @@ class Dir(BaseModel, Generic[T], SerializableType):
             The path to the downloaded directory
 
         Example:
-            ```python
-            local_dir = await directory.download('/tmp/my_data/')
-            ```
+        ```python
+        local_dir = await directory.download('/tmp/my_data/')
+        ```
         """
         local_dest = str(local_path) if local_path else str(storage.get_random_local_path())
         if not storage.is_remote(self.path):
@@ -230,9 +230,9 @@ class Dir(BaseModel, Generic[T], SerializableType):
             The path to the downloaded directory
 
         Example:
-            ```python
-            local_dir = directory.download_sync('/tmp/my_data/')
-            ```
+        ```python
+        local_dir = directory.download_sync('/tmp/my_data/')
+        ```
         """
         local_dest = str(local_path) if local_path else str(storage.get_random_local_path())
         if not storage.is_remote(self.path):
@@ -268,11 +268,11 @@ class Dir(BaseModel, Generic[T], SerializableType):
             A new Dir instance pointing to the uploaded directory
 
         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')
-            ```
+        ```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))
@@ -291,11 +291,11 @@ class Dir(BaseModel, Generic[T], SerializableType):
                             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")
-            ```
+        ```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)
 
@@ -312,9 +312,9 @@ class Dir(BaseModel, Generic[T], SerializableType):
             A new Dir instance pointing to the uploaded directory
 
         Example:
-            ```python
-            remote_dir = Dir[DataFrame].from_local_sync('/tmp/data_dir/', 's3://bucket/data/')
-            ```
+        ```python
+        remote_dir = Dir[DataFrame].from_local_sync('/tmp/data_dir/', 's3://bucket/data/')
+        ```
         """
         # Implement this after we figure out the final sync story
         raise NotImplementedError("Sync upload is not implemented for remote paths")
@@ -327,10 +327,10 @@ class Dir(BaseModel, Generic[T], SerializableType):
             True if the directory exists, False otherwise
 
         Example:
-            ```python
-            if await directory.exists():
-                # Process the directory
-            ```
+        ```python
+        if await directory.exists():
+            # Process the directory
+        ```
         """
         fs = storage.get_underlying_filesystem(path=self.path)
         if isinstance(fs, AsyncFileSystem):
@@ -346,10 +346,10 @@ class Dir(BaseModel, Generic[T], SerializableType):
             True if the directory exists, False otherwise
 
         Example:
-            ```python
-            if directory.exists_sync():
-                # Process the directory
-            ```
+        ```python
+        if directory.exists_sync():
+            # Process the directory
+        ```
         """
         fs = storage.get_underlying_filesystem(path=self.path)
         return fs.exists(self.path)
@@ -365,11 +365,11 @@ class Dir(BaseModel, Generic[T], SerializableType):
             A File instance if the file exists, None otherwise
 
         Example:
-            ```python
-            file = await directory.get_file("data.csv")
-            if file:
-                # Process the file
-            ```
+        ```python
+        file = await directory.get_file("data.csv")
+        if file:
+            # Process the file
+        ```
         """
         fs = storage.get_underlying_filesystem(path=self.path)
         file_path = fs.sep.join([self.path, file_name])
@@ -390,11 +390,11 @@ class Dir(BaseModel, Generic[T], SerializableType):
             A File instance if the file exists, None otherwise
 
         Example:
-            ```python
-            file = directory.get_file_sync("data.csv")
-            if file:
-                # Process the file
-            ```
+        ```python
+        file = directory.get_file_sync("data.csv")
+        if file:
+            # Process the file
+        ```
         """
         file_path = os.path.join(self.path, file_name)
         file = File[T](path=file_path)

flyte/models.py

@@ -3,13 +3,14 @@ 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:
@@ -329,7 +330,10 @@ 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)
+            if typing.get_origin(param.annotation) is Literal:
+                param_info[name] = (literal_to_enum(param.annotation), param.default)
+            else:
+                param_info[name] = (param.annotation, param.default)
 
         # Get return type
         outputs = extract_return_annotation(sig.return_annotation)

flyte/remote/_action.py

@@ -369,14 +369,15 @@ class Action(ToJSONMixin):
                     # If the action is done, handle the final state
                     if ad.done():
                         progress.stop_task(task_id)
-                        if ad.pb2.status.phase == run_definition_pb2.PHASE_SUCCEEDED:
-                            console.print(f"[bold green]Run '{self.run_name}' completed successfully.[/bold green]")
-                        else:
-                            error_message = ad.error_info.message if ad.error_info else ""
-                            console.print(
-                                f"[bold red]Run '{self.run_name}' exited unsuccessfully in state {ad.phase}"
-                                f" with error: {error_message}[/bold red]"
-                            )
+                        if not quiet:
+                            if ad.pb2.status.phase == run_definition_pb2.PHASE_SUCCEEDED:
+                                console.print(f"[bold green]Run '{self.run_name}' completed successfully.[/bold green]")
+                            else:
+                                error_message = ad.error_info.message if ad.error_info else ""
+                                console.print(
+                                    f"[bold red]Run '{self.run_name}' exited unsuccessfully in state {ad.phase}"
+                                    f" with error: {error_message}[/bold red]"
+                                )
                         break
         except asyncio.CancelledError:
             # Handle cancellation gracefully

flyte/remote/_client/auth/_authenticators/device_code.py

@@ -1,4 +1,4 @@
-import click
+from rich import print as rich_print
 
 from flyte._logging import logger
 from flyte.remote._client.auth import _token_client as token_client
@@ -94,10 +94,9 @@ class DeviceCodeAuthenticator(Authenticator):
 
         full_uri = f"{resp.verification_uri}?user_code={resp.user_code}"
         text = (
-            f"To Authenticate, navigate in a browser to the following URL: "
-            f"{click.style(full_uri, fg='blue', underline=True)}"
+            f"To Authenticate, navigate in a browser to the following URL: [blue link={full_uri}]{full_uri}[/blue link]"
         )
-        click.secho(text)
+        rich_print(text)
         try:
             token, refresh_token, expires_in = await token_client.poll_token_endpoint(
                 resp,