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..34385619c8 100644 --- a/sagemaker-core/src/sagemaker/core/processing.py +++ b/sagemaker-core/src/sagemaker/core/processing.py @@ -489,7 +489,10 @@ def _normalize_outputs(self, outputs=None): # If the output's s3_uri is not an s3_uri, create one. parse_result = urlparse(output.s3_output.s3_uri) if parse_result.scheme != "s3": - if getattr(self.sagemaker_session, "local_mode", False) and parse_result.scheme == "file": + if ( + getattr(self.sagemaker_session, "local_mode", False) + and parse_result.scheme == "file" + ): normalized_outputs.append(output) continue if _pipeline_config: @@ -1025,6 +1028,14 @@ class FrameworkProcessor(ScriptProcessor): framework_entrypoint_command = ["/bin/bash"] + # Container path where the source bundle is extracted. For a local ``source_dir`` the + # managed code prefix (which also holds ``install_requirements.py``) is mounted here. + _SOURCE_CODE_CONTAINER_DIR = "/opt/ml/processing/input/code" + # Container path where helper scripts are mounted when ``source_dir`` is an S3 URI. In + # that case ``/input/code`` maps to the user's (possibly read-only) S3 location, so the + # managed ``install_requirements.py`` must be mounted from a separate input instead. + _AUX_CONTAINER_DIR = "/opt/ml/processing/input/aux" + def __init__( self, image_uri: Union[str, PipelineVariable], @@ -1129,6 +1140,29 @@ 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, @@ -1136,11 +1170,36 @@ def _package_code( requirements, job_name, kms_key, + dependencies=None, ): - """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 +1214,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 +1237,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 +1253,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, @@ -1197,7 +1263,7 @@ def run( job_name: Optional[str] = None, experiment_config: Optional[Dict[str, str]] = None, kms_key: Optional[str] = None, - entry_point: Optional[str] = None + entry_point: Optional[str] = None, ): """Runs a processing job. @@ -1206,7 +1272,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 @@ -1239,7 +1311,8 @@ def run( job_name, inputs, kms_key, - entry_point + entry_point, + dependencies=dependencies, ) # Submit a processing job. @@ -1263,7 +1336,8 @@ def _pack_and_upload_code( job_name, inputs, kms_key=None, - entry_point=None + entry_point=None, + dependencies=None, ): """Pack local code bundle and upload to Amazon S3.""" if code.startswith("s3://"): @@ -1276,6 +1350,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,24 +1358,55 @@ 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 and where install_requirements.py will be + # available inside the container. When source_dir is S3, s3_payload points to the + # user's existing (possibly read-only) location, so helpers go to a separate managed + # prefix which is mounted at a dedicated aux path rather than /input/code. + 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") + install_requirements_dir = self._AUX_CONTAINER_DIR + # Mount the managed helper so install_requirements.py is present in the container. + inputs = (inputs or []) + [ + ProcessingInput( + input_name="aux", + s3_input=ProcessingS3Input( + s3_uri=install_req_s3_uri, + local_path=self._AUX_CONTAINER_DIR, + s3_data_type="S3Prefix", + s3_input_mode="File", + ), + ) + ] + 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") + # For a local source_dir the managed code prefix (with install_requirements.py) is + # mounted at /input/code alongside the extracted bundle. + install_requirements_dir = self._SOURCE_CODE_CONTAINER_DIR - # 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 + with open(_ir_mod.__file__, "r") as _ir_file: + _ir_body = _ir_file.read() s3.S3Uploader.upload_string_as_file_body( - body=open(_ir_mod.__file__, "r").read(), + body=_ir_body, desired_s3_uri=install_req_s3_uri, kms_key=evaluated_kms_key, sagemaker_session=self.sagemaker_session, ) 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 + script, + evaluated_kms_key, + entrypoint_s3_uri, + entry_point, + source_dir, + install_requirements_dir, ) return s3_runproc_sh, inputs, job_name @@ -1338,12 +1444,22 @@ def _set_entrypoint(self, command, user_script_name): ) self.entrypoint = self.framework_entrypoint_command + [user_script_location] - def _create_and_upload_runproc(self, user_script, kms_key, entrypoint_s3_uri, entry_point=None, source_dir=None): + def _create_and_upload_runproc( + self, + user_script, + kms_key, + entrypoint_s3_uri, + entry_point=None, + source_dir=None, + install_requirements_dir=None, + ): """Create runproc shell script and upload to S3 bucket.""" from sagemaker.core.workflow.utilities import _pipeline_config, hash_object if _pipeline_config and _pipeline_config.pipeline_name: - runproc_file_str = self._generate_framework_script(user_script, entry_point, source_dir) + runproc_file_str = self._generate_framework_script( + user_script, entry_point, source_dir, install_requirements_dir + ) runproc_file_hash = hash_object(runproc_file_str) s3_uri = s3.s3_path_join( self._s3_code_prefix(), @@ -1360,7 +1476,9 @@ def _create_and_upload_runproc(self, user_script, kms_key, entrypoint_s3_uri, en ) else: s3_runproc_sh = s3.S3Uploader.upload_string_as_file_body( - self._generate_framework_script(user_script, entry_point, source_dir), + self._generate_framework_script( + user_script, entry_point, source_dir, install_requirements_dir + ), desired_s3_uri=entrypoint_s3_uri, kms_key=kms_key, sagemaker_session=self.sagemaker_session, @@ -1368,24 +1486,34 @@ def _create_and_upload_runproc(self, user_script, kms_key, entrypoint_s3_uri, en return s3_runproc_sh - def _generate_framework_script(self, user_script: str, entry_point: str = None, source_dir: str = None) -> str: + def _generate_framework_script( + self, + user_script: str, + entry_point: str = None, + source_dir: str = None, + install_requirements_dir: str = None, + ) -> str: """Generate the framework entrypoint file (as text) for a processing job.""" if entry_point: - return self._generate_custom_framework_script(user_script, entry_point, source_dir) + return self._generate_custom_framework_script( + user_script, entry_point, source_dir, install_requirements_dir + ) + + install_requirements_dir = install_requirements_dir or self._SOURCE_CODE_CONTAINER_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/ - + # Debug: List files before extraction echo "Files in /opt/ml/processing/input/code/ before extraction:" ls -la - + # Extract source code if [ -f sourcedir.tar.gz ]; then tar -xzf sourcedir.tar.gz @@ -1400,18 +1528,23 @@ def _generate_framework_script(self, user_script: str, entry_point: str = None, # Some py3 containers has typing, which may breaks pip install pip uninstall --yes typing - python3 /opt/ml/processing/input/code/install_requirements.py requirements.txt + python3 {install_requirements_dir}/install_requirements.py requirements.txt fi {entry_point_command} {entry_point} "$@" """ ).format( + install_requirements_dir=install_requirements_dir, entry_point_command=" ".join(self.command), entry_point=user_script, ) def _generate_custom_framework_script( - self, user_script: str, entry_point: str, source_dir: str = None + self, + user_script: str, + entry_point: str, + source_dir: str = None, + install_requirements_dir: str = None, ) -> str: """ Generate a custom framework script with a user-provided entrypoint embedded. @@ -1424,10 +1557,52 @@ def _generate_custom_framework_script( entry_point (str): Path to the custom entrypoint script file source_dir (str): Path to the source directory. If provided and entry_point is relative, it will be combined with source_dir. + install_requirements_dir (str): Container directory that holds + ``install_requirements.py`` (default: the extracted source code dir). 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): + install_requirements_dir = install_requirements_dir or self._SOURCE_CODE_CONTAINER_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 {install_requirements_dir}/install_requirements.py requirements.txt + fi + + # Execute custom entrypoint + chmod +x {entry_point} + ./{entry_point} + + {entry_point_command} {user_script} "$@" + """ + ).format( + install_requirements_dir=install_requirements_dir, + 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) @@ -1439,11 +1614,13 @@ def _generate_custom_framework_script( entry_point_content = f.read() # Generate the script with embedded entry_point content - return dedent("""\ + return dedent( + """\ {entry_point_content} {entry_point_command} {entry_point} "$@" - """).format( + """ + ).format( entry_point_content=entry_point_content, entry_point_command=" ".join(self.command), entry_point=user_script, diff --git a/sagemaker-core/tests/unit/test_processing.py b/sagemaker-core/tests/unit/test_processing.py index f5349faea0..ca811424d1 100644 --- a/sagemaker-core/tests/unit/test_processing.py +++ b/sagemaker-core/tests/unit/test_processing.py @@ -14,7 +14,7 @@ import pytest import os import tempfile -from unittest.mock import Mock, patch, MagicMock, mock_open +from unittest.mock import Mock, patch from sagemaker.core.processing import ( Processor, ScriptProcessor, @@ -48,9 +48,6 @@ def mock_session(): class TestProcessorNormalizeArgs: def test_normalize_args_with_pipeline_variable_code(self, mock_session): - from sagemaker.core.workflow.pipeline_context import PipelineSession - from sagemaker.core.workflow import is_pipeline_variable - processor = Processor( role="arn:aws:iam::123456789012:role/SageMakerRole", image_uri="test-image:latest", @@ -238,8 +235,6 @@ def test_normalize_outputs_invalid_type(self, mock_session): processor._normalize_outputs(["invalid"]) - - class TestBugConditionFileUriReplacedInLocalMode: """Bug condition exploration test: file:// URIs should be preserved in local mode. @@ -444,7 +439,9 @@ def test_pipeline_variable_uri_skips_normalization(self, session_local_mode_fals [{"key": "value"}], ], ) - def test_non_processing_output_raises_type_error(self, invalid_output, session_local_mode_false): + def test_non_processing_output_raises_type_error( + self, invalid_output, session_local_mode_false + ): """Non-ProcessingOutput objects must raise TypeError. **Validates: Requirements 3.4** @@ -882,9 +879,7 @@ def test_generate_framework_script_with_custom_entry_point(self, mock_session): entry_point_path = f.name try: - script = processor._generate_framework_script( - "train.py", entry_point=entry_point_path - ) + script = processor._generate_framework_script("train.py", entry_point=entry_point_path) assert custom_script_content in script assert "python3 train.py" in script assert "tar -xzf sourcedir.tar.gz" not in script @@ -1042,7 +1037,7 @@ def test_start_new_submit_success(self, mock_session): }, ): with patch("sagemaker.core.processing.serialize", return_value={}): - with patch("sagemaker.core.processing.ProcessingJob") as mock_job: + with patch("sagemaker.core.processing.ProcessingJob"): with patch( "sagemaker.core.utils.code_injection.codec.transform", return_value={"processing_job_name": "test-job"}, @@ -1220,7 +1215,6 @@ def test_package_code_source_dir_not_exists(self, mock_session): kms_key=None, ) - def test_package_code_with_code_location(self, mock_session): processor = FrameworkProcessor( role="arn:aws:iam::123456789012:role/SageMakerRole", @@ -1477,8 +1471,6 @@ def test_generate_job_name_with_invalid_chars(self, mock_session): class TestProcessorWithPipelineVariable: def test_get_process_args_with_pipeline_variable_role(self, mock_session): - from sagemaker.core.workflow import is_pipeline_variable - role_var = Mock() with patch("sagemaker.core.processing.is_pipeline_variable", return_value=True): @@ -1524,11 +1516,17 @@ def test_start_new_removes_tags_from_processing_job(self, mock_session): "tags": [{"Key": "Key", "Value": "Value"}], }, ): - with patch("sagemaker.core.processing.serialize", return_value={"tags": [{"Key": "Key", "Value": "Value"}]}): + with patch( + "sagemaker.core.processing.serialize", + return_value={"tags": [{"Key": "Key", "Value": "Value"}]}, + ): with patch("sagemaker.core.processing.ProcessingJob") as mock_job_class: with patch( "sagemaker.core.utils.code_injection.codec.transform", - return_value={"processing_job_name": "test-job", "tags": [{"Key": "Key", "Value": "Value"}]}, + return_value={ + "processing_job_name": "test-job", + "tags": [{"Key": "Key", "Value": "Value"}], + }, ): processor._start_new([], [], None) # Verify ProcessingJob was called without tags @@ -1763,3 +1761,277 @@ 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 plus an 'aux' input + # that mounts the managed install_requirements.py helper into the container. + assert len(result_inputs) == 2 + inputs_by_name = {i.input_name: i for i in result_inputs} + + code_input = inputs_by_name["code"] + assert "s3://my-bucket/code/" in code_input.s3_input.s3_uri + + aux_input = inputs_by_name["aux"] + assert aux_input.s3_input.local_path == "/opt/ml/processing/input/aux" + assert aux_input.s3_input.s3_uri.endswith("install_requirements.py") + # The helper must NOT come from the user's (possibly read-only) S3 location. + assert not aux_input.s3_input.s3_uri.startswith("s3://my-bucket/code/") + + # --- _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_s3_source_dir_script_references_mounted_install_requirements(self, mock_session): + """The install_requirements.py path in the generated script must match the aux mount. + + For an S3 source_dir, /opt/ml/processing/input/code maps to the user's location and + does NOT contain install_requirements.py; the helper is mounted at the aux path. The + generated script must invoke it from there, not from the code dir. + """ + processor = self._make_processor(mock_session) + aux_dir = FrameworkProcessor._AUX_CONTAINER_DIR + + # Both the plain framework script and the custom-entrypoint script must point at + # the aux mount when told the helper lives there. + plain = processor._generate_framework_script( + user_script="train.py", + source_dir="s3://my-bucket/code/sourcedir.tar.gz", + install_requirements_dir=aux_dir, + ) + custom = processor._generate_custom_framework_script( + user_script="train.py", + entry_point="setup.sh", + source_dir="s3://my-bucket/code/sourcedir.tar.gz", + install_requirements_dir=aux_dir, + ) + for script in (plain, custom): + assert f"{aux_dir}/install_requirements.py" in script + assert "/opt/ml/processing/input/code/install_requirements.py" not in script + + def test_local_source_dir_script_references_code_dir_install_requirements(self, mock_session): + """For a local source_dir the helper is co-located under /input/code (unchanged).""" + processor = self._make_processor(mock_session) + script = processor._generate_framework_script(user_script="train.py") + assert "/opt/ml/processing/input/code/install_requirements.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}")