feature: add check for if TrialComponent is already associated with a Trial

Author: ryansteakley

src/sagemaker/experiments/run.py

@@ -25,7 +25,10 @@
 
 from sagemaker.apiutils import _utils
 from sagemaker.experiments import _api_types
-from sagemaker.experiments._api_types import TrialComponentArtifact, _TrialComponentStatusType
+from sagemaker.experiments._api_types import (
+    TrialComponentArtifact,
+    _TrialComponentStatusType,
+)
 from sagemaker.experiments._helper import (
     _ArtifactUploader,
     _LineageArtifactTracker,
@@ -192,7 +195,31 @@ def __init__(
                 self.run_name,
                 self.experiment_name,
             )
-        self._trial.add_trial_component(self._trial_component)
+
+        def search_trial_component_associated_trial():
+            search_results = sagemaker_session.sagemaker_client.search(
+                Resource="ExperimentTrialComponent",
+                SearchExpression={
+                    "Filters": [
+                        {
+                            "Name": "TrialComponentName",
+                            "Operator": "Equals",
+                            "Value": str(self._trial_component.trial_component_name),
+                        },
+                        {
+                            "Name": "Parents.TrialName",
+                            "Operator": "Equals",
+                            "Value": str(self._trial.trial_name),
+                        },
+                    ]
+                },
+            )
+            if search_results["Results"]:
+                return True
+            return False
+
+        if not search_trial_component_associated_trial():
+            self._trial.add_trial_component(self._trial_component)
 
         self._artifact_uploader = _ArtifactUploader(
             trial_component_name=self._trial_component.trial_component_name,
@@ -336,7 +363,10 @@ def log_precision_recall(
             "noSkill": no_skill,
         }
         self._log_graph_artifact(
-            artifact_name=title, data=data, graph_type="PrecisionRecallCurve", is_output=is_output
+            artifact_name=title,
+            data=data,
+            graph_type="PrecisionRecallCurve",
+            is_output=is_output,
         )
 
     @validate_invoked_inside_run_context
@@ -369,7 +399,9 @@ def log_roc_curve(
                 If set to False then represented as input association.
         """
         verify_length_of_true_and_predicted(
-            true_labels=y_true, predicted_attrs=y_score, predicted_attrs_name="predicted scores"
+            true_labels=y_true,
+            predicted_attrs=y_score,
+            predicted_attrs_name="predicted scores",
         )
 
         get_module("sklearn")
@@ -420,7 +452,9 @@ def log_confusion_matrix(
                 If set to False then represented as input association.
         """
         verify_length_of_true_and_predicted(
-            true_labels=y_true, predicted_attrs=y_pred, predicted_attrs_name="predicted labels"
+            true_labels=y_true,
+            predicted_attrs=y_pred,
+            predicted_attrs_name="predicted labels",
         )
 
         get_module("sklearn")
@@ -435,12 +469,19 @@ def log_confusion_matrix(
             "confusionMatrix": matrix.tolist(),
         }
         self._log_graph_artifact(
-            artifact_name=title, data=data, graph_type="ConfusionMatrix", is_output=is_output
+            artifact_name=title,
+            data=data,
+            graph_type="ConfusionMatrix",
+            is_output=is_output,
         )
 
     @validate_invoked_inside_run_context
     def log_artifact(
-        self, name: str, value: str, media_type: Optional[str] = None, is_output: bool = True
+        self,
+        name: str,
+        value: str,
+        media_type: Optional[str] = None,
+        is_output: bool = True,
     ):
         """Record a single artifact for this run.
 
@@ -563,11 +604,17 @@ def _log_graph_artifact(self, data, graph_type, is_output, artifact_name=None):
         # create an artifact and association for the table
         if is_output:
             self._lineage_artifact_tracker.add_output_artifact(
-                name=artifact_name, source_uri=s3_uri, etag=etag, artifact_type=graph_type
+                name=artifact_name,
+                source_uri=s3_uri,
+                etag=etag,
+                artifact_type=graph_type,
             )
         else:
             self._lineage_artifact_tracker.add_input_artifact(
-                name=artifact_name, source_uri=s3_uri, etag=etag, artifact_type=graph_type
+                name=artifact_name,
+                source_uri=s3_uri,
+                etag=etag,
+                artifact_type=graph_type,
             )
 
     def _verify_trial_component_artifacts_length(self, is_output):
@@ -707,7 +754,8 @@ def __exit__(self, exc_type, exc_value, exc_traceback):
         self._trial_component.end_time = end_time
         if exc_value:
             self._trial_component.status = _api_types.TrialComponentStatus(
-                primary_status=_TrialComponentStatusType.Failed.value, message=str(exc_value)
+                primary_status=_TrialComponentStatusType.Failed.value,
+                message=str(exc_value),
             )
         else:
             self._trial_component.status = _api_types.TrialComponentStatus(
@@ -816,7 +864,8 @@ def load_run(
         run_instance = _RunContext.get_current_run()
     elif environment:
         exp_config = get_tc_and_exp_config_from_job_env(
-            environment=environment, sagemaker_session=sagemaker_session or _utils.default_session()
+            environment=environment,
+            sagemaker_session=sagemaker_session or _utils.default_session(),
         )
         run_name = Run._extract_run_name_from_tc_name(
             trial_component_name=exp_config[RUN_NAME],

tests/unit/sagemaker/experiments/conftest.py

@@ -72,6 +72,7 @@ def run_obj(sagemaker_session):
                 "sagemaker.experiments.run._Trial._load_or_create",
                 MagicMock(side_effect=mock_trial_load_or_create_func),
             ):
+                sagemaker_session.sagemaker_client.search.return_value = {"Results": []}
                 run = Run(
                     experiment_name=TEST_EXP_NAME,
                     sagemaker_session=sagemaker_session,

tests/unit/sagemaker/experiments/test_run.py

@@ -70,6 +70,7 @@
 )
 @patch.object(_TrialComponent, "save")
 def test_run_init(mock_tc_save, sagemaker_session):
+    sagemaker_session.sagemaker_client.search.return_value = {"Results": []}
     with Run(
         experiment_name=TEST_EXP_NAME, run_name=TEST_RUN_NAME, sagemaker_session=sagemaker_session
     ) as run_obj:
@@ -93,6 +94,7 @@ def test_run_init(mock_tc_save, sagemaker_session):
 
     # trail_component.save is called when entering/ exiting the with block
     mock_tc_save.assert_called()
+    run_obj._trial.add_trial_component.assert_called()
 
 
 def test_run_init_name_length_exceed_limit(sagemaker_session):
@@ -158,6 +160,18 @@ def test_run_load_no_run_name_and_in_train_job(mock_run_env, sagemaker_session):
         # The Run object has been created else where
         "ExperimentConfig": exp_config,
     }
+    sagemaker_session.sagemaker_client.search.return_value = {
+        "Results": [
+            {
+                "TrialComponent": {
+                    "Parents": [
+                        {"ExperimentName": TEST_EXP_NAME, "TrialName": exp_config[TRIAL_NAME]}
+                    ],
+                    "TrialComponentName": expected_tc_name,
+                }
+            }
+        ]
+    }
     with load_run(sagemaker_session=sagemaker_session) as run_obj:
         assert run_obj._in_load
         assert not run_obj._inside_init_context
@@ -171,6 +185,7 @@ def test_run_load_no_run_name_and_in_train_job(mock_run_env, sagemaker_session):
         assert run_obj.experiment_config == exp_config
 
     client.describe_training_job.assert_called_once_with(TrainingJobName=job_name)
+    run_obj._trial.add_trial_component.assert_not_called()
 
 
 @patch("sagemaker.experiments.run._RunEnvironment")
@@ -230,6 +245,7 @@ def test_run_load_no_run_name_and_not_in_train_job_but_no_obj_in_context(sagemak
     MagicMock(side_effect=mock_tc_load_or_create_func),
 )
 def test_run_load_with_run_name_and_exp_name(sagemaker_session):
+    sagemaker_session.sagemaker_client.search.return_value = {"Results": []}
     with load_run(
         run_name=TEST_RUN_NAME,
         experiment_name=TEST_EXP_NAME,
@@ -250,6 +266,8 @@ def test_run_load_with_run_name_and_exp_name(sagemaker_session):
         assert run_obj._experiment
         assert run_obj.experiment_config == expected_exp_config
 
+    run_obj._trial.add_trial_component.assert_called()
+
 
 def test_run_load_with_run_name_but_no_exp_name(sagemaker_session):
     with pytest.raises(ValueError) as err:
@@ -296,11 +314,24 @@ def test_run_load_in_sm_processing_job(mock_run_env, sagemaker_session):
         # The Run object has been created else where
         "ExperimentConfig": exp_config,
     }
+    sagemaker_session.sagemaker_client.search.return_value = {
+        "Results": [
+            {
+                "TrialComponent": {
+                    "Parents": [
+                        {"ExperimentName": TEST_EXP_NAME, "TrialName": exp_config[TRIAL_NAME]}
+                    ],
+                    "TrialComponentName": expected_tc_name,
+                }
+            }
+        ]
+    }
 
     with load_run(sagemaker_session=sagemaker_session):
         pass
 
     client.describe_processing_job.assert_called_once_with(ProcessingJobName=job_name)
+    mock_run_env._trial.add_trial_component.assert_not_called()
 
 
 @patch(
@@ -337,11 +368,24 @@ def test_run_load_in_sm_transform_job(mock_run_env, sagemaker_session):
         # The Run object has been created else where
         "ExperimentConfig": exp_config,
     }
+    sagemaker_session.sagemaker_client.search.return_value = {
+        "Results": [
+            {
+                "TrialComponent": {
+                    "Parents": [
+                        {"ExperimentName": TEST_EXP_NAME, "TrialName": exp_config[TRIAL_NAME]}
+                    ],
+                    "TrialComponentName": expected_tc_name,
+                }
+            }
+        ]
+    }
 
     with load_run(sagemaker_session=sagemaker_session):
         pass
 
     client.describe_transform_job.assert_called_once_with(TransformJobName=job_name)
+    mock_run_env._trial.add_trial_component.assert_not_called()
 
 
 @patch(
@@ -359,6 +403,7 @@ def test_run_load_in_sm_transform_job(mock_run_env, sagemaker_session):
 )
 @patch.object(_TrialComponent, "save")
 def test_run_object_serialize_deserialize(mock_tc_save, sagemaker_session):
+    sagemaker_session.sagemaker_client.search.return_value = {"Results": []}
     run_obj = Run(
         experiment_name=TEST_EXP_NAME,
         run_name=TEST_RUN_NAME,

tests/unit/sagemaker/remote_function/core/test_stored_function.py

@@ -109,6 +109,7 @@ def test_save_with_parameter_of_run_type(
 ):
     session = Mock()
     s3_base_uri = random_s3_uri()
+    session.sagemaker_client.search.return_value = {"Results": []}
 
     run = Run(
         experiment_name=TEST_EXP_NAME,

tests/unit/sagemaker/remote_function/test_client.py

@@ -113,6 +113,7 @@ def run_obj(sagemaker_session):
                 "sagemaker.experiments.run._Trial._load_or_create",
                 MagicMock(side_effect=mock_trial_load_or_create_func),
             ):
+                sagemaker_session.sagemaker_client.search.return_value = {"Results": []}
                 run = Run(
                     experiment_name="test-exp",
                     sagemaker_session=sagemaker_session,