wafer-cli 0.2.6__py3-none-any.whl → 0.2.7__py3-none-any.whl

wafer/target_lock.py

@@ -1,198 +0,0 @@
-"""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)