Connect to a managed Slurm cluster

export ORNN_API_URL="https://api.ornn.com"
export RESERVATION_ID='11111111-1111-4111-8111-111111111111'
export NODE_ID='33333333-3333-4333-8333-333333333333'

# Optional: find the IDs of your organization's active SSH keys.
ORGANIZATION_ID="$(curl --fail-with-body -sS "$ORNN_API_URL/v1/auth/session" \
  -H "Authorization: Bearer $ORNN_API_KEY" | jq -er '.organizations[0].id')"
curl --fail-with-body -sS \
  "$ORNN_API_URL/v1/organizations/$ORGANIZATION_ID/ssh-keys?status=active" \
  -H "Authorization: Bearer $ORNN_API_KEY" | jq '.ssh_keys[] | {id, label, fingerprint}'

jq -n --arg reservation "$RESERVATION_ID" --arg node "$NODE_ID" \
  '{kind: "slurm", reservation_id: $reservation,
    network_mode: "public", node_ids: [$node]}' \
  | curl --fail-with-body -sS -X POST "$ORNN_API_URL/v1/orchestrator/controllers" \
      -H "Authorization: Bearer $ORNN_API_KEY" \
      -H 'Content-Type: application/json' \
      --data-binary @- > launch.json
export CLUSTER_ID="$(jq -er '.id' launch.json)"

import os
import time

import requests

API = "https://api.ornn.com"
http = requests.Session()
http.headers["Authorization"] = f"Bearer {os.environ['ORNN_API_KEY']}"

RESERVATION_ID = "11111111-1111-4111-8111-111111111111"
NODE_IDS = ["33333333-3333-4333-8333-333333333333"]


LIVE_STATES = ("pending", "provisioning", "active", "deleting")


def slurm_cluster(cluster_id=None):
    """Return the {controller, nodes} entry for a Slurm cluster, or None.

    Without an ID, return your organization's live Slurm cluster. An
    organization runs at most one at a time.
    """
    response = http.get(f"{API}/v1/orchestrator/controllers", params={"kind": "slurm"}, timeout=30)
    response.raise_for_status()
    for entry in response.json()["controllers"]:
        controller = entry["controller"]
        if controller["id"] == cluster_id or (
            cluster_id is None and controller["state"] in LIVE_STATES
        ):
            return entry
    return None


existing = slurm_cluster()
if existing and existing["controller"]["state"] in ("pending", "provisioning", "active"):
    CLUSTER_ID = existing["controller"]["id"]  # reuse it; do not launch again
else:
    response = http.post(
        f"{API}/v1/orchestrator/controllers",
        json={
            "kind": "slurm",
            "reservation_id": RESERVATION_ID,
            "network_mode": "public",
            "node_ids": NODE_IDS,
        },
        timeout=30,
    )
    if response.status_code != 202:
        raise RuntimeError(f"Launch refused: {response.status_code} {response.text[:300]}")
    CLUSTER_ID = response.json()["id"]

ornn launch --mode slurm --reservation "$RESERVATION_ID" --key-id "$SSH_KEY_ID" \
  --nodes 1 --machines "$NODE_ID" --network public \
  --organization "$ORGANIZATION_ID" --yes \
  --idempotency-key "slurm-launch-$RESERVATION_ID" --format json

until curl --fail-with-body -sS \
    "$ORNN_API_URL/v1/orchestrator/controllers?kind=slurm" \
    -H "Authorization: Bearer $ORNN_API_KEY" \
    | jq -e --arg id "$CLUSTER_ID" \
      '.controllers[] | select(.controller.id == $id) | .controller.state | IN("active", "failed")'; do
  sleep 10
done

curl --fail-with-body -sS \
  "$ORNN_API_URL/v1/orchestrator/controllers/$CLUSTER_ID/access" \
  -H "Authorization: Bearer $ORNN_API_KEY" > login.json
LOGIN_HOST="$(jq -er '.login_host' login.json)"
LOGIN_PORT="$(jq -er '.login_port' login.json)"
LOGIN_USER="$(jq -er '.login_user' login.json)"

# Pin the login node's host key before the first connection.
printf '[%s]:%s %s\n' "$LOGIN_HOST" "$LOGIN_PORT" \
  "$(jq -er '.ssh_host_key' login.json | cut -d' ' -f1,2)" >> ~/.ssh/known_hosts
ssh -i ~/.ssh/id_ed25519 -p "$LOGIN_PORT" "$LOGIN_USER@$LOGIN_HOST" sinfo -N -l

deadline = time.monotonic() + 3600
while True:
    entry = slurm_cluster(cluster_id=CLUSTER_ID)
    state = entry["controller"]["state"] if entry else "missing"
    if state == "active":
        break
    if state in ("failed", "missing"):
        raise RuntimeError(f"Launch did not finish: {entry and entry['controller']['progress']}")
    if time.monotonic() > deadline:
        raise TimeoutError("Still provisioning; contact Ornn with the cluster ID.")
    time.sleep(10)

while True:
    response = http.get(f"{API}/v1/orchestrator/controllers/{CLUSTER_ID}/access", timeout=30)
    code = (response.json().get("detail") or {}).get("code") if response.status_code == 404 else None
    if code == "slurm_credentials_not_ready":
        time.sleep(10)
        continue
    response.raise_for_status()
    login = response.json()
    break

print(login["ssh_command"])  # for example: ssh -p 30022 tenant@203.0.113.10
print("Host key:", login["ssh_host_key"])

ornn clusters wait "$CLUSTER_ID" --timeout 60m --format json
ornn connect "$CLUSTER_ID" --print