From a8f9562feea42d0392076cafb1921c28a9b07487 Mon Sep 17 00:00:00 2001 From: Akshay Shah Date: Sat, 10 Oct 2026 10:15:01 -0500 Subject: [PATCH] fix(train): remove the local-mode code dir with the root-aware _rmtree In Mode.LOCAL_CONTAINER, ModelTrainer copies the SDK drivers (and source code) into a TemporaryDirectory under local_container_root and mounts it into the training container. The container runs as root and writes files there (e.g. common/__pycache__), so the TemporaryDirectory.cleanup() after training failed with PermissionError: Operation not permitted. _LocalContainer already removes its own directories with _rmtree, which falls back to chmod-ing the tree from a short-lived container; use it for the code dir as well before the regular cleanup. Fixes #5542 --- .../src/sagemaker/train/model_trainer.py | 7 ++++- .../tests/unit/train/test_model_trainer.py | 26 +++++++++++++++++++ 2 files changed, 32 insertions(+), 1 deletion(-) diff --git a/sagemaker-train/src/sagemaker/train/model_trainer.py b/sagemaker-train/src/sagemaker/train/model_trainer.py index 8816ea68f9..7c8c2d32d5 100644 --- a/sagemaker-train/src/sagemaker/train/model_trainer.py +++ b/sagemaker-train/src/sagemaker/train/model_trainer.py @@ -122,7 +122,7 @@ from sagemaker.core.workflow.pipeline_context import PipelineSession, runnable_by_pipeline from sagemaker.core.helper.pipeline_variable import StrPipeVar -from sagemaker.train.local.local_container import _LocalContainer +from sagemaker.train.local.local_container import _LocalContainer, _rmtree class Mode(Enum): @@ -1000,6 +1000,11 @@ def train( environment=training_request["environment"], ) local_container.train(wait) + if self._temp_code_dir is not None: + # The code dir is mounted into the container, which can leave root-owned + # files (e.g. __pycache__) behind that TemporaryDirectory.cleanup() cannot + # remove, so use the same root-aware removal as the container root. + _rmtree(self._temp_code_dir.name, local_container.image, local_container.is_studio) if self._temp_code_dir is not None: self._temp_code_dir.cleanup() diff --git a/sagemaker-train/tests/unit/train/test_model_trainer.py b/sagemaker-train/tests/unit/train/test_model_trainer.py index f990b7ab88..73f771285d 100644 --- a/sagemaker-train/tests/unit/train/test_model_trainer.py +++ b/sagemaker-train/tests/unit/train/test_model_trainer.py @@ -626,6 +626,32 @@ def test_create_input_data_channel_custom_input_s3_key_prefix( assert f"{DEFAULT_BASE_NAME}/input/code" not in mock_upload_data.call_args.kwargs["key_prefix"] +@patch("sagemaker.train.model_trainer._rmtree") +@patch("sagemaker.train.model_trainer._LocalContainer") +def test_local_container_train_removes_temp_code_dir_with_root_owned_files( + mock_local_container, mock_rmtree, tmp_path +): + """The sm_drivers temp dir is mounted into the container, which can leave root-owned + files (e.g. __pycache__) in it, so it must be removed with the root-aware _rmtree.""" + mock_local_container.return_value.image = DEFAULT_IMAGE + mock_local_container.return_value.is_studio = False + + trainer = ModelTrainer( + training_mode=Mode.LOCAL_CONTAINER, + training_image=DEFAULT_IMAGE, + role=DEFAULT_ROLE, + source_code=DEFAULT_SOURCE_CODE, + compute=Compute(instance_type="local", instance_count=1), + local_container_root=str(tmp_path), + ) + + trainer.train(wait=True) + + temp_code_dir = trainer._temp_code_dir.name + mock_local_container.return_value.train.assert_called_once_with(True) + mock_rmtree.assert_called_once_with(temp_code_dir, DEFAULT_IMAGE, False) + + @patch("sagemaker.train.model_trainer.ModelTrainer._resolve_staging_bucket") @patch("sagemaker.train.model_trainer.Session.upload_data") @patch("sagemaker.train.model_trainer.Session.default_bucket")