task(RHOAIENG-33045): Centralised image logic + fixed tests

Signed-off-by: Pat O'Connor <[email protected]>

Author: chipspeak

.github/workflows/e2e_tests.yaml

@@ -57,7 +57,7 @@ jobs:
       - name: Set up specific Python version
         uses: actions/setup-python@v5
         with:
-          python-version: '3.11'
+          python-version: '3.12'
           cache: 'pip' # caching pip dependencies
 
       - name: Setup NVidia GPU environment for KinD

src/codeflare_sdk/common/utils/constants.py

@@ -1,8 +1,14 @@
 RAY_VERSION = "2.47.1"
-'''
+"""
 The below are used to define the default runtime image for the Ray Cluster.
 * For python 3.11:ray:2.47.1-py311-cu121
 * For python 3.12:ray:2.47.1-py312-cu121
-'''
+"""
 CUDA_PY311_RUNTIME_IMAGE = "quay.io/modh/ray@sha256:6d076aeb38ab3c34a6a2ef0f58dc667089aa15826fa08a73273c629333e12f1e"
 CUDA_PY312_RUNTIME_IMAGE = "quay.io/modh/ray@sha256:23860dfe2e47bb69709b3883b08fd1a4d836ce02eaf8d0afeeafe6986d0fc8fb"
+
+# Centralized image selection
+SUPPORTED_PYTHON_VERSIONS = {
+    "3.11": CUDA_PY311_RUNTIME_IMAGE,
+    "3.12": CUDA_PY312_RUNTIME_IMAGE,
+}

src/codeflare_sdk/common/utils/unit_test_support.py

@@ -15,6 +15,7 @@
 import string
 import sys
 from codeflare_sdk.common.utils import constants
+from codeflare_sdk.common.utils.utils import get_ray_image_for_python_version
 from codeflare_sdk.ray.cluster.cluster import (
     Cluster,
     ClusterConfiguration,
@@ -69,7 +70,7 @@ def create_cluster_wrong_type():
         worker_extended_resource_requests={"nvidia.com/gpu": 7},
         appwrapper=True,
         image_pull_secrets=["unit-test-pull-secret"],
-        image=constants.CUDA_RUNTIME_IMAGE,
+        image=constants.CUDA_PY312_RUNTIME_IMAGE,
         write_to_file=True,
         labels={1: 1},
     )
@@ -294,8 +295,8 @@ def apply_template(yaml_file_path, variables):
 
 
 def get_expected_image():
-    # TODO: Select image based on Python version
-    return constants.CUDA_RUNTIME_IMAGE
+    # Use centralized image selection logic (fallback to 3.12 for test consistency)
+    return get_ray_image_for_python_version(warn_on_unsupported=True)
 
 
 def get_template_variables():

src/codeflare_sdk/common/utils/utils.py

@@ -0,0 +1,47 @@
+# Copyright 2025 IBM, Red Hat
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License 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 codeflare_sdk.common.utils.constants import (
+    SUPPORTED_PYTHON_VERSIONS,
+    CUDA_PY312_RUNTIME_IMAGE,
+)
+
+
+def get_ray_image_for_python_version(python_version=None, warn_on_unsupported=True):
+    """
+    Get the appropriate Ray image for a given Python version.
+    If no version is provided, uses the current runtime Python version.
+    This prevents us needing to hard code image versions for tests.
+
+    Args:
+        python_version: Python version string (e.g. "3.11"). If None, detects current version.
+        warn_on_unsupported: If True, warns and returns None for unsupported versions.
+                           If False, silently falls back to Python 3.12 image.
+    """
+    if python_version is None:
+        import sys
+
+        python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
+
+    if python_version in SUPPORTED_PYTHON_VERSIONS:
+        return SUPPORTED_PYTHON_VERSIONS[python_version]
+    elif warn_on_unsupported:
+        import warnings
+
+        warnings.warn(
+            f"No default Ray image defined for {python_version}. Please provide your own image or use one of the following python versions: {', '.join(SUPPORTED_PYTHON_VERSIONS.keys())}."
+        )
+        return None
+    else:
+        return CUDA_PY312_RUNTIME_IMAGE

src/codeflare_sdk/ray/cluster/build_ray_cluster.py

@@ -21,6 +21,7 @@
 from ...common.kubernetes_cluster import get_api_client, config_check
 from kubernetes.client.exceptions import ApiException
 from ...common.utils.constants import RAY_VERSION
+from ...common.utils.utils import get_ray_image_for_python_version
 import codeflare_sdk
 import os
 
@@ -95,10 +96,8 @@
     ),
 ]
 
-SUPPORTED_PYTHON_VERSIONS = {
-    "3.11": constants.CUDA_PY311_RUNTIME_IMAGE,
-    "3.12": constants.CUDA_PY312_RUNTIME_IMAGE,
-}
+# Use centralized mapping from constants (so that we only have to update constants.py)
+SUPPORTED_PYTHON_VERSIONS = constants.SUPPORTED_PYTHON_VERSIONS
 
 
 # RayCluster/AppWrapper builder function
@@ -273,16 +272,11 @@ def with_nb_annotations(annotations: dict):
 def update_image(image) -> str:
     """
     The update_image() function automatically sets the image config parameter to a preset image based on Python version if not specified.
-    If no Ray image exists for the given Python version a warning is produced.
+    This now points to the centralized function in utils.py.
     """
     if not image:
-        python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
-        if python_version in SUPPORTED_PYTHON_VERSIONS:
-            image = SUPPORTED_PYTHON_VERSIONS[python_version]
-        else:
-            warnings.warn(
-                f"No default Ray image defined for {python_version}. Please provide your own image or use one of the following python versions: {', '.join(SUPPORTED_PYTHON_VERSIONS.keys())}."
-            )
+        # Pull the image based on the matching Python version (or output a warning if not supported)
+        image = get_ray_image_for_python_version(warn_on_unsupported=True)
     return image
 
 

src/codeflare_sdk/ray/cluster/test_build_ray_cluster.py

@@ -43,6 +43,7 @@ def test_update_image_without_supported_python_version(mocker):
         "codeflare_sdk.ray.cluster.build_ray_cluster.SUPPORTED_PYTHON_VERSIONS",
         {
             "3.11": "ray-py3.11",
+            "3.12": "ray-py3.12",
         },
     )
 
@@ -60,7 +61,7 @@ def test_update_image_without_supported_python_version(mocker):
 
     # Assert that the warning was called with the expected message
     warn_mock.assert_called_once_with(
-        "No default Ray image defined for 3.8. Please provide your own image or use one of the following python versions: 3.11."
+        "No default Ray image defined for 3.8. Please provide your own image or use one of the following python versions: 3.11, 3.12."
     )
 
     # Assert that no image was set since the Python version is not supported

tests/e2e/support.py

@@ -27,7 +27,10 @@ def get_ray_cluster(cluster_name, namespace):
 
 
 def get_ray_image():
-    return os.getenv("RAY_IMAGE", constants.CUDA_RUNTIME_IMAGE)
+    return os.getenv(
+        "RAY_IMAGE",
+        constants.get_ray_image_for_python_version(warn_on_unsupported=False),
+    )
 
 
 def get_setup_env_variables(**kwargs):