diff --git a/sagemaker-core/src/sagemaker/core/common_utils.py b/sagemaker-core/src/sagemaker/core/common_utils.py index e8704ef600..5b8b64d4b2 100644 --- a/sagemaker-core/src/sagemaker/core/common_utils.py +++ b/sagemaker-core/src/sagemaker/core/common_utils.py @@ -442,6 +442,7 @@ def download_folder(bucket_name, prefix, target, sagemaker_session): # Do this first, in case the object has broader permissions than the bucket. if not prefix.endswith("/"): try: + os.makedirs(target, exist_ok=True) file_destination = os.path.join(target, os.path.basename(prefix)) s3.Object(bucket_name, prefix).download_file(file_destination, ExtraArgs=extra_args) return diff --git a/sagemaker-core/src/sagemaker/core/modules/local_core/local_container.py b/sagemaker-core/src/sagemaker/core/modules/local_core/local_container.py index 7df9409b86..5d80a5ca1e 100644 --- a/sagemaker-core/src/sagemaker/core/modules/local_core/local_container.py +++ b/sagemaker-core/src/sagemaker/core/modules/local_core/local_container.py @@ -19,7 +19,7 @@ import re import shutil import subprocess -from tempfile import TemporaryDirectory +from tempfile import mkdtemp from typing import Any, Dict, List, Optional from pydantic import BaseModel, ConfigDict @@ -586,7 +586,8 @@ def _get_data_source_local_path(self, data_source: DataSource): if data_source.s3_data_source != Unassigned(): uri = data_source.s3_data_source.s3_uri parsed_uri = urlparse(uri) - local_dir = TemporaryDirectory(prefix=os.path.join(self.container_root + "/")).name + os.makedirs(self.container_root, exist_ok=True) + local_dir = mkdtemp(dir=self.container_root) self._temporary_folders.append(local_dir) download_folder(parsed_uri.netloc, parsed_uri.path, local_dir, self.sagemaker_session) return local_dir diff --git a/sagemaker-core/tests/unit/modules/local_core/test_local_container.py b/sagemaker-core/tests/unit/modules/local_core/test_local_container.py index cf5f4c847e..8528710adc 100644 --- a/sagemaker-core/tests/unit/modules/local_core/test_local_container.py +++ b/sagemaker-core/tests/unit/modules/local_core/test_local_container.py @@ -496,8 +496,8 @@ def test_ecr_login_if_needed_with_non_ecr_image(self, mock_session, basic_channe assert result is False @patch("sagemaker.core.modules.local_core.local_container.download_folder") - @patch("sagemaker.core.modules.local_core.local_container.TemporaryDirectory") - def test_get_data_source_local_path_s3(self, mock_temp_dir, mock_download, mock_session): + @patch("sagemaker.core.modules.local_core.local_container.mkdtemp") + def test_get_data_source_local_path_s3(self, mock_mkdtemp, mock_download, mock_session): """Test _get_data_source_local_path with S3 data source""" data_source = DataSource( s3_data_source=S3DataSource( @@ -522,13 +522,62 @@ def test_get_data_source_local_path_s3(self, mock_temp_dir, mock_download, mock_ sagemaker_session=mock_session, ) - mock_temp_dir.return_value.name = "/tmp/temp123" + mock_mkdtemp.return_value = "/tmp/temp123" result = container._get_data_source_local_path(data_source) assert result == "/tmp/temp123" mock_download.assert_called_once() + @patch("sagemaker.core.modules.local_core.local_container.download_folder") + def test_get_data_source_local_path_s3_single_file_dir_exists( + self, mock_download, mock_session, tmp_path + ): + """Regression test for #5033: channel dir must exist when download_folder runs. + + On unfixed code the TemporaryDirectory object is garbage-collected + immediately, deleting the directory before download_folder is called. + """ + data_source = DataSource( + s3_data_source=S3DataSource( + s3_uri="s3://bucket/data/train.csv", + s3_data_type="S3Prefix", + s3_data_distribution_type="FullyReplicated", + ) + ) + channel = Channel(channel_name="training", data_source=data_source) + + container = _LocalContainer( + training_job_name="test-job", + instance_type="local", + instance_count=1, + image="test-image:latest", + container_root=str(tmp_path), + input_data_config=[channel], + environment={}, + hyper_parameters={}, + container_entrypoint=[], + container_arguments=[], + sagemaker_session=mock_session, + ) + + def download_side_effect(bucket, prefix, target, session): + # The target directory MUST exist at this point + assert os.path.isdir(target), f"Target directory {target} does not exist" + # Simulate downloading a file into the target + filepath = os.path.join(target, os.path.basename(prefix)) + with open(filepath, "w") as f: + f.write("data") + + mock_download.side_effect = download_side_effect + + result = container._get_data_source_local_path(data_source) + + assert os.path.isdir(result) + assert str(tmp_path) == str(os.path.dirname(result)) + assert os.path.isfile(os.path.join(result, "train.csv")) + assert result in container._temporary_folders + def test_get_data_source_local_path_local(self, mock_session): """Test _get_data_source_local_path with local file system""" data_source = DataSource( diff --git a/sagemaker-core/tests/unit/test_common_utils.py b/sagemaker-core/tests/unit/test_common_utils.py index 58863335bb..3a532e35b3 100644 --- a/sagemaker-core/tests/unit/test_common_utils.py +++ b/sagemaker-core/tests/unit/test_common_utils.py @@ -961,6 +961,35 @@ def test_download_folder_with_traversal_error(self): with pytest.raises(ValueError): download_folder("bucket", "/../prefix/", tmpdir, mock_session) + def test_download_folder_single_file_creates_missing_target(self): + """Regression test for #5033: download_folder must create the target dir. + + When downloading a single S3 object, the target directory may not + exist yet (e.g. because the TemporaryDirectory that held it was + garbage-collected). download_folder should create it. + """ + from sagemaker.core.common_utils import download_folder + + mock_session = Mock() + mock_s3 = Mock() + mock_obj = Mock() + mock_session.s3_resource = mock_s3 + mock_s3.Object.return_value = mock_obj + + with tempfile.TemporaryDirectory() as tmpdir: + target = os.path.join(tmpdir, "missing", "dir") + + def fake_download(dest, ExtraArgs=None): + # Simulate what boto3 does: write to the given path + with open(dest, "w") as f: + f.write("data") + + mock_obj.download_file.side_effect = fake_download + + download_folder("bucket", "data/train.csv", target, mock_session) + + assert os.path.isfile(os.path.join(target, "train.csv")) + class TestDownloadFilesUnderPrefixPathTraversal: """Test _download_files_under_prefix blocks path traversal attacks.""" diff --git a/sagemaker-train/src/sagemaker/train/local/local_container.py b/sagemaker-train/src/sagemaker/train/local/local_container.py index ddc21a09d8..daa60355a9 100644 --- a/sagemaker-train/src/sagemaker/train/local/local_container.py +++ b/sagemaker-train/src/sagemaker/train/local/local_container.py @@ -20,7 +20,7 @@ import re import shutil import subprocess -from tempfile import TemporaryDirectory +from tempfile import mkdtemp from typing import Any, Dict, List, Optional from pydantic import BaseModel, ConfigDict @@ -594,7 +594,8 @@ def _get_data_source_local_path(self, data_source: DataSource): if data_source.s3_data_source != Unassigned(): uri = data_source.s3_data_source.s3_uri parsed_uri = urlparse(uri) - local_dir = TemporaryDirectory(prefix=os.path.join(self.container_root + "/")).name + os.makedirs(self.container_root, exist_ok=True) + local_dir = mkdtemp(dir=self.container_root) self._temporary_folders.append(local_dir) download_folder(parsed_uri.netloc, parsed_uri.path, local_dir, self.sagemaker_session) return local_dir diff --git a/sagemaker-train/tests/unit/train/local/test_local_container.py b/sagemaker-train/tests/unit/train/local/test_local_container.py index be3ed26cfa..77343a6003 100644 --- a/sagemaker-train/tests/unit/train/local/test_local_container.py +++ b/sagemaker-train/tests/unit/train/local/test_local_container.py @@ -223,3 +223,58 @@ def test_get_compose_cmd_prefix_standalone_fallback( mock_which.return_value = "/usr/local/bin/docker-compose" result = container._get_compose_cmd_prefix() assert result == ["docker-compose"] + + +class TestGetDataSourceLocalPathS3: + """Regression tests for #5033: channel temp dirs must survive until download_folder runs.""" + + @patch("sagemaker.train.local.local_container.download_folder") + def test_get_data_source_local_path_s3_single_file_dir_exists( + self, mock_download, tmp_path, _basic_channel + ): + """The channel directory must exist when download_folder is called. + + On unfixed code the TemporaryDirectory object is garbage-collected + immediately, deleting the directory before download_folder runs. + """ + import os + + data_source = DataSource( + s3_data_source=S3DataSource( + s3_uri="s3://bucket/data/train.csv", + s3_data_type="S3Prefix", + s3_data_distribution_type="FullyReplicated", + ) + ) + channel = Channel(channel_name="training", data_source=data_source) + + container = _LocalContainer( + training_job_name="test-job", + instance_type="local", + instance_count=1, + image=IMAGE, + container_root=str(tmp_path), + input_data_config=[channel], + environment={}, + hyper_parameters={}, + container_entrypoint=[], + container_arguments=[], + sagemaker_session=None, + ) + + def download_side_effect(bucket, prefix, target, session): + # The target directory MUST exist at this point + assert os.path.isdir(target), f"Target directory {target} does not exist" + # Simulate downloading a file into the target + filepath = os.path.join(target, os.path.basename(prefix)) + with open(filepath, "w") as f: + f.write("data") + + mock_download.side_effect = download_side_effect + + result = container._get_data_source_local_path(data_source) + + assert os.path.isdir(result) + assert str(tmp_path) == str(os.path.dirname(result)) + assert os.path.isfile(os.path.join(result, "train.csv")) + assert result in container._temporary_folders