ViperEkura
59248032dc
- Fix pre-existing ruff lint warnings (F401, F541, F841, E741) - Exclude .md/.json/.yml from ruff format check - Unblock SIGTERM/SIGINT via pthread_sigmask in early signal handler - Do not restore SIG_DFL on unregister to prevent pending signal kills
54 lines
1.3 KiB
Python
54 lines
1.3 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)
|
|
_unblock_signals()
|
|
|
|
|
|
def _unblock_signals():
|
|
try:
|
|
mask = signal.pthread_sigmask(signal.SIG_BLOCK, set())
|
|
blocked = {signal.SIGTERM, signal.SIGINT} & mask
|
|
if blocked:
|
|
signal.pthread_sigmask(signal.SIG_UNBLOCK, blocked)
|
|
except (AttributeError, OSError):
|
|
pass
|
|
|
|
|
|
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()
|