From 291abdb2498967f28f19253e48cc796cb2420770 Mon Sep 17 00:00:00 2001 From: Mufaddal Rohawala Date: Fri, 6 Nov 2026 20:37:52 +0000 Subject: [PATCH] fix(train): validate symlinks stay within source root when staging to S3 When ModelTrainer stages a local source_dir to S3 for the code channel, it copies/archives the directory (shutil.copytree, Session.upload_data), which follows nested symlinks: a symlink under source_dir that points elsewhere on the file system gets packaged as a regular file, pulling in unrelated content. Validate source_dir before staging the code channel in _create_training_job_args, reusing _validate_source_directory and _validate_symlinks_within_root from sagemaker.core.common_utils (the same guards ModelBuilder uses, added in #6545). The check is scoped to source_dir and skipped for LOCAL_CONTAINER (which mounts via EFS and does not upload). User InputData data channels are intentionally not covered, so shared-storage and local Hugging Face snapshot layouts that rely on symlinks keep working, consistent with #6545. Nested symlinks that resolve inside source_dir (relative or absolute) are still supported; S3 and .tar.gz inputs are unaffected. Adds unit tests driven through _create_training_job_args(): an in-root symlink is staged, an out-of-root symlink is rejected, and a data channel with an out-of-root symlink still uploads. Behavior change: a symlink under source_dir that resolves outside source_dir now raises ValueError during ModelTrainer.train(). To fix, copy the target into source_dir or point source_dir at a parent directory that contains both. --- .../src/sagemaker/train/model_trainer.py | 14 +++- .../tests/unit/train/test_model_trainer.py | 83 +++++++++++++++++++ 2 files changed, 96 insertions(+), 1 deletion(-) diff --git a/sagemaker-train/src/sagemaker/train/model_trainer.py b/sagemaker-train/src/sagemaker/train/model_trainer.py index 122bba9460..1e5a781875 100644 --- a/sagemaker-train/src/sagemaker/train/model_trainer.py +++ b/sagemaker-train/src/sagemaker/train/model_trainer.py @@ -33,7 +33,11 @@ from sagemaker.core.shapes import AlgorithmSpecification, ModelPackageConfig from sagemaker.core.utils.utils import serialize from sagemaker.core.apiutils._boto_functions import to_pascal_case -from sagemaker.core.common_utils import name_from_base +from sagemaker.core.common_utils import ( + name_from_base, + _validate_source_directory, + _validate_symlinks_within_root, +) from dataclasses import dataclass @@ -812,6 +816,14 @@ def _create_training_job_args( # If source code is provided, create a channel for the source code # The source code will be mounted at /opt/ml/input/data/code in the container if self.source_code.source_dir: + if self.training_mode != Mode.LOCAL_CONTAINER: + # The source directory is staged to S3 for the code channel by + # copying/archiving it, which follows nested symlinks. Validate the + # directory before staging so a nested symlink pointing outside it + # does not pull in unrelated content from elsewhere on the file + # system. S3 and .tar.gz inputs are unaffected. + _validate_source_directory(self.source_code.source_dir) + _validate_symlinks_within_root(self.source_code.source_dir, "source_dir") source_code_channel = self.create_input_data_channel( channel_name=SM_CODE, data_source=self.source_code.source_dir, diff --git a/sagemaker-train/tests/unit/train/test_model_trainer.py b/sagemaker-train/tests/unit/train/test_model_trainer.py index 2bb0e6246f..56e6a99fb5 100644 --- a/sagemaker-train/tests/unit/train/test_model_trainer.py +++ b/sagemaker-train/tests/unit/train/test_model_trainer.py @@ -2815,3 +2815,86 @@ def test_pipeline_and_interactive_emit_identical_tags(model_trainer): pipeline_tags = _pipeline_trainer([TAG])._create_training_job_args()["Tags"] assert pipeline_tags == _train_wire_tags(model_trainer) + + +def _make_source_dir(root, entry_script="train.py"): + """Create a minimal valid source_dir with an entry script.""" + with open(os.path.join(root, entry_script), "w") as f: + f.write("print('hi')\n") + return root + + +@patch("sagemaker.train.model_trainer.Session.upload_data") +@patch("sagemaker.train.model_trainer.Session.default_bucket") +def test_source_dir_with_in_root_symlink_is_staged( + mock_default_bucket, mock_upload_data, modules_session +): + """A source_dir whose nested symlinks resolve inside it is staged normally.""" + mock_default_bucket.return_value = DEFAULT_BUCKET + mock_upload_data.return_value = f"s3://{DEFAULT_BUCKET}/code" + with tempfile.TemporaryDirectory() as source_dir: + _make_source_dir(source_dir) + os.symlink(os.path.join(source_dir, "train.py"), os.path.join(source_dir, "entry_link.py")) + trainer = ModelTrainer( + training_image=DEFAULT_IMAGE, + role=DEFAULT_ROLE, + compute=DEFAULT_COMPUTE_CONFIG, + stopping_condition=DEFAULT_STOPPING_CONDITION, + output_data_config=DEFAULT_OUTPUT_DATA_CONFIG, + source_code=SourceCode(source_dir=source_dir, entry_script="train.py"), + sagemaker_session=modules_session, + ) + # Should not raise; the code channel is staged. + trainer._create_training_job_args(input_data_config=[]) + assert mock_upload_data.called + + +@patch("sagemaker.train.model_trainer.Session.upload_data") +@patch("sagemaker.train.model_trainer.Session.default_bucket") +def test_source_dir_with_out_of_root_symlink_is_rejected( + mock_default_bucket, mock_upload_data, modules_session +): + """A source_dir nested symlink resolving outside it is rejected before staging.""" + mock_default_bucket.return_value = DEFAULT_BUCKET + with tempfile.TemporaryDirectory() as outside_dir, tempfile.TemporaryDirectory() as source_dir: + _make_source_dir(source_dir) + nested = os.path.join(source_dir, "configs") + os.mkdir(nested) + outside_file = os.path.join(outside_dir, "outside.txt") + with open(outside_file, "w") as f: + f.write("data\n") + os.symlink(outside_file, os.path.join(nested, "linked.cfg")) + trainer = ModelTrainer( + training_image=DEFAULT_IMAGE, + role=DEFAULT_ROLE, + compute=DEFAULT_COMPUTE_CONFIG, + stopping_condition=DEFAULT_STOPPING_CONDITION, + output_data_config=DEFAULT_OUTPUT_DATA_CONFIG, + source_code=SourceCode(source_dir=source_dir, entry_script="train.py"), + sagemaker_session=modules_session, + ) + with pytest.raises(ValueError, match="resolves outside"): + trainer._create_training_job_args(input_data_config=[]) + mock_upload_data.assert_not_called() + + +@patch("sagemaker.train.model_trainer.Session.upload_data") +@patch("sagemaker.train.model_trainer.Session.default_bucket") +def test_data_channel_out_of_root_symlink_is_not_rejected( + mock_default_bucket, mock_upload_data, model_trainer +): + """The validation is scoped to source_dir; user data channels keep their symlinks. + + Shared-storage links (e.g. ``data/train.csv -> /mnt/datasets/train.csv``) and local + Hugging Face snapshot layouts must still upload, matching #6545's handling. + """ + mock_default_bucket.return_value = DEFAULT_BUCKET + mock_upload_data.return_value = f"s3://{DEFAULT_BUCKET}/data" + with tempfile.TemporaryDirectory() as outside_dir, tempfile.TemporaryDirectory() as data_dir: + shared = os.path.join(outside_dir, "train.csv") + with open(shared, "w") as f: + f.write("a,b\n1,2\n") + os.symlink(shared, os.path.join(data_dir, "train.csv")) + channel = model_trainer.create_input_data_channel("train", data_dir) + assert channel.channel_name == "train" + mock_upload_data.assert_called_once()