flyte 2.0.0b6__py3-none-any.whl → 2.0.0b8__py3-none-any.whl

flyte/cli/_deploy.py

@@ -109,8 +109,8 @@ class DeployEnvCommand(click.Command):
                 version=self.deploy_args.version,
             )
 
-        console.print(common.get_table("Environments", deployment[0].env_repr(), simple=obj.simple))
-        console.print(common.get_table("Tasks", deployment[0].task_repr(), simple=obj.simple))
+        console.print(common.format("Environments", deployment[0].env_repr(), obj.output_format))
+        console.print(common.format("Tasks", deployment[0].task_repr(), obj.output_format))
 
 
 class DeployEnvRecursiveCommand(click.Command):
@@ -139,7 +139,7 @@ class DeployEnvRecursiveCommand(click.Command):
         if failed_paths:
             console.print(f"Loaded {len(loaded_modules)} modules with, but failed to load {len(failed_paths)} paths:")
             console.print(
-                common.get_table("Modules", [[("Path", p), ("Err", e)] for p, e in failed_paths], simple=obj.simple)
+                common.format("Modules", [[("Path", p), ("Err", e)] for p, e in failed_paths], obj.output_format)
             )
         else:
             console.print(f"Loaded {len(loaded_modules)} modules")
@@ -149,9 +149,7 @@ class DeployEnvRecursiveCommand(click.Command):
         if not all_envs:
             console.print("No environments found to deploy")
             return
-        console.print(
-            common.get_table("Loaded Environments", [[("name", e.name)] for e in all_envs], simple=obj.simple)
-        )
+        console.print(common.format("Loaded Environments", [[("name", e.name)] for e in all_envs], obj.output_format))
 
         if not self.deploy_args.ignore_load_errors and len(failed_paths) > 0:
             raise click.ClickException(
@@ -168,11 +166,9 @@ class DeployEnvRecursiveCommand(click.Command):
             )
 
         console.print(
-            common.get_table("Environments", [env for d in deployments for env in d.env_repr()], simple=obj.simple)
-        )
-        console.print(
-            common.get_table("Tasks", [task for d in deployments for task in d.task_repr()], simple=obj.simple)
+            common.format("Environments", [env for d in deployments for env in d.env_repr()], obj.output_format)
         )
+        console.print(common.format("Tasks", [task for d in deployments for task in d.task_repr()], obj.output_format))
 
 
 class EnvPerFileGroup(common.ObjectsPerFileGroup):

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
             )
         )
 
@@ -194,7 +201,7 @@ def logs(
 
     async def _run_log_view(_obj):
         task = asyncio.create_task(
-            _obj.show_logs(
+            _obj.show_logs.aio(
                 max_lines=lines, show_ts=show_ts, raw=not pretty, attempt=attempt, filter_system=filter_system
             )
         )
@@ -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

@@ -116,8 +116,8 @@ class RunTaskCommand(click.Command):
                         "Run",
                         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,
+                        f"➡️  [blue bold][link={r.url}]{r.url}[/link][/blue bold]",
+                        obj.output_format,
                     )
                 )
                 if self.run_args.follow:
@@ -125,7 +125,7 @@ class RunTaskCommand(click.Command):
                         "[dim]Log streaming enabled, will wait for task to start running "
                         "and log stream to be available[/dim]"
                     )
-                    await r.show_logs(max_lines=30, show_ts=True, raw=False)
+                    await r.show_logs.aio(max_lines=30, show_ts=True, raw=False)
 
         asyncio.run(_run())
 
@@ -212,21 +212,27 @@ run = TaskFiles(
 Run a task from a python file.
 
 Example usage:
+
 ```bash
-flyte run --name examples/basics/hello.py my_task --arg1 value1 --arg2 value2
+flyte run --project my-project --domain development hello.py my_task --arg1 value1 --arg2 value2
 ```
-Note: all arguments for the run command are provided right after the `run` command and before the file name.
 
-You can also specify the project and domain using the `--project` and `--domain` options, respectively. These
-options can be set in the config file or passed as command line arguments.
+Arguments to the run command are provided right after the `run` command and before the file name.
+For example, the command above specifies the project and domain.
+
+To run a task locally, use the `--local` flag. This will run the task in the local environment instead of the remote
+Flyte environment:
 
-Note: The arguments for the task are provided after the task name and can be retrieved using `--help`
-Example:
 ```bash
-flyte run --name examples/basics/hello.py my_task --help
+flyte run --local hello.py my_task --arg1 value1 --arg2 value2
 ```
 
-To run a task locally, use the `--local` flag. This will run the task in the local environment instead of the remote
- Flyte environment.
+Other arguments to the run command are listed below.
+
+Arguments for the task itself are provided after the task name and can be retrieved using `--help`. For example:
+
+```bash
+flyte run hello.py my_task --help
+```
 """,
 )

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.
     """
@@ -257,6 +258,7 @@ class Action:
         """
         return self.pb2.id
 
+    @syncify
     async def show_logs(
         self,
         attempt: int | None = None,
@@ -411,7 +413,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 +694,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 +711,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/_logs.py

@@ -30,7 +30,7 @@ def _format_line(logline: payload_pb2.LogLine, show_ts: bool, filter_system: boo
         if logline.originator == payload_pb2.LogLineOriginator.SYSTEM:
             return None
     style = style_map.get(logline.originator, "")
-    if "flyte" in logline.message and "flyte.errors" not in logline.message:
+    if "[flyte]" in logline.message and "flyte.errors" not in logline.message:
         if filter_system:
             return None
         style = "dim"
@@ -101,7 +101,7 @@ class Logs:
         cls,
         action_id: identifier_pb2.ActionIdentifier,
         attempt: int = 1,
-        retry: int = 3,
+        retry: int = 5,
     ) -> AsyncGenerator[payload_pb2.LogLine, None]:
         """
         Tail the logs for a given action ID and attempt.
@@ -135,7 +135,7 @@ class Logs:
                             f"Log stream not available for action {action_id.name} in run {action_id.run.name}."
                         )
                 else:
-                    await asyncio.sleep(1)
+                    await asyncio.sleep(2)
 
     @classmethod
     async def create_viewer(

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
@@ -134,6 +141,7 @@ class Run:
         """
         return self.action.watch(cache_data_on_done=cache_data_on_done)
 
+    @syncify
     async def show_logs(
         self,
         attempt: int | None = None,
@@ -142,7 +150,7 @@ class Run:
         raw: bool = False,
         filter_system: bool = False,
     ):
-        await self.action.show_logs(attempt, max_lines, show_ts, raw, filter_system=filter_system)
+        await self.action.show_logs.aio(attempt, max_lines, show_ts, raw, filter_system=filter_system)
 
     @syncify
     async def details(self) -> RunDetails:
@@ -213,7 +221,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

@@ -3,21 +3,27 @@ from __future__ import annotations
 import functools
 from dataclasses import dataclass
 from threading import Lock
-from typing import Any, AsyncIterator, Callable, Coroutine, Dict, Iterator, Literal, Optional, Tuple, Union
+from typing import Any, AsyncIterator, Callable, Coroutine, Dict, Iterator, Literal, Optional, Tuple, Union, cast
 
 import rich.repr
+from flyteidl.core import literals_pb2
 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._internal.runtime.resources_serde import get_proto_resources
+from flyte._internal.runtime.task_serde import get_proto_retry_strategy, get_proto_timeout, get_security_context
 from flyte._logging import logger
 from flyte._protos.common import identifier_pb2, list_pb2
 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:
     """
@@ -61,6 +67,15 @@ class LazyEntity:
                 raise RuntimeError(f"Error downloading the task {self._name}, (check original exception...)")
             return self._task
 
+    @syncify
+    async def override(
+        self,
+        **kwargs: Any,
+    ) -> LazyEntity:
+        task_details = cast(TaskDetails, await self.fetch.aio())
+        task_details.override(**kwargs)
+        return self
+
     async def __call__(self, *args, **kwargs):
         """
         Forwards the call to the underlying task. The entity will be fetched if not already present
@@ -79,7 +94,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 +216,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
@@ -259,19 +283,33 @@ class TaskDetails:
     def override(
         self,
         *,
-        local: Optional[bool] = None,
-        ref: Optional[bool] = None,
         resources: Optional[flyte.Resources] = None,
-        cache: flyte.CacheRequest = "auto",
         retries: Union[int, flyte.RetryStrategy] = 0,
         timeout: Optional[flyte.TimeoutType] = None,
-        reusable: Union[flyte.ReusePolicy, Literal["auto"], None] = None,
         env: Optional[Dict[str, str]] = None,
         secrets: Optional[flyte.SecretRequest] = None,
-        max_inline_io_bytes: int | None = None,
         **kwargs: Any,
     ) -> TaskDetails:
-        raise NotImplementedError
+        if len(kwargs) > 0:
+            raise ValueError(
+                f"ReferenceTasks [{self.name}] do not support overriding with kwargs: {kwargs}, "
+                f"Check the parameters for override method."
+            )
+        template = self.pb2.spec.task_template
+        if secrets:
+            template.security_context.CopyFrom(get_security_context(secrets))
+        if template.HasField("container"):
+            if env:
+                template.container.env.clear()
+                template.container.env.extend([literals_pb2.KeyValuePair(key=k, value=v) for k, v in env.items()])
+            if resources:
+                template.container.resources.CopyFrom(get_proto_resources(resources))
+        if retries:
+            template.metadata.retries.CopyFrom(get_proto_retry_strategy(retries))
+        if timeout:
+            template.metadata.timeout.CopyFrom(get_proto_timeout(timeout))
+
+        return self
 
     def __rich_repr__(self) -> rich.repr.Result:
         """
@@ -294,7 +332,7 @@ class TaskDetails:
 
 
 @dataclass
-class Task:
+class Task(ToJSONMixin):
     pb2: task_definition_pb2.Task
 
     def __init__(self, pb2: task_definition_pb2.Task):