Every AI job in the house runs on one card, and a lot of them want it at the same time.

PROXY_LEASE_MAX_TTL_SEC = int(os.environ.get("PROXY_LEASE_MAX_TTL_SEC", "3600"))

# The lease dict is REPLACED wholesale (never mutated in place) so the hot-path
# read in _gpu_lease_active() is lock-free safe (CPython atomic ref read).
_gpu_lease_lock = threading.Lock()
_gpu_lease: dict | None = None


def _gpu_lease_active() -> dict | None:
    """Return the active lease dict if one is held and not expired, else None."""
    lease = _gpu_lease
    if lease is None:
        return None
    if time.time() * 1000 >= lease["expires_ms"]:
        return None  # expired; the monitor loop finalizes the release (restart vLLM)
    return lease


@app.post("/api/gpu/lease")
async def api_gpu_lease_acquire(request: Request):
    """Acquire an exclusive GPU lease: drain ALL VRAM (stop vLLM, evict Ollama,
    free ComfyUI) and park non-PROTECTED LLM callers on the hold queue until
    release or TTL expiry. Mode stays 'work' -- this is NOT a mode flip.

      ttl_sec : seconds until auto-release, capped at PROXY_LEASE_MAX_TTL_SEC.
                The lease ALWAYS expires.
    """
    global _gpu_lease
    ttl_sec = max(1, min(ttl_sec, PROXY_LEASE_MAX_TTL_SEC))

    # Guard 1: one lease at a time.
    if _gpu_lease_active() is not None:
        return JSONResponse({"error": "already_leased"}, status_code=409)
    # Guard 2: only from work mode (gaming/imagegen already own the card; a lease
    # on top would restart vLLM on release while the mode still says drained).
    if mode != "work":
        return JSONResponse({"error": "mode_not_work"}, status_code=409)
    # Guard 3: don't stomp an in-progress mode rollout.
    if _active_rollout is not None:
        return JSONResponse({"error": "rollout_active"}, status_code=409)

    now_ms = time.time() * 1000
    lease = {"id": f"lease-{int(now_ms)}-{secrets.token_hex(4)}",
             "holder": holder, "reason": reason,
             "acquired_ms": int(now_ms),
             "expires_ms": int(now_ms + ttl_sec * 1000),
             "ttl_sec": ttl_sec, "renew_count": 0}
    with _gpu_lease_lock:
        # Re-check under lock (TOCTOU with the active() check above).
        if _gpu_lease is not None and time.time() * 1000 < _gpu_lease["expires_ms"]:
            return JSONResponse({"error": "already_leased"}, status_code=409)
        _gpu_lease = lease

    _persist_lease_to_disk(lease)          # a proxy restart mid-lease must not drop it
    # Same drain routine game mode uses; the mode flag itself is untouched.
    vram_drain = await _drain_vram_for_gaming()
    return JSONResponse(payload)