diff --git a/platform/docker/flux-worker/flux_worker.py b/platform/docker/flux-worker/flux_worker.py index 09a2f4e..f19bbce 100644 --- a/platform/docker/flux-worker/flux_worker.py +++ b/platform/docker/flux-worker/flux_worker.py @@ -6,6 +6,7 @@ from __future__ import annotations import gc import json import os +import signal import time from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer from pathlib import Path @@ -21,6 +22,11 @@ LOAD_SECONDS = 0.0 if len(TOKEN) < 32: raise RuntimeError("WORKER_TOKEN is missing or too short") +# The container is intentionally disposable. After the router has received a +# completed response and saved image, an immediate process exit releases the +# CUDA context much faster than Python/PyTorch interpreter teardown. +signal.signal(signal.SIGTERM, lambda *_: os._exit(0)) + def load_pipeline() -> None: global PIPE, LOAD_SECONDS