change: Make repack step output path align with model repack path

Author: Dewen Qi

src/sagemaker/model.py

@@ -527,10 +527,10 @@ def _upload_code(self, key_prefix: str, repack: bool = False) -> None:
                 artifact should be repackaged into a new S3 object. (default: False).
         """
         local_code = utils.get_config_value("local.local_code", self.sagemaker_session.config)
+        bucket = self.bucket or self.sagemaker_session.default_bucket()
         if (self.sagemaker_session.local_mode and local_code) or self.entry_point is None:
             self.uploaded_code = None
         elif not repack:
-            bucket = self.bucket or self.sagemaker_session.default_bucket()
             self.uploaded_code = fw_utils.tar_and_upload_dir(
                 session=self.sagemaker_session.boto_session,
                 bucket=bucket,
@@ -557,6 +557,9 @@ def _upload_code(self, key_prefix: str, repack: bool = False) -> None:
                     )
                     return
                 self.sagemaker_session.context.need_runtime_repack.add(id(self))
+                self.sagemaker_session.context.runtime_repack_output_prefix = "s3://{}/{}".format(
+                    bucket, key_prefix
+                )
                 # Add the uploaded_code and repacked_model_data to update the container env
                 self.repacked_model_data = self.model_data
                 self.uploaded_code = fw_utils.UploadedCode(
@@ -567,7 +570,6 @@ def _upload_code(self, key_prefix: str, repack: bool = False) -> None:
             if local_code and self.model_data.startswith("file://"):
                 repacked_model_data = self.model_data
             else:
-                bucket = self.bucket or self.sagemaker_session.default_bucket()
                 repacked_model_data = "s3://" + "/".join([bucket, key_prefix, "model.tar.gz"])
                 self.uploaded_code = fw_utils.UploadedCode(
                     s3_prefix=repacked_model_data, script_name=os.path.basename(self.entry_point)

src/sagemaker/tensorflow/model.py

@@ -363,13 +363,10 @@ def prepare_container_def(
             instance_type, accelerator_type, serverless_inference_config=serverless_inference_config
         )
         env = self._get_container_env()
+        key_prefix = sagemaker.fw_utils.model_code_key_prefix(self.key_prefix, self.name, image_uri)
+        bucket = self.bucket or self.sagemaker_session.default_bucket()
 
         if self.entry_point and not is_pipeline_variable(self.model_data):
-            key_prefix = sagemaker.fw_utils.model_code_key_prefix(
-                self.key_prefix, self.name, image_uri
-            )
-
-            bucket = self.bucket or self.sagemaker_session.default_bucket()
             model_data = s3.s3_path_join("s3://", bucket, key_prefix, "model.tar.gz")
 
             sagemaker.utils.repack_model(
@@ -385,6 +382,9 @@ def prepare_container_def(
             # model is not yet there, defer repacking to later during pipeline execution
             if isinstance(self.sagemaker_session, PipelineSession):
                 self.sagemaker_session.context.need_runtime_repack.add(id(self))
+                self.sagemaker_session.context.runtime_repack_output_prefix = "s3://{}/{}".format(
+                    bucket, key_prefix
+                )
             else:
                 logging.warning(
                     "The model_data is a Pipeline variable of type %s, "

src/sagemaker/workflow/model_step.py

@@ -23,8 +23,6 @@
 from sagemaker.workflow.step_collections import StepCollection
 from sagemaker.workflow.steps import Step, CreateModelStep
 
-NEED_RUNTIME_REPACK = "need_runtime_repack"
-
 _CREATE_MODEL_RETRY_POLICIES = "create_model_retry_policies"
 _REGISTER_MODEL_RETRY_POLICIES = "register_model_retry_policies"
 _REPACK_MODEL_RETRY_POLICIES = "repack_model_retry_policies"
@@ -155,6 +153,7 @@ def __init__(
         self._create_model_args = self.step_args.create_model_request
         self._register_model_args = self.step_args.create_model_package_request
         self._need_runtime_repack = self.step_args.need_runtime_repack
+        self._runtime_repack_output_prefix = self.step_args.runtime_repack_output_prefix
         self._assign_and_validate_retry_policies(retry_policies)
 
         if self._need_runtime_repack:
@@ -268,6 +267,7 @@ def _append_repack_model_step(self):
                     ),
                     depends_on=self.depends_on,
                     retry_policies=self._repack_model_retry_policies,
+                    output_path=self._runtime_repack_output_prefix,
                 )
                 self.steps.append(repack_model_step)
 

src/sagemaker/workflow/pipeline_context.py

@@ -67,6 +67,7 @@ def __init__(self, model):
         self.create_model_package_request = None
         self.create_model_request = None
         self.need_runtime_repack = set()
+        self.runtime_repack_output_prefix = None
 
 
 class PipelineSession(Session):
@@ -139,14 +140,14 @@ def _intercept_create_request(self, request: Dict, create, func_name: str = None
         else:
             self.context = _JobStepArguments(func_name, request)
 
-    def init_step_arguments(self, model):
+    def init_model_step_arguments(self, model):
         """Create a `_ModelStepArguments` (if not exist) as pipeline context
 
         Args:
             model (Model or PipelineModel): A `sagemaker.model.Model`
                 or `sagemaker.pipeline.PipelineModel` instance
         """
-        if not self._context or not isinstance(self._context, _ModelStepArguments):
+        if not isinstance(self._context, _ModelStepArguments):
             self._context = _ModelStepArguments(model)
 
 
@@ -197,7 +198,7 @@ def wrapper(*args, **kwargs):
                 UserWarning,
             )
             if run_func.__name__ in ["register", "create"]:
-                self_instance.sagemaker_session.init_step_arguments(self_instance)
+                self_instance.sagemaker_session.init_model_step_arguments(self_instance)
                 run_func(*args, **kwargs)
                 context = self_instance.sagemaker_session.context
                 self_instance.sagemaker_session.context = None

tests/integ/sagemaker/workflow/test_model_steps.py

@@ -836,6 +836,7 @@ def test_tensorflow_model_register_and_deploy_with_runtime_repack(
         sagemaker_session=pipeline_session,
         entry_point=os.path.join(_TENSORFLOW_PATH, "inference.py"),
         dependencies=[os.path.join(_TENSORFLOW_PATH, "dependency.py")],
+        code_location=f"s3://{pipeline_session.default_bucket()}/model-code",
     )
     step_args = tf_model.register(
         content_types=["application/json"],

tests/unit/sagemaker/workflow/test_model_step.py

@@ -67,6 +67,8 @@
 _DIR_NAME = "/opt/ml/model/code"
 _XGBOOST_PATH = os.path.join(DATA_DIR, "xgboost_abalone")
 _TENSORFLOW_PATH = os.path.join(DATA_DIR, "tfs/tfs-test-entrypoint-and-dependencies")
+_REPACK_OUTPUT_KEY_PREFIX = "code-output"
+_MODEL_CODE_LOCATION = f"s3://{_BUCKET}/{_REPACK_OUTPUT_KEY_PREFIX}"
 
 
 @pytest.fixture
@@ -688,6 +690,7 @@ def test_conditional_model_create_and_regis(
                 entry_point=f"{DATA_DIR}/{_SCRIPT_NAME}",
                 role=_ROLE,
                 enable_network_isolation=True,
+                code_location=_MODEL_CODE_LOCATION,
             ),
             2,
         ),
@@ -711,6 +714,7 @@ def test_conditional_model_create_and_regis(
                 entry_point=f"{DATA_DIR}/{_SCRIPT_NAME}",
                 role=_ROLE,
                 framework_version="1.5.0",
+                code_location=_MODEL_CODE_LOCATION,
             ),
             2,
         ),
@@ -742,6 +746,7 @@ def test_conditional_model_create_and_regis(
                 image_uri=_IMAGE_URI,
                 entry_point=f"{DATA_DIR}/{_SCRIPT_NAME}",
                 role=_ROLE,
+                code_location=_MODEL_CODE_LOCATION,
             ),
             2,
         ),
@@ -758,21 +763,45 @@ def test_conditional_model_create_and_regis(
     ],
 )
 def test_create_model_among_different_model_types(test_input, pipeline_session, model_data_param):
+    def assert_test_result(steps: list):
+        # If expected_step_num is 2, it means a runtime repack step is appended
+        # If expected_step_num is 1, it means no runtime repack is needed
+        assert len(steps) == expected_step_num
+        if expected_step_num == 2:
+            assert steps[0]["Type"] == "Training"
+            if _REPACK_OUTPUT_KEY_PREFIX == model.key_prefix:
+                assert steps[0]["Arguments"]["OutputDataConfig"]["S3OutputPath"] == (
+                    f"{_MODEL_CODE_LOCATION}/{model.name}"
+                )
+            else:
+                assert steps[0]["Arguments"]["OutputDataConfig"]["S3OutputPath"] == (
+                    f"s3://{_BUCKET}/{model.name}"
+                )
+
     model, expected_step_num = test_input
     model.sagemaker_session = pipeline_session
     model.model_data = model_data_param
-    step_args = model.create(
+    create_model_step_args = model.create(
         instance_type="c4.4xlarge",
     )
-    model_steps = ModelStep(
+    create_model_steps = ModelStep(
         name="MyModelStep",
-        step_args=step_args,
+        step_args=create_model_step_args,
     )
-    steps = model_steps.request_dicts()
+    assert_test_result(create_model_steps.request_dicts())
 
-    # If expected_step_num is 2, it means a runtime repack step is appended
-    # If expected_step_num is 1, it means no runtime repack is needed
-    assert len(steps) == expected_step_num
+    register_model_step_args = model.register(
+        content_types=["text/csv"],
+        response_types=["text/csv"],
+        inference_instances=["ml.t2.medium", "ml.m5.xlarge"],
+        transform_instances=["ml.m5.xlarge"],
+        model_package_group_name="MyModelPackageGroup",
+    )
+    register_model_steps = ModelStep(
+        name="MyModelStep",
+        step_args=register_model_step_args,
+    )
+    assert_test_result(register_model_steps.request_dicts())
 
 
 @pytest.mark.parametrize(