gpu-dev 0.5.18__tar.gz → 0.5.19__tar.gz

{gpu_dev-0.5.18 → gpu_dev-0.5.19}/PKG-INFO

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

{gpu_dev-0.5.18 → gpu_dev-0.5.19}/cli-tools/gpu-dev-cli/gpu_dev.egg-info/PKG-INFO

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

{gpu_dev-0.5.18 → gpu_dev-0.5.19}/cli-tools/gpu-dev-cli/gpu_dev.egg-info/SOURCES.txt

@@ -112,4 +112,8 @@ terraform-gpu-devservers/ssh-proxy/requirements.txt
 terraform-gpu-devservers/templates/al2023-cpu-user-data.sh
 terraform-gpu-devservers/templates/al2023-user-data.sh
 terraform-gpu-devservers/templates/user-data-self-managed.sh
-terraform-gpu-devservers/templates/user-data.sh
+terraform-gpu-devservers/templates/user-data.sh
+tests/submit/README.md
+tests/submit/fail/run.sh
+tests/submit/multinode/run.sh
+tests/submit/success/run.sh

{gpu_dev-0.5.18 → gpu_dev-0.5.19}/cli-tools/gpu-dev-cli/gpu_dev_cli/cli.py

@@ -1357,18 +1357,26 @@ _SUBMIT_GPU_TYPES = ["b200", "b200-mig-1g", "b200-mig-2g", "b200-mig-3g", "h200"
 @main.command(context_settings={"ignore_unknown_options": True})
 @click.option("--gpu-type", type=click.Choice(_SUBMIT_GPU_TYPES, case_sensitive=False), default="a100", show_default=True)
 @click.option("--gpus", type=int, default=1, show_default=True, help="GPU count (multinode if > per-node max).")
-@click.option("--hours", type=float, default=1.0, show_default=True, help="Reservation duration ceiling (job auto-cancels on exit).")
+@click.option("--hours", type=float, default=1.0, show_default=True, help="Reservation lifetime ceiling — job auto-cancels well before this if it finishes.")
 @click.option("--disk", type=str, default=None, help="Persistent disk name (master node only). Omit for ephemeral storage.")
 @click.option("--no-persistent-disk", is_flag=True, help="Skip persistent disk entirely.")
+@click.option("--dockerfile", type=click.Path(exists=True, dir_okay=False, resolve_path=True), default=None,
+              help="Local Dockerfile to build into the pod image (build context = the Dockerfile's directory).")
+@click.option("--dockerimage", type=str, default=None,
+              help="Pre-built container image reference (e.g. ghcr.io/me/img:tag) to run instead of the default.")
+@click.option("--preserve-entrypoint", is_flag=True,
+              help="Keep the custom image's ENTRYPOINT/CMD instead of letting gpu-dev wrap with the SSH harness. Note: submit needs SSH to work.")
 @click.option("--runtime", type=click.Path(exists=True, file_okay=False, resolve_path=True), default=None,
               help="Local directory to rsync to /workspace/submit-<id>/ on master node before run.")
 @click.option("--no-pull", is_flag=True, help="Skip syncing the remote workspace back to --runtime after the job finishes.")
 @click.option("--keep-alive", is_flag=True, help="Don't cancel the reservation when the job exits.")
 @click.option("--name", type=str, default=None, help="Reservation name.")
-@click.option("--timeout", type=int, default=20, show_default=True, help="Minutes to wait for the reservation to become active.")
+@click.option("--timeout", type=int, default=24 * 60, show_default=True,
+              help="Minutes to wait for the reservation to become active. Defaults to 24h since GPU reservations may queue when the cluster is full.")
 @click.argument("command", nargs=-1, required=True)
 @click.pass_context
-def submit(ctx, gpu_type, gpus, hours, disk, no_persistent_disk, runtime, no_pull, keep_alive, name, timeout, command):
+def submit(ctx, gpu_type, gpus, hours, disk, no_persistent_disk, dockerfile, dockerimage, preserve_entrypoint,
+           runtime, no_pull, keep_alive, name, timeout, command):
     """Submit a job: reserve, sync code, run, sync results back, auto-cancel.
 
     \b
@@ -1421,12 +1429,47 @@ def submit(ctx, gpu_type, gpus, hours, disk, no_persistent_disk, runtime, no_pul
     # SSH into rank 0, so passing --disk is fine.
     disk_name = None if no_persistent_disk else disk
 
+    # Build dockerfile context if provided (mirrors the reserve-flow logic)
+    dockerfile_payload = None
+    if dockerfile:
+        import os, tarfile, tempfile, base64
+        if os.path.getsize(dockerfile) > 512 * 1024:
+            rprint("[red]❌ Dockerfile too large (max 512KB)[/red]")
+            sys.exit(2)
+        ctx_dir = os.path.dirname(os.path.abspath(dockerfile))
+        rprint(f"[cyan]📦 Building tar.gz context from {ctx_dir}[/cyan]")
+        with tempfile.NamedTemporaryFile(suffix=".tar.gz", delete=False) as tmp:
+            with tarfile.open(tmp.name, "w:gz") as tar:
+                for root, _, files in os.walk(ctx_dir):
+                    for f in files:
+                        full = os.path.join(root, f)
+                        tar.add(full, arcname=os.path.relpath(full, ctx_dir))
+                if os.path.basename(dockerfile).lower() != "dockerfile":
+                    tar.add(dockerfile, arcname="Dockerfile")
+            tar_size = os.path.getsize(tmp.name)
+            if tar_size > 700 * 1024:
+                os.unlink(tmp.name)
+                rprint(f"[red]❌ Build context too large: {tar_size}B (max ~700KB compressed)[/red]")
+                sys.exit(2)
+            with open(tmp.name, "rb") as fh:
+                dockerfile_payload = base64.b64encode(fh.read()).decode("utf-8")
+            os.unlink(tmp.name)
+        rprint(f"[green]✅ Dockerfile context: {tar_size}B compressed[/green]")
+
+    if dockerimage and not preserve_entrypoint:
+        rprint("[dim]Note: passing --dockerimage without --preserve-entrypoint, so gpu-dev wraps the image with the SSH harness.[/dim]")
+    if preserve_entrypoint and not (dockerfile or dockerimage):
+        rprint("[red]❌ --preserve-entrypoint requires --dockerfile or --dockerimage[/red]")
+        sys.exit(2)
+
     rprint(f"[cyan]🎫 Reserving {gpus}x {gpu_type.upper()} for up to {hours}h...[/cyan]")
     if is_multinode:
         reservation_ids = rm.create_multinode_reservation(
             user_id=user_info["user_id"], gpu_count=gpus, gpu_type=gt,
             duration_hours=hours, name=name, github_user=user_info["github_user"],
-            no_persistent_disk=no_persistent_disk, disk_name=disk_name)
+            no_persistent_disk=no_persistent_disk, disk_name=disk_name,
+            dockerfile=dockerfile_payload, dockerimage=dockerimage,
+            preserve_entrypoint=preserve_entrypoint)
         if not reservation_ids:
             rprint("[red]❌ Failed to create multinode reservation[/red]")
             sys.exit(2)
@@ -1435,7 +1478,9 @@ def submit(ctx, gpu_type, gpus, hours, disk, no_persistent_disk, runtime, no_pul
         primary_id = rm.create_reservation(
             user_id=user_info["user_id"], gpu_count=gpus, gpu_type=gt,
             duration_hours=hours, name=name, github_user=user_info["github_user"],
-            no_persistent_disk=no_persistent_disk, disk_name=disk_name)
+            no_persistent_disk=no_persistent_disk, disk_name=disk_name,
+            dockerfile=dockerfile_payload, dockerimage=dockerimage,
+            preserve_entrypoint=preserve_entrypoint)
         if not primary_id:
             rprint("[red]❌ Failed to create reservation[/red]")
             sys.exit(2)
@@ -1456,7 +1501,11 @@ def submit(ctx, gpu_type, gpus, hours, disk, no_persistent_disk, runtime, no_pul
                 rprint(f"[dim]   cancel {rid[:8]} failed: {ce}[/dim]")
 
     try:
-        rprint(f"[cyan]⏳ Waiting for reservation {short_id} to become active (up to {timeout}m)...[/cyan]")
+        if timeout >= 60:
+            wait_str = f"up to {timeout//60}h{(" " + str(timeout%60) + "m") if timeout%60 else ""}"
+        else:
+            wait_str = f"up to {timeout}m"
+        rprint(f"[cyan]⏳ Waiting for reservation {short_id} to become active ({wait_str}; can queue when cluster is full)...[/cyan]")
         if is_multinode:
             results = rm.wait_for_multinode_reservation_completion(reservation_ids, timeout_minutes=timeout)
         else:

{gpu_dev-0.5.18 → gpu_dev-0.5.19}/pyproject.toml

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

gpu_dev-0.5.19/tests/submit/README.md

@@ -0,0 +1,63 @@
+# `gpu-dev submit` smoke tests
+
+Three tests that exercise `gpu-dev submit` end-to-end. Each test lives in its
+own folder so you can `--runtime` it directly. Output files written by the
+script are pulled back into the same folder via the post-run rsync.
+
+> Requires `gpu-dev >= 0.5.19`. No Lambda update needed.
+
+## 1. success — single T4 GPU, exit 0
+
+```bash
+cd tests/submit/success
+gpu-dev submit --gpu-type t4 --gpus 1 --runtime ./ -- bash run.sh
+echo $?    # 0
+ls         # nvidia-info.txt, compute.txt, status.txt all created
+```
+
+## 2. fail — single T4 GPU, exit 7
+
+Writes a partial file before exploding so you can confirm rsync still pulls
+output on failure and the local exit code is the remote's.
+
+```bash
+cd tests/submit/fail
+gpu-dev submit --gpu-type t4 --gpus 1 --runtime ./ -- bash run.sh
+echo $?    # 7
+ls         # step1.txt, step2.txt, gpus-before-fail.txt — but no step3.txt
+```
+
+## 3. multinode — 2x H100 nodes, exit 0
+
+Reserves 16 H100s (= 2 nodes), verifies env vars + peer ssh + NCCL all_reduce
+across the whole cluster via mpirun (orchestrated entirely from rank 0).
+
+```bash
+cd tests/submit/multinode
+gpu-dev submit --gpu-type h100 --gpus 16 --runtime ./ -- bash run.sh
+echo $?    # 0
+cat multinode-env.txt resolved-ips.txt peer-ssh.txt nccl-all_reduce.log
+```
+
+## What each test proves
+
+| Test       | Proves                                                                        |
+|------------|-------------------------------------------------------------------------------|
+| success    | reserve → rsync up → exec → rsync back → cancel → exit 0                      |
+| fail       | exit code propagation; rsync-back still runs on non-zero exit; cancel fires   |
+| multinode  | MULTINODE_* env vars; peer DNS / passwordless ssh; cross-node NCCL via mpirun |
+
+After every run, `gpu-dev list` should show neither reservation — both auto-cancelled.
+Use `--keep-alive` on any of them if you want to debug interactively afterward.
+
+## Other submit flags (forwarded to `reserve`)
+
+- `--hours N` — reservation lifetime ceiling (default 1.0)
+- `--disk NAME` — attach a persistent disk to the master node
+- `--no-persistent-disk` — skip persistent disk
+- `--dockerfile PATH` — build a custom image from this Dockerfile
+- `--dockerimage REF` — use a pre-built container image
+- `--preserve-entrypoint` — keep the custom image's ENTRYPOINT (you must run sshd yourself for submit to work)
+- `--timeout MINUTES` — wait-for-active timeout (default 1440 = 24h, since reservations may queue)
+- `--no-pull` — skip the post-run sync-back
+- `--keep-alive` — skip auto-cancel

gpu_dev-0.5.19/tests/submit/fail/run.sh

@@ -0,0 +1,20 @@
+#!/usr/bin/env bash
+# Failure test for `gpu-dev submit`: writes a partial output, then exits 7.
+# Verifies the post-run rsync still pulls the partial files even on failure,
+# the auto-cancel runs on non-zero exit, and the local exit code is preserved.
+set -e
+
+echo "=== host ==="
+hostname
+date -u
+
+# Write a partial file so we can verify it was synced back
+echo "step1 done at $(date -u)" > step1.txt
+nvidia-smi -L > gpus-before-fail.txt
+
+# Now error out
+echo "About to fail..." > step2.txt
+python3 -c "import sys; sys.exit(7)"
+
+# Should not reach here
+echo "should-not-appear" > step3.txt

gpu_dev-0.5.19/tests/submit/multinode/run.sh

@@ -0,0 +1,65 @@
+#!/usr/bin/env bash
+# Multinode test for `gpu-dev submit`: runs on rank 0 only and orchestrates the
+# whole cluster via mpirun (uses passwordless ssh + the headless service DNS we
+# already set up). Verifies env vars, peer connectivity, and an actual NCCL
+# all_reduce across all nodes.
+set -euo pipefail
+cd "$(dirname "$0")"
+
+echo "=== rank 0 host: $(hostname) at $(date -u) ==="
+
+echo "=== multinode env ==="
+{
+  echo "MULTINODE_HOSTS=$MULTINODE_HOSTS"
+  echo "MULTINODE_PEER_PODS=$MULTINODE_PEER_PODS"
+  echo "MULTINODE_RANK=$MULTINODE_RANK"
+  echo "MULTINODE_SIZE=$MULTINODE_SIZE"
+  echo "MASTER_ADDR=$MASTER_ADDR"
+  echo "MASTER_PORT=$MASTER_PORT"
+  echo "MULTINODE_IPS=${MULTINODE_IPS:-(not set)}"
+} | tee multinode-env.txt
+
+if [[ -z "${MULTINODE_HOSTS:-}" ]]; then
+    echo "ERROR: MULTINODE_HOSTS empty — submit with --gpus >= 16 on h100" >&2
+    exit 2
+fi
+
+# Resolve IPs even if the bashrc helper didn't run (defensive)
+IPS=""
+for h in $(echo "$MULTINODE_HOSTS" | tr ',' ' '); do
+    ip=$(getent hosts "$h" | awk '{print $1}' | head -1)
+    [[ -n "$ip" ]] && IPS="${IPS:+$IPS,}$ip"
+done
+echo "Resolved IPS=$IPS" | tee resolved-ips.txt
+
+echo "=== peer ssh check (port 2222 inside cluster) ==="
+peer_host=$(echo "$MULTINODE_HOSTS" | cut -d, -f2)
+ssh -o StrictHostKeyChecking=no -p 2222 "$peer_host" 'hostname; nvidia-smi -L | wc -l' \
+    | tee peer-ssh.txt
+
+GPUS_PER_NODE=$(nvidia-smi -L | wc -l)
+echo "GPUS_PER_NODE=$GPUS_PER_NODE" | tee gpus-per-node.txt
+
+# Build --host arg: ip1:N,ip2:N,...
+HOST_ARG=$(echo "$IPS" | awk -v g="$GPUS_PER_NODE" -F, '{out=""; for(i=1;i<=NF;i++){out=out ($i ":" g) (i<NF?",":"")}; print out}')
+echo "HOST_ARG=$HOST_ARG"
+
+echo "=== NCCL all_reduce_perf via mpirun ==="
+# Note: -g 1 = 1 GPU per process, -n 20 iterations. Sweep 1M..1G in factor-of-2 steps.
+mpirun --host "$HOST_ARG" \
+    --mca plm_rsh_args "-p 2222 -o StrictHostKeyChecking=no" \
+    -x PATH -x LD_LIBRARY_PATH \
+    -x FI_PROVIDER -x FI_EFA_USE_DEVICE_RDMA \
+    -x NCCL_NET_GDR_LEVEL -x NCCL_ALGO \
+    -x NCCL_SOCKET_IFNAME -x NCCL_DEBUG -x NCCL_IB_HCA \
+    /opt/nccl-tests/build/all_reduce_perf -b 1M -e 1G -f 2 -g 1 -n 20 \
+    2>&1 | tee nccl-all_reduce.log
+
+echo "=== summary ==="
+{
+    echo "rank=$MULTINODE_RANK size=$MULTINODE_SIZE"
+    echo "host_arg=$HOST_ARG"
+    echo "completed at $(date -u)"
+} | tee summary.txt
+
+echo "DONE"

gpu_dev-0.5.19/tests/submit/success/run.sh

@@ -0,0 +1,23 @@
+#!/usr/bin/env bash
+# Smoke test for `gpu-dev submit`: runs on a single GPU, expected exit 0.
+set -euo pipefail
+
+echo "=== host ==="
+hostname
+date -u
+
+echo "=== nvidia-smi ==="
+nvidia-smi | tee nvidia-info.txt
+
+echo "=== compute ==="
+python3 - <<'PY' | tee compute.txt
+import torch
+assert torch.cuda.is_available(), "CUDA not available"
+n = torch.cuda.device_count()
+x = torch.arange(1_000_000, device="cuda", dtype=torch.float32)
+s = x.sum().item()
+print(f"devices={n} sum(0..999_999)={s}")
+PY
+
+echo "ok at $(date -u)" > status.txt
+echo "DONE"