ViperEkura
ceadc34ea9
- Register SIGTERM/SIGINT handlers in training loop, set stop flag on signal - Check stop_requested at each epoch/batch boundary, break and call on_error to save checkpoint - LocalStrategy parent forwards signal to child processes via terminate(), waits up to 600s for graceful exit - TrainContext gains threading.Event-based stop_requested/request_stop - Tests verify SIGTERM/SIGINT trigger checkpoint save with exit code 0, works on both CPU and GPU
45 lines
1.1 KiB
Python
45 lines
1.1 KiB
Python
import logging
|
|
import os
|
|
import signal
|
|
import threading
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
_early_stop = threading.Event()
|
|
_active_context = None
|
|
|
|
|
|
def _early_handler(signum: int, frame):
|
|
sig = signal.Signals(signum)
|
|
logger.warning(
|
|
"Received %s (pid=%d), requesting graceful training stop...",
|
|
sig.name,
|
|
os.getpid(),
|
|
)
|
|
_early_stop.set()
|
|
if _active_context is not None:
|
|
_active_context.request_stop()
|
|
|
|
|
|
def install_early_signal_handlers():
|
|
for sig in (signal.SIGTERM, signal.SIGINT):
|
|
signal.signal(sig, _early_handler)
|
|
|
|
|
|
def register_signal_handlers(context):
|
|
global _active_context
|
|
_active_context = context
|
|
for sig in (signal.SIGTERM, signal.SIGINT):
|
|
signal.signal(sig, _early_handler)
|
|
if _early_stop.is_set():
|
|
context.request_stop()
|
|
logger.warning("Signal was received during initialization, stopping...")
|
|
|
|
|
|
def unregister_signal_handlers():
|
|
global _active_context
|
|
_active_context = None
|
|
_early_stop.clear()
|
|
signal.signal(signal.SIGTERM, signal.SIG_DFL)
|
|
signal.signal(signal.SIGINT, signal.SIG_DFL)
|