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

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wafer-cli
-Version: 0.2.40
+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.40 → wafer_cli-0.2.41}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "wafer-cli"
-version = "0.2.40"
+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.40 → wafer_cli-0.2.41}/wafer/workspaces.py

@@ -844,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,
@@ -853,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)
@@ -868,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.40 → 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.40
+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