wafer-cli 0.2.14__py3-none-any.whl

wafer/global_config.py

@@ -0,0 +1,350 @@
+"""Global configuration for Wafer CLI.
+
+Handles API URLs, environment presets (staging/prod/local), and user preferences.
+Separate from config.py which handles Docker execution config for `wafer run`.
+"""
+
+import os
+import tomllib
+from dataclasses import dataclass, field
+from pathlib import Path
+from typing import Literal
+
+# Config file location
+CONFIG_DIR = Path.home() / ".wafer"
+CONFIG_FILE = CONFIG_DIR / "config.toml"
+
+# Environment type
+EnvironmentName = Literal["staging", "prod", "local"]
+
+
+@dataclass(frozen=True)
+class ApiEnvironment:
+    """API environment configuration. Immutable."""
+
+    name: str
+    api_url: str
+    supabase_url: str
+    supabase_anon_key: str
+
+    def __post_init__(self) -> None:
+        assert self.name, "environment name cannot be empty"
+        assert self.api_url, "api_url cannot be empty"
+        assert self.supabase_url, "supabase_url cannot be empty"
+        assert self.supabase_anon_key, "supabase_anon_key cannot be empty"
+
+
+# Built-in environment presets
+# Anon keys are public (used client-side) - safe to embed
+BUILTIN_ENVIRONMENTS: dict[str, ApiEnvironment] = {
+    "staging": ApiEnvironment(
+        name="staging",
+        api_url="https://wafer-api-staging.onrender.com",
+        supabase_url="https://xudshwhzytyfxwwyofli.supabase.co",
+        supabase_anon_key="eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6Inh1ZHNod2h6eXR5Znh3d3lvZmxpIiwicm9sZSI6ImFub24iLCJpYXQiOjE3NjU4Mzc5MDEsImV4cCI6MjA4MTQxMzkwMX0.JvuF4349z1ermKmrxEKGDHQj9I_ylLZYjjuouJleYhY",
+    ),
+    "prod": ApiEnvironment(
+        name="prod",
+        api_url="https://www.api.wafer.ai",
+        supabase_url="https://auth.wafer.ai",
+        supabase_anon_key="eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZSIsInJlZiI6Imh2bHB0aGNueGx5d2xxdWljaXFlIiwicm9sZSI6ImFub24iLCJpYXQiOjE3NjQ2MjQ1NTIsImV4cCI6MjA4MDIwMDU1Mn0.1ywPDp-QHgbqPOJocQvXEKKDjGt3BsoNluvVoQ7EW3o",
+    ),
+    "local": ApiEnvironment(
+        name="local",
+        api_url="http://localhost:8000",
+        supabase_url="http://localhost:54321",
+        supabase_anon_key="eyJhbGciOiJIUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJzdXBhYmFzZS1kZW1vIiwicm9sZSI6ImFub24iLCJleHAiOjE5ODM4MTI5OTZ9.CRXP1A7WOeoJeXxjNni43kdQwgnWNReilDMblYTn_I0",
+    ),
+}
+
+# Default environment when no config exists
+DEFAULT_ENVIRONMENT = "prod"
+
+
+@dataclass(frozen=True)
+class Preferences:
+    """User preferences. Immutable.
+
+    mode: "implicit" (default) = quiet output, use -v for status messages
+          "explicit" = verbose output, shows [wafer] status messages
+    analytics_enabled: True (default) = send anonymous usage analytics to PostHog
+                       False = disable all analytics tracking
+    """
+
+    mode: Literal["explicit", "implicit"] = "implicit"
+    analytics_enabled: bool = True
+
+
+@dataclass(frozen=True)
+class Defaults:
+    """Default values for commands. Immutable."""
+
+    workspace: str | None = None
+    gpu: str = "H100"
+    exec_timeout: int = 300  # seconds
+
+
+@dataclass(frozen=True)
+class GlobalConfig:
+    """Global Wafer CLI configuration. Immutable."""
+
+    environment: str = DEFAULT_ENVIRONMENT
+    environments: dict[str, ApiEnvironment] = field(
+        default_factory=lambda: BUILTIN_ENVIRONMENTS.copy()
+    )
+    preferences: Preferences = field(default_factory=Preferences)
+    defaults: Defaults = field(default_factory=Defaults)
+
+    def __post_init__(self) -> None:
+        # Validate environment exists
+        assert self.environment in self.environments, (
+            f"environment '{self.environment}' not found. "
+            f"Available: {list(self.environments.keys())}"
+        )
+
+    def get_api_environment(self) -> ApiEnvironment:
+        """Get the current API environment."""
+        return self.environments[self.environment]
+
+    @property
+    def api_url(self) -> str:
+        """Get current API URL."""
+        return self.get_api_environment().api_url
+
+    @property
+    def supabase_url(self) -> str:
+        """Get current Supabase URL."""
+        return self.get_api_environment().supabase_url
+
+
+def _parse_config_file(path: Path) -> GlobalConfig:
+    """Parse config from TOML file.
+
+    Args:
+        path: Path to config file
+
+    Returns:
+        GlobalConfig instance
+    """
+    with open(path, "rb") as f:
+        data = tomllib.load(f)
+
+    # Parse API section
+    api_data = data.get("api", {})
+    environment = api_data.get("environment", DEFAULT_ENVIRONMENT)
+
+    # Merge built-in environments with user-defined ones
+    environments = BUILTIN_ENVIRONMENTS.copy()
+    user_envs = api_data.get("environments", {})
+    for name, env_config in user_envs.items():
+        if isinstance(env_config, dict):
+            # User can override built-in or define new
+            base_env = environments.get(name, BUILTIN_ENVIRONMENTS["prod"])
+            environments[name] = ApiEnvironment(
+                name=name,
+                api_url=env_config.get("url", base_env.api_url),
+                supabase_url=env_config.get("supabase_url", base_env.supabase_url),
+                supabase_anon_key=env_config.get("supabase_anon_key", base_env.supabase_anon_key),
+            )
+
+    # Parse preferences
+    pref_data = data.get("preferences", {})
+    mode = pref_data.get("mode", "explicit")
+    assert mode in ("explicit", "implicit"), f"mode must be 'explicit' or 'implicit', got '{mode}'"
+    analytics_enabled = pref_data.get("analytics_enabled", True)
+    assert isinstance(analytics_enabled, bool), f"analytics_enabled must be true or false, got '{analytics_enabled}'"
+    preferences = Preferences(mode=mode, analytics_enabled=analytics_enabled)
+
+    # Parse defaults
+    defaults_data = data.get("defaults", {})
+    defaults = Defaults(
+        workspace=defaults_data.get("workspace"),
+        gpu=defaults_data.get("gpu", "H100"),
+        exec_timeout=defaults_data.get("exec_timeout", 300),
+    )
+
+    return GlobalConfig(
+        environment=environment,
+        environments=environments,
+        preferences=preferences,
+        defaults=defaults,
+    )
+
+
+# Cached config instance
+_cached_config: GlobalConfig | None = None
+
+
+def load_global_config() -> GlobalConfig:
+    """Load global config from file, with env var overrides.
+
+    Priority (highest to lowest):
+    1. Environment variables (WAFER_API_URL, SUPABASE_URL)
+    2. Config file (~/.wafer/config.toml)
+    3. Built-in defaults (prod environment)
+
+    Returns:
+        GlobalConfig instance
+    """
+    global _cached_config
+
+    if _cached_config is not None:
+        return _cached_config
+
+    # Start with defaults
+    if CONFIG_FILE.exists():
+        config = _parse_config_file(CONFIG_FILE)
+    else:
+        config = GlobalConfig()
+
+    _cached_config = config
+    return config
+
+
+def get_api_url() -> str:
+    """Get API URL with env var override.
+
+    Priority:
+    1. WAFER_API_URL env var
+    2. Config file
+    3. Default (prod)
+    """
+    env_url = os.environ.get("WAFER_API_URL")
+    if env_url:
+        return env_url
+    return load_global_config().api_url
+
+
+def get_supabase_url() -> str:
+    """Get Supabase URL with env var override.
+
+    Priority:
+    1. SUPABASE_URL env var
+    2. Config file
+    3. Default (prod)
+    """
+    env_url = os.environ.get("SUPABASE_URL")
+    if env_url:
+        return env_url
+    return load_global_config().supabase_url
+
+
+def get_supabase_anon_key() -> str:
+    """Get Supabase anon key for current environment.
+
+    The anon key is public and used for client-side auth operations
+    like token refresh.
+    """
+    return load_global_config().get_api_environment().supabase_anon_key
+
+
+def get_preferences() -> Preferences:
+    """Get user preferences."""
+    return load_global_config().preferences
+
+
+def get_defaults() -> Defaults:
+    """Get default values."""
+    return load_global_config().defaults
+
+
+def clear_config_cache() -> None:
+    """Clear cached config. Useful after config changes."""
+    global _cached_config
+    _cached_config = None
+
+
+def save_global_config(config: GlobalConfig) -> None:
+    """Save config to TOML file, preserving existing Docker config sections.
+
+    Merges the global API config with any existing [default] and [environments.*]
+    sections that are used by the Docker execution config (WaferConfig).
+    """
+    CONFIG_DIR.mkdir(parents=True, exist_ok=True)
+
+    # Load existing config to preserve Docker sections
+    existing_data: dict = {}
+    if CONFIG_FILE.exists():
+        with open(CONFIG_FILE, "rb") as f:
+            existing_data = tomllib.load(f)
+
+    # Build new content, preserving Docker config sections
+    lines = []
+
+    # Preserve [default] section (Docker config)
+    if "default" in existing_data:
+        lines.append("[default]")
+        for key, value in existing_data["default"].items():
+            if isinstance(value, str):
+                lines.append(f'{key} = "{value}"')
+            else:
+                lines.append(f"{key} = {value}")
+        lines.append("")
+
+    # API section (global config)
+    lines.append("[api]")
+    lines.append(f'environment = "{config.environment}"')
+    lines.append("")
+
+    # Only write non-builtin API environments
+    for name, env in config.environments.items():
+        if name not in BUILTIN_ENVIRONMENTS or env != BUILTIN_ENVIRONMENTS[name]:
+            lines.append(f"[api.environments.{name}]")
+            lines.append(f'url = "{env.api_url}"')
+            lines.append(f'supabase_url = "{env.supabase_url}"')
+            lines.append("")
+
+    # Preserve [environments.*] sections (Docker config)
+    if "environments" in existing_data:
+        for name, env_config in existing_data["environments"].items():
+            if isinstance(env_config, dict):
+                lines.append(f"[environments.{name}]")
+                for key, value in env_config.items():
+                    if isinstance(value, str):
+                        lines.append(f'{key} = "{value}"')
+                    else:
+                        lines.append(f"{key} = {value}")
+                lines.append("")
+
+    # Preferences section (only if non-default values)
+    pref_lines = []
+    if config.preferences.mode != "implicit":
+        pref_lines.append(f'mode = "{config.preferences.mode}"')
+    if not config.preferences.analytics_enabled:
+        pref_lines.append("analytics_enabled = false")
+
+    if pref_lines:
+        lines.append("[preferences]")
+        lines.extend(pref_lines)
+        lines.append("")
+
+    # Defaults section (only if non-default values)
+    defaults_lines = []
+    if config.defaults.workspace:
+        defaults_lines.append(f'workspace = "{config.defaults.workspace}"')
+    if config.defaults.gpu != "H100":
+        defaults_lines.append(f'gpu = "{config.defaults.gpu}"')
+    if config.defaults.exec_timeout != 300:
+        defaults_lines.append(f"exec_timeout = {config.defaults.exec_timeout}")
+
+    if defaults_lines:
+        lines.append("[defaults]")
+        lines.extend(defaults_lines)
+        lines.append("")
+
+    CONFIG_FILE.write_text("\n".join(lines))
+    clear_config_cache()
+
+
+def init_config(environment: str = DEFAULT_ENVIRONMENT) -> GlobalConfig:
+    """Initialize config file with defaults.
+
+    Args:
+        environment: Initial environment to use
+
+    Returns:
+        The created GlobalConfig
+    """
+    config = GlobalConfig(environment=environment)
+    save_global_config(config)
+    return config

wafer/gpu_run.py

@@ -0,0 +1,307 @@
+"""Remote execution on GPU targets.
+
+Provides push/run primitives for remote GPU execution via SSH and Docker.
+"""
+
+import sys
+from dataclasses import dataclass
+from pathlib import Path
+
+from wafer_core.ssh import SSHClient
+from wafer_core.utils.kernel_utils.targets.config import BaremetalTarget, VMTarget
+
+# Constants
+REMOTE_WORKSPACE_BASE = "~/.wafer/workspaces"
+CONTAINER_WORKSPACE = "/workspace"
+
+
+@dataclass(frozen=True)
+class PushResult:
+    """Result of pushing a directory to remote target."""
+
+    workspace_name: str  # Just the workspace name (e.g., "project")
+    workspace_path: (
+        str  # Full absolute path on remote (e.g., "/home/user/.wafer/workspaces/project")
+    )
+    files_uploaded: list[str]  # Relative paths of uploaded files
+
+
+def push_directory(
+    local_path: Path,
+    target: BaremetalTarget | VMTarget,
+) -> PushResult:
+    """Push local directory to remote target.
+
+    Uploads directory to ~/.wafer/workspaces/<dirname> on target.
+
+    Args:
+        local_path: Local directory to upload
+        target: Remote target configuration
+
+    Returns:
+        PushResult with workspace path and list of uploaded files
+
+    Raises:
+        FileNotFoundError: If local_path doesn't exist
+        ValueError: If local_path is not a directory
+    """
+    # Validate inputs
+    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}")
+    if target.ssh_target is None:
+        raise ValueError(f"Target '{target.name}' must have ssh_target configured")
+    if target.ssh_key is None:
+        raise ValueError(f"Target '{target.name}' must have ssh_key configured")
+
+    client = SSHClient(target.ssh_target, target.ssh_key)
+
+    workspace_name = local_path.name
+    remote_workspace = f"{REMOTE_WORKSPACE_BASE}/{workspace_name}"
+
+    # Create workspace directory
+    client.exec(f"mkdir -p {remote_workspace}")
+    expanded_workspace = client.expand_path(remote_workspace)
+
+    # Upload directory recursively
+    client.upload_files(str(local_path), expanded_workspace, recursive=True)
+
+    # Get list of uploaded files (relative paths)
+    files_uploaded = []
+    for file in local_path.rglob("*"):
+        if file.is_file():
+            files_uploaded.append(str(file.relative_to(local_path)))
+
+    return PushResult(
+        workspace_name=workspace_name,
+        workspace_path=expanded_workspace,
+        files_uploaded=files_uploaded,
+    )
+
+
+def run_command(
+    command: str,
+    workspace: str,
+    target: BaremetalTarget | VMTarget,
+    gpu_id: int | None = None,
+) -> int:
+    """Run command in Docker on remote target, streaming output.
+
+    Args:
+        command: Command to execute inside container
+        workspace: Workspace name (subdirectory under ~/.wafer/workspaces/)
+        target: Remote target configuration (must have docker_image)
+        gpu_id: GPU ID to use (defaults to first in target.gpu_ids)
+
+    Returns:
+        Exit code from command (0 = success)
+
+    Raises:
+        ValueError: If target is missing required configuration
+    """
+    if target.docker_image is None:
+        raise ValueError(f"Target '{target.name}' must have docker_image configured")
+    if target.ssh_target is None:
+        raise ValueError(f"Target '{target.name}' must have ssh_target configured")
+    if target.ssh_key is None:
+        raise ValueError(f"Target '{target.name}' must have ssh_key configured")
+
+    client = SSHClient(target.ssh_target, target.ssh_key)
+
+    effective_gpu_id = gpu_id if gpu_id is not None else target.gpu_ids[0]
+
+    # Get expanded workspace path
+    remote_workspace = f"{REMOTE_WORKSPACE_BASE}/{workspace}"
+    expanded_workspace = client.expand_path(remote_workspace)
+
+    # Build docker command with workspace mounted
+    volumes = {expanded_workspace: CONTAINER_WORKSPACE}
+    docker_cmd = _build_docker_command(
+        image=target.docker_image,
+        inner_cmd=command,
+        gpu_id=effective_gpu_id,
+        volumes=volumes,
+    )
+
+    # Stream execution
+    exit_code = 0
+    try:
+        for line in client.exec_stream(docker_cmd):
+            print(line)
+    except Exception as e:
+        print(f"\nExecution failed: {e}", file=sys.stderr)
+        exit_code = 1
+
+    return exit_code
+
+
+def run_command_capture(
+    command: str,
+    workspace: str,
+    target: BaremetalTarget | VMTarget,
+) -> tuple[int, str]:
+    """Run command on remote target (without Docker) and capture output.
+
+    This is useful for commands that don't need GPU access, like running
+    NCU to analyze a profile file.
+
+    Args:
+        command: Command to execute on the remote host
+        workspace: Workspace name (subdirectory under ~/.wafer/workspaces/)
+        target: Remote target configuration
+
+    Returns:
+        Tuple of (exit_code, output_text)
+
+    Raises:
+        ValueError: If target is missing required configuration
+    """
+    if target.ssh_target is None:
+        raise ValueError(f"Target '{target.name}' must have ssh_target configured")
+    if target.ssh_key is None:
+        raise ValueError(f"Target '{target.name}' must have ssh_key configured")
+
+    client = SSHClient(target.ssh_target, target.ssh_key)
+
+    # Get expanded workspace path
+    remote_workspace = f"{REMOTE_WORKSPACE_BASE}/{workspace}"
+    expanded_workspace = client.expand_path(remote_workspace)
+
+    # Run command in workspace directory
+    full_cmd = f"cd {expanded_workspace} && {command}"
+
+    # Capture output
+    output_lines = []
+    exit_code = 0
+    try:
+        for line in client.exec_stream(full_cmd):
+            output_lines.append(line)
+    except Exception as e:
+        print(f"\nExecution failed: {e}", file=sys.stderr)
+        exit_code = 1
+
+    return exit_code, "\n".join(output_lines)
+
+
+def _build_docker_command(
+    image: str,
+    inner_cmd: str,
+    gpu_id: int,
+    volumes: dict[str, str],
+) -> str:
+    """Build docker run command string."""
+    import shlex
+
+    parts = ["docker", "run", "--rm"]
+    parts.extend(["--gpus", f"'device={gpu_id}'"])
+
+    for host_path, container_path in volumes.items():
+        parts.extend(["-v", f"{host_path}:{container_path}"])
+
+    parts.extend(["-w", CONTAINER_WORKSPACE])
+    parts.append(image)
+    parts.append(f"bash -c {shlex.quote(inner_cmd)}")
+
+    return " ".join(parts)
+
+
+def _build_uv_install_cmd() -> str:
+    """Build command to ensure uv is available."""
+    ensure_curl = (
+        "which curl > /dev/null 2>&1 || "
+        "(apt-get update -qq && apt-get install -qq -y curl > /dev/null)"
+    )
+    install_uv = "which uv > /dev/null 2>&1 || curl -LsSf https://astral.sh/uv/install.sh | sh"
+    source_uv = "export PATH=$HOME/.local/bin:$PATH"
+    return f"{ensure_curl} && {install_uv} && {source_uv}"
+
+
+def run_python_docker(
+    file_path: Path,
+    args: list[str],
+    target: BaremetalTarget | VMTarget,
+    gpu_id: int,
+) -> int:
+    """Run Python file in Docker container on remote GPU."""
+    if target.docker_image is None:
+        raise ValueError(f"Target '{target.name}' must have docker_image configured")
+    if target.ssh_target is None:
+        raise ValueError(f"Target '{target.name}' must have ssh_target configured")
+    if target.ssh_key is None:
+        raise ValueError(f"Target '{target.name}' must have ssh_key configured")
+
+    print(f"Connecting to {target.ssh_target}...")
+    client = SSHClient(target.ssh_target, target.ssh_key)
+
+    # Setup workspace
+    remote_workspace = f"{REMOTE_WORKSPACE_BASE}/python_run"
+    client.exec(f"mkdir -p {remote_workspace}")
+    expanded_workspace = client.expand_path(remote_workspace)
+
+    # Upload file
+    remote_file = f"{expanded_workspace}/{file_path.name}"
+    print(f"Uploading {file_path.name}...")
+    client.upload_files(str(file_path), remote_file)
+
+    # Build inner command: install uv, run script with inline deps
+    script_args = " ".join(args) if args else ""
+    uv_setup = _build_uv_install_cmd()
+    inner_cmd = f"{uv_setup} && uv run --script {file_path.name} {script_args}"
+
+    # Build docker command
+    volumes = {expanded_workspace: CONTAINER_WORKSPACE}
+    docker_cmd = _build_docker_command(
+        image=target.docker_image,
+        inner_cmd=inner_cmd,
+        gpu_id=gpu_id,
+        volumes=volumes,
+    )
+
+    print(f"Running on GPU {gpu_id} with {target.docker_image}...")
+    print("-" * 60)
+
+    # Stream execution
+    exit_code = 0
+    try:
+        for line in client.exec_stream(docker_cmd):
+            print(line)
+    except Exception as e:
+        print(f"\nExecution failed: {e}", file=sys.stderr)
+        exit_code = 1
+
+    return exit_code
+
+
+def run_python_file(
+    file_path: Path,
+    args: list[str],
+    target: BaremetalTarget | VMTarget,
+    gpu_id: int | None = None,
+) -> int:
+    """Run Python file on remote GPU in Docker container.
+
+    Args:
+        file_path: Path to Python script
+        args: Arguments to pass to script
+        target: Remote target configuration (must have docker_image)
+        gpu_id: GPU ID to use (defaults to first in target.gpu_ids)
+
+    Returns:
+        Exit code from script
+
+    Raises:
+        FileNotFoundError: If file doesn't exist
+        ValueError: If path is not a file or target has no docker_image
+    """
+    # Validate inputs
+    if not file_path.exists():
+        raise FileNotFoundError(f"File not found: {file_path}")
+    if not file_path.is_file():
+        raise ValueError(f"Not a file: {file_path}")
+    if not target.docker_image:
+        raise ValueError(f"Target '{target.name}' has no docker_image configured")
+
+    effective_gpu_id = gpu_id if gpu_id is not None else target.gpu_ids[0]
+
+    return run_python_docker(file_path, args, target, effective_gpu_id)