add delete_all() method under experiment

Author: yzhu0

src/smexperiments/experiment.py

@@ -12,7 +12,8 @@
 # language governing permissions and limitations under the License.
 
 """Contains the SageMaker Experiment class."""
-from smexperiments import _base_types, api_types, trial, _utils
+from smexperiments import _base_types, api_types, trial, _utils, trial_component
+import time
 
 
 class Experiment(_base_types.Record):
@@ -227,3 +228,44 @@ def create_trial(self, trial_name=None, trial_name_prefix="SageMakerTrial"):
             experiment_name=self.experiment_name,
             sagemaker_boto_client=self.sagemaker_boto_client,
         )
+
+    def delete_all(self, action):
+        """
+        Force to delete the experiment and associated trials, trial components under the experiment.
+
+        Args:
+            action (str): pass in string '--force' to confirm recursively delete all the experiments, trials,
+            and trial components.
+        """
+        if action != "--force":
+            raise ValueError(
+                "Must confirm with string '--force' in order to delete the experiment and "
+                "associated trials, trial components."
+            )
+
+        delete_count = 0
+        last_exception_message = None
+        while True:
+            if delete_count == 3:
+                raise Exception("Fail to delete because" + last_exception_message + ", please try again.")
+            try:
+                for trial_summary in self.list_trials():
+                    t = trial.Trial.load(
+                        sagemaker_boto_client=self.sagemaker_boto_client, trial_name=trial_summary.trial_name
+                    )
+                    for trial_component_summary in t.list_trial_components():
+                        tc = trial_component.TrialComponent.load(
+                            sagemaker_boto_client=self.sagemaker_boto_client,
+                            trial_component_name=trial_component_summary.trial_component_name,
+                        )
+                        tc.delete(force_disassociate=True)
+                        t.remove_trial_component(tc)
+                        # to prevent throttling
+                        time.sleep(0.2)
+                    t.delete()
+                self.delete()
+                break
+            except Exception as ex:
+                last_exception_message = ex
+            finally:
+                delete_count = delete_count + 1

tests/conftest.py

@@ -94,12 +94,75 @@ def tempdir():
 def experiment_obj(sagemaker_boto_client):
     description = "{}-{}".format("description", str(uuid.uuid4()))
     boto3.set_stream_logger("", logging.INFO)
+    experiment_name = name()
     experiment_obj = experiment.Experiment.create(
-        experiment_name=name(), description=description, sagemaker_boto_client=sagemaker_boto_client, tags=TAGS
+        experiment_name=experiment_name, description=description, sagemaker_boto_client=sagemaker_boto_client, tags=TAGS
     )
     yield experiment_obj
     time.sleep(0.5)
     experiment_obj.delete()
+    with pytest.raises(sagemaker_boto_client.exceptions.ResourceNotFound):
+        sagemaker_boto_client.describe_experiment(ExperimentName=experiment_name)
+
+
+@pytest.fixture
+def complex_experiment_obj(sagemaker_boto_client):
+    description = "{}-{}".format("description", str(uuid.uuid4()))
+    boto3.set_stream_logger("", logging.INFO)
+
+    # create experiment
+    experiment_obj_name = name()
+    experiment_obj = experiment.Experiment.create(
+        experiment_name=experiment_obj_name, description=description, sagemaker_boto_client=sagemaker_boto_client
+    )
+
+    # create trials
+    trial_objs = []
+    trial_name1 = name()
+    trial_name2 = name()
+    trial_name3 = name()
+
+    next_trial1 = trial.Trial.create(
+        trial_name=trial_name1, experiment_name=experiment_obj_name, sagemaker_boto_client=sagemaker_boto_client,
+    )
+    trial_objs.append(next_trial1)
+    next_trial2 = trial.Trial.create(
+        trial_name=trial_name2, experiment_name=experiment_obj_name, sagemaker_boto_client=sagemaker_boto_client,
+    )
+    trial_objs.append(next_trial2)
+    next_trial3 = trial.Trial.create(
+        trial_name=trial_name3, experiment_name=experiment_obj_name, sagemaker_boto_client=sagemaker_boto_client,
+    )
+    trial_objs.append(next_trial3)
+
+    # create trial components
+    trial_component_name = name()
+    trial_component_obj = trial_component.TrialComponent.create(
+        trial_component_name=trial_component_name, sagemaker_boto_client=sagemaker_boto_client,
+    )
+
+    # associate the trials with trial components
+    for trial_obj in trial_objs:
+        sagemaker_boto_client.associate_trial_component(
+            TrialName=trial_obj.trial_name, TrialComponentName=trial_component_obj.trial_component_name
+        )
+    time.sleep(1.0)
+    yield experiment_obj
+    experiment_obj.delete_all(action="--force")
+
+    # load experiment and verify experiment got deleted
+    with pytest.raises(sagemaker_boto_client.exceptions.ResourceNotFound):
+        sagemaker_boto_client.describe_experiment(ExperimentName=experiment_obj_name)
+    # load trials and verify trials got deleted
+    with pytest.raises(sagemaker_boto_client.exceptions.ResourceNotFound):
+        sagemaker_boto_client.describe_trial(TrialName=trial_name1)
+    with pytest.raises(sagemaker_boto_client.exceptions.ResourceNotFound):
+        sagemaker_boto_client.describe_trial(TrialName=trial_name2)
+    with pytest.raises(sagemaker_boto_client.exceptions.ResourceNotFound):
+        sagemaker_boto_client.describe_trial(TrialName=trial_name3)
+    # load trial component and verify trial component got deleted
+    with pytest.raises(sagemaker_boto_client.exceptions.ResourceNotFound):
+        sagemaker_boto_client.describe_trial_component(TrialComponentName=trial_component_name)
 
 
 @pytest.fixture

tests/integ/test_experiment.py

@@ -11,6 +11,7 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 import datetime
+import pytest
 
 from tests.helpers import name
 from smexperiments import experiment, trial
@@ -125,3 +126,12 @@ def test_list_trials(experiment_obj, trials):
     trial_names = [trial_obj.trial_name for trial_obj in trials]
     assert set(trial_names) == set([s.trial_name for s in experiment_obj.list_trials()])
     assert trial_names  # sanity test
+
+
+def test_delete_all(complex_experiment_obj):
+    assert complex_experiment_obj.experiment_name
+
+
+def test_delete_all_fails(experiment_obj):
+    with pytest.raises(ValueError):
+        experiment_obj.delete_all(action="test")

tests/unit/test_experiment.py

@@ -210,3 +210,89 @@ def test_delete(sagemaker_boto_client):
     sagemaker_boto_client.delete_experiment.return_value = {}
     obj.delete()
     sagemaker_boto_client.delete_experiment.assert_called_with(ExperimentName="foo")
+
+
+def test_delete_all_with_incorrect_action_name(sagemaker_boto_client):
+    obj = experiment.Experiment(sagemaker_boto_client, experiment_name="foo", description="bar")
+    with pytest.raises(ValueError):
+        obj.delete_all(action="abc")
+
+
+def test_delete_all(sagemaker_boto_client):
+    obj = experiment.Experiment(sagemaker_boto_client, experiment_name="foo", description="bar")
+    sagemaker_boto_client.list_trials.return_value = {
+        "TrialSummaries": [
+            {"TrialName": "trial-1", "CreationTime": datetime_obj, "LastModifiedTime": datetime_obj},
+            {"TrialName": "trial-2", "CreationTime": datetime_obj, "LastModifiedTime": datetime_obj},
+        ]
+    }
+    sagemaker_boto_client.describe_trial.side_effect = [
+        {"Trialname": "trial-1", "ExperimentName": "experiment-name-value"},
+        {"Trialname": "trial-2", "ExperimentName": "experiment-name-value"},
+    ]
+    sagemaker_boto_client.list_trial_components.side_effect = [
+        {
+            "TrialComponentSummaries": [
+                {
+                    "TrialComponentName": "trial-component-1",
+                    "CreationTime": datetime_obj,
+                    "LastModifiedTime": datetime_obj,
+                },
+                {
+                    "TrialComponentName": "trial-component-2",
+                    "CreationTime": datetime_obj,
+                    "LastModifiedTime": datetime_obj,
+                },
+            ]
+        },
+        {
+            "TrialComponentSummaries": [
+                {
+                    "TrialComponentName": "trial-component-3",
+                    "CreationTime": datetime_obj,
+                    "LastModifiedTime": datetime_obj,
+                },
+                {
+                    "TrialComponentName": "trial-component-4",
+                    "CreationTime": datetime_obj,
+                    "LastModifiedTime": datetime_obj,
+                },
+            ]
+        },
+    ]
+
+    sagemaker_boto_client.describe_trial_component.side_effect = [
+        {"TrialComponentName": "trial-component-1"},
+        {"TrialComponentName": "trial-component-2"},
+        {"TrialComponentName": "trial-component-3"},
+        {"TrialComponentName": "trial-component-4"},
+    ]
+
+    sagemaker_boto_client.delete_trial_component.return_value = {}
+    sagemaker_boto_client.delete_trial.return_value = {}
+    sagemaker_boto_client.delete_experiment.return_value = {}
+
+    obj.delete_all(action="--force")
+
+    sagemaker_boto_client.delete_experiment.assert_called_with(ExperimentName="foo")
+
+    delete_trial_expected_calls = [
+        unittest.mock.call(TrialName="trial-1"),
+        unittest.mock.call(TrialName="trial-2"),
+    ]
+    assert delete_trial_expected_calls == sagemaker_boto_client.delete_trial.mock_calls
+
+    delete_trial_component_expected_calls = [
+        unittest.mock.call(TrialComponentName="trial-component-1"),
+        unittest.mock.call(TrialComponentName="trial-component-2"),
+        unittest.mock.call(TrialComponentName="trial-component-3"),
+        unittest.mock.call(TrialComponentName="trial-component-4"),
+    ]
+    assert delete_trial_component_expected_calls == sagemaker_boto_client.delete_trial_component.mock_calls
+
+
+def test_delete_all_fail(sagemaker_boto_client):
+    obj = experiment.Experiment(sagemaker_boto_client, experiment_name="foo", description="bar")
+    sagemaker_boto_client.list_trials.side_effect = Exception
+    with pytest.raises(Exception):
+        obj.delete_all(action="--force")