lambdalabs-cli-helper

llcloud

#!/usr/bin/env python3
"""llcloud - quick CLI for Lambda Labs cloud instances."""

import argparse
import base64
import json
import os
import subprocess
import sys
import time
import urllib.error
import urllib.request

API_BASE = "https://cloud.lambdalabs.com/api/v1"


def auth_header():
    key = os.environ.get("LAMBDA_API_KEY")
    if not key:
        sys.exit("error: LAMBDA_API_KEY not set")
    token = base64.b64encode(f"{key}:".encode()).decode()
    return f"Basic {token}"


def request(method, path, body=None):
    url = f"{API_BASE}{path}"
    data = json.dumps(body).encode() if body is not None else None
    req = urllib.request.Request(url, data=data, method=method)
    req.add_header("Authorization", auth_header())
    req.add_header("User-Agent", "llcloud/0.1")
    req.add_header("Accept", "application/json")
    if data is not None:
        req.add_header("Content-Type", "application/json")
    try:
        with urllib.request.urlopen(req) as resp:
            return json.loads(resp.read().decode())
    except urllib.error.HTTPError as e:
        raw = e.read().decode(errors="replace")
        try:
            err = json.loads(raw).get("error", {})
            detail = err.get("message") or raw
        except json.JSONDecodeError:
            detail = raw
        sys.exit(f"error: HTTP {e.code}: {detail}")
    except urllib.error.URLError as e:
        sys.exit(f"error: {e.reason}")


def print_table(headers, rows):
    widths = [max(len(h), *(len(str(r[i])) for r in rows)) for i, h in enumerate(headers)]
    fmt = "  ".join(f"{{:<{w}}}" for w in widths)
    print(fmt.format(*headers))
    for r in rows:
        print(fmt.format(*[str(c) for c in r]))


def cmd_status(args):
    data = request("GET", "/instances").get("data", [])
    if args.json:
        print(json.dumps(data, indent=2))
        return
    if not data:
        print("no instances")
        return
    rows = [
        (
            i.get("id", ""),
            i.get("name") or "-",
            (i.get("instance_type") or {}).get("name", ""),
            (i.get("region") or {}).get("name", ""),
            i.get("status", ""),
            i.get("ip") or "-",
        )
        for i in data
    ]
    print_table(("ID", "NAME", "TYPE", "REGION", "STATUS", "IP"), rows)


def _local_pub_keys():
    ssh_dir = os.path.expanduser("~/.ssh")
    if not os.path.isdir(ssh_dir):
        return []
    out = []
    for name in sorted(os.listdir(ssh_dir)):
        if not name.endswith(".pub"):
            continue
        path = os.path.join(ssh_dir, name)
        try:
            with open(path) as f:
                content = f.read().strip()
        except OSError:
            continue
        if content:
            out.append((path, content))
    return out


def _normalize_pubkey(s):
    parts = s.strip().split()
    return " ".join(parts[:2]) if len(parts) >= 2 else s.strip()


def _match_local_to_registered(registered):
    by_norm = {_normalize_pubkey(r.get("public_key", "")): r for r in registered}
    for path, content in _local_pub_keys():
        match = by_norm.get(_normalize_pubkey(content))
        if match:
            return match, path
    return None, None


DEFAULT_IMAGE_FAMILY = "lambda-stack-24-04"


def cmd_images(args):
    data = request("GET", "/images").get("data", [])
    if args.json:
        print(json.dumps(data, indent=2))
        return
    if not data:
        print("no images")
        return
    rows = [
        (
            i.get("id", ""),
            i.get("family", ""),
            i.get("version", ""),
            (i.get("region") or {}).get("name", ""),
            i.get("name", ""),
        )
        for i in data
    ]
    print_table(("ID", "FAMILY", "VERSION", "REGION", "NAME"), rows)


def cmd_ssh_keys(args):
    data = request("GET", "/ssh-keys").get("data", [])
    if args.json:
        print(json.dumps(data, indent=2))
        return
    if not data:
        print("no SSH keys")
        return
    rows = [(k.get("id", ""), k.get("name", "")) for k in data]
    print_table(("ID", "NAME"), rows)


def cmd_launch(args):
    ssh_key_name = args.ssh_key
    if not ssh_key_name:
        registered = request("GET", "/ssh-keys").get("data", [])
        match, path = _match_local_to_registered(registered)
        if not match:
            lines = [
                "error: no local ~/.ssh/*.pub matched a registered Lambda SSH key.",
                "  Re-run with --ssh-key <name>. Registered keys:",
            ]
            if registered:
                lines.extend(f"    {r.get('name')}" for r in registered)
            else:
                lines.append("    (none registered — add one in the Lambda dashboard)")
            sys.exit("\n".join(lines))
        ssh_key_name = match["name"]
        print(
            f"using SSH key '{ssh_key_name}' (matched {path})",
            file=sys.stderr,
        )
    body = {
        "region_name": args.region,
        "instance_type_name": args.type,
        "ssh_key_names": [ssh_key_name],
        "quantity": args.quantity,
    }
    if args.name:
        body["name"] = args.name
    if args.filesystem:
        body["file_system_names"] = [args.filesystem]

    if args.image == "latest":
        pass
    elif args.image and args.image.startswith("family:"):
        body["image"] = {"family": args.image[len("family:"):]}
    elif args.image:
        body["image"] = {"id": args.image}
    else:
        body["image"] = {"family": DEFAULT_IMAGE_FAMILY}
    result = request("POST", "/instance-operations/launch", body)
    ids = result.get("data", {}).get("instance_ids", [])

    pubkeys = []
    if args.add_keys and ids:
        registered = request("GET", "/ssh-keys").get("data", [])
        pubkeys = [
            r["public_key"].strip()
            for r in registered
            if r.get("public_key", "").strip()
        ]

    for iid in ids:
        print(iid)
        print(f"  waiting for {iid} to become active...", file=sys.stderr)
        info = _wait_for_active(iid)
        ip = info["ip"]
        if args.add_keys and pubkeys:
            print(
                f"  adding {len(pubkeys)} SSH key(s) to {ip}...",
                file=sys.stderr,
            )
            _ssh_add_keys(ip, args.ssh_user, pubkeys)
        print(f"  to connect:   ssh {args.ssh_user}@{ip}", file=sys.stderr)
        print(f"  to terminate: llcloud terminate {iid}", file=sys.stderr)


def _wait_for_active(instance_id, timeout=600, interval=10):
    deadline = time.time() + timeout
    while time.time() < deadline:
        info = request("GET", f"/instances/{instance_id}").get("data", {})
        if info.get("status") == "active" and info.get("ip"):
            return info
        time.sleep(interval)
    sys.exit(f"error: instance {instance_id} not active after {timeout}s")


def _ssh_ready(ip, user, timeout=300, interval=5):
    deadline = time.time() + timeout
    base = [
        "ssh",
        "-o", "StrictHostKeyChecking=accept-new",
        "-o", "BatchMode=yes",
        "-o", "ConnectTimeout=5",
        f"{user}@{ip}",
        "true",
    ]
    while time.time() < deadline:
        rc = subprocess.call(base, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
        if rc == 0:
            return
        time.sleep(interval)
    sys.exit(f"error: SSH to {user}@{ip} not ready after {timeout}s")


def _ssh_add_keys(ip, user, pubkeys):
    _ssh_ready(ip, user)
    payload = "\n".join(pubkeys) + "\n"
    remote = (
        "umask 077; mkdir -p ~/.ssh; touch ~/.ssh/authorized_keys; "
        "while IFS= read -r line; do "
        '[ -z "$line" ] && continue; '
        'grep -qxF -- "$line" ~/.ssh/authorized_keys || echo "$line" >> ~/.ssh/authorized_keys; '
        "done"
    )
    proc = subprocess.run(
        [
            "ssh",
            "-o", "StrictHostKeyChecking=accept-new",
            "-o", "BatchMode=yes",
            f"{user}@{ip}",
            remote,
        ],
        input=payload,
        text=True,
        capture_output=True,
    )
    if proc.returncode != 0:
        sys.exit(f"error: ssh key install failed on {ip}: {proc.stderr.strip()}")


def cmd_terminate(args):
    result = request(
        "POST",
        "/instance-operations/terminate",
        {"instance_ids": args.instance_ids},
    )
    for i in result.get("data", {}).get("terminated_instances", []):
        print(f"terminated {i.get('id')}")


def main():
    p = argparse.ArgumentParser(prog="llcloud", description="Lambda Labs cloud CLI")
    sub = p.add_subparsers(dest="cmd", required=True)

    s = sub.add_parser("status", aliases=["list", "ls"], help="list instances")
    s.add_argument("--json", action="store_true", help="emit JSON")
    s.set_defaults(func=cmd_status)

    l = sub.add_parser("launch", aliases=["start"], help="launch an instance")
    l.add_argument("--type", required=True, help="instance type (e.g. gpu_1x_a10)")
    l.add_argument("--region", required=True, help="region (e.g. us-east-1)")
    l.add_argument(
        "--ssh-key",
        help="registered SSH key name (auto-detected from ~/.ssh/*.pub if omitted)",
    )
    l.add_argument("--name", help="instance name")
    l.add_argument("--filesystem", help="filesystem to attach")
    l.add_argument("--quantity", type=int, default=1, help="number to launch")
    l.add_argument(
        "--image",
        help=(
            "image override: <id>, `family:<name>`, or `latest` (API default). "
            f"Defaults to `family:{DEFAULT_IMAGE_FAMILY}`."
        ),
    )
    l.add_argument(
        "--add-keys",
        action=argparse.BooleanOptionalAction,
        default=True,
        help="append all registered SSH keys to the instance's authorized_keys after it comes up (default: on; pass --no-add-keys to skip)",
    )
    l.add_argument(
        "--ssh-user",
        default="ubuntu",
        help="SSH user for key install + ssh hint (default: ubuntu)",
    )
    l.set_defaults(func=cmd_launch)

    k = sub.add_parser("ssh-keys", aliases=["keys"], help="list registered SSH keys")
    k.add_argument("--json", action="store_true", help="emit JSON")
    k.set_defaults(func=cmd_ssh_keys)

    im = sub.add_parser("images", help="list available machine images")
    im.add_argument("--json", action="store_true", help="emit JSON")
    im.set_defaults(func=cmd_images)

    t = sub.add_parser("terminate", aliases=["stop", "kill"], help="terminate instances")
    t.add_argument("instance_ids", nargs="+", help="instance ids")
    t.set_defaults(func=cmd_terminate)

    args = p.parse_args()
    args.func(args)


if __name__ == "__main__":
    main()