scripts.slot_lease

Reference › scripts › slot_lease

scripts/slot_lease.py reserves worker slots for parallel swarms that share the ORCHESTRATA_WORKER_SLOT numbering and the orchkg_w<N> per-worker test databases, so two swarms never use the same slot. A human or a swarm launcher runs it; it has no internal imports. See start_test_services for how the per-worker databases are created and swarm_supervisor for the companion supervision tool.

State is a file; _locked takes an exclusive cross-process lock on its sibling .lock file around every read-modify-write (_read_state, _write_state). Subcommands: acquire --count N --owner NAME [--ttl-minutes M] reserves a contiguous range of slots (find_contiguous_range) and prints export SLOT_LEASE_OWNER=... and export SLOT_LEASE_BASE=... lines for the caller to evaluate, with a lease lifetime of 60 minutes by default and no expiry when the value is zero or negative; release --owner NAME frees every slot that owner holds; list shows current leases; and check SLOT [--owner NAME] tells whether a slot is free or owned by that owner. reap_expired is a pure function that drops leases whose expiry has passed. A refused acquire (invalid count, or no free range) is printed as [slot_lease] error: with exit 1.

At a glance

Class diagram

No classes; see Functions.

Dependency diagram

No internal imports and no importers.

API

Functions

Source

Permalink: https://gitlab.com/MSamy/orchestrata-ultimate-architecture/-/blob/a4cbc2bdf/scripts/slot_lease.py

#!/usr/bin/env python3
"""slot_lease.py — worker-slot reservation for parallel swarms.

``ORCHESTRATA_WORKER_SLOT=slot<N>`` labels map to per-slot test databases
(``orchkg_w<N>``). When more than one swarm (Claude subagents + kimi +
minimax, per the IRC #agents delegation fleet) runs tests concurrently
against the same graph cluster, two swarms picking overlapping slot ranges
by hand silently corrupt each other's test data. This CLI is the missing
reservation layer: a flock-protected lease file, contiguous-range
allocation over slots ``0..31``, TTL-based expiry, and a plain owner-name
identity so ``scripts/sweep_test_pollution.py`` (see its ``--paths`` sibling
edit) can tell "my slot" from "someone else's live slot" before sweeping it.

State file: ``~/.local/state/orchestrata/slot-leases.json``, overridable via
the ``SLOT_LEASE_STATE`` environment variable (tests point this at a
``tmp_path`` file rather than touching the real one).

CLI:
    python scripts/slot_lease.py acquire --count 4 --owner swarmA --ttl-minutes 90
    python scripts/slot_lease.py list
    python scripts/slot_lease.py check 5 --owner swarmA
    python scripts/slot_lease.py release --owner swarmA

``acquire`` prints shell-exportable ``export KEY=value`` lines (including
``SLOT_LEASE_OWNER``, so a downstream ``sweep_test_pollution.py`` call that
``eval``s the acquire output automatically identifies itself).
"""

from __future__ import annotations

import argparse
import contextlib
import fcntl
import json
import os
import sys
import time
from pathlib import Path
from typing import Any

#: Slots this project's test infra recognizes (``ORCHESTRATA_WORKER_SLOT=
#: slot0``..``slot31``, per ``tests/conftest.py``'s per-worker database map).
TOTAL_SLOTS = 32

STATE_ENV = "SLOT_LEASE_STATE"
_DEFAULT_STATE = Path.home() / ".local" / "state" / "orchestrata" / "slot-leases.json"


def default_state_path() -> Path:
    override = os.environ.get(STATE_ENV, "").strip()
    return Path(override) if override else _DEFAULT_STATE


# ---------------------------------------------------------------------------
# State I/O — flock-protected read-modify-write
# ---------------------------------------------------------------------------


@contextlib.contextmanager
def _locked(path: Path):
    """Exclusive cross-PROCESS lock on ``path``'s sibling ``.lock`` file —
    real ``fcntl.flock``, so two real subprocesses racing an ``acquire``
    genuinely serialize rather than merely appearing to in one process.
    """
    path.parent.mkdir(parents=True, exist_ok=True)
    lock_path = path.with_name(path.name + ".lock")
    with open(lock_path, "a+") as lock_fh:
        fcntl.flock(lock_fh.fileno(), fcntl.LOCK_EX)
        try:
            yield
        finally:
            fcntl.flock(lock_fh.fileno(), fcntl.LOCK_UN)


def _read_state(path: Path) -> dict[str, Any]:
    if not path.exists():
        return {"leases": {}}
    try:
        with path.open(encoding="utf-8") as fh:
            data = json.load(fh)
    except json.JSONDecodeError, OSError:
        return {"leases": {}}
    if not isinstance(data, dict) or not isinstance(data.get("leases"), dict):
        return {"leases": {}}
    return data


def _write_state(path: Path, state: dict[str, Any]) -> None:
    tmp = path.with_name(path.name + ".tmp")
    with tmp.open("w", encoding="utf-8") as fh:
        json.dump(state, fh, indent=2, sort_keys=True)
    tmp.replace(path)


def reap_expired(state: dict[str, Any], now: float) -> dict[str, Any]:
    """Drop every lease whose ``expires_at`` has passed. Pure — no I/O — so
    TTL-expiry tests can drive it directly with a synthetic clock instead of
    sleeping minutes.
    """
    leases = state.get("leases", {})
    kept = {
        slot: info
        for slot, info in leases.items()
        if info.get("expires_at") is None or info["expires_at"] > now
    }
    return {"leases": kept}


# ---------------------------------------------------------------------------
# Allocation
# ---------------------------------------------------------------------------


def find_contiguous_range(
    state: dict[str, Any], count: int, total: int = TOTAL_SLOTS
) -> list[int] | None:
    """The lowest free contiguous ``count``-slot range in ``0..total-1``, or
    ``None`` when none exists."""
    leased = {int(s) for s in state.get("leases", {})}
    for start in range(0, total - count + 1):
        block = range(start, start + count)
        if all(s not in leased for s in block):
            return list(block)
    return None


def acquire(
    owner: str,
    count: int,
    *,
    ttl_minutes: float | None = 60.0,
    state_path: Path | None = None,
    now: float | None = None,
) -> list[int]:
    """Reserve the lowest free contiguous ``count``-slot range for ``owner``.

    Raises ``ValueError`` for a nonsensical ``count`` and ``RuntimeError``
    when no such range is currently free — never silently hands out a
    partial or overlapping range.
    """
    if count < 1 or count > TOTAL_SLOTS:
        raise ValueError(f"count must be between 1 and {TOTAL_SLOTS}, got {count}")
    path = state_path or default_state_path()
    ts = now if now is not None else time.time()
    with _locked(path):
        state = reap_expired(_read_state(path), ts)
        block = find_contiguous_range(state, count)
        if block is None:
            raise RuntimeError(
                f"no contiguous free range of {count} slot(s) available in "
                f"0..{TOTAL_SLOTS - 1} — run `slot_lease.py list` to see current leases"
            )
        expires_at = ts + ttl_minutes * 60.0 if ttl_minutes else None
        for slot in block:
            state["leases"][str(slot)] = {
                "owner": owner,
                "acquired_at": ts,
                "ttl_minutes": ttl_minutes,
                "expires_at": expires_at,
            }
        _write_state(path, state)
    return block


def release(
    owner: str, *, state_path: Path | None = None, now: float | None = None
) -> list[int]:
    """Release every slot currently leased to ``owner``. Returns the
    (sorted) list of slots that were released."""
    path = state_path or default_state_path()
    ts = now if now is not None else time.time()
    with _locked(path):
        state = reap_expired(_read_state(path), ts)
        removed = [
            int(slot)
            for slot, info in state["leases"].items()
            if info.get("owner") == owner
        ]
        for slot in removed:
            del state["leases"][str(slot)]
        _write_state(path, state)
    return sorted(removed)


def list_leases(
    *, state_path: Path | None = None, now: float | None = None
) -> dict[str, Any]:
    """Current leases, after reaping expired ones (persisted — "Expired
    TTLs are reaped on any command")."""
    path = state_path or default_state_path()
    ts = now if now is not None else time.time()
    with _locked(path):
        state = reap_expired(_read_state(path), ts)
        _write_state(path, state)
    return state


def active_owner(
    slot: int, *, state_path: Path | None = None, now: float | None = None
) -> str | None:
    """Owner currently holding ``slot`` (after reaping), or ``None`` if free.

    The read-only-feeling seam ``scripts/sweep_test_pollution.py`` imports
    directly (as a bare sibling-script import, matching ``_irc_send.py``'s
    existing ``_irc_conversation`` precedent) to decide whether an
    ``orchkg_w<N>`` database is safe to sweep.
    """
    path = state_path or default_state_path()
    ts = now if now is not None else time.time()
    with _locked(path):
        state = reap_expired(_read_state(path), ts)
        _write_state(path, state)
    info = state["leases"].get(str(slot))
    return info.get("owner") if info else None


def check(
    slot: int,
    owner: str,
    *,
    state_path: Path | None = None,
    now: float | None = None,
) -> bool:
    """``True`` when ``slot`` is free or already owned by ``owner``;
    ``False`` when it is leased to someone else."""
    holder = active_owner(slot, state_path=state_path, now=now)
    return holder is None or holder == owner


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def _fmt_ts(ts: float | None) -> str:
    if ts is None:
        return "never"
    return time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(ts))


def main(argv: list[str] | None = None) -> int:
    parser = argparse.ArgumentParser(
        description=(
            "Worker-slot reservation for parallel swarms sharing "
            "ORCHESTRATA_WORKER_SLOT / orchkg_w<N> test databases."
        )
    )
    sub = parser.add_subparsers(dest="command", required=True)

    acq = sub.add_parser("acquire", help="Reserve a contiguous range of slots")
    acq.add_argument("--count", type=int, required=True)
    acq.add_argument("--owner", required=True)
    acq.add_argument(
        "--ttl-minutes",
        type=float,
        default=60.0,
        help="Lease lifetime in minutes (default: 60; 0/negative disables expiry)",
    )

    rel = sub.add_parser("release", help="Release every slot owned by --owner")
    rel.add_argument("--owner", required=True)

    sub.add_parser("list", help="Show current leases")

    chk = sub.add_parser("check", help="Is a slot free, or owned by --owner?")
    chk.add_argument("slot", type=int)
    chk.add_argument("--owner", default="")

    args = parser.parse_args(argv)

    if args.command == "acquire":
        ttl = args.ttl_minutes if args.ttl_minutes > 0 else None
        try:
            block = acquire(args.owner, args.count, ttl_minutes=ttl)
        except (ValueError, RuntimeError) as exc:
            print(f"[slot_lease] error: {exc}", file=sys.stderr)
            return 1
        print(f"export SLOT_LEASE_OWNER={args.owner}")
        print(f"export SLOT_LEASE_BASE={block[0]}")
        print(f"export SLOT_LEASE_COUNT={len(block)}")
        print(f"export SLOT_LEASE_SLOTS={','.join(str(s) for s in block)}")
        return 0

    if args.command == "release":
        removed = release(args.owner)
        print(
            f"[slot_lease] released {len(removed)} slot(s) for {args.owner!r}: {removed}"
        )
        return 0

    if args.command == "list":
        state = list_leases()
        leases = state.get("leases", {})
        if not leases:
            print("(no active leases)")
            return 0
        print(f"{'SLOT':<6}{'OWNER':<24}{'ACQUIRED':<22}{'EXPIRES'}")
        for slot in sorted(leases, key=int):
            info = leases[slot]
            print(
                f"{slot:<6}{info.get('owner', ''):<24}"
                f"{_fmt_ts(info.get('acquired_at')):<22}{_fmt_ts(info.get('expires_at'))}"
            )
        return 0

    if args.command == "check":
        ok = check(args.slot, args.owner)
        holder = active_owner(args.slot)
        if ok:
            state_note = "free" if holder is None else f"owned by you ({holder})"
            print(f"[slot_lease] slot {args.slot} is {state_note}")
        else:
            print(f"[slot_lease] slot {args.slot} is leased to {holder!r}")
        return 0 if ok else 1

    return 1


if __name__ == "__main__":
    sys.exit(main())