scripts/keep_warm.py

"""Keep the Modal ttw-serve endpoint warm so the live Space has no cold start.

The Engine scales down 300s after the last call (serve.py). This pings it with a
tiny batch every PING_EVERY seconds so a judge clicking Step, or you recording the
demo, never waits on a cold start.

CREDIT NOTE: this holds an L4 GPU warm and burns Modal credits the whole time it
runs. Run it only while recording the demo or during the final judging window, and
stop it (Ctrl-C) when done. Do NOT leave it running for days.

Usage:
    PYTHONUTF8=1 PYTHONIOENCODING=utf-8 python scripts/keep_warm.py
    # custom cadence (default 240s, must stay under the 300s scaledown window):
    PYTHONUTF8=1 PYTHONIOENCODING=utf-8 python scripts/keep_warm.py --every 200
"""

from __future__ import annotations

import argparse
import time
from datetime import datetime

from ttw.llm import ModalLLM

PING_EVERY = 240  # seconds; under serve.py's 300s scaledown_window

_PING = [[{"role": "user", "content": "ping"}]]


def main() -> None:
    parser = argparse.ArgumentParser(description="Keep the Modal endpoint warm.")
    parser.add_argument(
        "--every", type=int, default=PING_EVERY,
        help="seconds between pings (keep under 300)",
    )
    args = parser.parse_args()

    llm = ModalLLM()
    print(f"Keeping ttw-serve warm: ping every {args.every}s. Ctrl-C to stop.")
    print("Credits burn while this runs; stop it when you're done.\n")
    try:
        while True:
            t0 = time.time()
            llm.chat_batch(_PING, max_tokens=1, temperature=0.0)
            dt = time.time() - t0
            stamp = datetime.now().strftime("%H:%M:%S")
            print(f"[{stamp}] warm (round trip {dt:.1f}s)")
            time.sleep(max(1, args.every))
    except KeyboardInterrupt:
        print("\nStopped. The endpoint will scale down on its own.")


if __name__ == "__main__":
    main()