776 lines
26 KiB
Python
776 lines
26 KiB
Python
#!/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 list_instances(self) -> list[dict[str, Any]]:
|
|
data = self.request(
|
|
"POST", "/api/v1/dev/instance/pro/list", {"page_index": 1, "page_size": 100}
|
|
)
|
|
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}
|
|
)
|
|
|
|
def power_on(self, instance_uuid: str) -> None:
|
|
self.request(
|
|
"POST",
|
|
"/api/v1/dev/instance/pro/power_on",
|
|
{"instance_uuid": instance_uuid, "payload": "gpu", "start_command": "sleep 1"},
|
|
)
|
|
|
|
|
|
@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 power_off_and_wait(client: AutoDLClient, instance_uuid: str, timeout: int = 300) -> None:
|
|
client.power_off(instance_uuid)
|
|
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 in {"shutdown", "stopped", "released"}:
|
|
return
|
|
time.sleep(5)
|
|
raise ControllerError(f"instance failed to stop; 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)}"
|
|
)
|
|
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")
|
|
|
|
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)
|
|
|
|
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 instances(client: AutoDLClient) -> None:
|
|
safe = [
|
|
{
|
|
"uuid": item.get("uuid"),
|
|
"name": item.get("name"),
|
|
"status": item.get("status"),
|
|
"region_name": item.get("region_name"),
|
|
"gpu_spec_uuid": item.get("gpu_spec_uuid"),
|
|
"req_gpu_amount": item.get("req_gpu_amount"),
|
|
"created_at": item.get("created_at"),
|
|
}
|
|
for item in client.list_instances()
|
|
]
|
|
print(json.dumps({"pro_instances": 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 {"shutdown", "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)
|
|
provision_existing_instance(client, instance_uuid, paths, runner_token)
|
|
except BaseException:
|
|
if instance_uuid:
|
|
print(f"provision failed; powering off {instance_uuid}", file=sys.stderr, flush=True)
|
|
try:
|
|
power_off_and_wait(client, instance_uuid)
|
|
state = read_state(paths)
|
|
state["phase"] = "stopped"
|
|
write_state(paths, state)
|
|
except Exception as cleanup_error:
|
|
print(f"power-off cleanup failed: {cleanup_error}", file=sys.stderr)
|
|
raise
|
|
|
|
|
|
def provision_existing_instance(
|
|
client: AutoDLClient, instance_uuid: str, paths: Paths, runner_token: str
|
|
) -> None:
|
|
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)
|
|
|
|
|
|
def resume(args: argparse.Namespace, client: AutoDLClient, paths: Paths) -> None:
|
|
if args.confirm_spend != "CREATE":
|
|
raise ControllerError("resume requires --confirm-spend CREATE")
|
|
runner_token = os.environ.get("TRAINING_RUNNER_REGISTRATION_TOKEN", "")
|
|
validate_local(paths, runner_token)
|
|
ensure_runner(paths)
|
|
state = read_state(paths)
|
|
instance_uuid = str(state["instance_uuid"])
|
|
current_status = client.status(instance_uuid)
|
|
if current_status not in {"shutdown", "stopped"}:
|
|
raise ControllerError(f"resume requires a stopped instance; current={current_status}")
|
|
try:
|
|
client.power_on(instance_uuid)
|
|
state["phase"] = "starting"
|
|
write_state(paths, state)
|
|
provision_existing_instance(client, instance_uuid, paths, runner_token)
|
|
except BaseException:
|
|
print(f"resume failed; powering off {instance_uuid}", file=sys.stderr, flush=True)
|
|
try:
|
|
power_off_and_wait(client, instance_uuid)
|
|
state = read_state(paths)
|
|
state["phase"] = "stopped"
|
|
write_state(paths, state)
|
|
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)
|
|
power_off_and_wait(client, 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", "instances", "provision", "resume", "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 == "instances":
|
|
instances(client)
|
|
elif args.operation == "provision":
|
|
provision(args, client, paths)
|
|
elif args.operation == "resume":
|
|
resume(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())
|