Use network flag

Author: mollyheamazon

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

@@ -57,14 +57,13 @@
 SM_STUDIO_LOCAL_MODE = "SM_STUDIO_LOCAL_MODE"
 
 
-def _rmtree(path, image=None):
+def _rmtree(path, image=None, is_studio=False):
     """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. "
@@ -73,11 +72,11 @@ def _rmtree(path, image=None):
             )
             raise
         try:
-            subprocess.run(
-                ["docker", "run", "--rm", "-v", f"{path}:/delete", image, "rm", "-rf", "/delete"],
-                check=True,
-                capture_output=True,
-            )
+            cmd = ["docker", "run", "--rm"]
+            if is_studio:
+                cmd += ["--network", "sagemaker"]
+            cmd += ["-v", f"{path}:/delete", image, "rm", "-rf", "/delete"]
+            subprocess.run(cmd, 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)
@@ -242,12 +241,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"), self.image)
-        _rmtree(os.path.join(self.container_root, "shared"), self.image)
+        _rmtree(os.path.join(self.container_root, "input"), self.image, self.is_studio)
+        _rmtree(os.path.join(self.container_root, "shared"), self.image, self.is_studio)
         for host in self.hosts:
-            _rmtree(os.path.join(self.container_root, host), self.image)
+            _rmtree(os.path.join(self.container_root, host), self.image, self.is_studio)
         for folder in self._temporary_folders:
-            _rmtree(os.path.join(self.container_root, folder), self.image)
+            _rmtree(os.path.join(self.container_root, folder), self.image, self.is_studio)
         return artifacts
 
     def retrieve_artifacts(

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

@@ -65,14 +65,13 @@
 SM_STUDIO_LOCAL_MODE = "SM_STUDIO_LOCAL_MODE"
 
 
-def _rmtree(path, image=None):
+def _rmtree(path, image=None, is_studio=False):
     """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. "
@@ -81,11 +80,11 @@ def _rmtree(path, image=None):
             )
             raise
         try:
-            subprocess.run(
-                ["docker", "run", "--rm", "-v", f"{path}:/delete", image, "rm", "-rf", "/delete"],
-                check=True,
-                capture_output=True,
-            )
+            cmd = ["docker", "run", "--rm"]
+            if is_studio:
+                cmd += ["--network", "sagemaker"]
+            cmd += ["-v", f"{path}:/delete", image, "rm", "-rf", "/delete"]
+            subprocess.run(cmd, 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)
@@ -250,12 +249,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"), self.image)
-        _rmtree(os.path.join(self.container_root, "shared"), self.image)
+        _rmtree(os.path.join(self.container_root, "input"), self.image, self.is_studio)
+        _rmtree(os.path.join(self.container_root, "shared"), self.image, self.is_studio)
         for host in self.hosts:
-            _rmtree(os.path.join(self.container_root, host), self.image)
+            _rmtree(os.path.join(self.container_root, host), self.image, self.is_studio)
         for folder in self._temporary_folders:
-            _rmtree(os.path.join(self.container_root, folder), self.image)
+            _rmtree(os.path.join(self.container_root, folder), self.image, self.is_studio)
         return artifacts
 
     def retrieve_artifacts(

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

@@ -42,6 +42,25 @@ def test_rmtree_permission_error_docker_fallback(self, mock_exists, mock_run, mo
             capture_output=True,
         )
 
+    @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_studio_adds_network(self, mock_exists, mock_run, mock_rmtree):
+        """In Studio, docker run includes --network sagemaker."""
+        mock_rmtree.side_effect = PermissionError("Permission denied")
+
+        _rmtree("/tmp/test", IMAGE, is_studio=True)
+
+        mock_run.assert_called_once_with(
+            [
+                "docker", "run", "--rm",
+                "--network", "sagemaker",
+                "-v", "/tmp/test:/delete", IMAGE, "rm", "-rf", "/delete",
+            ],
+            check=True,
+            capture_output=True,
+        )
+
     @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=True)