129 lines
4.6 KiB
Python
129 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 subprocess
|
|
import platform
|
|
import signal
|
|
|
|
# Add parent directory to path to import utils
|
|
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
|
|
|
try:
|
|
from ComfyUI_Distributed.utils.process import is_process_alive, terminate_process
|
|
from ComfyUI_Distributed.utils.constants import WORKER_CHECK_INTERVAL, PROCESS_TERMINATION_TIMEOUT
|
|
except ImportError:
|
|
# Fallback if running from different context
|
|
def is_process_alive(pid):
|
|
"""Check if a process with given PID is still alive."""
|
|
try:
|
|
if platform.system() == "Windows":
|
|
# Windows: use tasklist
|
|
result = subprocess.run(['tasklist', '/FI', f'PID eq {pid}'],
|
|
capture_output=True, text=True)
|
|
return str(pid) in result.stdout
|
|
else:
|
|
# Unix: send signal 0
|
|
os.kill(pid, 0)
|
|
return True
|
|
except (OSError, subprocess.SubprocessError):
|
|
return False
|
|
|
|
WORKER_CHECK_INTERVAL = 2.0
|
|
PROCESS_TERMINATION_TIMEOUT = 5.0
|
|
|
|
def monitor_and_run(master_pid, command):
|
|
"""Run command and monitor master process."""
|
|
# Start the actual worker process
|
|
worker_process = subprocess.Popen(command)
|
|
|
|
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:
|
|
print(f"[Distributed] Could not write PID file: {e}")
|
|
|
|
# Define cleanup function
|
|
def cleanup_worker(signum=None, frame=None):
|
|
"""Clean up worker process when monitor is terminated."""
|
|
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 NameError:
|
|
# Fallback if terminate_process wasn't imported
|
|
worker_process.terminate()
|
|
try:
|
|
worker_process.wait(timeout=PROCESS_TERMINATION_TIMEOUT)
|
|
except subprocess.TimeoutExpired:
|
|
print("[Distributed] Worker didn't terminate gracefully, forcing kill...")
|
|
worker_process.kill()
|
|
worker_process.wait()
|
|
|
|
print("[Distributed] Worker terminated.")
|
|
sys.exit(0)
|
|
|
|
# Register signal handlers for graceful shutdown
|
|
signal.signal(signal.SIGTERM, cleanup_worker)
|
|
signal.signal(signal.SIGINT, cleanup_worker)
|
|
if platform.system() != "Windows":
|
|
signal.signal(signal.SIGHUP, cleanup_worker)
|
|
|
|
# Monitor loop
|
|
check_interval = WORKER_CHECK_INTERVAL
|
|
|
|
try:
|
|
while True:
|
|
# Check if worker is still running
|
|
if worker_process.poll() is not None:
|
|
print(f"[Distributed] Worker process exited with code: {worker_process.returncode}")
|
|
sys.exit(worker_process.returncode)
|
|
|
|
# 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()
|
|
|
|
time.sleep(check_interval)
|
|
|
|
except KeyboardInterrupt:
|
|
cleanup_worker()
|
|
|
|
if __name__ == "__main__":
|
|
# 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")
|
|
sys.exit(1)
|
|
|
|
try:
|
|
master_pid = int(master_pid)
|
|
except ValueError:
|
|
print(f"[Distributed] Error: Invalid master PID: {master_pid}")
|
|
sys.exit(1)
|
|
|
|
# Get the actual command to run (all remaining arguments)
|
|
if len(sys.argv) < 2:
|
|
print("[Distributed] Error: No command specified")
|
|
sys.exit(1)
|
|
|
|
command = sys.argv[1:]
|
|
|
|
# Start monitoring
|
|
monitor_and_run(master_pid, command) |