Scope engine image builds and verification by release target

release_plan.py

@@ -0,0 +1,74 @@
+"""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()