Fix source_dir in FrameworkProcessor

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.

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,23 +1203,29 @@ 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,
                 "source",
                 "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,19 +1339,28 @@ 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,
         )
 
         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)