wafer-cli 0.2.39__tar.gz → 0.2.41__tar.gz

{wafer_cli-0.2.39 → wafer_cli-0.2.41}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wafer-cli
-Version: 0.2.39
+Version: 0.2.41
 Summary: CLI for running GPU workloads, managing remote workspaces, and evaluating/optimizing kernels
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown

{wafer_cli-0.2.39 → wafer_cli-0.2.41}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "wafer-cli"
-version = "0.2.39"
+version = "0.2.41"
 description = "CLI for running GPU workloads, managing remote workspaces, and evaluating/optimizing kernels"
 readme = "README.md"
 requires-python = ">=3.11"

{wafer_cli-0.2.39 → wafer_cli-0.2.41}/wafer/cli.py

@@ -220,16 +220,22 @@ workspaces_app = typer.Typer(
 
 Workspaces are on-demand cloud GPU environments. Requires authentication (wafer login).
 
+Environment Types:
+  modal      Serverless GPU execution (fast startup, no SSH)
+  baremetal  Dedicated GPU server (SSH access, hardware counters)
+
 Available GPUs:
-  MI300X  AMD Instinct MI300X (192GB HBM3, ROCm)
+  MI300X  AMD Instinct MI300X (192GB HBM3, ROCm) - baremetal only
   B200    NVIDIA Blackwell B200 (180GB HBM3e, CUDA)
+  H100    NVIDIA Hopper H100 (80GB HBM3, CUDA)
 
 Commands:
-  wafer workspaces create dev --gpu B200   # Create workspace
-  wafer workspaces exec dev -- python x.py # Run commands
-  wafer workspaces ssh dev                 # Interactive SSH
-  wafer workspaces sync dev ./project      # Sync files
-  wafer workspaces delete dev              # Clean up"""
+  wafer workspaces create dev -g B200 -e modal      # Create Modal workspace
+  wafer workspaces create dev -g B200 -e baremetal  # Create baremetal workspace
+  wafer workspaces exec dev -- python x.py          # Run commands
+  wafer workspaces ssh dev                          # Interactive SSH (baremetal only)
+  wafer workspaces sync dev ./project               # Sync files
+  wafer workspaces delete dev                       # Clean up"""
 )
 app.add_typer(workspaces_app, name="workspaces")
 
@@ -3272,7 +3278,7 @@ def demo_eval(
         # Step 1: Create workspace
         typer.echo(f"\n[1/4] Creating workspace '{workspace_name}'...")
         result = subprocess.run(
-            ["wafer", "workspaces", "create", workspace_name, "--gpu", "B200", "--json"],
+            ["wafer", "workspaces", "create", workspace_name, "--gpu", "B200", "--environment", "modal", "--json"],
             capture_output=True,
             text=True,
             check=True,
@@ -3359,7 +3365,7 @@ print(f"Performance: {(t1-t0)/100*1e6:.1f} us/iter")
             typer.echo("\n✓ Demo complete! To evaluate your own kernels:")
             typer.echo("")
             typer.echo("  # Using workspaces (no setup required):")
-            typer.echo("  wafer workspaces create dev --gpu B200")
+            typer.echo("  wafer workspaces create dev --gpu B200 --environment modal")
             typer.echo("  wafer workspaces exec --sync ./my-kernel dev -- python my_test.py")
             typer.echo("")
             typer.echo("  # Or using wafer evaluate with a configured target:")
@@ -4816,7 +4822,10 @@ def workspaces_list(
 def workspaces_create(
     name: str = typer.Argument(..., help="Workspace name"),
     gpu_type: str = typer.Option(
-        "B200", "--gpu", "-g", help="GPU type: MI300X (AMD) or B200 (NVIDIA, default)"
+        ..., "--gpu", "-g", help="GPU type: MI300X (AMD) or B200/H100 (NVIDIA)"
+    ),
+    environment: str = typer.Option(
+        ..., "--environment", "-e", help="Environment type: modal or baremetal"
     ),
     image: str | None = typer.Option(None, "--image", "-i", help="Docker image (optional)"),
     wait: bool = typer.Option(
@@ -4826,16 +4835,22 @@ def workspaces_create(
 ) -> None:
     """Create a new workspace.
 
+    Per-vendor architecture: each workspace has a single environment type.
+
+    Environment Types:
+        modal      Serverless GPU execution (fast startup, no SSH)
+        baremetal  Dedicated GPU server (SSH access, hardware counters)
+
     Available GPUs:
-        MI300X  AMD Instinct MI300X (192GB HBM3, ROCm)
-        B200    NVIDIA Blackwell B200 (180GB HBM3e, CUDA)
+        MI300X     AMD Instinct MI300X (192GB HBM3, ROCm) - baremetal only
+        B200       NVIDIA Blackwell B200 (180GB HBM3e, CUDA)
+        H100       NVIDIA Hopper H100 (80GB HBM3, CUDA)
 
     Example:
-        wafer workspaces create my-kernel                # B200 (default)
-        wafer workspaces create my-kernel --gpu MI300X   # AMD MI300X
-        wafer workspaces create my-kernel --gpu B200     # NVIDIA B200
-        wafer workspaces create my-kernel --image pytorch/pytorch:2.5.1-cuda12.4-cudnn9-devel
-        wafer workspaces create my-kernel --wait
+        wafer workspaces create my-kernel --gpu B200 --environment modal
+        wafer workspaces create my-kernel --gpu MI300X --environment baremetal
+        wafer workspaces create my-kernel -g B200 -e baremetal  # SSH + ncu profiling
+        wafer workspaces create my-kernel -g B200 -e modal --wait
     """
     from .workspaces import create_workspace
 
@@ -4843,6 +4858,7 @@ def workspaces_create(
         result = create_workspace(
             name,
             gpu_type=gpu_type,
+            environment_type=environment,
             image=image,
             wait=wait,
             json_output=json_output,

{wafer_cli-0.2.39 → wafer_cli-0.2.41}/wafer/workspaces.py

@@ -249,16 +249,20 @@ def list_workspaces(json_output: bool = False) -> str:
 
 def create_workspace(
     name: str,
-    gpu_type: str = "B200",
+    gpu_type: str,
+    environment_type: str,
     image: str | None = None,
     wait: bool = False,
     json_output: bool = False,
 ) -> str:
     """Create a new workspace.
 
+    Per-vendor architecture: each workspace has a single environment type.
+
     Args:
         name: Workspace name (must be unique)
-        gpu_type: GPU type (default: B200)
+        gpu_type: GPU type (required: B200, H100, MI300X)
+        environment_type: Environment type (required: modal, baremetal)
         image: Docker image (optional, uses default if not specified)
         wait: If True, stream provisioning progress and return SSH credentials
         json_output: If True, return raw JSON; otherwise return formatted text
@@ -272,6 +276,7 @@ def create_workspace(
     # Validate inputs
     assert name, "Workspace name must be non-empty"
     assert gpu_type, "GPU type must be non-empty"
+    assert environment_type, "Environment type must be non-empty"
 
     api_url, headers = _get_client()
 
@@ -296,6 +301,7 @@ def create_workspace(
     request_body: dict = {
         "name": name,
         "gpu_type": gpu_type,
+        "environment_type": environment_type,
     }
     if image:
         request_body["image"] = image
@@ -838,6 +844,73 @@ def _parse_sse_content(content: str) -> SSEEvent:
     return SSEEvent(output=content, exit_code=None, is_error=False)
 
 
+def _exec_via_ssh(ssh_host: str, ssh_port: int, ssh_user: str, command: str) -> int:
+    """Execute command via SSH, streaming output to stdout/stderr.
+
+    Used for baremetal workspaces. The workspace's zsh plugin handles GPU routing.
+
+    Returns:
+        Exit code from remote command
+    """
+    import selectors
+    import shlex
+    import subprocess
+    import sys
+
+    assert ssh_host, "SSH host required"
+    assert ssh_port > 0, "SSH port must be positive"
+    assert ssh_user, "SSH user required"
+    assert command, "Command required"
+
+    ssh_cmd = [
+        "ssh",
+        "-p", str(ssh_port),
+        "-t",  # Force TTY for zsh plugin to work
+        "-o", "StrictHostKeyChecking=no",
+        "-o", "UserKnownHostsFile=/dev/null",
+        "-o", "BatchMode=yes",
+        "-o", "LogLevel=ERROR",
+        f"{ssh_user}@{ssh_host}",
+        f"zsh -i -l -c {shlex.quote(command)}",  # Interactive login shell to load plugins
+    ]
+
+    process = subprocess.Popen(
+        ssh_cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.PIPE,
+        text=True,
+    )
+
+    assert process.stdout is not None
+    assert process.stderr is not None
+
+    sel = selectors.DefaultSelector()
+    sel.register(process.stdout, selectors.EVENT_READ)
+    sel.register(process.stderr, selectors.EVENT_READ)
+
+    while True:
+        for key, _ in sel.select(timeout=0.1):
+            line = key.fileobj.readline()  # type: ignore
+            if line:
+                if key.fileobj == process.stdout:
+                    print(line, end="", file=sys.stdout)
+                else:
+                    print(line, end="", file=sys.stderr)
+
+        if process.poll() is not None:
+            for line in process.stdout:
+                print(line, end="", file=sys.stdout)
+            for line in process.stderr:
+                print(line, end="", file=sys.stderr)
+            sys.stdout.flush()
+            sys.stderr.flush()
+            break
+
+    sel.close()
+    assert process.returncode is not None, "SSH process did not terminate properly"
+    return process.returncode
+
+
 def exec_command(
     workspace_id: str,
     command: str,
@@ -847,11 +920,18 @@ def exec_command(
 ) -> int:
     """Execute a command in workspace, streaming output.
 
+    For baremetal workspaces (with SSH access), commands are executed via SSH.
+    The workspace's zsh plugin handles GPU routing automatically, ensuring
+    packages installed via pip persist across commands.
+
+    For Modal workspaces (no SSH), commands are executed via the API.
+
     Args:
         workspace_id: Workspace ID or name
         command: Command to execute
         timeout_seconds: Execution timeout (default: 300, from config)
         routing: Routing hint - "auto", "gpu", "cpu", or "baremetal" (default: auto)
+        pull_image: Pull image on target if missing (only for API exec)
 
     Returns:
         Exit code (0 = success, non-zero = failure)
@@ -862,6 +942,20 @@ def exec_command(
     assert workspace_id, "Workspace ID must be non-empty"
     assert command, "Command must be non-empty"
 
+    # Get workspace details to check if SSH is available (baremetal)
+    workspace = get_workspace_raw(workspace_id)
+    ssh_host = workspace.get("ssh_host")
+    ssh_port = workspace.get("ssh_port")
+    ssh_user = workspace.get("ssh_user")
+
+    # Baremetal workspaces have SSH access - use SSH for stateful execution
+    # This ensures pip installs persist because we're in the workspace container
+    # The zsh plugin still handles GPU routing for GPU commands
+    if ssh_host and ssh_port and ssh_user:
+        return _exec_via_ssh(ssh_host, ssh_port, ssh_user, command)
+
+    # Modal workspaces (no SSH) - use API exec
+    # Modal Named Sandboxes persist state within their lifetime
     api_url, headers = _get_client()
 
     # Base64 encode command to avoid escaping issues

{wafer_cli-0.2.39 → wafer_cli-0.2.41}/wafer_cli.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wafer-cli
-Version: 0.2.39
+Version: 0.2.41
 Summary: CLI for running GPU workloads, managing remote workspaces, and evaluating/optimizing kernels
 Requires-Python: >=3.11
 Description-Content-Type: text/markdown