diff --git a/.buildkite/release/build.rayci.yml b/.buildkite/release/build.rayci.yml index bee001c4ef19..29e764aa9bd3 100644 --- a/.buildkite/release/build.rayci.yml +++ b/.buildkite/release/build.rayci.yml @@ -1,5 +1,55 @@ group: release build steps: + # Build ray core components (required for wheel builds) + - name: ray-core-build + label: "wanda: core binary parts py{{matrix}} (x86_64)" + wanda: ci/docker/ray-core.wanda.yaml + matrix: + - "3.10" + - "3.11" + - "3.12" + env: + PYTHON_VERSION: "{{matrix}}" + ARCH_SUFFIX: "" + HOSTTYPE: "x86_64" + MANYLINUX_VERSION: "260103.868e54c" + tags: oss + + - name: ray-dashboard-build + label: "wanda: dashboard" + wanda: ci/docker/ray-dashboard.wanda.yaml + env: + HOSTTYPE: "x86_64" + MANYLINUX_VERSION: "260103.868e54c" + tags: oss + + - name: ray-java-build + label: "wanda: java build (x86_64)" + wanda: ci/docker/ray-java.wanda.yaml + tags: oss + env: + ARCH_SUFFIX: "" + HOSTTYPE: "x86_64" + MANYLINUX_VERSION: "260103.868e54c" + + - name: ray-wheel-build + label: "wanda: wheel py{{matrix}} (x86_64)" + wanda: ci/docker/ray-wheel.wanda.yaml + matrix: + - "3.10" + - "3.11" + - "3.12" + env: + PYTHON_VERSION: "{{matrix}}" + ARCH_SUFFIX: "" + HOSTTYPE: "x86_64" + MANYLINUX_VERSION: "260103.868e54c" + tags: oss + depends_on: + - ray-core-build + - ray-java-build + - ray-dashboard-build + - name: raycpubaseextra-testdeps label: "wanda: ray.py{{matrix}}.cpu.base-extra-testdeps" wanda: docker/base-extra-testdeps/cpu.wanda.yaml @@ -67,26 +117,72 @@ steps: depends_on: - ray-mlcudabaseextra - - label: ":tapioca: build: ray py{{matrix.python}}-{{matrix.platform}} image for release tests" + # Build ray anyscale images using wanda (CPU) + - name: ray-anyscale-cpu-build + label: "wanda: ray-anyscale py{{matrix}} cpu" + wanda: ci/docker/ray-anyscale-cpu.wanda.yaml + matrix: + # This list should be kept in sync with the list of supported Python in + # release test suite + - "3.10" + - "3.11" + - "3.12" + env: + PYTHON_VERSION: "{{matrix}}" + ARCH_SUFFIX: "" + tags: + - oss + depends_on: + - ray-wheel-build + - raycpubaseextra-testdeps + + # Build ray anyscale images using wanda (CUDA) + - name: ray-anyscale-cuda-build + label: "wanda: ray-anyscale py{{matrix.python}} cu{{matrix.cuda}}" + wanda: ci/docker/ray-anyscale-cuda.wanda.yaml + matrix: + setup: + python: + - "3.10" + - "3.11" + - "3.12" + cuda: + - "12.3.2-cudnn9" + env: + PYTHON_VERSION: "{{matrix.python}}" + CUDA_VERSION: "{{matrix.cuda}}" + ARCH_SUFFIX: "" + tags: + - oss + depends_on: + - ray-wheel-build + - raycudabaseextra-testdeps + + # Push anyscale images to ECR/GCP/Azure registries + - label: ":docker: push: ray-anyscale py{{matrix.python}} {{matrix.platform}}" key: anyscalebuild instance_type: release-medium mount_buildkite_agent: true tags: - oss commands: - - bazel run //ci/ray_ci:build_in_docker -- anyscale - --python-version {{matrix.python}} --platform {{matrix.platform}} - --image-type ray --upload + # Authenticate with GCP and Azure before pushing + - bash release/gcloud_docker_login.sh release/aws2gce_iam.json + - bash release/azure_docker_login.sh + - az acr login --name rayreleasetest + # PATH must include gcloud SDK for crane to use docker credential helper + - export PATH=$(pwd)/google-cloud-sdk/bin:$PATH && + bazel run //ci/ray_ci/automation:push_anyscale_image -- + --python-version {{matrix.python}} + --platform {{matrix.platform}} + --image-type ray + --upload depends_on: - - manylinux-x86_64 - - forge - - raycpubaseextra-testdeps - - raycudabaseextra-testdeps + - ray-anyscale-cpu-build + - ray-anyscale-cuda-build matrix: setup: python: - # This list should be kept in sync with the list of supported Python in - # release test suite - "3.10" - "3.11" - "3.12" @@ -94,35 +190,93 @@ steps: - cu12.3.2-cudnn9 - cpu - - label: ":tapioca: build: ray-llm py{{matrix}} image for release tests" + # Build ray-llm anyscale image using wanda + - name: ray-llm-anyscale-cuda-build + label: "wanda: ray-llm-anyscale py{{matrix.python}} cu{{matrix.cuda}}" + wanda: ci/docker/ray-llm-anyscale-cuda.wanda.yaml + matrix: + setup: + python: + - "3.11" + cuda: + - "12.8.1-cudnn" + env: + PYTHON_VERSION: "{{matrix.python}}" + CUDA_VERSION: "{{matrix.cuda}}" + ARCH_SUFFIX: "" + tags: + - oss + depends_on: + - ray-wheel-build + - ray-llmbaseextra-testdeps + + # Push ray-llm anyscale images to ECR/GCP/Azure registries + - label: ":docker: push: ray-llm-anyscale py{{matrix}} cu12.8.1-cudnn" key: anyscalellmbuild instance_type: release-medium mount_buildkite_agent: true tags: - oss commands: - - bazel run //ci/ray_ci:build_in_docker -- anyscale --python-version {{matrix}} - --platform cu12.8.1-cudnn --image-type ray-llm --upload + # Authenticate with GCP and Azure before pushing + - bash release/gcloud_docker_login.sh release/aws2gce_iam.json + - bash release/azure_docker_login.sh + - az acr login --name rayreleasetest + # PATH must include gcloud SDK for crane to use docker credential helper + - export PATH=$(pwd)/google-cloud-sdk/bin:$PATH && + bazel run //ci/ray_ci/automation:push_anyscale_image -- + --python-version {{matrix}} + --platform cu12.8.1-cudnn + --image-type ray-llm + --upload depends_on: - - manylinux-x86_64 - - forge - - ray-llmbaseextra-testdeps + - ray-llm-anyscale-cuda-build matrix: - "3.11" - - label: ":tapioca: build: ray-ml py{{matrix}} image for release tests" + # Build ray-ml anyscale image using wanda + - name: ray-ml-anyscale-cuda-build + label: "wanda: ray-ml-anyscale py{{matrix.python}} cu{{matrix.cuda}}" + wanda: ci/docker/ray-ml-anyscale-cuda.wanda.yaml + matrix: + setup: + python: + # This list should be kept in sync with the list of supported Python in + # release test suite + - "3.10" + cuda: + - "12.1.1-cudnn8" + env: + PYTHON_VERSION: "{{matrix.python}}" + CUDA_VERSION: "{{matrix.cuda}}" + ARCH_SUFFIX: "" + tags: + - oss + depends_on: + - ray-wheel-build + - ray-mlcudabaseextra-testdeps + + # Push ray-ml anyscale images to ECR/GCP/Azure registries + - label: ":docker: push: ray-ml-anyscale py{{matrix}} cu12.1.1-cudnn8" key: anyscalemlbuild instance_type: release-medium mount_buildkite_agent: true tags: - oss commands: - - bazel run //ci/ray_ci:build_in_docker -- anyscale --python-version {{matrix}} - --platform cu12.1.1-cudnn8 --image-type ray-ml --upload + # Authenticate with GCP and Azure before pushing + - bash release/gcloud_docker_login.sh release/aws2gce_iam.json + - bash release/azure_docker_login.sh + - az acr login --name rayreleasetest + # PATH must include gcloud SDK for crane to use docker credential helper + - export PATH=$(pwd)/google-cloud-sdk/bin:$PATH && + bazel run //ci/ray_ci/automation:push_anyscale_image -- + --python-version {{matrix}} + --platform cu12.1.1-cudnn8 + --image-type ray-ml + --upload depends_on: - - manylinux-x86_64 - - forge - - ray-mlcudabaseextra-testdeps + - ray-ml-anyscale-cuda-build matrix: # This list should be kept in sync with the list of supported Python in # release test suite diff --git a/ci/docker/ray-anyscale-cpu.wanda.yaml b/ci/docker/ray-anyscale-cpu.wanda.yaml new file mode 100644 index 000000000000..a4a1280dda5e --- /dev/null +++ b/ci/docker/ray-anyscale-cpu.wanda.yaml @@ -0,0 +1,20 @@ +# Ray Anyscale CPU Image (for release tests) +# Installs ray wheel into the CPU base-extra-testdeps image +# +# This produces the anyscale test image for CPU-only release tests. +# Unlike the regular ray image, this uses base-extra-testdeps which +# includes additional test dependencies. +# +name: "ray-anyscale-py$PYTHON_VERSION-cpu$ARCH_SUFFIX" +disable_caching: true +froms: + - "cr.ray.io/rayproject/ray-py$PYTHON_VERSION-cpu-base-extra-testdeps" # CPU base with test deps + - "cr.ray.io/rayproject/ray-wheel-py$PYTHON_VERSION$ARCH_SUFFIX" # Ray wheel +dockerfile: ci/docker/ray-image.Dockerfile +srcs: + - python/requirements_compiled.txt +build_args: + - PYTHON_VERSION + - ARCH_SUFFIX + - BASE_IMAGE=cr.ray.io/rayproject/ray-py$PYTHON_VERSION-cpu-base-extra-testdeps + - RAY_WHEEL_IMAGE=cr.ray.io/rayproject/ray-wheel-py$PYTHON_VERSION$ARCH_SUFFIX diff --git a/ci/docker/ray-anyscale-cuda.wanda.yaml b/ci/docker/ray-anyscale-cuda.wanda.yaml new file mode 100644 index 000000000000..f85a73ad9818 --- /dev/null +++ b/ci/docker/ray-anyscale-cuda.wanda.yaml @@ -0,0 +1,21 @@ +# Ray Anyscale CUDA Image (for release tests) +# Installs ray wheel into the CUDA base-extra-testdeps image +# +# This produces the anyscale test image for GPU release tests. +# Unlike the regular ray image, this uses base-extra-testdeps which +# includes additional test dependencies. +# +name: "ray-anyscale-py$PYTHON_VERSION-cu$CUDA_VERSION$ARCH_SUFFIX" +disable_caching: true +froms: + - "cr.ray.io/rayproject/ray-py$PYTHON_VERSION-cu$CUDA_VERSION-base-extra-testdeps" # CUDA base with test deps + - "cr.ray.io/rayproject/ray-wheel-py$PYTHON_VERSION$ARCH_SUFFIX" # Ray wheel +dockerfile: ci/docker/ray-image.Dockerfile +srcs: + - python/requirements_compiled.txt +build_args: + - PYTHON_VERSION + - CUDA_VERSION + - ARCH_SUFFIX + - BASE_IMAGE=cr.ray.io/rayproject/ray-py$PYTHON_VERSION-cu$CUDA_VERSION-base-extra-testdeps + - RAY_WHEEL_IMAGE=cr.ray.io/rayproject/ray-wheel-py$PYTHON_VERSION$ARCH_SUFFIX diff --git a/ci/docker/ray-llm-anyscale-cuda.wanda.yaml b/ci/docker/ray-llm-anyscale-cuda.wanda.yaml new file mode 100644 index 000000000000..584121d4700e --- /dev/null +++ b/ci/docker/ray-llm-anyscale-cuda.wanda.yaml @@ -0,0 +1,20 @@ +# Ray-LLM Anyscale CUDA Image (for release tests) +# Installs ray wheel into the ray-llm CUDA base-extra-testdeps image +# +# This produces the anyscale test image for ray-llm GPU release tests. +# Uses ray-llm base with extra test dependencies. +# +name: "ray-llm-anyscale-py$PYTHON_VERSION-cu$CUDA_VERSION$ARCH_SUFFIX" +disable_caching: true +froms: + - "cr.ray.io/rayproject/ray-llm-py$PYTHON_VERSION-cu$CUDA_VERSION-base-extra-testdeps" # ray-llm base with test deps + - "cr.ray.io/rayproject/ray-wheel-py$PYTHON_VERSION$ARCH_SUFFIX" # Ray wheel +dockerfile: ci/docker/ray-image.Dockerfile +srcs: + - python/requirements_compiled.txt +build_args: + - PYTHON_VERSION + - CUDA_VERSION + - ARCH_SUFFIX + - BASE_IMAGE=cr.ray.io/rayproject/ray-llm-py$PYTHON_VERSION-cu$CUDA_VERSION-base-extra-testdeps + - RAY_WHEEL_IMAGE=cr.ray.io/rayproject/ray-wheel-py$PYTHON_VERSION$ARCH_SUFFIX diff --git a/ci/docker/ray-ml-anyscale-cuda.wanda.yaml b/ci/docker/ray-ml-anyscale-cuda.wanda.yaml new file mode 100644 index 000000000000..82efd3242eef --- /dev/null +++ b/ci/docker/ray-ml-anyscale-cuda.wanda.yaml @@ -0,0 +1,20 @@ +# Ray-ML Anyscale CUDA Image (for release tests) +# Installs ray wheel into the ray-ml CUDA base-extra-testdeps image +# +# This produces the anyscale test image for ray-ml GPU release tests. +# Uses ray-ml base with extra test dependencies. +# +name: "ray-ml-anyscale-py$PYTHON_VERSION-cu$CUDA_VERSION$ARCH_SUFFIX" +disable_caching: true +froms: + - "cr.ray.io/rayproject/ray-ml-py$PYTHON_VERSION-cu$CUDA_VERSION-base-extra-testdeps" # ray-ml base with test deps + - "cr.ray.io/rayproject/ray-wheel-py$PYTHON_VERSION$ARCH_SUFFIX" # Ray wheel +dockerfile: ci/docker/ray-image.Dockerfile +srcs: + - python/requirements_compiled.txt +build_args: + - PYTHON_VERSION + - CUDA_VERSION + - ARCH_SUFFIX + - BASE_IMAGE=cr.ray.io/rayproject/ray-ml-py$PYTHON_VERSION-cu$CUDA_VERSION-base-extra-testdeps + - RAY_WHEEL_IMAGE=cr.ray.io/rayproject/ray-wheel-py$PYTHON_VERSION$ARCH_SUFFIX diff --git a/ci/ray_ci/automation/BUILD.bazel b/ci/ray_ci/automation/BUILD.bazel index 5184e41df944..51920da51fef 100644 --- a/ci/ray_ci/automation/BUILD.bazel +++ b/ci/ray_ci/automation/BUILD.bazel @@ -331,3 +331,30 @@ py_test( ci_require("pytest"), ], ) + +py_binary( + name = "push_anyscale_image", + srcs = ["push_anyscale_image.py"], + exec_compatible_with = ["//bazel:py3"], + deps = [ + ":crane_lib", + "//ci/ray_ci:ray_ci_lib", + "//release:ray_release", + ci_require("click"), + ], +) + +py_test( + name = "test_push_anyscale_image", + size = "small", + srcs = ["test_push_anyscale_image.py"], + exec_compatible_with = ["//bazel:py3"], + tags = [ + "ci_unit", + "team:ci", + ], + deps = [ + ":push_anyscale_image", + ci_require("pytest"), + ], +) diff --git a/ci/ray_ci/automation/push_anyscale_image.py b/ci/ray_ci/automation/push_anyscale_image.py new file mode 100644 index 000000000000..9951adf1e58b --- /dev/null +++ b/ci/ray_ci/automation/push_anyscale_image.py @@ -0,0 +1,254 @@ +""" +Push Wanda-cached anyscale images to ECR, GCP, and Azure registries. + +This script copies anyscale images from the Wanda cache to the three cloud +registries used for release tests: +- AWS ECR: anyscale/{image_type}:{tag} +- GCP Artifact Registry: anyscale/{image_type}:{tag} +- Azure Container Registry: anyscale/{image_type}:{tag} + +Example: + bazel run //ci/ray_ci/automation:push_anyscale_image -- \\ + --python-version 3.10 \\ + --platform cpu \\ + --image-type ray \\ + --upload + +Run with --help to see all options. +""" + +import logging +import os +import sys +from typing import List + +import click + +from ci.ray_ci.automation.crane_lib import ( + call_crane_copy, + call_crane_manifest, +) +from ci.ray_ci.utils import ci_init, ecr_docker_login + +# Configure logging +logging.basicConfig( + level=logging.INFO, + format="%(message)s", + stream=sys.stdout, +) +logger = logging.getLogger(__name__) + +# Registry URLs +_DOCKER_ECR_REPO = os.environ.get( + "RAYCI_WORK_REPO", + "029272617770.dkr.ecr.us-west-2.amazonaws.com/rayproject", +) +_DOCKER_GCP_REGISTRY = os.environ.get( + "RAYCI_GCP_REGISTRY", + "us-west1-docker.pkg.dev/anyscale-oss-ci", +) +_DOCKER_AZURE_REGISTRY = os.environ.get( + "RAYCI_AZURE_REGISTRY", + "rayreleasetest.azurecr.io", +) + +# GPU_PLATFORM is the default GPU platform that gets aliased as "gpu" +# This must match the definition in ci/ray_ci/docker_container.py +GPU_PLATFORM = "cu12.1.1-cudnn8" + + +class PushAnyscaleImageError(Exception): + """Error raised when pushing anyscale images fails.""" + + +def _format_python_version_tag(python_version: str) -> str: + """Format python version as -py310 (no dots, with hyphen prefix).""" + return f"-py{python_version.replace('.', '')}" + + +def _format_platform_tag(platform: str) -> str: + """Format platform as -cpu or shortened CUDA version.""" + if platform == "cpu": + return "-cpu" + # cu12.3.2-cudnn9 -> -cu123 + versions = platform.split(".") + return f"-{versions[0]}{versions[1]}" + + +def _get_image_tags(python_version: str, platform: str) -> List[str]: + """ + Generate image tags matching the original docker_container.py format. + + Returns multiple tags for the image (canonical + aliases). + For GPU_PLATFORM, also generates -gpu alias tags to match release test expectations. + """ + branch = os.environ.get("BUILDKITE_BRANCH", "") + commit = os.environ.get("BUILDKITE_COMMIT", "")[:6] + rayci_build_id = os.environ.get("RAYCI_BUILD_ID", "") + + py_tag = _format_python_version_tag(python_version) + platform_tag = _format_platform_tag(platform) + + # For GPU_PLATFORM, also create -gpu alias (release tests use type: gpu) + platform_tags = [platform_tag] + if platform == GPU_PLATFORM: + platform_tags.append("-gpu") + + tags = [] + + if branch == "master": + # On master, use sha and build_id as tags + for ptag in platform_tags: + tags.append(f"{commit}{py_tag}{ptag}") + if rayci_build_id: + for ptag in platform_tags: + tags.append(f"{rayci_build_id}{py_tag}{ptag}") + elif branch.startswith("releases/"): + # On release branches, use release name + release_name = branch[len("releases/") :] + for ptag in platform_tags: + tags.append(f"{release_name}.{commit}{py_tag}{ptag}") + if rayci_build_id: + for ptag in platform_tags: + tags.append(f"{rayci_build_id}{py_tag}{ptag}") + else: + # For other branches (PRs, etc.) + pr = os.environ.get("BUILDKITE_PULL_REQUEST", "false") + if pr != "false": + for ptag in platform_tags: + tags.append(f"pr-{pr}.{commit}{py_tag}{ptag}") + else: + for ptag in platform_tags: + tags.append(f"{commit}{py_tag}{ptag}") + if rayci_build_id: + for ptag in platform_tags: + tags.append(f"{rayci_build_id}{py_tag}{ptag}") + + return tags + + +def _get_wanda_image_name(python_version: str, platform: str, image_type: str) -> str: + """Get the wanda-cached image name. + + Platform is passed with "cu" prefix (e.g., "cu12.3.2-cudnn9") or "cpu". + """ + if platform == "cpu": + return f"{image_type}-anyscale-py{python_version}-cpu" + else: + # Platform already includes "cu" prefix from pipeline matrix + return f"{image_type}-anyscale-py{python_version}-{platform}" + + +def _image_exists(tag: str) -> bool: + """Check if a container image manifest exists using crane.""" + return_code, _ = call_crane_manifest(tag) + return return_code == 0 + + +def _copy_image(source: str, destination: str, dry_run: bool = False) -> None: + """Copy a container image from source to destination using crane.""" + if dry_run: + logger.info(f"DRY RUN: Would copy {source} -> {destination}") + return + + logger.info(f"Copying {source} -> {destination}") + return_code, output = call_crane_copy(source, destination) + if return_code != 0: + raise PushAnyscaleImageError(f"Crane copy failed: {output}") + + +@click.command() +@click.option( + "--python-version", + type=str, + required=True, + help="Python version (e.g., '3.10')", +) +@click.option( + "--platform", + type=str, + required=True, + help="Platform (e.g., 'cpu', 'cu12.3.2-cudnn9')", +) +@click.option( + "--image-type", + type=str, + default="ray", + help="Image type (e.g., 'ray', 'ray-llm', 'ray-ml')", +) +@click.option( + "--upload", + is_flag=True, + default=False, + help="Actually push to registries. Without this flag, runs in dry-run mode.", +) +def main( + python_version: str, + platform: str, + image_type: str, + upload: bool, +) -> None: + """ + Push a Wanda-cached anyscale image to ECR, GCP, and Azure registries. + + NOTE: GCP and Azure authentication must be done BEFORE calling this script + (e.g., via gcloud_docker_login.sh and azure_docker_login.sh in the pipeline). + ECR authentication is handled internally. + """ + ci_init() + + dry_run = not upload + if dry_run: + logger.info("DRY RUN MODE - no images will be pushed") + + # Get required environment variables + rayci_work_repo = os.environ.get("RAYCI_WORK_REPO", _DOCKER_ECR_REPO) + rayci_build_id = os.environ.get("RAYCI_BUILD_ID") + + if not rayci_build_id: + raise PushAnyscaleImageError("RAYCI_BUILD_ID environment variable not set") + + # Construct source image from Wanda cache + wanda_image_name = _get_wanda_image_name(python_version, platform, image_type) + source_tag = f"{rayci_work_repo}:{rayci_build_id}-{wanda_image_name}" + + logger.info(f"Source image (Wanda): {source_tag}") + + # Authenticate with ECR (source registry) + ecr_registry = rayci_work_repo.split("/")[0] + ecr_docker_login(ecr_registry) + + # Verify source image exists + logger.info("Verifying source image in Wanda cache...") + if not _image_exists(source_tag): + raise PushAnyscaleImageError( + f"Source image not found in Wanda cache: {source_tag}" + ) + + # Get image tags + tags = _get_image_tags(python_version, platform) + canonical_tag = tags[0] + + logger.info(f"Canonical tag: {canonical_tag}") + logger.info(f"All tags: {tags}") + + # Push to all three registries (ECR, GCP, Azure) + # NOTE: Authentication for GCP/Azure must be done in the pipeline step BEFORE + # calling this script (e.g., via gcloud_docker_login.sh and azure_docker_login.sh). + registries = [ + (ecr_registry, "ECR"), + (_DOCKER_GCP_REGISTRY, "GCP"), + (_DOCKER_AZURE_REGISTRY, "Azure"), + ] + + for tag in tags: + for registry, name in registries: + dest_image = f"{registry}/anyscale/{image_type}:{tag}" + logger.info(f"Pushing to {name}: {dest_image}") + _copy_image(source_tag, dest_image, dry_run) + + logger.info("Successfully pushed anyscale images to all registries") + + +if __name__ == "__main__": + main() diff --git a/ci/ray_ci/automation/test_push_anyscale_image.py b/ci/ray_ci/automation/test_push_anyscale_image.py new file mode 100644 index 000000000000..2c0e33cd4b90 --- /dev/null +++ b/ci/ray_ci/automation/test_push_anyscale_image.py @@ -0,0 +1,203 @@ +import sys +from unittest import mock + +import pytest + +from ci.ray_ci.automation.push_anyscale_image import ( + GPU_PLATFORM, + _format_platform_tag, + _format_python_version_tag, + _get_image_tags, + _get_wanda_image_name, +) + + +class TestFormatPythonVersionTag: + @pytest.mark.parametrize( + ("python_version", "expected"), + [ + ("3.10", "-py310"), + ("3.11", "-py311"), + ("3.12", "-py312"), + ("3.9", "-py39"), + ], + ) + def test_format_python_version_tag(self, python_version, expected): + assert _format_python_version_tag(python_version) == expected + + +class TestFormatPlatformTag: + @pytest.mark.parametrize( + ("platform", "expected"), + [ + ("cpu", "-cpu"), + ("cu11.7.1-cudnn8", "-cu117"), + ("cu11.8.0-cudnn8", "-cu118"), + ("cu12.1.1-cudnn8", "-cu121"), + ("cu12.3.2-cudnn9", "-cu123"), + ("cu12.8.1-cudnn", "-cu128"), + ], + ) + def test_format_platform_tag(self, platform, expected): + assert _format_platform_tag(platform) == expected + + +class TestGetWandaImageName: + @pytest.mark.parametrize( + ("python_version", "platform", "image_type", "expected"), + [ + ("3.10", "cpu", "ray", "ray-anyscale-py3.10-cpu"), + ("3.11", "cu12.1.1-cudnn8", "ray", "ray-anyscale-py3.11-cu12.1.1-cudnn8"), + ("3.10", "cpu", "ray-llm", "ray-llm-anyscale-py3.10-cpu"), + ( + "3.11", + "cu12.8.1-cudnn", + "ray-llm", + "ray-llm-anyscale-py3.11-cu12.8.1-cudnn", + ), + ("3.10", "cpu", "ray-ml", "ray-ml-anyscale-py3.10-cpu"), + ( + "3.11", + "cu12.3.2-cudnn9", + "ray-ml", + "ray-ml-anyscale-py3.11-cu12.3.2-cudnn9", + ), + ], + ) + def test_get_wanda_image_name(self, python_version, platform, image_type, expected): + assert _get_wanda_image_name(python_version, platform, image_type) == expected + + +class TestGetImageTags: + @mock.patch.dict( + "os.environ", + { + "BUILDKITE_BRANCH": "master", + "BUILDKITE_COMMIT": "abc123def456", + "RAYCI_BUILD_ID": "build-123", + }, + ) + def test_master_branch_tags(self): + tags = _get_image_tags(python_version="3.10", platform="cpu") + + assert tags == [ + "abc123-py310-cpu", + "build-123-py310-cpu", + ] + + @mock.patch.dict( + "os.environ", + { + "BUILDKITE_BRANCH": "master", + "BUILDKITE_COMMIT": "abc123def456", + "RAYCI_BUILD_ID": "build-123", + }, + ) + def test_master_branch_gpu_platform_includes_alias(self): + tags = _get_image_tags(python_version="3.10", platform=GPU_PLATFORM) + + assert tags == [ + "abc123-py310-cu121", + "abc123-py310-gpu", + "build-123-py310-cu121", + "build-123-py310-gpu", + ] + + @mock.patch.dict( + "os.environ", + { + "BUILDKITE_BRANCH": "master", + "BUILDKITE_COMMIT": "abc123def456", + "RAYCI_BUILD_ID": "", + }, + ) + def test_master_branch_no_build_id(self): + tags = _get_image_tags(python_version="3.11", platform="cu12.3.2-cudnn9") + + assert tags == ["abc123-py311-cu123"] + + @mock.patch.dict( + "os.environ", + { + "BUILDKITE_BRANCH": "releases/2.44.0", + "BUILDKITE_COMMIT": "abc123def456", + "RAYCI_BUILD_ID": "build-456", + }, + ) + def test_release_branch_tags(self): + tags = _get_image_tags(python_version="3.10", platform="cpu") + + assert tags == [ + "2.44.0.abc123-py310-cpu", + "build-456-py310-cpu", + ] + + @mock.patch.dict( + "os.environ", + { + "BUILDKITE_BRANCH": "releases/2.44.0", + "BUILDKITE_COMMIT": "abc123def456", + "RAYCI_BUILD_ID": "build-456", + }, + ) + def test_release_branch_gpu_platform_includes_alias(self): + tags = _get_image_tags(python_version="3.10", platform=GPU_PLATFORM) + + assert tags == [ + "2.44.0.abc123-py310-cu121", + "2.44.0.abc123-py310-gpu", + "build-456-py310-cu121", + "build-456-py310-gpu", + ] + + @mock.patch.dict( + "os.environ", + { + "BUILDKITE_BRANCH": "feature-branch", + "BUILDKITE_COMMIT": "abc123def456", + "BUILDKITE_PULL_REQUEST": "123", + "RAYCI_BUILD_ID": "build-789", + }, + ) + def test_pr_branch_tags(self): + tags = _get_image_tags(python_version="3.12", platform="cpu") + + assert tags == [ + "pr-123.abc123-py312-cpu", + "build-789-py312-cpu", + ] + + @mock.patch.dict( + "os.environ", + { + "BUILDKITE_BRANCH": "feature-branch", + "BUILDKITE_COMMIT": "abc123def456", + "BUILDKITE_PULL_REQUEST": "false", + "RAYCI_BUILD_ID": "build-789", + }, + ) + def test_non_pr_feature_branch_tags(self): + tags = _get_image_tags(python_version="3.10", platform="cpu") + + assert tags == [ + "abc123-py310-cpu", + "build-789-py310-cpu", + ] + + @mock.patch.dict( + "os.environ", + { + "BUILDKITE_BRANCH": "feature-branch", + "BUILDKITE_COMMIT": "abc123def456", + "BUILDKITE_PULL_REQUEST": "false", + "RAYCI_BUILD_ID": "", + }, + ) + def test_feature_branch_no_build_id(self): + tags = _get_image_tags(python_version="3.10", platform="cpu") + + assert tags == ["abc123-py310-cpu"] + + +if __name__ == "__main__": + sys.exit(pytest.main(["-vv", __file__]))