[release-automation] More flexibility on the pipeline (ray-project#45254

)

- Only upload `linux_x86_64`, `mac x86_64`, and `mac arm64` wheels
unless it's a full platform release.
- Each of the wheel sanity check step has its own block step.
- Uploading wheels to PyPI no longer requires all sanity checks to pass.
This is up for release manager to judge whether the wheels are ready.
- Add key for each of the step. This is to support automation that
checks for BK job status (it's better to retrieve jobs from BK and
identify with key than job label)

---------

Signed-off-by: khluu <[email protected]>

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

@@ -43,6 +43,7 @@ steps:
  RAYCI_SCHEDULE: "nightly"
 
  - label: "Trigger Postmerge MacOS test"
+ key: trigger-postmerge-macos
  trigger: "postmerge-macos"
  depends_on: check-commit-hash
  build:
@@ -90,6 +91,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:
@@ -101,6 +103,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(".")[1]
+ full_platform = True if int(minor_version) % 10 == 0 else False
+ wheels = _get_wheel_names(ray_version=ray_version, full_platform=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,11 @@
 
 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(PLATFORMS) * len(RAY_TYPES)
+ 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
@@ -44,7 +47,32 @@ 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 FULL_PLATFORMS
+
+
+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
+ (
+ 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 PARTIAL_PLATFORMS
 
 
 def test_check_downloaded_wheels():
@@ -134,10 +162,46 @@ 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=ray_version, full_platform=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=ray_version, full_platform=False
+ )
  assert mock_download_wheel.call_count == len(SAMPLE_WHEELS)
  for i, call_args in enumerate(mock_download_wheel.call_args_list):
  assert (