Remove alpine

Author: mollyheamazon

sagemaker-core/src/sagemaker/core/modules/local_core/local_container.py

@@ -57,24 +57,36 @@
 SM_STUDIO_LOCAL_MODE = "SM_STUDIO_LOCAL_MODE"
 
 
-def _rmtree(path):
+def _rmtree(path, image=None):
     """Remove a directory tree, handling root-owned files from Docker containers."""
     try:
         shutil.rmtree(path)
     except PermissionError:
         # Files created by Docker containers are owned by root.
         # Use a Docker container to remove them since os.chmod will also fail.
+        # Use the training image (already pulled) to avoid needing alpine.
+        if image is None:
+            logger.warning(
+                "Failed to clean up root-owned files in %s. "
+                "You may need to remove them manually with: sudo rm -rf %s",
+                path, path,
+            )
+            raise
         try:
             subprocess.run(
-                ["docker", "run", "--rm", "-v", f"{path}:/delete", "alpine", "rm", "-rf", "/delete"],
+                ["docker", "run", "--rm", "-v", f"{path}:/delete", image, "rm", "-rf", "/delete"],
                 check=True,
                 capture_output=True,
             )
             # The mount point directory itself may remain — clean it up
             if os.path.exists(path):
                 shutil.rmtree(path, ignore_errors=True)
         except Exception:
-            logger.warning("Failed to clean up root-owned files in %s. You may need to remove them manually with: sudo rm -rf %s", path, path)
+            logger.warning(
+                "Failed to clean up root-owned files in %s. "
+                "You may need to remove them manually with: sudo rm -rf %s",
+                path, path,
+            )
             raise
 
 
@@ -230,12 +242,12 @@ def train(
         # Print our Job Complete line
         logger.info("Local training job completed, output artifacts saved to %s", artifacts)
 
-        _rmtree(os.path.join(self.container_root, "input"))
-        _rmtree(os.path.join(self.container_root, "shared"))
+        _rmtree(os.path.join(self.container_root, "input"), self.image)
+        _rmtree(os.path.join(self.container_root, "shared"), self.image)
         for host in self.hosts:
-            _rmtree(os.path.join(self.container_root, host))
+            _rmtree(os.path.join(self.container_root, host), self.image)
         for folder in self._temporary_folders:
-            _rmtree(os.path.join(self.container_root, folder))
+            _rmtree(os.path.join(self.container_root, folder), self.image)
         return artifacts
 
     def retrieve_artifacts(

sagemaker-train/src/sagemaker/train/local/local_container.py

@@ -65,24 +65,36 @@
 SM_STUDIO_LOCAL_MODE = "SM_STUDIO_LOCAL_MODE"
 
 
-def _rmtree(path):
+def _rmtree(path, image=None):
     """Remove a directory tree, handling root-owned files from Docker containers."""
     try:
         shutil.rmtree(path)
     except PermissionError:
         # Files created by Docker containers are owned by root.
         # Use a Docker container to remove them since os.chmod will also fail.
+        # Use the training image (already pulled) to avoid needing alpine.
+        if image is None:
+            logger.warning(
+                "Failed to clean up root-owned files in %s. "
+                "You may need to remove them manually with: sudo rm -rf %s",
+                path, path,
+            )
+            raise
         try:
             subprocess.run(
-                ["docker", "run", "--rm", "-v", f"{path}:/delete", "alpine", "rm", "-rf", "/delete"],
+                ["docker", "run", "--rm", "-v", f"{path}:/delete", image, "rm", "-rf", "/delete"],
                 check=True,
                 capture_output=True,
             )
             # The mount point directory itself may remain — clean it up
             if os.path.exists(path):
                 shutil.rmtree(path, ignore_errors=True)
         except Exception:
-            logger.warning("Failed to clean up root-owned files in %s. You may need to remove them manually with: sudo rm -rf %s", path, path)
+            logger.warning(
+                "Failed to clean up root-owned files in %s. "
+                "You may need to remove them manually with: sudo rm -rf %s",
+                path, path,
+            )
             raise
 
 
@@ -238,12 +250,12 @@ def train(
         # Print our Job Complete line
         logger.info("Local training job completed, output artifacts saved to %s", artifacts)
 
-        _rmtree(os.path.join(self.container_root, "input"))
-        _rmtree(os.path.join(self.container_root, "shared"))
+        _rmtree(os.path.join(self.container_root, "input"), self.image)
+        _rmtree(os.path.join(self.container_root, "shared"), self.image)
         for host in self.hosts:
-            _rmtree(os.path.join(self.container_root, host))
+            _rmtree(os.path.join(self.container_root, host), self.image)
         for folder in self._temporary_folders:
-            _rmtree(os.path.join(self.container_root, folder))
+            _rmtree(os.path.join(self.container_root, folder), self.image)
         return artifacts
 
     def retrieve_artifacts(

sagemaker-train/tests/unit/train/local/test_local_container.py

@@ -15,27 +15,29 @@
 
 from sagemaker.train.local.local_container import _rmtree
 
+IMAGE = "763104351884.dkr.ecr.us-east-1.amazonaws.com/pytorch-training:2.1-cpu-py310"
+
 
 class TestRmtree:
     """Test cases for _rmtree function."""
 
     @patch("sagemaker.train.local.local_container.shutil.rmtree")
     def test_rmtree_success(self, mock_rmtree):
         """Normal case — shutil.rmtree succeeds."""
-        _rmtree("/tmp/test")
+        _rmtree("/tmp/test", IMAGE)
         mock_rmtree.assert_called_once_with("/tmp/test")
 
     @patch("sagemaker.train.local.local_container.shutil.rmtree")
     @patch("sagemaker.train.local.local_container.subprocess.run")
     @patch("sagemaker.train.local.local_container.os.path.exists", return_value=False)
     def test_rmtree_permission_error_docker_fallback(self, mock_exists, mock_run, mock_rmtree):
-        """PermissionError triggers docker fallback to remove root-owned files."""
+        """PermissionError triggers docker fallback using the training image."""
         mock_rmtree.side_effect = PermissionError("Permission denied")
 
-        _rmtree("/tmp/test")
+        _rmtree("/tmp/test", IMAGE)
 
         mock_run.assert_called_once_with(
-            ["docker", "run", "--rm", "-v", "/tmp/test:/delete", "alpine", "rm", "-rf", "/delete"],
+            ["docker", "run", "--rm", "-v", "/tmp/test:/delete", IMAGE, "rm", "-rf", "/delete"],
             check=True,
             capture_output=True,
         )
@@ -47,7 +49,7 @@ def test_rmtree_cleans_up_mount_point(self, mock_exists, mock_run, mock_rmtree):
         """After docker cleanup, remaining mount point directory is removed."""
         mock_rmtree.side_effect = [PermissionError("Permission denied"), None]
 
-        _rmtree("/tmp/test")
+        _rmtree("/tmp/test", IMAGE)
 
         assert mock_rmtree.call_count == 2
         mock_rmtree.assert_has_calls([
@@ -60,7 +62,15 @@ def test_rmtree_cleans_up_mount_point(self, mock_exists, mock_run, mock_rmtree):
     def test_rmtree_docker_fallback_fails_raises(self, mock_run, mock_rmtree):
         """If docker fallback also fails, the exception propagates."""
         mock_rmtree.side_effect = PermissionError("Permission denied")
-        mock_run.side_effect = Exception("docker not available")
+        mock_run.side_effect = Exception("docker failed")
+
+        with pytest.raises(Exception, match="docker failed"):
+            _rmtree("/tmp/test", IMAGE)
+
+    @patch("sagemaker.train.local.local_container.shutil.rmtree")
+    def test_rmtree_no_image_raises(self, mock_rmtree):
+        """PermissionError without image raises immediately."""
+        mock_rmtree.side_effect = PermissionError("Permission denied")
 
-        with pytest.raises(Exception, match="docker not available"):
+        with pytest.raises(PermissionError):
             _rmtree("/tmp/test")