123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258 |
- from typing import List, Optional, Dict
- import boto3
- import hashlib
- import os
- import subprocess
- import sys
- import time
- from ray_release.config import RELEASE_PACKAGE_DIR
- from ray_release.configs.global_config import get_global_config
- from ray_release.logger import logger
- from ray_release.test import (
- Test,
- DATAPLANE_ECR_REPO,
- DATAPLANE_ECR_ML_REPO,
- )
- DATAPLANE_S3_BUCKET = "ray-release-automation-results"
- DATAPLANE_FILENAME = "dataplane_20240613.tar.gz"
- DATAPLANE_DIGEST = "b7c68dd3cf9ef05b2b1518a32729fc036f06ae83bae9b9ecd241e33b37868944"
- BASE_IMAGE_WAIT_TIMEOUT = 7200
- BASE_IMAGE_WAIT_DURATION = 30
- RELEASE_BYOD_DIR = os.path.join(RELEASE_PACKAGE_DIR, "ray_release/byod")
- REQUIREMENTS_BYOD = "requirements_byod"
- REQUIREMENTS_ML_BYOD = "requirements_ml_byod"
- def build_champagne_image(
- ray_version: str,
- python_version: str,
- image_type: str,
- ) -> str:
- """
- Builds the Anyscale champagne image.
- """
- _download_dataplane_build_file()
- env = os.environ.copy()
- env["DOCKER_BUILDKIT"] = "1"
- if image_type == "cpu":
- ray_project = "ray"
- anyscale_repo = DATAPLANE_ECR_REPO
- else:
- ray_project = "ray-ml"
- anyscale_repo = DATAPLANE_ECR_ML_REPO
- ray_image = f"rayproject/{ray_project}:{ray_version}-{python_version}-{image_type}"
- anyscale_image = (
- f"{get_global_config()['byod_ecr']}/{anyscale_repo}:champagne-{ray_version}"
- )
- logger.info(f"Building champagne anyscale image from {ray_image}")
- with open(DATAPLANE_FILENAME, "rb") as build_file:
- subprocess.check_call(
- [
- "docker",
- "build",
- "--progress=plain",
- "--build-arg",
- f"BASE_IMAGE={ray_image}",
- "-t",
- anyscale_image,
- "-",
- ],
- stdin=build_file,
- stdout=sys.stderr,
- env=env,
- )
- _validate_and_push(anyscale_image)
- return anyscale_image
- def build_anyscale_custom_byod_image(test: Test) -> None:
- if not test.require_custom_byod_image():
- logger.info(f"Test {test.get_name()} does not require a custom byod image")
- return
- byod_image = test.get_anyscale_byod_image()
- if _image_exist(byod_image):
- logger.info(f"Image {byod_image} already exists")
- return
- env = os.environ.copy()
- env["DOCKER_BUILDKIT"] = "1"
- subprocess.check_call(
- [
- "docker",
- "build",
- "--progress=plain",
- "--build-arg",
- f"BASE_IMAGE={test.get_anyscale_base_byod_image()}",
- "--build-arg",
- f"POST_BUILD_SCRIPT={test.get_byod_post_build_script()}",
- "-t",
- byod_image,
- "-f",
- os.path.join(RELEASE_BYOD_DIR, "byod.custom.Dockerfile"),
- RELEASE_BYOD_DIR,
- ],
- stdout=sys.stderr,
- env=env,
- )
- _validate_and_push(byod_image)
- def build_anyscale_base_byod_images(tests: List[Test]) -> None:
- """
- Builds the Anyscale BYOD images for the given tests.
- """
- _download_dataplane_build_file()
- to_be_built = {}
- built = set()
- for test in tests:
- to_be_built[test.get_anyscale_base_byod_image()] = test
- env = os.environ.copy()
- env["DOCKER_BUILDKIT"] = "1"
- start = int(time.time())
- # ray images are built on post-merge, so we can wait for them to be available
- while (
- len(built) < len(to_be_built)
- and int(time.time()) - start < BASE_IMAGE_WAIT_TIMEOUT
- ):
- for byod_image, test in to_be_built.items():
- py_version = test.get_python_version()
- if test.use_byod_ml_image():
- byod_requirements = f"{REQUIREMENTS_ML_BYOD}_{py_version}.txt"
- else:
- byod_requirements = f"{REQUIREMENTS_BYOD}_{py_version}.txt"
- if _image_exist(byod_image):
- logger.info(f"Image {byod_image} already exists")
- built.add(byod_image)
- continue
- ray_image = test.get_ray_image()
- if not _image_exist(ray_image):
- # TODO(can): instead of waiting for the base image to be built, we can
- # build it ourselves
- timeout = BASE_IMAGE_WAIT_TIMEOUT - (int(time.time()) - start)
- logger.info(
- f"Image {ray_image} does not exist yet. "
- f"Wait for another {timeout}s..."
- )
- time.sleep(BASE_IMAGE_WAIT_DURATION)
- continue
- logger.info(f"Building {byod_image} from {ray_image}")
- with open(DATAPLANE_FILENAME, "rb") as build_file:
- subprocess.check_call(
- [
- "docker",
- "build",
- "--progress=plain",
- "--build-arg",
- f"BASE_IMAGE={ray_image}",
- "-t",
- byod_image,
- "-",
- ],
- stdin=build_file,
- stdout=sys.stderr,
- env=env,
- )
- subprocess.check_call(
- [
- "docker",
- "build",
- "--progress=plain",
- "--build-arg",
- f"BASE_IMAGE={byod_image}",
- "--build-arg",
- f"PIP_REQUIREMENTS={byod_requirements}",
- "--build-arg",
- "DEBIAN_REQUIREMENTS=requirements_debian_byod.txt",
- "-t",
- byod_image,
- "-f",
- os.path.join(RELEASE_BYOD_DIR, "byod.Dockerfile"),
- RELEASE_BYOD_DIR,
- ],
- stdout=sys.stderr,
- env=env,
- )
- _validate_and_push(byod_image)
- built.add(byod_image)
- def _validate_and_push(byod_image: str) -> None:
- """
- Validates the given image and pushes it to ECR.
- """
- docker_ray_commit = (
- subprocess.check_output(
- [
- "docker",
- "run",
- "-ti",
- "--entrypoint",
- "python",
- byod_image,
- "-c",
- "import ray; print(ray.__commit__)",
- ],
- )
- .decode("utf-8")
- .strip()
- )
- if os.environ.get("RAY_IMAGE_TAG"):
- logger.info(f"Ray commit from image: {docker_ray_commit}")
- else:
- expected_ray_commit = _get_ray_commit()
- assert (
- docker_ray_commit == expected_ray_commit
- ), f"Expected ray commit {expected_ray_commit}, found {docker_ray_commit}"
- logger.info(f"Pushing image to registry: {byod_image}")
- subprocess.check_call(
- ["docker", "push", byod_image],
- stdout=sys.stderr,
- )
- def _get_ray_commit(envs: Optional[Dict[str, str]] = None) -> str:
- if envs is None:
- envs = os.environ
- for key in [
- "RAY_WANT_COMMIT_IN_IMAGE",
- "COMMIT_TO_TEST",
- "BUILDKITE_COMMIT",
- ]:
- commit = envs.get(key, "")
- if commit:
- return commit
- return ""
- def _download_dataplane_build_file() -> None:
- """
- Downloads the dataplane build file from S3.
- """
- s3 = boto3.client("s3")
- s3.download_file(
- Bucket=DATAPLANE_S3_BUCKET,
- Key=DATAPLANE_FILENAME,
- Filename=DATAPLANE_FILENAME,
- )
- with open(DATAPLANE_FILENAME, "rb") as build_context:
- digest = hashlib.sha256(build_context.read()).hexdigest()
- assert digest == DATAPLANE_DIGEST, "Mismatched dataplane digest found!"
- def _image_exist(image: str) -> bool:
- """
- Checks if the given image exists in Docker
- """
- p = subprocess.run(
- ["docker", "manifest", "inspect", image],
- stdout=sys.stderr,
- stderr=sys.stderr,
- )
- return p.returncode == 0
|