add task_resolver arg to @task decorator (#765)

Signed-off-by: Niels Bantilan <[email protected]>
Signed-off-by: Eduardo Apolinario <[email protected]>

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