feat: Add check for if TrialComponent is already associated with a Trial in Run (#3956)

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

* move _trial_component_is_associated_to_trial to private and move it as a method under trial_component, add doc-string and unit-tests

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,
@@ -200,7 +203,11 @@ def __init__(
                 self.run_name,
                 self.experiment_name,
             )
-        self._trial.add_trial_component(self._trial_component)
+
+        if not _TrialComponent._trial_component_is_associated_to_trial(
+            self._trial_component.trial_component_name, self._trial.trial_name, sagemaker_session
+        ):
+            self._trial.add_trial_component(self._trial_component)
 
         self._artifact_uploader = _ArtifactUploader(
             trial_component_name=self._trial_component.trial_component_name,
@@ -348,7 +355,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
@@ -381,7 +391,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")
@@ -432,7 +444,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")
@@ -447,12 +461,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.
 
@@ -575,11 +596,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):
@@ -719,7 +746,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(
@@ -837,7 +865,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],

src/sagemaker/experiments/trial_component.py

@@ -345,3 +345,42 @@ def _load_or_create(
             run_tc = _TrialComponent.load(trial_component_name, sagemaker_session)
             is_existed = True
         return run_tc, is_existed
+
+    @classmethod
+    def _trial_component_is_associated_to_trial(
+        cls, trial_component_name, trial_name=None, sagemaker_session=None
+    ):
+        """Returns a bool based on if trial_component is already associated with the trial.
+
+        Args:
+            trial_component_name (str): The name of the trial component.
+            trial_name: (str): The name of the trial.
+            sagemaker_session (sagemaker.session.Session): Session object which
+                manages interactions with Amazon SageMaker APIs and any other
+                AWS services needed.
+
+        Returns:
+            bool: A boolean variable indicating whether the trial component is already
+                  associated with the trial.
+
+        """
+        search_results = sagemaker_session.sagemaker_client.search(
+            Resource="ExperimentTrialComponent",
+            SearchExpression={
+                "Filters": [
+                    {
+                        "Name": "TrialComponentName",
+                        "Operator": "Equals",
+                        "Value": str(trial_component_name),
+                    },
+                    {
+                        "Name": "Parents.TrialName",
+                        "Operator": "Equals",
+                        "Value": str(trial_name),
+                    },
+                ]
+            },
+        )
+        if search_results["Results"]:
+            return True
+        return False

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

@@ -93,6 +93,7 @@ def test_run_init(
     expected_artifact_bucket,
     expected_artifact_prefix,
 ):
+    sagemaker_session.sagemaker_client.search.return_value = {"Results": []}
     with Run(
         experiment_name=TEST_EXP_NAME,
         run_name=TEST_RUN_NAME,
@@ -121,6 +122,7 @@ def test_run_init(
 
     # 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):
@@ -206,6 +208,18 @@ def test_run_load_no_run_name_and_in_train_job(
         # 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, **kwargs) as run_obj:
         assert run_obj._in_load
         assert not run_obj._inside_init_context
@@ -221,6 +235,7 @@ def test_run_load_no_run_name_and_in_train_job(
         assert run_obj._artifact_uploader.artifact_prefix == expected_artifact_prefix
 
     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")
@@ -296,6 +311,7 @@ def test_run_load_no_run_name_and_not_in_train_job_but_no_obj_in_context(sagemak
 def test_run_load_with_run_name_and_exp_name(
     sagemaker_session, kwargs, expected_artifact_bucket, expected_artifact_prefix
 ):
+    sagemaker_session.sagemaker_client.search.return_value = {"Results": []}
     with load_run(
         run_name=TEST_RUN_NAME,
         experiment_name=TEST_EXP_NAME,
@@ -319,6 +335,8 @@ def test_run_load_with_run_name_and_exp_name(
         assert run_obj._artifact_uploader.artifact_bucket == expected_artifact_bucket
         assert run_obj._artifact_uploader.artifact_prefix == expected_artifact_prefix
 
+    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:
@@ -365,11 +383,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(
@@ -406,11 +437,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(
@@ -428,6 +472,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/experiments/test_trial_component.py

@@ -396,3 +396,26 @@ def test_search(sagemaker_session):
         ),
     ]
     assert expected == list(_TrialComponent.search(sagemaker_session=sagemaker_session))
+
+
+def test_trial_component_is_associated_to_trial(sagemaker_session):
+    obj = _TrialComponent(sagemaker_session, trial_component_name="tc-1")
+    sagemaker_session.sagemaker_client.search.return_value = {
+        "Results": [
+            {
+                "TrialComponent": {
+                    "Parents": [{"ExperimentName": "e-1", "TrialName": "t-1"}],
+                    "TrialComponentName": "tc-1",
+                }
+            }
+        ]
+    }
+
+    assert obj._trial_component_is_associated_to_trial("tc-1", "t-1", sagemaker_session) is True
+
+
+def test_trial_component_is_not_associated_to_trial(sagemaker_session):
+    obj = _TrialComponent(sagemaker_session, trial_component_name="tc-1")
+    sagemaker_session.sagemaker_client.search.return_value = {"Results": []}
+
+    assert obj._trial_component_is_associated_to_trial("tc-1", "t-1", sagemaker_session) is False

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

@@ -114,6 +114,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,