From dd23503883fdf319e8beb6e7a190445ecf89973c Mon Sep 17 00:00:00 2001 From: robertvoy Date: Wed, 16 Jul 2025 14:07:52 +1000 Subject: [PATCH] better shutdown and security --- distributed.py | 126 +++++++++++++++++++++++++++++++++++++++---------- 1 file changed, 102 insertions(+), 24 deletions(-) diff --git a/distributed.py b/distributed.py index 2b4eaca..9664862 100644 --- a/distributed.py +++ b/distributed.py @@ -14,6 +14,8 @@ import platform import time import atexit import signal +import sys +import shlex from comfy.utils import ProgressBar # Import shared utilities @@ -781,9 +783,20 @@ class WorkerProcessManager: raise RuntimeError(error_msg) - # Add any extra arguments + # Add any extra arguments safely if worker_config.get('extra_args'): - cmd.extend(worker_config['extra_args'].split()) + raw_args = worker_config['extra_args'].strip() + if raw_args: + # Safely split using shlex to handle quotes and spaces + extra_args_list = shlex.split(raw_args) + + # Validate: Block dangerous shell meta-characters (e.g., ; | & > < ` $) + forbidden_chars = set(';|>&<`$()[]{}*!?') + for arg in extra_args_list: + if any(c in forbidden_chars for c in arg): + raise ValueError(f"Invalid characters in extra_args: {arg}. Forbidden: {''.join(forbidden_chars)}") + + cmd.extend(extra_args_list) return cmd @@ -1210,15 +1223,9 @@ def delayed_auto_launch(): timer.daemon = True timer.start() -# Call delayed auto-launch only if we're the master (not a worker) -if not os.environ.get('COMFYUI_MASTER_PID'): - delayed_auto_launch() -else: - debug_log("Running as worker, skipping auto-launch") - -# Register cleanup on exit - only clean up if setting is enabled -def cleanup_on_exit(signum=None, frame=None): - """Handle cleanup on exit or signal""" +# Async cleanup function for proper shutdown +async def async_cleanup_and_exit(signum=None): + """Async-friendly cleanup and exit.""" try: config = load_config() if config.get('settings', {}).get('stop_workers_on_master_exit', True): @@ -1226,24 +1233,95 @@ def cleanup_on_exit(signum=None, frame=None): worker_manager.cleanup_all() else: print("\n[Distributed] Master shutting down, workers will continue running") - # Still save the current state + worker_manager.save_processes() + except Exception as e: + print(f"[Distributed] Error during cleanup: {e}") + + # On Windows, we need to exit differently + if platform.system() == "Windows": + # Force exit on Windows + sys.exit(0) + else: + # On Unix, stop the event loop gracefully + loop = asyncio.get_running_loop() + loop.stop() + +def register_async_signals(): + """Register async signal handlers for graceful shutdown.""" + if platform.system() == "Windows": + # Windows doesn't support add_signal_handler, use traditional signal handling + def signal_handler(signum, frame): + # Schedule the async cleanup in the event loop + loop = server.PromptServer.instance.loop + if loop and loop.is_running(): + asyncio.run_coroutine_threadsafe(async_cleanup_and_exit(signum), loop) + else: + # Fallback to sync cleanup if loop isn't running + try: + config = load_config() + if config.get('settings', {}).get('stop_workers_on_master_exit', True): + print("\n[Distributed] Master shutting down, stopping all managed workers...") + worker_manager.cleanup_all() + else: + print("\n[Distributed] Master shutting down, workers will continue running") + worker_manager.save_processes() + except Exception as e: + print(f"[Distributed] Error during cleanup: {e}") + sys.exit(0) + + signal.signal(signal.SIGINT, signal_handler) + signal.signal(signal.SIGTERM, signal_handler) + else: + # Unix-like systems support add_signal_handler + loop = server.PromptServer.instance.loop + + for sig in (signal.SIGINT, signal.SIGTERM): + loop.add_signal_handler(sig, lambda s=sig: asyncio.create_task(async_cleanup_and_exit(s))) + + # SIGHUP is Unix-only + loop.add_signal_handler(signal.SIGHUP, lambda: asyncio.create_task(async_cleanup_and_exit(signal.SIGHUP))) + +def sync_cleanup(): + """Synchronous wrapper for atexit.""" + try: + # For atexit, we don't want to stop the loop or exit + config = load_config() + if config.get('settings', {}).get('stop_workers_on_master_exit', True): + print("\n[Distributed] Master shutting down, stopping all managed workers...") + worker_manager.cleanup_all() + else: + print("\n[Distributed] Master shutting down, workers will continue running") worker_manager.save_processes() except Exception as e: print(f"[Distributed] Error during cleanup: {e}") -# Register cleanup handlers -atexit.register(cleanup_on_exit) +# Register atexit handler for normal exits +atexit.register(sync_cleanup) -# Handle terminal window closing and Ctrl+C -try: - signal.signal(signal.SIGINT, cleanup_on_exit) - signal.signal(signal.SIGTERM, cleanup_on_exit) - - if platform.system() != "Windows": - # SIGHUP is sent when terminal closes on Unix - signal.signal(signal.SIGHUP, cleanup_on_exit) -except Exception as e: - print(f"[Distributed] Warning: Could not set signal handlers: {e}") +# Call delayed auto-launch only if we're the master (not a worker) +if not os.environ.get('COMFYUI_MASTER_PID'): + delayed_auto_launch() + try: + register_async_signals() # Register async signal handlers + except Exception as e: + print(f"[Distributed] Warning: Could not register async signal handlers: {e}") + # Fall back to basic signal handling + def basic_signal_handler(signum, frame): + print("\n[Distributed] Received signal, shutting down...") + try: + config = load_config() + if config.get('settings', {}).get('stop_workers_on_master_exit', True): + worker_manager.cleanup_all() + else: + worker_manager.save_processes() + except Exception as cleanup_error: + print(f"[Distributed] Error during cleanup: {cleanup_error}") + sys.exit(0) + + signal.signal(signal.SIGINT, basic_signal_handler) + signal.signal(signal.SIGTERM, basic_signal_handler) +else: + debug_log("Running as worker, skipping auto-launch") # --- Persistent State Storage --- # Store job queue on the persistent server instance to survive script reloads