list trials under trial component

Author: yzhu0

src/smexperiments/trial_component.py

@@ -11,7 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 """Contains the TrialComponent class."""
-from smexperiments import _base_types, api_types
+from smexperiments import _base_types, api_types, trial
 import time
 
 
@@ -130,11 +130,11 @@ def delete(self, force_disassociate=None):
                     )
 
                 # Disassociate the trials and trial components
-                for trial in list_trials_response["TrialSummaries"]:
+                for per_trial in list_trials_response["TrialSummaries"]:
                     # to prevent DisassociateTrialComponent throttling
                     time.sleep(1.2)
                     self.sagemaker_boto_client.disassociate_trial_component(
-                        TrialName=trial["TrialName"], TrialComponentName=self.trial_component_name
+                        TrialName=per_trial["TrialName"], TrialComponentName=self.trial_component_name
                     )
 
                 if "NextToken" in list_trials_response:
@@ -144,6 +144,17 @@ def delete(self, force_disassociate=None):
 
         return self._invoke_api(self._boto_delete_method, self._boto_delete_members)
 
+    def list_trials(self):
+        """
+        Load a list of trials that contains the same trial component name
+
+        Returns:
+            A list of trials that contains the same trial component name
+        """
+        return trial.Trial.list(
+            trial_component_name=self.trial_component_name, sagemaker_boto_client=self.sagemaker_boto_client
+        )
+
     @classmethod
     def load(cls, trial_component_name, sagemaker_boto_client=None):
         """Load an existing trial component and return an ``TrialComponent`` object representing it.

tests/integ/test_trial_component.py

@@ -112,6 +112,27 @@ def test_list_sort(trial_components, sagemaker_boto_client):
     assert trial_component_names  # sanity test
 
 
+def test_list_empty_trials(experiment_obj, trial_component_obj):
+    actual_value = trial_component_obj.list_trials()
+    assert len(list(actual_value)) == 0
+
+
+def test_list_trials(trial_obj, trial_component_obj, sagemaker_boto_client):
+    trial_obj.add_trial_component(trial_component_obj)
+    actual_value = list(trial_component_obj.list_trials())
+
+    assert len(actual_value) == 1
+    assert actual_value[0].trial_name == trial_obj.trial_name
+    trial_obj.remove_trial_component(trial_component_obj)
+
+    trial_components = list(
+        trial_component.TrialComponent.list(
+            sagemaker_boto_client=sagemaker_boto_client, trial_name=trial_obj.trial_name
+        )
+    )
+    assert 0 == len(trial_components)
+
+
 def test_list_trial_components_by_experiment(experiment_obj, trial_component_obj, sagemaker_boto_client):
     trial_obj = experiment_obj.create_trial()
     trial_obj.add_trial_component(trial_component_obj)

tests/unit/test_trial_component.py

@@ -10,7 +10,7 @@
 # distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
-from smexperiments import trial_component, api_types
+from smexperiments import trial_component, api_types, trial
 
 import datetime
 import pytest
@@ -285,6 +285,31 @@ def test_delete_with_force_disassociate(sagemaker_boto_client):
     sagemaker_boto_client.delete_trial_component.assert_called_with(TrialComponentName="foo")
 
 
+def test_list_trials(sagemaker_boto_client):
+    sagemaker_boto_client.list_trials.return_value = {
+        "TrialSummaries": [
+            {
+                "TrialName": "trial-1",
+                "CreationTime": None,
+                "LastModifiedTime": None,
+            },
+            {
+                "TrialName": "trial-2",
+                "CreationTime": None,
+                "LastModifiedTime": None,
+            },
+        ]
+    }
+    obj = trial_component.TrialComponent(sagemaker_boto_client, trial_component_name="foo", discompplay_name="bar")
+
+    expected = [
+        api_types.TrialSummary(trial_name="trial-1", creation_time=None, last_modified_time=None),
+        api_types.TrialSummary(trial_name="trial-2", creation_time=None, last_modified_time=None),
+    ]
+    assert expected == list(obj.list_trials())
+    sagemaker_boto_client.list_trials.assert_called_with(TrialComponentName="foo")
+
+
 def test_boto_ignore():
     obj = trial_component.TrialComponent(sagemaker_boto_client, trial_component_name="foo", display_name="bar")
     assert obj._boto_ignore() == ["ResponseMetadata", "CreatedBy"]