modify output handler

Signed-off-by: Samhita Alla <[email protected]>
flyteorg · Jul 22, 2024 · 610f134 · 610f134
1 parent 7293bb8
commit 610f134
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 5 deletions.
diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/agent.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/agent.py
@@ -83,7 +83,7 @@ async def create(
 
     async def get(self, resource_meta: SageMakerEndpointMetadata, **kwargs) -> Resource:
         try:
-            endpoint_status, idempotence_token = await self._call(
+            endpoint_status, _ = await self._call(
                 method="describe_endpoint",
                 config={"EndpointName": resource_meta.config.get("EndpointName")},
                 inputs=resource_meta.inputs,
@@ -97,7 +97,7 @@ async def get(self, resource_meta: SageMakerEndpointMetadata, **kwargs) -> Resou
             if error_code == "ValidationException" and "Could not find endpoint" in error_message:
                 raise Exception(
                     "This might be due to resource limits being exceeded, preventing the creation of a new endpoint. Please check your resource usage and limits."
-                ) from e
+                )
             raise e
 
         current_state = endpoint_status.get("EndpointStatus")

diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/boto3_agent.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/boto3_agent.py
@@ -78,15 +78,16 @@ async def do(
                     error_message,
                 ).group(0)
                 if arn:
+                    arn_result = None
                     if method == "create_model":
-                        result = {"ModelArn": arn}
+                        arn_result = {"ModelArn": arn}
                     elif method == "create_endpoint_config":
-                        result = {"EndpointConfigArn": arn}
+                        arn_result = {"EndpointConfigArn": arn}
 
                     return Resource(
                         phase=TaskExecution.SUCCEEDED,
                         outputs={
-                            "result": result,
+                            "result": arn_result if arn_result else {"result": f"Entity already exists {arn}."},
                             "idempotence_token": e.idempotence_token,
                         },
                     )