Async Agent Interface Refactor (flyteorg#2467)

Signed-off-by: bugra.gedik <[email protected]>
bgedik · Jul 3, 2024 · 1facfd6 · 1facfd6
1 parent e414861
commit 1facfd6
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/flytekit/extend/backend/base_agent.py b/flytekit/extend/backend/base_agent.py
@@ -152,8 +152,8 @@ def metadata_type(self) -> ResourceMeta:
     def create(
         self,
         task_template: TaskTemplate,
+        output_prefix: str,
         inputs: Optional[LiteralMap],
-        output_prefix: Optional[str],
         task_execution_metadata: Optional[TaskExecutionMetadata],
         **kwargs,
     ) -> ResourceMeta:
@@ -297,7 +297,9 @@ def execute(self: PythonTask, **kwargs) -> LiteralMap:
         task_template = get_serializable(OrderedDict(), ss, self).template
         self._agent = AgentRegistry.get_agent(task_template.type, task_template.task_type_version)
 
-        resource_mata = asyncio.run(self._create(task_template, output_prefix, kwargs))
+        resource_mata = asyncio.run(
+            self._create(task_template=task_template, output_prefix=output_prefix, inputs=kwargs)
+        )
         resource = asyncio.run(self._get(resource_meta=resource_mata))
 
         if resource.phase != TaskExecution.SUCCEEDED: