don't return triton image as a functools.partial (#2326)

* move triton image code from mixin to task Signed-off-by: Samhita Alla <[email protected]> * don't return triton image partial Signed-off-by: Samhita Alla <[email protected]> --------- Signed-off-by: Samhita Alla <[email protected]> Signed-off-by: Jan Fiedler <[email protected]>
flyteorg · Jul 25, 2024 · 688f0af · 688f0af
1 parent 8caf617
commit 688f0af
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 9 deletions.
diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/__init__.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/__init__.py
@@ -19,8 +19,6 @@
    delete_sagemaker_deployment
 """
 
-from functools import partial
-
 from .agent import SageMakerEndpointAgent
 from .boto3_agent import BotoAgent
 from .boto3_task import BotoConfig, BotoTask
@@ -37,7 +35,5 @@
 
 
 def triton_image_uri(version: str = "23.12"):
-    return partial(
-        "{account_id}.dkr.ecr.{region}.{base}/sagemaker-tritonserver:{version}-py3".format,
-        version=version,
-    )
+    image = "{account_id}.dkr.ecr.{region}.{base}/sagemaker-tritonserver:{version}-py3"
+    return image.replace("{version}", version)
diff --git a/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/boto3_mixin.py b/plugins/flytekit-aws-sagemaker/flytekitplugins/awssagemaker_inference/boto3_mixin.py
@@ -1,4 +1,3 @@
-from functools import partial
 from typing import Any, Dict, Optional
 
 import aioboto3
@@ -152,12 +151,12 @@ async def _call(
             base = "amazonaws.com.cn" if final_region.startswith("cn-") else "amazonaws.com"
             images = {
                 image_name: (
-                    image(
+                    image.format(
                         account_id=account_id_map[final_region],
                         region=final_region,
                         base=base,
                     )
-                    if isinstance(image, partial)
+                    if isinstance(image, str) and "sagemaker-tritonserver" in image
                     else image
                 )
                 for image_name, image in images.items()