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