Add cleanup mechanism for MC dataset integ test

Author: mollyheamazon

sagemaker-train/src/sagemaker/ai_registry/dataset.py

@@ -389,7 +389,6 @@ def get_versions(self) -> List["DataSet"]:
 
         return datasets
 
-    @classmethod
     @classmethod
     @_telemetry_emitter(feature=Feature.MODEL_CUSTOMIZATION, func_name="DataSet.get_all")
     def get_all(cls, max_results: Optional[int] = None, sagemaker_session=None):

sagemaker-train/tests/integ/ai_registry/conftest.py

@@ -106,13 +106,8 @@ def cleanup_list():
     """Track resources for cleanup."""
     resources = []
     yield resources
-    for evaluator in resources:
+    for resource in resources:
         try:
-            from sagemaker.ai_registry.air_hub import AIRHub
-            AIRHub.delete_hub_content(
-                hub_content_type=evaluator.hub_content_type,
-                hub_content_name=evaluator.name,
-                hub_content_version=evaluator.version
-            )
+            resource.delete()
         except Exception:
             pass

sagemaker-train/tests/integ/ai_registry/test_dataset.py

@@ -129,9 +129,10 @@ def test_create_dataset_from_s3_nova_eval(self, unique_name, test_bucket, cleanu
         cleanup_list.append(dataset)
         assert dataset.name == unique_name
 
-    def test_get_dataset(self, unique_name, sample_jsonl_file):
+    def test_get_dataset(self, unique_name, sample_jsonl_file, cleanup_list):
         """Test retrieving dataset by name."""
         created = DataSet.create(name=unique_name, source=sample_jsonl_file, wait=False)
+        cleanup_list.append(created)
         retrieved = DataSet.get(unique_name)
         assert retrieved.name == created.name
         assert retrieved.arn == created.arn
@@ -141,16 +142,18 @@ def test_get_all_datasets(self):
         datasets = list(DataSet.get_all(max_results=5))
         assert isinstance(datasets, list)
 
-    def test_dataset_refresh(self, unique_name, sample_jsonl_file):
+    def test_dataset_refresh(self, unique_name, sample_jsonl_file, cleanup_list):
         """Test refreshing dataset status."""
         dataset = DataSet.create(name=unique_name, source=sample_jsonl_file, wait=False)
+        cleanup_list.append(dataset)
         dataset.refresh()
         time.sleep(3)
         assert dataset.status in [HubContentStatus.IMPORTING.value, HubContentStatus.AVAILABLE.value]
 
-    def test_dataset_get_versions(self, unique_name, sample_jsonl_file):
+    def test_dataset_get_versions(self, unique_name, sample_jsonl_file, cleanup_list):
         """Test getting dataset versions."""
         dataset = DataSet.create(name=unique_name, source=sample_jsonl_file, wait=False)
+        cleanup_list.append(dataset)
         versions = dataset.get_versions()
         assert len(versions) >= 1
         assert all(isinstance(v, DataSet) for v in versions)
@@ -178,7 +181,7 @@ def test_create_dataset_version(self, unique_name, sample_jsonl_file, cleanup_li
         """Test creating new dataset version."""
         dataset = DataSet.create(name=unique_name, source=sample_jsonl_file, wait=False)
         result = dataset.create_version(sample_jsonl_file)
-        cleanup_list.append(cleanup_list)
+        cleanup_list.append(dataset)
         assert result is True
 
     def test_dataset_validation_invalid_extension(self, unique_name):