Adjust demo notebooks with new ray image defaults

ChristianZaccaria · openshift-merge-bot[bot] · commit 2c7bcfa5ac5a · 2024-10-17T13:48:37.000Z
diff --git a/demo-notebooks/additional-demos/hf_interactive.ipynb b/demo-notebooks/additional-demos/hf_interactive.ipynb
@@ -68,8 +68,12 @@
    "source": [
     "Here, we want to define our cluster by specifying the resources we require for our batch workload. Below, we define our cluster object (which generates a corresponding Ray Cluster).\n",
     "\n",
-    "NOTE: 'quay.io/modh/ray:2.35.0-py39-cu121' is the default image used by the CodeFlare SDK for creating a RayCluster resource. \n",
-    "If you have your own Ray image which suits your purposes, specify it in image field to override the default image."
+    "NOTE: The default images used by the CodeFlare SDK for creating a RayCluster resource depend on the installed Python version:\n",
+    "\n",
+    "- For Python 3.9: 'quay.io/modh/ray:2.35.0-py39-cu121'\n",
+    "- For Python 3.11: 'quay.io/modh/ray:2.35.0-py311-cu121'\n",
+    "\n",
+    "If you prefer to use a custom Ray image that better suits your needs, you can specify it in the image field to override the default."
    ]
   },
   {
diff --git a/demo-notebooks/additional-demos/local_interactive.ipynb b/demo-notebooks/additional-demos/local_interactive.ipynb
@@ -35,8 +35,12 @@
    "metadata": {},
    "source": [
     "\n",
-    "NOTE: 'quay.io/modh/ray:2.35.0-py39-cu121' is the default image used by the CodeFlare SDK for creating a RayCluster resource. \n",
-    "If you have your own Ray image which suits your purposes, specify it in image field to override the default image."
+    "NOTE: The default images used by the CodeFlare SDK for creating a RayCluster resource depend on the installed Python version:\n",
+    "\n",
+    "- For Python 3.9: 'quay.io/modh/ray:2.35.0-py39-cu121'\n",
+    "- For Python 3.11: 'quay.io/modh/ray:2.35.0-py311-cu121'\n",
+    "\n",
+    "If you prefer to use a custom Ray image that better suits your needs, you can specify it in the image field to override the default."
    ]
   },
   {
diff --git a/demo-notebooks/additional-demos/ray_job_client.ipynb b/demo-notebooks/additional-demos/ray_job_client.ipynb
@@ -41,8 +41,12 @@
    "metadata": {},
    "source": [
     "\n",
-    "NOTE: 'quay.io/modh/ray:2.35.0-py39-cu121' is the default image used by the CodeFlare SDK for creating a RayCluster resource. \n",
-    "If you have your own Ray image which suits your purposes, specify it in image field to override the default image."
+    "NOTE: The default images used by the CodeFlare SDK for creating a RayCluster resource depend on the installed Python version:\n",
+    "\n",
+    "- For Python 3.9: 'quay.io/modh/ray:2.35.0-py39-cu121'\n",
+    "- For Python 3.11: 'quay.io/modh/ray:2.35.0-py311-cu121'\n",
+    "\n",
+    "If you prefer to use a custom Ray image that better suits your needs, you can specify it in the image field to override the default."
    ]
   },
   {
diff --git a/demo-notebooks/guided-demos/0_basic_ray.ipynb b/demo-notebooks/guided-demos/0_basic_ray.ipynb
@@ -47,8 +47,12 @@
    "source": [
     "Here, we want to define our cluster by specifying the resources we require for our batch workload. Below, we define our cluster object (which generates a corresponding RayCluster).\n",
     "\n",
-    "NOTE: 'quay.io/modh/ray:2.35.0-py39-cu121' is the default image used by the CodeFlare SDK for creating a RayCluster resource. \n",
-    "If you have your own Ray image which suits your purposes, specify it in image field to override the default image."
+    "NOTE: The default images used by the CodeFlare SDK for creating a RayCluster resource depend on the installed Python version:\n",
+    "\n",
+    "- For Python 3.9: 'quay.io/modh/ray:2.35.0-py39-cu121'\n",
+    "- For Python 3.11: 'quay.io/modh/ray:2.35.0-py311-cu121'\n",
+    "\n",
+    "If you prefer to use a custom Ray image that better suits your needs, you can specify it in the image field to override the default."
    ]
   },
   {
diff --git a/demo-notebooks/guided-demos/1_cluster_job_client.ipynb b/demo-notebooks/guided-demos/1_cluster_job_client.ipynb
@@ -41,8 +41,12 @@
    "source": [
     "Here, we want to define our cluster by specifying the resources we require for our batch workload. Below, we define our cluster object (which generates a corresponding RayCluster).\n",
     "\n",
-    "NOTE: 'quay.io/modh/ray:2.35.0-py39-cu121' is the default image used by the CodeFlare SDK for creating a RayCluster resource. \n",
-    "If you have your own Ray image which suits your purposes, specify it in image field to override the default image."
+    "NOTE: The default images used by the CodeFlare SDK for creating a RayCluster resource depend on the installed Python version:\n",
+    "\n",
+    "- For Python 3.9: 'quay.io/modh/ray:2.35.0-py39-cu121'\n",
+    "- For Python 3.11: 'quay.io/modh/ray:2.35.0-py311-cu121'\n",
+    "\n",
+    "If you prefer to use a custom Ray image that better suits your needs, you can specify it in the image field to override the default."
    ]
   },
   {
diff --git a/demo-notebooks/guided-demos/2_basic_interactive.ipynb b/demo-notebooks/guided-demos/2_basic_interactive.ipynb
@@ -44,8 +44,12 @@
    "source": [
     "Once again, let's start by running through the same cluster setup as before:\n",
     "\n",
-    "NOTE: 'quay.io/modh/ray:2.35.0-py39-cu121' is the default image used by the CodeFlare SDK for creating a RayCluster resource. \n",
-    "If you have your own Ray image which suits your purposes, specify it in image field to override the default image."
+    "NOTE: The default images used by the CodeFlare SDK for creating a RayCluster resource depend on the installed Python version:\n",
+    "\n",
+    "- For Python 3.9: 'quay.io/modh/ray:2.35.0-py39-cu121'\n",
+    "- For Python 3.11: 'quay.io/modh/ray:2.35.0-py311-cu121'\n",
+    "\n",
+    "If you prefer to use a custom Ray image that better suits your needs, you can specify it in the image field to override the default."
    ]
   },
   {
diff --git a/demo-notebooks/guided-demos/3_widget_example.ipynb b/demo-notebooks/guided-demos/3_widget_example.ipynb
@@ -47,8 +47,12 @@
    "source": [
     "Here, we want to define our cluster by specifying the resources we require for our batch workload. Below, we define our cluster object (which generates a corresponding RayCluster).\n",
     "\n",
-    "NOTE: 'quay.io/modh/ray:2.35.0-py39-cu121' is the default image used by the CodeFlare SDK for creating a RayCluster resource. \n",
-    "If you have your own Ray image which suits your purposes, specify it in image field to override the default image."
+    "NOTE: The default images used by the CodeFlare SDK for creating a RayCluster resource depend on the installed Python version:\n",
+    "\n",
+    "- For Python 3.9: 'quay.io/modh/ray:2.35.0-py39-cu121'\n",
+    "- For Python 3.11: 'quay.io/modh/ray:2.35.0-py311-cu121'\n",
+    "\n",
+    "If you prefer to use a custom Ray image that better suits your needs, you can specify it in the image field to override the default."
    ]
   },
   {
diff --git a/demo-notebooks/guided-demos/notebook-ex-outputs/0_basic_ray.ipynb b/demo-notebooks/guided-demos/notebook-ex-outputs/0_basic_ray.ipynb
@@ -47,8 +47,12 @@
    "source": [
     "Here, we want to define our cluster by specifying the resources we require for our batch workload. Below, we define our cluster object (which generates a corresponding RayCluster).\n",
     "\n",
-    "NOTE: 'quay.io/modh/ray:2.35.0-py39-cu121' is the default image used by the CodeFlare SDK for creating a RayCluster resource. \n",
-    "If you have your own Ray image which suits your purposes, specify it in image field to override the default image."
+    "NOTE: The default images used by the CodeFlare SDK for creating a RayCluster resource depend on the installed Python version:\n",
+    "\n",
+    "- For Python 3.9: 'quay.io/modh/ray:2.35.0-py39-cu121'\n",
+    "- For Python 3.11: 'quay.io/modh/ray:2.35.0-py311-cu121'\n",
+    "\n",
+    "If you prefer to use a custom Ray image that better suits your needs, you can specify it in the image field to override the default."
    ]
   },
   {
diff --git a/demo-notebooks/guided-demos/notebook-ex-outputs/1_cluster_job_client.ipynb b/demo-notebooks/guided-demos/notebook-ex-outputs/1_cluster_job_client.ipynb
@@ -41,8 +41,12 @@
    "source": [
     "Here, we want to define our cluster by specifying the resources we require for our batch workload. Below, we define our cluster object (which generates a corresponding RayCluster).\n",
     "\n",
-    "NOTE: 'quay.io/modh/ray:2.35.0-py39-cu121' is the default image used by the CodeFlare SDK for creating a RayCluster resource. \n",
-    "If you have your own Ray image which suits your purposes, specify it in image field to override the default image."
+    "NOTE: The default images used by the CodeFlare SDK for creating a RayCluster resource depend on the installed Python version:\n",
+    "\n",
+    "- For Python 3.9: 'quay.io/modh/ray:2.35.0-py39-cu121'\n",
+    "- For Python 3.11: 'quay.io/modh/ray:2.35.0-py311-cu121'\n",
+    "\n",
+    "If you prefer to use a custom Ray image that better suits your needs, you can specify it in the image field to override the default."
    ]
   },
   {
diff --git a/demo-notebooks/guided-demos/notebook-ex-outputs/2_basic_interactive.ipynb b/demo-notebooks/guided-demos/notebook-ex-outputs/2_basic_interactive.ipynb
@@ -44,8 +44,12 @@
    "source": [
     "Once again, let's start by running through the same cluster setup as before:\n",
     "\n",
-    "NOTE: 'quay.io/modh/ray:2.35.0-py39-cu121' is the default image used by the CodeFlare SDK for creating a RayCluster resource. \n",
-    "If you have your own Ray image which suits your purposes, specify it in image field to override the default image."
+    "NOTE: The default images used by the CodeFlare SDK for creating a RayCluster resource depend on the installed Python version:\n",
+    "\n",
+    "- For Python 3.9: 'quay.io/modh/ray:2.35.0-py39-cu121'\n",
+    "- For Python 3.11: 'quay.io/modh/ray:2.35.0-py311-cu121'\n",
+    "\n",
+    "If you prefer to use a custom Ray image that better suits your needs, you can specify it in the image field to override the default."
    ]
   },
   {
diff --git a/demo-notebooks/guided-demos/preview_nbs/0_basic_ray.ipynb b/demo-notebooks/guided-demos/preview_nbs/0_basic_ray.ipynb
@@ -47,8 +47,12 @@
    "source": [
     "Here, we want to define our cluster by specifying the resources we require for our batch workload. Below, we define our cluster object (which generates a corresponding RayCluster).\n",
     "\n",
-    "NOTE: 'quay.io/modh/ray:2.35.0-py39-cu121' is the default image used by the CodeFlare SDK for creating a RayCluster resource. \n",
-    "If you have your own Ray image which suits your purposes, specify it in image field to override the default image."
+    "NOTE: The default images used by the CodeFlare SDK for creating a RayCluster resource depend on the installed Python version:\n",
+    "\n",
+    "- For Python 3.9: 'quay.io/modh/ray:2.35.0-py39-cu121'\n",
+    "- For Python 3.11: 'quay.io/modh/ray:2.35.0-py311-cu121'\n",
+    "\n",
+    "If you prefer to use a custom Ray image that better suits your needs, you can specify it in the image field to override the default."
    ]
   },
   {
diff --git a/demo-notebooks/guided-demos/preview_nbs/1_cluster_job_client.ipynb b/demo-notebooks/guided-demos/preview_nbs/1_cluster_job_client.ipynb
@@ -41,8 +41,12 @@
    "source": [
     "Here, we want to define our cluster by specifying the resources we require for our batch workload. Below, we define our cluster object (which generates a corresponding RayCluster).\n",
     "\n",
-    "NOTE: 'quay.io/modh/ray:2.35.0-py39-cu121' is the default image used by the CodeFlare SDK for creating a RayCluster resource. \n",
-    "If you have your own Ray image which suits your purposes, specify it in image field to override the default image."
+    "NOTE: The default images used by the CodeFlare SDK for creating a RayCluster resource depend on the installed Python version:\n",
+    "\n",
+    "- For Python 3.9: 'quay.io/modh/ray:2.35.0-py39-cu121'\n",
+    "- For Python 3.11: 'quay.io/modh/ray:2.35.0-py311-cu121'\n",
+    "\n",
+    "If you prefer to use a custom Ray image that better suits your needs, you can specify it in the image field to override the default."
    ]
   },
   {
diff --git a/demo-notebooks/guided-demos/preview_nbs/2_basic_interactive.ipynb b/demo-notebooks/guided-demos/preview_nbs/2_basic_interactive.ipynb
@@ -44,8 +44,12 @@
    "source": [
     "Once again, let's start by running through the same cluster setup as before:\n",
     "\n",
-    "NOTE: 'quay.io/modh/ray:2.35.0-py39-cu121' is the default image used by the CodeFlare SDK for creating a RayCluster resource. \n",
-    "If you have your own Ray image which suits your purposes, specify it in image field to override the default image."
+    "NOTE: The default images used by the CodeFlare SDK for creating a RayCluster resource depend on the installed Python version:\n",
+    "\n",
+    "- For Python 3.9: 'quay.io/modh/ray:2.35.0-py39-cu121'\n",
+    "- For Python 3.11: 'quay.io/modh/ray:2.35.0-py311-cu121'\n",
+    "\n",
+    "If you prefer to use a custom Ray image that better suits your needs, you can specify it in the image field to override the default."
    ]
   },
   {
diff --git a/src/codeflare_sdk/ray/cluster/generate_yaml.py b/src/codeflare_sdk/ray/cluster/generate_yaml.py
@@ -20,6 +20,7 @@
 import json
 import sys
 import typing
+import warnings
 import yaml
 import os
 import uuid
@@ -34,7 +35,7 @@
 
 SUPPORTED_PYTHON_VERSIONS = {
     "3.9": "quay.io/modh/ray@sha256:0d715f92570a2997381b7cafc0e224cfa25323f18b9545acfd23bc2b71576d06",
-    "3.11": "quay.io/modh/ray:2.35.0-py311-cu121",
+    "3.11": "quay.io/modh/ray@sha256:db667df1bc437a7b0965e8031e905d3ab04b86390d764d120e05ea5a5c18d1b4",
 }
 
 
@@ -96,12 +97,11 @@ def update_image(spec, image):
     containers = spec.get("containers")
     if not image:
         python_version = f"{sys.version_info.major}.{sys.version_info.minor}"
-        try:
-            if python_version in SUPPORTED_PYTHON_VERSIONS:
-                image = SUPPORTED_PYTHON_VERSIONS[python_version]
-        except Exception:  # pragma: no cover
-            print(
-                f"Python version '{python_version}' is not supported. Only {', '.join(SUPPORTED_PYTHON_VERSIONS.keys())} are supported."
+        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())}."
             )
     for container in containers:
         container["image"] = image
diff --git a/src/codeflare_sdk/ray/cluster/test_generate_yaml.py b/src/codeflare_sdk/ray/cluster/test_generate_yaml.py
@@ -11,7 +11,9 @@
 # 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.ray.cluster.generate_yaml import gen_names
+from collections import namedtuple
+import sys
+from .generate_yaml import gen_names, update_image
 import uuid
 
 
@@ -32,3 +34,37 @@ def test_gen_names_without_name(mocker):
     appwrapper_name, cluster_name = gen_names(None)
     assert appwrapper_name.startswith("appwrapper-")
     assert cluster_name.startswith("cluster-")
+
+
+def test_update_image_without_supported_python_version(mocker):
+    # Mock SUPPORTED_PYTHON_VERSIONS
+    mocker.patch.dict(
+        "codeflare_sdk.ray.cluster.generate_yaml.SUPPORTED_PYTHON_VERSIONS",
+        {
+            "3.9": "ray-py3.9",
+            "3.11": "ray-py3.11",
+        },
+    )
+
+    # Create a namedtuple to mock sys.version_info
+    VersionInfo = namedtuple(
+        "version_info", ["major", "minor", "micro", "releaselevel", "serial"]
+    )
+    mocker.patch.object(sys, "version_info", VersionInfo(3, 8, 0, "final", 0))
+
+    # Mock warnings.warn to check if it gets called
+    warn_mock = mocker.patch("warnings.warn")
+
+    # Create a sample spec
+    spec = {"containers": [{"image": None}]}
+
+    # Call the update_image function with no image provided
+    update_image(spec, None)
+
+    # 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.9, 3.11."
+    )
+
+    # Assert that no image was set in the containers since the Python version is not supported
+    assert spec["containers"][0]["image"] is None
