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

@@ -196,29 +196,9 @@ def __init__(
                 self.experiment_name,
             )
 
-        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():
+        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(

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/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