Add disassociate param under the TC delete method

Author: yzhu0

src/smexperiments/trial_component.py

@@ -103,8 +103,39 @@ def save(self):
         """Save the state of this TrialComponent to SageMaker."""
         return self._invoke_api(self._boto_update_method, self._boto_update_members)
 
-    def delete(self):
-        """Delete this TrialComponent from SageMaker."""
+    def delete(self, force_disassociate=None):
+        """Delete this TrialComponent from SageMaker.
+
+        Args:
+            force_disassociate (boolean): Indicates whether to force disassociate the trial component with the trials
+            before deletion. If set to true, force disassociate the trial component with associated trials first, then
+            delete the trial component. If it's not set or set to false, it will delete the trial component directory
+             without disassociation.
+        """
+        if force_disassociate:
+            next_token = None
+
+            while True:
+                if next_token:
+                    list_trials_response = self.sagemaker_boto_client.list_trials(
+                        TrialComponentName=self.trial_component_name, NextToken=next_token
+                    )
+                else:
+                    list_trials_response = self.sagemaker_boto_client.list_trials(
+                        TrialComponentName=self.trial_component_name
+                    )
+
+                # Disassociate the trials and trial components
+                for trial in list_trials_response["TrialSummaries"]:
+                    self.sagemaker_boto_client.disassociate_trial_component(
+                        TrialName=trial["TrialName"], TrialComponentName=self.trial_component_name
+                    )
+
+                if "NextToken" in list_trials_response:
+                    next_token = list_trials_response["NextToken"]
+                else:
+                    break
+
         self._invoke_api(self._boto_delete_method, self._boto_delete_members)
 
     @classmethod

tests/conftest.py

@@ -125,6 +125,20 @@ def trial_component_obj(sagemaker_boto_client):
     trial_component_obj.delete()
 
 
+@pytest.fixture
+def trial_component_with_force_disassociation_obj(trials, sagemaker_boto_client):
+    trial_component_obj = trial_component.TrialComponent.create(
+        trial_component_name=name(), sagemaker_boto_client=sagemaker_boto_client
+    )
+    for trial in trials:
+        sagemaker_boto_client.associate_trial_component(
+            TrialName=trial.trial_name, TrialComponentName=trial_component_obj.trial_component_name
+        )
+    yield trial_component_obj
+    time.sleep(0.5)
+    trial_component_obj.delete(force_disassociate=True)
+
+
 @pytest.fixture
 def trials(experiment_obj, sagemaker_boto_client):
     trial_objs = []

tests/integ/test_trial_component.py

@@ -31,6 +31,18 @@ def test_create_tags(trial_component_obj, sagemaker_boto_client):
     assert actual_tags == trial_component_obj.tags
 
 
+def test_delete_with_disassociate(trial_component_with_disassociation_obj, sagemaker_boto_client):
+    assert trial_component_with_disassociation_obj.trial_component_name
+
+
+def test_delete_with_force_disassociate(trial_component_with_force_disassociation_obj, sagemaker_boto_client):
+    assert trial_component_with_force_disassociation_obj.trial_component_name
+    trials = sagemaker_boto_client.list_trials(
+        TrialComponentName=trial_component_with_force_disassociation_obj.trial_component_name
+    )["TrialSummaries"]
+    assert len(trials) == 3
+
+
 def test_save(trial_component_obj, sagemaker_boto_client):
     trial_component_obj.display_name = str(uuid.uuid4())
     trial_component_obj.status = api_types.TrialComponentStatus(primary_status="InProgress", message="Message")

tests/unit/test_trial_component.py

@@ -253,6 +253,26 @@ def test_delete(sagemaker_boto_client):
     sagemaker_boto_client.delete_trial_component.assert_called_with(TrialComponentName="foo")
 
 
+def test_delete_with_force_disassociate(sagemaker_boto_client):
+    obj = trial_component.TrialComponent(sagemaker_boto_client, trial_component_name="foo", display_name="bar")
+    sagemaker_boto_client.delete_trial_component.return_value = {}
+
+    sagemaker_boto_client.list_trials.side_effect = [
+        {"TrialSummaries": [{"TrialName": "trial-1"}, {"TrialName": "trial-2"}], "NextToken": "a"},
+        {"TrialSummaries": [{"TrialName": "trial-3"}, {"TrialName": "trial-4"}]},
+    ]
+
+    obj.delete(force_disassociate=True)
+    expected_calls = [
+        unittest.mock.call(TrialName="trial-1", TrialComponentName="foo"),
+        unittest.mock.call(TrialName="trial-2", TrialComponentName="foo"),
+        unittest.mock.call(TrialName="trial-3", TrialComponentName="foo"),
+        unittest.mock.call(TrialName="trial-4", TrialComponentName="foo"),
+    ]
+    assert expected_calls == sagemaker_boto_client.disassociate_trial_component.mock_calls
+    sagemaker_boto_client.delete_trial_component.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"]