flyte 2.0.0b17__py3-none-any.whl → 2.0.0b19__py3-none-any.whl

flyte/_internal/runtime/convert.py

@@ -308,15 +308,82 @@ def generate_inputs_hash(serialized_inputs: str | bytes) -> str:
     return hash_data(serialized_inputs)
 
 
+def generate_inputs_repr_for_literal(literal: literals_pb2.Literal) -> bytes:
+    """
+    Generate a byte representation for a single literal that is meant to be hashed as part of the cache key
+    computation for an Action. This function should just serialize the literal deterministically, but will
+    use an existing hash value if present in the Literal.  This is trivial, except we need to handle nested literals
+    (inside collections and maps), that may have the hash property set.
+
+    :param literal: The literal to get a hashable representation for.
+    :return: byte representation of the literal that can be fed into a hash function.
+    """
+    # If the literal has a hash value, use that instead of serializing the full literal
+    if literal.hash:
+        return literal.hash.encode("utf-8")
+
+    if literal.HasField("collection"):
+        buf = bytearray()
+        for nested_literal in literal.collection.literals:
+            if nested_literal.hash:
+                buf += nested_literal.hash.encode("utf-8")
+            else:
+                buf += generate_inputs_repr_for_literal(nested_literal)
+
+        b = bytes(buf)
+        return b
+
+    elif literal.HasField("map"):
+        buf = bytearray()
+        # Sort keys to ensure deterministic ordering
+        for key in sorted(literal.map.literals.keys()):
+            nested_literal = literal.map.literals[key]
+            buf += key.encode("utf-8")
+            if nested_literal.hash:
+                buf += nested_literal.hash.encode("utf-8")
+            else:
+                buf += generate_inputs_repr_for_literal(nested_literal)
+
+        b = bytes(buf)
+        return b
+
+    # For all other cases (scalars, etc.), just serialize the literal normally
+    return literal.SerializeToString(deterministic=True)
+
+
+def generate_inputs_hash_for_named_literals(inputs: list[run_definition_pb2.NamedLiteral]) -> str:
+    """
+    Generate a hash for the inputs using the new literal representation approach that respects
+    hash values already present in literals. This is used to uniquely identify the inputs for a task
+    when some literals may have precomputed hash values.
+
+    :param inputs: List of NamedLiteral inputs to hash.
+    :return: A base64-encoded string representation of the hash.
+    """
+    if not inputs:
+        return ""
+
+    # Build the byte representation by concatenating each literal's representation
+    combined_bytes = b""
+    for named_literal in inputs:
+        # Add the name to ensure order matters
+        name_bytes = named_literal.name.encode("utf-8")
+        literal_bytes = generate_inputs_repr_for_literal(named_literal.value)
+        # Combine name and literal bytes with a separator to avoid collisions
+        combined_bytes += name_bytes + b":" + literal_bytes + b";"
+
+    return hash_data(combined_bytes)
+
+
 def generate_inputs_hash_from_proto(inputs: run_definition_pb2.Inputs) -> str:
     """
     Generate a hash for the inputs. This is used to uniquely identify the inputs for a task.
     :param inputs: The inputs to hash.
     :return: A hexadecimal string representation of the hash.
     """
-    if not inputs:
+    if not inputs or not inputs.literals:
         return ""
-    return generate_inputs_hash(inputs.SerializeToString(deterministic=True))
+    return generate_inputs_hash_for_named_literals(list(inputs.literals))
 
 
 def generate_interface_hash(task_interface: interface_pb2.TypedInterface) -> str:

flyte/_internal/runtime/taskrunner.py

@@ -4,6 +4,7 @@ invoked within a context tree.
 """
 
 import pathlib
+import time
 from typing import Any, Dict, List, Optional, Tuple
 
 import flyte.report
@@ -172,6 +173,8 @@ async def extract_download_run_upload(
     This method is invoked from the CLI (urun) and is used to run a task. This assumes that the context tree
     has already been created, and the task has been loaded. It also handles the loading of the task.
     """
+    t = time.time()
+    logger.warning(f"Task {action.name} started at {t}")
     outputs, err = await convert_and_run(
         task=task,
         input_path=input_path,
@@ -194,4 +197,4 @@ async def extract_download_run_upload(
         logger.info(f"Task {task.name} completed successfully, no outputs")
         return
     await upload_outputs(outputs, output_path) if output_path else None
-    logger.info(f"Task {task.name} completed successfully, uploaded outputs to {output_path}")
+    logger.warning(f"Task {task.name} completed successfully, uploaded outputs to {output_path} in {time.time() - t}s")

flyte/_logging.py

@@ -75,39 +75,52 @@ def get_rich_handler(log_level: int) -> Optional[logging.Handler]:
     return handler
 
 
-def get_default_handler(log_level: int) -> logging.Handler:
-    handler = logging.StreamHandler()
-    handler.setLevel(log_level)
-    formatter = logging.Formatter(fmt="[%(name)s] %(message)s")
-    if log_format_from_env() == "json":
-        pass
-        # formatter = jsonlogger.JsonFormatter(fmt="%(asctime)s %(name)s %(levelname)s %(message)s")
-    handler.setFormatter(formatter)
-    return handler
-
-
 def initialize_logger(log_level: int = get_env_log_level(), enable_rich: bool = False):
     """
     Initializes the global loggers to the default configuration.
+    When enable_rich=True, upgrades to Rich handler for local CLI usage.
     """
     global logger  # noqa: PLW0603
-    logger = _create_logger("flyte", log_level, enable_rich)
 
+    # Clear existing handlers to reconfigure
+    root = logging.getLogger()
+    root.handlers.clear()
 
-def _create_logger(name: str, log_level: int = DEFAULT_LOG_LEVEL, enable_rich: bool = False) -> logging.Logger:
-    """
-    Creates a logger with the given name and log level.
-    """
-    logger = logging.getLogger(name)
-    logger.setLevel(log_level)
-    handler = None
-    logger.handlers = []
+    flyte_logger = logging.getLogger("flyte")
+    flyte_logger.handlers.clear()
+
+    # Set up root logger handler
+    root_handler = None
+    if enable_rich:
+        root_handler = get_rich_handler(log_level)
+
+    if root_handler is None:
+        root_handler = logging.StreamHandler()
+
+    # Add context filter to root handler for all logging
+    root_handler.addFilter(ContextFilter())
+    root.addHandler(root_handler)
+
+    # Set up Flyte logger handler
+    flyte_handler = None
     if enable_rich:
-        handler = get_rich_handler(log_level)
-    if handler is None:
-        handler = get_default_handler(log_level)
-    logger.addHandler(handler)
-    return logger
+        flyte_handler = get_rich_handler(log_level)
+
+    if flyte_handler is None:
+        flyte_handler = logging.StreamHandler()
+        flyte_handler.setLevel(log_level)
+        formatter = logging.Formatter(fmt="%(message)s")
+        flyte_handler.setFormatter(formatter)
+
+    # Add both filters to Flyte handler
+    flyte_handler.addFilter(FlyteInternalFilter())
+    flyte_handler.addFilter(ContextFilter())
+
+    flyte_logger.addHandler(flyte_handler)
+    flyte_logger.setLevel(log_level)
+    flyte_logger.propagate = False  # Prevent double logging
+
+    logger = flyte_logger
 
 
 def log(fn=None, *, level=logging.DEBUG, entry=True, exit=True):
@@ -135,4 +148,75 @@ def log(fn=None, *, level=logging.DEBUG, entry=True, exit=True):
     return decorator(fn)
 
 
-logger = _create_logger("flyte", get_env_log_level())
+class ContextFilter(logging.Filter):
+    """
+    A logging filter that adds the current action's run name and name to all log records.
+    Applied globally to capture context for both user and Flyte internal logging.
+    """
+
+    def filter(self, record):
+        from flyte._context import ctx
+
+        c = ctx()
+        if c:
+            action = c.action
+            record.msg = f"[{action.run_name}][{action.name}] {record.msg}"
+        return True
+
+
+class FlyteInternalFilter(logging.Filter):
+    """
+    A logging filter that adds [flyte] prefix to internal Flyte logging only.
+    """
+
+    def filter(self, record):
+        if record.name.startswith("flyte"):
+            record.msg = f"[flyte] {record.msg}"
+        return True
+
+
+def _setup_root_logger():
+    """
+    Configure the root logger to capture all logging with context information.
+    This ensures both user code and Flyte internal logging get the context.
+    """
+    root = logging.getLogger()
+    root.handlers.clear()  # Remove any existing handlers to prevent double logging
+
+    # Create a basic handler for the root logger
+    handler = logging.StreamHandler()
+    # Add context filter to ALL logging
+    handler.addFilter(ContextFilter())
+
+    # Simple formatter since filters handle prefixes
+    root.addHandler(handler)
+
+
+def _create_flyte_logger() -> logging.Logger:
+    """
+    Create the internal Flyte logger with [flyte] prefix.
+    """
+    flyte_logger = logging.getLogger("flyte")
+    flyte_logger.setLevel(get_env_log_level())
+
+    # Add a handler specifically for flyte logging with the prefix filter
+    handler = logging.StreamHandler()
+    handler.setLevel(get_env_log_level())
+    handler.addFilter(FlyteInternalFilter())
+    handler.addFilter(ContextFilter())
+
+    formatter = logging.Formatter(fmt="%(message)s")
+    handler.setFormatter(formatter)
+
+    # Prevent propagation to root to avoid double logging
+    flyte_logger.propagate = False
+    flyte_logger.addHandler(handler)
+
+    return flyte_logger
+
+
+# Initialize root logger for global context
+_setup_root_logger()
+
+# Create the Flyte internal logger
+logger = _create_flyte_logger()

flyte/_map.py

@@ -1,4 +1,6 @@
 import asyncio
+import functools
+import logging
 from typing import Any, AsyncGenerator, AsyncIterator, Generic, Iterable, Iterator, List, Union, cast
 
 from flyte.syncify import syncify
@@ -11,7 +13,14 @@ from ._task import P, R, TaskTemplate
 class MapAsyncIterator(Generic[P, R]):
     """AsyncIterator implementation for the map function results"""
 
-    def __init__(self, func: TaskTemplate[P, R], args: tuple, name: str, concurrency: int, return_exceptions: bool):
+    def __init__(
+        self,
+        func: TaskTemplate[P, R] | functools.partial[R],
+        args: tuple,
+        name: str,
+        concurrency: int,
+        return_exceptions: bool,
+    ):
         self.func = func
         self.args = args
         self.name = name
@@ -49,13 +58,16 @@ class MapAsyncIterator(Generic[P, R]):
             return result
         except Exception as e:
             self._exception_count += 1
-            logger.debug(f"Task {self._current_index - 1} failed with exception: {e}")
+            logger.debug(
+                f"Task {self._current_index - 1} failed with exception: {e}, return_exceptions={self.return_exceptions}"
+            )
             if self.return_exceptions:
                 return e
             else:
                 # Cancel remaining tasks
                 for remaining_task in self._tasks[self._current_index + 1 :]:
                     remaining_task.cancel()
+                logger.warning("Exception raising is `ON`, raising exception and cancelling remaining tasks")
                 raise e
 
     async def _initialize(self):
@@ -64,10 +76,26 @@ class MapAsyncIterator(Generic[P, R]):
         tasks = []
         task_count = 0
 
-        for arg_tuple in zip(*self.args):
-            task = asyncio.create_task(self.func.aio(*arg_tuple))
-            tasks.append(task)
-            task_count += 1
+        if isinstance(self.func, functools.partial):
+            # Handle partial functions by merging bound args/kwargs with mapped args
+            base_func = cast(TaskTemplate, self.func.func)
+            bound_args = self.func.args
+            bound_kwargs = self.func.keywords or {}
+
+            for arg_tuple in zip(*self.args):
+                # Merge bound positional args with mapped args
+                merged_args = bound_args + arg_tuple
+                if logger.isEnabledFor(logging.DEBUG):
+                    logger.debug(f"Running {base_func.name} with args: {merged_args} and kwargs: {bound_kwargs}")
+                task = asyncio.create_task(base_func.aio(*merged_args, **bound_kwargs))
+                tasks.append(task)
+                task_count += 1
+        else:
+            # Handle regular TaskTemplate functions
+            for arg_tuple in zip(*self.args):
+                task = asyncio.create_task(self.func.aio(*arg_tuple))
+                tasks.append(task)
+                task_count += 1
 
         if task_count == 0:
             logger.info(f"Group '{self.name}' has no tasks to process")
@@ -107,9 +135,46 @@ class _Mapper(Generic[P, R]):
         """Get the name of the group, defaulting to 'map' if not provided."""
         return f"{task_name}_{group_name or 'map'}"
 
+    @staticmethod
+    def validate_partial(func: functools.partial[R]):
+        """
+        This method validates that the provided partial function is valid for mapping, i.e. only the one argument
+        is left for mapping and the rest are provided as keywords or args.
+
+        :param func: partial function to validate
+        :raises TypeError: if the partial function is not valid for mapping
+        """
+        f = cast(TaskTemplate, func.func)
+        inputs = f.native_interface.inputs
+        params = list(inputs.keys())
+        total_params = len(params)
+        provided_args = len(func.args)
+        provided_kwargs = len(func.keywords or {})
+
+        # Calculate how many parameters are left unspecified
+        unspecified_count = total_params - provided_args - provided_kwargs
+
+        # Exactly one parameter should be left for mapping
+        if unspecified_count != 1:
+            raise TypeError(
+                f"Partial function must leave exactly one parameter unspecified for mapping. "
+                f"Found {unspecified_count} unspecified parameters in {f.name}, "
+                f"params: {inputs.keys()}"
+            )
+
+        # Validate that no parameter is both in args and keywords
+        if func.keywords:
+            param_names = list(inputs.keys())
+            for i, arg_name in enumerate(param_names[: provided_args + 1]):
+                if arg_name in func.keywords:
+                    raise TypeError(
+                        f"Parameter '{arg_name}' is provided both as positional argument and keyword argument "
+                        f"in partial function {f.name}."
+                    )
+
     def __call__(
         self,
-        func: TaskTemplate[P, R],
+        func: TaskTemplate[P, R] | functools.partial[R],
         *args: Iterable[Any],
         group_name: str | None = None,
         concurrency: int = 0,
@@ -128,7 +193,13 @@ class _Mapper(Generic[P, R]):
         if not args:
             return
 
-        name = self._get_name(func.name, group_name)
+        if isinstance(func, functools.partial):
+            f = cast(TaskTemplate, func.func)
+            self.validate_partial(func)
+        else:
+            f = cast(TaskTemplate, func)
+
+        name = self._get_name(f.name, group_name)
         logger.debug(f"Blocking Map for {name}")
         with group(name):
             import flyte
@@ -154,7 +225,7 @@ class _Mapper(Generic[P, R]):
                     *args,
                     name=name,
                     concurrency=concurrency,
-                    return_exceptions=True,
+                    return_exceptions=return_exceptions,
                 ),
             ):
                 logger.debug(f"Mapped {x}, task {i}")
@@ -163,7 +234,7 @@ class _Mapper(Generic[P, R]):
 
     async def aio(
         self,
-        func: TaskTemplate[P, R],
+        func: TaskTemplate[P, R] | functools.partial[R],
         *args: Iterable[Any],
         group_name: str | None = None,
         concurrency: int = 0,
@@ -171,7 +242,14 @@ class _Mapper(Generic[P, R]):
     ) -> AsyncGenerator[Union[R, Exception], None]:
         if not args:
             return
-        name = self._get_name(func.name, group_name)
+
+        if isinstance(func, functools.partial):
+            f = cast(TaskTemplate, func.func)
+            self.validate_partial(func)
+        else:
+            f = cast(TaskTemplate, func)
+
+        name = self._get_name(f.name, group_name)
         with group(name):
             import flyte
 
@@ -199,7 +277,7 @@ class _Mapper(Generic[P, R]):
 
 @syncify
 async def _map(
-    func: TaskTemplate[P, R],
+    func: TaskTemplate[P, R] | functools.partial[R],
     *args: Iterable[Any],
     name: str = "map",
     concurrency: int = 0,

flyte/_pod.py

@@ -3,7 +3,7 @@ from typing import TYPE_CHECKING, Dict, Optional
 
 if TYPE_CHECKING:
     from flyteidl.core.tasks_pb2 import K8sPod
-    from kubernetes.client import ApiClient, V1PodSpec
+    from kubernetes.client import V1PodSpec
 
 
 _PRIMARY_CONTAINER_NAME_FIELD = "primary_container_name"
@@ -21,6 +21,7 @@ class PodTemplate(object):
 
     def to_k8s_pod(self) -> "K8sPod":
         from flyteidl.core.tasks_pb2 import K8sObjectMetadata, K8sPod
+        from kubernetes.client import ApiClient
 
         return K8sPod(
             metadata=K8sObjectMetadata(labels=self.labels, annotations=self.annotations),

flyte/_run.py

@@ -161,7 +161,12 @@ class _Runner:
                 code_bundle = cached_value.code_bundle
                 image_cache = cached_value.image_cache
             else:
-                image_cache = await build_images.aio(cast(Environment, obj.parent_env()))
+                if not self._dry_run:
+                    image_cache = await build_images.aio(cast(Environment, obj.parent_env()))
+                else:
+                    from ._internal.imagebuild.image_builder import ImageCache
+
+                    image_cache = ImageCache(image_lookup={})
 
                 if self._interactive_mode:
                     code_bundle = await build_pkl_bundle(

flyte/_task.py

@@ -258,6 +258,9 @@ class TaskTemplate(Generic[P, R]):
             else:
                 raise RuntimeSystemError("BadContext", "Controller is not initialized.")
         else:
+            from flyte._logging import logger
+
+            logger.warning(f"Task {self.name} running aio outside of a task context.")
             # Local execute, just stay out of the way, but because .aio is used, we want to return an awaitable,
             # even for synchronous tasks. This is to support migration.
             return self.forward(*args, **kwargs)

flyte/_task_environment.py

@@ -18,7 +18,7 @@ from typing import (
 
 import rich.repr
 
-from ._cache import CacheRequest
+from ._cache import Cache, CacheRequest
 from ._doc import Documentation
 from ._environment import Environment
 from ._image import Image
@@ -74,6 +74,10 @@ class TaskEnvironment(Environment):
         super().__post_init__()
         if self.reusable is not None and self.plugin_config is not None:
             raise ValueError("Cannot set plugin_config when environment is reusable.")
+        if self.reusable and not isinstance(self.reusable, ReusePolicy):
+            raise TypeError(f"Expected reusable to be of type ReusePolicy, got {type(self.reusable)}")
+        if self.cache and not isinstance(self.cache, (str, Cache)):
+            raise TypeError(f"Expected cache to be of type str or Cache, got {type(self.cache)}")
 
     def clone_with(
         self,

flyte/_trace.py

@@ -3,6 +3,7 @@ import inspect
 import time
 from typing import Any, AsyncGenerator, AsyncIterator, Awaitable, Callable, TypeGuard, TypeVar, Union, cast
 
+from flyte._logging import logger
 from flyte.models import NativeInterface
 
 T = TypeVar("T")
@@ -33,10 +34,13 @@ def trace(func: Callable[..., T]) -> Callable[..., T]:
             iface = NativeInterface.from_callable(func)
             info, ok = await controller.get_action_outputs(iface, func, *args, **kwargs)
             if ok:
+                logger.info(f"Found existing trace info for {func}, {info}")
                 if info.output:
                     return info.output
                 elif info.error:
                     raise info.error
+            else:
+                logger.debug(f"No existing trace info found for {func}, proceeding to execute.")
             start_time = time.time()
             try:
                 # Cast to Awaitable to satisfy mypy
@@ -44,6 +48,7 @@ def trace(func: Callable[..., T]) -> Callable[..., T]:
                 results = await coroutine_result
                 info.add_outputs(results, start_time=start_time, end_time=time.time())
                 await controller.record_trace(info)
+                logger.debug(f"Finished trace for {func}, {info}")
                 return results
             except Exception as e:
                 # If there is an error, we need to record it

flyte/_version.py

@@ -28,7 +28,7 @@ version_tuple: VERSION_TUPLE
 commit_id: COMMIT_ID
 __commit_id__: COMMIT_ID
 
-__version__ = version = '2.0.0b17'
-__version_tuple__ = version_tuple = (2, 0, 0, 'b17')
+__version__ = version = '2.0.0b19'
+__version_tuple__ = version_tuple = (2, 0, 0, 'b19')
 
-__commit_id__ = commit_id = 'gfe0ca1266'
+__commit_id__ = commit_id = 'g172a8c6b7'

flyte/cli/_create.py

@@ -98,7 +98,7 @@ def secret(
     "-o",
     "--output",
     type=click.Path(exists=False, writable=True),
-    default=Path.cwd() / "config.yaml",
+    default=Path.cwd() / ".flyte" / "config.yaml",
     help="Path to the output directory where the configuration will be saved. Defaults to current directory.",
     show_default=True,
 )
@@ -147,6 +147,9 @@ def config(
 
     output_path = Path(output)
 
+    if not output_path.parent.exists():
+        output_path.parent.mkdir(parents=True)
+
     if output_path.exists() and not force:
         force = click.confirm(f"Overwrite [{output_path}]?", default=False)
         if not force:

flyte/cli/_deploy.py

@@ -4,8 +4,7 @@ from pathlib import Path
 from types import ModuleType
 from typing import Any, Dict, List, cast, get_args
 
-import click
-from click import Context
+import rich_click as click
 
 import flyte
 
@@ -87,14 +86,14 @@ class DeployArguments:
         return [common.get_option_from_metadata(f.metadata) for f in fields(cls) if f.metadata]
 
 
-class DeployEnvCommand(click.Command):
+class DeployEnvCommand(click.RichCommand):
     def __init__(self, env_name: str, env: Any, deploy_args: DeployArguments, *args, **kwargs):
         self.env_name = env_name
         self.env = env
         self.deploy_args = deploy_args
         super().__init__(*args, **kwargs)
 
-    def invoke(self, ctx: Context):
+    def invoke(self, ctx: click.Context):
         from rich.console import Console
 
         console = Console()
@@ -125,7 +124,7 @@ class DeployEnvRecursiveCommand(click.Command):
         self.deploy_args = deploy_args
         super().__init__(*args, **kwargs)
 
-    def invoke(self, ctx: Context):
+    def invoke(self, ctx: click.Context):
         from rich.console import Console
 
         from flyte._environment import list_loaded_environments

flyte/cli/_params.py

@@ -283,13 +283,17 @@ class UnionParamType(click.ParamType):
     A composite type that allows for multiple types to be specified. This is used for union types.
     """
 
-    def __init__(self, types: typing.List[click.ParamType]):
+    def __init__(self, types: typing.List[click.ParamType | None]):
         super().__init__()
         self._types = self._sort_precedence(types)
-        self.name = "|".join([t.name for t in self._types])
+        self.name = "|".join([t.name for t in self._types if t is not None])
+        self.optional = False
+        if None in types:
+            self.name = f"Optional[{self.name}]"
+            self.optional = True
 
     @staticmethod
-    def _sort_precedence(tp: typing.List[click.ParamType]) -> typing.List[click.ParamType]:
+    def _sort_precedence(tp: typing.List[click.ParamType | None]) -> typing.List[click.ParamType]:
         unprocessed = []
         str_types = []
         others = []
@@ -311,6 +315,8 @@ class UnionParamType(click.ParamType):
         """
         for p in self._types:
             try:
+                if p is None and value is None:
+                    return None
                 return p.convert(value, param, ctx)
             except Exception as e:
                 logger.debug(f"Ignoring conversion error for type {p} trying other variants in Union. Error: {e}")
@@ -433,7 +439,10 @@ def literal_type_to_click_type(lt: LiteralType, python_type: typing.Type) -> cli
         for i in range(len(lt.union_type.variants)):
             variant = lt.union_type.variants[i]
             variant_python_type = typing.get_args(python_type)[i]
-            cts.append(literal_type_to_click_type(variant, variant_python_type))
+            if variant_python_type is type(None):
+                cts.append(None)
+            else:
+                cts.append(literal_type_to_click_type(variant, variant_python_type))
         return UnionParamType(cts)
 
     if lt.HasField("enum_type"):
@@ -461,6 +470,9 @@ class FlyteLiteralConverter(object):
     def is_bool(self) -> bool:
         return self.click_type == click.BOOL
 
+    def is_optional(self) -> bool:
+        return isinstance(self.click_type, UnionParamType) and self.click_type.optional
+
     def convert(
         self, ctx: click.Context, param: typing.Optional[click.Parameter], value: typing.Any
     ) -> typing.Union[Literal, typing.Any]:
@@ -525,6 +537,8 @@ def to_click_option(
     if literal_converter.is_bool():
         required = False
         is_flag = True
+    if literal_converter.is_optional():
+        required = False
 
     return click.Option(
         param_decls=[f"--{input_name}"],

flyte/cli/_run.py

@@ -112,7 +112,7 @@ class RunArguments:
         return [common.get_option_from_metadata(f.metadata) for f in fields(cls) if f.metadata]
 
 
-class RunTaskCommand(click.Command):
+class RunTaskCommand(click.RichCommand):
     def __init__(self, obj_name: str, obj: Any, run_args: RunArguments, *args, **kwargs):
         self.obj_name = obj_name
         self.obj = cast(TaskTemplate, obj)
@@ -196,7 +196,7 @@ class TaskPerFileGroup(common.ObjectsPerFileGroup):
         )
 
 
-class RunReferenceTaskCommand(click.Command):
+class RunReferenceTaskCommand(click.RichCommand):
     def __init__(self, task_name: str, run_args: RunArguments, version: str | None, *args, **kwargs):
         self.task_name = task_name
         self.run_args = run_args