From 1c9f3e23d151c98f5002808ef328c2db49535cd2 Mon Sep 17 00:00:00 2001 From: Niels Bantilan Date: Tue, 7 Dec 2021 19:26:51 -0500 Subject: [PATCH] add task_resolver arg to @task decorator (#765) Signed-off-by: Niels Bantilan Signed-off-by: Eduardo Apolinario --- flytekit/core/task.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/flytekit/core/task.py b/flytekit/core/task.py index 24fe283399c..62f982a11bd 100644 --- a/flytekit/core/task.py +++ b/flytekit/core/task.py @@ -2,7 +2,7 @@ import inspect from typing import Any, Callable, Dict, List, Optional, Type, Union -from flytekit.core.base_task import TaskMetadata +from flytekit.core.base_task import TaskMetadata, TaskResolverMixin from flytekit.core.interface import transform_signature_to_interface from flytekit.core.python_function_task import PythonFunctionTask from flytekit.core.reference_entity import ReferenceEntity, TaskReference @@ -87,6 +87,7 @@ def task( limits: Optional[Resources] = None, secret_requests: Optional[List[Secret]] = None, execution_mode: Optional[PythonFunctionTask.ExecutionBehavior] = PythonFunctionTask.ExecutionBehavior.DEFAULT, + task_resolver: Optional[TaskResolverMixin] = None, ) -> Union[Callable, PythonFunctionTask]: """ This is the core decorator to use for any task type in flytekit. @@ -170,6 +171,7 @@ def foo2(): Refer to :py:class:`Secret` to understand how to specify the request for a secret. It may change based on the backend provider. :param execution_mode: This is mainly for internal use. Please ignore. It is filled in automatically. + :param task_resolver: Provide a custom task resolver. """ def wrapper(fn) -> PythonFunctionTask: @@ -192,6 +194,7 @@ def wrapper(fn) -> PythonFunctionTask: limits=limits, secret_requests=secret_requests, execution_mode=execution_mode, + task_resolver=task_resolver, ) return task_instance