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

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

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

{gpu_dev-0.5.16 → gpu_dev-0.5.18}/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.16
+Version: 0.5.18
 Summary: CLI tool for PyTorch GPU developer server reservations
 Author: PyTorch Team
 Requires-Python: >=3.10

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

@@ -1349,6 +1349,205 @@ def reserve(
         rprint(f"[red]❌ Error: {str(e)}[/red]")
 
 
+_SUBMIT_GPU_TYPES = ["b200", "b200-mig-1g", "b200-mig-2g", "b200-mig-3g", "h200", "h100",
+                     "h100-mig-1g", "h100-mig-2g", "h100-mig-3g", "a100", "rtxpro6000",
+                     "a10g", "t4", "l4", "t4-small", "cpu-arm", "cpu-x86"]
+
+
+@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("--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("--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.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):
+    """Submit a job: reserve, sync code, run, sync results back, auto-cancel.
+
+    \b
+    Examples:
+      gpu-dev submit --runtime ./ -- python train.py
+      gpu-dev submit --gpus 16 --gpu-type h100 --runtime . -- bash run.sh
+      gpu-dev submit --keep-alive -- nvidia-smi
+
+    The job runs on rank 0 (master pod). For multinode jobs, MULTINODE_HOSTS / RANK /
+    SIZE / MASTER_ADDR / MASTER_PORT are exported on every pod so torchrun and friends
+    work without manual wiring. Exit code mirrors the remote command's exit code.
+    """
+    import subprocess
+    import shlex
+    import sys
+    from pathlib import Path
+
+    if not command:
+        rprint("[red]❌ Provide a command after --, e.g. gpu-dev submit --runtime ./ -- python train.py[/red]")
+        sys.exit(2)
+
+    gt = gpu_type.lower()
+    # Per-type max GPUs (mirrors gpu_configs in reserve flow)
+    max_per_node = {
+        "t4": 4, "l4": 4, "a10g": 4, "rtxpro6000": 4, "t4-small": 1,
+        "a100": 8, "h100": 8, "h200": 8, "b200": 8,
+        "h100-mig-1g": 16, "h100-mig-2g": 8, "h100-mig-3g": 8,
+        "b200-mig-1g": 4, "b200-mig-2g": 2, "b200-mig-3g": 2,
+        "cpu-arm": 0, "cpu-x86": 0,
+    }.get(gt)
+    if max_per_node is None:
+        rprint(f"[red]❌ Unknown gpu-type '{gpu_type}'[/red]")
+        sys.exit(2)
+
+    is_multinode = gt not in ("cpu-arm", "cpu-x86") and gpus > max_per_node
+    if is_multinode and gpus % max_per_node != 0:
+        rprint(f"[red]❌ For multinode {gt}, --gpus must be a multiple of {max_per_node}[/red]")
+        sys.exit(2)
+
+    config = load_config()
+    try:
+        user_info = authenticate_user(config)
+    except RuntimeError as e:
+        rprint(f"[red]❌ {str(e)}[/red]")
+        sys.exit(2)
+
+    rm = ReservationManager(config)
+
+    # Determine effective disk handling. Multinode: only master gets persistent disk; we always
+    # SSH into rank 0, so passing --disk is fine.
+    disk_name = None if no_persistent_disk else disk
+
+    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)
+        if not reservation_ids:
+            rprint("[red]❌ Failed to create multinode reservation[/red]")
+            sys.exit(2)
+        primary_id = reservation_ids[0]
+    else:
+        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)
+        if not primary_id:
+            rprint("[red]❌ Failed to create reservation[/red]")
+            sys.exit(2)
+        reservation_ids = [primary_id]
+
+    short_id = primary_id[:8]
+    cancelled = {"done": False}
+
+    def maybe_cancel(reason: str):
+        if cancelled["done"] or keep_alive:
+            return
+        cancelled["done"] = True
+        rprint(f"[yellow]🛑 Cancelling reservation {short_id} ({reason})[/yellow]")
+        for rid in reservation_ids:
+            try:
+                rm.cancel_reservation(rid, user_info["user_id"])
+            except Exception as ce:
+                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 is_multinode:
+            results = rm.wait_for_multinode_reservation_completion(reservation_ids, timeout_minutes=timeout)
+        else:
+            single = rm.wait_for_reservation_completion(primary_id, timeout_minutes=timeout)
+            results = [single] if single else None
+        if not results:
+            rprint("[red]❌ Reservation never became active[/red]")
+            maybe_cancel("activation timeout")
+            sys.exit(1)
+
+        # Resolve master pod (rank 0)
+        conn = rm.get_connection_info(primary_id, user_info["user_id"])
+        if not conn:
+            rprint("[red]❌ Could not fetch connection info[/red]")
+            maybe_cancel("no connection info")
+            sys.exit(1)
+        if conn.get("is_multinode"):
+            nodes = sorted(conn["nodes"], key=lambda n: n.get("node_index", 0))
+            master = nodes[0]
+            master_id, master_pod, master_fqdn, master_name = (
+                master["reservation_id"], master["pod_name"],
+                master.get("fqdn"), master.get("name"))
+        else:
+            master_id, master_pod, master_fqdn, master_name = (
+                primary_id, conn["pod_name"], conn.get("fqdn"), conn.get("name"))
+
+        # Ensure SSH config exists
+        gpu_dev_dir = Path.home() / ".gpu-dev"
+        config_file = gpu_dev_dir / f"{master_id[:8]}-sshconfig"
+        if not config_file.exists():
+            if not (master_fqdn and master_pod):
+                rprint("[red]❌ Master pod has no FQDN yet — can't SSH[/red]")
+                maybe_cancel("no fqdn")
+                sys.exit(1)
+            create_ssh_config_for_reservation(master_fqdn, master_pod, master_id, master_name)
+
+        ssh_alias = master_pod
+        ssh_base = ["ssh", "-F", str(config_file), "-o", "StrictHostKeyChecking=accept-new"]
+        rsync_e = " ".join(shlex.quote(x) for x in ssh_base)
+
+        # Working directory and rsync up
+        if runtime:
+            workdir = f"/workspace/submit-{master_id[:8]}"
+            rprint(f"[cyan]📦 Syncing {runtime} → {ssh_alias}:{workdir}[/cyan]")
+            r = subprocess.run(ssh_base + [ssh_alias, f"mkdir -p {shlex.quote(workdir)}"])
+            if r.returncode != 0:
+                rprint("[red]❌ Failed to create remote workspace[/red]")
+                maybe_cancel("mkdir failed"); sys.exit(2)
+            r = subprocess.run([
+                "rsync", "-az", "--delete", "-e", rsync_e,
+                f"{runtime.rstrip('/')}/", f"{ssh_alias}:{workdir}/",
+            ])
+            if r.returncode != 0:
+                rprint("[red]❌ Upload rsync failed[/red]")
+                maybe_cancel("upload failed"); sys.exit(2)
+        else:
+            workdir = "/home/dev"
+
+        # Run remote command via login shell so MULTINODE_* etc. are loaded
+        remote_cmd = " ".join(shlex.quote(c) for c in command)
+        rprint(f"[cyan]🚀 Running on {ssh_alias}: {remote_cmd}[/cyan]\n")
+        ssh_run = ssh_base + [ssh_alias,
+                              f"cd {shlex.quote(workdir)} && bash -lc {shlex.quote(remote_cmd)}"]
+        rc = subprocess.call(ssh_run)
+        rprint(f"\n[dim]Job exited with code {rc}[/dim]")
+
+        # Sync back results before cancelling
+        if runtime and not no_pull:
+            rprint(f"[cyan]📥 Syncing {ssh_alias}:{workdir}/ → {runtime}[/cyan]")
+            pull = subprocess.run([
+                "rsync", "-az", "-e", rsync_e,
+                f"{ssh_alias}:{workdir}/", f"{runtime.rstrip('/')}/",
+            ])
+            if pull.returncode != 0:
+                rprint(f"[yellow]⚠️  Result rsync exited with {pull.returncode} — your output may be incomplete[/yellow]")
+
+        maybe_cancel("job complete")
+        sys.exit(rc)
+
+    except KeyboardInterrupt:
+        rprint("\n[yellow]Interrupted — cancelling[/yellow]")
+        maybe_cancel("user interrupt")
+        sys.exit(130)
+    except SystemExit:
+        raise
+    except Exception as e:
+        rprint(f"[red]❌ Submit error: {e}[/red]")
+        maybe_cancel("submit error")
+        sys.exit(2)
+
+
 @main.command()
 @click.option(
     "--user",
@@ -1542,13 +1741,12 @@ def list(ctx: click.Context, user: Optional[str], status: Optional[str], details
                     if "@" in user_id:
                         user_display = user_id.split("@")[0]
 
-                    # Format GPU information
+                    # Format GPU information (MIG-friendly via _format_gpu_display)
                     if gpu_type and gpu_type not in ["unknown", "Unknown"]:
-                        # For CPU nodes (gpu_count = 0), show just the type
                         if gpu_count == 0:
                             gpu_display = gpu_type
                         else:
-                            gpu_display = f"{gpu_count}x {gpu_type}"
+                            gpu_display = _format_gpu_display(gpu_count, gpu_type)
                     else:
                         gpu_display = str(gpu_count)
 
@@ -1844,7 +2042,7 @@ def list(ctx: click.Context, user: Optional[str], status: Optional[str], details
                                         if gpu_count == 0:
                                             gpu_display = gpu_type
                                         else:
-                                            gpu_display = f"{gpu_count}x {gpu_type}"
+                                            gpu_display = _format_gpu_display(gpu_count, gpu_type)
                                     else:
                                         gpu_display = str(gpu_count)
 
@@ -2417,6 +2615,9 @@ def _format_gpu_display(gpu_count, gpu_type):
         "h100-mig-3g": "40GB H100 (MIG)",
         "h100-mig-4g": "40GB H100 (MIG)",
         "h100-mig-7g": "80GB H100 (MIG)",
+        "b200-mig-1g": "23GB B200 (MIG)",
+        "b200-mig-2g": "45GB B200 (MIG)",
+        "b200-mig-3g": "90GB B200 (MIG)",
     }
     if gt_lower in mig_friendly:
         return f"{gpu_count}× {mig_friendly[gt_lower]}"

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

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "gpu-dev"
-version = "0.5.16"
+version = "0.5.18"
 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.16 → gpu_dev-0.5.18}/terraform-gpu-devservers/lambda/reservation_processor/index.py

@@ -308,30 +308,35 @@ def get_target_az_for_reservation(gpu_type, gpus_requested):
                     f"Node {node.metadata.name} in {node_az}: {available_gpus} available GPUs")
 
         if candidate_nodes:
-            # Return the AZ of the first suitable node (Kubernetes scheduler will make the final decision)
+            # Binpacking: pack into the most-loaded node that still fits the request.
+            # Sort by free GPUs ASC so the fullest node comes first; ties broken by node name
+            # so the choice is deterministic across Lambda invocations.
+            candidate_nodes.sort(key=lambda n: (n['available_gpus'], n['node_name']))
             selected_node = candidate_nodes[0]
             target_az = selected_node['az']
+            target_node = selected_node['node_name']
             logger.info(
-                f"Target AZ for {gpu_type} reservation: {target_az} (node: {selected_node['node_name']})")
-            return target_az
+                f"Binpacked target for {gpu_type} {gpus_requested}gpu: "
+                f"node={target_node} az={target_az} free={selected_node['available_gpus']} "
+                f"(candidates considered: {len(candidate_nodes)})")
+            return target_az, target_node
 
         if all_ready_nodes:
-            # No single node has enough GPUs, but nodes exist — return AZ of the node
-            # with the most available GPUs so the disk is created in the right AZ
+            # No single node has enough GPUs — return AZ of the node with the most available GPUs
+            # so disk lands in the right AZ. No node hint (pod will Pending until something frees up).
             best_node = max(all_ready_nodes, key=lambda n: n['available_gpus'])
             target_az = best_node['az']
             logger.info(
                 f"No single node has {gpus_requested} {gpu_type} GPUs, "
                 f"but {len(all_ready_nodes)} nodes exist. Using AZ {target_az} "
                 f"from node {best_node['node_name']} ({best_node['available_gpus']} GPUs available)")
-            return target_az
+            return target_az, None
 
         logger.warning(f"No ready/schedulable {gpu_type} nodes found in cluster")
         return None, None
 
     except Exception as e:
         logger.error(f"Error determining target AZ for {gpu_type}: {str(e)}")
-        # Fallback to primary AZ if detection fails (no node hint — let k8s pick).
         return PRIMARY_AVAILABILITY_ZONE, None
 
 
@@ -1418,6 +1423,11 @@ def coordinate_multinode_reservation(master_reservation_id: str, total_nodes: in
             logger.info(
                 f"Starting parallel processing for {total_nodes} nodes")
 
+            # Deterministic peer pod names by node_index so MULTINODE_RANK aligns with the
+            # position of this pod in MULTINODE_HOSTS across all replicas.
+            nodes_sorted = sorted(nodes, key=lambda n: int(n.get("node_index", 0)))
+            peer_pod_names = [f"gpu-dev-{n['reservation_id'][:8]}" for n in nodes_sorted]
+
             def process_single_node(node_data):
                 """Process a single node - to be run in parallel"""
                 i, node = node_data
@@ -1430,7 +1440,8 @@ def coordinate_multinode_reservation(master_reservation_id: str, total_nodes: in
                         'action': 'process_multinode_individual',
                         'node_index': int(node_index),
                         'total_nodes': int(total_nodes),
-                        'master_reservation_id': str(master_reservation_id)
+                        'master_reservation_id': str(master_reservation_id),
+                        'multinode_peer_pods': peer_pod_names,
                     }
 
                     logger.info(
@@ -1536,6 +1547,12 @@ def process_multinode_individual_node(message_body: dict) -> bool:
 
         node_data = response["Item"]
 
+        # Forward peer pod list from coordinator into request dict so create_pod can
+        # bake MULTINODE_HOSTS / MASTER_ADDR / MULTINODE_RANK env vars into the pod.
+        peer_pods = message_body.get("multinode_peer_pods")
+        if peer_pods:
+            node_data["multinode_peer_pods"] = peer_pods
+
         # Update status to preparing pod
         update_multinode_pod_status(
             reservation_id, "preparing pod", node_index, total_nodes)
@@ -2722,6 +2739,7 @@ def allocate_gpu_resources(reservation_id: str, request: dict[str, Any], trace_d
         persistent_volume_id = None
         device_name = None
         target_az = None  # Initialize target_az for use in connection info update
+        target_node = None  # Initialize target_node (binpacking hostname pin) for create_pod
         is_new_disk = False  # Initialize is_new_disk for all code paths
 
         # If we're using persistent disk, immediately mark this reservation as having a volume
@@ -2749,8 +2767,8 @@ def allocate_gpu_resources(reservation_id: str, request: dict[str, Any], trace_d
                     detailed_status="Setting up persistent disk" + (f" '{disk_name}'" if disk_name else "")
                 )
 
-                # Determine target AZ for this reservation
-                target_az = get_target_az_for_reservation(gpu_type, gpu_count)
+                # Determine target AZ + node for this reservation (binpacking)
+                target_az, target_node = get_target_az_for_reservation(gpu_type, gpu_count)
                 if not target_az:
                     raise ValueError(f"No {gpu_type} nodes found in cluster")
 
@@ -2881,6 +2899,9 @@ def allocate_gpu_resources(reservation_id: str, request: dict[str, Any], trace_d
             dockerfile_base64_data=dockerfile_base64_data,
             dockerimage=dockerimage,
             target_az=target_az,
+            target_node=target_node,
+            multinode_peer_pods=request.get("multinode_peer_pods"),
+            multinode_rank=int(request.get("node_index", 0)) if is_multinode else 0,
             preserve_entrypoint=preserve_entrypoint,
             node_labels=node_labels,
             trace_data=trace_data,
@@ -3421,6 +3442,9 @@ def create_kubernetes_resources(
     recreate_env: bool = False,
     efs_filesystem_id: str = None,
     is_multinode: bool = False,
+    target_node: str = None,
+    multinode_peer_pods: list = None,
+    multinode_rank: int = 0,
     dockerfile_base64_data: str = None,
     dockerimage: str = None,
     target_az: str = None,
@@ -3524,6 +3548,9 @@ def create_kubernetes_resources(
                         dockerfile_base64_data=dockerfile_base64_data,
                         dockerimage=dockerimage,
                         target_az=target_az,
+                        target_node=target_node,
+                        multinode_peer_pods=multinode_peer_pods,
+                        multinode_rank=multinode_rank,
                         preserve_entrypoint=preserve_entrypoint,
                         node_labels=node_labels,
                         trace_data=trace_data,
@@ -3610,6 +3637,9 @@ def create_kubernetes_resources(
                         dockerfile_base64_data=dockerfile_base64_data,
                         dockerimage=dockerimage,
                         target_az=target_az,
+                        target_node=target_node,
+                        multinode_peer_pods=multinode_peer_pods,
+                        multinode_rank=multinode_rank,
                         preserve_entrypoint=preserve_entrypoint,
                         node_labels=node_labels,
                         trace_data=trace_data,
@@ -3712,6 +3742,30 @@ def find_available_node_port(k8s_client) -> int:
         return random.randint(30000, 32767)
 
 
+def _mig_slice_fraction(gpu_type: str) -> float:
+    """For MIG SKUs return slice fraction of a single GPU (1g=1/7, 2g=2/7, ..., 7g=1).
+
+    Slice naming counts GPCs (compute slices). H100 and B200 both have 7 GPCs per GPU
+    in the typical all-balanced profile, so a 1g slice is 1/7 of a GPU regardless of
+    family. Used to size CPU/memory requests proportional to the GPU fraction the pod
+    actually consumes — the older `gpu_count/max_gpus` ratio over-claimed node resources
+    (a 1g slice would claim 1/4 or 1/16 of the host instead of 1/56).
+    """
+    if "mig" not in gpu_type:
+        return 1.0
+    try:
+        slices = int(gpu_type.split("-mig-")[1].rstrip("g"))
+    except (IndexError, ValueError):
+        return 1.0
+    return slices / 7.0
+
+
+# Number of full GPUs on the underlying instance — used to convert the slice fraction
+# into a fraction of the host's CPU/memory. Both p5.48xlarge (H100) and p6-b200.48xlarge
+# (B200) have 8 GPUs, which matches every MIG-capable instance type we currently run.
+_FULL_GPUS_PER_MIG_NODE = 8
+
+
 def get_pod_resource_limits(gpu_count: int, gpu_type: str, is_multinode: bool = False) -> dict:
     """Get resource limits for pod based on GPU type and deployment mode"""
     gpu_count = int(gpu_count)
@@ -3731,13 +3785,19 @@ def get_pod_resource_limits(gpu_count: int, gpu_type: str, is_multinode: bool =
             resource_name = config.get("k8s_resource", "nvidia.com/gpu")
             limits[resource_name] = str(gpu_count)
 
-            gpu_ratio = gpu_count / max_gpus if max_gpus > 0 else 1.0
-
-            # Calculate proportional limits with CPU overprovisioning for burst capacity
-            # Give 1.5x CPU limit to allow burst, capped at node total
-            fractional_cpu = config["cpus"] * gpu_ratio
-            proportional_cpu_limit = min(config["cpus"], int(fractional_cpu * 1.5))
-            proportional_memory_limit = int(config["memory_gb"] * gpu_ratio)
+            if "mig" in gpu_type:
+                # Scale by GPC fraction (slice of one GPU), not slice count over max slices.
+                slice_fraction = _mig_slice_fraction(gpu_type)
+                cpu_per_full_gpu = config["cpus"] / _FULL_GPUS_PER_MIG_NODE
+                mem_per_full_gpu = config["memory_gb"] / _FULL_GPUS_PER_MIG_NODE
+                fractional_cpu = cpu_per_full_gpu * slice_fraction * gpu_count
+                proportional_cpu_limit = max(1, min(config["cpus"], int(fractional_cpu * 1.5)))
+                proportional_memory_limit = max(1, int(mem_per_full_gpu * slice_fraction * gpu_count))
+            else:
+                gpu_ratio = gpu_count / max_gpus if max_gpus > 0 else 1.0
+                fractional_cpu = config["cpus"] * gpu_ratio
+                proportional_cpu_limit = min(config["cpus"], int(fractional_cpu * 1.5))
+                proportional_memory_limit = int(config["memory_gb"] * gpu_ratio)
 
             limits.update({
                 "cpu": str(proportional_cpu_limit),
@@ -3777,13 +3837,16 @@ def get_pod_resource_requests(gpu_count: int, gpu_type: str, is_multinode: bool
         if gpu_count > 0:
             resource_name = config.get("k8s_resource", "nvidia.com/gpu")
             requests[resource_name] = str(gpu_count)
-            gpu_ratio = gpu_count / max_gpus if max_gpus > 0 else 1.0
-
-            # Calculate proportional requests (reserve 10% for system overhead)
-            # This ensures requests don't exceed node allocatable resources
-            # Limits can be higher for burst capacity (Burstable QoS)
-            proportional_cpu_request = int(config["cpus"] * gpu_ratio * 0.9)
-            proportional_memory_request = int(config["memory_gb"] * gpu_ratio * 0.9)
+            if "mig" in gpu_type:
+                slice_fraction = _mig_slice_fraction(gpu_type)
+                cpu_per_full_gpu = config["cpus"] / _FULL_GPUS_PER_MIG_NODE
+                mem_per_full_gpu = config["memory_gb"] / _FULL_GPUS_PER_MIG_NODE
+                proportional_cpu_request = max(1, int(cpu_per_full_gpu * slice_fraction * gpu_count * 0.9))
+                proportional_memory_request = max(1, int(mem_per_full_gpu * slice_fraction * gpu_count * 0.9))
+            else:
+                gpu_ratio = gpu_count / max_gpus if max_gpus > 0 else 1.0
+                proportional_cpu_request = int(config["cpus"] * gpu_ratio * 0.9)
+                proportional_memory_request = int(config["memory_gb"] * gpu_ratio * 0.9)
 
             requests.update({
                 "cpu": str(proportional_cpu_request),
@@ -3886,6 +3949,30 @@ def get_nccl_env_vars(gpu_type: str) -> list:
     return env_vars
 
 
+def _get_multinode_env_vars(peer_pods: list, rank: int) -> list:
+    """Build env vars exposing peer hostnames/rank/master to the pod.
+
+    Hostnames use the per-pod headless service we already create elsewhere, so they
+    resolve to the current pod IP via cluster DNS even if a pod is recreated. We
+    don\'t inject IPs at pod-creation time (they aren\'t known until kube schedules
+    everyone) — the bashrc/zshrc helper resolves and exports MULTINODE_IPS at shell
+    start, and a /usr/local/bin/multinode-ips helper is available for non-interactive
+    callers.
+    """
+    if not peer_pods or len(peer_pods) <= 1:
+        return []
+    namespace = "gpu-dev"
+    hosts = [f"{p}-headless.{namespace}.svc.cluster.local" for p in peer_pods]
+    return [
+        client.V1EnvVar(name="MULTINODE_HOSTS", value=",".join(hosts)),
+        client.V1EnvVar(name="MULTINODE_PEER_PODS", value=",".join(peer_pods)),
+        client.V1EnvVar(name="MULTINODE_RANK", value=str(rank)),
+        client.V1EnvVar(name="MULTINODE_SIZE", value=str(len(peer_pods))),
+        client.V1EnvVar(name="MASTER_ADDR", value=hosts[0]),
+        client.V1EnvVar(name="MASTER_PORT", value="29500"),
+    ]
+
+
 def create_pod(
     k8s_client,
     pod_name: str,
@@ -3902,6 +3989,9 @@ def create_pod(
     dockerfile_base64_data: str = None,
     dockerimage: str = None,
     target_az: str = None,
+    target_node: str = None,
+    multinode_peer_pods: list = None,
+    multinode_rank: int = 0,
     preserve_entrypoint: bool = False,
     node_labels: dict = None,
     trace_data: dict = None,
@@ -4386,6 +4476,16 @@ EOF_PROFILE
 # User identification
 export GPU_DEV_USER_ID="{user_id or '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.
+# Skipped (empty exports) for single-node reservations where MULTINODE_* aren't set.
+export MULTINODE_HOSTS="$MULTINODE_HOSTS"
+export MULTINODE_PEER_PODS="$MULTINODE_PEER_PODS"
+export MULTINODE_RANK="$MULTINODE_RANK"
+export MULTINODE_SIZE="$MULTINODE_SIZE"
+export MASTER_ADDR="$MASTER_ADDR"
+export MASTER_PORT="$MASTER_PORT"
+
 # Function to check for GPU reservation expiry warnings and startup script status
 check_warnings() {{
     # Check for startup script still running
@@ -4404,6 +4504,22 @@ check_warnings() {{
 
 # Run warning check before every command prompt
 PROMPT_COMMAND="check_warnings; \$PROMPT_COMMAND"
+
+# Multinode peer IP resolution: MULTINODE_HOSTS is baked at pod creation, but per-pod
+# IPs are only known once kube schedules them. Resolve at shell start so users can do
+# torchrun --master_addr=\$MASTER_ADDR or mpirun -H "\$MULTINODE_IPS" without extra steps.
+if [ -n "\$MULTINODE_HOSTS" ]; then
+    _MULTINODE_IPS=""
+    for _h in \$(echo "\$MULTINODE_HOSTS" | tr ',' ' '); do
+        _ip=\$(getent hosts "\$_h" 2>/dev/null | awk '{{print \$1}}' | head -1)
+        if [ -n "\$_ip" ]; then
+            _MULTINODE_IPS="\${{_MULTINODE_IPS:+\$_MULTINODE_IPS,}}\$_ip"
+        fi
+    done
+    export MULTINODE_IPS="\$_MULTINODE_IPS"
+    [ -n "\$MULTINODE_IPS" ] && export MASTER_IP=\$(echo "\$MULTINODE_IPS" | cut -d, -f1)
+    unset _MULTINODE_IPS _h _ip
+fi
 EOF_BASHRC_EXT
 
                         cat > /home/dev/.zshrc_ext << EOF_ZSHRC_EXT
@@ -4414,6 +4530,15 @@ EOF_BASHRC_EXT
 # User identification
 export GPU_DEV_USER_ID="{user_id or '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.
+export MULTINODE_HOSTS="$MULTINODE_HOSTS"
+export MULTINODE_PEER_PODS="$MULTINODE_PEER_PODS"
+export MULTINODE_RANK="$MULTINODE_RANK"
+export MULTINODE_SIZE="$MULTINODE_SIZE"
+export MASTER_ADDR="$MASTER_ADDR"
+export MASTER_PORT="$MASTER_PORT"
+
 # Function to check for GPU reservation expiry warnings and startup script status
 check_warnings() {{
     # Check for startup script still running
@@ -4433,6 +4558,20 @@ check_warnings() {{
 
 # Run warning check before every command prompt (zsh hook)
 precmd() {{ check_warnings }}
+
+# Multinode peer IP resolution (see .bashrc_ext for rationale)
+if [[ -n "\$MULTINODE_HOSTS" ]]; then
+    _MULTINODE_IPS=""
+    for _h in \${{(s:,:)MULTINODE_HOSTS}}; do
+        _ip=\$(getent hosts "\$_h" 2>/dev/null | awk '{{print \$1}}' | head -1)
+        if [[ -n "\$_ip" ]]; then
+            _MULTINODE_IPS="\${{_MULTINODE_IPS:+\$_MULTINODE_IPS,}}\$_ip"
+        fi
+    done
+    export MULTINODE_IPS="\$_MULTINODE_IPS"
+    [[ -n "\$MULTINODE_IPS" ]] && export MASTER_IP="\${{MULTINODE_IPS%%,*}}"
+    unset _MULTINODE_IPS _h _ip
+fi
 EOF_ZSHRC_EXT
 
                         chown 1081:1081 /home/dev/.bashrc_ext /home/dev/.zshrc_ext
@@ -5163,7 +5302,7 @@ EOF
                         client.V1EnvVar(
                             name="NVIDIA_DRIVER_CAPABILITIES", value="compute,utility"
                         )
-                    ] + get_nccl_env_vars(gpu_type) + get_cpu_thread_env_vars(gpu_count, gpu_type),
+                    ] + get_nccl_env_vars(gpu_type) + get_cpu_thread_env_vars(gpu_count, gpu_type) + _get_multinode_env_vars(multinode_peer_pods, multinode_rank),
                     resources=client.V1ResourceRequirements(
                         limits=get_pod_resource_limits(
                             gpu_count, gpu_type, is_multinode),
@@ -5309,7 +5448,12 @@ EOF
             ] if _pod_uses_efa(gpu_count, gpu_type, is_multinode) else []),
             node_selector={
                 "GpuType": get_node_gpu_type(gpu_type),
-                **({} if target_az is None else {"topology.kubernetes.io/zone": target_az})
+                **({} if target_az is None else {"topology.kubernetes.io/zone": target_az}),
+                # Hard-pin to the binpacked node when Lambda picked one. Lambda runs
+                # serialized (reserved_concurrent_executions=1), so allocations seen by the
+                # next invocation include this pod. If the node is unavailable, the pod
+                # stays Pending and surfaces the error rather than spreading.
+                **({} if target_node is None else {"kubernetes.io/hostname": target_node}),
             },
             # Node affinity for profiling-dedicated preference
             # If user requests nsight=true, prefer profiling-dedicated nodes
@@ -6303,6 +6447,9 @@ def get_instance_type_and_gpu_info(k8s_client, pod_name: str) -> tuple[str, str]
             "nvidia.com/mig-3g.40gb": "h100-mig-3g",
             "nvidia.com/mig-4g.40gb": "h100-mig-4g",
             "nvidia.com/mig-7g.80gb": "h100-mig-7g",
+            "nvidia.com/mig-1g.23gb": "b200-mig-1g",
+            "nvidia.com/mig-2g.45gb": "b200-mig-2g",
+            "nvidia.com/mig-3g.90gb": "b200-mig-3g",
         }
         if pod.spec.containers:
             for c in pod.spec.containers:

{gpu_dev-0.5.16 → gpu_dev-0.5.18}/terraform-gpu-devservers/lambda/shared/k8s_client.py

@@ -31,9 +31,14 @@ def get_bearer_token() -> str:
     """
     Create a k8s-aws-v1 bearer token by presigning STS:GetCallerIdentity.
     IMPORTANT: base64url-encode the FULL presigned URL, then strip padding.
+
+    expires_in must match _EFFECTIVE_TOKEN_TTL: previously this was 60s while the cache
+    held the token for 14 min, so warm Lambda containers handed EKS expired URLs and got
+    401s for ~13 min until the next refresh. 900s is the typical EKS get-token default
+    and the max for IAM-role-derived presigned URLs.
     """
     logger.info("Starting bearer token generation")
-    STS_TOKEN_EXPIRES_IN = 60
+    STS_TOKEN_EXPIRES_IN = 900
     session = boto3.session.Session(region_name=REGION)
     logger.info(f"Created boto3 session for region {REGION}")
 

{gpu_dev-0.5.16 → gpu_dev-0.5.18}/terraform-gpu-devservers/lambda.tf

@@ -180,8 +180,8 @@ resource "aws_lambda_function" "reservation_processor" {
       HOSTED_ZONE_ID                     = local.effective_domain_name != "" ? local.hosted_zone_id : ""
       SSH_DOMAIN_MAPPINGS_TABLE          = local.effective_domain_name != "" ? aws_dynamodb_table.ssh_domain_mappings.name : ""
       SSL_CERTIFICATE_ARN                = local.effective_domain_name != "" ? aws_acm_certificate.wildcard[0].arn : ""
-      LAMBDA_VERSION                     = "0.5.13"
-      MIN_CLI_VERSION                    = "0.5.9"
+      LAMBDA_VERSION                     = "0.5.22"
+      MIN_CLI_VERSION                    = "0.5.16"
       DISK_CONTENTS_BUCKET               = aws_s3_bucket.disk_contents.bucket
       OPERATIONS_TABLE                   = aws_dynamodb_table.operations.name
     }, local.alb_env_vars)