wehub-resource-sync
559 lines
17 KiB
Python
559 lines
17 KiB
Python
#!/usr/bin/env python3
|
|
# /// script
|
|
# requires-python = ">=3.9"
|
|
# dependencies = []
|
|
# ///
|
|
"""
|
|
Build Ray Docker images locally using raymake.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
from ci.build.build_common import (
|
|
BuildError,
|
|
detect_host_arch,
|
|
find_ray_root,
|
|
get_git_commit,
|
|
log,
|
|
parse_file,
|
|
)
|
|
|
|
DEFAULT_ARCHITECTURE = "x86_64"
|
|
RAYCI_REGISTRY = "cr.ray.io/rayproject"
|
|
|
|
# Maps image type to its base name (YAML config key and Docker Hub repo).
|
|
_BASE_TYPE_MAP: dict[str, str] = {
|
|
"ray": "ray",
|
|
"ray-extra": "ray",
|
|
"ray-llm": "ray-llm",
|
|
"ray-llm-extra": "ray-llm",
|
|
}
|
|
|
|
|
|
def _load_ray_images() -> dict:
|
|
path = find_ray_root() / "ray-images.json"
|
|
if not path.exists():
|
|
raise BuildError(f"Missing {path}")
|
|
return json.loads(path.read_text())
|
|
|
|
|
|
def _build_image_type_config() -> dict[str, dict]:
|
|
"""Build IMAGE_TYPE_CONFIG from ray-images.yaml."""
|
|
raw = _load_ray_images()
|
|
config: dict[str, dict] = {}
|
|
for image_type, yaml_key in _BASE_TYPE_MAP.items():
|
|
yaml_cfg = raw[yaml_key]
|
|
defaults = yaml_cfg["defaults"]
|
|
platforms = yaml_cfg["platforms"]
|
|
config[image_type] = {
|
|
"python_versions": yaml_cfg["python"],
|
|
"platforms": platforms,
|
|
"architectures": yaml_cfg["architectures"],
|
|
"default_python": defaults["python"],
|
|
"default_platform": "cpu"
|
|
if "cpu" in platforms
|
|
else defaults["gpu_platform"],
|
|
}
|
|
return config
|
|
|
|
|
|
IMAGE_TYPE_CONFIG: dict[str, dict] = _build_image_type_config()
|
|
|
|
|
|
class RayImageError(Exception):
|
|
"""Raised when a RayImage field combination is invalid."""
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RayImage:
|
|
"""Immutable identity of a Ray Docker image variant."""
|
|
|
|
image_type: str
|
|
python_version: str
|
|
platform: str
|
|
architecture: str = DEFAULT_ARCHITECTURE
|
|
|
|
@property
|
|
def wanda_image_name(self) -> str:
|
|
"""Wanda output image name (without registry prefix)."""
|
|
if self.platform == "cpu":
|
|
return f"{self.image_type}-py{self.python_version}-cpu{self.arch_suffix}"
|
|
return f"{self.image_type}-py{self.python_version}-{self.platform}{self.arch_suffix}"
|
|
|
|
@property
|
|
def arch_suffix(self) -> str:
|
|
if self.architecture == DEFAULT_ARCHITECTURE:
|
|
return ""
|
|
return f"-{self.architecture}"
|
|
|
|
@property
|
|
def repo(self) -> str:
|
|
return _BASE_TYPE_MAP[self.image_type]
|
|
|
|
@property
|
|
def variation_suffix(self) -> str:
|
|
if self.image_type.endswith("-extra"):
|
|
return "-extra"
|
|
return ""
|
|
|
|
@property
|
|
def platform_class(self) -> str:
|
|
"""Returns 'cpu', 'cuda', 'tpu', etc."""
|
|
if self.platform == "cpu":
|
|
return "cpu"
|
|
if self.platform == "tpu":
|
|
return "tpu"
|
|
if self.platform.startswith("cu"):
|
|
return "cuda"
|
|
raise RayImageError(f"Could not determine platform class for {self.platform!r}")
|
|
|
|
def get_wanda_spec_path(self) -> str:
|
|
"""Returns the path to the wanda spec file for this image."""
|
|
key: WandaSpecKey = (self.image_type, self.platform_class)
|
|
try:
|
|
return WANDA_SPEC_PATHS[key]
|
|
except KeyError:
|
|
raise RayImageError(
|
|
f"No wanda spec for image_type={self.image_type!r}, "
|
|
f"platform={self.platform!r}"
|
|
) from None
|
|
|
|
def validate(self) -> None:
|
|
if self.image_type not in IMAGE_TYPE_CONFIG:
|
|
valid = ", ".join(IMAGE_TYPE_CONFIG)
|
|
raise RayImageError(
|
|
f"Unknown image type {self.image_type!r}. Valid types: {valid}"
|
|
)
|
|
cfg = IMAGE_TYPE_CONFIG[self.image_type]
|
|
if self.python_version not in cfg["python_versions"]:
|
|
raise RayImageError(
|
|
f"Invalid python version {self.python_version} "
|
|
f"for {self.image_type}. "
|
|
f"Valid versions: {', '.join(cfg['python_versions'])}"
|
|
)
|
|
if self.platform not in cfg["platforms"]:
|
|
raise RayImageError(
|
|
f"Invalid platform {self.platform} "
|
|
f"for {self.image_type}. "
|
|
f"Valid platforms: {', '.join(cfg['platforms'])}"
|
|
)
|
|
if self.architecture not in cfg["architectures"]:
|
|
raise RayImageError(
|
|
f"Invalid architecture {self.architecture} "
|
|
f"for {self.image_type}. "
|
|
f"Valid architectures: {', '.join(cfg['architectures'])}"
|
|
)
|
|
|
|
|
|
# (image_type, platform_class) where platform_class is "cpu" or "cuda".
|
|
WandaSpecKey = tuple[str, str]
|
|
|
|
WANDA_SPEC_PATHS: dict[WandaSpecKey, str] = {
|
|
("ray", "cpu"): "ci/docker/ray-image-cpu.wanda.yaml",
|
|
("ray", "cuda"): "ci/docker/ray-image-cuda.wanda.yaml",
|
|
("ray", "tpu"): "ci/docker/ray-image-tpu.wanda.yaml",
|
|
("ray-extra", "cpu"): "ci/docker/ray-extra-image-cpu.wanda.yaml",
|
|
("ray-extra", "cuda"): "ci/docker/ray-extra-image-cuda.wanda.yaml",
|
|
("ray-llm", "cuda"): "ci/docker/ray-llm-image-cuda.wanda.yaml",
|
|
("ray-llm-extra", "cuda"): "ci/docker/ray-llm-extra-image-cuda.wanda.yaml",
|
|
}
|
|
|
|
SUPPORTED_IMAGE_TYPES: list[str] = list(dict.fromkeys(t for t, _ in WANDA_SPEC_PATHS))
|
|
REGISTRY_PREFIX = "cr.ray.io/rayproject/"
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class ImageBuildConfig:
|
|
ray_image: RayImage
|
|
ray_root: Path
|
|
raymake_version: str
|
|
manylinux_version: str
|
|
ray_version: str
|
|
commit: str
|
|
wanda_spec_path: str
|
|
wanda_image_tag: str
|
|
nightly_alias: str | None
|
|
build_env: dict[str, str]
|
|
|
|
@classmethod
|
|
def from_args(
|
|
cls,
|
|
image_type: str,
|
|
python_version: str,
|
|
image_platform: str,
|
|
*,
|
|
no_java: bool = False,
|
|
no_dashboard: bool = False,
|
|
) -> ImageBuildConfig:
|
|
ray_image = RayImage(
|
|
image_type=image_type,
|
|
python_version=python_version,
|
|
platform=image_platform,
|
|
architecture=cls._detect_host_arch(),
|
|
)
|
|
try:
|
|
ray_image.validate()
|
|
ray_image.get_wanda_spec_path()
|
|
except RayImageError as e:
|
|
raise BuildError(str(e)) from e
|
|
|
|
root = find_ray_root()
|
|
manylinux_version = parse_file(
|
|
root / "rayci.env", r'MANYLINUX_VERSION=["\']?([^"\'\s]+)'
|
|
)
|
|
ray_version = parse_file(root / "rayci.env", r'RAY_VERSION=["\']?([^"\'\s]+)')
|
|
commit = get_git_commit(root)
|
|
|
|
cfg = IMAGE_TYPE_CONFIG[image_type]
|
|
nightly_alias = (
|
|
f"{REGISTRY_PREFIX}{ray_image.repo}:nightly{ray_image.variation_suffix}"
|
|
if (
|
|
ray_image.python_version == cfg["default_python"]
|
|
and ray_image.platform == cfg["default_platform"]
|
|
and ray_image.architecture == DEFAULT_ARCHITECTURE
|
|
)
|
|
else None
|
|
)
|
|
|
|
java_image = (
|
|
"scratch"
|
|
if no_java
|
|
else f"{RAYCI_REGISTRY}/ray-java-build{ray_image.arch_suffix}"
|
|
)
|
|
dashboard_image = (
|
|
"scratch"
|
|
if no_dashboard
|
|
else f"{RAYCI_REGISTRY}/ray-dashboard{ray_image.arch_suffix}"
|
|
)
|
|
|
|
build_env: dict[str, str] = {
|
|
"PYTHON_VERSION": ray_image.python_version,
|
|
"MANYLINUX_VERSION": manylinux_version,
|
|
"HOSTTYPE": ray_image.architecture,
|
|
"ARCH_SUFFIX": ray_image.arch_suffix,
|
|
"JDK_SUFFIX": "" if no_java else "-jdk",
|
|
"BUILDKITE_COMMIT": commit,
|
|
"RAY_VERSION": ray_version,
|
|
"IS_LOCAL_BUILD": "true",
|
|
"IMAGE_TYPE": ray_image.repo,
|
|
"RAY_JAVA_IMAGE": java_image,
|
|
"RAY_DASHBOARD_IMAGE": dashboard_image,
|
|
}
|
|
if ray_image.platform.startswith("cu"):
|
|
build_env["CUDA_VERSION"] = ray_image.platform.removeprefix("cu")
|
|
|
|
return cls(
|
|
ray_image=ray_image,
|
|
ray_root=root,
|
|
raymake_version=(root / ".rayciversion").read_text().strip(),
|
|
manylinux_version=manylinux_version,
|
|
ray_version=ray_version,
|
|
commit=commit,
|
|
wanda_spec_path=ray_image.get_wanda_spec_path(),
|
|
wanda_image_tag=f"{REGISTRY_PREFIX}{ray_image.wanda_image_name}",
|
|
nightly_alias=nightly_alias,
|
|
build_env=build_env,
|
|
)
|
|
|
|
@staticmethod
|
|
def _detect_host_arch() -> str:
|
|
return detect_host_arch()
|
|
|
|
|
|
class ImageBuilder:
|
|
def __init__(self, config: ImageBuildConfig):
|
|
self.config = config
|
|
|
|
def build(self) -> str:
|
|
"""Build the image and return the primary image tag."""
|
|
if not shutil.which("raymake"):
|
|
raise BuildError("raymake not found. Run via ./build-image.sh")
|
|
|
|
log.info("Build configuration:")
|
|
summary = {
|
|
"Image Type": self.config.ray_image.image_type,
|
|
"Python": self.config.ray_image.python_version,
|
|
"Platform": self.config.ray_image.platform,
|
|
"Arch": self.config.ray_image.architecture,
|
|
"Commit": self.config.commit,
|
|
"Raymake": self.config.raymake_version,
|
|
"Ray Version": self.config.ray_version,
|
|
"Wanda Spec": self.config.wanda_spec_path,
|
|
}
|
|
print("-" * 50)
|
|
for k, v in summary.items():
|
|
print(f"{k:<12}: {v}")
|
|
print("-" * 50)
|
|
|
|
cmd = [
|
|
"raymake",
|
|
str(self.config.wanda_spec_path),
|
|
]
|
|
|
|
log.info(f"Running raymake: {self.config.wanda_spec_path}")
|
|
log.info(f"Build environment: {self.config.build_env}")
|
|
proc = subprocess.Popen(
|
|
cmd,
|
|
cwd=self.config.ray_root,
|
|
env={**os.environ, **self.config.build_env},
|
|
)
|
|
try:
|
|
if proc.wait() != 0:
|
|
raise BuildError("raymake failed")
|
|
except KeyboardInterrupt:
|
|
try:
|
|
proc.wait(timeout=2)
|
|
except subprocess.TimeoutExpired:
|
|
proc.kill()
|
|
raise
|
|
|
|
image_tag = self.config.wanda_image_tag
|
|
log.info(f"Built image: {image_tag}")
|
|
|
|
alias = self.config.nightly_alias
|
|
if alias:
|
|
self._docker_tag(image_tag, alias)
|
|
log.info(f"Tagged alias: {alias}")
|
|
|
|
return image_tag
|
|
|
|
@staticmethod
|
|
def _docker_tag(source: str, dest: str) -> None:
|
|
result = subprocess.run(
|
|
["docker", "tag", source, dest],
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
if result.returncode != 0:
|
|
raise BuildError(f"docker tag failed: {result.stderr.strip()}")
|
|
|
|
|
|
def _image_type_table() -> str:
|
|
"""Build an ASCII table of all image types from IMAGE_TYPE_CONFIG."""
|
|
|
|
def _tag_default(items, default):
|
|
return [f"{v} (default)" if v == default else v for v in items]
|
|
|
|
def _wrap(items, width):
|
|
"""Join items with ', ' and wrap into lines that fit in *width*."""
|
|
lines: list[str] = []
|
|
cur = ""
|
|
for item in items:
|
|
entry = item if not cur else f", {item}"
|
|
if cur and len(cur) + len(entry) > width:
|
|
lines.append(cur)
|
|
cur = item
|
|
else:
|
|
cur += entry
|
|
if cur:
|
|
lines.append(cur)
|
|
return lines
|
|
|
|
headers = ("IMAGE TYPE", "PYTHON VERSIONS (-p)", "PLATFORMS (--platform)")
|
|
plat_width = 52 # wrap platforms to this width
|
|
|
|
def _build_rows(type_list):
|
|
rows: list[tuple[str, str, list[str]]] = []
|
|
for name in type_list:
|
|
cfg = IMAGE_TYPE_CONFIG[name]
|
|
py_cell = ", ".join(
|
|
_tag_default(cfg["python_versions"], cfg["default_python"])
|
|
)
|
|
plat_lines = _wrap(
|
|
_tag_default(cfg["platforms"], cfg["default_platform"]),
|
|
plat_width,
|
|
)
|
|
rows.append((name, py_cell, plat_lines))
|
|
return rows
|
|
|
|
rows = _build_rows(SUPPORTED_IMAGE_TYPES)
|
|
|
|
# Compute column widths.
|
|
widths = [len(h) for h in headers]
|
|
for label, py_cell, plat_lines in rows:
|
|
widths[0] = max(widths[0], len(label))
|
|
widths[1] = max(widths[1], len(py_cell))
|
|
widths[2] = max(widths[2], *(len(l) for l in plat_lines))
|
|
|
|
sep = " "
|
|
|
|
def fmt(cells):
|
|
return sep.join(c.ljust(widths[i]) for i, c in enumerate(cells))
|
|
|
|
lines = [fmt(headers), sep.join("-" * w for w in widths)]
|
|
for label, py_cell, plat_lines in rows:
|
|
lines.append(fmt((label, py_cell, plat_lines[0])))
|
|
for cont in plat_lines[1:]:
|
|
lines.append(fmt(("", "", cont)))
|
|
|
|
return "\n".join(lines)
|
|
|
|
|
|
def _nightly_table() -> str:
|
|
"""Build a table mapping build commands to nightly docker run commands."""
|
|
headers = ("BUILD", "RUN")
|
|
rows = []
|
|
for name in SUPPORTED_IMAGE_TYPES:
|
|
cfg = IMAGE_TYPE_CONFIG[name]
|
|
img = RayImage(name, cfg["default_python"], cfg["default_platform"])
|
|
tag = f"{REGISTRY_PREFIX}{img.repo}:nightly{img.variation_suffix}"
|
|
rows.append((f"./build-image.sh {name}", f"docker run -it {tag}"))
|
|
|
|
widths = [len(h) for h in headers]
|
|
for row in rows:
|
|
for i, cell in enumerate(row):
|
|
widths[i] = max(widths[i], len(cell))
|
|
|
|
sep = " "
|
|
|
|
def fmt(cells):
|
|
return sep.join(c.ljust(widths[i]) for i, c in enumerate(cells))
|
|
|
|
lines = [fmt(headers), sep.join("-" * w for w in widths)]
|
|
lines.extend(fmt(row) for row in rows)
|
|
return "\n".join(lines)
|
|
|
|
|
|
def _build_examples() -> str:
|
|
"""Generate non-default example commands from IMAGE_TYPE_CONFIG."""
|
|
ray_cfg = IMAGE_TYPE_CONFIG["ray"]
|
|
ray_py = ray_cfg["default_python"]
|
|
ray_plat = ray_cfg["default_platform"]
|
|
alt_py = next((v for v in ray_cfg["python_versions"] if v != ray_py), None)
|
|
alt_plat = next((p for p in ray_cfg["platforms"] if p != ray_plat), None)
|
|
|
|
examples: list[tuple[str, str]] = []
|
|
if alt_py:
|
|
examples.append(
|
|
(f"ray -p {alt_py}", f"ray {ray_plat} py{alt_py}"),
|
|
)
|
|
if alt_plat:
|
|
examples.append(
|
|
(f"ray --platform {alt_plat}", f"ray {alt_plat} py{ray_py}"),
|
|
)
|
|
if alt_py and alt_plat:
|
|
examples.append(
|
|
(f"ray -p {alt_py} --platform {alt_plat}", f"ray {alt_plat} py{alt_py}"),
|
|
)
|
|
|
|
if not examples:
|
|
return ""
|
|
|
|
cmd_prefix = " ./build-image.sh "
|
|
cmds = [cmd_prefix + args for args, _ in examples]
|
|
max_cmd = max(len(c) for c in cmds)
|
|
lines = []
|
|
for cmd, (_, comment) in zip(cmds, examples):
|
|
lines.append(f"{cmd:<{max_cmd}} # {comment}")
|
|
return "\n".join(lines)
|
|
|
|
|
|
def main():
|
|
epilog_parts = [
|
|
"Supported image types:",
|
|
_image_type_table(),
|
|
"",
|
|
"Nightly images (built with defaults):",
|
|
_nightly_table(),
|
|
"",
|
|
"Examples:",
|
|
_build_examples(),
|
|
]
|
|
|
|
parser = argparse.ArgumentParser(
|
|
prog="./build-image.sh",
|
|
description="Build Ray Docker images locally using raymake.",
|
|
epilog="\n".join(epilog_parts),
|
|
formatter_class=argparse.RawDescriptionHelpFormatter,
|
|
)
|
|
|
|
parser.add_argument(
|
|
"image_type",
|
|
choices=SUPPORTED_IMAGE_TYPES,
|
|
metavar="IMAGE_TYPE",
|
|
help=f"Image type: {', '.join(SUPPORTED_IMAGE_TYPES)}",
|
|
)
|
|
parser.add_argument(
|
|
"-p",
|
|
"--python-version",
|
|
default=None,
|
|
metavar="VERSION",
|
|
help="Python version (default depends on image type; see table below)",
|
|
)
|
|
parser.add_argument(
|
|
"--platform",
|
|
default=None,
|
|
metavar="PLATFORM",
|
|
help="Target platform (default depends on image type; see table below)",
|
|
)
|
|
parser.add_argument(
|
|
"--list-platforms",
|
|
action="store_true",
|
|
help="List valid platforms for the given image type and exit",
|
|
)
|
|
parser.add_argument(
|
|
"--no-java",
|
|
action="store_true",
|
|
default=False,
|
|
help="Build wheel without Java JARs",
|
|
)
|
|
parser.add_argument(
|
|
"--no-dashboard",
|
|
action="store_true",
|
|
default=False,
|
|
help="Build wheel without the dashboard",
|
|
)
|
|
|
|
if len(sys.argv) == 1:
|
|
parser.print_help()
|
|
sys.exit(0)
|
|
|
|
args = parser.parse_args()
|
|
|
|
cfg = IMAGE_TYPE_CONFIG[args.image_type]
|
|
|
|
if args.list_platforms:
|
|
print(f"Valid platforms for {args.image_type}:")
|
|
for p in cfg["platforms"]:
|
|
marker = " (default)" if p == cfg["default_platform"] else ""
|
|
print(f" {p}{marker}")
|
|
sys.exit(0)
|
|
|
|
python_version = args.python_version or cfg["default_python"]
|
|
platform_choice = args.platform or cfg["default_platform"]
|
|
|
|
try:
|
|
config = ImageBuildConfig.from_args(
|
|
args.image_type,
|
|
python_version,
|
|
platform_choice,
|
|
no_java=args.no_java,
|
|
no_dashboard=args.no_dashboard,
|
|
)
|
|
builder = ImageBuilder(config)
|
|
image_tag = builder.build()
|
|
|
|
print()
|
|
log.info("Success!")
|
|
print(f" docker run -it {image_tag}")
|
|
alias = config.nightly_alias
|
|
if alias:
|
|
print(f" docker run -it {alias}")
|
|
except BuildError as e:
|
|
log.error(e)
|
|
sys.exit(1)
|
|
except KeyboardInterrupt:
|
|
sys.exit(130)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|