[release auto] support uploading wheels from arbitrary branch (ray-project#58344)

so that we can perform surgeries in the last minute if required.

Signed-off-by: Lonnie Liu <[email protected]>

Author: aslonnie

ci/ray_ci/automation/ray_wheels_lib.py

@@ -1,5 +1,5 @@
 import os
-from typing import List
+from typing import List, Optional
 
 import boto3
 
@@ -82,6 +82,7 @@ def download_ray_wheels_from_s3(
     commit_hash: str,
     ray_version: str,
     directory_path: str,
+    branch: Optional[str] = None,
 ) -> None:
     """
     Download Ray wheels from S3 to the given directory.
@@ -93,8 +94,10 @@ def download_ray_wheels_from_s3(
     """
     full_directory_path = os.path.join(bazel_workspace_dir, directory_path)
     wheels = _get_wheel_names(ray_version=ray_version)
+    if not branch:
+        branch = f"releases/{ray_version}"
     for wheel in wheels:
-        s3_key = f"releases/{ray_version}/{commit_hash}/{wheel}.whl"
+        s3_key = f"{branch}/{commit_hash}/{wheel}.whl"
         download_wheel_from_s3(s3_key, full_directory_path)
 
     _check_downloaded_wheels(full_directory_path, wheels)

ci/ray_ci/automation/test_ray_wheels_lib.py

@@ -153,6 +153,36 @@ def test_download_ray_wheels_from_s3(
         mock_check_wheels.assert_called_with(tmp_dir, SAMPLE_WHEELS)
 
 
+@mock.patch("ci.ray_ci.automation.ray_wheels_lib.download_wheel_from_s3")
+@mock.patch("ci.ray_ci.automation.ray_wheels_lib._check_downloaded_wheels")
+@mock.patch("ci.ray_ci.automation.ray_wheels_lib._get_wheel_names")
+def test_download_ray_wheels_from_s3_with_branch(
+    mock_get_wheel_names, mock_check_wheels, mock_download_wheel
+):
+    commit_hash = "1234567"
+    ray_version = "1.0.0"
+
+    mock_get_wheel_names.return_value = SAMPLE_WHEELS
+
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        download_ray_wheels_from_s3(
+            commit_hash=commit_hash,
+            ray_version=ray_version,
+            directory_path=tmp_dir,
+            branch="custom_branch",
+        )
+
+        mock_get_wheel_names.assert_called_with(ray_version=ray_version)
+        assert mock_download_wheel.call_count == len(SAMPLE_WHEELS)
+        for i, call_args in enumerate(mock_download_wheel.call_args_list):
+            assert (
+                call_args[0][0] == f"custom_branch/{commit_hash}/{SAMPLE_WHEELS[i]}.whl"
+            )
+            assert call_args[0][1] == tmp_dir
+
+        mock_check_wheels.assert_called_with(tmp_dir, SAMPLE_WHEELS)
+
+
 @mock.patch("ci.ray_ci.automation.ray_wheels_lib.download_wheel_from_s3")
 @mock.patch("ci.ray_ci.automation.ray_wheels_lib._check_downloaded_wheels")
 @mock.patch("ci.ray_ci.automation.ray_wheels_lib._get_wheel_names")

ci/ray_ci/automation/upload_wheels_pypi.py

@@ -14,15 +14,21 @@
 @click.option("--ray_version", required=True, type=str)
 @click.option("--commit_hash", required=True, type=str)
 @click.option("--pypi_env", required=True, type=click.Choice(["test", "prod"]))
+@click.option("--branch", required=False, type=str)
 @click.option("--build_tag", required=False, type=str)
 def main(
-    ray_version: str, commit_hash: str, pypi_env: str, build_tag: Optional[str] = None
+    ray_version: str,
+    commit_hash: str,
+    pypi_env: str,
+    branch: Optional[str] = None,
+    build_tag: Optional[str] = None,
 ):
     with tempfile.TemporaryDirectory() as temp_dir:
         download_ray_wheels_from_s3(
             commit_hash=commit_hash,
             ray_version=ray_version,
             directory_path=temp_dir,
+            branch=branch,
         )
         if build_tag:
             add_build_tag_to_wheels(directory_path=temp_dir, build_tag=build_tag)