75 lines
2.3 KiB
Python
75 lines
2.3 KiB
Python
|
"""Select the image builds and GPU verifications for a scoped release tag."""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import json
|
||
|
|
import re
|
||
|
|
from dataclasses import dataclass
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
|
||
|
|
GPU_GROUPS = {
|
||
|
|
"nvidia-a100": ("Nvidia-A100",),
|
||
|
|
"cambricon-mlu370": ("Cambricon_mlu-370-x4", "Cambricon_mlu-370-x8"),
|
||
|
|
"metax-c500": ("MetaX_c-500",),
|
||
|
|
"vastai-va16": ("Vastai_va16",),
|
||
|
|
}
|
||
|
|
RELEASE_TAG = re.compile(
|
||
|
|
r"v[0-9]+\.[0-9]+\.[0-9]+-(nvidia-a100|cambricon-mlu370|metax-c500|vastai-va16|all)\Z"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass(frozen=True)
|
||
|
|
class ReleasePlan:
|
||
|
|
target: str
|
||
|
|
image_targets: tuple[str, ...]
|
||
|
|
gpu_types: tuple[str, ...]
|
||
|
|
verification_pairs: tuple[tuple[str, str], ...]
|
||
|
|
|
||
|
|
|
||
|
|
def plan_for_tag(tag: str) -> ReleasePlan:
|
||
|
|
match = RELEASE_TAG.fullmatch(tag)
|
||
|
|
if match is None:
|
||
|
|
raise ValueError(
|
||
|
|
"release tag must be vX.Y.Z-<target>, where target is "
|
||
|
|
"nvidia-a100, cambricon-mlu370, metax-c500, vastai-va16, or all"
|
||
|
|
)
|
||
|
|
target = match.group(1)
|
||
|
|
image_targets = tuple(GPU_GROUPS) if target == "all" else (target,)
|
||
|
|
verification_pairs = tuple(
|
||
|
|
(gpu, group) for group in image_targets for gpu in GPU_GROUPS[group]
|
||
|
|
)
|
||
|
|
gpu_types = tuple(gpu for gpu, _ in verification_pairs)
|
||
|
|
return ReleasePlan(target, image_targets, gpu_types, verification_pairs)
|
||
|
|
|
||
|
|
|
||
|
|
def write_gitea_env(plan: ReleasePlan, path: Path) -> None:
|
||
|
|
with path.open("a", encoding="utf-8") as env_file:
|
||
|
|
env_file.write(f"RELEASE_TARGET={plan.target}\n")
|
||
|
|
env_file.write(f"IMAGE_TARGETS={' '.join(plan.image_targets)}\n")
|
||
|
|
env_file.write(f"GPU_TYPES={' '.join(plan.gpu_types)}\n")
|
||
|
|
env_file.write(
|
||
|
|
"VERIFY_PAIRS="
|
||
|
|
+ " ".join(f"{gpu}:{group}" for gpu, group in plan.verification_pairs)
|
||
|
|
+ "\n"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
||
|
|
parser.add_argument("tag", help="scoped release tag, such as v0.3.0-metax-c500")
|
||
|
|
parser.add_argument("--gitea-env", type=Path, help="append the plan to GITEA_ENV")
|
||
|
|
args = parser.parse_args()
|
||
|
|
try:
|
||
|
|
plan = plan_for_tag(args.tag)
|
||
|
|
except ValueError as exc:
|
||
|
|
parser.error(str(exc))
|
||
|
|
if args.gitea_env:
|
||
|
|
write_gitea_env(plan, args.gitea_env)
|
||
|
|
print(json.dumps(plan.__dict__, ensure_ascii=False))
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|