flyte 2.0.0b5__py3-none-any.whl → 2.0.0b7__py3-none-any.whl

flyte/cli/_get.py

@@ -48,13 +48,20 @@ def project(cfg: common.CLIConfig, name: str | None = None):
     if name:
         console.print(pretty_repr(Project.get(name)))
     else:
-        console.print(common.get_table("Projects", Project.listall(), simple=cfg.simple))
+        console.print(common.format("Projects", Project.listall(), cfg.output_format))
 
 
 @get.command(cls=common.CommandBase)
 @click.argument("name", type=str, required=False)
+@click.option("--limit", type=int, default=100, help="Limit the number of runs to fetch when listing.")
 @click.pass_obj
-def run(cfg: common.CLIConfig, name: str | None = None, project: str | None = None, domain: str | None = None):
+def run(
+    cfg: common.CLIConfig,
+    name: str | None = None,
+    project: str | None = None,
+    domain: str | None = None,
+    limit: int = 100,
+):
     """
     Get a list of all runs, or details of a specific run by name.
 
@@ -71,13 +78,13 @@ def run(cfg: common.CLIConfig, name: str | None = None, project: str | None = No
         details = RunDetails.get(name=name)
         console.print(pretty_repr(details))
     else:
-        console.print(common.get_table("Runs", Run.listall(), simple=cfg.simple))
+        console.print(common.format("Runs", Run.listall(limit=limit), cfg.output_format))
 
 
 @get.command(cls=common.CommandBase)
 @click.argument("name", type=str, required=False)
 @click.argument("version", type=str, required=False)
-@click.option("--limit", type=int, default=100, help="Limit the number of tasks to show.")
+@click.option("--limit", type=int, default=100, help="Limit the number of tasks to fetch.")
 @click.pass_obj
 def task(
     cfg: common.CLIConfig,
@@ -105,9 +112,9 @@ def task(
             t = v.fetch()
             console.print(pretty_repr(t))
         else:
-            console.print(common.get_table("Tasks", Task.listall(by_task_name=name, limit=limit), simple=cfg.simple))
+            console.print(common.format("Tasks", Task.listall(by_task_name=name, limit=limit), cfg.output_format))
     else:
-        console.print(common.get_table("Tasks", Task.listall(limit=limit), simple=cfg.simple))
+        console.print(common.format("Tasks", Task.listall(limit=limit), cfg.output_format))
 
 
 @get.command(cls=common.CommandBase)
@@ -133,8 +140,8 @@ def action(
     else:
         # List all actions for the run
         console.print(
-            common.get_table(
-                f"Actions for {run_name}", flyte.remote._action.Action.listall(for_run_name=run_name), simple=cfg.simple
+            common.format(
+                f"Actions for {run_name}", flyte.remote._action.Action.listall(for_run_name=run_name), cfg.output_format
             )
         )
 
@@ -230,7 +237,7 @@ def secret(
     if name:
         console.print(pretty_repr(remote.Secret.get(name)))
     else:
-        console.print(common.get_table("Secrets", remote.Secret.listall(), simple=cfg.simple))
+        console.print(common.format("Secrets", remote.Secret.listall(), cfg.output_format))
 
 
 @get.command(cls=common.CommandBase)
@@ -299,7 +306,7 @@ def io(
         common.get_panel(
             "Inputs & Outputs",
             f"[green bold]Inputs[/green bold]\n{inputs}\n\n[blue bold]Outputs[/blue bold]\n{outputs}",
-            simple=cfg.simple,
+            cfg.output_format,
         )
     )
 

flyte/cli/_option.py

@@ -0,0 +1,33 @@
+from click import Option, UsageError
+
+
+class MutuallyExclusiveMixin:
+    def __init__(self, *args, **kwargs):
+        self.mutually_exclusive = set(kwargs.pop("mutually_exclusive", []))
+        self.error_format = kwargs.pop(
+            "error_msg", "Illegal usage: options '{name}' and '{invalid}' are mutually exclusive"
+        )
+        super().__init__(*args, **kwargs)
+
+    def handle_parse_result(self, ctx, opts, args):
+        self_present = self.name in opts and opts[self.name] is not None
+        others_intersect = self.mutually_exclusive.intersection(opts)
+        others_present = others_intersect and any(opts[value] is not None for value in others_intersect)
+
+        if others_present:
+            if self_present:
+                raise UsageError(self.error_format.format(name=self.name, invalid=", ".join(self.mutually_exclusive)))
+            else:
+                self.prompt = None
+
+        return super().handle_parse_result(ctx, opts, args)
+
+
+# See https://stackoverflow.com/a/37491504/499285 and https://stackoverflow.com/a/44349292/499285
+class MutuallyExclusiveOption(MutuallyExclusiveMixin, Option):
+    def __init__(self, *args, **kwargs):
+        mutually_exclusive = kwargs.get("mutually_exclusive", [])
+        help = kwargs.get("help", "")
+        if mutually_exclusive:
+            kwargs["help"] = help + f" Mutually exclusive with {', '.join(mutually_exclusive)}."
+        super().__init__(*args, **kwargs)

flyte/cli/_run.py

@@ -117,7 +117,7 @@ class RunTaskCommand(click.Command):
                         f"[green bold]Created Run: {r.name} [/green bold] "
                         f"(Project: {r.action.action_id.run.project}, Domain: {r.action.action_id.run.domain})\n"
                         f"➡️  [blue bold]{r.url}[/blue bold]",
-                        simple=obj.simple,
+                        obj.output_format,
                     )
                 )
                 if self.run_args.follow:

flyte/cli/main.py

@@ -1,4 +1,5 @@
 import rich_click as click
+from typing_extensions import get_args
 
 from flyte._logging import initialize_logger, logger
 
@@ -107,10 +108,13 @@ def _verbosity_to_loglevel(verbosity: int) -> int | None:
     help="Path to the configuration file to use. If not specified, the default configuration file is used.",
 )
 @click.option(
-    "--simple",
-    is_flag=True,
-    default=False,
-    help="Use a simple output format for commands that support it. This is useful for copying, pasting, and scripting.",
+    "--output-format",
+    "-of",
+    type=click.Choice(get_args(common.OutputFormat), case_sensitive=False),
+    default="table",
+    help="Output format for commands that support it. Defaults to 'table'.",
+    show_default=True,
+    required=False,
 )
 @click.rich_config(help_config=help_config)
 @click.pass_context
@@ -121,8 +125,8 @@ def main(
     verbose: int,
     org: str | None,
     config_file: str | None,
-    simple: bool = False,
     auth_type: str | None = None,
+    output_format: common.OutputFormat = "table",
 ):
     """
     The Flyte CLI is the command line interface for working with the Flyte SDK and backend.
@@ -176,8 +180,8 @@ def main(
         org=org,
         config=cfg,
         ctx=ctx,
-        simple=simple,
         auth_type=auth_type,
+        output_format=output_format,
     )
 
 

flyte/errors.py

@@ -10,6 +10,16 @@ from typing import Literal
 ErrorKind = Literal["system", "unknown", "user"]
 
 
+def silence_grpc_polling_error(loop, context):
+    """
+    Suppress specific gRPC polling errors in the event loop.
+    """
+    exc = context.get("exception")
+    if isinstance(exc, BlockingIOError):
+        return  # suppress
+    loop.default_exception_handler(context)
+
+
 class BaseRuntimeError(RuntimeError):
     """
     Base class for all Union runtime errors. These errors are raised when the underlying task execution fails, either
@@ -86,6 +96,9 @@ class TaskTimeoutError(RuntimeUserError):
     This error is raised when the underlying task execution runs for longer than the specified timeout.
     """
 
+    def __init__(self, message: str):
+        super().__init__("TaskTimeoutError", message, "user")
+
 
 class RetriesExhaustedError(RuntimeUserError):
     """
@@ -199,3 +212,12 @@ class InlineIOMaxBytesBreached(RuntimeUserError):
 
     def __init__(self, message: str):
         super().__init__("InlineIOMaxBytesBreached", message, "user")
+
+
+class RunAbortedError(RuntimeUserError):
+    """
+    This error is raised when the run is aborted by the user.
+    """
+
+    def __init__(self, message: str):
+        super().__init__("RunAbortedError", message, "user")

flyte/remote/_action.py

@@ -28,6 +28,7 @@ from flyte._initialize import ensure_client, get_client, get_common_config
 from flyte._protos.common import identifier_pb2, list_pb2
 from flyte._protos.workflow import run_definition_pb2, run_service_pb2
 from flyte._protos.workflow.run_service_pb2 import WatchActionDetailsResponse
+from flyte.remote._common import ToJSONMixin
 from flyte.remote._logs import Logs
 from flyte.syncify import syncify
 
@@ -120,7 +121,7 @@ def _action_done_check(phase: run_definition_pb2.Phase) -> bool:
 
 
 @dataclass
-class Action:
+class Action(ToJSONMixin):
     """
     A class representing an action. It is used to manage the run of a task and its state on the remote Union API.
     """
@@ -411,7 +412,7 @@ class Action:
 
 
 @dataclass
-class ActionDetails:
+class ActionDetails(ToJSONMixin):
     """
     A class representing an action. It is used to manage the run of a task and its state on the remote Union API.
     """
@@ -692,7 +693,7 @@ class ActionDetails:
 
 
 @dataclass
-class ActionInputs(UserDict):
+class ActionInputs(UserDict, ToJSONMixin):
     """
     A class representing the inputs of an action. It is used to manage the inputs of a task and its state on the
     remote Union API.
@@ -709,7 +710,7 @@ class ActionInputs(UserDict):
         return rich.pretty.pretty_repr(types.literal_string_repr(self.pb2))
 
 
-class ActionOutputs(tuple):
+class ActionOutputs(tuple, ToJSONMixin):
     """
     A class representing the outputs of an action. It is used to manage the outputs of a task and its state on the
     remote Union API.

flyte/remote/_common.py

@@ -0,0 +1,30 @@
+import json
+
+from google.protobuf.json_format import MessageToDict, MessageToJson
+
+
+class ToJSONMixin:
+    """
+    A mixin class that provides a method to convert an object to a JSON-serializable dictionary.
+    """
+
+    def to_dict(self) -> dict:
+        """
+        Convert the object to a JSON-serializable dictionary.
+
+        Returns:
+            dict: A dictionary representation of the object.
+        """
+        if hasattr(self, "pb2"):
+            return MessageToDict(self.pb2)
+        else:
+            return {k: v for k, v in self.__dict__.items() if not k.startswith("_")}
+
+    def to_json(self) -> str:
+        """
+        Convert the object to a JSON string.
+
+        Returns:
+            str: A JSON string representation of the object.
+        """
+        return MessageToJson(self.pb2) if hasattr(self, "pb2") else json.dumps(self.to_dict())

flyte/remote/_project.py

@@ -9,15 +9,17 @@ from flyteidl.admin import common_pb2, project_pb2
 from flyte._initialize import ensure_client, get_client
 from flyte.syncify import syncify
 
+from ._common import ToJSONMixin
+
 
 # TODO Add support for orgs again
 @dataclass
-class Project:
+class Project(ToJSONMixin):
     """
     A class representing a project in the Union API.
     """
 
-    _pb2: project_pb2.Project
+    pb2: project_pb2.Project
 
     @syncify
     @classmethod
@@ -76,11 +78,11 @@ class Project:
                 break
 
     def __rich_repr__(self) -> rich.repr.Result:
-        yield "name", self._pb2.name
-        yield "id", self._pb2.id
-        yield "description", self._pb2.description
-        yield "state", project_pb2.Project.ProjectState.Name(self._pb2.state)
+        yield "name", self.pb2.name
+        yield "id", self.pb2.id
+        yield "description", self.pb2.description
+        yield "state", project_pb2.Project.ProjectState.Name(self.pb2.state)
         yield (
             "labels",
-            ", ".join([f"{k}: {v}" for k, v in self._pb2.labels.values.items()]) if self._pb2.labels else None,
+            ", ".join([f"{k}: {v}" for k, v in self.pb2.labels.values.items()]) if self.pb2.labels else None,
         )

flyte/remote/_run.py

@@ -13,11 +13,12 @@ from flyte.syncify import syncify
 
 from . import Action, ActionDetails, ActionInputs, ActionOutputs
 from ._action import _action_details_rich_repr, _action_rich_repr
+from ._common import ToJSONMixin
 from ._console import get_run_url
 
 
 @dataclass
-class Run:
+class Run(ToJSONMixin):
     """
     A class representing a run of a task. It is used to manage the run of a task and its state on the remote
     Union API.
@@ -41,12 +42,14 @@ class Run:
         cls,
         filters: str | None = None,
         sort_by: Tuple[str, Literal["asc", "desc"]] | None = None,
+        limit: int = 100,
     ) -> AsyncIterator[Run]:
         """
         Get all runs for the current project and domain.
 
         :param filters: The filters to apply to the project list.
         :param sort_by: The sorting criteria for the project list, in the format (field, order).
+        :param limit: The maximum number of runs to return.
         :return: An iterator of runs.
         """
         ensure_client()
@@ -57,9 +60,10 @@ class Run:
             direction=(list_pb2.Sort.ASCENDING if sort_by[1] == "asc" else list_pb2.Sort.DESCENDING),
         )
         cfg = get_common_config()
+        i = 0
         while True:
             req = list_pb2.ListRequest(
-                limit=100,
+                limit=min(100, limit),
                 token=token,
                 sort_by=sort_pb2,
             )
@@ -76,6 +80,9 @@ class Run:
             )
             token = resp.token
             for r in resp.runs:
+                i += 1
+                if i > limit:
+                    return
                 yield cls(r)
             if not token:
                 break
@@ -213,7 +220,7 @@ class Run:
 
 
 @dataclass
-class RunDetails:
+class RunDetails(ToJSONMixin):
     """
     A class representing a run of a task. It is used to manage the run of a task and its state on the remote
     Union API.

flyte/remote/_secret.py

@@ -7,13 +7,14 @@ import rich.repr
 
 from flyte._initialize import ensure_client, get_client, get_common_config
 from flyte._protos.secret import definition_pb2, payload_pb2
+from flyte.remote._common import ToJSONMixin
 from flyte.syncify import syncify
 
 SecretTypes = Literal["regular", "image_pull"]
 
 
 @dataclass
-class Secret:
+class Secret(ToJSONMixin):
     pb2: definition_pb2.Secret
 
     @syncify

flyte/remote/_task.py

@@ -10,6 +10,7 @@ from google.protobuf import timestamp
 
 import flyte
 import flyte.errors
+from flyte._cache.cache import CacheBehavior
 from flyte._context import internal_ctx
 from flyte._initialize import ensure_client, get_client, get_common_config
 from flyte._logging import logger
@@ -18,6 +19,8 @@ from flyte._protos.workflow import task_definition_pb2, task_service_pb2
 from flyte.models import NativeInterface
 from flyte.syncify import syncify
 
+from ._common import ToJSONMixin
+
 
 def _repr_task_metadata(metadata: task_definition_pb2.TaskMetadata) -> rich.repr.Result:
     """
@@ -79,7 +82,7 @@ AutoVersioning = Literal["latest", "current"]
 
 
 @dataclass
-class TaskDetails:
+class TaskDetails(ToJSONMixin):
     pb2: task_definition_pb2.TaskDetails
     max_inline_io_bytes: int = 10 * 1024 * 1024  # 10 MB
 
@@ -201,11 +204,20 @@ class TaskDetails:
         """
         The cache policy of the task.
         """
+        metadata = self.pb2.spec.task_template.metadata
+        behavior: CacheBehavior
+        if not metadata.discoverable:
+            behavior = "disable"
+        elif metadata.discovery_version:
+            behavior = "override"
+        else:
+            behavior = "auto"
+
         return flyte.Cache(
-            behavior="enabled" if self.pb2.spec.task_template.metadata.discoverable else "disable",
-            version_override=self.pb2.spec.task_template.metadata.discovery_version,
-            serialize=self.pb2.spec.task_template.metadata.cache_serializable,
-            ignored_inputs=tuple(self.pb2.spec.task_template.metadata.cache_ignore_input_vars),
+            behavior=behavior,
+            version_override=metadata.discovery_version if metadata.discovery_version else None,
+            serialize=metadata.cache_serializable,
+            ignored_inputs=tuple(metadata.cache_ignore_input_vars),
         )
 
     @property
@@ -294,7 +306,7 @@ class TaskDetails:
 
 
 @dataclass
-class Task:
+class Task(ToJSONMixin):
     pb2: task_definition_pb2.Task
 
     def __init__(self, pb2: task_definition_pb2.Task):

flyte/syncify/_api.py

@@ -64,6 +64,10 @@ class _BackgroundLoop:
         atexit.register(self.stop)
 
     def _run(self):
+        import flyte.errors
+
+        # Set the exception handler to silence specific gRPC polling errors
+        self.loop.set_exception_handler(flyte.errors.silence_grpc_polling_error)
         asyncio.set_event_loop(self.loop)
         self.loop.run_forever()
 

{flyte-2.0.0b5.data → flyte-2.0.0b7.data}/scripts/runtime.py

@@ -26,7 +26,6 @@ DOMAIN_NAME = "FLYTE_INTERNAL_TASK_DOMAIN"
 ORG_NAME = "_U_ORG_NAME"
 ENDPOINT_OVERRIDE = "_U_EP_OVERRIDE"
 RUN_OUTPUT_BASE_DIR = "_U_RUN_BASE"
-ENABLE_REF_TASKS = "_REF_TASKS"  # This is a temporary flag to enable reference tasks in the runtime.
 
 # TODO: Remove this after proper auth is implemented
 _UNION_EAGER_API_KEY_ENV_VAR = "_UNION_EAGER_API_KEY"
@@ -87,6 +86,7 @@ def main(
 
     import flyte
     import flyte._utils as utils
+    import flyte.errors
     from flyte._initialize import init
     from flyte._internal.controllers import create_controller
     from flyte._internal.imagebuild.image_builder import ImageCache
@@ -123,22 +123,7 @@ def main(
         logger.debug(f"Using controller endpoint: {ep} with kwargs: {controller_kwargs}")
 
     bundle = CodeBundle(tgz=tgz, pkl=pkl, destination=dest, computed_version=version)
-    enable_ref_tasks = os.getenv(ENABLE_REF_TASKS, "false").lower() in ("true", "1", "yes")
-    # We init regular client here so that reference tasks can work
-    # Current reference tasks will not work with remote controller, because we create 2 different
-    # channels on different threads and this is not supported by grpcio or the auth system. It ends up leading
-    # File "src/python/grpcio/grpc/_cython/_cygrpc/aio/completion_queue.pyx.pxi", line 147,
-    # in grpc._cython.cygrpc.PollerCompletionQueue._handle_events
-    # BlockingIOError: [Errno 11] Resource temporarily unavailable
-    # TODO solution is to use a single channel for both controller and reference tasks, but this requires a refactor
-    if enable_ref_tasks:
-        logger.warning(
-            "Reference tasks are enabled. This will initialize client and you will see a BlockIOError. "
-            "This is harmless, but a nuisance. We are working on a fix."
-        )
-        init(org=org, project=project, domain=domain, **controller_kwargs)
-    else:
-        init()
+    init(org=org, project=project, domain=domain, **controller_kwargs)
     # Controller is created with the same kwargs as init, so that it can be used to run tasks
     controller = create_controller(ct="remote", **controller_kwargs)
 
@@ -164,6 +149,8 @@ def main(
 
     # Run both coroutines concurrently and wait for first to finish and cancel the other
     async def _run_and_stop():
+        loop = asyncio.get_event_loop()
+        loop.set_exception_handler(flyte.errors.silence_grpc_polling_error)
         await utils.run_coros(controller_failure, task_coroutine)
         await controller.stop()
 

{flyte-2.0.0b5.dist-info → flyte-2.0.0b7.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flyte
-Version: 2.0.0b5
+Version: 2.0.0b7
 Summary: Add your description here
 Author-email: Ketan Umare <kumare3@users.noreply.github.com>
 Requires-Python: >=3.10