520 lines
19 KiB
Python
Executable File
520 lines
19 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""Create an AutoDL Pro worker, attach an ephemeral Gitea runner, and tear it down.
|
|
|
|
This program is intended to run on the persistent relay/controller host. Project
|
|
source is never copied by this program: Gitea Actions checks it out after the
|
|
runner accepts a job. Only model data, the runner binary, and runner credentials
|
|
are transferred to the temporary AutoDL worker.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import hashlib
|
|
import json
|
|
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
|
|
|
|
|
|
DEFAULT_RUNNER_URL = (
|
|
"https://dl.gitea.com/gitea-runner/0.2.13/act_runner-0.2.13-linux-amd64"
|
|
)
|
|
DEFAULT_RUNNER_SHA256 = "3acac8b506ac8cadc88a55155b5d6378f0fab0b8f62d1e0c0450f4ccd69733e2"
|
|
DEFAULT_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
|
|
|
|
|
|
def read_secret(path: str) -> str:
|
|
secret_path = Path(path).expanduser()
|
|
mode = stat.S_IMODE(secret_path.stat().st_mode)
|
|
if mode & 0o077:
|
|
raise ControllerError(f"secret file must not be group/world accessible: {secret_path}")
|
|
value = secret_path.read_text(encoding="utf-8").strip()
|
|
if not value:
|
|
raise ControllerError(f"secret file is empty: {secret_path}")
|
|
return value
|
|
|
|
|
|
class AutoDLProClient:
|
|
def __init__(self, base_url: str, token: str, timeout: int = 30):
|
|
self.base_url = base_url.rstrip("/")
|
|
self.token = token
|
|
self.timeout = timeout
|
|
|
|
def request(self, method: str, path: str, body: dict[str, Any]) -> Any:
|
|
payload = json.dumps(body).encode("utf-8")
|
|
request = urllib.request.Request(
|
|
f"{self.base_url}{path}",
|
|
data=payload,
|
|
method=method,
|
|
headers={"Authorization": self.token, "Content-Type": "application/json"},
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(request, timeout=self.timeout) as response:
|
|
parsed = json.loads(response.read().decode("utf-8"))
|
|
except (urllib.error.URLError, json.JSONDecodeError) as exc:
|
|
raise ControllerError(f"AutoDL API request failed: {method} {path}: {exc}") from exc
|
|
if parsed.get("code") != "Success":
|
|
raise ControllerError(
|
|
f"AutoDL API rejected {method} {path}: "
|
|
f"code={parsed.get('code')!r}, msg={parsed.get('msg')!r}"
|
|
)
|
|
return parsed.get("data")
|
|
|
|
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}
|
|
)
|
|
|
|
def release(self, instance_uuid: str) -> None:
|
|
self.request(
|
|
"POST", "/api/v1/dev/instance/pro/release", {"instance_uuid": instance_uuid}
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class SshTarget:
|
|
host: str
|
|
port: int
|
|
key: Path
|
|
known_hosts: Path
|
|
|
|
@classmethod
|
|
def from_snapshot(cls, snapshot: dict[str, Any], ssh_config: dict[str, Any]) -> "SshTarget":
|
|
host = str(snapshot.get("proxy_host", "")).strip()
|
|
port = int(snapshot.get("ssh_port", 0))
|
|
if not host or port <= 0:
|
|
raise ControllerError("AutoDL snapshot did not contain a usable SSH endpoint")
|
|
return cls(
|
|
host=host,
|
|
port=port,
|
|
key=Path(ssh_config["private_key"]).expanduser(),
|
|
known_hosts=Path(ssh_config["known_hosts"]).expanduser(),
|
|
)
|
|
|
|
def 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 ssh(self, command: str, *, check: bool = True) -> subprocess.CompletedProcess[str]:
|
|
return subprocess.run(
|
|
["ssh", *self.options(), f"root@{self.host}", command],
|
|
check=check,
|
|
text=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
)
|
|
|
|
def rsync_to(self, source: Path, destination: str) -> None:
|
|
ssh_transport = ["ssh", *self.options()]
|
|
subprocess.run(
|
|
[
|
|
"rsync",
|
|
"-ah",
|
|
"--partial",
|
|
"--append-verify",
|
|
"-e",
|
|
shlex.join(ssh_transport),
|
|
f"{source}/",
|
|
f"root@{self.host}:{destination.rstrip('/')}/",
|
|
],
|
|
check=True,
|
|
)
|
|
|
|
def rsync_from(self, source: str, destination: Path) -> None:
|
|
destination.mkdir(parents=True, exist_ok=True)
|
|
ssh_transport = ["ssh", *self.options()]
|
|
subprocess.run(
|
|
[
|
|
"rsync",
|
|
"-ah",
|
|
"--partial",
|
|
"-e",
|
|
shlex.join(ssh_transport),
|
|
f"root@{self.host}:{source.rstrip('/')}/",
|
|
f"{destination}/",
|
|
],
|
|
check=True,
|
|
)
|
|
|
|
def scp_to(self, source: Path, destination: str) -> None:
|
|
options = self.options()
|
|
# scp uses uppercase -P for its port option.
|
|
options[options.index("-p")] = "-P"
|
|
subprocess.run(
|
|
["scp", *options, str(source), f"root@{self.host}:{destination}"], check=True
|
|
)
|
|
|
|
|
|
def load_config(path: Path) -> dict[str, Any]:
|
|
config = json.loads(path.read_text(encoding="utf-8"))
|
|
for section in ("autodl", "runner", "ssh", "transfer"):
|
|
if section not in config:
|
|
raise ControllerError(f"missing configuration section: {section}")
|
|
return config
|
|
|
|
|
|
def preflight(config: dict[str, Any]) -> None:
|
|
autodl = config["autodl"]
|
|
if str(autodl["image_uuid"]).startswith("REPLACE_"):
|
|
raise ControllerError("replace the placeholder AutoDL private image UUID")
|
|
read_secret(autodl["token_file"])
|
|
read_secret(config["runner"]["registration_token_file"])
|
|
if config.get("dispatch", {}).get("enabled", False):
|
|
read_secret(config["dispatch"]["token_file"])
|
|
|
|
private_key = Path(config["ssh"]["private_key"]).expanduser()
|
|
if not private_key.is_file():
|
|
raise ControllerError(f"AutoDL SSH private key does not exist: {private_key}")
|
|
if stat.S_IMODE(private_key.stat().st_mode) & 0o077:
|
|
raise ControllerError(f"SSH private key must have mode 0600: {private_key}")
|
|
|
|
model_source = Path(config["transfer"]["model_source"]).expanduser()
|
|
missing = [
|
|
name
|
|
for name in config["transfer"].get("model_required_files", DEFAULT_MODEL_FILES)
|
|
if not (model_source / name).is_file() or (model_source / name).stat().st_size == 0
|
|
]
|
|
if missing:
|
|
raise ControllerError(f"model download is incomplete; file check failed: {missing}")
|
|
|
|
|
|
def download_runner(config: dict[str, Any]) -> Path:
|
|
runner = config["runner"]
|
|
cache_path = Path(runner["cache_path"]).expanduser()
|
|
expected = runner.get("sha256", DEFAULT_RUNNER_SHA256)
|
|
if cache_path.exists() and sha256(cache_path) == expected:
|
|
return cache_path
|
|
cache_path.parent.mkdir(parents=True, exist_ok=True)
|
|
url = runner.get("download_url", DEFAULT_RUNNER_URL)
|
|
print(f"Downloading pinned act_runner from {url}", flush=True)
|
|
with urllib.request.urlopen(url, timeout=300) as response, tempfile.NamedTemporaryFile(
|
|
dir=cache_path.parent, delete=False
|
|
) as temporary:
|
|
while chunk := response.read(1024 * 1024):
|
|
temporary.write(chunk)
|
|
temporary_path = Path(temporary.name)
|
|
if sha256(temporary_path) != expected:
|
|
temporary_path.unlink(missing_ok=True)
|
|
raise ControllerError("downloaded act_runner checksum did not match the pinned checksum")
|
|
temporary_path.chmod(0o755)
|
|
temporary_path.replace(cache_path)
|
|
return cache_path
|
|
|
|
|
|
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 wait_for_status(
|
|
client: AutoDLProClient, instance_uuid: str, desired: str, timeout_seconds: int
|
|
) -> None:
|
|
deadline = time.monotonic() + timeout_seconds
|
|
last = ""
|
|
while time.monotonic() < deadline:
|
|
last = client.status(instance_uuid)
|
|
print(f"AutoDL instance {instance_uuid}: {last}", flush=True)
|
|
if last == desired:
|
|
return
|
|
if last in {"error", "released"}:
|
|
break
|
|
time.sleep(10)
|
|
raise ControllerError(f"instance did not reach {desired!r}; last status was {last!r}")
|
|
|
|
|
|
def wait_for_ssh(target: SshTarget, timeout_seconds: int) -> None:
|
|
target.known_hosts.parent.mkdir(parents=True, exist_ok=True)
|
|
target.known_hosts.touch(mode=0o600, exist_ok=True)
|
|
deadline = time.monotonic() + timeout_seconds
|
|
while time.monotonic() < deadline:
|
|
result = target.ssh("true", check=False)
|
|
if result.returncode == 0:
|
|
return
|
|
time.sleep(10)
|
|
raise ControllerError("SSH did not become ready before the timeout")
|
|
|
|
|
|
def runner_config(labels: list[str], timeout: str) -> str:
|
|
rendered_labels = "\n".join(f' - "{label}"' for label in labels)
|
|
return f"""log:
|
|
level: info
|
|
runner:
|
|
file: /opt/ces-runner/.runner
|
|
capacity: 1
|
|
timeout: {timeout}
|
|
fetch_timeout: 5s
|
|
fetch_interval: 2s
|
|
labels:
|
|
{rendered_labels}
|
|
container:
|
|
docker_host: "-"
|
|
"""
|
|
|
|
|
|
def provision_worker(config: dict[str, Any], target: SshTarget, runner_binary: Path) -> None:
|
|
transfer = config["transfer"]
|
|
runner = config["runner"]
|
|
model_source = Path(transfer["model_source"]).expanduser()
|
|
if not model_source.is_dir():
|
|
raise ControllerError(f"model source directory does not exist: {model_source}")
|
|
token_file = Path(runner["registration_token_file"]).expanduser()
|
|
read_secret(str(token_file)) # Validate permissions and contents before copying.
|
|
|
|
target.ssh(
|
|
"install -d -m 700 /opt/ces-runner /run/ces-runner-secrets "
|
|
f"{shlex.quote(transfer['model_destination'])}"
|
|
)
|
|
target.rsync_to(model_source, transfer["model_destination"])
|
|
target.scp_to(runner_binary, "/opt/ces-runner/act_runner")
|
|
target.scp_to(token_file, "/run/ces-runner-secrets/registration-token")
|
|
|
|
labels = list(runner["labels"])
|
|
with tempfile.NamedTemporaryFile("w", encoding="utf-8", delete=False) as temporary:
|
|
temporary.write(runner_config(labels, runner.get("job_timeout", "24h")))
|
|
config_path = Path(temporary.name)
|
|
try:
|
|
target.scp_to(config_path, "/opt/ces-runner/config.yaml")
|
|
finally:
|
|
config_path.unlink(missing_ok=True)
|
|
|
|
name = f"{runner['name_prefix']}-{int(time.time())}"
|
|
labels_argument = ",".join(labels)
|
|
register = " ".join(
|
|
[
|
|
"cd /opt/ces-runner && chmod 700 act_runner",
|
|
"&& chmod 600 /run/ces-runner-secrets/registration-token config.yaml",
|
|
"&& ./act_runner -c config.yaml register --no-interactive --ephemeral",
|
|
f"--instance {shlex.quote(runner['gitea_instance'])}",
|
|
"--token-file /run/ces-runner-secrets/registration-token",
|
|
f"--name {shlex.quote(name)}",
|
|
f"--labels {shlex.quote(labels_argument)}",
|
|
]
|
|
)
|
|
result = target.ssh(register, check=False)
|
|
if result.returncode != 0:
|
|
raise ControllerError(f"runner registration failed:\n{result.stdout}")
|
|
|
|
start = (
|
|
"cd /opt/ces-runner && "
|
|
"nohup sh -c './act_runner -c config.yaml daemon; "
|
|
'code=$?; printf %s "$code" > runner.exit; '
|
|
"rm -f /run/ces-runner-secrets/registration-token' "
|
|
"> runner.log 2>&1 < /dev/null & echo $! > runner.pid"
|
|
)
|
|
target.ssh(start)
|
|
print(f"Registered ephemeral runner {name}", flush=True)
|
|
|
|
|
|
def dispatch_workflow(config: dict[str, Any]) -> None:
|
|
dispatch = config.get("dispatch", {})
|
|
if not dispatch.get("enabled", False):
|
|
print("Workflow dispatch disabled; trigger Formal GPU training in Gitea now.", flush=True)
|
|
return
|
|
token = read_secret(dispatch["token_file"])
|
|
body = json.dumps({"ref": dispatch.get("ref", "main")}).encode("utf-8")
|
|
url = (
|
|
f"{dispatch['gitea_instance'].rstrip('/')}/api/v1/repos/"
|
|
f"{dispatch['repository']}/actions/workflows/{dispatch['workflow']}/dispatches"
|
|
)
|
|
request = urllib.request.Request(
|
|
url,
|
|
data=body,
|
|
method="POST",
|
|
headers={"Authorization": f"token {token}", "Content-Type": "application/json"},
|
|
)
|
|
try:
|
|
with urllib.request.urlopen(request, timeout=30) as response:
|
|
if response.status not in {200, 201, 204}:
|
|
raise ControllerError(f"unexpected workflow dispatch status: {response.status}")
|
|
except urllib.error.URLError as exc:
|
|
raise ControllerError(f"Gitea workflow dispatch failed: {exc}") from exc
|
|
print(f"Dispatched {dispatch['workflow']} at {dispatch.get('ref', 'main')}", flush=True)
|
|
|
|
|
|
def wait_for_runner(target: SshTarget, max_runtime_hours: float) -> int:
|
|
deadline = time.monotonic() + max_runtime_hours * 3600
|
|
while time.monotonic() < deadline:
|
|
result = target.ssh(
|
|
"test -f /opt/ces-runner/runner.exit && cat /opt/ces-runner/runner.exit",
|
|
check=False,
|
|
)
|
|
if result.returncode == 0 and result.stdout.strip().isdigit():
|
|
return int(result.stdout.strip())
|
|
time.sleep(30)
|
|
raise ControllerError("runner exceeded the configured maximum runtime")
|
|
|
|
|
|
def create_spec(config: dict[str, Any]) -> dict[str, Any]:
|
|
autodl = config["autodl"]
|
|
spec = {
|
|
"req_gpu_amount": autodl.get("req_gpu_amount", 1),
|
|
"expand_system_disk_by_gb": autodl["expand_system_disk_by_gb"],
|
|
"gpu_spec_uuid": autodl["gpu_spec_uuid"],
|
|
"image_uuid": autodl["image_uuid"],
|
|
"cuda_v_from": autodl["cuda_v_from"],
|
|
"instance_name": autodl.get("instance_name", "ces-gitea-runner"),
|
|
}
|
|
if autodl.get("data_center_list"):
|
|
spec["data_center_list"] = autodl["data_center_list"]
|
|
return spec
|
|
|
|
|
|
def print_plan(config: dict[str, Any]) -> None:
|
|
safe = {
|
|
"create": create_spec(config),
|
|
"model": {
|
|
"source": config["transfer"]["model_source"],
|
|
"destination": config["transfer"]["model_destination"],
|
|
},
|
|
"runner": {
|
|
"ephemeral": True,
|
|
"labels": config["runner"]["labels"],
|
|
"gitea_instance": config["runner"]["gitea_instance"],
|
|
},
|
|
"dispatch": {
|
|
key: value
|
|
for key, value in config.get("dispatch", {}).items()
|
|
if key != "token_file"
|
|
},
|
|
"cleanup": {
|
|
"power_off": True,
|
|
"release": config["autodl"].get("release_after_run", False),
|
|
},
|
|
}
|
|
print(json.dumps(safe, indent=2, ensure_ascii=False))
|
|
|
|
|
|
def run(config: dict[str, Any]) -> None:
|
|
preflight(config)
|
|
token = read_secret(config["autodl"]["token_file"])
|
|
client = AutoDLProClient(config["autodl"].get("base_url", "https://api.autodl.com"), token)
|
|
runner_binary = download_runner(config)
|
|
instance_uuid: str | None = None
|
|
target: SshTarget | None = None
|
|
try:
|
|
instance_uuid = client.create(create_spec(config))
|
|
print(f"Created AutoDL Pro instance {instance_uuid}", flush=True)
|
|
wait_for_status(
|
|
client,
|
|
instance_uuid,
|
|
"running",
|
|
int(config["autodl"].get("startup_timeout_seconds", 1800)),
|
|
)
|
|
target = SshTarget.from_snapshot(client.snapshot(instance_uuid), config["ssh"])
|
|
wait_for_ssh(target, int(config["ssh"].get("connect_timeout_seconds", 600)))
|
|
provision_worker(config, target, runner_binary)
|
|
dispatch_workflow(config)
|
|
runner_exit = wait_for_runner(
|
|
target, float(config["autodl"].get("max_runtime_hours", 24))
|
|
)
|
|
print(f"Ephemeral runner exited with code {runner_exit}", flush=True)
|
|
target.rsync_from(
|
|
config["transfer"]["artifact_source"],
|
|
Path(config["transfer"]["artifact_destination"]).expanduser() / instance_uuid,
|
|
)
|
|
target.rsync_from(
|
|
"/opt/ces-runner",
|
|
Path(config["transfer"]["artifact_destination"]).expanduser()
|
|
/ instance_uuid
|
|
/ "runner",
|
|
)
|
|
finally:
|
|
if instance_uuid is not None:
|
|
if target is not None:
|
|
target.ssh("rm -f /run/ces-runner-secrets/registration-token", check=False)
|
|
print(f"Powering off AutoDL Pro instance {instance_uuid}", flush=True)
|
|
client.power_off(instance_uuid)
|
|
if config["autodl"].get("release_after_run", False):
|
|
wait_for_status(
|
|
client,
|
|
instance_uuid,
|
|
"stopped",
|
|
int(config["autodl"].get("shutdown_timeout_seconds", 600)),
|
|
)
|
|
client.release(instance_uuid)
|
|
print(f"Released AutoDL Pro instance {instance_uuid}", flush=True)
|
|
|
|
|
|
def main() -> int:
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("command", choices=("plan", "run"))
|
|
parser.add_argument("--config", type=Path, required=True)
|
|
parser.add_argument(
|
|
"--confirm-create",
|
|
action="store_true",
|
|
help="required with 'run'; acknowledges that a paid AutoDL instance will be created",
|
|
)
|
|
args = parser.parse_args()
|
|
try:
|
|
config = load_config(args.config)
|
|
if args.command == "plan":
|
|
print_plan(config)
|
|
else:
|
|
if not args.confirm_create:
|
|
raise ControllerError("run requires --confirm-create because it starts billing")
|
|
run(config)
|
|
except (ControllerError, OSError, subprocess.CalledProcessError) as exc:
|
|
print(f"error: {exc}", file=sys.stderr)
|
|
return 1
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|