diff --git a/release/ray_release/command_runner/client_runner.py b/release/ray_release/command_runner/client_runner.py
index 53cade0b7..b6407f001 100644
--- a/release/ray_release/command_runner/client_runner.py
+++ b/release/ray_release/command_runner/client_runner.py
@@ -9,8 +9,6 @@ import time
 from collections import deque
 from typing import Optional, Dict, Any
 
-import ray
-
 from ray_release.anyscale_util import LAST_LOGS_LENGTH
 
 from ray_release.cluster_manager.cluster_manager import ClusterManager
@@ -67,6 +65,8 @@ class ClientRunner(CommandRunner):
             raise LocalEnvSetupError(f"Error setting up local environment: {e}") from e
 
     def wait_for_nodes(self, num_nodes: int, timeout: float = 900):
+        import ray
+
         ray_address = self.cluster_manager.get_cluster_address()
         try:
             if ray.is_initialized:
diff --git a/release/run_release_test.sh b/release/run_release_test.sh
index 804e02835..d7398b091 100755
--- a/release/run_release_test.sh
+++ b/release/run_release_test.sh
@@ -31,11 +31,35 @@ RAY_TEST_REPO=${RAY_TEST_REPO-https://github.com/ray-project/ray.git}
 RAY_TEST_BRANCH=${RAY_TEST_BRANCH-master}
 RELEASE_RESULTS_DIR=${RELEASE_RESULTS_DIR-/tmp/artifacts}
 
+# This is not a great idea if your OS is different to the one
+# used in the product clusters. However, we need this in CI as reloading
+# Ray within the python process does not work for protobuf changes.
+INSTALL_MATCHING_RAY=${BUILDKITE-false}
+
 export RAY_TEST_REPO RAY_TEST_BRANCH RELEASE_RESULTS_DIR
 
 if [ -z "${NO_INSTALL}" ]; then
   pip install -q -r requirements.txt
   pip install -q -U boto3 botocore
+
+  if [ "${INSTALL_MATCHING_RAY-false}" == "true" ]; then
+    # Find ray-wheels parameter and install locally
+    i=1
+    for arg in "$@"; do
+      j=$((i+1))
+      if [ "$arg" == "--ray-wheels" ]; then
+        PARSED_RAY_WHEELS="${!j}"
+      fi
+      i=$j
+    done
+
+    if [ -n "${PARSED_RAY_WHEELS}" ]; then
+      echo "Installing Ray wheels locally: ${PARSED_RAY_WHEELS}"
+      pip install -U --force-reinstall "${PARSED_RAY_WHEELS}"
+    else
+      echo "Warning: No Ray wheels found to install locally"
+    fi
+  fi
 fi
 
 if [ -z "${NO_CLONE}" ]; then