flyte 0.2.0b1__py3-none-any.whl → 0.2.0b3__py3-none-any.whl

flyte/__init__.py

@@ -17,15 +17,13 @@ async def my_task():
 """
 
 __all__ = [
-    "ABFS",
-    "GCS",
     "GPU",
-    "S3",
     "TPU",
     "Cache",
     "CachePolicy",
     "CacheRequest",
     "Device",
+    "Environment",
     "Image",
     "Resources",
     "RetryStrategy",
@@ -48,9 +46,10 @@ __all__ = [
 from ._cache import Cache, CachePolicy, CacheRequest
 from ._context import ctx
 from ._deploy import deploy
+from ._environment import Environment
 from ._group import group
 from ._image import Image
-from ._initialize import ABFS, GCS, S3, init
+from ._initialize import init
 from ._resources import GPU, TPU, Device, Resources
 from ._retry import RetryStrategy
 from ._reusable_environment import ReusePolicy

flyte/_bin/runtime.py

@@ -8,7 +8,7 @@ Refrain from importing any modules here. If you need to import any modules, do i
 import asyncio
 import os
 import sys
-from typing import List
+from typing import Any, List
 
 import click
 
@@ -27,6 +27,9 @@ ORG_NAME = "_U_ORG_NAME"
 ENDPOINT_OVERRIDE = "_U_EP_OVERRIDE"
 RUN_OUTPUT_BASE_DIR = "_U_RUN_BASE"
 
+# TODO: Remove this after proper auth is implemented
+_UNION_EAGER_API_KEY_ENV_VAR = "_UNION_EAGER_API_KEY"
+
 
 @click.command("a0")
 @click.option("--inputs", "-i", required=True)
@@ -74,11 +77,11 @@ def main(
     sys.path.insert(0, ".")
 
     import flyte._utils as utils
-    from flyte._datastructures import ActionID, Checkpoints, CodeBundle, RawDataPath
-    from flyte._initialize import S3, initialize_in_cluster
+    from flyte._initialize import initialize_in_cluster
     from flyte._internal.controllers import create_controller
     from flyte._internal.imagebuild.image_builder import ImageCache
     from flyte._internal.runtime.entrypoints import load_and_run_task
+    from flyte.models import ActionID, Checkpoints, CodeBundle, RawDataPath
 
     assert org, "Org is required for now"
     assert project, "Project is required"
@@ -91,12 +94,23 @@ def main(
     if name.startswith("{{"):
         name = os.getenv("ACTION_NAME", "")
 
-    ep = os.environ.get(ENDPOINT_OVERRIDE, "host.docker.internal:8090")
+    # Figure out how to connect
+    # This detection of api key is a hack for now.
+    controller_kwargs: dict[str, Any] = {"insecure": False}
+    if api_key := os.getenv(_UNION_EAGER_API_KEY_ENV_VAR):
+        from flyte._logging import logger
+
+        logger.warning(f"Using api key {api_key}")
+        controller_kwargs["api_key"] = api_key
+    else:
+        ep = os.environ.get(ENDPOINT_OVERRIDE, "host.docker.internal:8090")
+        controller_kwargs["endpoint"] = ep
+        if "localhost" in ep or "docker" in ep:
+            controller_kwargs["insecure"] = True
 
     bundle = CodeBundle(tgz=tgz, pkl=pkl, destination=dest, computed_version=version)
-    # TODO configure storage correctly for cluster
-    initialize_in_cluster(storage=S3.auto())
-    controller = create_controller(ct="remote", endpoint=ep, insecure=True)
+    initialize_in_cluster()
+    controller = create_controller(ct="remote", **controller_kwargs)
 
     ic = ImageCache.from_transport(image_cache) if image_cache else None
 

flyte/_cache/cache.py

@@ -14,10 +14,9 @@ from typing import (
 import rich.repr
 from typing_extensions import Literal, ParamSpec, TypeVar, get_args
 
-from flyte._datastructures import CodeBundle
-
 # if TYPE_CHECKING:
 from flyte._image import Image
+from flyte.models import CodeBundle
 
 P = ParamSpec("P")
 FuncOut = TypeVar("FuncOut")

flyte/_cli/_common.py

@@ -11,6 +11,7 @@ from types import MappingProxyType, ModuleType
 from typing import Any, Dict, Iterable, List, Optional
 
 import rich.box
+import rich.repr
 import rich_click as click
 from rich.console import Console
 from rich.panel import Panel
@@ -28,7 +29,7 @@ PROJECT_OPTION = click.Option(
     param_decls=["-p", "--project"],
     required=False,
     type=str,
-    default="default",
+    default=None,
     help="Project to operate on",
     show_default=True,
 )
@@ -37,7 +38,7 @@ DOMAIN_OPTION = click.Option(
     param_decls=["-d", "--domain"],
     required=False,
     type=str,
-    default="development",
+    default=None,
     help="Domain to operate on",
     show_default=True,
 )
@@ -63,6 +64,7 @@ def _common_options() -> List[click.Option]:
 # This is global state for the CLI, it is manipulated by the main command
 
 
+@rich.repr.auto
 @dataclass(frozen=True)
 class CLIConfig:
     """
@@ -84,13 +86,17 @@ class CLIConfig:
     def init(self, project: str | None = None, domain: str | None = None):
         import flyte
 
+        if self.config and self.config.task:
+            # If the config is already set, we can use it to initialize flyte
+            project = project or self.config.task.project
+            domain = domain or self.config.task.domain
+
         flyte.init(
             endpoint=self.endpoint,
             insecure=self.insecure,
             org=self.org_override,
             project=project,
             domain=domain,
-            log_level=self.log_level,
             config=self.config,
         )
 
@@ -113,12 +119,26 @@ class InvokeBaseMixin:
                 raise click.ClickException(f"Requested object NOT FOUND. Please check your input. Error: {e.details()}")
             if e.code() == grpc.StatusCode.ALREADY_EXISTS:
                 raise click.ClickException("Resource already exists.")
+            if e.code() == grpc.StatusCode.INTERNAL:
+                raise click.ClickException(f"Internal server error: {e.details()}")
+            if e.code() == grpc.StatusCode.UNAVAILABLE:
+                raise click.ClickException(
+                    f"Service is currently unavailable. Please try again later. Error: {e.details()}"
+                )
+            if e.code() == grpc.StatusCode.PERMISSION_DENIED:
+                raise click.ClickException(f"Permission denied. Please check your access rights. Error: {e.details()}")
+            if e.code() == grpc.StatusCode.INVALID_ARGUMENT:
+                raise click.ClickException(f"Invalid argument provided. Please check your input. Error: {e.details()}")
             raise click.ClickException(f"RPC error invoking command: {e!s}") from e
         except flyte.errors.InitializationError:
             raise click.ClickException("Initialize the CLI with a remote configuration. For example, pass --endpoint")
         except click.exceptions.Exit as e:
             # This is a normal exit, do nothing
             raise e
+        except click.exceptions.NoArgsIsHelpError:
+            # Do not raise an error if no arguments are passed, just show the help message.
+            # https://github.com/pallets/click/pull/1489
+            return None
         except Exception as e:
             if ctx.obj and ctx.obj.log_level and ctx.obj.log_level <= logging.DEBUG:
                 # If the user has requested verbose output, print the full traceback
@@ -185,8 +205,10 @@ class ObjectsPerFileGroup(GroupBase):
     Group that creates a command for each object in a python file.
     """
 
-    def __init__(self, filename: Path, *args, **kwargs):
+    def __init__(self, filename: Path | None = None, *args, **kwargs):
         super().__init__(*args, **kwargs)
+        if filename is None:
+            raise ValueError("filename must be provided")
         if not filename.exists():
             raise click.ClickException(f"{filename} does not exists")
         self.filename = filename

flyte/_cli/_create.py

@@ -1,3 +1,4 @@
+from pathlib import Path
 from typing import get_args
 
 import rich_click as click
@@ -40,3 +41,50 @@ def secret(
         with open(from_file, "rb") as f:
             value = f.read()
     Secret.create(name=name, value=value, type=type)
+
+
+@create.command(cls=common.CommandBase)
+@click.option("--endpoint", type=str, help="Endpoint of the Flyte backend.")
+@click.option("--insecure", is_flag=True, help="Use insecure connection to the Flyte backend.")
+@click.option(
+    "--org",
+    type=str,
+    required=False,
+    help="Organization to use, this will override the organization in the config file.",
+)
+@click.option(
+    "-o",
+    "--output",
+    type=click.Path(),
+    default=Path.cwd(),
+    help="Path to the output dir where the config will be saved, defaults to current directory.",
+)
+def config(
+    output: Path,
+    endpoint: str | None = None,
+    insecure: bool = False,
+    org: str | None = None,
+    project: str | None = None,
+    domain: str | None = None,
+):
+    """
+    Create a new config file.
+    """
+    import yaml
+
+    output_file = output / "config.yaml"
+    with open(output_file, "w") as f:
+        d = {
+            "admin": {
+                "endpoint": endpoint,
+                "insecure": insecure,
+            },
+            "task": {
+                "org": org,
+                "project": project,
+                "domain": domain,
+            },
+        }
+        yaml.dump(d, f)
+
+    click.echo(f"Config file created at {output_file}")

flyte/_cli/_deploy.py

@@ -91,15 +91,17 @@ class EnvPerFileGroup(common.ObjectsPerFileGroup):
     """
 
     def __init__(self, filename: Path, deploy_args: DeployArguments, *args, **kwargs):
+        args = (filename, *args)
         super().__init__(*args, **kwargs)
         self.deploy_args = deploy_args
 
     def _filter_objects(self, module: ModuleType) -> Dict[str, Any]:
-        return {k: v for k, v in module.__dict__.items() if isinstance(v, flyte.TaskEnvironment)}
+        return {k: v for k, v in module.__dict__.items() if isinstance(v, flyte.Environment)}
 
     def _get_command_for_obj(self, ctx: click.Context, obj_name: str, obj: Any) -> click.Command:
-        obj = cast(flyte.TaskEnvironment, obj)
+        obj = cast(flyte.Environment, obj)
         return DeployEnvCommand(
+            name=obj_name,
             obj_name=obj_name,
             obj=obj,
             help=obj.description,

flyte/_cli/_get.py

@@ -112,17 +112,17 @@ def action(
 
 
 @get.command(cls=common.CommandBase)
-@click.argument("run_name", type=str, required=False)
+@click.argument("run_name", type=str, required=True)
 @click.argument("action_name", type=str, required=False)
-@click.option("--lines", "-l", type=int, default=30, help="Number of lines to show")
+@click.option("--lines", "-l", type=int, default=30, help="Number of lines to show, only useful for --pretty")
 @click.option("--show-ts", is_flag=True, help="Show timestamps")
 @click.option(
-    "--follow",
-    "-f",
+    "--pretty",
     is_flag=True,
     default=False,
-    help="Follow the logs (like tail -f), will open a panel with logs scrolling.",
+    help="Show logs in a auto scrolling box, where number of lines is limited to `--lines`",
 )
+@click.option("--attempt", "-a", type=int, default=None, help="Attempt number to show logs for")
 @click.pass_obj
 def logs(
     cfg: common.CLIConfig,
@@ -132,7 +132,8 @@ def logs(
     domain: str | None = None,
     lines: int = 30,
     show_ts: bool = False,
-    follow: bool = True,
+    pretty: bool = True,
+    attempt: int | None = None,
 ):
     """
     Get the current run.
@@ -142,7 +143,7 @@ def logs(
     cfg.init(project=project, domain=domain)
 
     async def _run_log_view(_obj):
-        task = asyncio.create_task(_obj.show_logs(max_lines=lines, show_ts=show_ts, raw=not follow))
+        task = asyncio.create_task(_obj.show_logs(max_lines=lines, show_ts=show_ts, raw=not pretty, attempt=attempt))
         try:
             await task
         except KeyboardInterrupt:
@@ -233,3 +234,13 @@ def io(
             f"[green bold]Inputs[/green bold]\n{inputs}\n\n[blue bold]Outputs[/blue bold]\n{outputs}",
         )
     )
+
+
+@get.command(cls=click.RichCommand)
+@click.pass_obj
+def config(cfg: common.CLIConfig):
+    """
+    Shows the automatically detected configuration to connect with remote Flyte services.
+    """
+    console = Console()
+    console.print(cfg)

flyte/_cli/_run.py

@@ -73,6 +73,7 @@ class RunTaskCommand(click.Command):
 
     def invoke(self, ctx: Context):
         obj: CLIConfig = ctx.obj or CLIConfig()
+        assert obj.endpoint, "CLI Config should have an endpoint"
         obj.init(self.run_args.project, self.run_args.domain)
 
         r = flyte.with_runcontext(

flyte/_cli/main.py

@@ -1,5 +1,7 @@
 import rich_click as click
 
+from flyte._logging import initialize_logger, logger
+
 from ..config import Config
 from ._common import CLIConfig
 from ._create import create
@@ -41,7 +43,7 @@ def _verbosity_to_loglevel(verbosity: int) -> int | None:
     required=False,
     help="insecure",
     type=bool,
-    default=False,
+    default=None,
 )
 @click.option(
     "-v",
@@ -53,10 +55,10 @@ def _verbosity_to_loglevel(verbosity: int) -> int | None:
     type=int,
 )
 @click.option(
-    "--org-override",
+    "--org",
     type=str,
     required=False,
-    help="Override for org",
+    help="Organization to use",
 )
 @click.option(
     "-c",
@@ -73,23 +75,27 @@ def main(
     endpoint: str | None,
     insecure: bool,
     verbose: int,
-    org_override: str | None,
+    org: str | None,
     config_file: str | None,
 ):
     """
     v2 cli. Root command, please use one of the subcommands.
     """
     log_level = _verbosity_to_loglevel(verbose)
+    if log_level is not None:
+        initialize_logger(log_level)
 
     config = Config.auto(config_file=config_file)
+    logger.debug(f"Using config file discovered at location {config.source}")
 
     ctx.obj = CLIConfig(
         log_level=log_level,
         endpoint=endpoint or config.platform.endpoint,
         insecure=insecure or config.platform.insecure,
-        org_override=org_override or config.task.org,
+        org_override=org or config.task.org,
         config=config,
     )
+    logger.debug(f"Final materialized Cli config: {ctx.obj}")
 
 
 main.add_command(run)

flyte/_code_bundle/bundle.py

@@ -5,13 +5,13 @@ import os
 import pathlib
 import tempfile
 from pathlib import Path
-from typing import Type
+from typing import ClassVar, Type
 
 from flyteidl.core.tasks_pb2 import TaskTemplate
 
-import flyte.storage as storage
-from flyte._datastructures import CodeBundle
 from flyte._logging import log, logger
+from flyte._utils import AsyncLRUCache
+from flyte.models import CodeBundle
 
 from ._ignore import GitIgnore, Ignore, StandardIgnore
 from ._packaging import create_bundle, list_files_to_bundle, print_ls_tree
@@ -21,10 +21,34 @@ _pickled_file_extension = ".pkl.gz"
 _tar_file_extension = ".tar.gz"
 
 
+class _PklCache:
+    _pkl_cache: ClassVar[AsyncLRUCache[str, str]] = AsyncLRUCache[str, str](maxsize=100)
+
+    @classmethod
+    async def put(cls, digest: str, upload_to_path: str, from_path: pathlib.Path) -> str:
+        """
+        Get the pickled code bundle from the cache or build it if not present.
+
+        :param digest: The hash digest of the task template.
+        :param upload_to_path: The path to upload the pickled file to.
+        :param from_path: The path to read the pickled file from.
+        :return: CodeBundle object containing the pickled file path and the computed version.
+        """
+        import flyte.storage as storage
+
+        async def put_data() -> str:
+            return await storage.put(str(from_path), to_path=str(upload_to_path))
+
+        return await cls._pkl_cache.get(
+            key=digest,
+            value_func=put_data,
+        )
+
+
 async def build_pkl_bundle(
     o: TaskTemplate,
     upload_to_controlplane: bool = True,
-    upload_from_dataplane_path: str | None = None,
+    upload_from_dataplane_base_path: str | None = None,
     copy_bundle_to: pathlib.Path | None = None,
 ) -> CodeBundle:
     """
@@ -36,16 +60,14 @@ async def build_pkl_bundle(
 
     :param o: Object to be pickled. This is the task template.
     :param upload_to_controlplane: Whether to upload the pickled file to the control plane or not
-    :param upload_from_dataplane_path: If we are on the dataplane, this is the path where the
+    :param upload_from_dataplane_base_path: If we are on the dataplane, this is the path where the
         pickled file should be uploaded to. upload_to_controlplane has to be False in this case.
     :param copy_bundle_to: If set, the bundle will be copied to this path. This is used for testing purposes.
     :return: CodeBundle object containing the pickled file path and the computed version.
     """
     import cloudpickle
 
-    import flyte.storage as storage
-
-    if upload_to_controlplane and upload_from_dataplane_path:
+    if upload_to_controlplane and upload_from_dataplane_base_path:
         raise ValueError("Cannot upload to control plane and upload from dataplane path at the same time.")
 
     logger.debug("Building pickled code bundle.")
@@ -61,10 +83,17 @@ async def build_pkl_bundle(
             hash_digest, remote_path = await upload_file(dest)
             return CodeBundle(pkl=remote_path, computed_version=hash_digest)
 
-        elif upload_from_dataplane_path:
-            logger.debug(f"Uploading pickled code bundle to dataplane path {upload_from_dataplane_path}.")
+        elif upload_from_dataplane_base_path:
+            from flyte._internal.runtime import io
+
             _, str_digest, _ = hash_file(file_path=dest)
-            final_path = await storage.put(str(dest), upload_from_dataplane_path)
+            upload_path = io.pkl_path(upload_from_dataplane_base_path, str_digest)
+            logger.debug(f"Uploading pickled code bundle to dataplane path {upload_path}.")
+            final_path = await _PklCache.put(
+                digest=str_digest,
+                upload_to_path=upload_path,
+                from_path=dest,
+            )
             return CodeBundle(pkl=final_path, computed_version=str_digest)
 
         else:
@@ -138,6 +167,8 @@ async def download_bundle(bundle: CodeBundle) -> pathlib.Path:
 
     :return: The path to the downloaded code bundle.
     """
+    import flyte.storage as storage
+
     dest = pathlib.Path(bundle.destination)
     if not dest.is_dir():
         raise ValueError(f"Destination path should be a directory, found {dest}, {dest.stat()}")

flyte/_context.py

@@ -4,7 +4,7 @@ import contextvars
 from dataclasses import dataclass, replace
 from typing import TYPE_CHECKING, Awaitable, Callable, Optional, ParamSpec, TypeVar
 
-from flyte._datastructures import GroupData, RawDataPath, TaskContext
+from flyte.models import GroupData, RawDataPath, TaskContext
 
 if TYPE_CHECKING:
     from flyte.report import Report

flyte/_deploy.py

@@ -6,8 +6,9 @@ from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
 
 import rich.repr
 
+from flyte.models import SerializationContext
+
 from ._api_commons import syncer
-from ._datastructures import SerializationContext
 from ._environment import Environment
 from ._image import Image
 from ._initialize import get_client, get_common_config, requires_client, requires_initialization
@@ -141,6 +142,7 @@ async def apply(deployment: DeploymentPlan, copy_style: CopyFiles, dryrun: bool
     tasks = []
     for env_name, env in deployment.envs.items():
         logger.info(f"Deploying environment {env_name}")
+        # TODO Make this pluggable based on the environment type
         if isinstance(env, TaskEnvironment):
             for task in env.tasks.values():
                 tasks.append(_deploy_task(task, dryrun=dryrun, serialization_context=sc))

flyte/_group.py

@@ -1,7 +1,7 @@
 from contextlib import contextmanager
 
 from ._context import internal_ctx
-from ._datastructures import GroupData
+from .models import GroupData
 
 
 @contextmanager