diff --git a/sagemaker-core/src/sagemaker/core/modules/configs.py b/sagemaker-core/src/sagemaker/core/modules/configs.py index 09e2671894..a23c5e14c4 100644 --- a/sagemaker-core/src/sagemaker/core/modules/configs.py +++ b/sagemaker-core/src/sagemaker/core/modules/configs.py @@ -91,7 +91,10 @@ class SourceCode(BaseConfig): Parameters: source_dir (Optional[str]): - The local directory containing the source code to be used in the training job container. + The local directory, S3 URI, or path to a tar.gz file stored locally or in S3 + that contains the source code to be used in the training job container. + When an S3 URI is provided, the source code is used directly from S3 + without local packaging. requirements (Optional[str]): The path within ``source_dir`` to a ``requirements.txt`` file. If specified, the listed requirements will be installed in the training job container. diff --git a/sagemaker-core/src/sagemaker/core/processing.py b/sagemaker-core/src/sagemaker/core/processing.py index 7ceeb9c19f..3f48383a47 100644 --- a/sagemaker-core/src/sagemaker/core/processing.py +++ b/sagemaker-core/src/sagemaker/core/processing.py @@ -1129,18 +1129,66 @@ def _s3_code_prefix(self): self.sagemaker_session.default_bucket_prefix or "", ) + @staticmethod + def _is_s3_uri(path: Optional[str]) -> bool: + """Check whether the given path is an S3 URI.""" + return bool(path) and path.lower().startswith("s3://") + + def _resolve_s3_source_dir(self, source_dir: str) -> str: + """Resolve an S3 source_dir to a sourcedir.tar.gz URI. + + If the URI already points to a .tar.gz file, return it unchanged. + Otherwise treat it as an S3 prefix and append sourcedir.tar.gz. + """ + if source_dir.lower().endswith(".tar.gz"): + return source_dir + return source_dir.rstrip("/") + "/sourcedir.tar.gz" + + def _resolve_helper_scripts_prefix(self, job_name: str) -> str: + """Return an S3 prefix for uploading helper scripts (runproc.sh, install_requirements.py).""" + return s3.s3_path_join( + self._s3_code_prefix(), + job_name, + "source", + ) + def _package_code( self, entry_point, source_dir, + dependencies, requirements, job_name, kms_key, ): - """Package and upload code to S3.""" + """Package and upload code to S3. + + If source_dir is an S3 URI, it is used directly as the code payload + (no local packaging or upload is performed). The S3 URI should point to + either a tar.gz archive or an S3 prefix containing the source code. + + Args: + entry_point (str): Path to the entry point script. + source_dir (str): Local directory, S3 URI, or None. + dependencies (list[str]): Additional local directories to include + in the tar.gz bundle (default: None). Not supported with S3 source_dir. + requirements (str): Path to requirements.txt relative to source_dir. + job_name (str): Processing job name (used in S3 key). + kms_key (str): KMS key for S3 upload encryption. + """ import tarfile import tempfile + # S3 source_dir: use it directly without local packaging. + # This restores v2 behavior where S3 paths with tar.gz archives were supported. + if self._is_s3_uri(source_dir): + if dependencies: + raise ValueError( + "dependencies is not supported when source_dir is an S3 URI. " + "Bundle dependencies into the S3 tar.gz archive instead." + ) + return self._resolve_s3_source_dir(source_dir) + # If source_dir is not provided, use the directory containing entry_point if source_dir is None: if os.path.isabs(entry_point): @@ -1155,15 +1203,22 @@ def _package_code( if not os.path.exists(source_dir): raise ValueError(f"source_dir does not exist: {source_dir}") - # Create tar.gz with source_dir contents + # Create tar.gz with source_dir contents + dependencies with tempfile.NamedTemporaryFile(suffix=".tar.gz", delete=False) as tmp: with tarfile.open(tmp.name, "w:gz") as tar: - # Add all files from source_dir to the root of the tar + # Add all files from source_dir for item in os.listdir(source_dir): item_path = os.path.join(source_dir, item) tar.add(item_path, arcname=item) - # Upload to S3 + # Add dependency directories at the root of the archive + for dep_path in (dependencies or []): + if not os.path.isabs(dep_path): + dep_path = os.path.abspath(dep_path) + if not os.path.exists(dep_path): + raise ValueError(f"Dependency path does not exist: {dep_path}") + tar.add(dep_path, arcname=os.path.basename(dep_path)) + s3_uri = s3.s3_path_join( self._s3_code_prefix(), job_name, @@ -1171,7 +1226,6 @@ def _package_code( "sourcedir.tar.gz", ) - # Upload the tar file directly to S3 s3.S3Uploader.upload_string_as_file_body( body=open(tmp.name, "rb").read(), desired_s3_uri=s3_uri, @@ -1188,6 +1242,7 @@ def run( self, code: str, source_dir: Optional[str] = None, + dependencies: Optional[List[str]] = None, requirements: Optional[str] = None, inputs: Optional[List[ProcessingInput]] = None, outputs: Optional[List["ProcessingOutput"]] = None, @@ -1206,7 +1261,13 @@ def run( framework script to run. source_dir (str): Path (absolute, relative or an S3 URI) to a directory with any other processing source code dependencies aside from the entry - point file (default: None). + point file (default: None). If ``source_dir`` is an S3 URI, it must + point to a tar.gz file named ``sourcedir.tar.gz``. + dependencies (list[str]): A list of paths to directories (absolute or + relative) with any additional libraries that will be exported to the + container (default: None). The library folders are copied into the + same tar.gz bundle as source_dir. Not supported when source_dir is + an S3 URI. requirements (str): Path to a requirements.txt file relative to source_dir (default: None). inputs (list[:class:`~sagemaker.processing.ProcessingInput`]): Input files for @@ -1235,6 +1296,7 @@ def run( s3_runproc_sh, inputs, job_name = self._pack_and_upload_code( code, source_dir, + dependencies, requirements, job_name, inputs, @@ -1259,6 +1321,7 @@ def _pack_and_upload_code( self, code, source_dir, + dependencies, requirements, job_name, inputs, @@ -1276,6 +1339,7 @@ def _pack_and_upload_code( s3_payload = self._package_code( entry_point=code, source_dir=source_dir, + dependencies=dependencies, requirements=requirements, job_name=job_name, kms_key=kms_key, @@ -1283,12 +1347,20 @@ def _pack_and_upload_code( inputs = self._patch_inputs_with_payload(inputs, s3_payload) - entrypoint_s3_uri = s3_payload.replace("sourcedir.tar.gz", "runproc.sh") + # Determine where to upload helper scripts. + # When source_dir is S3, s3_payload points to the user's existing location, + # so helpers go to a separate managed prefix. + if self._is_s3_uri(source_dir): + helper_prefix = self._resolve_helper_scripts_prefix(job_name) + entrypoint_s3_uri = s3.s3_path_join(helper_prefix, "runproc.sh") + install_req_s3_uri = s3.s3_path_join(helper_prefix, "install_requirements.py") + else: + entrypoint_s3_uri = s3_payload.replace("sourcedir.tar.gz", "runproc.sh") + install_req_s3_uri = s3_payload.replace("sourcedir.tar.gz", "install_requirements.py") - # Upload the CodeArtifact-aware install_requirements script alongside the source code + # Upload install_requirements helper import sagemaker.core.utils.install_requirements as _ir_mod - install_req_s3_uri = s3_payload.replace("sourcedir.tar.gz", "install_requirements.py") evaluated_kms_key = kms_key if kms_key else self.output_kms_key s3.S3Uploader.upload_string_as_file_body( body=open(_ir_mod.__file__, "r").read(), @@ -1298,7 +1370,6 @@ def _pack_and_upload_code( ) script = os.path.basename(code) - evaluated_kms_key = kms_key if kms_key else self.output_kms_key s3_runproc_sh = self._create_and_upload_runproc( script, evaluated_kms_key, entrypoint_s3_uri, entry_point, source_dir ) @@ -1428,6 +1499,44 @@ def _generate_custom_framework_script( Returns: str: The generated script content """ + # When source_dir is an S3 URI, we cannot read the entry_point file locally. + # Instead, generate a script that executes the entry_point from the extracted + # source bundle on the container. + if self._is_s3_uri(source_dir): + return dedent( + """\ + #!/bin/bash + + # Exit on any error. SageMaker uses error code to mark failed job. + set -e + + cd /opt/ml/processing/input/code/ + + # Extract source code + if [ -f sourcedir.tar.gz ]; then + tar -xzf sourcedir.tar.gz + else + echo "ERROR: sourcedir.tar.gz not found!" + exit 1 + fi + + if [[ -f 'requirements.txt' ]]; then + pip uninstall --yes typing + python3 /opt/ml/processing/input/code/install_requirements.py requirements.txt + fi + + # Execute custom entrypoint + chmod +x {entry_point} + ./{entry_point} + + {entry_point_command} {user_script} "$@" + """ + ).format( + entry_point=entry_point, + entry_point_command=" ".join(self.command), + user_script=user_script, + ) + # Resolve the full path to the entry_point file if source_dir and not os.path.isabs(entry_point): full_entry_point_path = os.path.join(source_dir, entry_point) diff --git a/sagemaker-core/tests/unit/test_processing.py b/sagemaker-core/tests/unit/test_processing.py index f5349faea0..dd619808ed 100644 --- a/sagemaker-core/tests/unit/test_processing.py +++ b/sagemaker-core/tests/unit/test_processing.py @@ -1763,3 +1763,236 @@ def test_latest_job_updated_after_run(self, mock_session): assert processor.latest_job == mock_job2 assert len(processor.jobs) == 2 + + +class TestFrameworkProcessorS3SourceDir: + """Tests for FrameworkProcessor S3 source_dir support (v2 parity regression fix).""" + + def _make_processor(self, mock_session, **kwargs): + return FrameworkProcessor( + role="arn:aws:iam::123456789012:role/SageMakerRole", + image_uri="test-image:latest", + instance_count=1, + instance_type="ml.m5.xlarge", + sagemaker_session=mock_session, + **kwargs, + ) + + # --- _is_s3_uri helper --- + + def test_is_s3_uri_with_s3_path(self, mock_session): + processor = self._make_processor(mock_session) + assert processor._is_s3_uri("s3://bucket/path") is True + + def test_is_s3_uri_with_uppercase(self, mock_session): + processor = self._make_processor(mock_session) + assert processor._is_s3_uri("S3://bucket/path") is True + + def test_is_s3_uri_with_local_path(self, mock_session): + processor = self._make_processor(mock_session) + assert processor._is_s3_uri("/local/path") is False + + def test_is_s3_uri_with_none(self, mock_session): + processor = self._make_processor(mock_session) + assert processor._is_s3_uri(None) is False + + def test_is_s3_uri_with_empty_string(self, mock_session): + processor = self._make_processor(mock_session) + assert processor._is_s3_uri("") is False + + # --- _resolve_s3_source_dir helper --- + + def test_resolve_s3_source_dir_with_tar_gz(self, mock_session): + processor = self._make_processor(mock_session) + result = processor._resolve_s3_source_dir("s3://bucket/code/sourcedir.tar.gz") + assert result == "s3://bucket/code/sourcedir.tar.gz" + + def test_resolve_s3_source_dir_with_prefix(self, mock_session): + processor = self._make_processor(mock_session) + result = processor._resolve_s3_source_dir("s3://bucket/code/") + assert result == "s3://bucket/code/sourcedir.tar.gz" + + def test_resolve_s3_source_dir_with_prefix_no_trailing_slash(self, mock_session): + processor = self._make_processor(mock_session) + result = processor._resolve_s3_source_dir("s3://bucket/code") + assert result == "s3://bucket/code/sourcedir.tar.gz" + + # --- _package_code with S3 source_dir --- + + def test_package_code_with_s3_tar_gz(self, mock_session): + processor = self._make_processor(mock_session) + result = processor._package_code( + entry_point="train.py", + source_dir="s3://my-bucket/code/sourcedir.tar.gz", + requirements=None, + job_name="test-job", + kms_key=None, + ) + assert result == "s3://my-bucket/code/sourcedir.tar.gz" + + def test_package_code_with_s3_prefix(self, mock_session): + processor = self._make_processor(mock_session) + result = processor._package_code( + entry_point="train.py", + source_dir="s3://my-bucket/code/", + requirements=None, + job_name="test-job", + kms_key=None, + ) + assert result == "s3://my-bucket/code/sourcedir.tar.gz" + + def test_package_code_with_s3_prefix_no_trailing_slash(self, mock_session): + processor = self._make_processor(mock_session) + result = processor._package_code( + entry_point="train.py", + source_dir="s3://my-bucket/code", + requirements=None, + job_name="test-job", + kms_key=None, + ) + assert result == "s3://my-bucket/code/sourcedir.tar.gz" + + def test_package_code_with_s3_does_not_upload(self, mock_session): + """S3 source_dir should not trigger any S3 upload.""" + processor = self._make_processor(mock_session) + with patch("sagemaker.core.s3.S3Uploader.upload_string_as_file_body") as mock_upload: + processor._package_code( + entry_point="train.py", + source_dir="s3://my-bucket/code/sourcedir.tar.gz", + requirements=None, + job_name="test-job", + kms_key=None, + ) + mock_upload.assert_not_called() + + # --- _pack_and_upload_code with S3 source_dir --- + + def test_pack_and_upload_code_with_s3_source_dir(self, mock_session): + processor = self._make_processor(mock_session) + + with patch( + "sagemaker.core.s3.S3Uploader.upload_string_as_file_body", + return_value="s3://test-bucket/sagemaker/test-job/source/runproc.sh", + ) as mock_upload: + result_uri, result_inputs, result_job_name = processor._pack_and_upload_code( + code="train.py", + source_dir="s3://my-bucket/code/sourcedir.tar.gz", + requirements=None, + job_name="test-job", + inputs=None, + kms_key=None, + ) + + # Should have uploaded install_requirements.py and runproc.sh + assert mock_upload.call_count == 2 + upload_uris = [ + call.kwargs.get("desired_s3_uri") or call.args[1] + for call in mock_upload.call_args_list + ] + assert any("install_requirements.py" in uri for uri in upload_uris) + assert any("runproc.sh" in uri for uri in upload_uris) + + # Helpers should go to the managed prefix, not the user's S3 location + for uri in upload_uris: + assert not uri.startswith("s3://my-bucket/code/") + + def test_pack_and_upload_code_with_s3_source_dir_creates_code_input(self, mock_session): + processor = self._make_processor(mock_session) + + with patch( + "sagemaker.core.s3.S3Uploader.upload_string_as_file_body", + return_value="s3://test-bucket/sagemaker/test-job/source/runproc.sh", + ): + _, result_inputs, _ = processor._pack_and_upload_code( + code="train.py", + source_dir="s3://my-bucket/code/sourcedir.tar.gz", + requirements=None, + job_name="test-job", + inputs=None, + kms_key=None, + ) + + # Should have a 'code' input pointing to the S3 source dir + assert len(result_inputs) == 1 + code_input = result_inputs[0] + assert code_input.input_name == "code" + assert "s3://my-bucket/code/" in code_input.s3_input.s3_uri + + # --- _generate_custom_framework_script with S3 source_dir --- + + def test_generate_custom_framework_script_with_s3_source_dir(self, mock_session): + processor = self._make_processor(mock_session) + script = processor._generate_custom_framework_script( + user_script="train.py", + entry_point="setup.sh", + source_dir="s3://my-bucket/code/sourcedir.tar.gz", + ) + + # Should not try to read local files — generates a container script instead + assert "#!/bin/bash" in script + assert "tar -xzf sourcedir.tar.gz" in script + assert "chmod +x setup.sh" in script + assert "./setup.sh" in script + assert "python train.py" in script + + def test_generate_custom_framework_script_with_local_source_dir(self, mock_session): + processor = self._make_processor(mock_session) + + with tempfile.TemporaryDirectory() as tmpdir: + entry_point_path = os.path.join(tmpdir, "setup.sh") + with open(entry_point_path, "w") as f: + f.write("#!/bin/bash\necho setup") + + script = processor._generate_custom_framework_script( + user_script="train.py", + entry_point="setup.sh", + source_dir=tmpdir, + ) + + # Should have read and embedded the local entry_point file content + assert "echo setup" in script + + # --- _generate_framework_script (no custom entry_point) --- + + def test_generate_framework_script_contains_extraction(self, mock_session): + processor = self._make_processor(mock_session) + script = processor._generate_framework_script(user_script="train.py") + + assert "tar -xzf sourcedir.tar.gz" in script + assert "python train.py" in script + + # --- Full run() integration with S3 source_dir --- + + def test_run_with_s3_source_dir(self, mock_session): + """End-to-end: run() with S3 source_dir should not raise ValueError.""" + processor = self._make_processor(mock_session) + mock_job = Mock() + + with patch.object(processor, "_start_new", return_value=mock_job): + with patch( + "sagemaker.core.s3.S3Uploader.upload_string_as_file_body", + return_value="s3://test-bucket/sagemaker/test-job/source/runproc.sh", + ): + processor.run( + code="train.py", + source_dir="s3://my-bucket/code/sourcedir.tar.gz", + wait=False, + ) + assert processor.latest_job == mock_job + + def test_run_with_s3_source_dir_prefix(self, mock_session): + """run() with S3 prefix (no .tar.gz) should also work.""" + processor = self._make_processor(mock_session) + mock_job = Mock() + + with patch.object(processor, "_start_new", return_value=mock_job): + with patch( + "sagemaker.core.s3.S3Uploader.upload_string_as_file_body", + return_value="s3://test-bucket/sagemaker/test-job/source/runproc.sh", + ): + processor.run( + code="train.py", + source_dir="s3://my-bucket/code/", + wait=False, + ) + assert processor.latest_job == mock_job diff --git a/sagemaker-mlops/tests/integ/code/s3_source_dir_processing/helpers.py b/sagemaker-mlops/tests/integ/code/s3_source_dir_processing/helpers.py new file mode 100644 index 0000000000..46711e1d95 --- /dev/null +++ b/sagemaker-mlops/tests/integ/code/s3_source_dir_processing/helpers.py @@ -0,0 +1,6 @@ +"""Helper module to verify cross-file imports work from S3 source_dir.""" + + +def get_greeting(name: str) -> str: + """Return a greeting string.""" + return f"Hello from S3 source_dir, {name}!" diff --git a/sagemaker-mlops/tests/integ/code/s3_source_dir_processing/process.py b/sagemaker-mlops/tests/integ/code/s3_source_dir_processing/process.py new file mode 100644 index 0000000000..a215c079de --- /dev/null +++ b/sagemaker-mlops/tests/integ/code/s3_source_dir_processing/process.py @@ -0,0 +1,27 @@ +"""Simple processing script for S3 source_dir integ test. + +This script validates that: +1. It can be executed from an S3-based source_dir +2. It can import from a sibling module in the same source bundle +""" +import os +import json + +from helpers import get_greeting + + +if __name__ == "__main__": + output_dir = "/opt/ml/processing/output" + os.makedirs(output_dir, exist_ok=True) + + result = { + "status": "success", + "greeting": get_greeting("integration-test"), + "source_dir_type": "s3", + } + + output_path = os.path.join(output_dir, "result.json") + with open(output_path, "w") as f: + json.dump(result, f) + + print(f"Processing complete. Output written to {output_path}")