fix: safer destructor in ModelTrainer

Author: benieric

src/sagemaker/modules/train/model_trainer.py

@@ -205,7 +205,9 @@ class ModelTrainer(BaseModel):
             "LOCAL_CONTAINER" mode.
     """
 
-    model_config = ConfigDict(arbitrary_types_allowed=True, extra="forbid")
+    model_config = ConfigDict(
+        arbitrary_types_allowed=True, validate_assignment=True, extra="forbid"
+    )
 
     training_mode: Mode = Mode.SAGEMAKER_TRAINING_JOB
     sagemaker_session: Optional[Session] = None
@@ -363,9 +365,10 @@ def _populate_intelligent_defaults_from_model_trainer_space(self):
 
     def __del__(self):
         """Destructor method to clean up the temporary directory."""
-        # Clean up the temporary directory if it exists
-        if self._temp_recipe_train_dir is not None:
-            self._temp_recipe_train_dir.cleanup()
+        # Clean up the temporary directory if it exists and class was initialized
+        if hasattr(self, "__pydantic_fields_set__"):
+            if self._temp_recipe_train_dir is not None:
+                self._temp_recipe_train_dir.cleanup()
 
     def _validate_training_image_and_algorithm_name(
         self, training_image: Optional[str], algorithm_name: Optional[str]

tests/unit/sagemaker/modules/train/test_model_trainer.py

@@ -18,6 +18,7 @@
 import json
 import os
 import pytest
+from pydantic import ValidationError
 from unittest.mock import patch, MagicMock, ANY
 
 from sagemaker import image_uris
@@ -1068,3 +1069,27 @@ def test_safe_configs():
     # Test invalid type fails
     with pytest.raises(ValueError):
         SourceCode(entry_script=1)
+
+
+@patch("sagemaker.modules.train.model_trainer.TemporaryDirectory")
+def test_destructor_cleanup(mock_tmp_dir, modules_session):
+    
+    with pytest.raises(ValidationError):
+        model_trainer = ModelTrainer(
+            training_image=DEFAULT_IMAGE,
+            role=DEFAULT_ROLE,
+            sagemaker_session=modules_session,
+            compute="test"
+        )
+    mock_tmp_dir.cleanup.assert_not_called()
+
+    model_trainer = ModelTrainer(
+        training_image=DEFAULT_IMAGE,
+        role=DEFAULT_ROLE,
+        sagemaker_session=modules_session,
+        compute=DEFAULT_COMPUTE_CONFIG,
+    )
+    model_trainer._temp_recipe_train_dir = mock_tmp_dir
+    mock_tmp_dir.assert_not_called()
+    del model_trainer
+    mock_tmp_dir.cleanup.assert_called_once()