feature: add delete_all method to Trial

Author: danabens

scripts/version.py

@@ -49,7 +49,11 @@ def parse(version):
     if not match:
         raise ValueError(f"invalid version: {version}")
 
-    return Version(int(match.group("major") or 0), int(match.group("minor") or 0), int(match.group("patch") or 0),)
+    return Version(
+        int(match.group("major") or 0),
+        int(match.group("minor") or 0),
+        int(match.group("patch") or 0),
+    )
 
 
 def next_version_from_current_version(current_version, increment_type):

src/smexperiments/_base_types.py

@@ -21,7 +21,7 @@ class ApiObject(object):
     to dicts into/from a Python object with standard python members. Clients invoke to_boto on an instance
      of ApiObject to transform the ApiObject into a boto representation. Clients invoke from_boto on a sub-class of
      ApiObject to instantiate an instance of that class from a boto representation.
-     """
+    """
 
     # A map from boto 'UpperCamelCase' name to member name. If a boto name does not appear in this dict then
     # it is converted to lower_snake_case.
@@ -81,7 +81,8 @@ def __hash__(self):
     def __repr__(self):
         """Returns a string representation of this ApiObject."""
         return "{}({})".format(
-            type(self).__name__, ",".join(["{}={}".format(k, repr(v)) for k, v in vars(self).items()]),
+            type(self).__name__,
+            ",".join(["{}={}".format(k, repr(v)) for k, v in vars(self).items()]),
         )
 
 

src/smexperiments/_boto_functions.py

@@ -83,12 +83,12 @@ def from_boto(boto_dict, boto_name_to_member_name, member_name_to_type):
 def to_boto(member_vars, member_name_to_boto_name, member_name_to_type):
     """Convert a dict of of snake case names to values into a boto UpperCamelCase representation.
 
-       Args:
-           member_vars dict[str, ?]: A map from snake case name to value.
-           member_name_to_boto_name dict[str, ?]: A map from snake_case name to boto name.
+    Args:
+        member_vars dict[str, ?]: A map from snake case name to value.
+        member_name_to_boto_name dict[str, ?]: A map from snake_case name to boto name.
 
-        Returns:
-            dict: boto dict converted to snake case
+     Returns:
+         dict: boto dict converted to snake case
 
     """
     to_boto_values = {}

src/smexperiments/experiment.py

@@ -59,6 +59,8 @@ class Experiment(_base_types.Record):
     _boto_update_members = ["experiment_name", "description", "display_name"]
     _boto_delete_members = ["experiment_name"]
 
+    MAX_DELETE_ALL_ATTEMPTS = 3
+
     def save(self):
         """Save the state of this Experiment to SageMaker.
 
@@ -91,7 +93,9 @@ def load(cls, experiment_name, sagemaker_boto_client=None):
             sagemaker.experiments.experiment.Experiment: A SageMaker ``Experiment`` object
         """
         return cls._construct(
-            cls._boto_load_method, experiment_name=experiment_name, sagemaker_boto_client=sagemaker_boto_client,
+            cls._boto_load_method,
+            experiment_name=experiment_name,
+            sagemaker_boto_client=sagemaker_boto_client,
         )
 
     @classmethod
@@ -119,7 +123,12 @@ def create(cls, experiment_name=None, description=None, tags=None, sagemaker_bot
 
     @classmethod
     def list(
-        cls, created_before=None, created_after=None, sort_by=None, sort_order=None, sagemaker_boto_client=None,
+        cls,
+        created_before=None,
+        created_after=None,
+        sort_by=None,
+        sort_order=None,
+        sagemaker_boto_client=None,
     ):
         """
         List experiments. Returns experiments in the account matching the specified criteria.
@@ -152,7 +161,12 @@ def list(
 
     @classmethod
     def search(
-        cls, search_expression=None, sort_by=None, sort_order=None, max_results=None, sagemaker_boto_client=None,
+        cls,
+        search_expression=None,
+        sort_by=None,
+        sort_order=None,
+        max_results=None,
+        sagemaker_boto_client=None,
     ):
         """
         Search experiments. Returns SearchResults in the account matching the search criteria.
@@ -243,11 +257,11 @@ def delete_all(self, action):
                 "associated trials, trial components."
             )
 
-        delete_count = 0
+        delete_attempt_count = 0
         last_exception = None
         while True:
-            if delete_count == 3:
-                raise Exception("Fail to delete, please try again.") from last_exception
+            if delete_attempt_count == self.MAX_DELETE_ALL_ATTEMPTS:
+                raise Exception("Failed to delete, please try again.") from last_exception
             try:
                 for trial_summary in self.list_trials():
                     t = trial.Trial.load(
@@ -269,4 +283,4 @@ def delete_all(self, action):
             except Exception as ex:
                 last_exception = ex
             finally:
-                delete_count = delete_count + 1
+                delete_attempt_count = delete_attempt_count + 1

src/smexperiments/metrics.py

@@ -137,5 +137,6 @@ def __str__(self):
 
     def __repr__(self):
         return "{}({})".format(
-            type(self).__name__, ",".join(["{}={}".format(k, repr(v)) for k, v in vars(self).items()]),
+            type(self).__name__,
+            ",".join(["{}={}".format(k, repr(v)) for k, v in vars(self).items()]),
         )

src/smexperiments/tracker.py

@@ -24,7 +24,7 @@
 
 
 class Tracker(object):
-    """"A SageMaker Experiments Tracker.
+    """A SageMaker Experiments Tracker.
 
     Use a tracker object to record experiment information to a SageMaker trial component.
 

src/smexperiments/training_job.py

@@ -18,7 +18,12 @@
 class TrainingJob(_base_types.Record):
     @classmethod
     def search(
-        cls, search_expression=None, sort_by=None, sort_order=None, max_results=None, sagemaker_boto_client=None,
+        cls,
+        search_expression=None,
+        sort_by=None,
+        sort_order=None,
+        max_results=None,
+        sagemaker_boto_client=None,
     ):
         """
         Search Training Job. Returns SearchResults in the account matching the search criteria.

src/smexperiments/trial.py

@@ -13,6 +13,7 @@
 """Contains the Trial class."""
 
 from smexperiments import api_types, _base_types, trial_component, _utils, tracker
+import time
 
 
 class Trial(_base_types.Record):
@@ -58,6 +59,8 @@ class Trial(_base_types.Record):
     _boto_update_members = ["trial_name", "display_name"]
     _boto_delete_members = ["trial_name"]
 
+    MAX_DELETE_ALL_ATTEMPTS = 3
+
     @classmethod
     def _boto_ignore(cls):
         return super(Trial, cls)._boto_ignore() + ["CreatedBy"]
@@ -170,7 +173,12 @@ def list(
 
     @classmethod
     def search(
-        cls, search_expression=None, sort_by=None, sort_order=None, max_results=None, sagemaker_boto_client=None,
+        cls,
+        search_expression=None,
+        sort_by=None,
+        sort_order=None,
+        max_results=None,
+        sagemaker_boto_client=None,
     ):
         """
         Search experiments. Returns SearchResults in the account matching the search criteria.
@@ -270,3 +278,36 @@ def list_trial_components(
             next_token=next_token,
             sagemaker_boto_client=self.sagemaker_boto_client,
         )
+
+    def delete_all(self, action):
+        """
+        Force to delete the trial and associated trial components under.
+
+        Args:
+            action (str): pass in string '--force' to confirm delete the trial and all associated trial components.
+        """
+        if action != "--force":
+            raise ValueError(
+                "Must confirm with string '--force' in order to delete the trial and " "associated trial components."
+            )
+
+        delete_attempt_count = 0
+        last_exception = None
+        while True:
+            if delete_attempt_count == self.MAX_DELETE_ALL_ATTEMPTS:
+                raise Exception("Failed to delete, please try again.") from last_exception
+            try:
+                for trial_component_summary in self.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)
+                    # to prevent throttling
+                    time.sleep(1.2)
+                self.delete()
+                break
+            except Exception as ex:
+                last_exception = ex
+            finally:
+                delete_attempt_count = delete_attempt_count + 1

src/smexperiments/trial_component.py

@@ -228,7 +228,12 @@ def list(
 
     @classmethod
     def search(
-        cls, search_expression=None, sort_by=None, sort_order=None, max_results=None, sagemaker_boto_client=None,
+        cls,
+        search_expression=None,
+        sort_by=None,
+        sort_order=None,
+        max_results=None,
+        sagemaker_boto_client=None,
     ):
         """
         Search experiments. Returns SearchResults in the account matching the search criteria.

tests/conftest.py

@@ -123,22 +123,29 @@ def complex_experiment_obj(sagemaker_boto_client):
     trial_name3 = name()
 
     next_trial1 = trial.Trial.create(
-        trial_name=trial_name1, experiment_name=experiment_obj_name, sagemaker_boto_client=sagemaker_boto_client,
+        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_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_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,
+        trial_component_name=trial_component_name,
+        sagemaker_boto_client=sagemaker_boto_client,
     )
 
     # associate the trials with trial components
@@ -181,7 +188,9 @@ def trial_obj(sagemaker_boto_client, experiment_obj):
 @pytest.fixture
 def trial_component_obj(sagemaker_boto_client):
     trial_component_obj = trial_component.TrialComponent.create(
-        trial_component_name=name(), sagemaker_boto_client=sagemaker_boto_client, tags=TAGS,
+        trial_component_name=name(),
+        sagemaker_boto_client=sagemaker_boto_client,
+        tags=TAGS,
     )
     yield trial_component_obj
     time.sleep(0.5)
@@ -346,7 +355,10 @@ def training_job_name(sagemaker_boto_client, training_role_arn, docker_image, tr
                 "DataSource": {"S3DataSource": {"S3Uri": training_s3_uri, "S3DataType": "S3Prefix"}},
             }
         ],
-        AlgorithmSpecification={"TrainingImage": docker_image, "TrainingInputMode": "File",},
+        AlgorithmSpecification={
+            "TrainingImage": docker_image,
+            "TrainingInputMode": "File",
+        },
         RoleArn=training_role_arn,
         ResourceConfig={"InstanceType": "ml.m5.large", "InstanceCount": 1, "VolumeSizeInGB": 10},
         StoppingCondition={"MaxRuntimeInSeconds": 900},