wafer-cli 0.2.9__py3-none-any.whl → 0.2.11__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: