feat(feature-processor): Add Lake Formation credential vending and Spark 3.5/Python 3.12 support

sagemaker-core/src/sagemaker/core/remote_function/job.py

@@ -670,6 +670,25 @@ def __init__(
             sagemaker_session=self.sagemaker_session,
         )
 
+        # When using Spark, ensure sagemaker-feature-store-pyspark is installed
+        # and its version-matched JAR is on Spark's classpath before spark-submit
+        if spark_config:
+            install_cmd = (
+                "pip install --root-user-action=ignore"
+                " 'sagemaker-feature-store-pyspark>=2,<3'"
+            )
+            copy_jar_cmd = (
+                "python3 -c \"import feature_store_pyspark, shutil, os, glob; "
+                "jars_dir = os.path.join(os.path.dirname(feature_store_pyspark.__file__), 'jars'); "
+                "[shutil.copy(j, '/usr/lib/spark/jars/') "
+                "for j in glob.glob(os.path.join(jars_dir, '*3.5*.jar'))]\""
+            )
+            if self.pre_execution_commands is None:
+                self.pre_execution_commands = [install_cmd, copy_jar_cmd]
+            elif install_cmd not in self.pre_execution_commands:
+                self.pre_execution_commands.append(install_cmd)
+                self.pre_execution_commands.append(copy_jar_cmd)
+
         self.pre_execution_script = resolve_value_from_config(
             direct_input=pre_execution_script,
             config_path=REMOTE_FUNCTION_PRE_EXECUTION_SCRIPT,
@@ -822,15 +841,23 @@ def _get_default_spark_image(session):
 
         py_version = str(sys.version_info[0]) + str(sys.version_info[1])
 
-        if py_version not in ["39"]:
+        if py_version not in ["39", "312"]:
             raise ValueError(
-                "The SageMaker Spark image for remote job only supports Python version 3.9. "
+                "The SageMaker Spark image for remote job only supports Python versions 3.9 and 3.12."
             )
 
+        # Detect Spark version from installed pyspark, fall back to default
+        spark_version = DEFAULT_SPARK_VERSION
+        try:
+            import pyspark
+            spark_version = ".".join(pyspark.__version__.split(".")[:2])
+        except ImportError:
+            pass
+
         image_uri = image_uris.retrieve(
             framework=SPARK_NAME,
             region=region,
-            version=DEFAULT_SPARK_VERSION,
+            version=spark_version,
             instance_type=None,
             py_version=f"py{py_version}",
             container_version=DEFAULT_SPARK_CONTAINER_VERSION,

sagemaker-core/tests/unit/remote_function/test_job.py

@@ -133,7 +133,7 @@ def test_get_default_spark_image_unsupported_python_raises_error(self, mock_sess
         with patch.object(sys, "version_info", (3, 8, 0)):
             with pytest.raises(
                 ValueError,
-                match="SageMaker Spark image for remote job only supports Python version 3.9",
+                match="SageMaker Spark image for remote job only supports Python versions 3.9 and 3.12",
             ):
                 _JobSettings._get_default_spark_image(mock_session)
 

sagemaker-core/tests/unit/remote_function/test_job_comprehensive.py

@@ -131,10 +131,21 @@ def test_get_default_spark_image_unsupported_python(self, mock_session):
         with patch.object(sys, "version_info", (3, 8, 0)):
             with pytest.raises(
                 ValueError,
-                match="SageMaker Spark image for remote job only supports Python version 3.9",
+                match="SageMaker Spark image for remote job only supports Python versions 3.9 and 3.12",
             ):
                 _JobSettings._get_default_spark_image(mock_session)
 
+    @patch("sagemaker.core.remote_function.job.image_uris.retrieve", return_value="mock-image-uri")
+    def test_get_default_spark_image_auto_detects_pyspark_version(
+        self, mock_retrieve, mock_session
+    ):
+        mock_pyspark = MagicMock(__version__="3.5.1")
+        with patch.object(sys, "version_info", (3, 9, 0)):
+            with patch.dict("sys.modules", {"pyspark": mock_pyspark}):
+                result = _JobSettings._get_default_spark_image(mock_session)
+        assert result == "mock-image-uri"
+        assert mock_retrieve.call_args.kwargs["version"] == "3.5"
+
 
 class TestJobMethods:
     """Test _Job class methods for uncovered lines."""

sagemaker-mlops/pyproject.toml

@@ -35,8 +35,8 @@ dependencies = [
 
 [project.optional-dependencies]
 feature-processor = [
-    "pyspark==3.3.2",
-    "sagemaker-feature-store-pyspark-3.3",
+    "pyspark==3.5.1",
+    "sagemaker-feature-store-pyspark==2.0.0",
     "setuptools<82",
 ]
 
@@ -45,8 +45,8 @@ test = [
     "pytest-cov",
     "mock",
     "setuptools<82",
-    "pyspark==3.3.2",
-    "sagemaker-feature-store-pyspark-3.3",
+    "pyspark==3.5.1",
+    "sagemaker-feature-store-pyspark==2.0.0",
     "pandas<3.0",
     "numpy<3.0",
 ]

sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_processor/_config_uploader.py

@@ -15,6 +15,8 @@
 from typing import Callable, Dict, Optional, Tuple, List, Union
 
 import attr
+from cryptography.hazmat.primitives.asymmetric import ec
+from cryptography.hazmat.primitives import serialization as crypto_serialization
 
 from sagemaker.core.helper.session_helper import Session
 from sagemaker.mlops.feature_store.feature_processor._constants import (
@@ -52,13 +54,13 @@ class ConfigUploader:
 
     def prepare_step_input_channel_for_spark_mode(
         self, func: Callable, s3_base_uri: str, sagemaker_session: Session
-    ) -> Tuple[List[Channel], Dict]:
+    ) -> Tuple[List[Channel], Dict, str]:
         """Prepares input channels for SageMaker Pipeline Step.
 
         Returns:
-            Tuple of (List[Channel], spark_dependency_paths dict)
+            Tuple of (List[Channel], spark_dependency_paths dict, public_key_pem str)
         """
-        self._prepare_and_upload_callable(func, s3_base_uri, sagemaker_session)
+        public_key_pem = self._prepare_and_upload_callable(func, s3_base_uri, sagemaker_session)
         bootstrap_scripts_s3uri = self._prepare_and_upload_runtime_scripts(
             self.remote_decorator_config.spark_config,
             s3_base_uri,
@@ -139,18 +141,33 @@ def prepare_step_input_channel_for_spark_mode(
             SPARK_JAR_FILES_PATH: submit_jars_s3_paths,
             SPARK_PY_FILES_PATH: submit_py_files_s3_paths,
             SPARK_FILES_PATH: submit_files_s3_path,
-        }
+        }, public_key_pem
 
     def _prepare_and_upload_callable(
         self, func: Callable, s3_base_uri: str, sagemaker_session: Session
-    ) -> None:
-        """Prepares and uploads callable to S3"""
+    ) -> str:
+        """Prepares and uploads callable to S3.
+
+        Returns:
+            str: The public key PEM string for signature verification on the remote side.
+        """
+        private_key = ec.generate_private_key(ec.SECP256R1())
+        public_key_pem = (
+            private_key.public_key()
+            .public_bytes(
+                crypto_serialization.Encoding.PEM,
+                crypto_serialization.PublicFormat.SubjectPublicKeyInfo,
+            )
+            .decode("utf-8")
+        )
         stored_function = StoredFunction(
             sagemaker_session=sagemaker_session,
             s3_base_uri=s3_base_uri,
+            signing_key=private_key,
             s3_kms_key=self.remote_decorator_config.s3_kms_key,
         )
         stored_function.save(func)
+        return public_key_pem
 
     def _prepare_and_upload_workspace(
         self,

sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_processor/_feature_processor_config.py

@@ -47,6 +47,7 @@ class FeatureProcessorConfig:
     parameters: Optional[Dict[str, Union[str, Dict]]] = attr.ib()
     enable_ingestion: bool = attr.ib()
     spark_config: Dict[str, str] = attr.ib()
+    use_lake_formation_credentials: bool = attr.ib()
 
     @staticmethod
     def create(
@@ -59,6 +60,7 @@ def create(
         parameters: Optional[Dict[str, Union[str, Dict]]],
         enable_ingestion: bool,
         spark_config: Dict[str, str],
+        use_lake_formation_credentials: bool = False,
     ) -> "FeatureProcessorConfig":
         """Static initializer."""
         return FeatureProcessorConfig(
@@ -69,4 +71,5 @@ def create(
             parameters=parameters,
             enable_ingestion=enable_ingestion,
             spark_config=spark_config,
+            use_lake_formation_credentials=use_lake_formation_credentials,
         )

sagemaker-mlops/src/sagemaker/mlops/feature_store/feature_processor/_image_resolver.py

@@ -0,0 +1,65 @@
+# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License"). You
+# may not use this file except in compliance with the License. A copy of
+# the License is located at
+#
+#     http://aws.amazon.com/apache2.0/
+#
+# or in the "license" file accompanying this file. This file is
+# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
+# ANY KIND, either express or implied. See the License for the specific
+# language governing permissions and limitations under the License.
+"""Resolves SageMaker Spark container image URIs based on installed PySpark and Python versions."""
+from __future__ import absolute_import
+
+import sys
+
+from sagemaker.core import image_uris
+
+SPARK_IMAGE_SUPPORT_MATRIX = {
+    "3.1": ["py37"],
+    "3.2": ["py39"],
+    "3.3": ["py39"],
+    "3.5": ["py39", "py312"],
+}
+
+
+def _get_spark_image_uri(session):
+    """Resolve the SageMaker Spark container image URI for the installed PySpark and Python versions.
+
+    Args:
+        session: SageMaker Session with boto_region_name attribute.
+
+    Returns:
+        str: The ECR image URI for the matching Spark container.
+
+    Raises:
+        ValueError: If the Spark/Python version combination is not supported.
+    """
+    import pyspark
+
+    spark_version = ".".join(pyspark.__version__.split(".")[:2])
+    py_version = f"py{sys.version_info[0]}{sys.version_info[1]}"
+
+    supported_py = SPARK_IMAGE_SUPPORT_MATRIX.get(spark_version)
+    if supported_py is None:
+        supported = ", ".join(sorted(SPARK_IMAGE_SUPPORT_MATRIX.keys()))
+        raise ValueError(
+            f"No SageMaker Spark container image available for Spark {spark_version}. "
+            f"Supported versions for remote execution: {supported}."
+        )
+
+    if py_version not in supported_py:
+        raise ValueError(
+            f"SageMaker Spark {spark_version} container images support "
+            f"{', '.join(supported_py)}. Current Python version: {py_version}."
+        )
+
+    return image_uris.retrieve(
+        framework="spark",
+        region=session.boto_region_name,
+        version=spark_version,
+        py_version=py_version,
+        container_version="v1",
+    )