gpu-dev 0.5.32__tar.gz → 0.6.2__tar.gz

{gpu_dev-0.5.32 → gpu_dev-0.6.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: gpu-dev
-Version: 0.5.32
+Version: 0.6.2
 Summary: CLI tool for PyTorch GPU developer server reservations
 Author: PyTorch Team
 Requires-Python: >=3.10
@@ -16,9 +16,29 @@ Requires-Dist: websockets>=12.0
 Requires-Dist: certifi>=2023.7.22
 Requires-Dist: mcp>=1.0.0
 
-# GPU Developer CLI
+# GPU Developer CLI & SDK
 
-A command-line tool for reserving and managing GPU development servers on AWS EKS.
+A command-line tool and Python SDK for reserving and managing GPU development servers.
+
+## Python SDK
+
+For programmatic access, use the [Python SDK](../../sdk/python/README.md):
+
+```python
+from gpu_dev import GpuDev
+
+client = GpuDev()
+sandbox = client.reserve(gpu_type="h100", gpu_count=2, hours=4)
+result = sandbox.exec("nvidia-smi")
+print(result.stdout)
+sandbox.cancel()
+```
+
+Install: `pip install -e sdk/python/` — see [SDK docs](../../sdk/python/README.md) and [quickstart notebook](../../sdk/python/examples/quickstart.ipynb).
+
+---
+
+## CLI
 
 ## Table of Contents
 

{gpu_dev-0.5.32 → gpu_dev-0.6.2}/cli-tools/gpu-dev-cli/README.md

@@ -1,6 +1,26 @@
-# GPU Developer CLI
+# GPU Developer CLI & SDK
 
-A command-line tool for reserving and managing GPU development servers on AWS EKS.
+A command-line tool and Python SDK for reserving and managing GPU development servers.
+
+## Python SDK
+
+For programmatic access, use the [Python SDK](../../sdk/python/README.md):
+
+```python
+from gpu_dev import GpuDev
+
+client = GpuDev()
+sandbox = client.reserve(gpu_type="h100", gpu_count=2, hours=4)
+result = sandbox.exec("nvidia-smi")
+print(result.stdout)
+sandbox.cancel()
+```
+
+Install: `pip install -e sdk/python/` — see [SDK docs](../../sdk/python/README.md) and [quickstart notebook](../../sdk/python/examples/quickstart.ipynb).
+
+---
+
+## CLI
 
 ## Table of Contents
 

{gpu_dev-0.5.32 → gpu_dev-0.6.2}/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.32
+Version: 0.6.2
 Summary: CLI tool for PyTorch GPU developer server reservations
 Author: PyTorch Team
 Requires-Python: >=3.10
@@ -16,9 +16,29 @@ Requires-Dist: websockets>=12.0
 Requires-Dist: certifi>=2023.7.22
 Requires-Dist: mcp>=1.0.0
 
-# GPU Developer CLI
+# GPU Developer CLI & SDK
 
-A command-line tool for reserving and managing GPU development servers on AWS EKS.
+A command-line tool and Python SDK for reserving and managing GPU development servers.
+
+## Python SDK
+
+For programmatic access, use the [Python SDK](../../sdk/python/README.md):
+
+```python
+from gpu_dev import GpuDev
+
+client = GpuDev()
+sandbox = client.reserve(gpu_type="h100", gpu_count=2, hours=4)
+result = sandbox.exec("nvidia-smi")
+print(result.stdout)
+sandbox.cancel()
+```
+
+Install: `pip install -e sdk/python/` — see [SDK docs](../../sdk/python/README.md) and [quickstart notebook](../../sdk/python/examples/quickstart.ipynb).
+
+---
+
+## CLI
 
 ## Table of Contents
 

{gpu_dev-0.5.32 → gpu_dev-0.6.2}/cli-tools/gpu-dev-cli/gpu_dev.egg-info/SOURCES.txt

@@ -30,6 +30,27 @@ docs/USER_GUIDE.md
 docs/devgpu-features.html
 docs/docker-mark-blue.svg
 docs/icons8-cursor-ai.svg
+sdk/python/README.md
+sdk/python/pyproject.toml
+sdk/python/examples/quickstart.ipynb
+sdk/python/src/gpu_dev/__init__.py
+sdk/python/src/gpu_dev/py.typed
+sdk/python/src/gpu_dev/_async/__init__.py
+sdk/python/src/gpu_dev/_backend/__init__.py
+sdk/python/src/gpu_dev/_backend/aws.py
+sdk/python/src/gpu_dev/_backend/protocol.py
+sdk/python/src/gpu_dev/_sync/__init__.py
+sdk/python/src/gpu_dev/_sync/client.py
+sdk/python/src/gpu_dev/_sync/sandbox.py
+sdk/python/src/gpu_dev/_transport/__init__.py
+sdk/python/src/gpu_dev/_transport/ssh.py
+sdk/python/src/gpu_dev/common/__init__.py
+sdk/python/src/gpu_dev/common/config.py
+sdk/python/src/gpu_dev/common/enums.py
+sdk/python/src/gpu_dev/common/errors.py
+sdk/python/src/gpu_dev/common/models.py
+sdk/python/tests/__init__.py
+sdk/python/tests/test_models.py
 terraform-gpu-devservers/.terraform.lock.hcl
 terraform-gpu-devservers/README.md
 terraform-gpu-devservers/alb.tf

{gpu_dev-0.5.32 → gpu_dev-0.6.2}/cli-tools/gpu-dev-cli/gpu_dev_cli/auth.py

@@ -13,7 +13,7 @@ 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_TTL_SECONDS = 14 * 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

{gpu_dev-0.5.32 → gpu_dev-0.6.2}/cli-tools/gpu-dev-cli/gpu_dev_cli/cli.py

@@ -41,33 +41,51 @@ from .interactive import (
 console = Console()
 
 
+_east1_table = None
+
 def _fetch_reservations_cross_region(reservation_mgr, user_filter, statuses, config=None):
     """Fetch reservations from current region + prod-east1 if on prod."""
-    reservations = reservation_mgr.list_reservations(
-        user_filter=user_filter, statuses_to_include=statuses)
-    # Cross-region fetch
-    try:
+    global _east1_table
+
+    from concurrent.futures import ThreadPoolExecutor
+
+    def _fetch_primary():
+        return reservation_mgr.list_reservations(
+            user_filter=user_filter, statuses_to_include=statuses)
+
+    def _fetch_east1():
+        global _east1_table
         cfg = config or load_config()
-        if cfg.user_config.get("environment") == "prod":
-            east1_env = Config.ENVIRONMENTS.get("prod-east1", {})
-            if east1_env:
-                import boto3 as _b3
-                east1_ddb = _b3.resource("dynamodb", region_name=east1_env["region"])
-                east1_table = east1_ddb.Table("pytorch-gpu-dev-reservations")
-                for st in (statuses or ["active"]):
-                    resp = east1_table.query(
-                        IndexName="StatusIndex",
-                        KeyConditionExpression="#s = :status",
-                        ExpressionAttributeNames={"#s": "status"},
-                        ExpressionAttributeValues={":status": st},
-                    )
-                    for item in resp.get("Items", []):
-                        if user_filter and item.get("user_id") != user_filter:
-                            continue
-                        item["_region"] = "us-east-1"
-                        reservations.append(item)
+        if cfg.user_config.get("environment") != "prod":
+            return []
+        east1_env = Config.ENVIRONMENTS.get("prod-east1", {})
+        if not east1_env or not user_filter:
+            return []
+        if _east1_table is None:
+            _east1_table = cfg.session.resource(
+                "dynamodb", region_name=east1_env["region"]
+            ).Table("pytorch-gpu-dev-reservations")
+        results = []
+        for st in (statuses or ["active"]):
+            resp = _east1_table.query(
+                IndexName="UserStatusIndex",
+                KeyConditionExpression="user_id = :uid AND #s = :status",
+                ExpressionAttributeNames={"#s": "status"},
+                ExpressionAttributeValues={":uid": user_filter, ":status": st},
+            )
+            for item in resp.get("Items", []):
+                item["_region"] = "us-east-1"
+                results.append(item)
+        return results
+
+    try:
+        with ThreadPoolExecutor(max_workers=2) as ex:
+            f1 = ex.submit(_fetch_primary)
+            f2 = ex.submit(_fetch_east1)
+            reservations = f1.result()
+            reservations.extend(f2.result())
     except Exception:
-        pass
+        reservations = _fetch_primary()
     return reservations
 
 
@@ -608,6 +626,8 @@ def main(ctx: click.Context) -> None:
 )
 @click.option("--spot", is_flag=True, default=False,
               help="Acknowledge spot instance (~1/3 cost, may be preempted with 2-min notice). Required for spot-only types.")
+@click.option("--fast-cache", is_flag=True, default=False, hidden=True,
+              help="Use NVMe local cache for faster session restore (experimental).")
 @click.pass_context
 def reserve(
     ctx: click.Context,
@@ -629,6 +649,7 @@ def reserve(
     disk: Optional[str],
     node_label: tuple,
     spot: bool = False,
+    fast_cache: bool = False,
 ) -> None:
     """Reserve GPU development server(s)
 
@@ -746,7 +767,10 @@ def reserve(
                     else:
                         f_ssh = ex.submit(validate_ssh_key_matches_github_user, config, None)
                         ssh_result = None
-                    f_avail = ex.submit(reservation_mgr.get_gpu_availability_by_type)
+                    # Only fetch availability if we need the interactive picker
+                    need_interactive = gpu_type is None
+                    if need_interactive:
+                        f_avail = ex.submit(reservation_mgr.get_gpu_availability_by_type)
 
                     # Surface auth failure first (most actionable).
                     try:
@@ -758,7 +782,7 @@ def reserve(
 
                     if ssh_result is None:
                         ssh_result = f_ssh.result()
-                    availability_info = f_avail.result()
+                    availability_info = f_avail.result() if need_interactive else None
 
             # Surface SSH validation failure with the same UX as before.
             if not ssh_result.get("valid"):
@@ -897,6 +921,13 @@ def reserve(
 
         else:
             # Non-interactive mode - use defaults and validate
+            # Route --spot to east1 when on prod (env vars override config region)
+            if spot and load_config().user_config.get("environment") == "prod":
+                east1_cfg = Config.ENVIRONMENTS.get("prod-east1", {})
+                if east1_cfg:
+                    import os as _os
+                    _os.environ["AWS_REGION"] = east1_cfg["region"]
+
             if gpu_type is None:
                 gpu_type = "a100"
             if hours is None:
@@ -1101,11 +1132,13 @@ def reserve(
                     rprint(f"[red]❌ {str(e)}[/red]")
                     return
 
-                # Validate SSH key matches configured GitHub username
-                live.update(Spinner("dots", text="🔐 Validating SSH key..."))
+                # Validate SSH key matches configured GitHub username (cached, ~0ms)
                 if not _validate_ssh_key_or_exit(config, live):
                     return
 
+                live.update(Spinner("dots", text="📡 Preparing reservation..."))
+                reservation_mgr = ReservationManager(config)
+
                 # Track if user explicitly requests no persistent disk
                 explicit_no_disk = explicit_no_disk_from_param
 
@@ -1217,11 +1250,6 @@ def reserve(
                                 rprint(f"[yellow]Use a different disk or wait for the reservation to end[/yellow]")
                                 return
 
-                live.update(
-                    Spinner("dots", text="📡 Setting up reservation manager...")
-                )
-                reservation_mgr = ReservationManager(config)
-
             # Submit reservation request
             live.update(
                 Spinner("dots", text="📡 Submitting reservation request...")
@@ -1357,6 +1385,7 @@ def reserve(
                     spot=spot,
                     node_labels=node_labels if node_labels else None,
                     trace=trace,
+                    fast_cache=fast_cache,
                 )
                 reservation_ids = [reservation_id] if reservation_id else None
 
@@ -2568,10 +2597,21 @@ def cancel(
             with Live(
                 Spinner("dots", text="📡 Cancelling reservations..."), console=console
             ) as live:
+                # Build east1 reservation manager for cross-region cancellations
+                east1_mgr = None
+                east1_env = Config.ENVIRONMENTS.get("prod-east1", {})
+                if east1_env:
+                    import os as _os
+                    _east1_config = Config()
+                    _east1_config.aws_region = east1_env["region"]
+                    east1_mgr = ReservationManager(_east1_config)
+
                 for reservation in reservations:
                     res_id = reservation.get("reservation_id", "")
                     if res_id:
-                        success = reservation_mgr.cancel_reservation(
+                        # Use east1 manager for east1 reservations
+                        mgr = east1_mgr if reservation.get("_region") in ("east1", "us-east-1") and east1_mgr else reservation_mgr
+                        success = mgr.cancel_reservation(
                             res_id, user_info["user_id"]
                         )
                         if success:
@@ -2869,36 +2909,42 @@ def _show_availability() -> None:
         ) as live:
             config = load_config()
 
-            # Authenticate using AWS credentials
+            # Authenticate and fetch availability (both regions in parallel)
             try:
                 user_info = authenticate_user(config)
                 reservation_mgr = ReservationManager(config)
-                availability_info = reservation_mgr.get_gpu_availability_by_type()
+
+                from concurrent.futures import ThreadPoolExecutor
+                _env_name = config.user_config.get("environment", "prod")
+                _east1_spot_types = frozenset(Config.ENVIRONMENTS.get("prod-east1", {}).get("spot_types", []))
+
+                def _fetch_east1_spot():
+                    if _env_name != "prod" or not _east1_spot_types:
+                        return {}
+                    east1_r = Config.ENVIRONMENTS["prod-east1"]["region"]
+                    east1_table = config.session.resource("dynamodb", region_name=east1_r).Table("pytorch-gpu-dev-gpu-availability")
+                    result = {}
+                    for item in east1_table.scan().get("Items", []):
+                        gt = item.get("gpu_type", "")
+                        if gt in _east1_spot_types:
+                            result[gt] = {
+                                "available": int(item.get("available_gpus", 0)),
+                                "total": int(item.get("total_gpus", 0)),
+                                "max_reservable": int(item.get("max_reservable", 0)),
+                                "spot_info": item.get("spot_info", {}),
+                            }
+                    return result
+
+                with ThreadPoolExecutor(max_workers=2) as ex:
+                    f_avail = ex.submit(reservation_mgr.get_gpu_availability_by_type)
+                    f_spot = ex.submit(_fetch_east1_spot)
+                    availability_info = f_avail.result()
+                    spot_region_info = f_spot.result()
             except RuntimeError as e:
                 live.stop()
                 rprint(f"[red]❌ {str(e)}[/red]")
                 return
 
-        # Cross-region: fetch spot availability from prod-east1
-        spot_region_info = {}
-        _env_name = config.user_config.get("environment", "prod")
-        _east1_spot_types = frozenset(Config.ENVIRONMENTS.get("prod-east1", {}).get("spot_types", []))
-        if _env_name == "prod" and _east1_spot_types:
-            try:
-                import boto3 as _b3
-                east1_r = Config.ENVIRONMENTS["prod-east1"]["region"]
-                for item in _b3.resource("dynamodb", region_name=east1_r).Table("pytorch-gpu-dev-gpu-availability").scan().get("Items", []):
-                    gt = item.get("gpu_type", "")
-                    if gt in _east1_spot_types:
-                        spot_region_info[gt] = {
-                            "available": int(item.get("available_gpus", 0)),
-                            "total": int(item.get("total_gpus", 0)),
-                            "max_reservable": int(item.get("max_reservable", 0)),
-                            "spot_info": item.get("spot_info", {}),
-                        }
-            except Exception:
-                pass
-
         if availability_info:
             # GPU architecture mapping (for display)
             gpu_architectures = {
@@ -3255,8 +3301,19 @@ def connect(ctx: click.Context, reservation_id: Optional[str]) -> None:
     For VS Code Remote or manual SSH, use 'gpu-dev show' to see full SSH command.
     """
     import subprocess
+    from pathlib import Path
 
     try:
+        # Fast path: if reservation ID given, check local SSH config first (no network)
+        if reservation_id:
+            ssh_config_dir = Path.home() / ".gpu-dev"
+            matches = list(ssh_config_dir.glob(f"{reservation_id}*-sshconfig")) if ssh_config_dir.exists() else []
+            if matches:
+                pod_name = f"gpu-dev-{reservation_id[:8]}"
+                rprint(f"[cyan]Connecting to {pod_name}...[/cyan]\n")
+                os.execvp("ssh", ["ssh", pod_name])
+                return
+
         with Live(
             Spinner("dots", text="📡 Fetching reservation details..."), console=console
         ) as live:
@@ -3301,21 +3358,30 @@ def connect(ctx: click.Context, reservation_id: Optional[str]) -> None:
 
                 live.start()
 
-            # If the selected reservation is from east1, switch to east1 reservation_mgr
-            _sel = next((r for r in (locals().get("reservations") or []) if r.get("reservation_id", "").startswith(reservation_id)), None)
-            if _sel and _sel.get("_region") == "us-east-1":
-                import os as _os
-                east1_cfg = Config.ENVIRONMENTS.get("prod-east1", {})
-                _os.environ["AWS_DEFAULT_REGION"] = east1_cfg["region"]
-                _east1_config = Config()
-                _east1_config.aws_region = east1_cfg["region"]
-                reservation_mgr = ReservationManager(_east1_config)
-
-            # Get connection info
+            # Try current region first, then cross-region if not found
             connection_info = reservation_mgr.get_connection_info(
                 reservation_id, user_info["user_id"]
             )
 
+            # If not found, try the other region
+            if not connection_info:
+                import os as _os
+                current_env = config.user_config.get("environment", "prod")
+                other_envs = {"prod": "prod-east1", "prod-east1": "prod"}
+                other_env_name = other_envs.get(current_env)
+                if other_env_name:
+                    other_env = Config.ENVIRONMENTS.get(other_env_name, {})
+                    if other_env:
+                        _os.environ["AWS_DEFAULT_REGION"] = other_env["region"]
+                        _other_config = Config()
+                        _other_config.aws_region = other_env["region"]
+                        other_mgr = ReservationManager(_other_config)
+                        connection_info = other_mgr.get_connection_info(
+                            reservation_id, user_info["user_id"]
+                        )
+                        if connection_info:
+                            reservation_mgr = other_mgr
+
         live.stop()
 
         if not connection_info:
@@ -3864,7 +3930,7 @@ def set(key: str, value: str) -> None:
 
 
 @config.command()
-@click.argument("env_name", type=click.Choice(["test", "prod", "prod-east1"]))
+@click.argument("env_name", type=click.Choice(["test", "prod"]))
 def environment(env_name: str) -> None:
     """Set the environment
 
@@ -3876,7 +3942,7 @@ def environment(env_name: str) -> None:
     \b
     Examples:
         gpu-dev config environment prod        # Production (us-east-2)
-        gpu-dev config environment prod-east1  # Spot-only us-east-1
+        gpu-dev config environment prod         # Production (spot accessible via interactive picker)
         gpu-dev config environment test        # Test (us-west-1)
 
     Environment configurations:

{gpu_dev-0.5.32 → gpu_dev-0.6.2}/cli-tools/gpu-dev-cli/gpu_dev_cli/config.py

@@ -3,6 +3,7 @@
 import os
 import json
 import boto3
+import botocore.exceptions
 from pathlib import Path
 from typing import Dict, Any, Optional
 
@@ -42,13 +43,14 @@ class Config:
         # Load unified config (handles migration from legacy files)
         self.user_config = self._load_config()
 
-        # Get region from config, then AWS env vars, or default
-        if self.user_config.get("region"):
+        # Get region: env vars take priority (for spot routing), then config, then default
+        env_region = os.getenv("AWS_REGION") or os.getenv("AWS_DEFAULT_REGION")
+        if env_region and env_region != self.user_config.get("region"):
+            self.aws_region = env_region
+        elif self.user_config.get("region"):
             self.aws_region = self.user_config["region"]
         else:
-            self.aws_region = os.getenv(
-                "AWS_REGION", os.getenv("AWS_DEFAULT_REGION", "us-east-2")
-            )
+            self.aws_region = "us-east-2"
 
         os.environ["AWS_DEFAULT_REGION"] = self.aws_region
 
@@ -71,17 +73,63 @@ class Config:
         self._sqs_client = None
         self._dynamodb = None
 
+    _CRED_CACHE = Path.home() / ".config" / "gpu-dev" / "aws-cred-cache.json"
+
     def _create_aws_session(self):
-        """Create AWS session with profile support"""
-        available_profiles = boto3.Session().available_profiles
-        if "gpu-dev" in available_profiles:
-            try:
-                session = boto3.Session(profile_name="gpu-dev")
-                session.get_credentials()
-                return session
-            except Exception:
-                pass
-        return boto3.Session()
+        """Create AWS session, caching resolved credentials to skip SSO resolution (~900ms)."""
+        import time as _time
+
+        # Try cached credentials first (avoids 900ms SSO resolution)
+        try:
+            if self._CRED_CACHE.exists():
+                cached = json.loads(self._CRED_CACHE.read_text())
+                if _time.time() < cached.get("expires", 0):
+                    return boto3.Session(
+                        aws_access_key_id=cached["access_key"],
+                        aws_secret_access_key=cached["secret_key"],
+                        aws_session_token=cached["token"],
+                        region_name=self.aws_region,
+                    )
+        except Exception:
+            pass
+
+        # Resolve credentials from SSO/profile (slow path, ~900ms)
+        try:
+            session = boto3.Session(profile_name="gpu-dev")
+            creds = session.get_credentials()
+            if not creds:
+                raise Exception("no credentials")
+        except Exception:
+            session = boto3.Session()
+            creds = session.get_credentials()
+
+        # Cache resolved credentials (safe — they're short-lived STS tokens)
+        try:
+            frozen = creds.get_frozen_credentials()
+            if frozen.token:
+                self._CRED_CACHE.parent.mkdir(parents=True, exist_ok=True)
+                self._CRED_CACHE.write_text(json.dumps({
+                    "access_key": frozen.access_key,
+                    "secret_key": frozen.secret_key,
+                    "token": frozen.token,
+                    "expires": _time.time() + 2700,  # cache 45min (SSO tokens last ~1h)
+                }))
+                self._CRED_CACHE.chmod(0o600)
+        except Exception:
+            pass
+
+        return session
+
+    def refresh_session(self):
+        """Clear cached credentials and re-resolve. Called on ExpiredTokenException."""
+        try:
+            self._CRED_CACHE.unlink(missing_ok=True)
+        except Exception:
+            pass
+        self.session = self._create_aws_session()
+        self._sts_client = None
+        self._sqs_client = None
+        self._dynamodb = None
 
     @property
     def sts_client(self):

{gpu_dev-0.5.32 → gpu_dev-0.6.2}/cli-tools/gpu-dev-cli/gpu_dev_cli/disks.py

@@ -355,8 +355,21 @@ def unlock_disk(disk_name: str, user_id: str, config: Config) -> bool:
         return False
 
     if not disk['in_use']:
-        print(f"Disk '{disk_name}' is not locked")
-        return False
+        # DDB says not locked — but check if EBS volume is still physically attached
+        try:
+            ec2 = config.session.client('ec2', region_name=config.aws_region)
+            vols = ec2.describe_volumes(Filters=[
+                {"Name": "tag:gpu-dev-user", "Values": [user_id]},
+                {"Name": "tag:disk_name", "Values": [disk_name]},
+                {"Name": "status", "Values": ["in-use"]},
+            ]).get("Volumes", [])
+            if not vols:
+                print(f"Disk '{disk_name}' is not locked")
+                return False
+            print(f"Disk '{disk_name}' DDB lock is clear but EBS volume is still attached — sending force-detach request")
+        except Exception:
+            print(f"Disk '{disk_name}' is not locked")
+            return False
 
     operation_id = str(uuid.uuid4())