wafer-cli 0.2.8__py3-none-any.whl → 0.2.10__py3-none-any.whl

wafer/GUIDE.md

@@ -8,12 +8,17 @@ Run code on cloud GPUs instantly with workspaces:
 
 ```bash
 wafer login                              # One-time auth
-wafer workspaces create dev --gpu B200   # Create workspace
+wafer workspaces create dev --gpu B200   # Create workspace (NVIDIA B200)
 wafer workspaces exec dev -- python -c "import torch; print(torch.cuda.get_device_name(0))"
 wafer workspaces sync dev ./my-project   # Sync files
 wafer workspaces exec dev -- python train.py
 ```
 
+**Available GPUs:**
+
+- `MI300X` - AMD Instinct MI300X (192GB HBM3, ROCm)
+- `B200` - NVIDIA Blackwell B200 (180GB HBM3e, CUDA) - default
+
 ## Documentation Lookup
 
 Answer GPU programming questions from indexed documentation.
@@ -83,13 +88,19 @@ wafer agent -t optimize-kernel \
 
 Cloud GPU environments with no setup required.
 
+**Available GPUs:**
+
+- `MI300X` - AMD Instinct MI300X (192GB HBM3, ROCm)
+- `B200` - NVIDIA Blackwell B200 (180GB HBM3e, CUDA) - default
+
 ```bash
-wafer workspaces create dev --gpu B200     # Create
-wafer workspaces list                       # List all
-wafer workspaces sync dev ./project         # Sync files
-wafer workspaces exec dev -- ./run.sh       # Run commands
-wafer workspaces ssh dev                    # Interactive SSH
-wafer workspaces delete dev                 # Cleanup
+wafer workspaces create dev --gpu B200 --wait     # NVIDIA B200
+wafer workspaces create amd-dev --gpu MI300X      # AMD MI300X
+wafer workspaces list                             # List all
+wafer workspaces sync dev ./project               # Sync files
+wafer workspaces exec dev -- ./run.sh             # Run commands
+wafer workspaces ssh dev                          # Interactive SSH
+wafer workspaces delete dev                       # Cleanup
 ```
 
 See `wafer workspaces --help` for details.

wafer/api_client.py

@@ -108,6 +108,7 @@ def run_command_stream(
     upload_dir: Path | None = None,
     workspace_id: str | None = None,
     gpu_id: int | None = None,
+    gpu_count: int = 1,
     docker_image: str | None = None,
     docker_entrypoint: str | None = None,
     pull_image: bool = False,
@@ -125,6 +126,7 @@ def run_command_stream(
         upload_dir: Directory to upload (stateless mode)
         workspace_id: Workspace ID from push (low-level mode)
         gpu_id: GPU ID to use (optional)
+        gpu_count: Number of GPUs needed (1-8, default 1)
         docker_image: Docker image override (optional)
         docker_entrypoint: Docker entrypoint override (optional, e.g., "bash")
         pull_image: Pull image if not available (optional, default False)
@@ -152,6 +154,8 @@ def run_command_stream(
 
     if gpu_id is not None:
         request_body["gpu_id"] = gpu_id
+    if gpu_count > 1:
+        request_body["gpu_count"] = gpu_count
     if docker_image is not None:
         request_body["docker_image"] = docker_image
     if docker_entrypoint is not None:

wafer/auth.py

@@ -345,3 +345,88 @@ def browser_login(timeout: int = 120, port: int | None = None) -> tuple[str, str
 
     server.server_close()
     raise TimeoutError(f"No response within {timeout} seconds")
+
+
+def device_code_login(timeout: int = 600) -> tuple[str, str | None]:
+    """Authenticate using state-based flow (no browser/port forwarding needed).
+
+    This is the SSH-friendly auth flow similar to GitHub CLI:
+    1. Request a state token from the API
+    2. Display the auth URL with state parameter
+    3. User visits URL on any device and signs in normally
+    4. Poll API until user completes authentication
+
+    Args:
+        timeout: Seconds to wait for authentication (default 600 = 10 minutes)
+
+    Returns:
+        Tuple of (access_token, refresh_token). refresh_token may be None.
+
+    Raises:
+        TimeoutError: If user doesn't authenticate within timeout
+        RuntimeError: If auth flow failed
+    """
+    api_url = get_api_url()
+
+    # Request state and auth URL
+    with httpx.Client(timeout=10.0) as client:
+        response = client.post(f"{api_url}/v1/auth/cli-auth/start", json={})
+        response.raise_for_status()
+        data = response.json()
+
+    state = data["state"]
+    auth_url = data["auth_url"]
+    expires_in = data["expires_in"]
+
+    # Display instructions to user
+    print("\n" + "=" * 60)
+    print("  WAFER CLI - Authentication")
+    print("=" * 60)
+    print(f"\n  Visit: {auth_url}")
+    print("\n  Sign in with GitHub to complete authentication")
+    print("\n" + "=" * 60 + "\n")
+
+    # Poll for authentication
+    start = time.time()
+    poll_interval = 5  # Poll every 5 seconds
+    last_poll = 0.0
+
+    print("Waiting for authentication", end="", flush=True)
+
+    while time.time() - start < min(timeout, expires_in):
+        # Show progress dots
+        if time.time() - last_poll >= poll_interval:
+            print(".", end="", flush=True)
+
+            # Poll the API
+            with httpx.Client(timeout=10.0) as client:
+                try:
+                    response = client.post(f"{api_url}/v1/auth/cli-auth/token", json={"state": state})
+
+                    if response.status_code == 200:
+                        # Success!
+                        data = response.json()
+                        print(f" {CHECK}\n")
+                        return data["access_token"], data.get("refresh_token")
+
+                    if response.status_code == 428:
+                        # Still waiting
+                        last_poll = time.time()
+                        time.sleep(1)
+                        continue
+
+                    # Some other error
+                    print(f" {CROSS}\n")
+                    raise RuntimeError(f"CLI auth flow failed: {response.status_code} {response.text}")
+
+                except httpx.RequestError as e:
+                    # Network error, retry
+                    print("!", end="", flush=True)
+                    last_poll = time.time()
+                    time.sleep(1)
+                    continue
+
+        time.sleep(0.5)  # Small sleep to avoid busy loop
+
+    print(f" {CROSS}\n")
+    raise TimeoutError(f"Authentication not completed within {expires_in} seconds")