Skip to content

fix(train): remove the local-mode code dir with the root-aware _rmtree - #6427

Open
AkshayShah03 wants to merge 1 commit into
aws:masterfrom
AkshayShah03:fix/local-container-temp-code-dir-cleanup
Open

AkshayShah03 wants to merge 1 commit into
aws:masterfrom
AkshayShah03:fix/local-container-temp-code-dir-cleanup

Conversation

@AkshayShah03

Copy link
Copy Markdown

Issue #, if available: Fixes #5542

Description of changes:

In Mode.LOCAL_CONTAINER, ModelTrainer copies the SDK drivers and the source code into a TemporaryDirectory under local_container_root (self._temp_code_dir) and mounts it into the training container. The container runs as root and writes into that directory, for example common/__pycache__. The plain self._temp_code_dir.cleanup() after training then fails.

This is the traceback the reporter posted on 3.21. dj_thcji is the temporary directory's random name:

PermissionError: [Errno 1] Operation not permitted:
'/home/sagemaker-user/amazon-sagemaker-from-idea-to-production/dj_thcji/common/__pycache__'

The 3.21 fix made _LocalContainer remove its own directories (input, shared, the host dirs) through _rmtree. On a PermissionError, _rmtree chmods the tree from a short-lived container and retries. The code directory owned by ModelTrainer wasn't covered.

This change removes the code directory with the same _rmtree(path, image, is_studio) after local training, before the regular TemporaryDirectory.cleanup(). That cleanup and __del__ are then no-ops on the already-removed path; I checked this on Python 3.10–3.13.

Testing:

  • New test_local_container_train_removes_temp_code_dir_with_root_owned_files. It checks that _rmtree gets the code directory and the container's image and Studio flag. It fails before this change.
  • sagemaker-train/tests/unit/train: no new failures. Master and this branch have the same 5 failures, which need AWS credentials or depend on test order.

By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice.

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 aws#5542

This branch is waiting to be deployed

1 waiting deployment
manual-approval — a8f9562f Waiting Oct 10, 2026 by AkshayShah03 via wait-for-approval #1979
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

ModelTrainer in Mode.LOCAL_CONTAINER throws an error when cleaning up folders upon exiting the container

1 participant