feat(feature-processor): Add Spark container image version resolver

Create _image_resolver.py with SPARK_IMAGE_SUPPORT_MATRIX and
_get_spark_image_uri(session) that auto-detects PySpark + Python
versions and returns the correct SageMaker Spark container image URI.
Raises ValueError for unsupported combinations (e.g. Spark 3.4 has
no container image, Python 3.12 only with Spark 3.5).

---
X-AI-Prompt: create image resolver module for multi-spark version container image resolution
X-AI-Tool: kiro

Author: BassemHalim

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

sagemaker-mlops/tests/unit/sagemaker/mlops/feature_store/feature_processor/test_image_resolver.py

@@ -0,0 +1,104 @@
+# 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.
+from __future__ import absolute_import
+
+import sys
+
+import pyspark
+import pytest
+from mock import Mock, patch
+
+from sagemaker.mlops.feature_store.feature_processor._image_resolver import _get_spark_image_uri
+
+
+@patch("sagemaker.mlops.feature_store.feature_processor._image_resolver.image_uris.retrieve")
+def test_spark_33_py39(mock_retrieve):
+    mock_retrieve.return_value = "123456.dkr.ecr.us-west-2.amazonaws.com/sagemaker-spark-processing:3.3-cpu-py39-v1"
+    session = Mock(boto_region_name="us-west-2")
+    with patch.object(pyspark, "__version__", "3.3.2"), \
+         patch.object(sys, "version_info", (3, 9, 0)):
+        result = _get_spark_image_uri(session)
+    mock_retrieve.assert_called_once_with(
+        framework="spark",
+        region="us-west-2",
+        version="3.3",
+        py_version="py39",
+        container_version="v1",
+    )
+    assert result == mock_retrieve.return_value
+
+
+@patch("sagemaker.mlops.feature_store.feature_processor._image_resolver.image_uris.retrieve")
+def test_spark_35_py39(mock_retrieve):
+    mock_retrieve.return_value = "123456.dkr.ecr.us-west-2.amazonaws.com/sagemaker-spark-processing:3.5-cpu-py39-v1"
+    session = Mock(boto_region_name="us-west-2")
+    with patch.object(pyspark, "__version__", "3.5.1"), \
+         patch.object(sys, "version_info", (3, 9, 0)):
+        result = _get_spark_image_uri(session)
+    mock_retrieve.assert_called_once_with(
+        framework="spark",
+        region="us-west-2",
+        version="3.5",
+        py_version="py39",
+        container_version="v1",
+    )
+    assert result == mock_retrieve.return_value
+
+
+@patch("sagemaker.mlops.feature_store.feature_processor._image_resolver.image_uris.retrieve")
+def test_spark_35_py312(mock_retrieve):
+    mock_retrieve.return_value = "123456.dkr.ecr.us-west-2.amazonaws.com/sagemaker-spark-processing:3.5-cpu-py312-v1"
+    session = Mock(boto_region_name="us-west-2")
+    with patch.object(pyspark, "__version__", "3.5.1"), \
+         patch.object(sys, "version_info", (3, 12, 0)):
+        result = _get_spark_image_uri(session)
+    mock_retrieve.assert_called_once_with(
+        framework="spark",
+        region="us-west-2",
+        version="3.5",
+        py_version="py312",
+        container_version="v1",
+    )
+    assert result == mock_retrieve.return_value
+
+
+def test_spark_34_raises():
+    session = Mock(boto_region_name="us-west-2")
+    with patch.object(pyspark, "__version__", "3.4.1"), \
+         patch.object(sys, "version_info", (3, 9, 0)):
+        with pytest.raises(ValueError, match="No SageMaker Spark container image available for Spark 3.4"):
+            _get_spark_image_uri(session)
+
+
+def test_spark_35_py310_raises():
+    session = Mock(boto_region_name="us-west-2")
+    with patch.object(pyspark, "__version__", "3.5.1"), \
+         patch.object(sys, "version_info", (3, 10, 0)):
+        with pytest.raises(ValueError, match="SageMaker Spark 3.5 container images support"):
+            _get_spark_image_uri(session)
+
+
+def test_spark_33_py312_raises():
+    session = Mock(boto_region_name="us-west-2")
+    with patch.object(pyspark, "__version__", "3.3.2"), \
+         patch.object(sys, "version_info", (3, 12, 0)):
+        with pytest.raises(ValueError, match="SageMaker Spark 3.3 container images support"):
+            _get_spark_image_uri(session)
+
+
+def test_unknown_spark_version_raises():
+    session = Mock(boto_region_name="us-west-2")
+    with patch.object(pyspark, "__version__", "3.6.0"), \
+         patch.object(sys, "version_info", (3, 9, 0)):
+        with pytest.raises(ValueError, match="No SageMaker Spark container image available for Spark 3.6"):
+            _get_spark_image_uri(session)