add output_prefix back as an arg to agent create

Signed-off-by: noahjax <[email protected]>
flyteorg · Mar 15, 2024 · 4f090e6 · 4f090e6
1 parent 4208da2
commit 4f090e6
Show file tree

Hide file tree

Showing 8 changed files with 16 additions and 9 deletions.
diff --git a/flytekit/extend/backend/agent_service.py b/flytekit/extend/backend/agent_service.py
@@ -117,7 +117,12 @@ async def CreateTask(self, request: CreateTaskRequest, context: grpc.ServicerCon
         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,
+            output_prefix=request.output_prefix,
+            inputs=inputs,
+        )
         return CreateTaskResponse(resource_meta=resource_mata.encode())
 
     @record_agent_metrics

diff --git a/flytekit/extend/backend/base_agent.py b/flytekit/extend/backend/base_agent.py
@@ -145,7 +145,7 @@ 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, output_prefix: str, inputs: Optional[LiteralMap], **kwargs) -> ResourceMeta:
         """
         Return a resource meta that can be used to get the status of the task.
         """
@@ -311,6 +311,7 @@ async def _create(
         resource_meta = await mirror_async_methods(
             self._agent.create,
             task_template=task_template,
+            output_prefix=output_prefix,
             inputs=literal_map,
         )
 

diff --git a/flytekit/sensor/sensor_engine.py b/flytekit/sensor/sensor_engine.py
@@ -17,7 +17,7 @@ class SensorEngine(AsyncAgentBase):
     def __init__(self):
         super().__init__(task_type_name="sensor", metadata_type=SensorMetadata)
 
-    async def create(self, task_template: TaskTemplate, inputs: Optional[LiteralMap] = None, **kwarg) -> SensorMetadata:
+    async def create(self, task_template: TaskTemplate, output_prefix: str, inputs: Optional[LiteralMap], **kwarg) -> SensorMetadata:
         sensor_metadata = SensorMetadata(**task_template.custom)
 
         if inputs:

diff --git a/plugins/flytekit-airflow/flytekitplugins/airflow/agent.py b/plugins/flytekit-airflow/flytekitplugins/airflow/agent.py
@@ -66,7 +66,7 @@ def __init__(self):
         super().__init__(task_type_name="airflow", metadata_type=AirflowMetadata)
 
     async def create(
-        self, task_template: TaskTemplate, inputs: Optional[LiteralMap] = None, **kwargs
+        self, task_template: TaskTemplate, inputs: Optional[LiteralMap] = None, output_prefix: str = "",  **kwargs
     ) -> AirflowMetadata:
         airflow_obj = jsonpickle.decode(task_template.custom["task_config_pkl"])
         airflow_instance = _get_airflow_instance(airflow_obj)

diff --git a/plugins/flytekit-bigquery/flytekitplugins/bigquery/agent.py b/plugins/flytekit-bigquery/flytekitplugins/bigquery/agent.py
@@ -40,7 +40,8 @@ def __init__(self):
     def create(
         self,
         task_template: TaskTemplate,
-        inputs: Optional[LiteralMap] = None,
+        output_prefix: str,
+        inputs: Optional[LiteralMap] = None, 
         **kwargs,
     ) -> BigQueryMetadata:
         job_config = None

diff --git a/plugins/flytekit-mmcloud/flytekitplugins/mmcloud/agent.py b/plugins/flytekit-mmcloud/flytekitplugins/mmcloud/agent.py
@@ -56,7 +56,7 @@ async def async_login(self):
             logger.info("Logged in to OpCenter")
 
     async def create(
-        self, task_template: TaskTemplate, inputs: Optional[LiteralMap] = None, **kwargs
+        self, task_template: TaskTemplate, output_prefix: str, inputs: Optional[LiteralMap] = None, **kwargs
     ) -> MMCloudMetadata:
         """
         Submit a Flyte task as MMCloud job to the OpCenter, and return the job UID for the task.

diff --git a/plugins/flytekit-snowflake/flytekitplugins/snowflake/agent.py b/plugins/flytekit-snowflake/flytekitplugins/snowflake/agent.py
@@ -63,7 +63,7 @@ def __init__(self):
         super().__init__(task_type_name=TASK_TYPE, metadata_type=SnowflakeJobMetadata)
 
     async def create(
-        self, task_template: TaskTemplate, inputs: Optional[LiteralMap] = None, **kwargs
+        self, task_template: TaskTemplate, output_prefix: str = "", inputs: Optional[LiteralMap] = None, **kwargs
     ) -> SnowflakeJobMetadata:
         ctx = FlyteContextManager.current_context()
         literal_types = task_template.interface.inputs

diff --git a/tests/flytekit/unit/extend/test_agent.py b/tests/flytekit/unit/extend/test_agent.py
@@ -55,7 +55,7 @@ class DummyAgent(AsyncAgentBase):
     def __init__(self):
         super().__init__(task_type_name="dummy", metadata_type=DummyMetadata)
 
-    def create(self, task_template: TaskTemplate, inputs: typing.Optional[LiteralMap], **kwargs) -> DummyMetadata:
+    def create(self, task_template: TaskTemplate, output_prefix: str, inputs: typing.Optional[LiteralMap], **kwargs) -> DummyMetadata:
         return DummyMetadata(job_id=dummy_id)
 
     def get(self, resource_meta: DummyMetadata, **kwargs) -> Resource:
@@ -72,7 +72,7 @@ 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, output_prefix: str, inputs: typing.Optional[LiteralMap] = None, **kwargs
     ) -> DummyMetadata:
         return DummyMetadata(job_id=dummy_id)