[release] Init step for custom BYOD image build (ray-project#55398)

- Add `custom_byod_build_init` to scan the tests and list out all custom
images to build, then create a yaml file (that will be included when
rayci runs) to launch the build jobs.
- Add `custom_byod_build` as a python script that the jobs can call to
build & push the images

This needs to be merged after
ray-project#55397

---------

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

Author: khluu

ci/ray_ci/oss_config.yaml

@@ -4,6 +4,7 @@ release_byod:
   ray_ml_cr_repo: ray-ml
   ray_llm_cr_repo: ray-llm
   byod_ecr: 029272617770.dkr.ecr.us-west-2.amazonaws.com
+  byod_ecr_region: us-west-2
   aws_cr: 029272617770.dkr.ecr.us-west-2.amazonaws.com
   gcp_cr: us-west1-docker.pkg.dev/anyscale-oss-ci
   aws2gce_credentials: release/aws2gce_iam.json

release/BUILD.bazel

@@ -456,6 +456,24 @@ py_test(
     ],
 )
 
+py_test(
+    name = "test_custom_byod_build_init_helper",
+    size = "small",
+    srcs = ["ray_release/tests/test_custom_byod_build_init_helper.py"],
+    data = [
+        "ray_release/configs/oss_config.yaml",
+    ],
+    exec_compatible_with = ["//:hermetic_python"],
+    tags = [
+        "release_unit",
+        "team:ci",
+    ],
+    deps = [
+        ":ray_release",
+        bk_require("pytest"),
+    ],
+)
+
 py_test(
     name = "test_cluster_manager",
     size = "small",
@@ -729,3 +747,12 @@ py_binary(
         ":ray_release",
     ],
 )
+
+py_binary(
+    name = "custom_byod_build_init",
+    srcs = ["ray_release/scripts/custom_byod_build_init.py"],
+    exec_compatible_with = ["//:hermetic_python"],
+    deps = [
+        ":ray_release",
+    ],
+)

release/ray_release/configs/global_config.py

@@ -11,6 +11,7 @@ class GlobalConfig(TypedDict):
     byod_ray_ml_cr_repo: str
     byod_ray_llm_cr_repo: str
     byod_ecr: str
+    byod_ecr_region: str
     byod_aws_cr: str
     byod_gcp_cr: str
     state_machine_pr_aws_bucket: str
@@ -67,6 +68,10 @@ def _init_global_config(config_file: str):
             config_content.get("byod", {}).get("byod_ecr")
             or config_content.get("release_byod", {}).get("byod_ecr")
         ),
+        byod_ecr_region=(
+            config_content.get("byod", {}).get("byod_ecr_region")
+            or config_content.get("release_byod", {}).get("byod_ecr_region")
+        ),
         byod_aws_cr=(
             config_content.get("byod", {}).get("aws_cr")
             or config_content.get("release_byod", {}).get("aws_cr")

release/ray_release/custom_byod_build_init_helper.py

@@ -0,0 +1,67 @@
+from typing import List, Tuple
+import yaml
+from ray_release.configs.global_config import get_global_config
+from ray_release.logger import logger
+from ray_release.test import Test
+
+
+def _generate_custom_build_step_key(image: str) -> str:
+    # Buildkite step key cannot contain special characters, so they need to be replaced.
+    # Buildkite also limits step key length to 80 characters.
+    return (
+        "custom_build_"
+        + image.replace("/", "_")
+        .replace(":", "_")
+        .replace(".", "_")
+        .replace("-", "_")[-40:]
+    )
+
+
+def get_images_from_tests(tests: List[Test]) -> List[Tuple[str, str, str]]:
+    """Get a list of custom BYOD images to build from a list of tests."""
+    custom_byod_images = set()
+    for test in tests:
+        if not test.require_custom_byod_image():
+            continue
+        custom_byod_image_build = (
+            test.get_anyscale_byod_image(),
+            test.get_anyscale_base_byod_image(),
+            test.get_byod_post_build_script(),
+        )
+        logger.info(f"To be built: {custom_byod_image_build[0]}")
+        custom_byod_images.add(custom_byod_image_build)
+    return list(custom_byod_images)
+
+
+def create_custom_build_yaml(destination_file: str, tests: List[Test]) -> None:
+    config = get_global_config()
+    if not config or not config.get("byod_ecr_region") or not config.get("byod_ecr"):
+        raise ValueError("byod_ecr_region and byod_ecr must be set in the config")
+    """Create a yaml file for building custom BYOD images"""
+    custom_byod_images = get_images_from_tests(tests)
+    if not custom_byod_images:
+        return
+    build_config = {"group": "Custom images build", "steps": []}
+
+    for image, base_image, post_build_script in custom_byod_images:
+        if not post_build_script:
+            continue
+        step = {
+            "label": f":tapioca: build custom: {image}",
+            "key": _generate_custom_build_step_key(image),
+            "instance_type": "release-medium",
+            "commands": [
+                f"aws ecr get-login-password --region {config['byod_ecr_region']} | docker login --username AWS --password-stdin {config['byod_ecr']}",
+                f"bazelisk run //release:custom_byod_build -- --image-name {image} --base-image {base_image} --post-build-script {post_build_script}",
+            ],
+        }
+        if "ray-ml" in image:
+            step["depends_on"] = "anyscalemlbuild"
+        elif "ray-llm" in image:
+            step["depends_on"] = "anyscalellmbuild"
+        else:
+            step["depends_on"] = "anyscalebuild"
+        build_config["steps"].append(step)
+
+    with open(destination_file, "w") as f:
+        yaml.dump(build_config, f, default_flow_style=False, sort_keys=False)

release/ray_release/scripts/custom_byod_build_init.py

@@ -0,0 +1,112 @@
+import os
+from typing import Tuple
+from pathlib import Path
+import sys
+
+import click
+
+from ray_release.buildkite.filter import filter_tests
+from ray_release.buildkite.settings import get_pipeline_settings
+from ray_release.config import (
+    read_and_validate_release_test_collection,
+    RELEASE_TEST_CONFIG_FILES,
+)
+from ray_release.configs.global_config import init_global_config
+from ray_release.exception import ReleaseTestConfigError, ReleaseTestCLIError
+from ray_release.logger import logger
+from ray_release.custom_byod_build_init_helper import create_custom_build_yaml
+
+
+@click.command(
+    help="Create a rayci yaml file for building custom BYOD images based on tests."
+)
+@click.option(
+    "--test-collection-file",
+    type=str,
+    multiple=True,
+    help="Test collection file, relative path to ray repo.",
+)
+@click.option(
+    "--run-jailed-tests",
+    is_flag=True,
+    show_default=True,
+    default=False,
+    help=("Will run jailed tests."),
+)
+@click.option(
+    "--run-unstable-tests",
+    is_flag=True,
+    show_default=True,
+    default=False,
+    help=("Will run unstable tests."),
+)
+@click.option(
+    "--global-config",
+    default="oss_config.yaml",
+    type=click.Choice(
+        [x.name for x in (Path(__file__).parent.parent / "configs").glob("*.yaml")]
+    ),
+    help="Global config to use for test execution.",
+)
+@click.option(
+    "--frequency",
+    default=None,
+    type=click.Choice(["manual", "nightly", "nightly-3x", "weekly"]),
+    help="Run frequency of the test",
+)
+@click.option(
+    "--test-filters",
+    default=None,
+    type=str,
+    help="Test filters by prefix/regex.",
+)
+def main(
+    test_collection_file: Tuple[str],
+    run_jailed_tests: bool = False,
+    run_unstable_tests: bool = False,
+    global_config: str = "oss_config.yaml",
+    frequency: str = None,
+    test_filters: str = None,
+):
+    global_config_file = os.path.join(
+        os.path.dirname(__file__), "..", "configs", global_config
+    )
+    init_global_config(global_config_file)
+    settings = get_pipeline_settings()
+
+    frequency = frequency or settings["frequency"]
+    prefer_smoke_tests = settings["prefer_smoke_tests"]
+    test_filters = test_filters or settings["test_filters"]
+
+    try:
+        test_collection = read_and_validate_release_test_collection(
+            test_collection_file or RELEASE_TEST_CONFIG_FILES
+        )
+    except ReleaseTestConfigError as e:
+        raise ReleaseTestConfigError(
+            "Cannot load test yaml file.\nHINT: If you're kicking off tests for a "
+            "specific commit on Buildkite to test Ray wheels, after clicking "
+            "'New build', leave the commit at HEAD, and only specify the commit "
+            "in the dialog that asks for the Ray wheels."
+        ) from e
+
+    filtered_tests = filter_tests(
+        test_collection,
+        frequency=frequency,
+        test_filters=test_filters,
+        prefer_smoke_tests=prefer_smoke_tests,
+        run_jailed_tests=run_jailed_tests,
+        run_unstable_tests=run_unstable_tests,
+    )
+    logger.info(f"Found {len(filtered_tests)} tests to run.")
+    if len(filtered_tests) == 0:
+        raise ReleaseTestCLIError(
+            "Empty test collection. The selected frequency or filter did "
+            "not return any tests to run. Adjust your filters."
+        )
+    tests = [test for test, _ in filtered_tests]
+    create_custom_build_yaml(".buildkite/release/custom_byod_build.rayci.yml", tests)
+
+
+if __name__ == "__main__":
+    sys.exit(main())

release/ray_release/tests/test_custom_byod_build_init_helper.py

@@ -0,0 +1,77 @@
+import os
+import tempfile
+import sys
+import pytest
+from unittest import mock
+import yaml
+
+from ray_release.custom_byod_build_init_helper import create_custom_build_yaml
+from ray_release.configs.global_config import init_global_config
+from ray_release.bazel import bazel_runfile
+from ray_release.test import Test
+from ray_release.configs.global_config import get_global_config
+
+
+init_global_config(bazel_runfile("release/ray_release/configs/oss_config.yaml"))
+
+
+@mock.patch("ray_release.custom_byod_build_init_helper.get_images_from_tests")
+def test_create_custom_build_yaml(mock_get_images_from_tests):
+    config = get_global_config()
+    custom_byod_images = [
+        (
+            "ray-project/ray-ml:abc123-custom",
+            "ray-project/ray-ml:abc123-base",
+            "custom_script.sh",
+        ),
+        ("ray-project/ray-ml:abc123-custom", "ray-project/ray-ml:abc123-base", ""),
+        (
+            "ray-project/ray-ml:nightly-py37-cpu-custom-abcdef123456789abc123456789",
+            "ray-project/ray-ml:nightly-py37-cpu-base",
+            "custom_script.sh",
+        ),  # longer than 40 chars
+    ]
+    mock_get_images_from_tests.return_value = custom_byod_images
+
+    # List of dummy tests
+    tests = [
+        Test(
+            name="test_1",
+            frequency="manual",
+            group="test_group",
+            team="test_team",
+            working_dir="test_working_dir",
+        ),
+        Test(
+            name="test_2",
+            frequency="manual",
+            group="test_group",
+            team="test_team",
+            working_dir="test_working_dir",
+        ),
+    ]
+    with tempfile.TemporaryDirectory() as tmpdir:
+        create_custom_build_yaml(
+            os.path.join(tmpdir, "custom_byod_build.rayci.yml"), tests
+        )
+        with open(os.path.join(tmpdir, "custom_byod_build.rayci.yml"), "r") as f:
+            content = yaml.safe_load(f)
+            assert content["group"] == "Custom images build"
+            assert len(content["steps"]) == 2
+            assert (
+                f"--region {config['byod_ecr_region']}"
+                in content["steps"][0]["commands"][0]
+            )
+            assert f"{config['byod_ecr']}" in content["steps"][0]["commands"][0]
+            assert (
+                f"--image-name {custom_byod_images[0][0]}"
+                in content["steps"][0]["commands"][1]
+            )
+            assert (
+                f"--image-name {custom_byod_images[2][0]}"
+                in content["steps"][1]["commands"][1]
+            )
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main(["-v", __file__]))