From 8a9512bf629bcdd947d2b99c55b5c47e73cebfad Mon Sep 17 00:00:00 2001
From: Kai Fricke <krfricke@users.noreply.github.com>
Date: Thu, 2 Jun 2022 11:32:51 +0200
Subject: [PATCH] [ci/release] Install local wheels in release test shell
 script (#25227)

We're currently installing matching wheels on the fly in the python script for Ray client tests. However, we can't reload modules with changed protobuf configurations, and thus can't reload ray completely. Since the `anyscale` pacakge depends on Ray, this effectively prevents us from installing matching wheels within the python script.

There are a few possible solutions to this. First, we could separate out the local environment preparation from the test running - this will duplicate some logic and is thus a bit more involved, but should be considered in the future. For now, we adjust the `run_release_tests.sh` shell script to install any passed `--ray-wheels` wheels locally. We only do this in CI instances per default as these wheels will not be compatible with e.g. MacOS.

Link to successful build: https://buildkite.com/ray-project/release-tests-branch/builds/619#_
---
 .../command_runner/client_runner.py           |  4 ++--
 release/run_release_test.sh                   | 24 +++++++++++++++++++
 2 files changed, 26 insertions(+), 2 deletions(-)

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