Move SageMakerEndpointMetadata to agent.py

flytekit/core/base_task.py

@@ -457,7 +457,7 @@ def __init__(
         self,
         task_type: str,
         name: str,
-        task_config: Optional[T],
+        task_config: Optional[T] = None,
         interface: Optional[Interface] = None,
         environment: Optional[Dict[str, str]] = None,
         disable_deck: Optional[bool] = None,

plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/agent.py

@@ -1,18 +1,36 @@
 import json
+from dataclasses import dataclass
 from datetime import datetime
-from typing import Optional
+from typing import Any, Dict, Optional
+
+import cloudpickle
 
 from flytekit.extend.backend.base_agent import (
     AgentRegistry,
     AsyncAgentBase,
     Resource,
+    ResourceMeta,
 )
 from flytekit.extend.backend.utils import convert_to_flyte_phase
 from flytekit.models.literals import LiteralMap
 from flytekit.models.task import TaskTemplate
 
 from .boto3_mixin import Boto3AgentMixin
-from .task import SageMakerEndpointMetadata
+
+
+@dataclass
+class SageMakerEndpointMetadata(ResourceMeta):
+    config: Dict[str, Any]
+    region: Optional[str] = None
+    inputs: Optional[LiteralMap] = None
+
+    def encode(self) -> bytes:
+        return cloudpickle.dumps(self)
+
+    @classmethod
+    def decode(cls, data: bytes) -> "SageMakerEndpointMetadata":
+        return cloudpickle.loads(data)
+
 
 states = {
     "Creating": "Running",

plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/task.py

@@ -1,12 +1,10 @@
-from dataclasses import dataclass
 from typing import Any, Dict, Optional, Type, Union
 
 from flytekit import ImageSpec, kwtypes
 from flytekit.configuration import SerializationSettings
 from flytekit.core.base_task import PythonTask
 from flytekit.core.interface import Interface
 from flytekit.extend.backend.base_agent import AsyncAgentExecutorMixin
-from flytekit.models.literals import LiteralMap
 
 from .boto3_task import BotoConfig, BotoTask
 
@@ -75,14 +73,7 @@ def __init__(
         )
 
 
-@dataclass
-class SageMakerEndpointMetadata(object):
-    config: Dict[str, Any]
-    region: Optional[str] = None
-    inputs: Optional[LiteralMap] = None
-
-
-class SageMakerEndpointTask(AsyncAgentExecutorMixin, PythonTask[SageMakerEndpointMetadata]):
+class SageMakerEndpointTask(AsyncAgentExecutorMixin, PythonTask):
     _TASK_TYPE = "sagemaker-endpoint"
 
     def __init__(
@@ -103,17 +94,15 @@ def __init__(
         """
         super().__init__(
             name=name,
-            task_config=SageMakerEndpointMetadata(
-                config=config,
-                region=region,
-            ),
             task_type=self._TASK_TYPE,
             interface=Interface(inputs=inputs, outputs=kwtypes(result=str)),
             **kwargs,
         )
+        self._config = config
+        self._region = region
 
     def get_custom(self, settings: SerializationSettings) -> Dict[str, Any]:
-        return {"config": self.task_config.config, "region": self.task_config.region}
+        return {"config": self._config, "region": self._region}
 
 
 class SageMakerDeleteEndpointTask(BotoTask):