gpu-dev 0.3.8__tar.gz → 0.3.9__tar.gz

{gpu_dev-0.3.8/cli-tools/gpu-dev-cli/gpu_dev.egg-info → gpu_dev-0.3.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpu-dev
-Version: 0.3.8
+Version: 0.3.9
 Summary: CLI tool for PyTorch GPU developer server reservations
 Author: PyTorch Team
 Requires-Python: >=3.10
@@ -12,7 +12,7 @@ Requires-Dist: pydantic>=2.5.0
 Requires-Dist: rich>=13.7.0
 Requires-Dist: pyyaml>=6.0.1
 Requires-Dist: questionary>=2.1.1
-Requires-Dist: websockets<13.0,>=12.0
+Requires-Dist: websockets>=12.0
 Requires-Dist: certifi>=2023.7.22
 Requires-Dist: mcp>=1.0.0
 

{gpu_dev-0.3.8 → gpu_dev-0.3.9/cli-tools/gpu-dev-cli/gpu_dev.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpu-dev
-Version: 0.3.8
+Version: 0.3.9
 Summary: CLI tool for PyTorch GPU developer server reservations
 Author: PyTorch Team
 Requires-Python: >=3.10
@@ -12,7 +12,7 @@ Requires-Dist: pydantic>=2.5.0
 Requires-Dist: rich>=13.7.0
 Requires-Dist: pyyaml>=6.0.1
 Requires-Dist: questionary>=2.1.1
-Requires-Dist: websockets<13.0,>=12.0
+Requires-Dist: websockets>=12.0
 Requires-Dist: certifi>=2023.7.22
 Requires-Dist: mcp>=1.0.0
 

{gpu_dev-0.3.8 → gpu_dev-0.3.9}/cli-tools/gpu-dev-cli/gpu_dev.egg-info/requires.txt

@@ -5,6 +5,6 @@ pydantic>=2.5.0
 rich>=13.7.0
 pyyaml>=6.0.1
 questionary>=2.1.1
-websockets<13.0,>=12.0
+websockets>=12.0
 certifi>=2023.7.22
 mcp>=1.0.0

{gpu_dev-0.3.8 → gpu_dev-0.3.9}/cli-tools/gpu-dev-cli/gpu_dev_cli/ssh_proxy.py

@@ -4,6 +4,7 @@ SSH ProxyCommand helper for tunneling SSH through WebSocket
 Used by ssh with: ssh -o ProxyCommand='gpu-dev-ssh-proxy %h %p' user@host
 """
 
+import os
 import sys
 import asyncio
 import websockets
@@ -18,6 +19,11 @@ async def tunnel_ssh(target_host: str, target_port: int):
         target_host: Target SSH hostname
         target_port: Target SSH port
     """
+    # Bypass corporate/local HTTP proxies for devservers.io - we connect
+    # directly to our ALB, and proxies can cause WebSocket handshake timeouts
+    for var in ("HTTPS_PROXY", "https_proxy", "HTTP_PROXY", "http_proxy", "ALL_PROXY", "all_proxy"):
+        os.environ.pop(var, None)
+
     # Determine proxy URL based on target host
     if ".test.devservers.io" in target_host:
         proxy_host = "ssh.test.devservers.io"
@@ -31,8 +37,8 @@ async def tunnel_ssh(target_host: str, target_port: int):
     ws_url = f"wss://{proxy_host}/tunnel/{target_host}"
 
     try:
-        # Connect to WebSocket proxy
-        async with websockets.connect(ws_url) as websocket:
+        # Connect to WebSocket proxy (20s timeout, generous for cold DNS/TLS)
+        async with websockets.connect(ws_url, open_timeout=20) as websocket:
             # Set up stdin/stdout for SSH
             loop = asyncio.get_event_loop()
             reader = asyncio.StreamReader()

{gpu_dev-0.3.8 → gpu_dev-0.3.9}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "gpu-dev"
-version = "0.3.8"
+version = "0.3.9"
 description = "CLI tool for PyTorch GPU developer server reservations"
 authors = [{name = "PyTorch Team"}]
 readme = "cli-tools/gpu-dev-cli/README.md"
@@ -17,7 +17,7 @@ dependencies = [
     "rich>=13.7.0",
     "pyyaml>=6.0.1",
     "questionary>=2.1.1",
-    "websockets>=12.0,<13.0",
+    "websockets>=12.0",
     "certifi>=2023.7.22",
     "mcp>=1.0.0",
 ]

{gpu_dev-0.3.8 → gpu_dev-0.3.9}/terraform-gpu-devservers/eks.tf

@@ -344,16 +344,31 @@ resource "aws_launch_template" "gpu_dev_launch_template" {
     }
   }
 
-  # Network interface (EFA enabled for supported instance types only)
+  # Primary network interface (card 0) - EFA+ENA for GPU instances, regular for CPU/T4-small
   network_interfaces {
+    network_card_index          = 0
+    device_index                = 0
     associate_public_ip_address = true
     security_groups             = [aws_security_group.gpu_dev_sg.id]
     subnet_id                   = each.value.gpu_config.use_placement_group ? null : (local.gpu_subnet_assignments[terraform.workspace][each.value.gpu_type] == "secondary" ? aws_subnet.gpu_dev_subnet_secondary.id : aws_subnet.gpu_dev_subnet.id)
-    # EFA is not supported on g4dn.2xlarge (t4-small) and CPU instances
-    interface_type              = (each.value.gpu_type == "t4-small" || each.value.gpu_config.gpus_per_instance == 0) ? "interface" : "efa"
+    interface_type              = try(each.value.gpu_config.efa_network_cards, 0) > 0 ? "efa" : "interface"
     delete_on_termination       = true
   }
 
+  # Additional EFA-only interfaces (cards 1-N) for multi-card instances (p5, p5e, p6, p4d)
+  dynamic "network_interfaces" {
+    for_each = try(each.value.gpu_config.efa_network_cards, 0) > 1 ? range(1, each.value.gpu_config.efa_network_cards) : []
+    content {
+      network_card_index          = network_interfaces.value
+      device_index                = 1
+      associate_public_ip_address = false
+      security_groups             = [aws_security_group.gpu_dev_sg.id]
+      subnet_id                   = each.value.gpu_config.use_placement_group ? null : (local.gpu_subnet_assignments[terraform.workspace][each.value.gpu_type] == "secondary" ? aws_subnet.gpu_dev_subnet_secondary.id : aws_subnet.gpu_dev_subnet.id)
+      interface_type              = "efa-only"
+      delete_on_termination       = true
+    }
+  }
+
   # Conditionally add instance_market_options for capacity block instances (only when capacity reservation exists)
   dynamic "instance_market_options" {
     for_each = each.value.capacity_reservation_id != null ? [1] : []
@@ -381,6 +396,7 @@ resource "aws_launch_template" "gpu_dev_launch_template" {
     region              = local.current_config.aws_region
     gpu_type            = local.gpu_type_kubernetes_labels[each.value.gpu_type]
     profiling_dedicated = try(each.value.gpu_config.profiling_dedicated, false)
+    container_image     = local.latest_image_uri
   }))
 
   tag_specifications {

{gpu_dev-0.3.8 → gpu_dev-0.3.9}/terraform-gpu-devservers/lambda/availability_updater/index.py

@@ -176,8 +176,8 @@ def update_gpu_availability(gpu_type: str, k8s_client=None) -> None:
         max_reservable = 0  # Maximum GPUs reservable (considering multinode for high-end GPUs)
         if k8s_client is not None and not is_cpu_type:
             try:
-                from kubernetes import client
-                v1 = client.CoreV1Api(k8s_client)
+                from kubernetes import client as k8s_client_lib
+                v1 = k8s_client_lib.CoreV1Api(k8s_client)
                 nodes = v1.list_node(label_selector=f"GpuType={gpu_type}")
 
                 single_node_max = 0  # Max available on any single node
@@ -216,8 +216,9 @@ def update_gpu_availability(gpu_type: str, k8s_client=None) -> None:
                 logger.info(f"Found {full_nodes_available} full nodes available for {gpu_type}, max reservable: {max_reservable} (single node max: {single_node_max})")
             except Exception as e:
                 logger.warning(f"Could not calculate full nodes available for {gpu_type}: {str(e)}")
-                full_nodes_available = 0
-                max_reservable = 0
+                # Fallback: use available_gpus so max_reservable isn't misleadingly 0
+                full_nodes_available = available_gpus // gpus_per_instance if gpus_per_instance > 0 else 0
+                max_reservable = available_gpus
         elif is_cpu_type:
             # For CPU nodes, each node supports 1 reservation
             full_nodes_available = available_gpus  # Each "GPU" represents one CPU node slot

{gpu_dev-0.3.8 → gpu_dev-0.3.9}/terraform-gpu-devservers/lambda/reservation_processor/index.py

@@ -63,19 +63,19 @@ OPERATIONS_TABLE = os.environ.get("OPERATIONS_TABLE", "pytorch-gpu-dev-operation
 
 # GPU Configuration - single source of truth for all GPU type mappings
 GPU_CONFIG = {
-    "t4": {"instance_type": "g4dn.12xlarge", "max_gpus": 4, "cpus": 48, "memory_gb": 192},
-    "l4": {"instance_type": "g6.12xlarge", "max_gpus": 4, "cpus": 48, "memory_gb": 192},
-    "a10g": {"instance_type": "g5.12xlarge", "max_gpus": 4, "cpus": 48, "memory_gb": 192},
-    "t4-small": {"instance_type": "g4dn.2xlarge", "max_gpus": 1, "cpus": 8, "memory_gb": 32},
-    "g5g": {"instance_type": "g5g.2xlarge", "max_gpus": 2, "cpus": 8, "memory_gb": 32},
-    "a100": {"instance_type": "p4d.24xlarge", "max_gpus": 8, "cpus": 96, "memory_gb": 1152},
-    "h100": {"instance_type": "p5.48xlarge", "max_gpus": 8, "cpus": 192, "memory_gb": 2048},
-    "h200": {"instance_type": "p5e.48xlarge", "max_gpus": 8, "cpus": 192, "memory_gb": 2048},
-    "b200": {"instance_type": "p6-b200.48xlarge", "max_gpus": 8, "cpus": 192, "memory_gb": 2048},
-    "cpu-arm": {"instance_type": "c7g.8xlarge", "max_gpus": 0, "cpus": 32, "memory_gb": 64},
-    "cpu-x86": {"instance_type": "c7i.8xlarge", "max_gpus": 0, "cpus": 32, "memory_gb": 64},
+    "t4": {"instance_type": "g4dn.12xlarge", "max_gpus": 4, "cpus": 48, "memory_gb": 192, "efa_count": 0},
+    "l4": {"instance_type": "g6.12xlarge", "max_gpus": 4, "cpus": 48, "memory_gb": 192, "efa_count": 1},
+    "a10g": {"instance_type": "g5.12xlarge", "max_gpus": 4, "cpus": 48, "memory_gb": 192, "efa_count": 1},
+    "t4-small": {"instance_type": "g4dn.2xlarge", "max_gpus": 1, "cpus": 8, "memory_gb": 32, "efa_count": 0},
+    "g5g": {"instance_type": "g5g.2xlarge", "max_gpus": 2, "cpus": 8, "memory_gb": 32, "efa_count": 0},
+    "a100": {"instance_type": "p4d.24xlarge", "max_gpus": 8, "cpus": 96, "memory_gb": 1152, "efa_count": 1},
+    "h100": {"instance_type": "p5.48xlarge", "max_gpus": 8, "cpus": 192, "memory_gb": 2048, "efa_count": 1},
+    "h200": {"instance_type": "p5e.48xlarge", "max_gpus": 8, "cpus": 192, "memory_gb": 2048, "efa_count": 1},
+    "b200": {"instance_type": "p6-b200.48xlarge", "max_gpus": 8, "cpus": 192, "memory_gb": 2048, "efa_count": 1},
+    "cpu-arm": {"instance_type": "c7g.8xlarge", "max_gpus": 0, "cpus": 32, "memory_gb": 64, "efa_count": 0},
+    "cpu-x86": {"instance_type": "c7i.8xlarge", "max_gpus": 0, "cpus": 32, "memory_gb": 64, "efa_count": 0},
 }
-GPU_CONFIG_DEFAULT = {"instance_type": "g4dn.12xlarge", "max_gpus": 4, "cpus": 48, "memory_gb": 192}
+GPU_CONFIG_DEFAULT = {"instance_type": "g4dn.12xlarge", "max_gpus": 4, "cpus": 48, "memory_gb": 192, "efa_count": 0}
 
 
 def retry_with_backoff(func, *args, max_retries=5, initial_delay=1, max_delay=32, **kwargs):
@@ -748,48 +748,36 @@ def restore_ebs_from_existing_snapshot(snapshot_id, target_az, user_id):
         raise
 
 
+_efs_cache = {}  # Module-level cache: user_id -> efs_id (shared across threads in same invocation)
+
+
 def create_or_find_user_efs(user_id: str) -> str:
     """Create or find existing EFS filesystem for user shared storage"""
+    if user_id in _efs_cache:
+        cached_id = _efs_cache[user_id]
+        logger.info(f"Using cached EFS {cached_id} for user {user_id}")
+        ensure_efs_mount_target(cached_id)
+        return cached_id
+
     try:
         logger.info(f"Looking for existing EFS filesystem for user {user_id}")
 
-        # Check for existing EFS with user tag
+        # Tags are included inline in describe_file_systems response - no
+        # need for separate describe_tags calls (which get throttled heavily).
+        matching_efs = []
         response = efs_client.describe_file_systems()
-
-        throttle_failures = 0
-        total_filesystems = len(response.get("FileSystems", []))
-        matching_efs = []  # Collect all matching EFS, sorted by creation time
-
-        for fs in response.get("FileSystems", []):
-            fs_id = fs["FileSystemId"]
-
-            # Get tags for this filesystem
-            try:
-                tags_response = retry_with_backoff(efs_client.describe_tags, FileSystemId=fs_id)
-                tags = {tag["Key"]: tag["Value"]
-                        for tag in tags_response.get("Tags", [])}
-
+        while True:
+            for fs in response.get("FileSystems", []):
+                tags = {tag["Key"]: tag["Value"] for tag in fs.get("Tags", [])}
                 if tags.get("gpu-dev-user") == user_id:
                     logger.info(
-                        f"Found existing EFS {fs_id} for user {user_id} (created {fs.get('CreationTime')})")
+                        f"Found existing EFS {fs['FileSystemId']} for user {user_id} (created {fs.get('CreationTime')})")
                     matching_efs.append(fs)
+            if "NextMarker" not in response:
+                break
+            response = efs_client.describe_file_systems(Marker=response["NextMarker"])
 
-            except Exception as tag_error:
-                error_str = str(tag_error)
-                # Track throttling failures separately
-                if "Throttling" in error_str or "RequestLimitExceeded" in error_str or "TooManyRequests" in error_str:
-                    throttle_failures += 1
-                    logger.warning(
-                        f"EFS DescribeTags throttled for {fs_id} ({throttle_failures}/{total_filesystems}): {tag_error}")
-                else:
-                    logger.warning(
-                        f"Could not get tags for EFS {fs_id}: {tag_error}")
-                continue
-
-        # If we found matching EFS, return the NEWEST one (by CreationTime)
-        # Do this BEFORE checking throttling - if we found EFS, throttling doesn't matter
         if matching_efs:
-            # Sort by CreationTime descending (newest first)
             matching_efs.sort(key=lambda x: x.get('CreationTime'), reverse=True)
             newest_efs = matching_efs[0]
             fs_id = newest_efs["FileSystemId"]
@@ -803,24 +791,10 @@ def create_or_find_user_efs(user_id: str) -> str:
             else:
                 logger.info(f"Using EFS {fs_id} for user {user_id}")
 
-            # Log throttling as warning but proceed anyway (we have valid EFS)
-            if throttle_failures > 0:
-                logger.warning(
-                    f"Had {throttle_failures}/{total_filesystems} throttling errors during scan, "
-                    f"but found valid EFS {fs_id} - proceeding"
-                )
-
-            # Ensure mount target exists
             ensure_efs_mount_target(fs_id)
+            _efs_cache[user_id] = fs_id
             return fs_id
 
-        # No matching EFS found - check if throttling prevented complete scan
-        if throttle_failures > 0:
-            raise Exception(
-                f"EFS DescribeTags API throttled ({throttle_failures}/{total_filesystems} filesystems). "
-                f"Cannot safely create new EFS - retry later to avoid duplicates."
-            )
-
         # Create new EFS filesystem
         logger.info(f"Creating new EFS filesystem for user {user_id}")
 
@@ -886,6 +860,7 @@ def create_or_find_user_efs(user_id: str) -> str:
             # Don't fail EFS creation for this
 
         logger.info(f"Created new EFS filesystem {fs_id} for user {user_id}")
+        _efs_cache[user_id] = fs_id
         return fs_id
 
     except Exception as e:
@@ -3519,8 +3494,9 @@ def get_pod_resource_limits(gpu_count: int, gpu_type: str, is_multinode: bool =
     )
 
     if use_efa:
-        limits["vpc.amazonaws.com/efa"] = "1"
-        logger.info(f"Using EFA for multinode full-node deployment: {gpu_count}/{max_gpus} GPUs")
+        efa_count = config.get("efa_count", 1)
+        limits["vpc.amazonaws.com/efa"] = str(efa_count)
+        logger.info(f"Using EFA ({efa_count} interfaces) for multinode full-node deployment: {gpu_count}/{max_gpus} GPUs")
     else:
         logger.info(f"Skipping EFA: multinode={is_multinode}, gpu_count={gpu_count}/{max_gpus}, gpu_type={gpu_type}")
 
@@ -3560,7 +3536,8 @@ def get_pod_resource_requests(gpu_count: int, gpu_type: str, is_multinode: bool
         gpu_count == max_gpus
     )
     if use_efa:
-        requests["vpc.amazonaws.com/efa"] = "1"
+        efa_count = config.get("efa_count", 1)
+        requests["vpc.amazonaws.com/efa"] = str(efa_count)
 
     return requests
 
@@ -3709,8 +3686,8 @@ def create_pod(
             init_containers=[
                 client.V1Container(
                     name="ssh-setup",
-                    image="alpine:latest",
-                    image_pull_policy="Always",  # Fail fast if image doesn't exist
+                    image="alpine:3.21",
+                    image_pull_policy="IfNotPresent",
                     command=["/bin/sh"],
                     args=[
                         "-c",
@@ -3790,7 +3767,7 @@ def create_pod(
                 client.V1Container(
                     name="gpu-dev",
                     image=container_image,
-                    image_pull_policy="Always",  # Always pull to check if image exists, fail fast if not
+                    image_pull_policy="IfNotPresent",
                     **({
                         "command": ["/bin/bash"],
                         "args": [

{gpu_dev-0.3.8 → gpu_dev-0.3.9}/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.3.7"
-      MIN_CLI_VERSION                    = "0.3.7"
+      LAMBDA_VERSION                     = "0.3.9"
+      MIN_CLI_VERSION                    = "0.3.9"
       DISK_CONTENTS_BUCKET               = aws_s3_bucket.disk_contents.bucket
       OPERATIONS_TABLE                   = aws_dynamodb_table.operations.name
     }, local.alb_env_vars)

{gpu_dev-0.3.8 → gpu_dev-0.3.9}/terraform-gpu-devservers/templates/al2023-user-data.sh

@@ -117,4 +117,10 @@ net.core.wmem_max=262144000
 EOF
 sysctl --system
 
+# Pre-pull GPU dev container image and refresh every 30 minutes
+# ECR credentials are handled by kubelet's credential provider
+ECR_IMAGE="${container_image}"
+crictl pull "$ECR_IMAGE" || echo "Initial image pre-pull failed (node may not be ready yet)"
+echo "*/30 * * * * ECR_LOGIN=\$(aws ecr get-login-password --region ${region}) && echo \$ECR_LOGIN | crictl pull --creds AWS:\$ECR_LOGIN $ECR_IMAGE 2>&1 | logger -t gpu-dev-image-pull" | crontab -
+
 echo "Amazon Linux 2023 EKS GPU node setup completed"