fix: local mode fails when an S3 input channel points to a single file

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

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

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(

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."""

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

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