add helper function to get tuning step top performing model s3 uri

Author: jerrypeng7773

src/sagemaker/workflow/steps.py

@@ -44,6 +44,7 @@
     PropertyFile,
     Properties,
 )
+from sagemaker.workflow.functions import Join
 
 
 class StepTypeEnum(Enum, metaclass=DefaultEnumMeta):
@@ -525,3 +526,27 @@ def to_request(self) -> RequestType:
             request_dict.update(self.cache_config.config)
 
         return request_dict
+
+    def get_top_model_s3_uri(self, top_k: int, s3_bucket: str, prefix: str = ""):
+        """Get the model artifact s3 uri from the top performing training jobs.
+
+        Args:
+            top_k (int): the index of the top performing training job
+                tuning step stores up to 50 top performing training jobs, hence
+                a valid top_k value is from 0 to 49. The best training job
+                model is at index 0
+            s3_bucket (str): the s3 bucket to store the training job output artifact
+            prefix (str): the s3 key prefix to store the training job output artifact
+        """
+        values = ["s3:/", s3_bucket]
+        if prefix != "" and prefix is not None:
+            values.append(prefix)
+
+        return Join(
+            on="/",
+            values=values
+            + [
+                self.properties.TrainingJobSummaries[top_k].TrainingJobName,
+                "output/model.tar.gz",
+            ],
+        )

tests/integ/test_workflow.py

@@ -884,6 +884,8 @@ def test_tuning(
         objective_type="Maximize",
         hyperparameter_ranges=hyperparameter_ranges,
         metric_definitions=[{"Name": "test:acc", "Regex": "Overall test accuracy: (.*?);"}],
+        max_jobs=2,
+        max_parallel_jobs=2,
     )
 
     step_tune = TuningStep(
@@ -892,13 +894,48 @@ def test_tuning(
         inputs=inputs,
     )
 
+    best_model = Model(
+        image_uri=pytorch_estimator.training_image_uri(),
+        model_data=step_tune.get_top_model_s3_uri(
+            top_k=0,
+            s3_bucket=sagemaker_session.default_bucket(),
+        ),
+        sagemaker_session=sagemaker_session,
+        role=role,
+    )
+    model_inputs = CreateModelInput(
+        instance_type="ml.m5.large",
+        accelerator_type="ml.eia1.medium",
+    )
+    step_best_model = CreateModelStep(
+        name="1st-model",
+        model=best_model,
+        inputs=model_inputs,
+    )
+
+    second_best_model = Model(
+        image_uri=pytorch_estimator.training_image_uri(),
+        model_data=step_tune.get_top_model_s3_uri(
+            top_k=1,
+            s3_bucket=sagemaker_session.default_bucket(),
+        ),
+        sagemaker_session=sagemaker_session,
+        role=role,
+    )
+
+    step_second_best_model = CreateModelStep(
+        name="2nd-best-model",
+        model=second_best_model,
+        inputs=model_inputs,
+    )
+
     pipeline = Pipeline(
         name=pipeline_name,
         parameters=[instance_count, instance_type],
-        steps=[step_tune],
+        steps=[step_tune, step_best_model, step_second_best_model],
         sagemaker_session=sagemaker_session,
     )
-    print(pipeline.definition())
+
     try:
         response = pipeline.create(role)
         create_arn = response["PipelineArn"]

tests/unit/sagemaker/workflow/test_steps.py

@@ -551,7 +551,7 @@ def test_single_algo_tuning_step(sagemaker_session):
         tuner=tuner,
         inputs=inputs,
     )
-    print(tuning_step.to_request())
+
     assert tuning_step.to_request() == {
         "Name": "MyTuningStep",
         "Type": "Tuning",
@@ -644,6 +644,18 @@ def test_single_algo_tuning_step(sagemaker_session):
     assert tuning_step.properties.TrainingJobSummaries[0].TrainingJobName.expr == {
         "Get": "Steps.MyTuningStep.TrainingJobSummaries[0].TrainingJobName"
     }
+    assert tuning_step.get_top_model_s3_uri(0, "my-bucket", "my-prefix").expr == {
+        "Std:Join": {
+            "On": "/",
+            "Values": [
+                "s3:/",
+                "my-bucket",
+                "my-prefix",
+                {"Get": "Steps.MyTuningStep.TrainingJobSummaries[0].TrainingJobName"},
+                "output/model.tar.gz",
+            ],
+        }
+    }
 
 
 def test_multi_algo_tuning_step(sagemaker_session):