[release test] add support for llm based images (#50593)

so that we can use the image for running related release tests.

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

ci/ray_ci/oss_config.yaml

@@ -2,6 +2,7 @@ release_byod:
   ray_ecr: rayproject
   ray_cr_repo: ray
   ray_ml_cr_repo: ray-ml
+  ray_llm_cr_repo: ray-llm
   byod_ecr: 029272617770.dkr.ecr.us-west-2.amazonaws.com
   aws_cr: 029272617770.dkr.ecr.us-west-2.amazonaws.com
   gcp_cr: us-west1-docker.pkg.dev/anyscale-oss-ci

release/ray_release/byod/build.py

@@ -23,6 +23,7 @@
 BASE_IMAGE_WAIT_DURATION = 30
 RELEASE_BYOD_DIR = os.path.join(RELEASE_PACKAGE_DIR, "ray_release/byod")
 REQUIREMENTS_BYOD = "requirements_byod"
+REQUIREMENTS_LLM_BYOD = "requirements_llm_byod"
 REQUIREMENTS_ML_BYOD = "requirements_ml_byod"
 
 
@@ -124,6 +125,8 @@ def build_anyscale_base_byod_images(tests: List[Test]) -> None:
             py_version = test.get_python_version()
             if test.use_byod_ml_image():
                 byod_requirements = f"{REQUIREMENTS_ML_BYOD}_{py_version}.txt"
+            elif test.use_byod_llm_image():
+                byod_requirements = f"{REQUIREMENTS_LLM_BYOD}_{py_version}.txt"
             else:
                 byod_requirements = f"{REQUIREMENTS_BYOD}_{py_version}.txt"
 

release/ray_release/byod/requirements_llm_byod_3.11.txt

@@ -0,0 +1 @@
+# Just a stub for now; please keep this empty.

release/ray_release/configs/global_config.py

@@ -9,6 +9,7 @@ class GlobalConfig(TypedDict):
     byod_ray_ecr: str
     byod_ray_cr_repo: str
     byod_ray_ml_cr_repo: str
+    byod_ray_llm_cr_repo: str
     byod_ecr: str
     byod_aws_cr: str
     byod_gcp_cr: str
@@ -58,6 +59,10 @@ def _init_global_config(config_file: str):
             config_content.get("byod", {}).get("ray_ml_cr_repo")
             or config_content.get("release_byod", {}).get("ray_ml_cr_repo")
         ),
+        byod_ray_llm_cr_repo=(
+            config_content.get("byod", {}).get("ray_llm_cr_repo")
+            or config_content.get("release_byod", {}).get("ray_llm_cr_repo")
+        ),
         byod_ecr=(
             config_content.get("byod", {}).get("byod_ecr")
             or config_content.get("release_byod", {}).get("byod_ecr")

release/ray_release/test.py

@@ -37,6 +37,7 @@
 )
 DATAPLANE_ECR_REPO = "anyscale/ray"
 DATAPLANE_ECR_ML_REPO = "anyscale/ray-ml"
+DATAPLANE_ECR_LLM_REPO = "anyscale/ray-llm"
 
 MACOS_TEST_PREFIX = "darwin:"
 LINUX_TEST_PREFIX = "linux:"
@@ -393,12 +394,21 @@ def get_bisect_daily_rate_limit(self) -> int:
             return WINDOWS_BISECT_DAILY_RATE_LIMIT
         return BISECT_DAILY_RATE_LIMIT
 
-    def get_byod_type(self) -> Optional[str]:
+    def get_byod_type(self) -> str:
         """
         Returns the type of the BYOD cluster.
         """
         return self["cluster"]["byod"].get("type", "cpu")
 
+    def get_tag_suffix(self) -> str:
+        """
+        Returns the tag suffix for the BYOD image.
+        """
+        byod_type = self.get_byod_type()
+        if byod_type.startswith("llm-"):
+            return byod_type[len("llm-") :]
+        return byod_type
+
     def get_byod_post_build_script(self) -> Optional[str]:
         """
         Returns the post-build script for the BYOD cluster.
@@ -528,7 +538,7 @@ def get_byod_base_image_tag(self) -> str:
             release_name = branch[len("releases/") :]
             ray_version = f"{release_name}.{ray_version}"
         python_version = f"py{self.get_python_version().replace('.',   '')}"
-        return f"{ray_version}-{python_version}-{self.get_byod_type()}"
+        return f"{ray_version}-{python_version}-{self.get_tag_suffix()}"
 
     def get_byod_image_tag(self) -> str:
         """
@@ -545,12 +555,17 @@ def use_byod_ml_image(self) -> bool:
         """Returns whether to use the ML image for this test."""
         return self.get_byod_type() == "gpu"
 
+    def use_byod_llm_image(self) -> bool:
+        return self.get_byod_type().startswith("llm-")
+
     def get_byod_repo(self) -> str:
         """
         Returns the byod repo to use for this test.
         """
         if self.use_byod_ml_image():
             return DATAPLANE_ECR_ML_REPO
+        if self.use_byod_llm_image():
+            return DATAPLANE_ECR_LLM_REPO
         return DATAPLANE_ECR_REPO
 
     def get_byod_ecr(self) -> str:
@@ -572,6 +587,8 @@ def get_ray_image(self) -> str:
         repo = self.get_byod_repo()
         if repo == DATAPLANE_ECR_REPO:
             repo_name = config["byod_ray_cr_repo"]
+        elif repo == DATAPLANE_ECR_LLM_REPO:
+            repo_name = config["byod_ray_llm_cr_repo"]
         elif repo == DATAPLANE_ECR_ML_REPO:
             repo_name = config["byod_ray_ml_cr_repo"]
         else:

release/ray_release/tests/test_global_config.py

@@ -15,6 +15,7 @@
   ray_cr_repo: ray
 release_byod:
   ray_ml_cr_repo: ray-ml
+  ray_llm_cr_repo: ray-llm
   byod_ecr: 029272617770.dkr.ecr.us-west-2.amazonaws.com
   aws_cr: 029272617770.dkr.ecr.us-west-2.amazonaws.com
   gcp_cr: us-west1-docker.pkg.dev/anyscale-oss-ci
@@ -52,6 +53,9 @@ def test_init_global_config() -> None:
         )
         assert config["state_machine_pr_aws_bucket"] == "ray-ci-pr-results"
         assert config["state_machine_disabled"] is True
+        assert config["byod_ray_cr_repo"] == "ray"
+        assert config["byod_ray_ml_cr_repo"] == "ray-ml"
+        assert config["byod_ray_llm_cr_repo"] == "ray-llm"
 
 
 if __name__ == "__main__":

release/ray_release/tests/test_test.py

@@ -113,6 +113,19 @@ def test_get_ray_image():
         ).get_ray_image()
         == "rayproject/ray-ml:123456-py39-gpu"
     )
+    assert (
+        _stub_test(
+            {
+                "python": "3.11",
+                "cluster": {
+                    "byod": {
+                        "type": "llm-cu124",
+                    }
+                },
+            }
+        ).get_ray_image()
+        == "rayproject/ray-llm:123456-py311-cu124"
+    )
     os.environ["BUILDKITE_BRANCH"] = "releases/1.0.0"
     assert (
         _stub_test({"cluster": {"byod": {}}}).get_ray_image()