#!/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.parse 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_USER_AGENT = "ces-autodl-controller/1.0" 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", ] PUBLIC_BASE_IMAGES = [ { "image_uuid": "base-image-12be412037", "frameworks": {"pytorch": "1.9.0"}, "cuda_version": "11.1", "python_version": "3.8", "os": "ubuntu18.04", }, { "image_uuid": "base-image-u9r24vthlk", "frameworks": {"pytorch": "1.10.0"}, "cuda_version": "11.3", "python_version": "3.8", "os": "ubuntu20.04", }, { "image_uuid": "base-image-l374uiucui", "frameworks": {"pytorch": "1.11.0"}, "cuda_version": "11.3", "python_version": "3.8", "os": "ubuntu20.04", }, { "image_uuid": "base-image-l2t43iu6uk", "frameworks": {"pytorch": "2.0.0", "tensorrt": "8.5.1"}, "cuda_version": "11.8", "python_version": "3.8", "os": "ubuntu20.04", }, { "image_uuid": "base-image-0gxqmciyth", "frameworks": {"tensorflow": "2.5.0"}, "cuda_version": "11.2", "python_version": "3.8", "os": "ubuntu18.04", }, { "image_uuid": "base-image-uxeklgirir", "frameworks": {"tensorflow": "2.9.0"}, "cuda_version": "11.2", "python_version": "3.8", "os": "ubuntu20.04", }, { "image_uuid": "base-image-4bpg0tt88l", "frameworks": {"tensorflow": "1.15.5"}, "cuda_version": "11.4", "python_version": "3.8", "os": None, }, { "image_uuid": "base-image-mbr2n4urrc", "frameworks": {"miniconda": "conda3"}, "cuda_version": "11.6", "python_version": "3.8", "os": "ubuntu20.04", }, { "image_uuid": "base-image-qkkhitpik5", "frameworks": {"miniconda": "conda3"}, "cuda_version": "10.2", "python_version": "3.8", "os": "ubuntu18.04", }, { "image_uuid": "base-image-h041hn36yt", "frameworks": {"miniconda": "conda3"}, "cuda_version": "11.1", "python_version": "3.8", "os": "ubuntu18.04", }, { "image_uuid": "base-image-7bn8iqhkb5", "frameworks": {"miniconda": "conda3", "cudagl": "11.3"}, "cuda_version": "11.3", "python_version": "3.8", "os": "ubuntu20.04", }, { "image_uuid": "base-image-k0vep6kyq8", "frameworks": {"miniconda": "conda3"}, "cuda_version": "9.0", "python_version": "3.6", "os": "ubuntu16.04", }, { "image_uuid": "base-image-l2843iu23k", "frameworks": {"tensorrt": "8.5.1"}, "cuda_version": "11.8", "python_version": "3.8", "os": "ubuntu20.04", }, ] 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: url = f"{API_BASE}{path}" data: bytes | None if method == "GET": url = f"{url}?{urllib.parse.urlencode(body)}" data = None else: data = json.dumps(body).encode() request = urllib.request.Request( url, data=data, 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}, ) # AutoDL returns JSON null instead of [] when the account has no # private images. Treat both representations as an empty result. return list(data.get("list") or []) 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 select_public_image(framework: str, framework_version: str, cuda_version: str) -> str: normalized_framework = framework.strip().casefold() matches = [ image for image in PUBLIC_BASE_IMAGES if image["frameworks"].get(normalized_framework) == framework_version.strip() and image["cuda_version"] == cuda_version.strip() ] if len(matches) != 1: raise ControllerError( "no unique public base image matches " f"framework={framework!r}, version={framework_version!r}, " f"cuda={cuda_version!r}; provide --image-uuid or update the catalog" ) return str(matches[0]["image_uuid"]) def requested_image_uuid(args: argparse.Namespace) -> str: if args.image_uuid: return args.image_uuid profile = (args.framework, args.framework_version, args.cuda_version) if any(profile): if not all(profile): raise ControllerError( "--framework, --framework-version, and --cuda-version must be used together" ) return select_public_image(*profile) raise ControllerError("an image UUID or a complete public image profile is required") 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) request = urllib.request.Request(RUNNER_URL, headers={"User-Agent": RUNNER_USER_AGENT}) with urllib.request.urlopen(request, 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, "public_base_images": PUBLIC_BASE_IMAGES}, 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(requested_image_uuid(args), 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("--framework", default="") parser.add_argument("--framework-version", default="") parser.add_argument("--cuda-version", 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())