From 3046988a9661096505ca118bb4861bbe307090f1 Mon Sep 17 00:00:00 2001 From: Leo Vasanko Date: Thu, 3 Sep 2026 03:24:46 +0000 Subject: [PATCH] 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 | 77 ++++++++++++++++++++++++++++++++++++------- 1 file changed, 65 insertions(+), 12 deletions(-) diff --git a/scripts/translator.py b/scripts/translator.py index 6a8599a..bc8efbd 100644 --- a/scripts/translator.py +++ b/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__":