From 9bd8b7cc1e535ee2b6b358967eeb1ec3d1758630 Mon Sep 17 00:00:00 2001 From: Codex Bot Date: Fri, 2 Oct 2026 01:51:33 +0800 Subject: [PATCH] Add AutoDL Pro runner controller --- .gitea/workflows/ci.yaml | 28 ++ .gitea/workflows/manage.yaml | 57 ++++ .gitignore | 5 + README.md | 59 ++++ controller.py | 545 +++++++++++++++++++++++++++++++++++ tests/test_controller.py | 62 ++++ 6 files changed, 756 insertions(+) create mode 100644 .gitea/workflows/ci.yaml create mode 100644 .gitea/workflows/manage.yaml create mode 100644 .gitignore create mode 100644 README.md create mode 100644 controller.py create mode 100644 tests/test_controller.py diff --git a/.gitea/workflows/ci.yaml b/.gitea/workflows/ci.yaml new file mode 100644 index 0000000..6ec2c72 --- /dev/null +++ b/.gitea/workflows/ci.yaml @@ -0,0 +1,28 @@ +name: Controller CI + +on: + push: + branches: [main] + +jobs: + test: + runs-on: ces-control + timeout-minutes: 10 + env: + GITEA_JOB_TOKEN: ${{ secrets.GITEA_TOKEN }} + steps: + - name: Checkout controller source + run: | + git init . + git remote add origin "$GITHUB_SERVER_URL/$GITHUB_REPOSITORY.git" + git -c http.extraHeader="Authorization: token $GITEA_JOB_TOKEN" fetch --depth 1 origin "$GITHUB_SHA" + git checkout --detach FETCH_HEAD + + - name: Verify controller + run: | + python3 -m py_compile controller.py + python3 -m unittest discover -s tests -v + + - name: Confirm no paid action ran + run: | + echo "Controller CI only validates code; it never creates an AutoDL instance." diff --git a/.gitea/workflows/manage.yaml b/.gitea/workflows/manage.yaml new file mode 100644 index 0000000..651dd28 --- /dev/null +++ b/.gitea/workflows/manage.yaml @@ -0,0 +1,57 @@ +name: Manage AutoDL GPU runner + +on: + workflow_dispatch: + inputs: + operation: + description: discover, provision, status, or shutdown + required: true + default: discover + image_uuid: + description: Required only for provision + required: false + default: "" + confirm_spend: + description: Type CREATE to provision a billed instance + required: false + default: "" + +jobs: + manage: + runs-on: ces-control + timeout-minutes: 180 + env: + AUTODL_API_TOKEN: ${{ secrets.AUTODL_API_TOKEN }} + TRAINING_RUNNER_REGISTRATION_TOKEN: ${{ secrets.TRAINING_RUNNER_REGISTRATION_TOKEN }} + GITEA_JOB_TOKEN: ${{ secrets.GITEA_TOKEN }} + OPERATION: ${{ inputs.operation }} + IMAGE_UUID: ${{ inputs.image_uuid }} + CONFIRM_SPEND: ${{ inputs.confirm_spend }} + steps: + - name: Checkout controller source + run: | + git init . + git remote add origin "$GITHUB_SERVER_URL/$GITHUB_REPOSITORY.git" + git -c http.extraHeader="Authorization: token $GITEA_JOB_TOKEN" fetch --depth 1 origin "$GITHUB_SHA" + git checkout --detach FETCH_HEAD + + - name: Validate requested operation + run: | + case "$OPERATION" in + discover|provision|status|shutdown) ;; + *) echo "Unsupported operation" >&2; exit 2 ;; + esac + + - name: Manage AutoDL Pro instance + run: | + python3 controller.py "$OPERATION" \ + --image-uuid "$IMAGE_UUID" \ + --confirm-spend "$CONFIRM_SPEND" + + - name: Controller host details + if: always() + run: | + echo "host=$(hostname)" + date -Iseconds + test -f /root/.local/state/ces-autodl-controller/state.json && \ + python3 -c 'import json; p="/root/.local/state/ces-autodl-controller/state.json"; d=json.load(open(p)); d.pop("root_password", None); print(json.dumps(d, indent=2))' || true diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..10f4262 --- /dev/null +++ b/.gitignore @@ -0,0 +1,5 @@ +__pycache__/ +*.pyc +.runner +config.yaml +state.json diff --git a/README.md b/README.md new file mode 100644 index 0000000..ef52e5e --- /dev/null +++ b/README.md @@ -0,0 +1,59 @@ +# AutoDL GPU Runner Controller + +该仓库只管理算力，不包含 Evidence Extractor 研究代码。控制流程运行在阿里云的常驻 +`ces-control:host` Gitea Runner 上，并通过 AutoDL Pro API 创建和关闭 GPU 实例。 + +## 边界 + +- 本仓库：AutoDL Pro 生命周期、模型传输、GPU Runner 注册和产物回收。 +- 阿里云：常驻控制 Runner、SigLIP2 权重以及训练产物的持久存储。 +- AutoDL：临时 GPU 执行节点。 +- 研究仓库：训练代码、配置和 Gitea Actions workflow 的唯一来源。 + +## Repository Secrets + +在本仓库 `Settings → Actions → Secrets` 配置： + +- `AUTODL_API_TOKEN`：AutoDL 控制台的开发者 Token。 +- `TRAINING_RUNNER_REGISTRATION_TOKEN`：研究仓库 `Create new Runner` 页面显示的 + Registration Token。 + +Secret 只注入阿里云控制任务，不落盘、不打印。AutoDL SSH 私钥保存在阿里云 +`/root/.config/ces-controller/autodl_ed25519`，不进入 Gitea。 + +## 首次启动 + +阿里云需要先作为本仓库的常驻 Runner 注册一次，且只使用标签： + +```text +ces-control:host +``` + +该 bootstrap 使用的是“本控制器仓库”的 Runner Registration Token，不是 Secret 中用于 +AutoDL GPU Runner 的 `TRAINING_RUNNER_REGISTRATION_TOKEN`。 + +Runner online 后，第一次 push 只执行 `Controller CI`，不会创建实例或产生 AutoDL 费用。 + +## 管理 workflow + +从 Actions 手动运行 `Manage AutoDL GPU runner`： + +- `discover`：列出可供 API 使用的私有镜像，不产生 GPU 费用； +- `provision`：创建 PRO6000、传输模型并注册 GPU Runner；必须填写镜像 UUID，并把确认值 + 填为 `CREATE`； +- `status`：查询实例和 Runner 进程状态； +- `shutdown`：回收训练产物并关闭实例，不释放系统盘。 + +`provision` 在创建实例前会检查 SigLIP2 的两个权重分片是否完整、AutoDL SSH 私钥权限是否为 +`0600`，并预下载及校验 `gitea-runner`。任何配置失败都会尝试关闭刚创建的实例，避免继续计费。 + +控制器固定使用官方 `gitea-runner 3.5.0` Linux AMD64 二进制及其 SHA-256，下载和校验均在 +创建付费实例之前完成。 + +AutoDL GPU Runner 注册为 host runner，并提供： + +```text +ubuntu-latest:host +ces-ci:host +gpu:host +``` diff --git a/controller.py b/controller.py new file mode 100644 index 0000000..0f0bc6a --- /dev/null +++ b/controller.py @@ -0,0 +1,545 @@ +#!/usr/bin/env python3 +"""Run AutoDL Pro GPU workers from a persistent Alibaba control host.""" + +from __future__ import annotations + +import argparse +import hashlib +import json +import os +import shlex +import stat +import subprocess +import sys +import tempfile +import time +import urllib.error +import urllib.request +from dataclasses import dataclass +from pathlib import Path +from typing import Any + + +API_BASE = "https://api.autodl.com" +RUNNER_VERSION = "3.5.0" +RUNNER_URL = ( + "https://dl.gitea.com/gitea-runner/3.5.0/gitea-runner-3.5.0-linux-amd64" +) +RUNNER_SHA256 = "16dc5c89ad5ac5496369642f01ed29c9bb130a2512e16a8cc3b400d9d0bfea73" +RUNNER_LABELS = ["ubuntu-latest:host", "ces-ci:host", "gpu:host"] +MODEL_FILES = [ + "config.json", + "model.safetensors.index.json", + "model-00001-of-00002.safetensors", + "model-00002-of-00002.safetensors", + "preprocessor_config.json", + "tokenizer.json", + "tokenizer.model", + "tokenizer_config.json", +] + + +class ControllerError(RuntimeError): + pass + + +class AutoDLClient: + def __init__(self, token: str): + if not token: + raise ControllerError("AUTODL_API_TOKEN is missing") + self.token = token + + def request(self, method: str, path: str, body: dict[str, Any]) -> Any: + request = urllib.request.Request( + f"{API_BASE}{path}", + data=json.dumps(body).encode(), + method=method, + headers={"Authorization": self.token, "Content-Type": "application/json"}, + ) + try: + with urllib.request.urlopen(request, timeout=45) as response: + result = json.loads(response.read().decode()) + except urllib.error.HTTPError as exc: + message = exc.read().decode(errors="replace")[:500] + raise ControllerError(f"AutoDL API {method} {path} returned {exc.code}: {message}") from exc + except (urllib.error.URLError, json.JSONDecodeError) as exc: + raise ControllerError(f"AutoDL API {method} {path} failed: {exc}") from exc + if result.get("code") != "Success": + raise ControllerError( + f"AutoDL API rejected {method} {path}: " + f"code={result.get('code')!r}, msg={result.get('msg')!r}" + ) + return result.get("data") + + def list_images(self) -> list[dict[str, Any]]: + data = self.request( + "POST", + "/api/v1/dev/instance/pro/image/private/list", + {"page_index": 1, "page_size": 100}, + ) + return list(data.get("list", [])) + + def create(self, spec: dict[str, Any]) -> str: + return str(self.request("POST", "/api/v1/dev/instance/pro/create", spec)) + + def status(self, instance_uuid: str) -> str: + return str( + self.request( + "GET", "/api/v1/dev/instance/pro/status", {"instance_uuid": instance_uuid} + ) + ) + + def snapshot(self, instance_uuid: str) -> dict[str, Any]: + return dict( + self.request( + "GET", "/api/v1/dev/instance/pro/snapshot", {"instance_uuid": instance_uuid} + ) + ) + + def power_off(self, instance_uuid: str) -> None: + self.request( + "POST", "/api/v1/dev/instance/pro/power_off", {"instance_uuid": instance_uuid} + ) + + +@dataclass(frozen=True) +class Paths: + model_source: Path + model_destination: str + artifact_source: str + artifact_destination: Path + ssh_key: Path + known_hosts: Path + state_file: Path + runner_cache: Path + + +@dataclass(frozen=True) +class SshTarget: + host: str + port: int + key: Path + known_hosts: Path + + @classmethod + def from_snapshot(cls, snapshot: dict[str, Any], paths: Paths) -> "SshTarget": + host = str(snapshot.get("proxy_host", "")).strip() + port = int(snapshot.get("ssh_port", 0)) + if not host or port <= 0: + raise ControllerError("AutoDL did not return a usable SSH endpoint") + return cls(host, port, paths.ssh_key, paths.known_hosts) + + def ssh_options(self) -> list[str]: + return [ + "-F", + "/dev/null", + "-o", + "BatchMode=yes", + "-o", + "IdentitiesOnly=yes", + "-o", + "StrictHostKeyChecking=accept-new", + "-o", + f"UserKnownHostsFile={self.known_hosts}", + "-i", + str(self.key), + "-p", + str(self.port), + ] + + def run(self, command: str, check: bool = True) -> subprocess.CompletedProcess[str]: + return subprocess.run( + ["ssh", *self.ssh_options(), f"root@{self.host}", command], + check=check, + text=True, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + ) + + def copy_to(self, source: Path, destination: str) -> None: + options = self.ssh_options() + options[options.index("-p")] = "-P" + subprocess.run( + ["scp", *options, str(source), f"root@{self.host}:{destination}"], check=True + ) + + def sync_to(self, source: Path, destination: str) -> None: + transport = shlex.join(["ssh", *self.ssh_options()]) + subprocess.run( + [ + "rsync", + "-ah", + "--partial", + "--append-verify", + "-e", + transport, + f"{source}/", + f"root@{self.host}:{destination.rstrip('/')}/", + ], + check=True, + ) + + def sync_from(self, source: str, destination: Path) -> None: + destination.mkdir(parents=True, exist_ok=True) + transport = shlex.join(["ssh", *self.ssh_options()]) + subprocess.run( + [ + "rsync", + "-ah", + "--partial", + "-e", + transport, + f"root@{self.host}:{source.rstrip('/')}/", + f"{destination}/", + ], + check=True, + ) + + +def default_paths() -> Paths: + return Paths( + model_source=Path("/data/ces-models/siglip2-giant"), + model_destination="/root/autodl-tmp/ces-models/siglip2-giant", + artifact_source="/root/autodl-tmp/ces-artifacts", + artifact_destination=Path("/data/ces-artifacts"), + ssh_key=Path("/root/.config/ces-controller/autodl_ed25519"), + known_hosts=Path("/root/.config/ces-controller/known_hosts"), + state_file=Path("/root/.local/state/ces-autodl-controller/state.json"), + runner_cache=Path(f"/var/cache/ces-controller/act_runner-{RUNNER_VERSION}"), + ) + + +def instance_spec(image_uuid: str, disk_gb: int, data_centers: list[str]) -> dict[str, Any]: + if not image_uuid or image_uuid.startswith("REPLACE_"): + raise ControllerError("a real AutoDL image UUID is required") + spec: dict[str, Any] = { + "req_gpu_amount": 1, + "expand_system_disk_by_gb": disk_gb, + "gpu_spec_uuid": "pro6000-p", + "image_uuid": image_uuid, + "cuda_v_from": 128, + "instance_name": "ces-gitea-gpu-runner", + } + if data_centers: + spec["data_center_list"] = data_centers + return spec + + +def validate_local(paths: Paths, runner_token: str) -> None: + if not runner_token: + raise ControllerError("TRAINING_RUNNER_REGISTRATION_TOKEN is missing") + missing = [ + name + for name in MODEL_FILES + if not (paths.model_source / name).is_file() + or (paths.model_source / name).stat().st_size == 0 + ] + if missing: + raise ControllerError(f"SigLIP2 download is incomplete: {missing}") + if not paths.ssh_key.is_file(): + raise ControllerError(f"AutoDL SSH key does not exist: {paths.ssh_key}") + if stat.S_IMODE(paths.ssh_key.stat().st_mode) & 0o077: + raise ControllerError(f"AutoDL SSH key must have mode 0600: {paths.ssh_key}") + + +def sha256(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as source: + while chunk := source.read(1024 * 1024): + digest.update(chunk) + return digest.hexdigest() + + +def ensure_runner(paths: Paths) -> Path: + if paths.runner_cache.is_file() and sha256(paths.runner_cache) == RUNNER_SHA256: + return paths.runner_cache + paths.runner_cache.parent.mkdir(parents=True, exist_ok=True) + with urllib.request.urlopen(RUNNER_URL, timeout=300) as response, tempfile.NamedTemporaryFile( + dir=paths.runner_cache.parent, delete=False + ) as temporary: + while chunk := response.read(1024 * 1024): + temporary.write(chunk) + downloaded = Path(temporary.name) + if sha256(downloaded) != RUNNER_SHA256: + downloaded.unlink(missing_ok=True) + raise ControllerError("act_runner checksum mismatch") + downloaded.chmod(0o755) + downloaded.replace(paths.runner_cache) + return paths.runner_cache + + +def write_state(paths: Paths, data: dict[str, Any]) -> None: + paths.state_file.parent.mkdir(parents=True, exist_ok=True) + temporary = paths.state_file.with_suffix(".tmp") + temporary.write_text(json.dumps(data, indent=2), encoding="utf-8") + temporary.chmod(0o600) + temporary.replace(paths.state_file) + + +def read_state(paths: Paths) -> dict[str, Any]: + if not paths.state_file.is_file(): + raise ControllerError(f"state file does not exist: {paths.state_file}") + return json.loads(paths.state_file.read_text(encoding="utf-8")) + + +def wait_running(client: AutoDLClient, instance_uuid: str, timeout: int = 1800) -> None: + deadline = time.monotonic() + timeout + last = "" + while time.monotonic() < deadline: + last = client.status(instance_uuid) + print(f"instance={instance_uuid} status={last}", flush=True) + if last == "running": + return + if last in {"error", "released"}: + break + time.sleep(10) + raise ControllerError(f"instance failed to reach running state; last={last!r}") + + +def wait_ssh(target: SshTarget, timeout: int = 600) -> None: + target.known_hosts.parent.mkdir(parents=True, exist_ok=True) + target.known_hosts.touch(mode=0o600, exist_ok=True) + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + if target.run("true", check=False).returncode == 0: + return + time.sleep(10) + raise ControllerError("AutoDL SSH did not become ready") + + +def runner_yaml() -> str: + labels = "\n".join(f' - "{label}"' for label in RUNNER_LABELS) + return f"""log: + level: info +runner: + file: /opt/ces-runner/.runner + capacity: 1 + timeout: 24h + fetch_timeout: 5s + fetch_interval: 2s + labels: +{labels} +container: + docker_host: "-" +""" + + +def temp_secret(value: str) -> Path: + descriptor, name = tempfile.mkstemp(prefix="ces-secret-") + path = Path(name) + try: + os.write(descriptor, value.encode()) + finally: + os.close(descriptor) + path.chmod(0o600) + return path + + +def provision_worker(target: SshTarget, paths: Paths, runner_token: str) -> None: + binary = ensure_runner(paths) + target.run( + "install -d -m 700 /opt/ces-runner /run/ces-runner-secrets " + f"{shlex.quote(paths.model_destination)}" + ) + print("Transferring SigLIP2 model to AutoDL", flush=True) + target.sync_to(paths.model_source, paths.model_destination) + target.copy_to(binary, "/opt/ces-runner/act_runner") + + with tempfile.NamedTemporaryFile("w", delete=False, encoding="utf-8") as temporary: + temporary.write(runner_yaml()) + runner_config = Path(temporary.name) + token_file = temp_secret(runner_token) + try: + target.copy_to(runner_config, "/opt/ces-runner/config.yaml") + target.copy_to(token_file, "/run/ces-runner-secrets/registration-token") + finally: + runner_config.unlink(missing_ok=True) + token_file.unlink(missing_ok=True) + + verification = target.run( + "nvidia-smi && python -c 'import torch; " + "print(torch.__version__, torch.version.cuda, torch.cuda.is_available(), " + "torch.cuda.get_device_name(), torch.cuda.get_device_capability())'", + check=False, + ) + print(verification.stdout, flush=True) + if verification.returncode != 0 or "True" not in verification.stdout: + raise ControllerError("AutoDL image failed the CUDA/PyTorch verification") + + name = f"ces-autodl-pro6000-{int(time.time())}" + labels = ",".join(RUNNER_LABELS) + register = " ".join( + [ + "cd /opt/ces-runner && chmod 700 act_runner", + "&& chmod 600 config.yaml /run/ces-runner-secrets/registration-token", + "&& rm -f .runner runner.exit", + "&& REGISTRATION_TOKEN=$(cat /run/ces-runner-secrets/registration-token)", + "&& ./act_runner -c config.yaml register --no-interactive", + "--instance https://gitea.entroia.com/", + '--token "$REGISTRATION_TOKEN"', + f"--name {shlex.quote(name)}", + f"--labels {shlex.quote(labels)}", + "&& unset REGISTRATION_TOKEN", + ] + ) + result = target.run(register, check=False) + target.run("rm -f /run/ces-runner-secrets/registration-token", check=False) + if result.returncode != 0: + raise ControllerError(f"GPU runner registration failed:\n{result.stdout}") + target.run( + "cd /opt/ces-runner && " + "nohup sh -c './act_runner -c config.yaml daemon; " + "code=$?; printf %s \"$code\" > runner.exit' " + "> runner.log 2>&1 < /dev/null & echo $! > runner.pid" + ) + time.sleep(3) + health = target.run( + "cd /opt/ces-runner && test -s runner.pid && kill -0 $(cat runner.pid) " + "&& tail -30 runner.log", + check=False, + ) + print(health.stdout, flush=True) + if health.returncode != 0: + raise ControllerError("GPU runner daemon did not stay running") + + +def discover(client: AutoDLClient) -> None: + images = client.list_images() + safe = [ + { + "image_uuid": item.get("image_uuid"), + "name": item.get("name"), + "status": item.get("status"), + "image_size": item.get("image_size"), + } + for item in images + ] + print(json.dumps({"private_images": safe}, indent=2, ensure_ascii=False)) + + +def provision(args: argparse.Namespace, client: AutoDLClient, paths: Paths) -> None: + if args.confirm_spend != "CREATE": + raise ControllerError("provision requires --confirm-spend CREATE") + runner_token = os.environ.get("TRAINING_RUNNER_REGISTRATION_TOKEN", "") + validate_local(paths, runner_token) + ensure_runner(paths) + if paths.state_file.is_file(): + previous = read_state(paths) + previous_uuid = previous.get("instance_uuid") + if previous_uuid: + previous_status = client.status(str(previous_uuid)) + if previous_status not in {"stopped", "released", "error"}: + raise ControllerError( + f"refusing to create a second billed instance; " + f"{previous_uuid} is {previous_status}" + ) + spec = instance_spec(args.image_uuid, args.disk_gb, args.data_center) + instance_uuid: str | None = None + try: + instance_uuid = client.create(spec) + write_state(paths, {"instance_uuid": instance_uuid, "phase": "creating"}) + print(f"created={instance_uuid}", flush=True) + wait_running(client, instance_uuid) + snapshot = client.snapshot(instance_uuid) + target = SshTarget.from_snapshot(snapshot, paths) + write_state( + paths, + { + "instance_uuid": instance_uuid, + "phase": "provisioning", + "ssh_host": target.host, + "ssh_port": target.port, + }, + ) + wait_ssh(target) + provision_worker(target, paths, runner_token) + write_state( + paths, + { + "instance_uuid": instance_uuid, + "phase": "runner-online", + "ssh_host": target.host, + "ssh_port": target.port, + }, + ) + print(f"runner_online=true instance={instance_uuid}", flush=True) + except BaseException: + if instance_uuid: + print(f"provision failed; powering off {instance_uuid}", file=sys.stderr, flush=True) + try: + client.power_off(instance_uuid) + except Exception as cleanup_error: + print(f"power-off cleanup failed: {cleanup_error}", file=sys.stderr) + raise + + +def status(client: AutoDLClient, paths: Paths) -> None: + state = read_state(paths) + instance_uuid = state["instance_uuid"] + result: dict[str, Any] = { + "instance_uuid": instance_uuid, + "controller_phase": state.get("phase"), + "autodl_status": client.status(instance_uuid), + } + if state.get("ssh_host") and state.get("ssh_port"): + target = SshTarget(state["ssh_host"], int(state["ssh_port"]), paths.ssh_key, paths.known_hosts) + health = target.run( + "cd /opt/ces-runner && test -s runner.pid && kill -0 $(cat runner.pid)", check=False + ) + result["runner_process_alive"] = health.returncode == 0 + print(json.dumps(result, indent=2)) + + +def shutdown(client: AutoDLClient, paths: Paths) -> None: + state = read_state(paths) + instance_uuid = state["instance_uuid"] + if state.get("ssh_host") and state.get("ssh_port"): + target = SshTarget(state["ssh_host"], int(state["ssh_port"]), paths.ssh_key, paths.known_hosts) + target.run( + "test ! -s /opt/ces-runner/runner.pid || kill $(cat /opt/ces-runner/runner.pid); " + "rm -f /run/ces-runner-secrets/registration-token", + check=False, + ) + destination = paths.artifact_destination / instance_uuid + try: + target.sync_from(paths.artifact_source, destination) + target.sync_from("/opt/ces-runner", destination / "runner") + except subprocess.CalledProcessError as exc: + print(f"artifact collection failed: {exc}", file=sys.stderr) + client.power_off(instance_uuid) + state["phase"] = "stopped" + write_state(paths, state) + print(f"powered_off={instance_uuid}") + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser() + parser.add_argument("operation", choices=("discover", "provision", "status", "shutdown")) + parser.add_argument("--image-uuid", default="") + parser.add_argument("--disk-gb", type=int, default=120) + parser.add_argument("--data-center", action="append", default=[]) + parser.add_argument("--confirm-spend", default="") + return parser.parse_args() + + +def main() -> int: + args = parse_args() + try: + paths = default_paths() + client = AutoDLClient(os.environ.get("AUTODL_API_TOKEN", "")) + if args.operation == "discover": + discover(client) + elif args.operation == "provision": + provision(args, client, paths) + elif args.operation == "status": + status(client, paths) + else: + shutdown(client, paths) + except (ControllerError, OSError, subprocess.CalledProcessError) as exc: + print(f"error: {exc}", file=sys.stderr) + return 1 + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/tests/test_controller.py b/tests/test_controller.py new file mode 100644 index 0000000..6e7beec --- /dev/null +++ b/tests/test_controller.py @@ -0,0 +1,62 @@ +import json +import os +import tempfile +import unittest +from pathlib import Path + +import controller + + +class ControllerTests(unittest.TestCase): + def test_runner_release_is_pinned(self): + self.assertEqual(controller.RUNNER_VERSION, "3.5.0") + self.assertEqual(len(controller.RUNNER_SHA256), 64) + + def test_instance_spec_uses_pro6000_and_cuda_128(self): + spec = controller.instance_spec("image-test", 120, []) + self.assertEqual(spec["gpu_spec_uuid"], "pro6000-p") + self.assertEqual(spec["cuda_v_from"], 128) + self.assertNotIn("data_center_list", spec) + + def test_instance_spec_preserves_data_center_preference(self): + spec = controller.instance_spec("image-test", 100, ["westDC3"]) + self.assertEqual(spec["data_center_list"], ["westDC3"]) + + def test_instance_spec_rejects_missing_image(self): + with self.assertRaises(controller.ControllerError): + controller.instance_spec("", 120, []) + + def test_runner_has_required_host_labels_and_no_docker(self): + rendered = controller.runner_yaml() + self.assertIn('"ubuntu-latest:host"', rendered) + self.assertIn('"ces-ci:host"', rendered) + self.assertIn('"gpu:host"', rendered) + self.assertIn('docker_host: "-"', rendered) + + def test_state_is_private_and_round_trips(self): + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + paths = controller.Paths( + root / "model", + "/model", + "/artifacts", + root / "artifacts", + root / "key", + root / "known_hosts", + root / "state.json", + root / "runner", + ) + controller.write_state(paths, {"instance_uuid": "pro-test"}) + self.assertEqual(controller.read_state(paths)["instance_uuid"], "pro-test") + self.assertEqual(os.stat(paths.state_file).st_mode & 0o777, 0o600) + + def test_example_workflow_files_are_valid_json_free(self): + # Guards against accidentally treating workflow YAML as controller state. + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "value.json" + path.write_text(json.dumps({"ok": True}), encoding="utf-8") + self.assertTrue(json.loads(path.read_text())["ok"]) + + +if __name__ == "__main__": + unittest.main()