Supporting tbac in load_run

Author: ananth102

src/sagemaker/experiments/run.py

@@ -771,6 +771,7 @@ def load_run(
     sagemaker_session: Optional["Session"] = None,
     artifact_bucket: Optional[str] = None,
     artifact_prefix: Optional[str] = None,
+    tags: Optional[List[Dict[str, str]]] = None,
 ) -> Run:
     """Load an existing run.
 
@@ -839,6 +840,8 @@ def load_run(
                 will be used.
         artifact_prefix (str): The S3 key prefix used to generate the S3 path
             to upload the artifact to (default: "trial-component-artifacts").
+        tags (List[Dict[str, str]]): A list of tags to be used for all create calls,
+            e.g. to create an experiment, a run group, etc. (default: None).
 
     Returns:
         Run: The loaded Run object.
@@ -860,6 +863,7 @@ def load_run(
             sagemaker_session=sagemaker_session or _utils.default_session(),
             artifact_bucket=artifact_bucket,
             artifact_prefix=artifact_prefix,
+            tags=tags,
         )
     elif _RunContext.get_current_run():
         run_instance = _RunContext.get_current_run()
@@ -879,6 +883,7 @@ def load_run(
             sagemaker_session=sagemaker_session or _utils.default_session(),
             artifact_bucket=artifact_bucket,
             artifact_prefix=artifact_prefix,
+            tags=tags,
         )
     else:
         raise RuntimeError(

tests/unit/sagemaker/experiments/test_run.py

@@ -55,6 +55,7 @@
     TEST_RUN_DISPLAY_NAME,
     TEST_ARTIFACT_BUCKET,
     TEST_ARTIFACT_PREFIX,
+    TEST_TAGS
 )
 
 
@@ -155,24 +156,22 @@ def test_run_init_name_length_exceed_limit(sagemaker_session):
 
 
 @pytest.mark.parametrize(
-    ("kwargs", "expected_artifact_bucket", "expected_artifact_prefix"),
+    ("kwargs", "expected_artifact_bucket", "expected_artifact_prefix","expected_tags"),
     [
-        ({}, None, _DEFAULT_ARTIFACT_PREFIX),
+        ({}, None, _DEFAULT_ARTIFACT_PREFIX, None),
         (
             {
                 "artifact_bucket": TEST_ARTIFACT_BUCKET,
                 "artifact_prefix": TEST_ARTIFACT_PREFIX,
+                "tags": TEST_TAGS
             },
             TEST_ARTIFACT_BUCKET,
             TEST_ARTIFACT_PREFIX,
+            TEST_TAGS
         ),
     ],
 )
 @patch.object(_TrialComponent, "save", MagicMock(return_value=None))
-@patch(
-    "sagemaker.experiments.run.Experiment._load_or_create",
-    MagicMock(return_value=Experiment(experiment_name=TEST_EXP_NAME)),
-)
 @patch(
     "sagemaker.experiments.run._Trial._load_or_create",
     MagicMock(side_effect=mock_trial_load_or_create_func),
@@ -189,6 +188,7 @@ def test_run_load_no_run_name_and_in_train_job(
     kwargs,
     expected_artifact_bucket,
     expected_artifact_prefix,
+    expected_tags
 ):
     client = sagemaker_session.sagemaker_client
     job_name = "my-train-job"
@@ -220,19 +220,22 @@ def test_run_load_no_run_name_and_in_train_job(
             }
         ]
     }
-    with load_run(sagemaker_session=sagemaker_session, **kwargs) as run_obj:
-        assert run_obj._in_load
-        assert not run_obj._inside_init_context
-        assert run_obj._inside_load_context
-        assert run_obj.run_name == TEST_RUN_NAME
-        assert run_obj._trial_component.trial_component_name == expected_tc_name
-        assert run_obj.run_group_name == Run._generate_trial_name(TEST_EXP_NAME)
-        assert run_obj._trial
-        assert run_obj.experiment_name == TEST_EXP_NAME
-        assert run_obj._experiment
-        assert run_obj.experiment_config == exp_config
-        assert run_obj._artifact_uploader.artifact_bucket == expected_artifact_bucket
-        assert run_obj._artifact_uploader.artifact_prefix == expected_artifact_prefix
+    expmock = MagicMock(return_value=Experiment(experiment_name=TEST_EXP_NAME,tags=expected_tags))
+    with patch("sagemaker.experiments.run.Experiment._load_or_create", expmock):
+        with load_run(sagemaker_session=sagemaker_session, **kwargs) as run_obj:
+            assert run_obj._in_load
+            assert not run_obj._inside_init_context
+            assert run_obj._inside_load_context
+            assert run_obj.run_name == TEST_RUN_NAME
+            assert run_obj._trial_component.trial_component_name == expected_tc_name
+            assert run_obj.run_group_name == Run._generate_trial_name(TEST_EXP_NAME)
+            assert run_obj._trial
+            assert run_obj.experiment_name == TEST_EXP_NAME
+            assert run_obj._experiment
+            assert run_obj.experiment_config == exp_config
+            assert run_obj._artifact_uploader.artifact_bucket == expected_artifact_bucket
+            assert run_obj._artifact_uploader.artifact_prefix == expected_artifact_prefix
+            assert run_obj._experiment.tags == expected_tags
 
     client.describe_training_job.assert_called_once_with(TrainingJobName=job_name)
     run_obj._trial.add_trial_component.assert_not_called()