fix: support image_uri being property ref for model

src/sagemaker/fw_utils.py

@@ -16,6 +16,7 @@
 import logging
 import os
 import re
+import time
 import shutil
 import tempfile
 from collections import namedtuple
@@ -24,6 +25,7 @@
 import sagemaker.image_uris
 from sagemaker.session_settings import SessionSettings
 import sagemaker.utils
+from sagemaker.workflow import is_pipeline_variable
 
 from sagemaker.deprecations import renamed_warning
 
@@ -395,8 +397,10 @@ def model_code_key_prefix(code_location_key_prefix, model_name, image):
     Returns:
         str: the key prefix to be used in uploading code
     """
-    training_job_name = sagemaker.utils.name_from_image(image)
-    return "/".join(filter(None, [code_location_key_prefix, model_name or training_job_name]))
+    name_from_image = f"/model_code/{int(time.time())}"
+    if not is_pipeline_variable(image):
+        name_from_image = sagemaker.utils.name_from_image(image)
+    return "/".join(filter(None, [code_location_key_prefix, model_name or name_from_image]))
 
 
 def warn_if_parameter_server_with_multi_gpu(training_instance_type, distribution):

tests/unit/sagemaker/workflow/test_model_step.py

@@ -46,6 +46,8 @@
     SageMakerJobStepRetryPolicy,
 )
 from sagemaker.xgboost import XGBoostModel
+from sagemaker.lambda_helper import Lambda
+from sagemaker.workflow.lambda_step import LambdaStep, LambdaOutput, LambdaOutputTypeEnum
 from tests.unit import DATA_DIR
 
 _IMAGE_URI = "fakeimage"
@@ -839,3 +841,44 @@ def _verify_register_model_container_definition(
     if submit_dir and not submit_dir.startswith("s3://"):
         # exclude the s3 path assertion as it contains timestamp
         assert submit_dir == expected_submit_dir
+
+
+def test_model_step_with_lambda_property_reference(pipeline_session):
+    lambda_step = LambdaStep(
+        name="MyLambda",
+        lambda_func=Lambda(
+            function_arn="arn:aws:lambda:us-west-2:123456789012:function:sagemaker_test_lambda"
+        ),
+        outputs=[
+            LambdaOutput(output_name="model_image", output_type=LambdaOutputTypeEnum.String),
+            LambdaOutput(output_name="model_artifact", output_type=LambdaOutputTypeEnum.String),
+        ],
+    )
+
+    model = PyTorchModel(
+        name="MyModel",
+        framework_version="1.8.0",
+        py_version="py3",
+        image_uri=lambda_step.properties.Outputs["model_image"],
+        model_data=lambda_step.properties.Outputs["model_artifact"],
+        sagemaker_session=pipeline_session,
+        entry_point=f"{DATA_DIR}/{_SCRIPT_NAME}",
+        role=_ROLE,
+    )
+
+    step_create_model = ModelStep(name="mymodelstep", step_args=model.create())
+
+    pipeline = Pipeline(
+        name="MyPipeline",
+        steps=[lambda_step, step_create_model],
+        sagemaker_session=pipeline_session,
+    )
+    steps = json.loads(pipeline.definition())["Steps"]
+    repack_step = steps[1]
+    assert repack_step["Arguments"]["InputDataConfig"][0]["DataSource"]["S3DataSource"][
+        "S3Uri"
+    ] == {"Get": "Steps.MyLambda.OutputParameters['model_artifact']"}
+    register_step = steps[2]
+    assert register_step["Arguments"]["PrimaryContainer"]["Image"] == {
+        "Get": "Steps.MyLambda.OutputParameters['model_image']"
+    }