gpu-dev 0.7.3__tar.gz → 0.7.5__tar.gz

{gpu_dev-0.7.3 → gpu_dev-0.7.5}/CLAUDE.md

@@ -79,7 +79,7 @@ Big push on warm pools + instant claims + prebuilt pytorch. Tracking state here
 - [ ] **Prebuilt built WITHOUT cuDNN** — `import torch` warns "compiled without cuDNN/MIOpen". CI/nightly build with cudnn9. Add libcudnn to the gpu-dev image + `USE_CUDNN=1` to the build recipe for fidelity (conv/cudnn-dependent ops + tests). Irrelevant for flex-attention int64 test; matters generally.
 - [ ] **`--ref pr/N` uses `pull/N/head`, not `/merge`** — `/head` is the PR author's raw branch tip (often based on old trunk, missing trunk-added tests); CI tests `/merge` (PR merged onto current trunk). For CI-repro fidelity, `pr/N` should fetch `pull/N/merge` (fall back to `/head` if no merge ref). `stage-pytorch` REF case in `index.py`. (This is why `pull/185479/head` lacked `test_large_kv_int64_pointer_math_cuda`.)
 - [ ] **Misleading disconnect/expiry message** — on `gpu-dev connect` connection loss OR reservation expiry, the CLI prints "❌ Authentication failed. You don't have SSH access... ask the primary user to add you" even for the PRIMARY user's own expired/cancelled reservation. Distinguish: (a) reservation expired -> "Reservation <id> expired at <time>"; (b) cancelled -> "Reservation was cancelled"; (c) connection dropped but still active -> "Connection lost, reconnect with gpu-dev connect <id>"; (d) genuine auth failure -> the current add-user message. Check reservation status before assuming auth failure.
-- [ ] **`gpu-dev cancel` from inside the pod** — show "Shutting down this reservation..." (graceful message) instead of an abrupt SSH drop, so the user knows the disconnect was intentional.
+- [x] **`gpu-dev cancel` from inside the pod** (DONE, 0.7.5) — two bugs: (1) cancel inside a **warm-claimed** pod failed with "GitHub username not configured" because the warm pod was pre-booted with `user_id="warm"` and the claim never stamped the real identity → `GPU_DEV_USER_ID/GPU_DEV_GITHUB_USER/AWS_ROLE_SESSION_NAME` stayed `"warm"`/empty. Fix: `try_claim_warm_pod` now seds the real `user_id`/`github_user` into both `.bashrc_ext`/`.zshrc_ext` + writes `GPU_DEV_RESERVATION_ID` (full id). Cold `_ext` derives `GPU_DEV_RESERVATION_ID` from the hostname (8-char prefix; cancellation resolves by prefix). (2) `gpu-dev cancel` (no id) inside a pod now fast-paths: cancels THIS reservation directly via `GPU_DEV_RESERVATION_ID`+`GPU_DEV_USER_ID` (no github_user/interactive) with the graceful "🛑 Shutting down..." message. Needs `tf apply` (lambda) + image rebuild (CLI in pods).
 - [ ] SSH CA certs to drop the ~0.33s `kubectl exec` key injection on warm claim (auth-model change).
 - [ ] AMI baker re-bakes on every base-EKS-AMI roll (5 baked AMIs in 2 days): pin the base AMI version + clean up old `gpu-dev-baked-*`.
 - [ ] **Warm pods: gate `warm-state=ready` on staging completion** (NOW MORE IMPORTANT — the built tree is ~30GB, and on GPU nodes it's a `cp` not reflink, so staging takes ~1-3min; a claim in that window hands over a half-copied tree). Two options: (a) claim-time check — exec `[ -f /home/dev/.pytorch-staging ]` in `try_claim_warm_pod`, skip pods still staging (simple, but adds ~0.5s exec to every warm claim); (b) label-flip — create with `warm-state=provisioning`, reconciler exec-checks staging + flips to `ready` (no claim latency, but 4 interacting changes: create label + reconciler flip + eviction must also target `provisioning` + claim already filters `ready`). Prefer (b). Marker: `.pytorch-staging` present during, removed when done; `.pytorch-ready` written at end.

{gpu_dev-0.7.3 → gpu_dev-0.7.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpu-dev
-Version: 0.7.3
+Version: 0.7.5
 Summary: CLI + Python SDK for PyTorch GPU developer server reservations
 Author: PyTorch Team
 Requires-Python: >=3.10

{gpu_dev-0.7.3 → gpu_dev-0.7.5}/cli-tools/gpu-dev-cli/gpu_dev_cli/cli.py

@@ -2668,6 +2668,24 @@ def cancel(
             rprint("[red]❌ Cannot specify both --all and a reservation ID[/red]")
             return
 
+        # Inside a gpu-dev pod, `gpu-dev cancel` (no id) shuts down THIS reservation
+        # directly. The pod knows its own reservation (GPU_DEV_RESERVATION_ID) and
+        # owner (GPU_DEV_USER_ID), so we skip the github_user / interactive list —
+        # which can't work in a pod that has no `gpu-dev config set github_user`.
+        pod_rid = os.environ.get("GPU_DEV_RESERVATION_ID", "").strip()
+        pod_uid = os.environ.get("GPU_DEV_USER_ID", "").strip()
+        if pod_rid and pod_uid and pod_uid != "warm" and not reservation_id and not all:
+            rprint("[yellow]🛑 Shutting down this reservation — if you're connected to this pod, your session will close shortly.[/yellow]")
+            try:
+                reservation_mgr = ReservationManager(load_config())
+                ok = reservation_mgr.cancel_reservation(pod_rid, pod_uid)
+            except Exception as e:
+                ok = False
+                rprint(f"[red]❌ Could not cancel from inside the pod: {e}[/red]")
+            if not ok:
+                rprint(f"[dim]If that didn't work, cancel from your laptop: gpu-dev cancel {pod_rid[:8]}[/dim]")
+            return
+
         # Handle --all flag (non-interactive)
         if all:
             with Live(

{gpu_dev-0.7.3 → gpu_dev-0.7.5}/cli-tools/gpu-dev-cli/gpu_dev_cli/interactive.py

@@ -159,37 +159,51 @@ def select_gpu_type_interactive(
             parts.append(f"{a}×{cgt.rsplit('-', 1)[-1].upper()}")
         return parts, tot_a, tot_c
 
-    # ── The selectable list IS the table ──────────────────────────────────────
-    # questionary indents Separators and Choices identically, so a Separator
-    # header + aligned column text line up with the selectable rows. Arrow keys
-    # move through the table; Enter picks the highlighted row. No separate print.
-    def _row_cells(gt, info, is_spot=False):
+    # ── The selectable list IS the table (boxed + colored) ────────────────────
+    # Box-drawing borders are non-selectable Separators; each GPU is a Choice whose
+    # title is FormattedText so cells are individually colored. questionary indents
+    # Separators and Choices identically, so the borders line up with the rows.
+    # Arrow keys move through the table; the » pointer marks the row, Enter picks it.
+    # Emoji are kept OUT of cells (double-width → would ragged the right border);
+    # status is conveyed with color instead.
+    G, R, Y, BL, CY, MG, DIM = (
+        "fg:ansigreen", "fg:ansired", "fg:ansiyellow",
+        "fg:ansiblue", "fg:ansicyan bold", "fg:ansimagenta", "fg:#808080")
+
+    def _status(info):
         avail = int(info.get("available", 0))
-        wd, emoji = _format_wait(avail, info.get("estimated_wait_minutes", 0))
+        est = info.get("estimated_wait_minutes", 0)
         ql = int(info.get("queue_length", 0))
+        if avail > 0:
+            text, style = "available now", G
+        elif est:
+            text, style = _format_wait(avail, est)[0], Y
+        else:
+            text, style = "queued", Y
         if ql > 0:
-            wd += f" · {ql} queued"
-        typ = f"{gt.upper()} *" if is_spot else gt.upper()
-        return [typ, str(avail), str(int(info.get("max_reservable", 0))),
-                str(int(info.get("total", 0)))], f"{emoji} {wd}"
+            text += f" · {ql} queued"
+        return text, style
 
-    # Rows: (cells[type, avail, maxres, total], status, value, kind).
-    data_rows = []
+    # rows: (cells[type, avail, maxres, total, status], styles|None, value, kind)
+    rows = []
     for gt, info in full_gpus.items():
         if info.get("maintenance", False):
-            data_rows.append((
-                [gt.upper(), "-", "-", str(int(info.get("total", 0)))],
-                f"MAINTENANCE: {info.get('maintenance_reason', '')}", gt, "maint"))
+            reason = (info.get("maintenance_reason", "") or "maintenance")[:18]
+            rows.append(([gt.upper(), "-", "-", str(int(info.get("total", 0))),
+                          f"MAINT: {reason}"], None, gt, "maint"))
             continue
-        cells, status = _row_cells(gt, info)
-        data_rows.append((cells, status, gt, "gpu"))
+        a = int(info.get("available", 0))
+        st_text, st_style = _status(info)
+        cells = [gt.upper(), str(a), str(int(info.get("max_reservable", 0))),
+                 str(int(info.get("total", 0))), st_text]
+        styles = [CY, G if a > 0 else R, G, BL, st_style]
+        rows.append((cells, styles, gt, "gpu"))
         parts, mig_a, mig_c = _mig_breakdown(gt)
         if parts:
-            data_rows.append((
-                [" └─ MIG", str(mig_a), "-", str(mig_c)],
-                f"{' '.join(parts)} · pick {gt.upper()} ↑", None, "mig"))
+            rows.append(([" └─ MIG", str(mig_a), "-", str(mig_c), " ".join(parts)],
+                         None, None, "mig"))
 
-    spot_data = []
+    spot_rows = []
     if spot_gpus:
         _pn = {"b300": 8, "b200": 8, "h200": 8, "h100": 8, "a100": 8, "t4": 4, "l4": 4}
         _od = {"b300": 95, "b200": 95, "h200": 55, "h100": 98, "a100": 32, "t4": 4.5, "l4": 7}
@@ -205,40 +219,48 @@ def select_gpu_type_interactive(
                 try:
                     disc = f"~{int((1 - float(sp) / _od.get(gt, 50)) * 100)}% off"
                 except (ValueError, TypeError):
-                    disc = "spot price n/a"
-            status = ("✅ node up" if avail > 0 else "⚡ spins up ~10min") + f" · {disc}"
-            spot_data.append((
-                [f"{gt.upper()} *", str(avail), f"{_pn.get(gt, 8)}/node", "-"],
-                status, f"spot:{gt}", "spot"))
+                    disc = "spot n/a"
+            st_text = ("node up · " if avail > 0 else "spins up · ") + disc
+            cells = [f"{gt.upper()} *", str(avail), f"{_pn.get(gt, 8)}/node", "-", st_text]
+            styles = [MG, G if avail > 0 else DIM, G, DIM, G if avail > 0 else Y]
+            spot_rows.append((cells, styles, f"spot:{gt}", "spot"))
+
+    headers = ["GPU Type", "Avail", "MaxRes", "Total", "Status"]
+    all_cells = [headers] + [r[0] for r in rows] + [s[0] for s in spot_rows]
+    W = [max(len(str(rc[i])) for rc in all_cells) for i in range(5)]
+
+    def _bar(left, mid, right):
+        return left + mid.join("─" * (w + 2) for w in W) + right
 
-    # Column widths over the 4 text columns (header + all rows).
-    headers = ["GPU Type", "Avail", "MaxRes", "Total"]
-    _all_cells = [headers] + [r[0] for r in data_rows] + [s[0] for s in spot_data]
-    widths = [max(len(str(row[i])) for row in _all_cells) for i in range(4)]
+    def _line(cells):  # plain string row (header / mig / maint), inside the box
+        return "│" + "│".join(f" {str(c):<{W[i]}} " for i, c in enumerate(cells)) + "│"
 
-    def _fmt(cells, status=""):
-        body = "  ".join(str(c).ljust(widths[i]) for i, c in enumerate(cells))
-        return f"{body}   {status}".rstrip()
+    def _ft(cells, styles):  # colored row -> FormattedText for a Choice
+        toks = [("class:separator", "│")]
+        for i, c in enumerate(cells):
+            toks.append((styles[i], f" {str(c):<{W[i]}} "))
+            toks.append(("class:separator", "│"))
+        return toks
 
     console.print()
-    choices = [questionary.Separator(_fmt(headers, "Status"))]
-    if not data_rows:
-        choices.append(questionary.Separator("(no GPU types available)"))
-    for cells, status, value, kind in data_rows:
-        title = _fmt(cells, status)
-        if kind == "mig":
-            choices.append(questionary.Separator(title))
-        elif kind == "maint":
-            choices.append(questionary.Choice(title=title, value=value, disabled="maintenance"))
+    choices = [questionary.Separator(_bar("┌", "┬", "┐")),
+               questionary.Separator(_line(headers)),
+               questionary.Separator(_bar("├", "┼", "┤"))]
+    if not rows:
+        choices.append(questionary.Separator(_line(["(none)", "", "", "", ""])))
+    for cells, styles, value, kind in rows:
+        if kind in ("mig", "maint"):
+            choices.append(questionary.Separator(_line(cells)))
         else:
-            choices.append(questionary.Choice(title=title, value=value))
+            choices.append(questionary.Choice(title=_ft(cells, styles), value=value))
 
-    if spot_data:
-        choices.append(questionary.Separator("⚡ Spot — us-east-1, ~70% cheaper, may be preempted:"))
-        for cells, status, value, _kind in spot_data:
-            choices.append(questionary.Choice(title=_fmt(cells, status), value=value))
+    if spot_rows:
+        choices.append(questionary.Separator(_bar("├", "┼", "┤")))
+        for cells, styles, value, _k in spot_rows:
+            choices.append(questionary.Choice(title=_ft(cells, styles), value=value))
 
-    choices.append(questionary.Separator("───"))
+    choices.append(questionary.Separator(_bar("└", "┴", "┘")))
+    choices.append(questionary.Separator(" "))
     if _hide_spot:
         choices.append(questionary.Choice(
             title="⚡ Show spot options (us-east-1, ~70% cheaper, may be preempted)",
@@ -251,7 +273,8 @@ def select_gpu_type_interactive(
     while True:
         try:
             answer = questionary.select(
-                "Select GPU type (↑/↓, Enter):", choices=choices, style=custom_style
+                "Select GPU type — ↑/↓ then Enter (MIG: pick its parent GPU):",
+                choices=choices, style=custom_style
             ).ask()
 
             if answer == "_refresh":

{gpu_dev-0.7.3 → gpu_dev-0.7.5}/gpu_dev.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpu-dev
-Version: 0.7.3
+Version: 0.7.5
 Summary: CLI + Python SDK for PyTorch GPU developer server reservations
 Author: PyTorch Team
 Requires-Python: >=3.10

{gpu_dev-0.7.3 → gpu_dev-0.7.5}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "gpu-dev"
-version = "0.7.3"
+version = "0.7.5"
 description = "CLI + Python SDK for PyTorch GPU developer server reservations"
 authors = [{name = "PyTorch Team"}]
 readme = "cli-tools/gpu-dev-cli/README.md"

{gpu_dev-0.7.3 → gpu_dev-0.7.5}/sdk/python/src/gpu_dev/__init__.py

@@ -63,4 +63,4 @@ try:
 
     __version__ = _pkg_version("gpu-dev")
 except Exception:
-    __version__ = "0.7.3"
+    __version__ = "0.7.5"

{gpu_dev-0.7.3 → gpu_dev-0.7.5}/terraform-gpu-devservers/lambda/reservation_processor/index.py

@@ -1633,7 +1633,18 @@ def try_claim_warm_pod(body: dict) -> bool:
             "while IFS= read -r k; do [ -n \"$k\" ] && ! grep -Fq \"$k\" /home/dev/.ssh/authorized_keys && echo \"$k\" >> /home/dev/.ssh/authorized_keys; done <<'KEOF'\n"
             f"{github_public_key}\n"
             "KEOF\n"
-            "chmod 700 /home/dev/.ssh && chmod 600 /home/dev/.ssh/authorized_keys && chown -R 1081:1081 /home/dev/.ssh"
+            "chmod 700 /home/dev/.ssh && chmod 600 /home/dev/.ssh/authorized_keys && chown -R 1081:1081 /home/dev/.ssh\n"
+            # Warm pods were pre-booted with user_id='warm'; stamp the real claimant's
+            # identity into the managed shell-ext files so `gpu-dev` inside the pod
+            # (cancel/list/...) authenticates as the user and IRSA assumes the right
+            # session. The user connects AFTER the claim, so their login shell picks
+            # these up. Also record the reservation id for `gpu-dev cancel`.
+            "for f in /home/dev/.bashrc_ext /home/dev/.zshrc_ext; do [ -f \"$f\" ] || continue\n"
+            f"  sed -i -e 's|^export GPU_DEV_USER_ID=.*|export GPU_DEV_USER_ID=\"{user_id}\"|'"
+            f" -e 's|^export GPU_DEV_GITHUB_USER=.*|export GPU_DEV_GITHUB_USER=\"{github_user}\"|'"
+            f" -e 's|^export AWS_ROLE_SESSION_NAME=.*|export AWS_ROLE_SESSION_NAME=\"{user_id}\"|' \"$f\"\n"
+            f"  grep -q '^export GPU_DEV_RESERVATION_ID=' \"$f\" && sed -i 's|^export GPU_DEV_RESERVATION_ID=.*|export GPU_DEV_RESERVATION_ID=\"{reservation_id}\"|' \"$f\" || echo 'export GPU_DEV_RESERVATION_ID=\"{reservation_id}\"' >> \"$f\"\n"
+            "done"
         )
         stream(
             v1.connect_get_namespaced_pod_exec, pod_name, "gpu-dev",
@@ -5272,6 +5283,10 @@ EOF_PROFILE
 
 # User identification
 export GPU_DEV_USER_ID="{user_id or 'dev'}"
+# Reservation id — from the pod hostname (gpu-dev-<id>). Warm claims overwrite this
+# in place with the full id; cold pods carry the 8-char prefix (cancellation resolves
+# by prefix). Lets `gpu-dev cancel` with no args inside the pod stop this reservation.
+export GPU_DEV_RESERVATION_ID="$(hostname | sed -e 's/^gpu-dev-//')"
 
 # Multinode peer info — inlined from container env at pod startup. sshd strips
 # container env vars from login shells, so we materialize the values into rc files.
@@ -5338,6 +5353,10 @@ EOF_BASHRC_EXT
 
 # User identification
 export GPU_DEV_USER_ID="{user_id or 'dev'}"
+# Reservation id — from the pod hostname (gpu-dev-<id>). Warm claims overwrite this
+# in place with the full id; cold pods carry the 8-char prefix (cancellation resolves
+# by prefix). Lets `gpu-dev cancel` with no args inside the pod stop this reservation.
+export GPU_DEV_RESERVATION_ID="$(hostname | sed -e 's/^gpu-dev-//')"
 
 # Multinode peer info — inlined from container env at pod startup. sshd strips
 # container env vars from login shells, so we materialize the values into rc files.