From 3c1038fd0a9ab87a3e31cd7fdf51684d5e2f26d3 Mon Sep 17 00:00:00 2001
From: Abhinav Maurya <owner@wegenz.in>
Date: Thu, 29 Oct 2026 07:34:28 +0530
Subject: [PATCH] fix: validate python version in get_base_python_image_uri

get_base_python_image_uri previously accepted any string for py_version
without validation, returning image URIs for non-existent base images in
ECR (such as 311 or 313) that failed downstream at job execution.

Add py_versions configuration to sagemaker-base-python.json (310, 38)
and validate py_version via _validate_py_version_and_set_if_needed,
raising a clear ValueError when an unsupported version is requested.
Also add unit test coverage for invalid python versions.

Fixes #6539
---
 .../core/image_uri_config/sagemaker-base-python.json     | 4 ++++
 sagemaker-core/src/sagemaker/core/image_uris.py          | 8 ++++++--
 sagemaker-core/tests/unit/image_uris/test_base_python.py | 9 +++++++++
 3 files changed, 19 insertions(+), 2 deletions(-)

diff --git a/sagemaker-core/src/sagemaker/core/image_uri_config/sagemaker-base-python.json b/sagemaker-core/src/sagemaker/core/image_uri_config/sagemaker-base-python.json
index 42749fbbee..34450798ea 100644
--- a/sagemaker-core/src/sagemaker/core/image_uri_config/sagemaker-base-python.json
+++ b/sagemaker-core/src/sagemaker/core/image_uri_config/sagemaker-base-python.json
@@ -47,6 +47,10 @@
                 "us-west-1": "742091327244",
                 "us-west-2": "236514542706"
             },
+            "py_versions": [
+                "310",
+                "38"
+            ],
             "repository": "sagemaker-base-python"
         }
     }
diff --git a/sagemaker-core/src/sagemaker/core/image_uris.py b/sagemaker-core/src/sagemaker/core/image_uris.py
index 760cfce54f..e360b59d62 100644
--- a/sagemaker-core/src/sagemaker/core/image_uris.py
+++ b/sagemaker-core/src/sagemaker/core/image_uris.py
@@ -972,11 +972,14 @@ def get_base_python_image_uri(region, py_version="310") -> str:
 
     Args:
         region (str): The AWS region to use for image URI.
-        py_version (str): The python version to use for the image. Can be 310 or 38
-        Default to 310
+        py_version (str): The python version to use for the image. Can be 310 or 38.
+            Defaults to 310.
 
     Returns:
         str: The image URI string.
+
+    Raises:
+        ValueError: If the python version is unsupported.
     """
 
     framework = "sagemaker-base-python"
@@ -988,6 +991,7 @@ def get_base_python_image_uri(region, py_version="310") -> str:
     config = config_for_framework(framework)
     version_config = config["versions"][_version_for_config(version, config)]
 
+    py_version = _validate_py_version_and_set_if_needed(py_version, version_config, framework)
     registry = _registry_from_region(region, version_config["registries"])
 
     repo = version_config["repository"] + "-" + py_version
diff --git a/sagemaker-core/tests/unit/image_uris/test_base_python.py b/sagemaker-core/tests/unit/image_uris/test_base_python.py
index 546b72bc72..cd52d297af 100644
--- a/sagemaker-core/tests/unit/image_uris/test_base_python.py
+++ b/sagemaker-core/tests/unit/image_uris/test_base_python.py
@@ -31,3 +31,12 @@ def test_get_base_python_image_uri(py_version, load_config):
             repo=repo, account=REGISTRIES[region], region=region
         )
         assert expected == uri
+
+
+@pytest.mark.parametrize("unsupported_py_version", ["311", "312", "313", "39", "invalid"])
+def test_get_base_python_image_uri_unsupported_version(unsupported_py_version):
+    with pytest.raises(ValueError, match="Unsupported Python version"):
+        image_uris.get_base_python_image_uri(
+            region="us-west-2",
+            py_version=unsupported_py_version,
+        )