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

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

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

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

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

@@ -16,6 +16,53 @@ from rich.spinner import Spinner
 _SSH_CACHE_TTL_SECONDS = 24 * 60 * 60
 _SSH_CACHE_PATH = Path(os.path.expanduser("~/.config/gpu-dev/ssh-validation-cache.json"))
 
+# 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_PATH = Path(os.path.expanduser("~/.config/gpu-dev/auth-cache.json"))
+
+
+def _auth_cache_key() -> str:
+    return os.environ.get("AWS_PROFILE", "default")
+
+
+def _load_auth_cache(github_user: str) -> Optional[Dict[str, Any]]:
+    try:
+        if not _AUTH_CACHE_PATH.exists():
+            return None
+        with open(_AUTH_CACHE_PATH) as f:
+            data = json.load(f)
+        entry = data.get(_auth_cache_key())
+        if not entry or entry.get("github_user") != github_user:
+            return None
+        if time.time() - float(entry.get("ts", 0)) > _AUTH_CACHE_TTL_SECONDS:
+            return None
+        return entry.get("result")
+    except Exception:
+        return None
+
+
+def _save_auth_cache(github_user: str, result: Dict[str, Any]) -> None:
+    try:
+        _AUTH_CACHE_PATH.parent.mkdir(parents=True, exist_ok=True)
+        data = {}
+        if _AUTH_CACHE_PATH.exists():
+            try:
+                with open(_AUTH_CACHE_PATH) as f:
+                    data = json.load(f)
+            except Exception:
+                data = {}
+        data[_auth_cache_key()] = {
+            "github_user": github_user,
+            "ts": int(time.time()),
+            "result": result,
+        }
+        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."""
@@ -50,31 +97,33 @@ def _save_ssh_cache(github_user: str, result: Dict[str, Any]) -> None:
 
 
 def authenticate_user(config: Config) -> Dict[str, Any]:
-    """Authenticate using AWS credentials - if you can call AWS, you're authorized"""
-    try:
-        # Test AWS access by getting caller identity
-        identity = config.get_user_identity()
+    """Authenticate using AWS credentials - if you can call AWS, you're authorized.
 
-        # Test specific resource access by trying to get queue URL
-        config.get_queue_url()
-
-        # Extract user info from AWS ARN
-        arn = identity["arn"]
-        user_name = arn.split("/")[-1]  # Extract username from ARN
+    Cached for 24h per AWS profile. The previous SQS get_queue_url probe was dropped:
+    it's a redundant permission check; reserve/cancel call SQS directly and surface
+    failures themselves, while list/show/avail don't touch SQS at all.
+    """
+    github_user = config.get_github_username()
+    if not github_user:
+        raise RuntimeError(
+            "GitHub username not configured. Please run: gpu-dev config set github_user <your-github-username>"
+        )
 
-        # Get GitHub username from config
-        github_user = config.get_github_username()
-        if not github_user:
-            raise RuntimeError(
-                f"GitHub username not configured. Please run: gpu-dev config set github_user <your-github-username>"
-            )
+    cached = _load_auth_cache(github_user)
+    if cached is not None:
+        return cached
 
-        return {
-            "user_id": user_name,  # AWS username for reservation ownership
-            "github_user": github_user,  # GitHub username for SSH keys
+    try:
+        identity = config.get_user_identity()
+        arn = identity["arn"]
+        user_name = arn.split("/")[-1]
+        result = {
+            "user_id": user_name,
+            "github_user": github_user,
             "arn": arn,
         }
-
+        _save_auth_cache(github_user, result)
+        return result
     except Exception as e:
         raise RuntimeError(f"AWS authentication failed: {e}")
 

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

@@ -688,6 +688,7 @@ def reserve(
             # and total wall-clock time drops from sum to max(each).
             from concurrent.futures import ThreadPoolExecutor
             config = load_config()
+            reservation_mgr = ReservationManager(config)
 
             with Live(
                 Spinner("dots", text="🚀 Loading…"), console=console
@@ -704,9 +705,7 @@ def reserve(
                     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()
-                    )
+                    f_avail = ex.submit(reservation_mgr.get_gpu_availability_by_type)
 
                     # Surface auth failure first (most actionable).
                     try:
@@ -2496,10 +2495,10 @@ def _show_availability() -> None:
             table = Table(
                 title="GPU Availability by Type (numbers are GPUs, not nodes)")
             table.add_column("GPU Type", style="cyan")
-            table.add_column("Available", style="green")
-            table.add_column("Max Reservable", style="bright_green")
+            table.add_column("Avail", style="green")
+            table.add_column("Max\nReservable", style="bright_green")
             table.add_column("Total", style="blue")
-            table.add_column("Queue Length", style="yellow")
+            table.add_column("Queue\nLength", style="yellow")
             table.add_column("Architecture", style="dim")
             table.add_column("Est. Wait Time", style="magenta")
 
@@ -2657,10 +2656,10 @@ def _show_availability_watch(interval: int) -> None:
                         table = Table(
                             title="GPU Availability by Type (numbers are GPUs, not nodes)")
                         table.add_column("GPU Type", style="cyan")
-                        table.add_column("Available", style="green")
-                        table.add_column("Max Reservable", style="blue")
+                        table.add_column("Avail", style="green")
+                        table.add_column("Max\nReservable", style="blue")
                         table.add_column("Total", style="blue")
-                        table.add_column("Queue Length", style="yellow")
+                        table.add_column("Queue\nLength", style="yellow")
                         table.add_column("Architecture", style="dim")
                         table.add_column("Est. Wait Time", style="magenta")
 

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

@@ -88,9 +88,9 @@ def select_gpu_type_interactive(
     console.print("\n[cyan]🖥️  GPU Availability:[/cyan]")
     table = Table()
     table.add_column("GPU Type", style="cyan")
-    table.add_column("Available", style="green")
+    table.add_column("Avail", style="green")
     table.add_column("Total", style="blue")
-    table.add_column("Queue Length", style="yellow")
+    table.add_column("Queue\nLength", style="yellow")
     table.add_column("Est. Wait Time", style="magenta")
 
     choices = []

{gpu_dev-0.5.13 → gpu_dev-0.5.15}/pyproject.toml

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

@@ -25,23 +25,40 @@ resource "kubernetes_config_map" "gpu_dev_mig_parted_config" {
   depends_on = [helm_release.nvidia_gpu_operator]
 }
 
-# Optional declarative B200 MIG node label. Set b200_mig_node_name in tfvars (or override the
-# variable's default below) to dedicate a specific B200 node to the mixed profile. Empty string
-# means "no node currently labelled" — the existing all-disabled stays in effect.
+# Declarative B200 MIG node label. Set b200_mig_node_name (per workspace via the locals lookup
+# below, or override via tfvars / -var) to dedicate a specific B200 node to the mixed profile.
+# Empty string means "no node labelled" — every B200 stays full.
+#
+# Future cleanup: when we split a B200 CR into two ASGs (one with mig_profile, one without),
+# the user_data path will set this label at boot for any instance in the MIG-dedicated ASG —
+# matching the H100 cr3 pattern. Until then, this declarative label pins the role to a hostname.
+locals {
+  # Workspace-scoped defaults so the resource is a no-op in non-prod and no apply ever tries to
+  # label a node that doesn't exist.
+  default_b200_mig_node_by_workspace = {
+    prod = "ip-10-0-67-125.us-east-2.compute.internal"
+  }
+  b200_mig_node_effective = (
+    var.b200_mig_node_name != ""
+    ? var.b200_mig_node_name
+    : lookup(local.default_b200_mig_node_by_workspace, terraform.workspace, "")
+  )
+}
+
 variable "b200_mig_node_name" {
-  description = "Hostname of the B200 node to label with nvidia.com/mig.config=b200-6full-2mig-balanced. Leave empty to skip."
+  description = "Hostname of the B200 node to label with nvidia.com/mig.config=b200-6full-2mig-balanced. Leave empty to use the per-workspace default in mig-config.tf."
   type        = string
   default     = ""
 }
 
 resource "kubernetes_labels" "b200_mig_node" {
-  count = var.b200_mig_node_name == "" ? 0 : 1
+  count = local.b200_mig_node_effective == "" ? 0 : 1
 
   api_version = "v1"
   kind        = "Node"
 
   metadata {
-    name = var.b200_mig_node_name
+    name = local.b200_mig_node_effective
   }
 
   labels = {