Robert Wojciechowski
3058b11ff2
- Remove DistributedListSplitter and DistributedListCollector nodes - Remove /distributed/init_list_queue and /distributed/request_list_item endpoints - Clean up all references from prompt_transform, queue_orchestration, constants, and tests - Refactor and reorganize distributed modules for quality improvements
136 lines
4.6 KiB
Python
136 lines
4.6 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Worker process monitor - monitors if the master process is still alive
|
|
and terminates the worker if the master dies.
|
|
"""
|
|
import os
|
|
import sys
|
|
import time
|
|
import platform
|
|
import signal
|
|
from typing import Any
|
|
|
|
try:
|
|
from ..utils.process import is_process_alive, launch_process_with_timeout, terminate_process
|
|
from ..utils.constants import WORKER_CHECK_INTERVAL, PROCESS_TERMINATION_TIMEOUT
|
|
except ImportError:
|
|
# Fallback when executed as a standalone script.
|
|
from utils.process import is_process_alive, launch_process_with_timeout, terminate_process
|
|
from utils.constants import WORKER_CHECK_INTERVAL, PROCESS_TERMINATION_TIMEOUT
|
|
|
|
def monitor_and_run(master_pid: int, command: list[str]) -> int:
|
|
"""Run command and monitor master process; return process exit code."""
|
|
# Start the actual worker process
|
|
print(f"[Distributed] Launching worker command: {' '.join(command)}")
|
|
worker_process = launch_process_with_timeout(command, timeout_seconds=10.0)
|
|
|
|
print(f"[Distributed] Started worker PID: {worker_process.pid}")
|
|
print(f"[Distributed] Monitoring master PID: {master_pid}")
|
|
|
|
# Write worker PID to a file so parent can track it
|
|
monitor_pid = os.getpid()
|
|
pid_info_file = os.environ.get('WORKER_PID_FILE')
|
|
if pid_info_file:
|
|
try:
|
|
with open(pid_info_file, 'w') as f:
|
|
f.write(f"{monitor_pid},{worker_process.pid}")
|
|
print(f"[Distributed] Wrote PID info to {pid_info_file}")
|
|
except Exception as e:
|
|
raise RuntimeError(f"Could not write PID file '{pid_info_file}': {e}") from e
|
|
|
|
# Define cleanup function
|
|
should_stop = False
|
|
monitor_exit_code = 0
|
|
|
|
def cleanup_worker(
|
|
signum: int | None = None,
|
|
frame: Any | None = None,
|
|
exit_code: int = 0,
|
|
) -> None:
|
|
"""Clean up worker process when monitor is terminated."""
|
|
nonlocal should_stop
|
|
nonlocal monitor_exit_code
|
|
_ = frame
|
|
should_stop = True
|
|
monitor_exit_code = int(exit_code)
|
|
|
|
if signum:
|
|
print(f"\n[Distributed] Received signal {signum}, terminating worker...")
|
|
else:
|
|
print("\n[Distributed] Terminating worker...")
|
|
|
|
if worker_process.poll() is None: # Still running
|
|
try:
|
|
terminate_process(worker_process, timeout=PROCESS_TERMINATION_TIMEOUT)
|
|
except TimeoutError:
|
|
print("[Distributed] Worker did not terminate within timeout.")
|
|
|
|
print("[Distributed] Worker terminated.")
|
|
|
|
def _signal_handler(signum, frame):
|
|
cleanup_worker(signum=signum, frame=frame, exit_code=0)
|
|
|
|
# Register signal handlers for graceful shutdown
|
|
signal.signal(signal.SIGTERM, _signal_handler)
|
|
signal.signal(signal.SIGINT, _signal_handler)
|
|
if platform.system() != "Windows":
|
|
signal.signal(signal.SIGHUP, _signal_handler)
|
|
|
|
# Monitor loop
|
|
check_interval = WORKER_CHECK_INTERVAL
|
|
|
|
try:
|
|
while True:
|
|
if should_stop:
|
|
break
|
|
|
|
# Check if worker is still running
|
|
if worker_process.poll() is not None:
|
|
print(f"[Distributed] Worker process exited with code: {worker_process.returncode}")
|
|
monitor_exit_code = int(worker_process.returncode or 0)
|
|
break
|
|
|
|
# Check if master is still alive
|
|
if not is_process_alive(master_pid):
|
|
print(f"[Distributed] Master process {master_pid} is no longer running. Terminating worker...")
|
|
cleanup_worker(exit_code=0)
|
|
break
|
|
|
|
time.sleep(check_interval)
|
|
|
|
except KeyboardInterrupt:
|
|
cleanup_worker(exit_code=0)
|
|
|
|
return monitor_exit_code
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
"""CLI entrypoint for worker monitor script."""
|
|
args = list(sys.argv[1:] if argv is None else argv)
|
|
|
|
# Get master PID from environment
|
|
master_pid = os.environ.get('COMFYUI_MASTER_PID')
|
|
if not master_pid:
|
|
print("[Distributed] Error: COMFYUI_MASTER_PID not set")
|
|
return 1
|
|
|
|
try:
|
|
master_pid = int(master_pid)
|
|
except ValueError:
|
|
print(f"[Distributed] Error: Invalid master PID: {master_pid}")
|
|
return 1
|
|
|
|
# Get the actual command to run (all remaining arguments)
|
|
if not args:
|
|
print("[Distributed] Error: No command specified")
|
|
return 1
|
|
|
|
command = args
|
|
|
|
# Start monitoring
|
|
return monitor_and_run(master_pid, command)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|