1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889 |
- import click
- import os
- from pathlib import Path
- import boto3
- from anyscale.sdk.anyscale_client.models.create_cluster_environment import (
- CreateClusterEnvironment,
- )
- from anyscale.sdk.anyscale_client.sdk import AnyscaleSDK
- from ray_release.bazel import bazel_runfile
- from ray_release.byod.build import build_champagne_image
- from ray_release.logger import logger
- from ray_release.configs.global_config import init_global_config
- CHAMPAGNE_IMAGE_TYPES = [
- # python_version, image_type, cluster_env_name
- ("py38", "cpu", "ray-champagne-cpu"),
- ("py38", "gpu", "ray-champagne-gpu"),
- ]
- ANYSCALE_SECRET_ARM = (
- "arn:aws:secretsmanager:us-west-2:029272617770:secret:release-automation"
- )
- ANYSCALE_SECRETS = {
- "anyscale-staging-token20231008005227440600000001-JTgxb0": (
- "https://console.anyscale-staging.com"
- ),
- "anyscale-demos-FaVACh": "https://console.anyscale.com",
- }
- CONFIG_CHOICES = click.Choice(
- [x.name for x in (Path(__file__).parent.parent / "configs").glob("*.yaml")]
- )
- @click.option(
- "--global-config",
- default="oss_config.yaml",
- type=CONFIG_CHOICES,
- help="Global config to use for test execution.",
- )
- def main(global_config: str = "oss_config.yaml") -> None:
- """
- Builds the Ray champagne image.
- """
- init_global_config(
- bazel_runfile("release/ray_release/configs", global_config),
- )
- branch = os.environ.get("BRANCH_TO_TEST", os.environ["BUILDKITE_BRANCH"])
- commit = os.environ.get("COMMIT_TO_TEST", os.environ["BUILDKITE_COMMIT"])
- assert branch.startswith(
- "releases/"
- ), f"Champagne building only supported on release branch, found {branch}"
- ray_version = f"{branch[len('releases/') :]}.{commit[:6]}"
- for python_version, image_type, cluster_env_name in CHAMPAGNE_IMAGE_TYPES:
- logger.info(f"Building champagne image for {python_version} {image_type}")
- anyscale_image = build_champagne_image(
- ray_version,
- python_version,
- image_type,
- )
- logger.info(f"Updating cluster environment {cluster_env_name}")
- _build_champaign_cluster_environment(anyscale_image, cluster_env_name)
- def _build_champaign_cluster_environment(
- anyscale_image: str,
- cluster_env_name: str,
- ) -> None:
- boto = boto3.client("secretsmanager", region_name="us-west-2")
- for secret_name, host in ANYSCALE_SECRETS.items():
- logger.info(
- f"\tUpdating cluster environment for secret: {secret_name}, host: {host}"
- )
- os.environ["ANYSCALE_CLI_TOKEN"] = boto.get_secret_value(
- SecretId=f"{ANYSCALE_SECRET_ARM}/{secret_name}"
- )["SecretString"]
- AnyscaleSDK(host=host).build_cluster_environment(
- create_cluster_environment=CreateClusterEnvironment(
- name=cluster_env_name,
- config_json=dict(
- docker_image=anyscale_image,
- ray_version="nightly",
- env_vars={},
- ),
- ),
- )
- if __name__ == "__main__":
- main()
|