gpu-dev 0.5.23__tar.gz → 0.5.25__tar.gz

{gpu_dev-0.5.23 → gpu_dev-0.5.25}/PKG-INFO

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

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

{gpu_dev-0.5.23 → gpu_dev-0.5.25}/cli-tools/gpu-dev-cli/gpu_dev_cli/auth.py

@@ -38,6 +38,19 @@ def _load_auth_cache(github_user: str) -> Optional[Dict[str, Any]]:
             return None
         if time.time() - float(entry.get("ts", 0)) > _AUTH_CACHE_TTL_SECONDS:
             return None
+        # Defense against stale cache on a persistent disk that pre-dates the IRSA fix:
+        # if AWS_ROLE_ARN points at a role the cached ARN doesn\'t reference, the cache
+        # is from a different identity (e.g. IMDS-fallback before fs_group=1081 landed)
+        # and should be ignored.
+        expected_role_arn = os.environ.get("AWS_ROLE_ARN", "")
+        cached_arn = (entry.get("result") or {}).get("arn", "")
+        if expected_role_arn:
+            try:
+                role_name = expected_role_arn.rsplit("/", 1)[-1]
+                if role_name and role_name not in cached_arn:
+                    return None
+            except Exception:
+                pass
         return entry.get("result")
     except Exception:
         return None

{gpu_dev-0.5.23 → gpu_dev-0.5.25}/cli-tools/gpu-dev-cli/gpu_dev_cli/cli.py

@@ -4,6 +4,7 @@ Reserve and manage GPU development servers
 """
 
 import click
+import os
 from typing import Optional
 from rich.console import Console
 from rich.table import Table
@@ -3218,6 +3219,20 @@ def connect(ctx: click.Context, reservation_id: Optional[str]) -> None:
         if "-A" not in ssh_command and "-o ForwardAgent=yes" not in ssh_command:
             ssh_command = ssh_command.replace("ssh ", "ssh -A ", 1)
 
+        # When running from inside a gpu-dev pod (=GPU_DEV_USER_ID env var set) and the
+        # forwarded SSH agent is reachable but empty, the next hop is going to fail with
+        # 'Permission denied (publickey)'. Warn upfront so the user knows to ssh-add on
+        # their laptop instead of debugging an opaque auth failure on the remote side.
+        if os.environ.get("GPU_DEV_USER_ID") and os.environ.get("SSH_AUTH_SOCK"):
+            try:
+                import subprocess as _sp
+                r = _sp.run(["ssh-add", "-L"], capture_output=True, text=True, timeout=3)
+                if r.returncode != 0 or not r.stdout.strip() or "no identities" in r.stdout.lower():
+                    rprint("[yellow]⚠️  Forwarded SSH agent is empty — second-hop SSH from a pod will fail auth.[/yellow]")
+                    rprint("[yellow]   On your laptop: `ssh-add ~/.ssh/id_ed25519` (or your GitHub key), then reconnect to this pod with `gpu-dev connect`.[/yellow]\n")
+            except Exception:
+                pass
+
         # Parse and execute the command, capturing exit code for auth failures
         rprint(f"[dim]Executing: {ssh_command}[/dim]\n")
         result = subprocess.run(ssh_command, shell=True)

{gpu_dev-0.5.23 → gpu_dev-0.5.25}/cli-tools/gpu-dev-cli/gpu_dev_cli/reservations.py

@@ -162,11 +162,20 @@ def _generate_ssh_config(hostname: str, pod_name: str) -> str:
     Returns:
         SSH config content as string
     """
+    import sys
+    # AddKeysToAgent makes SSH stash the IdentityFile into ssh-agent the first time it
+    # uses it for this Host, so the next agent-forwarding hop (pod → pod) actually has
+    # something to forward. On macOS, UseKeychain persists the passphrase via Keychain
+    # so users aren\'t prompted on every shell restart. Linux ssh errors on UseKeychain,
+    # so guard it with IgnoreUnknown.
+    extra = "    AddKeysToAgent yes\n"
+    if sys.platform == "darwin":
+        extra += "    IgnoreUnknown UseKeychain\n    UseKeychain yes\n"
     config_content = f"""Host {pod_name}
     HostName {hostname}
     User dev
     ForwardAgent yes
-    ProxyCommand gpu-dev-ssh-proxy %h %p
+{extra}    ProxyCommand gpu-dev-ssh-proxy %h %p
     StrictHostKeyChecking no
     UserKnownHostsFile /dev/null
 """

{gpu_dev-0.5.23 → gpu_dev-0.5.25}/pyproject.toml

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

@@ -2888,6 +2888,7 @@ def allocate_gpu_resources(reservation_id: str, request: dict[str, Any], trace_d
             gpu_count=gpu_count,
             gpu_type=gpu_type,
             github_public_key=github_public_key,
+            github_user=github_user,
             reservation_id=reservation_id,
             jupyter_enabled=jupyter_enabled,
             persistent_volume_id=persistent_volume_id,
@@ -3434,7 +3435,8 @@ def create_kubernetes_resources(
     gpu_count: int,
     gpu_type: str,
     github_public_key: str,
-    reservation_id: str,
+    github_user: str = "",
+    reservation_id: str = None,
     jupyter_enabled: bool = False,
     persistent_volume_id: str = None,
     user_id: str = None,
@@ -3538,6 +3540,7 @@ def create_kubernetes_resources(
                         gpu_count,
                         gpu_type,
                         github_public_key,
+                        github_user=github_user,
                         jupyter_enabled=True,
                         persistent_volume_id=persistent_volume_id,
                         user_id=user_id,
@@ -3627,6 +3630,7 @@ def create_kubernetes_resources(
                         gpu_count,
                         gpu_type,
                         github_public_key,
+                        github_user=github_user,
                         jupyter_enabled=False,
                         persistent_volume_id=persistent_volume_id,
                         user_id=user_id,
@@ -3979,6 +3983,7 @@ def create_pod(
     gpu_count: int,
     gpu_type: str,
     github_public_key: str,
+    github_user: str = "",
     jupyter_enabled: bool = False,
     persistent_volume_id: str = None,
     user_id: str = None,

{gpu_dev-0.5.23 → gpu_dev-0.5.25}/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.24"
+      LAMBDA_VERSION                     = "0.5.25"
       MIN_CLI_VERSION                    = "0.5.16"
       DISK_CONTENTS_BUCKET               = aws_s3_bucket.disk_contents.bucket
       OPERATIONS_TABLE                   = aws_dynamodb_table.operations.name