diff --git a/release/ray_release/byod/build.py b/release/ray_release/byod/build.py
index 841bd2ad96c1f..2816bfd12b577 100644
--- a/release/ray_release/byod/build.py
+++ b/release/ray_release/byod/build.py
@@ -25,7 +25,6 @@
 RELEASE_BYOD_DIR = os.path.join(RELEASE_PACKAGE_DIR, "ray_release/byod")
 REQUIREMENTS_BYOD = "requirements_byod"
 REQUIREMENTS_ML_BYOD = "requirements_ml_byod"
-PYTHON_VERSION = "3.8"
 
 
 def build_champagne_image(
@@ -124,11 +123,12 @@ def build_anyscale_base_byod_images(tests: List[Test]) -> None:
         and int(time.time()) - start < BASE_IMAGE_WAIT_TIMEOUT
     ):
         for byod_image, test in to_be_built.items():
-            byod_requirements = (
-                f"{REQUIREMENTS_BYOD}_{test.get('python', PYTHON_VERSION)}.txt"
-                if test.get_byod_type() == "cpu"
-                else f"{REQUIREMENTS_ML_BYOD}_{test.get('python', PYTHON_VERSION)}.txt"
-            )
+            py_version = test.get_python_version()
+            if test.use_byod_ml_image():
+                byod_requirements = f"{REQUIREMENTS_ML_BYOD}_{py_version}.txt"
+            else:
+                byod_requirements = f"{REQUIREMENTS_BYOD}_{py_version}.txt"
+
             if _byod_image_exist(test):
                 logger.info(f"Image {byod_image} already exists")
                 built.add(byod_image)
@@ -210,7 +210,7 @@ def _validate_and_push(byod_image: str) -> None:
         assert (
             docker_ray_commit == expected_ray_commit
         ), f"Expected ray commit {expected_ray_commit}, found {docker_ray_commit}"
-    logger.info(f"Pushing image to ECR: {byod_image}")
+    logger.info(f"Pushing image to registry: {byod_image}")
     subprocess.check_call(
         ["docker", "push", byod_image],
         stdout=sys.stderr,
diff --git a/release/ray_release/test.py b/release/ray_release/test.py
index 8aca16a428cfb..62df819c6318d 100644
--- a/release/ray_release/test.py
+++ b/release/ray_release/test.py
@@ -251,7 +251,7 @@ def get_byod_image_tag(self) -> str:
         }
         return f"{self.get_byod_base_image_tag()}-{dict_hash(custom_info)}"
 
-    def _use_byod_ml_image(self) -> bool:
+    def use_byod_ml_image(self) -> bool:
         """Returns whether to use the ML image for this test."""
         return self.get_byod_type() == "gpu"
 
@@ -259,7 +259,7 @@ def get_byod_repo(self) -> str:
         """
         Returns the byod repo to use for this test.
         """
-        if self._use_byod_ml_image():
+        if self.use_byod_ml_image():
             return DATAPLANE_ECR_ML_REPO
         return DATAPLANE_ECR_REPO