gpu-dev 0.5.15__tar.gz → 0.5.17__tar.gz

{gpu_dev-0.5.15 → gpu_dev-0.5.17}/PKG-INFO

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

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

{gpu_dev-0.5.15 → gpu_dev-0.5.17}/cli-tools/gpu-dev-cli/gpu_dev_cli/auth.py

@@ -19,7 +19,7 @@ _SSH_CACHE_PATH = Path(os.path.expanduser("~/.config/gpu-dev/ssh-validation-cach
 # Cache for authenticate_user. STS GetCallerIdentity is stable per AWS profile and slow under SSO
 # (~500ms-1.5s). Cache for 24h keyed by AWS_PROFILE; if creds rotate the user_id rarely changes,
 # and the next AWS call (DDB/SQS) will surface a credential error if it does.
-_AUTH_CACHE_TTL_SECONDS = 24 * 60 * 60
+_AUTH_CACHE_TTL_SECONDS = 60 * 60
 _AUTH_CACHE_PATH = Path(os.path.expanduser("~/.config/gpu-dev/auth-cache.json"))
 
 
@@ -64,6 +64,22 @@ def _save_auth_cache(github_user: str, result: Dict[str, Any]) -> None:
         pass
 
 
+def clear_auth_cache() -> None:
+    """Drop the cached auth entry for the current AWS profile. Call this after a credential
+    error to force the next authenticate_user() to re-hit STS."""
+    try:
+        if not _AUTH_CACHE_PATH.exists():
+            return
+        with open(_AUTH_CACHE_PATH) as f:
+            data = json.load(f)
+        if _auth_cache_key() in data:
+            del data[_auth_cache_key()]
+            with open(_AUTH_CACHE_PATH, "w") as f:
+                json.dump(data, f)
+    except Exception:
+        pass
+
+
 def _load_ssh_cache(github_user: str) -> Optional[Dict[str, Any]]:
     """Return cached validation if it's fresh and matches the configured github_user, else None."""
     try:
@@ -125,6 +141,7 @@ def authenticate_user(config: Config) -> Dict[str, Any]:
         _save_auth_cache(github_user, result)
         return result
     except Exception as e:
+        clear_auth_cache()
         raise RuntimeError(f"AWS authentication failed: {e}")
 
 

{gpu_dev-0.5.15 → gpu_dev-0.5.17}/cli-tools/gpu-dev-cli/gpu_dev_cli/cli.py

@@ -1542,13 +1542,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 +1843,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 +2416,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.15 → gpu_dev-0.5.17}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "gpu-dev"
-version = "0.5.15"
+version = "0.5.17"
 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.15 → gpu_dev-0.5.17}/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
 
 
@@ -2722,6 +2727,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 +2755,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 +2887,7 @@ 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,
             preserve_entrypoint=preserve_entrypoint,
             node_labels=node_labels,
             trace_data=trace_data,
@@ -3421,6 +3428,7 @@ def create_kubernetes_resources(
     recreate_env: bool = False,
     efs_filesystem_id: str = None,
     is_multinode: bool = False,
+    target_node: str = None,
     dockerfile_base64_data: str = None,
     dockerimage: str = None,
     target_az: str = None,
@@ -3524,6 +3532,7 @@ def create_kubernetes_resources(
                         dockerfile_base64_data=dockerfile_base64_data,
                         dockerimage=dockerimage,
                         target_az=target_az,
+                        target_node=target_node,
                         preserve_entrypoint=preserve_entrypoint,
                         node_labels=node_labels,
                         trace_data=trace_data,
@@ -3610,6 +3619,7 @@ def create_kubernetes_resources(
                         dockerfile_base64_data=dockerfile_base64_data,
                         dockerimage=dockerimage,
                         target_az=target_az,
+                        target_node=target_node,
                         preserve_entrypoint=preserve_entrypoint,
                         node_labels=node_labels,
                         trace_data=trace_data,
@@ -3902,6 +3912,7 @@ def create_pod(
     dockerfile_base64_data: str = None,
     dockerimage: str = None,
     target_az: str = None,
+    target_node: str = None,
     preserve_entrypoint: bool = False,
     node_labels: dict = None,
     trace_data: dict = None,
@@ -5309,7 +5320,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

{gpu_dev-0.5.15 → gpu_dev-0.5.17}/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.17"
+      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)