Translator: unload the model after 60s idle, reload on demand

The service stays connected full time; the GPU is held only while
translating. Loads at startup (backlog is likely after a downtime).

scripts/translator.py

@@ -22,7 +22,10 @@ languages (``Data.translate_langs``) ∩ the announced capabilities; a
 dropped connection's in-flight item is simply re-offered
 (docs/localization.md). For parallelism, run multiple instances.
 
-Seed-X-PPO-7B (bf16, ~15 GB) is the only supported model.
+Seed-X-PPO-7B (bf16, ~15 GB) is the only supported model. The script stays
+running and connected full time; the model loads at startup (a backlog is
+likely after a downtime) and is unloaded after 60 s idle, re-loading on the
+next job — the GPU is held only while actually translating.
 
 Usage:
     uv run scripts/translator.py ws://localhost:8410/_translate/KEY
@@ -31,6 +34,7 @@ Usage:
 
 import argparse
 import asyncio
+import gc
 import sys
 import time
 
@@ -148,12 +152,59 @@ class Result(msgspec.Struct, tag="result"):
     texts: list[str]  #: the job's segments translated, same order and count
 
 
-def load_seed_x():
-    t0 = time.monotonic()
-    tokenizer = AutoTokenizer.from_pretrained(SEED_X)
-    model = AutoModelForCausalLM.from_pretrained(SEED_X, dtype=torch.bfloat16, device_map="auto")
-    print(f"[seed-x loaded in {time.monotonic() - t0:.0f}s]", file=sys.stderr)
-    return tokenizer, model
+#: Idle seconds after the last job before the model is unloaded (the GPU
+#: is released; the WebSocket connection and tokenizer stay).
+IDLE_UNLOAD_S = 60
+
+
+class SeedX:
+    """The Seed-X model, loaded at startup and re-loaded on demand.
+
+    Loading up front covers the likely backlog after a downtime (and any
+    first-run model download) before the server starts dispatching. After
+    IDLE_UNLOAD_S without a job the model is dropped and re-loaded on the
+    next one — the script stays connected the whole time, holding the GPU
+    only while translating. The tokenizer (small, CPU) loads once.
+    """
+
+    def __init__(self):
+        self.tokenizer = AutoTokenizer.from_pretrained(SEED_X)
+        self.model = None
+        self._unload_task = None
+        self._load()
+
+    def _load(self):
+        t0 = time.monotonic()
+        self.model = AutoModelForCausalLM.from_pretrained(
+            SEED_X, dtype=torch.bfloat16, device_map="auto")
+        print(f"[seed-x loaded in {time.monotonic() - t0:.0f}s]", file=sys.stderr)
+
+    def get(self):
+        """The (tokenizer, model) pair, re-loading the model if it was
+        idle-unloaded, and cancelling any pending idle unload."""
+        if self._unload_task:
+            self._unload_task.cancel()
+            self._unload_task = None
+        if self.model is None:
+            self._load()
+        return self.tokenizer, self.model
+
+    def idle(self):
+        """Re-arm the idle unload after a job completes (arming it at job
+        START could unload under a >IDLE_UNLOAD_S generation)."""
+        self._unload_task = asyncio.create_task(self._unload_later())
+
+    async def _unload_later(self):
+        try:
+            await asyncio.sleep(IDLE_UNLOAD_S)
+        except asyncio.CancelledError:
+            return
+        if self.model is not None:
+            self.model = None
+            gc.collect()
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+            print(f"[seed-x unloaded after {IDLE_UNLOAD_S}s idle]", file=sys.stderr)
 
 
 def seed_x_chunk(tokenizer, model, text: str, target_lang: str, tag: str,
@@ -210,9 +261,10 @@ def seed_x_chunk(tokenizer, model, text: str, target_lang: str, tag: str,
     return again, n + n2, dt + dt2
 
 
-async def do_job(ws, job: Job, tokenizer, model) -> None:
+async def do_job(ws, job: Job, seed_x: SeedX) -> None:
     """Translate the job's segments (one model call each) and send them back."""
     lang_name = SEED_X_NAMES[job.lang].capitalize()
+    tokenizer, model = seed_x.get()
     # Deliberately blocking: nothing else needs the loop while the job is
     # being answered, and the reconnect loop recovers a dropped connection
     # (the in-flight item is simply re-offered).
@@ -228,10 +280,12 @@ async def do_job(ws, job: Job, tokenizer, model) -> None:
     print(f"[{job.lang} {job.kind} {job.path or '/'}: {len(texts)} segments, "
           f"{tokens} tokens in {dt:.1f}s = {tokens / dt:.1f} tok/s]", file=sys.stderr)
     await ws.send(msgspec.json.encode(Result(lang=job.lang, key=job.key, texts=texts)).decode())
+    seed_x.idle()
 
 
-async def serve(url: str, tokenizer, model) -> None:
+async def serve(url: str, seed_x: SeedX) -> None:
     """Connect, announce capabilities, answer jobs; reconnect with backoff."""
+    seed_x.idle()  # the startup load also unloads when no work arrives
     backoff = 1
     while True:
         try:
@@ -241,7 +295,7 @@ async def serve(url: str, tokenizer, model) -> None:
                 print(f"[connected; announced {len(SEED_X_NAMES)} language capabilities]",
                       file=sys.stderr)
                 async for raw in ws:
-                    await do_job(ws, msgspec.json.decode(raw, type=Job), tokenizer, model)
+                    await do_job(ws, msgspec.json.decode(raw, type=Job), seed_x)
         except websockets.exceptions.InvalidHandshake:
             sys.exit("handshake rejected; check the URL (including the key)")
         except (OSError, websockets.exceptions.ConnectionClosed) as e:
@@ -258,8 +312,7 @@ def main():
     if not args.url.startswith(("ws://", "wss://")):
         p.error("url must start with ws:// or wss://")
 
-    tokenizer, model = load_seed_x()  # once, before the (re)connect loop
-    asyncio.run(serve(args.url, tokenizer, model))
+    asyncio.run(serve(args.url, SeedX()))
 
 
 if __name__ == "__main__":