wafer-cli 0.2.4__py3-none-any.whl → 0.2.5__py3-none-any.whl

wafer/target_lock.py

@@ -0,0 +1,198 @@
+"""Target locking for concurrent access control.
+
+Uses file locks (fcntl.flock) to ensure only one process uses a target at a time.
+Locks are automatically released when the process exits or crashes.
+
+Usage:
+    # Try to acquire a single target
+    with try_acquire_target("mi300x-1") as acquired:
+        if acquired:
+            # Got the lock, run eval
+            ...
+        else:
+            # Target busy
+            ...
+
+    # Acquire first available from a pool
+    with acquire_from_pool(["mi300x-1", "mi300x-2", "mi300x-3"]) as target:
+        if target:
+            # Got a target, run eval
+            ...
+        else:
+            # All targets busy
+            ...
+"""
+
+from __future__ import annotations
+
+import fcntl
+import os
+from contextlib import contextmanager
+from pathlib import Path
+from typing import Iterator
+
+# Lock directory
+LOCKS_DIR = Path.home() / ".wafer" / "locks"
+
+
+def _ensure_locks_dir() -> None:
+    """Ensure locks directory exists."""
+    LOCKS_DIR.mkdir(parents=True, exist_ok=True)
+
+
+def _lock_path(target_name: str) -> Path:
+    """Get path to lock file for a target."""
+    return LOCKS_DIR / f"{target_name}.lock"
+
+
+@contextmanager
+def try_acquire_target(target_name: str) -> Iterator[bool]:
+    """Try to acquire exclusive lock on a target.
+
+    Args:
+        target_name: Name of the target to lock
+
+    Yields:
+        True if lock was acquired, False if target is busy
+
+    The lock is automatically released when the context exits,
+    or if the process crashes.
+    """
+    _ensure_locks_dir()
+    lock_file = _lock_path(target_name)
+
+    # Open or create lock file
+    fd = os.open(str(lock_file), os.O_CREAT | os.O_RDWR)
+
+    try:
+        # Try non-blocking exclusive lock
+        fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        # Write PID to lock file for debugging
+        os.ftruncate(fd, 0)
+        os.write(fd, f"{os.getpid()}\n".encode())
+        try:
+            yield True
+        finally:
+            # Release lock
+            fcntl.flock(fd, fcntl.LOCK_UN)
+    except BlockingIOError:
+        # Lock is held by another process
+        yield False
+    finally:
+        os.close(fd)
+
+
+@contextmanager
+def acquire_from_pool(target_names: list[str]) -> Iterator[str | None]:
+    """Acquire first available target from a list.
+
+    Tries each target in order, returns the first one that's available.
+
+    Args:
+        target_names: List of target names to try
+
+    Yields:
+        Name of acquired target, or None if all are busy
+
+    Example:
+        with acquire_from_pool(["gpu-1", "gpu-2", "gpu-3"]) as target:
+            if target:
+                print(f"Got {target}")
+                run_eval(target)
+            else:
+                print("All targets busy")
+    """
+    _ensure_locks_dir()
+
+    # Try each target in order
+    for target_name in target_names:
+        lock_file = _lock_path(target_name)
+        fd = os.open(str(lock_file), os.O_CREAT | os.O_RDWR)
+
+        try:
+            fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+            # Got the lock - write PID and yield
+            os.ftruncate(fd, 0)
+            os.write(fd, f"{os.getpid()}\n".encode())
+            try:
+                yield target_name
+                return  # Success - exit after context
+            finally:
+                fcntl.flock(fd, fcntl.LOCK_UN)
+                os.close(fd)
+        except BlockingIOError:
+            # This target is busy, try next
+            os.close(fd)
+            continue
+
+    # All targets busy
+    yield None
+
+
+def is_target_locked(target_name: str) -> bool:
+    """Check if a target is currently locked.
+
+    Note: This is a point-in-time check - the lock status can change
+    immediately after this returns.
+
+    Args:
+        target_name: Name of the target to check
+
+    Returns:
+        True if target is locked, False if available
+    """
+    _ensure_locks_dir()
+    lock_file = _lock_path(target_name)
+
+    if not lock_file.exists():
+        return False
+
+    fd = os.open(str(lock_file), os.O_RDONLY)
+    try:
+        # Try non-blocking lock
+        fcntl.flock(fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        # Got it - so it wasn't locked
+        fcntl.flock(fd, fcntl.LOCK_UN)
+        return False
+    except BlockingIOError:
+        return True
+    finally:
+        os.close(fd)
+
+
+def get_lock_holder(target_name: str) -> int | None:
+    """Get PID of process holding lock on a target.
+
+    Args:
+        target_name: Name of the target
+
+    Returns:
+        PID of lock holder, or None if not locked or unknown
+    """
+    lock_file = _lock_path(target_name)
+
+    if not lock_file.exists():
+        return None
+
+    try:
+        content = lock_file.read_text().strip()
+        return int(content)
+    except (ValueError, OSError):
+        return None
+
+
+def list_locked_targets() -> list[str]:
+    """List all currently locked targets.
+
+    Returns:
+        List of target names that are currently locked
+    """
+    _ensure_locks_dir()
+
+    locked = []
+    for lock_file in LOCKS_DIR.glob("*.lock"):
+        target_name = lock_file.stem
+        if is_target_locked(target_name):
+            locked.append(target_name)
+
+    return sorted(locked)

wafer/targets.py

@@ -257,6 +257,164 @@ def get_default_target() -> str | None:
     return data.get("default_target")
 
 
+# ── Pool Management ─────────────────────────────────────────────────────────
+
+
+def get_pool(name: str) -> list[str]:
+    """Get list of targets in a named pool.
+
+    Pools are defined in ~/.wafer/config.toml:
+        [pools.my-pool]
+        targets = ["target-1", "target-2", "target-3"]
+
+    Args:
+        name: Pool name
+
+    Returns:
+        List of target names in the pool
+
+    Raises:
+        FileNotFoundError: If pool doesn't exist
+    """
+    if not CONFIG_FILE.exists():
+        raise FileNotFoundError(f"Pool not found: {name} (no config file)")
+
+    with open(CONFIG_FILE, "rb") as f:
+        data = tomllib.load(f)
+
+    pools = data.get("pools", {})
+    if name not in pools:
+        raise FileNotFoundError(
+            f"Pool not found: {name}\n"
+            f"  Define pools in ~/.wafer/config.toml:\n"
+            f"  [pools.{name}]\n"
+            f'  targets = ["target-1", "target-2"]'
+        )
+
+    pool_config = pools[name]
+    targets = pool_config.get("targets", [])
+
+    if not targets:
+        raise ValueError(f"Pool '{name}' has no targets defined")
+
+    return targets
+
+
+def list_pools() -> list[str]:
+    """List all configured pool names.
+
+    Returns:
+        Sorted list of pool names
+    """
+    if not CONFIG_FILE.exists():
+        return []
+
+    with open(CONFIG_FILE, "rb") as f:
+        data = tomllib.load(f)
+
+    return sorted(data.get("pools", {}).keys())
+
+
+def save_pool(name: str, targets: list[str]) -> None:
+    """Save or update a pool configuration.
+
+    Args:
+        name: Pool name
+        targets: List of target names (must all exist)
+
+    Raises:
+        FileNotFoundError: If any target doesn't exist
+    """
+    # Verify all targets exist
+    existing_targets = list_targets()
+    missing = [t for t in targets if t not in existing_targets]
+    if missing:
+        raise FileNotFoundError(f"Targets not found: {', '.join(missing)}")
+
+    _ensure_dirs()
+
+    # Load existing config
+    if CONFIG_FILE.exists():
+        with open(CONFIG_FILE, "rb") as f:
+            data = tomllib.load(f)
+    else:
+        data = {}
+
+    # Update pools section
+    if "pools" not in data:
+        data["pools"] = {}
+
+    data["pools"][name] = {"targets": targets}
+
+    # Write back - need custom handling for nested structure
+    _write_config_with_pools(data)
+
+
+def _write_config_with_pools(data: dict) -> None:
+    """Write config file with pools support.
+
+    Handles the nested [pools.name] TOML structure and preserves
+    existing nested sections like [default], [api], [environments.*].
+    """
+    lines = []
+
+    # Collect nested sections to write after top-level keys
+    nested_sections: dict[str, dict] = {}
+
+    # Write top-level keys first (except pools and nested dicts)
+    for key, value in data.items():
+        if key == "pools":
+            continue
+        if value is None:
+            continue
+        if isinstance(value, dict):
+            # Save nested sections for later
+            nested_sections[key] = value
+        elif isinstance(value, str):
+            lines.append(f'{key} = "{value}"')
+        elif isinstance(value, bool):
+            lines.append(f"{key} = {str(value).lower()}")
+        elif isinstance(value, int | float):
+            lines.append(f"{key} = {value}")
+        elif isinstance(value, list):
+            if all(isinstance(v, int) for v in value):
+                lines.append(f"{key} = {value}")
+            else:
+                formatted = ", ".join(f'"{v}"' if isinstance(v, str) else str(v) for v in value)
+                lines.append(f"{key} = [{formatted}]")
+
+    # Write nested sections (e.g., [default], [api], [environments.foo])
+    for section_name, section_data in nested_sections.items():
+        lines.append("")
+        lines.append(f"[{section_name}]")
+        for key, value in section_data.items():
+            if value is None:
+                continue
+            if isinstance(value, str):
+                lines.append(f'{key} = "{value}"')
+            elif isinstance(value, bool):
+                lines.append(f"{key} = {str(value).lower()}")
+            elif isinstance(value, int | float):
+                lines.append(f"{key} = {value}")
+            elif isinstance(value, list):
+                if all(isinstance(v, int) for v in value):
+                    lines.append(f"{key} = {value}")
+                else:
+                    formatted = ", ".join(f'"{v}"' if isinstance(v, str) else str(v) for v in value)
+                    lines.append(f"{key} = [{formatted}]")
+
+    # Write pools
+    pools = data.get("pools", {})
+    for pool_name, pool_config in pools.items():
+        lines.append("")
+        lines.append(f"[pools.{pool_name}]")
+        targets = pool_config.get("targets", [])
+        formatted = ", ".join(f'"{t}"' for t in targets)
+        lines.append(f"targets = [{formatted}]")
+
+    CONFIG_FILE.write_text("\n".join(lines) + "\n")
+
+
 def set_default_target(name: str) -> None:
     """Set default target.
 

{wafer_cli-0.2.4.dist-info → wafer_cli-0.2.5.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: wafer-cli
-Version: 0.2.4
+Version: 0.2.5
 Summary: CLI tool for running commands on remote GPUs and GPU kernel optimization agent
 Requires-Python: >=3.11
 Requires-Dist: typer>=0.12.0

{wafer_cli-0.2.4.dist-info → wafer_cli-0.2.5.dist-info}/RECORD

@@ -5,10 +5,10 @@ wafer/api_client.py,sha256=cPULiTxqOAYYSfDTNJgd-6Pqrt3IM4Gm9903U7yGIwY,6163
 wafer/auth.py,sha256=ZLsXZ73GDLD8GL7Rij1ELtuLqyJ5EU_uPBUMPVKwExA,10703
 wafer/autotuner.py,sha256=6gH0Ho7T58EFerMQcHQxshWe3DF4qU7fb5xthAh5SPM,44364
 wafer/billing.py,sha256=jbLB2lI4_9f2KD8uEFDi_ixLlowe5hasC0TIZJyIXRg,7163
-wafer/cli.py,sha256=19BwYcZo6QfJWfGwkSXNFx__63ZjZM22woL66fNQhIM,200512
+wafer/cli.py,sha256=rdR84w5ubXO1W2xnrURTrX3AtXy3VeUFVekeGw0djyA,211114
 wafer/config.py,sha256=h5Eo9_yfWqWGoPNdVQikI9GoZVUeysunSYiixf1mKcw,3411
 wafer/corpus.py,sha256=yTF3UA5bOa8BII2fmcXf-3WsIsM5DX4etysv0AzVknE,8912
-wafer/evaluate.py,sha256=sGtEkLmZ8wEeC2BRYMoCbCPIlQvuv6tG77910gbIygc,135553
+wafer/evaluate.py,sha256=D_0WqIw1HysH7XXiyjxRpv6BUuOoPljN9-1vjPt4xFo,166765
 wafer/global_config.py,sha256=fhaR_RU3ufMksDmOohH1OLeQ0JT0SDW1hEip_zaP75k,11345
 wafer/gpu_run.py,sha256=TwqXy72T7f2I7e6n5WWod3xgxCPnDhU0BgLsB4CUoQY,9716
 wafer/inference.py,sha256=tZCO5i05FKY27ewis3CSBHFBeFbXY3xwj0DSjdoMY9s,4314
@@ -18,7 +18,8 @@ wafer/problems.py,sha256=ce2sy10A1nnNUG3VGsseTS8jL7LZsku4dE8zVf9JHQ4,11296
 wafer/rocprof_compute.py,sha256=Tu16Vb05b2grvheFWi1XLGlAr6m48NEDeZoDyw_4Uzw,19885
 wafer/rocprof_sdk.py,sha256=fAYCxpfJa5BZTTkIMBOXg4KsYK4i_wNOKrJJn1ZfypM,10086
 wafer/rocprof_systems.py,sha256=4IWbMcbYk1x_8iS7P3FC_u5sgH6EXADCtR2lV9id80M,18629
-wafer/targets.py,sha256=WE5TJgFPGtEIh7VaTQHZ4wB2t4kW0c5K8-UmQ_39Ock,10254
+wafer/target_lock.py,sha256=QW0NMlu9Paa28O5iSAvGtN11j3kU2di0lADBrfwr2js,5160
+wafer/targets.py,sha256=JlLvi18IHtOkgtBdkv_nUrzBweVmFoOQH-9tQW5s1yQ,15250
 wafer/tracelens.py,sha256=g9ZIeFyNojZn4uTd3skPqIrRiL7aMJOz_-GOd3aiyy4,7998
 wafer/wevin_cli.py,sha256=1_o2P47namZmPkbt47TnyYDmwhEzQYbSg5zjHffu2JQ,16802
 wafer/workspaces.py,sha256=92LG1mtkzNz-ap3XzcqY6KnQ9SUCFG8VBIOUj1Who64,25757
@@ -27,8 +28,8 @@ wafer/templates/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 wafer/templates/ask_docs.py,sha256=Lxs-faz9v5m4Qa4NjF2X_lE8KwM9ES9MNJkxo7ep56o,2256
 wafer/templates/optimize_kernel.py,sha256=u6AL7Q3uttqlnBLzcoFdsiPq5lV2TV3bgqwCYYlK9gk,2357
 wafer/templates/trace_analyze.py,sha256=XE1VqzVkIUsZbXF8EzQdDYgg-AZEYAOFpr6B_vnRELc,2880
-wafer_cli-0.2.4.dist-info/METADATA,sha256=OqNyd57hfOUZvBWFd8iMdrfH1VICqcqQf_K1sHZVhfA,559
-wafer_cli-0.2.4.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
-wafer_cli-0.2.4.dist-info/entry_points.txt,sha256=WqB7hB__WhtPY8y1cO2sZiUz7fCq6Ik-usAigpeFvWE,41
-wafer_cli-0.2.4.dist-info/top_level.txt,sha256=2MK1IVMWfpLL8BZCQ3E9aG6L6L666gSA_teYlwan4fs,6
-wafer_cli-0.2.4.dist-info/RECORD,,
+wafer_cli-0.2.5.dist-info/METADATA,sha256=O1JyAJtPtr5j4sz5jKiDfyB4d14Mb3tSSDEnQftyzwE,559
+wafer_cli-0.2.5.dist-info/WHEEL,sha256=qELbo2s1Yzl39ZmrAibXA2jjPLUYfnVhUNTlyF1rq0Y,92
+wafer_cli-0.2.5.dist-info/entry_points.txt,sha256=WqB7hB__WhtPY8y1cO2sZiUz7fCq6Ik-usAigpeFvWE,41
+wafer_cli-0.2.5.dist-info/top_level.txt,sha256=2MK1IVMWfpLL8BZCQ3E9aG6L6L666gSA_teYlwan4fs,6
+wafer_cli-0.2.5.dist-info/RECORD,,