gpu-dev 0.5.7__tar.gz → 0.5.9__tar.gz

{gpu_dev-0.5.7 → gpu_dev-0.5.9}/PKG-INFO

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

{gpu_dev-0.5.7 → gpu_dev-0.5.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.5.7
+Version: 0.5.9
 Summary: CLI tool for PyTorch GPU developer server reservations
 Author: PyTorch Team
 Requires-Python: >=3.10

{gpu_dev-0.5.7 → gpu_dev-0.5.9}/cli-tools/gpu-dev-cli/gpu_dev_cli/auth.py

@@ -1,11 +1,53 @@
 """Minimal AWS-only authentication for GPU Dev CLI"""
 
+import json
+import os
 import subprocess
 import re
-from typing import Dict, Any
+import time
+from pathlib import Path
+from typing import Dict, Any, Optional
 from .config import Config
 from rich.spinner import Spinner
 
+# SSH validation result is cached locally for 24h. New keys pushed to GitHub still take effect
+# at reservation time (pods fetch live keys via init container) — caching only skips the
+# pre-flight "are you who you say you are" check.
+_SSH_CACHE_TTL_SECONDS = 24 * 60 * 60
+_SSH_CACHE_PATH = Path(os.path.expanduser("~/.config/gpu-dev/ssh-validation-cache.json"))
+
+
+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:
+        if not _SSH_CACHE_PATH.exists():
+            return None
+        with open(_SSH_CACHE_PATH) as f:
+            data = json.load(f)
+        if data.get("configured_user") != github_user:
+            return None
+        if time.time() - float(data.get("ts", 0)) > _SSH_CACHE_TTL_SECONDS:
+            return None
+        return data.get("result")
+    except Exception:
+        return None
+
+
+def _save_ssh_cache(github_user: str, result: Dict[str, Any]) -> None:
+    """Persist a successful validation result. Failures are not cached (so they can recover)."""
+    if not result.get("valid"):
+        return
+    try:
+        _SSH_CACHE_PATH.parent.mkdir(parents=True, exist_ok=True)
+        with open(_SSH_CACHE_PATH, "w") as f:
+            json.dump({
+                "configured_user": github_user,
+                "ts": int(time.time()),
+                "result": result,
+            }, f)
+    except Exception:
+        pass
+
 
 def authenticate_user(config: Config) -> Dict[str, Any]:
     """Authenticate using AWS credentials - if you can call AWS, you're authorized"""
@@ -59,6 +101,13 @@ def validate_ssh_key_matches_github_user(config: Config, live=None) -> Dict[str,
                 "error": "GitHub username not configured. Run: gpu-dev config set github_user <username>",
             }
 
+        # Cache short-circuit — skip the SSH handshake (~1-3s) if we recently validated this user.
+        # Cache TTL is 24h. New keys pushed to GitHub still take effect at reservation time
+        # (pods fetch live keys via init container), so caching the pre-flight check is safe.
+        cached = _load_ssh_cache(github_user)
+        if cached is not None:
+            return cached
+
         # Run ssh git@github.com with interactive host verification support
         ssh_output = None
 
@@ -139,7 +188,7 @@ def validate_ssh_key_matches_github_user(config: Config, live=None) -> Dict[str,
         # Compare usernames (case-insensitive)
         is_valid = ssh_detected_user.lower() == github_user.lower()
 
-        return {
+        result = {
             "valid": is_valid,
             "configured_user": github_user,
             "ssh_user": ssh_detected_user,
@@ -147,6 +196,8 @@ def validate_ssh_key_matches_github_user(config: Config, live=None) -> Dict[str,
             if is_valid
             else f"SSH key belongs to '{ssh_detected_user}' but configured user is '{github_user}'",
         }
+        _save_ssh_cache(github_user, result)
+        return result
 
     except Exception as e:
         return {

{gpu_dev-0.5.7 → gpu_dev-0.5.9}/cli-tools/gpu-dev-cli/gpu_dev_cli/cli.py

@@ -681,29 +681,55 @@ def reserve(
             rprint(
                 "[dim]Use --no-interactive flag to disable interactive mode[/dim]\n")
 
-            # Setup config early for availability check
+            # Run auth + SSH validation + availability fetch in parallel — they're independent
+            # and total wall-clock time drops from sum to max(each).
+            from concurrent.futures import ThreadPoolExecutor
+            config = load_config()
+
             with Live(
-                Spinner("dots", text="📡 Loading GPU availability..."), console=console
+                Spinner("dots", text="🚀 Loading…"), console=console
             ) as live:
-                config = load_config()
-                try:
-                    user_info = authenticate_user(config)
-                except RuntimeError as e:
-                    live.stop()
-                    rprint(f"[red]❌ {str(e)}[/red]")
-                    return
-
-                # Validate SSH key matches configured GitHub username
-                live.update(Spinner("dots", text="🔐 Validating SSH key..."))
-                if not _validate_ssh_key_or_exit(config, live):
-                    return
+                with ThreadPoolExecutor(max_workers=3) as ex:
+                    f_auth = ex.submit(authenticate_user, config)
+                    # SSH validation may invoke `ssh git@github.com` interactively for password-protected keys;
+                    # do it on the main thread when the cache is cold so prompts work. Probe cache first.
+                    from .auth import _load_ssh_cache, validate_ssh_key_matches_github_user
+                    cached_ssh = _load_ssh_cache(config.get_github_username() or "")
+                    if cached_ssh is not None:
+                        f_ssh = None
+                        ssh_result = cached_ssh
+                    else:
+                        f_ssh = ex.submit(validate_ssh_key_matches_github_user, config, None)
+                        ssh_result = None
+                    f_avail = ex.submit(
+                        lambda: ReservationManager(config).get_gpu_availability_by_type()
+                    )
 
-                live.update(
-                    Spinner("dots", text="📡 Loading GPU availability..."))
-                reservation_mgr = ReservationManager(config)
-                availability_info = reservation_mgr.get_gpu_availability_by_type()
+                    # Surface auth failure first (most actionable).
+                    try:
+                        user_info = f_auth.result()
+                    except RuntimeError as e:
+                        live.stop()
+                        rprint(f"[red]❌ {str(e)}[/red]")
+                        return
 
-            live.stop()
+                    if ssh_result is None:
+                        ssh_result = f_ssh.result()
+                    availability_info = f_avail.result()
+
+            # Surface SSH validation failure with the same UX as before.
+            if not ssh_result.get("valid"):
+                rprint("[red]❌ Github SSH key validation failed[/red]")
+                if ssh_result.get("ssh_user") and ssh_result.get("configured_user"):
+                    rprint("\n[yellow]💡 Fix by updating your config:[/yellow]")
+                    rprint(f"   [cyan]gpu-dev config set github_user {ssh_result['ssh_user']}[/cyan]")
+                elif not ssh_result.get("configured_user"):
+                    rprint("\n[yellow]💡 Fix by configuring your GitHub username:[/yellow]")
+                    rprint("   [cyan]gpu-dev config set github_user <your-github-username>[/cyan]")
+                else:
+                    rprint("\n[yellow]💡 gpu-dev utilizes Github keys for auth![/yellow]")
+                    rprint("[yellow]💡 Check https://fburl.com/gh-ssh for info on how to add your ssh key to Github[/yellow]")
+                return
 
             if not availability_info:
                 rprint("[red]❌ Could not get GPU availability information[/red]")

{gpu_dev-0.5.7 → gpu_dev-0.5.9}/cli-tools/gpu-dev-cli/gpu_dev_cli/reservations.py

@@ -1032,6 +1032,11 @@ class ReservationManager:
                 queue_length = self._get_queue_length_for_gpu_type(gpu_type)
                 estimated_wait = queue_length * 15 if queue_length > 0 else 0
 
+                # size_etas is a DDB Map of {size_str: epoch_seconds (Decimal)} — pass through
+                # so the interactive count menu can render "[available in 1h24m]" labels.
+                raw_etas = item.get("size_etas", {}) or {}
+                size_etas = {str(k): int(v) for k, v in raw_etas.items()} if raw_etas else {}
+
                 availability_info[gpu_type] = {
                     "available": int(item.get("available_gpus", 0)),
                     "total": int(item.get("total_gpus", 0)),
@@ -1045,6 +1050,7 @@ class ReservationManager:
                     "last_updated": item.get("last_updated_timestamp", 0),
                     "maintenance": bool(item.get("maintenance", False)),
                     "maintenance_reason": item.get("maintenance_reason", ""),
+                    "size_etas": size_etas,
                 }
 
             return availability_info

{gpu_dev-0.5.7 → gpu_dev-0.5.9}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "gpu-dev"
-version = "0.5.7"
+version = "0.5.9"
 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.7 → gpu_dev-0.5.9}/terraform-gpu-devservers/availability.tf

@@ -26,8 +26,15 @@ resource "aws_lambda_function" "availability_updater" {
   role             = aws_iam_role.availability_updater_role.arn
   handler          = "index.handler"
   runtime          = "python3.11"
-  timeout          = 300
-  source_code_hash = null_resource.availability_updater_build.triggers.code_hash
+  timeout                        = 300
+  # 1769 MB is the sweet spot — Lambda allocates one full vCPU at this threshold.
+  # Beyond 1769 MB you get fractional second vCPUs (less linear gain), and our work is single-threaded.
+  memory_size                    = 1769
+  # Cap concurrent invocations at 1: each run does ~30 EKS API calls per gpu_type, and
+  # uncapped concurrency was hammering the cluster API into throttling, leaving later
+  # gpu_types in each run timing out and never producing size_etas.
+  reserved_concurrent_executions = 1
+  source_code_hash               = null_resource.availability_updater_build.triggers.code_hash
 
   environment {
     variables = {

{gpu_dev-0.5.7 → gpu_dev-0.5.9}/terraform-gpu-devservers/lambda/availability_updater/index.py

@@ -24,6 +24,45 @@ RESERVATIONS_TABLE = os.environ.get("RESERVATIONS_TABLE", "pytorch-gpu-dev-reser
 SUPPORTED_GPU_TYPES = json.loads(os.environ["SUPPORTED_GPU_TYPES"])
 
 
+def _parse_expires_at(value):
+    """Parse the reservations table's `expires_at` field to a unix epoch (int).
+
+    DDB stores it as either an ISO-8601 datetime string ("2026-05-02T00:12:03.674845") OR
+    occasionally a numeric epoch — handle both. Returns None if unparseable.
+    """
+    if value is None:
+        return None
+    # Numeric (Decimal/int/float) → epoch seconds directly.
+    if not isinstance(value, str):
+        try:
+            return int(float(value))
+        except (ValueError, TypeError):
+            return None
+    s = value.strip()
+    if not s:
+        return None
+    # ISO-8601 first (the actual production format).
+    try:
+        from datetime import datetime, timezone
+        # `fromisoformat` accepts microseconds; tolerate optional 'Z' suffix.
+        if s.endswith("Z"):
+            s2 = s[:-1] + "+00:00"
+        else:
+            s2 = s
+        dt = datetime.fromisoformat(s2)
+        if dt.tzinfo is None:
+            # Convention in this codebase: timestamps written via datetime.utcnow().isoformat() are UTC.
+            dt = dt.replace(tzinfo=timezone.utc)
+        return int(dt.timestamp())
+    except (ValueError, TypeError):
+        pass
+    # Numeric-as-string fallback.
+    try:
+        return int(float(s))
+    except (ValueError, TypeError):
+        return None
+
+
 def get_gpu_resource_name(gpu_type: str) -> str:
     return SUPPORTED_GPU_TYPES.get(gpu_type, {}).get("k8s_resource", "nvidia.com/gpu")
 
@@ -542,9 +581,8 @@ def compute_size_etas(v1, gpu_type, node_label_value, resource_name, gpus_per_in
             continue
         if pod_name not in pod_to_info:
             continue
-        try:
-            ts = int(float(expires_at))
-        except (ValueError, TypeError):
+        ts = _parse_expires_at(expires_at)
+        if ts is None:
             continue
         node_name, gpus = pod_to_info[pod_name]
         node_state[node_name]["expirations"].append((ts, gpus))

{gpu_dev-0.5.7 → gpu_dev-0.5.9}/terraform-gpu-devservers/lambda.tf

@@ -180,7 +180,7 @@ 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.7"
+      LAMBDA_VERSION                     = "0.5.9"
       MIN_CLI_VERSION                    = "0.5.5"
       DISK_CONTENTS_BUCKET               = aws_s3_bucket.disk_contents.bucket
       OPERATIONS_TABLE                   = aws_dynamodb_table.operations.name