Run batch GPU jobs from a script

curl -s https://api.deploygpu.ai/api/v1/users/me/ssh-keys -H "X-API-Key: $DEPLOYGPU_API_KEY"
export DEPLOYGPU_SSH_KEY_ID=ID_FROM_THE_LIST
export DEPLOYGPU_SSH_KEY_FILE=~/.ssh/id_ed25519

#!/usr/bin/env bash
set -euo pipefail
mkdir -p results
nvidia-smi > results/gpu.txt
python3 run.py > results/output.txt

"""Run one batch job on a rented GPU VM, copy the results back and always stop the VM.

Not yet run against the real API.
"""
import math
import os
import signal
import subprocess
import sys
import tempfile
import time

import requests

API = "https://api.deploygpu.ai/api/v1"
HEADERS = {"X-API-Key": os.environ["DEPLOYGPU_API_KEY"]}
KEY_FILE = os.environ["DEPLOYGPU_SSH_KEY_FILE"]  # private half of an SSH key on your account
KEY_ID = os.environ["DEPLOYGPU_SSH_KEY_ID"]      # its id from GET /api/v1/users/me/ssh-keys

GPU = "rtx-a6000"         # a variant slug from GET /api/v1/gpu-families
JOB_HOURS = 2             # how long you expect the job to take
JOB_TIMEOUT_S = 3 * 3600  # the job is killed after this
DISK_GB = 100             # used only where the disk size is adjustable
# Offers call Massed Compute "massed-compute", the create call wants "masscompute".
PROVIDER = {"massed-compute": "masscompute"}


def vm_offers():
    params = {"gpu": GPU, "max_gpu_count": 1, "sort_by": "priceHourly", "sort_order": "asc", "size": 20}
    resp = requests.get(f"{API}/offers", params=params, timeout=30)
    resp.raise_for_status()
    # Vast.ai containers are "dedicated", and the API refuses them.
    return [o for o in resp.json()["offers"] if o["deploymentType"] == "virtual" and o["isAvailable"]]


def check_balance(offer):
    resp = requests.get(f"{API}/account/balance", headers=HEADERS, timeout=30)
    resp.raise_for_status()
    hours = math.ceil(JOB_HOURS) + 1  # paid time starts at provisioning and ends after the copy
    if float(resp.json()["balance"]) < offer["priceHourly"] * hours:
        sys.exit(f"Balance is below the price of {hours} hours on this offer")


def create(offer):
    slug = offer["provider"]["slug"]
    body = {"provider": PROVIDER.get(slug, slug), "offer_id": offer["id"], "ssh_key_ids": [KEY_ID]}
    low, high = offer.get("storageMinGB"), offer.get("storageMaxGB")
    if low is not None and high is not None and low != high:
        body["disk_gb"] = min(max(DISK_GB, low), high)
    try:
        resp = requests.post(f"{API}/deployments", headers=HEADERS, json=body, timeout=60)
    except requests.RequestException as exc:
        # No answer, so a GPU may already be provisioning. There is no idempotency key:
        # never send the create again until you have listed your deployments.
        sys.exit(f"Create got no answer ({exc}). List your deployments before running again.")
    if resp.status_code == 202:
        return resp.json()["id"]
    try:
        code = resp.json()["error"]["code"]
    except (ValueError, KeyError, TypeError):
        code = None
    if code in ("offer_unavailable", "offer_cooling_down"):
        return None  # nothing was created or charged, so try the next offer
    sys.exit(f"Create failed: HTTP {resp.status_code} {code}")


def launch():
    for offer in vm_offers():
        check_balance(offer)
        deployment_id = create(offer)
        if deployment_id:
            return deployment_id
    sys.exit(f"No {GPU} offer accepted a create call right now")


def wait_active(deployment_id, limit_s=900):
    start = time.monotonic()
    while time.monotonic() - start < limit_s:
        resp = requests.get(f"{API}/deployments/{deployment_id}", headers=HEADERS, timeout=30)
        if resp.status_code == 404 and time.monotonic() - start < 30:
            time.sleep(2)  # right after the 202, a GET can 404 for about a second
            continue
        resp.raise_for_status()
        deployment = resp.json()
        if deployment["status"] == "active":
            return deployment["connection"]
        if deployment["status"] in ("failed", "interrupted", "stopping", "terminated"):
            sys.exit(
                f"{deployment['status']}: {deployment.get('failure_code')}, "
                f"retryable={deployment.get('retryable')}, {deployment.get('suggested_action')}"
            )
        time.sleep(10)
    sys.exit("Timed out waiting for active")


def stop(deployment_id):
    for attempt in range(5):
        try:
            resp = requests.post(f"{API}/deployments/{deployment_id}/stop", headers=HEADERS, timeout=60)
            if resp.status_code in (200, 400):  # 400: already stopping, terminated or failed
                print(f"Stop answered {resp.status_code} for {deployment_id}")
                return
            print(f"Stop answered {resp.status_code}, retrying")
        except requests.RequestException as exc:
            print(f"Stop failed ({exc}), retrying")
        time.sleep(5 * (attempt + 1))
    sys.exit(f"STOP FAILED: stop {deployment_id} in the dashboard now")


def main():
    for sig in (signal.SIGTERM, signal.SIGHUP):  # by default these skip finally: a CI cancel, a closed terminal
        signal.signal(sig, lambda *_: sys.exit(1))
    deployment_id = launch()
    print(f"Launched {deployment_id}")
    conn = None
    try:
        conn = wait_active(deployment_id)
        known_hosts = os.path.join(tempfile.mkdtemp(), "known_hosts")
        opts = [
            "-i", KEY_FILE, "-o", "IdentitiesOnly=yes", "-o", "BatchMode=yes", "-o", "ConnectTimeout=10",
            "-o", "StrictHostKeyChecking=accept-new", "-o", f"UserKnownHostsFile={known_hosts}",
        ]
        target = f"{conn['ssh_user']}@{conn['public_ip']}"
        ssh = ["ssh", "-p", str(conn["ssh_port"]), *opts, target]
        scp = ["scp", "-P", str(conn["ssh_port"]), *opts]

        for _ in range(30):  # active does not mean SSH answers yet
            try:
                if subprocess.run([*ssh, "true"], capture_output=True, timeout=30).returncode == 0:
                    break
            except subprocess.TimeoutExpired:
                pass
            time.sleep(10)
        else:
            sys.exit("SSH never answered")

        subprocess.run([*scp, "-r", "job", f"{target}:"], check=True, timeout=600)
        subprocess.run([*ssh, "cd job && bash run.sh"], check=True, timeout=JOB_TIMEOUT_S)
    finally:
        if conn:  # copy the results out before the disk is deleted, even after a failed job
            try:
                subprocess.run([*scp, "-r", f"{target}:job/results", f"results-{deployment_id}"],
                               check=True, timeout=600)
            except Exception as exc:
                print(f"Results NOT copied: {exc}", file=sys.stderr)
        stop(deployment_id)


if __name__ == "__main__":
    main()