Support safe AutoDL instance recovery

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: