From b58cc1f5f8bae930eb0fba061d1bfe93d3574eaa Mon Sep 17 00:00:00 2001 From: Codex Bot Date: Fri, 2 Oct 2026 15:23:33 +0800 Subject: [PATCH] Support safe AutoDL instance recovery --- .gitea/workflows/manage.yaml | 4 +- README.md | 2 + controller.py | 162 +++++++++++++++++++++++++++-------- 3 files changed, 128 insertions(+), 40 deletions(-) diff --git a/.gitea/workflows/manage.yaml b/.gitea/workflows/manage.yaml index 4068d95..0f4d486 100644 --- a/.gitea/workflows/manage.yaml +++ b/.gitea/workflows/manage.yaml @@ -4,7 +4,7 @@ on: workflow_dispatch: inputs: operation: - description: discover, provision, status, or shutdown + description: discover, instances, provision, resume, status, or shutdown required: true default: discover image_uuid: @@ -51,7 +51,7 @@ jobs: - name: Validate requested operation run: | case "$OPERATION" in - discover|provision|status|shutdown) ;; + discover|instances|provision|resume|status|shutdown) ;; *) echo "Unsupported operation" >&2; exit 2 ;; esac diff --git a/README.md b/README.md index 0f8714c..6557d77 100644 --- a/README.md +++ b/README.md @@ -59,8 +59,10 @@ ces-controller --help 从 Actions 手动运行 `Manage AutoDL GPU runner`: - `discover`:列出账号私有镜像和内置公共基础镜像,不产生 GPU 费用; +- `instances`:只读列出当前 Token 账号下的 Pro 实例; - `provision`:创建 PRO6000、传输模型并注册 GPU Runner;必须填写镜像 UUID,并把确认值 填为 `CREATE`; +- `resume`:对失败后已关机的现有实例继续配置;同样必须填写 `CREATE`; - `status`:查询实例和 Runner 进程状态; - `shutdown`:回收训练产物并关闭实例,不释放系统盘。 diff --git a/controller.py b/controller.py index ea99d48..f57e03f 100644 --- a/controller.py +++ b/controller.py @@ -183,6 +183,12 @@ class AutoDLClient: # 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)) @@ -205,6 +211,13 @@ class AutoDLClient: "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: @@ -431,6 +444,19 @@ def wait_running(client: AutoDLClient, instance_uuid: str, timeout: int = 1800) 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) @@ -476,6 +502,16 @@ def provision_worker(target: SshTarget, paths: Paths, runner_token: str) -> None "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") @@ -491,16 +527,6 @@ def provision_worker(target: SshTarget, paths: Paths, runner_token: str) -> None 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( @@ -558,6 +584,22 @@ def discover(client: AutoDLClient) -> None: ) +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") @@ -569,7 +611,7 @@ def provision(args: argparse.Namespace, client: AutoDLClient, paths: Paths) -> N previous_uuid = previous.get("instance_uuid") if previous_uuid: previous_status = client.status(str(previous_uuid)) - if previous_status not in {"stopped", "released", "error"}: + 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}" @@ -580,40 +622,77 @@ def provision(args: argparse.Namespace, client: AutoDLClient, paths: Paths) -> N 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) + 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: - client.power_off(instance_uuid) + 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"] @@ -647,7 +726,7 @@ def shutdown(client: AutoDLClient, paths: Paths) -> None: 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) + power_off_and_wait(client, instance_uuid) state["phase"] = "stopped" write_state(paths, state) print(f"powered_off={instance_uuid}") @@ -655,7 +734,10 @@ def shutdown(client: AutoDLClient, paths: Paths) -> None: def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser() - parser.add_argument("operation", choices=("discover", "provision", "status", "shutdown")) + 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="") @@ -673,8 +755,12 @@ def main() -> int: 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: