[release-automation] More flexibility on the pipeline

.buildkite/release-automation/pre_release.rayci.yml

@@ -45,6 +45,7 @@ steps:
  RAYCI_SCHEDULE: "nightly"
 
  - label: "Trigger Postmerge MacOS test"
+ key: trigger-postmerge-macos
  trigger: "postmerge-macos"
  depends_on: check-commit-hash
  build:
@@ -92,6 +93,7 @@ steps:
  depends_on: check-commit-hash
 
  - label: "Trigger Release weekly test"
+ key: trigger-release-weekly
  trigger: "release"
  depends_on: block-release-weekly-test
  build:
@@ -103,6 +105,7 @@ steps:
  RELEASE_FREQUENCY: "weekly"
 
  - label: "Check Ray commit in {{matrix}} nightly images"
+ key: check-ray-commit
  if: build.branch !~ /^releases\// && build.env("RAYCI_WEEKLY_RELEASE_NIGHTLY") == "1"
  depends_on: trigger-postmerge-nightly
  allow_dependency_failure: true

.buildkite/release-automation/wheels.rayci.yml

@@ -1,6 +1,7 @@
 group: Upload & Validate wheels
 steps:
  - block: "Upload wheels from S3 to TestPyPI"
+ key: block-upload-wheels-testpypi
  depends_on: 
  - forge
 
@@ -14,10 +15,10 @@ steps:
  --commit_hash="$BUILDKITE_COMMIT"
  --pypi_env=test
 
- - block: "Download & validate Ray wheels from TestPyPI"
+ - block: "Download & validate Ray wheels from TestPyPI Linux x86_64"
+ key: block-validate-linux-x86_64-wheels
  depends_on:
  - forge
- - forge_arm64
 
  - label: "Linux x86_64 Python {{matrix}}"
  key: validate-linux-x86_64-wheels
@@ -31,6 +32,11 @@ steps:
  - "3.10"
  - "3.11"
 
+ - block: "Download & validate Ray wheels from TestPyPI Linux arm64"
+ key: block-validate-linux-arm64-wheels
+ depends_on:
+ - forge-arm64
+
  - label: "Linux arm64 Python {{matrix}}"
  key: validate-linux-arm64-wheels
  instance_type: medium-arm64
@@ -44,6 +50,9 @@ steps:
  - "3.10"
  - "3.11"
 
+ - block: "Download & validate Ray wheels from TestPyPI Mac"
+ key: block-validate-macos-wheels
+
  - label: "MacOS x86_64"
  key: validate-macos-x86_64-wheels
  job_env: MACOS
@@ -59,11 +68,7 @@ steps:
  - ./.buildkite/release-automation/verify-macos-wheels.sh arm64
 
  - block: "Upload wheels to PyPI"
- depends_on: 
- - validate-linux-x86-64-wheels
- - validate-linux-arm64-wheels
- - validate-macos-x86-64-wheels
- - validate-macos-arm64-wheels
+ key: block-upload-wheels-pypi
 
  - label: "Upload wheels to PyPI"
  key: upload-wheels-pypi

ci/ray_ci/automation/ray_wheels_lib.py

@@ -1,19 +1,24 @@
 import boto3
-from typing import List
+from typing import List, Optional
 import os
 
 from ci.ray_ci.utils import logger
 
 bazel_workspace_dir = os.environ.get("BUILD_WORKSPACE_DIRECTORY", "")
 
 PYTHON_VERSIONS = ["cp39-cp39", "cp310-cp310", "cp311-cp311"]
-PLATFORMS = [
+FULL_PLATFORMS = [
  "manylinux2014_x86_64",
  "manylinux2014_aarch64",
  "macosx_10_15_x86_64",
  "macosx_11_0_arm64",
  "win_amd64",
 ]
+PARTIAL_PLATFORMS = [
+ "manylinux2014_x86_64",
+ "macosx_10_15_x86_64",
+ "macosx_11_0_arm64",
+]
 RAY_TYPES = ["ray", "ray_cpp"]
 
 
@@ -41,11 +46,14 @@ def _check_downloaded_wheels(directory_path: str, wheels: List[str]) -> None:
  )
 
 
-def _get_wheel_names(ray_version: str) -> List[str]:
+def _get_wheel_names(
+ ray_version: str, full_platform: Optional[bool] = False
+) -> List[str]:
  """List all wheel names for the given ray version."""
  wheel_names = []
+ platforms = FULL_PLATFORMS if full_platform else PARTIAL_PLATFORMS
  for python_version in PYTHON_VERSIONS:
- for platform in PLATFORMS:
+ for platform in platforms:
  for ray_type in RAY_TYPES:
  wheel_name = f"{ray_type}-{ray_version}-{python_version}-{platform}"
  wheel_names.append(wheel_name)
@@ -71,7 +79,9 @@ def download_wheel_from_s3(key: str, directory_path: str) -> None:
 
 
 def download_ray_wheels_from_s3(
- commit_hash: str, ray_version: str, directory_path: str
+ commit_hash: str,
+ ray_version: str,
+ directory_path: str,
 ) -> None:
  """
  Download Ray wheels from S3 to the given directory.
@@ -82,8 +92,9 @@ def download_ray_wheels_from_s3(
  directory_path: The directory to download the wheels to.
  """
  full_directory_path = os.path.join(bazel_workspace_dir, directory_path)
-
- wheels = _get_wheel_names(ray_version)
+ _, minor_version, _ = ray_version.split(".")
+ full_platform = True if int(minor_version) % 10 == 0 else False
+ wheels = _get_wheel_names(ray_version, full_platform)
  for wheel in wheels:
  s3_key = f"releases/{ray_version}/{commit_hash}/{wheel}.whl"
  download_wheel_from_s3(s3_key, full_directory_path)

ci/ray_ci/automation/test_ray_wheels_lib.py

@@ -11,7 +11,8 @@
  download_ray_wheels_from_s3,
  _check_downloaded_wheels,
  PYTHON_VERSIONS,
- PLATFORMS,
+ FULL_PLATFORMS,
+ PARTIAL_PLATFORMS,
  RAY_TYPES,
 )
 
@@ -26,9 +27,36 @@
 
 def test_get_wheel_names():
  ray_version = "1.0.0"
- wheel_names = _get_wheel_names(ray_version)
+ wheel_names = _get_wheel_names(ray_version, full_platform=True)
+
+ assert len(wheel_names) == len(PYTHON_VERSIONS) * len(FULL_PLATFORMS) * len(
+ RAY_TYPES
+ )
+
+ for wheel_name in wheel_names:
+ assert len(wheel_name.split("-")) == 5
+ (
+ ray_type,
+ ray_version,
+ python_version,
+ python_version2,
+ platform,
+ ) = wheel_name.split("-")
+ platform = platform.split(".")[0] # Remove the .whl suffix
+
+ assert ray_type in RAY_TYPES
+ assert ray_version == ray_version
+ assert f"{python_version}-{python_version2}" in PYTHON_VERSIONS
+ assert platform in FULL_PLATFORMS
 
- assert len(wheel_names) == len(PYTHON_VERSIONS) * len(PLATFORMS) * len(RAY_TYPES)
+
+def test_get_wheel_names_partial_platform():
+ ray_version = "1.1.0"
+ wheel_names = _get_wheel_names(ray_version, full_platform=False)
+
+ assert len(wheel_names) == len(PYTHON_VERSIONS) * len(PARTIAL_PLATFORMS) * len(
+ RAY_TYPES
+ )
 
  for wheel_name in wheel_names:
  assert len(wheel_name.split("-")) == 5
@@ -44,7 +72,7 @@ def test_get_wheel_names():
  assert ray_type in RAY_TYPES
  assert ray_version == ray_version
  assert f"{python_version}-{python_version2}" in PYTHON_VERSIONS
- assert platform in PLATFORMS
+ assert platform in PARTIAL_PLATFORMS
 
 
 def test_check_downloaded_wheels():
@@ -134,10 +162,42 @@ def test_download_ray_wheels_from_s3(
 
  with tempfile.TemporaryDirectory() as tmp_dir:
  download_ray_wheels_from_s3(
- commit_hash=commit_hash, ray_version=ray_version, directory_path=tmp_dir
+ commit_hash=commit_hash,
+ ray_version=ray_version,
+ directory_path=tmp_dir,
+ )
+
+ mock_get_wheel_names.assert_called_with(ray_version, True)
+ 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"releases/{ray_version}/{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")
+def test_download_ray_wheels_from_s3_partial_platform(
+ mock_get_wheel_names, mock_check_wheels, mock_download_wheel
+):
+ commit_hash = "1234567"
+ ray_version = "1.1.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,
  )
 
- mock_get_wheel_names.assert_called_with(ray_version)
+ mock_get_wheel_names.assert_called_with(ray_version, False)
  assert mock_download_wheel.call_count == len(SAMPLE_WHEELS)
  for i, call_args in enumerate(mock_download_wheel.call_args_list):
  assert (