wafer-cli 0.1.0__py3-none-any.whl

wafer/__init__.py

@@ -0,0 +1,3 @@
+"""Wafer - Standalone tool for Docker execution on remote GPUs."""
+
+__version__ = "0.1.0"

wafer/api_client.py

@@ -0,0 +1,201 @@
+"""Wafer API client for remote GPU operations.
+
+Thin client that calls wafer-api endpoints instead of direct SSH.
+"""
+
+import base64
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+
+import httpx
+
+# Default API URL (can be overridden via environment variable)
+DEFAULT_API_URL = "http://localhost:8000"
+
+
+@dataclass(frozen=True)
+class PushResult:
+    """Result of pushing files to GPU."""
+
+    workspace_id: str
+    workspace_path: str
+    files_uploaded: list[str]
+
+
+@dataclass(frozen=True)
+class ApiConfig:
+    """API client configuration."""
+
+    base_url: str = DEFAULT_API_URL
+    timeout: float = 60.0
+
+
+def get_api_url() -> str:
+    """Get API URL from environment or default."""
+    import os
+
+    return os.environ.get("WAFER_API_URL", DEFAULT_API_URL)
+
+
+def _get_auth_headers() -> dict[str, str]:
+    """Get auth headers from stored credentials (lazy import to avoid circular)."""
+    from .auth import get_auth_headers
+
+    return get_auth_headers()
+
+
+def push_directory(local_path: Path, workspace_name: str | None = None) -> PushResult:
+    """Push local directory to GPU via wafer-api.
+
+    Args:
+        local_path: Local directory to upload
+        workspace_name: Optional workspace name (defaults to directory name)
+
+    Returns:
+        PushResult with workspace_id and uploaded files
+
+    Raises:
+        FileNotFoundError: If local_path doesn't exist
+        ValueError: If local_path is not a directory
+        httpx.HTTPError: If API request fails
+    """
+    if not local_path.exists():
+        raise FileNotFoundError(f"Path not found: {local_path}")
+    if not local_path.is_dir():
+        raise ValueError(f"Not a directory: {local_path}")
+
+    # Collect files and encode as base64
+    files = []
+    for file_path in local_path.rglob("*"):
+        if file_path.is_file():
+            relative_path = file_path.relative_to(local_path)
+            content = file_path.read_bytes()
+            files.append({
+                "path": str(relative_path),
+                "content": base64.b64encode(content).decode(),
+            })
+
+    # Build request
+    request_body = {
+        "files": files,
+        "workspace_name": workspace_name or local_path.name,
+    }
+
+    # Call API
+    api_url = get_api_url()
+    headers = _get_auth_headers()
+    with httpx.Client(timeout=60.0, headers=headers) as client:
+        response = client.post(f"{api_url}/v1/gpu/push", json=request_body)
+        response.raise_for_status()
+        data = response.json()
+
+    return PushResult(
+        workspace_id=data["workspace_id"],
+        workspace_path=data["workspace_path"],
+        files_uploaded=data["files_uploaded"],
+    )
+
+
+def _collect_files(local_path: Path) -> list[dict]:
+    """Collect files from directory as base64-encoded dicts."""
+    files = []
+    for file_path in local_path.rglob("*"):
+        if file_path.is_file():
+            relative_path = file_path.relative_to(local_path)
+            content = file_path.read_bytes()
+            files.append({
+                "path": str(relative_path),
+                "content": base64.b64encode(content).decode(),
+            })
+    return files
+
+
+def run_command_stream(
+    command: str,
+    upload_dir: Path | None = None,
+    workspace_id: str | None = None,
+    gpu_id: int | None = None,
+    docker_image: str | None = None,
+    docker_entrypoint: str | None = None,
+    pull_image: bool = False,
+    require_hardware_counters: bool = False,
+    target: str | None = None,
+) -> int:
+    """Run command on GPU via wafer-api, streaming output.
+
+    Two modes (mutually exclusive):
+    - upload_dir: Upload files and run (stateless, high-level)
+    - workspace_id: Use existing workspace (low-level)
+
+    Args:
+        command: Command to execute inside container
+        upload_dir: Directory to upload (stateless mode)
+        workspace_id: Workspace ID from push (low-level mode)
+        gpu_id: GPU ID to use (optional)
+        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)
+        require_hardware_counters: Require baremetal for ncu profiling (optional)
+        target: Target name to use (optional, defaults to user's default)
+
+    Returns:
+        Exit code (0 = success, non-zero = failure)
+
+    Raises:
+        httpx.HTTPError: If API request fails
+    """
+    request_body: dict = {
+        "command": command,
+    }
+
+    # Add files or workspace_id (mutually exclusive)
+    if upload_dir is not None:
+        files = _collect_files(upload_dir)
+        request_body["files"] = files
+        request_body["workspace_name"] = upload_dir.name
+    elif workspace_id is not None:
+        request_body["workspace_id"] = workspace_id
+    # else: no files, no workspace (run command in temp workspace)
+
+    if gpu_id is not None:
+        request_body["gpu_id"] = gpu_id
+    if docker_image is not None:
+        request_body["docker_image"] = docker_image
+    if docker_entrypoint is not None:
+        request_body["docker_entrypoint"] = docker_entrypoint
+    if pull_image:
+        request_body["pull_image"] = True
+    if require_hardware_counters:
+        request_body["require_hardware_counters"] = True
+    if target is not None:
+        request_body["target"] = target
+
+    api_url = get_api_url()
+    headers = _get_auth_headers()
+    exit_code = 0
+
+    with httpx.Client(timeout=None, headers=headers) as client:  # No timeout for streaming
+        with client.stream("POST", f"{api_url}/v1/gpu/jobs", json=request_body) as response:
+            response.raise_for_status()
+
+            for line in response.iter_lines():
+                if not line:
+                    continue
+
+                # Parse SSE format: "data: <content>"
+                if line.startswith("data: "):
+                    content = line[6:]  # Strip "data: " prefix
+
+                    if content == "[DONE]":
+                        break
+                    elif content.startswith("[ERROR]"):
+                        print(content[8:], file=sys.stderr)  # Strip "[ERROR] "
+                        exit_code = 1
+                        break
+                    else:
+                        print(content)
+
+    return exit_code
+
+

wafer/auth.py

@@ -0,0 +1,254 @@
+"""CLI authentication and credential management.
+
+Handles storing/loading credentials from ~/.wafer/credentials.json
+and verifying tokens against the wafer-api.
+"""
+
+import json
+import socket
+import webbrowser
+from dataclasses import dataclass
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from pathlib import Path
+from urllib.parse import parse_qs, urlparse
+
+import httpx
+
+from .api_client import get_api_url
+
+# Default Supabase project URL (can be overridden)
+DEFAULT_SUPABASE_URL = "https://auth.wafer.ai"
+
+CREDENTIALS_DIR = Path.home() / ".wafer"
+CREDENTIALS_FILE = CREDENTIALS_DIR / "credentials.json"
+
+
+@dataclass
+class Credentials:
+    """Stored credentials."""
+
+    access_token: str
+    email: str | None = None
+
+
+@dataclass
+class UserInfo:
+    """User info from token verification."""
+
+    user_id: str
+    email: str | None
+
+
+def save_credentials(token: str, email: str | None = None) -> None:
+    """Save credentials to ~/.wafer/credentials.json."""
+    CREDENTIALS_DIR.mkdir(parents=True, exist_ok=True)
+    data = {"access_token": token}
+    if email:
+        data["email"] = email
+    CREDENTIALS_FILE.write_text(json.dumps(data, indent=2))
+    # Set restrictive permissions (owner read/write only)
+    CREDENTIALS_FILE.chmod(0o600)
+
+
+def load_credentials() -> Credentials | None:
+    """Load credentials from ~/.wafer/credentials.json.
+
+    Returns None if file doesn't exist or is invalid.
+    """
+    if not CREDENTIALS_FILE.exists():
+        return None
+    try:
+        data = json.loads(CREDENTIALS_FILE.read_text())
+        return Credentials(
+            access_token=data["access_token"],
+            email=data.get("email"),
+        )
+    except (json.JSONDecodeError, KeyError):
+        return None
+
+
+def clear_credentials() -> bool:
+    """Remove credentials file.
+
+    Returns True if file was removed, False if it didn't exist.
+    """
+    if CREDENTIALS_FILE.exists():
+        CREDENTIALS_FILE.unlink()
+        return True
+    return False
+
+
+def get_auth_headers() -> dict[str, str]:
+    """Get Authorization headers if credentials exist.
+
+    Returns empty dict if not logged in.
+    """
+    creds = load_credentials()
+    if creds:
+        return {"Authorization": f"Bearer {creds.access_token}"}
+    return {}
+
+
+def verify_token(token: str) -> UserInfo:
+    """Verify token with wafer-api and return user info.
+
+    Raises:
+        httpx.HTTPStatusError: If token is invalid (401) or other HTTP error
+        httpx.RequestError: If API is unreachable
+    """
+    api_url = get_api_url()
+    with httpx.Client(timeout=10.0) as client:
+        response = client.post(
+            f"{api_url}/v1/auth/verify",
+            json={"token": token},
+        )
+        response.raise_for_status()
+        data = response.json()
+        return UserInfo(
+            user_id=data["user_id"],
+            email=data.get("email"),
+        )
+
+
+def _find_free_port() -> int:
+    """Find a free port for the callback server."""
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
+        s.bind(("", 0))
+        return s.getsockname()[1]
+
+
+def get_supabase_url() -> str:
+    """Get Supabase URL from environment or default."""
+    import os
+
+    return os.environ.get("SUPABASE_URL", DEFAULT_SUPABASE_URL)
+
+
+class OAuthCallbackHandler(BaseHTTPRequestHandler):
+    """HTTP handler that catches the OAuth callback with access token."""
+
+    access_token: str | None = None
+    error: str | None = None
+
+    def log_message(self, format: str, *args: object) -> None:
+        """Suppress default logging."""
+        pass
+
+    def do_GET(self) -> None:
+        """Handle GET request - catch the callback or serve the HTML page."""
+        parsed = urlparse(self.path)
+
+        if parsed.path == "/callback":
+            # This is the redirect from Supabase with hash fragment
+            # But hash fragments aren't sent to server, so serve a page that extracts it
+            html = """<!DOCTYPE html>
+<html>
+<head><title>Wafer CLI Login</title></head>
+<body>
+<h2>Completing login...</h2>
+<script>
+// Extract token from hash fragment
+const hash = window.location.hash.substring(1);
+const params = new URLSearchParams(hash);
+const accessToken = params.get('access_token');
+const error = params.get('error_description') || params.get('error');
+
+if (accessToken) {
+    // Send token to our local server
+    fetch('/token?access_token=' + encodeURIComponent(accessToken))
+        .then(() => {
+            document.body.innerHTML = '<h2>✓ Login successful!</h2><p>You can close this window.</p>';
+        });
+} else if (error) {
+    fetch('/token?error=' + encodeURIComponent(error));
+    document.body.innerHTML = '<h2>✗ Login failed</h2><p>' + error + '</p>';
+} else {
+    document.body.innerHTML = '<h2>✗ No token received</h2>';
+}
+</script>
+</body>
+</html>"""
+            self.send_response(200)
+            self.send_header("Content-Type", "text/html")
+            self.end_headers()
+            self.wfile.write(html.encode())
+
+        elif parsed.path == "/token":
+            # JavaScript sends us the token
+            params = parse_qs(parsed.query)
+            if "access_token" in params:
+                OAuthCallbackHandler.access_token = params["access_token"][0]
+            elif "error" in params:
+                OAuthCallbackHandler.error = params["error"][0]
+
+            self.send_response(200)
+            self.send_header("Content-Type", "text/plain")
+            self.end_headers()
+            self.wfile.write(b"OK")
+
+        else:
+            self.send_response(404)
+            self.end_headers()
+
+
+def browser_login(timeout: int = 120) -> str:
+    """Open browser for GitHub OAuth and return access token.
+
+    Starts a local HTTP server, opens browser to Supabase OAuth,
+    and waits for the callback with the access token.
+
+    Args:
+        timeout: Seconds to wait for callback (default 120)
+
+    Returns:
+        Access token string
+
+    Raises:
+        TimeoutError: If no callback received within timeout
+        RuntimeError: If OAuth flow failed
+    """
+    import time
+
+    port = _find_free_port()
+    redirect_uri = f"http://localhost:{port}/callback"
+    supabase_url = get_supabase_url()
+
+    # Build OAuth URL
+    auth_url = (
+        f"{supabase_url}/auth/v1/authorize"
+        f"?provider=github"
+        f"&redirect_to={redirect_uri}"
+    )
+
+    # Reset state
+    OAuthCallbackHandler.access_token = None
+    OAuthCallbackHandler.error = None
+
+    # Start local server
+    server = HTTPServer(("localhost", port), OAuthCallbackHandler)
+    server.timeout = 1  # Check for token every second
+
+    # Open browser
+    print("Opening browser for GitHub authentication...")
+    print(f"If browser doesn't open, visit: {auth_url}")
+    webbrowser.open(auth_url)
+
+    # Wait for callback
+    start = time.time()
+    print("Waiting for authentication...", end="", flush=True)
+
+    while time.time() - start < timeout:
+        server.handle_request()
+
+        if OAuthCallbackHandler.access_token:
+            print(" ✓")
+            server.server_close()
+            return OAuthCallbackHandler.access_token
+
+        if OAuthCallbackHandler.error:
+            print(" ✗")
+            server.server_close()
+            raise RuntimeError(f"OAuth failed: {OAuthCallbackHandler.error}")
+
+    server.server_close()
+    raise TimeoutError(f"No response within {timeout} seconds")