[ci] add missing python versions

and add version/platform/arch assertions. this makes sure that
we remember to add/change the list

Signed-off-by: Lonnie Liu <lonnie@anyscale.com>

ci/ray_ci/docker_container.py

@@ -20,8 +20,8 @@
 ]
 GPU_PLATFORM = "cu12.1.1-cudnn8"
 
-PYTHON_VERSIONS_RAY = ["3.9", "3.10", "3.11"]
-PYTHON_VERSIONS_RAY_ML = ["3.9", "3.10"]
+PYTHON_VERSIONS_RAY = ["3.9", "3.10", "3.11", "3.12"]
+PYTHON_VERSIONS_RAY_ML = ["3.9", "3.10", "3.11"]
 ARCHITECTURES_RAY = ["x86_64", "aarch64"]
 ARCHITECTURES_RAY_ML = ["x86_64"]
 
@@ -46,6 +46,15 @@ def __init__(
         upload: bool = False,
     ) -> None:
         assert "RAYCI_CHECKOUT_DIR" in os.environ, "RAYCI_CHECKOUT_DIR not set"
+
+        assert python_version in PYTHON_VERSIONS_RAY
+        assert platform in PLATFORMS_RAY
+        assert architecture in ARCHITECTURES_RAY
+        if image_type == RayType.RAY_ML:
+            assert python_version in PYTHON_VERSIONS_RAY_ML
+            assert platform in PLATFORMS_RAY_ML
+            assert architecture in ARCHITECTURES_RAY_ML
+
         rayci_checkout_dir = os.environ["RAYCI_CHECKOUT_DIR"]
         self.python_version = python_version
         self.platform = platform
@@ -122,13 +131,13 @@ def _get_image_tags(self, external: bool = False) -> List[str]:
         versions = self._get_image_version_tags(external)
 
         platforms = [self.get_platform_tag()]
-        if self.platform == "cpu" and self.image_type == "ray":
+        if self.platform == "cpu" and self.image_type == RayType.RAY:
             # no tag is alias to cpu for ray image
             platforms.append("")
         elif self.platform == GPU_PLATFORM:
             # gpu is alias to cu118 for ray image
             platforms.append("-gpu")
-            if self.image_type == "ray-ml":
+            if self.image_type == RayType.RAY_ML:
                 # no tag is alias to gpu for ray-ml image
                 platforms.append("")
 

ci/ray_ci/test_ray_docker_container.py

@@ -7,6 +7,7 @@
 
 from ci.ray_ci.builder_container import DEFAULT_PYTHON_VERSION
 from ci.ray_ci.container import _DOCKER_ECR_REPO
+from ci.ray_ci.docker_container import GPU_PLATFORM
 from ci.ray_ci.ray_docker_container import RayDockerContainer
 from ci.ray_ci.test_base import RayCITestBase
 from ci.ray_ci.utils import RAY_VERSION
@@ -203,8 +204,8 @@ def test_canonical_tag(self) -> None:
         container = RayDockerContainer(v, "cpu", "ray", "aarch64")
         assert container._get_canonical_tag() == f"{sha}-{pv}-cpu-aarch64"
 
-        container = RayDockerContainer(v, "cu11.8.0-cudnn8", "ray-ml")
-        assert container._get_canonical_tag() == f"{sha}-{pv}-cu118"
+        container = RayDockerContainer(v, GPU_PLATFORM, "ray-ml")
+        assert container._get_canonical_tag() == f"{sha}-{pv}-cu121"
 
         with mock.patch.dict(os.environ, {"BUILDKITE_BRANCH": "releases/1.0.0"}):
             container = RayDockerContainer(v, "cpu", "ray")