Pass additional fields to agent create

flytekit/extend/backend/agent_service.py

@@ -117,7 +117,12 @@
         agent = AgentRegistry.get_agent(template.type, template.task_type_version)
 
         logger.info(f"{agent.name} start creating the job")
-        resource_mata = await mirror_async_methods(agent.create, task_template=template, inputs=inputs)
+        resource_mata = await mirror_async_methods(
+            agent.create,
+            task_template=template,
+            inputs=inputs,
+            output_prefix=request.output_prefix,
+        )
         return CreateTaskResponse(resource_meta=resource_mata.encode())
 
     @record_agent_metrics

flytekit/extend/backend/base_agent.py

@@ -145,7 +145,9 @@ def metadata_type(self) -> ResourceMeta:
         return self._metadata_type
 
     @abstractmethod
-    def create(self, task_template: TaskTemplate, inputs: Optional[LiteralMap], **kwargs) -> ResourceMeta:
+    def create(
+        self, task_template: TaskTemplate, inputs: Optional[LiteralMap], output_prefix: Optional[str], **kwargs
+    ) -> ResourceMeta:
         """
         Return a resource meta that can be used to get the status of the task.
         """
@@ -312,6 +314,7 @@ async def _create(
             self._agent.create,
             task_template=task_template,
             inputs=literal_map,
+            output_prefix=output_prefix,
         )
 
         signal.signal(signal.SIGINT, partial(self.signal_handler, resource_meta))  # type: ignore

tests/flytekit/unit/extend/test_agent.py

@@ -47,6 +47,7 @@
 @dataclass
 class DummyMetadata(ResourceMeta):
     job_id: str
+    output_path: typing.Optional[str] = None
 
 
 class DummyAgent(AsyncAgentBase):
@@ -72,9 +73,14 @@ def __init__(self):
         super().__init__(task_type_name="async_dummy", metadata_type=DummyMetadata)
 
     async def create(
-        self, task_template: TaskTemplate, inputs: typing.Optional[LiteralMap] = None, **kwargs
+        self,
+        task_template: TaskTemplate,
+        inputs: typing.Optional[LiteralMap] = None,
+        output_prefix: typing.Optional[str] = None,
+        **kwargs,
     ) -> DummyMetadata:
-        return DummyMetadata(job_id=dummy_id)
+        output_path = f"{output_prefix}/{dummy_id}" if output_prefix else None
+        return DummyMetadata(job_id=dummy_id, output_path=output_path)
 
     async def get(self, resource_meta: DummyMetadata, **kwargs) -> Resource:
         return Resource(phase=TaskExecution.SUCCEEDED, log_links=[TaskLog(name="console", uri="localhost:3000")])
@@ -164,7 +170,7 @@ async def test_async_agent_service(agent):
 
     inputs_proto = task_inputs.to_flyte_idl()
     output_prefix = "/tmp"
-    metadata_bytes = DummyMetadata(job_id=dummy_id).encode()
+    metadata_bytes = DummyMetadata(job_id=dummy_id, output_path=f"{output_prefix}/{dummy_id}").encode()
 
     tmp = get_task_template(agent.task_category.name).to_flyte_idl()
     task_category = TaskCategory(name=agent.task_category.name, version=0)