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,31 @@ 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, re; "
+ "release_file = os.path.join(os.environ.get('SPARK_HOME', '/usr/lib/spark'), 'RELEASE'); "
+ "spark_ver = '3.5'; "
+ "rf = open(release_file).read() if os.path.exists(release_file) else ''; "
+ "m = re.search(r'Spark (\\d+\\.\\d+)', rf); "
+ "spark_ver = m.group(1) if m else spark_ver; "
+ "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, '*' + spark_ver + '*.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 +847,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

@@ -26,6 +26,7 @@ dependencies = [
  "sagemaker-core>=2.11.0",
  "sagemaker-train>=1.11.0",
  "sagemaker-serve>=1.11.0",
+ "cryptography>=46.0.0",
  "boto3>=1.42.2,<2.0",
  "botocore>=1.42.2,<2.0",
  "pyiceberg[glue]>=0.8.0",
@@ -35,8 +36,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 +46,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/__init__.py

@@ -4,7 +4,7 @@
 
 # FeatureGroup with additional operational support
 from sagemaker.core.resources import FeatureGroup, FeatureMetadata
-from sagemaker.mlops.feature_store.feature_group_manager import FeatureGroupManager, LakeFormationConfig
+from sagemaker.mlops.feature_store.feature_group_manager import FeatureGroupManager, LakeFormationConfig, IcebergProperties
 
 # Shapes from core (Pydantic - no to_dict() needed)
 from sagemaker.core.shapes import (
@@ -82,6 +82,7 @@
  "FeatureValue",
  "Filter",
  "LakeFormationConfig",
+ "IcebergProperties",
  "OfflineStoreConfig",
  "OnlineStoreConfig",
  "OnlineStoreSecurityConfig",

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",
+ )