wehub-resource-sync
194 lines
6.8 KiB
Python
194 lines
6.8 KiB
Python
import os
|
|
from datetime import datetime
|
|
from enum import Enum
|
|
from typing import Dict, List
|
|
|
|
from ci.ray_ci.configs import DEFAULT_ARCHITECTURE, DEFAULT_PYTHON_TAG_VERSION
|
|
from ci.ray_ci.linux_container import LinuxContainer
|
|
from ci.ray_ci.supported_images import (
|
|
get_architectures,
|
|
get_default,
|
|
get_platforms,
|
|
get_python_versions,
|
|
)
|
|
|
|
PLATFORMS_RAY = get_platforms("ray")
|
|
PLATFORMS_RAY_ML = get_platforms("ray-ml")
|
|
PLATFORMS_RAY_LLM = get_platforms("ray-llm")
|
|
GPU_PLATFORM = get_default("ray", "gpu_platform")
|
|
|
|
PYTHON_VERSIONS_RAY = get_python_versions("ray")
|
|
PYTHON_VERSIONS_RAY_ML = get_python_versions("ray-ml")
|
|
PYTHON_VERSIONS_RAY_LLM = get_python_versions("ray-llm")
|
|
ARCHITECTURES_RAY = get_architectures("ray")
|
|
ARCHITECTURES_RAY_ML = get_architectures("ray-ml")
|
|
ARCHITECTURES_RAY_LLM = get_architectures("ray-llm")
|
|
|
|
|
|
class RayType(str, Enum):
|
|
RAY = "ray"
|
|
RAY_EXTRA = "ray-extra"
|
|
RAY_ML = "ray-ml"
|
|
RAY_ML_EXTRA = "ray-ml-extra"
|
|
RAY_LLM = "ray-llm"
|
|
RAY_LLM_EXTRA = "ray-llm-extra"
|
|
|
|
|
|
RAY_REPO_MAP: Dict[str, str] = {
|
|
RayType.RAY.value: RayType.RAY.value,
|
|
RayType.RAY_ML.value: RayType.RAY_ML.value,
|
|
RayType.RAY_LLM.value: RayType.RAY_LLM.value,
|
|
RayType.RAY_EXTRA.value: RayType.RAY.value,
|
|
RayType.RAY_ML_EXTRA.value: RayType.RAY_ML.value,
|
|
RayType.RAY_LLM_EXTRA.value: RayType.RAY_LLM.value,
|
|
}
|
|
|
|
|
|
class DockerContainer(LinuxContainer):
|
|
"""
|
|
Container for building and publishing ray docker images
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
python_version: str,
|
|
platform: str,
|
|
image_type: str,
|
|
architecture: str = DEFAULT_ARCHITECTURE,
|
|
canonical_tag: str = None,
|
|
upload: bool = False,
|
|
) -> None:
|
|
assert "RAYCI_CHECKOUT_DIR" in os.environ, "RAYCI_CHECKOUT_DIR not set"
|
|
rayci_checkout_dir = os.environ["RAYCI_CHECKOUT_DIR"]
|
|
|
|
super().__init__(
|
|
"forge" if architecture == "x86_64" else "forge-aarch64",
|
|
python_version=python_version,
|
|
volumes=[
|
|
f"{rayci_checkout_dir}:/rayci",
|
|
"/var/run/docker.sock:/var/run/docker.sock",
|
|
],
|
|
)
|
|
|
|
if image_type in [RayType.RAY_ML, RayType.RAY_ML_EXTRA]:
|
|
assert python_version in PYTHON_VERSIONS_RAY_ML
|
|
assert platform in PLATFORMS_RAY_ML
|
|
assert architecture in ARCHITECTURES_RAY_ML
|
|
elif image_type in [RayType.RAY_LLM, RayType.RAY_LLM_EXTRA]:
|
|
assert python_version in PYTHON_VERSIONS_RAY_LLM
|
|
assert platform in PLATFORMS_RAY_LLM
|
|
assert architecture in ARCHITECTURES_RAY_LLM
|
|
else:
|
|
# ray or ray-extra
|
|
assert python_version in PYTHON_VERSIONS_RAY
|
|
assert platform in PLATFORMS_RAY
|
|
assert architecture in ARCHITECTURES_RAY
|
|
|
|
self.platform = platform
|
|
self.image_type = image_type
|
|
self.architecture = architecture
|
|
self.canonical_tag = canonical_tag
|
|
self.upload = upload
|
|
|
|
def _get_image_version_tags(self, external: bool) -> List[str]:
|
|
"""
|
|
Get version tags.
|
|
|
|
Args:
|
|
external: If True, return the external image tags. If False, return the
|
|
internal image tags.
|
|
"""
|
|
branch = os.environ.get("BUILDKITE_BRANCH", "")
|
|
sha_tag = os.environ["BUILDKITE_COMMIT"][:6]
|
|
rayci_build_id = os.environ["RAYCI_BUILD_ID"]
|
|
pr = os.environ.get("BUILDKITE_PULL_REQUEST", "false")
|
|
formatted_date = datetime.now().strftime("%y%m%d")
|
|
|
|
if branch == "master":
|
|
if external and os.environ.get("RAYCI_SCHEDULE") == "nightly":
|
|
return [f"nightly.{formatted_date}.{sha_tag}", "nightly"]
|
|
return [sha_tag, rayci_build_id]
|
|
|
|
if branch and branch.startswith("releases/"):
|
|
release_name = branch[len("releases/") :]
|
|
release_tag = f"{release_name}.{sha_tag}"
|
|
if external:
|
|
# Avoid saving build ID ones when saving it on public registries.
|
|
return [release_tag]
|
|
return [release_tag, rayci_build_id]
|
|
|
|
if pr != "false":
|
|
return [f"pr-{pr}.{sha_tag}", rayci_build_id]
|
|
|
|
return [sha_tag, rayci_build_id]
|
|
|
|
def _get_canonical_tag(self) -> str:
|
|
# The canonical tag is the first tag in the list of tags. The list of tag is
|
|
# never empty because the image is always tagged with at least the sha tag.
|
|
#
|
|
# The canonical tag is the most complete tag with no abbreviation,
|
|
# e.g. sha-pyversion-platform
|
|
return self.canonical_tag if self.canonical_tag else self._get_image_tags()[0]
|
|
|
|
def _get_python_version_tag(self) -> str:
|
|
return f"-py{self.python_version.replace('.', '')}" # 3.x -> py3x
|
|
|
|
def _get_platform_tag(self) -> str:
|
|
if self.platform == "cpu":
|
|
return "-cpu"
|
|
if self.platform == "tpu":
|
|
return "-tpu"
|
|
versions = self.platform.split(".")
|
|
return f"-{versions[0]}{versions[1]}" # cu11.8.0 -> cu118
|
|
|
|
def _get_image_tags(self, external: bool = False) -> List[str]:
|
|
"""
|
|
Get image tags & alias for the container image.
|
|
|
|
An image tag is composed by ray version tag, python version and platform.
|
|
See https://docs.ray.io/en/latest/ray-overview/installation.html for
|
|
more information on the image tags.
|
|
|
|
Args:
|
|
external: If True, return the external image tags. If False, return the
|
|
internal image tags.
|
|
"""
|
|
|
|
versions = self._get_image_version_tags(external)
|
|
|
|
platforms = [self._get_platform_tag()]
|
|
if self.platform == "cpu" and self.image_type in [
|
|
RayType.RAY,
|
|
RayType.RAY_EXTRA,
|
|
]:
|
|
# no tag is alias to cpu for ray image
|
|
platforms.append("")
|
|
elif self.platform == GPU_PLATFORM:
|
|
# gpu is alias to cu118 for ray image
|
|
platforms.append("-gpu")
|
|
if self.image_type in [RayType.RAY_ML, RayType.RAY_ML_EXTRA]:
|
|
# no tag is alias to gpu for ray-ml image
|
|
platforms.append("")
|
|
|
|
py_versions = [self._get_python_version_tag()]
|
|
if self.python_version == DEFAULT_PYTHON_TAG_VERSION:
|
|
py_versions.append("")
|
|
|
|
variation = ""
|
|
if self.image_type in [
|
|
RayType.RAY_EXTRA,
|
|
RayType.RAY_ML_EXTRA,
|
|
RayType.RAY_LLM_EXTRA,
|
|
]:
|
|
variation = "-extra"
|
|
|
|
tags = []
|
|
for version in versions:
|
|
for platform in platforms:
|
|
for py_version in py_versions:
|
|
tag = f"{version}{variation}{py_version}{platform}"
|
|
if self.architecture != DEFAULT_ARCHITECTURE:
|
|
tag += f"-{self.architecture}"
|
|
tags.append(tag)
|
|
return tags
|