robertvoy-ComfyUI-Distributed/workers/worker_monitor.py

#!/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
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 if running from different context
    def is_process_alive(pid: int | str) -> bool:
        """Best-effort fallback process liveness probe."""
        try:
            pid = int(pid)
        except (TypeError, ValueError):
            return False

        if os.name == "nt":
            try:
                output = subprocess.check_output(
                    ["tasklist", "/FI", f"PID eq {pid}"],
                    stderr=subprocess.DEVNULL,
                    text=True,
                    timeout=5.0,
                )
            except Exception:
                return False
            return str(pid) in output

        try:
            os.kill(pid, 0)
            return True
        except OSError:
            return False
    
    def launch_process_with_timeout(
        command: list[str],
        timeout_seconds: float = 10.0,
        **kwargs: Any,
    ):
        _ = kwargs
        raise RuntimeError(
            "launch_process_with_timeout unavailable: utils.process import failed"
        )

    WORKER_CHECK_INTERVAL = 2.0
    PROCESS_TERMINATION_TIMEOUT = 5.0

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 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.")

    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())