522 lines
22 KiB
Python
522 lines
22 KiB
Python
import os
|
|
import json
|
|
import re
|
|
import asyncio
|
|
import subprocess
|
|
import platform
|
|
import time
|
|
import atexit
|
|
import signal
|
|
import sys
|
|
import shlex
|
|
import uuid
|
|
import torch
|
|
|
|
from ..utils.logging import debug_log, log
|
|
from ..utils.config import load_config, save_config
|
|
from ..utils.process import is_process_alive, terminate_process, get_python_executable
|
|
from ..utils.constants import PROCESS_TERMINATION_TIMEOUT, WORKER_CHECK_INTERVAL, PROCESS_WAIT_TIMEOUT
|
|
|
|
# Try to import psutil for better process management
|
|
try:
|
|
import psutil
|
|
PSUTIL_AVAILABLE = True
|
|
except ImportError:
|
|
log("psutil not available, using fallback process management")
|
|
PSUTIL_AVAILABLE = False
|
|
|
|
|
|
class WorkerProcessManager:
|
|
def __init__(self):
|
|
self.processes = {} # worker_id -> process info
|
|
self.load_processes()
|
|
|
|
def find_comfy_root(self):
|
|
"""Find the ComfyUI root directory."""
|
|
# Start from current file location
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
# Method 1: Check for environment variable override
|
|
env_root = os.environ.get('COMFYUI_ROOT')
|
|
if env_root and os.path.exists(os.path.join(env_root, "main.py")):
|
|
debug_log(f"Found ComfyUI root via COMFYUI_ROOT environment variable: {env_root}")
|
|
return env_root
|
|
|
|
# Method 2: Try going up from custom_nodes directory
|
|
# This file should be in ComfyUI/custom_nodes/ComfyUI-Distributed/
|
|
potential_root = os.path.dirname(os.path.dirname(current_dir))
|
|
if os.path.exists(os.path.join(potential_root, "main.py")):
|
|
debug_log(f"Found ComfyUI root via directory traversal: {potential_root}")
|
|
return potential_root
|
|
|
|
# Method 3: Look for common Docker paths
|
|
docker_paths = ["/basedir", "/ComfyUI", "/app", "/workspace/ComfyUI", "/comfyui", "/opt/ComfyUI", "/workspace"]
|
|
for path in docker_paths:
|
|
if os.path.exists(path) and os.path.exists(os.path.join(path, "main.py")):
|
|
debug_log(f"Found ComfyUI root in Docker path: {path}")
|
|
return path
|
|
|
|
# Method 4: Search upwards for main.py
|
|
search_dir = current_dir
|
|
for _ in range(5): # Limit search depth
|
|
if os.path.exists(os.path.join(search_dir, "main.py")):
|
|
debug_log(f"Found ComfyUI root via upward search: {search_dir}")
|
|
return search_dir
|
|
parent = os.path.dirname(search_dir)
|
|
if parent == search_dir: # Reached root
|
|
break
|
|
search_dir = parent
|
|
|
|
# Method 5: Try to import and use folder_paths
|
|
try:
|
|
import folder_paths
|
|
# folder_paths.base_path should point to ComfyUI root
|
|
if hasattr(folder_paths, 'base_path') and os.path.exists(os.path.join(folder_paths.base_path, "main.py")):
|
|
debug_log(f"Found ComfyUI root via folder_paths: {folder_paths.base_path}")
|
|
return folder_paths.base_path
|
|
except Exception:
|
|
pass
|
|
|
|
# If all methods fail, log detailed information and return the best guess
|
|
log(f"Warning: Could not reliably determine ComfyUI root directory")
|
|
log(f"Current directory: {current_dir}")
|
|
log(f"Initial guess was: {potential_root}")
|
|
|
|
# Return the initial guess as fallback
|
|
return potential_root
|
|
|
|
def _find_windows_terminal(self):
|
|
"""Find Windows Terminal executable."""
|
|
# Common locations for Windows Terminal
|
|
possible_paths = [
|
|
os.path.expandvars(r"%LOCALAPPDATA%\Microsoft\WindowsApps\wt.exe"),
|
|
os.path.expandvars(r"%PROGRAMFILES%\WindowsApps\Microsoft.WindowsTerminal_*\wt.exe"),
|
|
"wt.exe" # Try PATH
|
|
]
|
|
|
|
for path in possible_paths:
|
|
if os.path.exists(path):
|
|
return path
|
|
# Handle wildcard for WindowsApps
|
|
if '*' in path:
|
|
import glob
|
|
matches = glob.glob(path)
|
|
if matches:
|
|
return matches[0]
|
|
|
|
# Try to find it in PATH
|
|
import shutil
|
|
wt_path = shutil.which("wt")
|
|
if wt_path:
|
|
return wt_path
|
|
|
|
return None
|
|
|
|
def build_launch_command(self, worker_config, comfy_root):
|
|
"""Build the command to launch a worker."""
|
|
# Use main.py directly - it's the most reliable method
|
|
main_py = os.path.join(comfy_root, "main.py")
|
|
|
|
if os.path.exists(main_py):
|
|
cmd = [
|
|
get_python_executable(),
|
|
main_py,
|
|
"--port", str(worker_config['port']),
|
|
"--enable-cors-header"
|
|
]
|
|
debug_log(f"Using main.py: {main_py}")
|
|
else:
|
|
# Provide detailed error message
|
|
error_msg = f"Could not find main.py in {comfy_root}\n"
|
|
error_msg += f"Searched for: {main_py}\n"
|
|
error_msg += f"Directory contents of {comfy_root}:\n"
|
|
try:
|
|
if os.path.exists(comfy_root):
|
|
files = os.listdir(comfy_root)[:20] # List first 20 files
|
|
error_msg += " " + "\n ".join(files)
|
|
if len(os.listdir(comfy_root)) > 20:
|
|
error_msg += f"\n ... and {len(os.listdir(comfy_root)) - 20} more files"
|
|
else:
|
|
error_msg += f" Directory {comfy_root} does not exist!"
|
|
except Exception as e:
|
|
error_msg += f" Error listing directory: {e}"
|
|
|
|
# Try to suggest the correct path
|
|
error_msg += "\n\nPossible solutions:\n"
|
|
error_msg += "1. Check if ComfyUI is installed in a different location\n"
|
|
error_msg += "2. For Docker: ComfyUI might be in /ComfyUI or /app\n"
|
|
error_msg += "3. Ensure the custom node is installed in the correct location\n"
|
|
|
|
raise RuntimeError(error_msg)
|
|
|
|
# Add any extra arguments safely
|
|
if worker_config.get('extra_args'):
|
|
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
|
|
|
|
def launch_worker(self, worker_config, show_window=False):
|
|
"""Launch a worker process with logging."""
|
|
comfy_root = self.find_comfy_root()
|
|
|
|
# Set up environment
|
|
env = os.environ.copy()
|
|
env['CUDA_VISIBLE_DEVICES'] = str(worker_config.get('cuda_device', 0))
|
|
env['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
|
|
|
|
# Pass master PID to worker so it can monitor if master is still alive
|
|
env['COMFYUI_MASTER_PID'] = str(os.getpid())
|
|
|
|
cmd = self.build_launch_command(worker_config, comfy_root)
|
|
|
|
# Change to ComfyUI root directory for the process
|
|
cwd = comfy_root
|
|
|
|
# Create log directory and file
|
|
log_dir = os.path.join(comfy_root, "logs", "workers")
|
|
os.makedirs(log_dir, exist_ok=True)
|
|
|
|
# Use daily log files instead of timestamp
|
|
date_stamp = time.strftime("%Y%m%d")
|
|
worker_name = worker_config.get('name', f'Worker{worker_config["id"]}')
|
|
# Clean worker name for filename
|
|
safe_name = "".join(c if c.isalnum() or c in ('-', '_') else '_' for c in worker_name)
|
|
log_file = os.path.join(log_dir, f"{safe_name}_{date_stamp}.log")
|
|
|
|
# Launch process with logging (append mode for daily logs)
|
|
with open(log_file, 'a') as log_handle:
|
|
# Write startup info to log with timestamp
|
|
log_handle.write(f"\n\n{'='*50}\n")
|
|
log_handle.write(f"=== ComfyUI Worker Session Started ===\n")
|
|
log_handle.write(f"Worker: {worker_name}\n")
|
|
log_handle.write(f"Port: {worker_config['port']}\n")
|
|
log_handle.write(f"CUDA Device: {worker_config.get('cuda_device', 0)}\n")
|
|
log_handle.write(f"Started: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")
|
|
log_handle.write(f"Command: {' '.join(cmd)}\n")
|
|
|
|
# Note about worker behavior
|
|
config = load_config()
|
|
stop_on_master_exit = config.get('settings', {}).get('stop_workers_on_master_exit', True)
|
|
|
|
if stop_on_master_exit:
|
|
log_handle.write("Note: Worker will stop when master shuts down\n")
|
|
else:
|
|
log_handle.write("Note: Worker will continue running after master shuts down\n")
|
|
|
|
log_handle.write("=" * 30 + "\n\n")
|
|
log_handle.flush()
|
|
|
|
# Wrap command with monitor if needed
|
|
if stop_on_master_exit and env.get('COMFYUI_MASTER_PID'):
|
|
# Use the monitor wrapper
|
|
monitor_script = os.path.join(os.path.dirname(__file__), 'worker_monitor.py')
|
|
monitored_cmd = [get_python_executable(), monitor_script] + cmd
|
|
log_handle.write(f"[Worker Monitor] Monitoring master PID: {env['COMFYUI_MASTER_PID']}\n")
|
|
log_handle.flush()
|
|
else:
|
|
monitored_cmd = cmd
|
|
|
|
# Platform-specific process creation - always hidden with logging
|
|
if platform.system() == "Windows":
|
|
CREATE_NO_WINDOW = 0x08000000
|
|
process = subprocess.Popen(
|
|
monitored_cmd, env=env, cwd=cwd,
|
|
stdout=log_handle,
|
|
stderr=subprocess.STDOUT,
|
|
creationflags=CREATE_NO_WINDOW
|
|
)
|
|
else:
|
|
# Unix-like systems
|
|
process = subprocess.Popen(
|
|
monitored_cmd, env=env, cwd=cwd,
|
|
stdout=log_handle,
|
|
stderr=subprocess.STDOUT,
|
|
start_new_session=True # Detach from parent
|
|
)
|
|
|
|
# Track the process with log file info - use string ID for consistency
|
|
worker_id = str(worker_config['id'])
|
|
self.processes[worker_id] = {
|
|
'pid': process.pid,
|
|
'process': process,
|
|
'started_at': time.time(),
|
|
'config': worker_config,
|
|
'log_file': log_file,
|
|
'is_monitor': stop_on_master_exit and env.get('COMFYUI_MASTER_PID'), # Track if using monitor
|
|
'launching': True # Mark as launching until confirmed running
|
|
}
|
|
|
|
# Save process info for persistence
|
|
self.save_processes()
|
|
|
|
if stop_on_master_exit and env.get('COMFYUI_MASTER_PID'):
|
|
debug_log(f"Launched worker {worker_name} via monitor (Monitor PID: {process.pid})")
|
|
else:
|
|
log(f"Launched worker {worker_name} directly (PID: {process.pid})")
|
|
debug_log(f"Log file: {log_file}")
|
|
return process.pid
|
|
|
|
def stop_worker(self, worker_id):
|
|
"""Stop a worker process."""
|
|
# Ensure worker_id is string
|
|
worker_id = str(worker_id)
|
|
if worker_id not in self.processes:
|
|
return False, "Worker not managed by UI"
|
|
|
|
proc_info = self.processes[worker_id]
|
|
process = proc_info.get('process')
|
|
pid = proc_info['pid']
|
|
|
|
debug_log(f"Attempting to stop worker {worker_id} (PID: {pid})")
|
|
|
|
# For restored processes without subprocess object
|
|
if not process:
|
|
try:
|
|
debug_log(f"[Distributed] Stopping restored process (no subprocess object)")
|
|
if self._kill_process_tree(pid):
|
|
del self.processes[worker_id]
|
|
self.save_processes()
|
|
debug_log(f"Successfully stopped worker {worker_id} and all child processes")
|
|
return True, "Worker stopped"
|
|
else:
|
|
return False, "Failed to stop worker process"
|
|
except Exception as e:
|
|
log(f"[Distributed] Exception during stop: {e}")
|
|
return False, f"Error stopping worker: {str(e)}"
|
|
|
|
# Normal case with subprocess object
|
|
# Check if still running
|
|
if process.poll() is not None:
|
|
# Already stopped
|
|
log(f"[Distributed] Worker {worker_id} already stopped")
|
|
del self.processes[worker_id]
|
|
self.save_processes()
|
|
return False, "Worker already stopped"
|
|
|
|
# Try to kill the entire process tree
|
|
try:
|
|
debug_log(f"Using process tree kill for worker {worker_id}")
|
|
if self._kill_process_tree(pid):
|
|
# Clean up tracking
|
|
del self.processes[worker_id]
|
|
self.save_processes()
|
|
debug_log(f"Successfully stopped worker {worker_id} and all child processes")
|
|
return True, "Worker stopped"
|
|
else:
|
|
# Fallback to normal termination
|
|
log(f"[Distributed] Process tree kill failed, trying normal termination")
|
|
if process:
|
|
terminate_process(process, timeout=PROCESS_TERMINATION_TIMEOUT)
|
|
|
|
del self.processes[worker_id]
|
|
self.save_processes()
|
|
return True, "Worker stopped (fallback)"
|
|
|
|
except Exception as e:
|
|
log(f"[Distributed] Exception during stop: {e}")
|
|
return False, f"Error stopping worker: {str(e)}"
|
|
|
|
def get_managed_workers(self):
|
|
"""Get list of workers managed by this process."""
|
|
managed = {}
|
|
for worker_id, proc_info in list(self.processes.items()):
|
|
# Check if process is still running
|
|
is_running, _ = self._check_worker_process(worker_id, proc_info)
|
|
|
|
if is_running:
|
|
managed[worker_id] = {
|
|
'pid': proc_info['pid'],
|
|
'started_at': proc_info['started_at'],
|
|
'log_file': proc_info.get('log_file'),
|
|
'launching': proc_info.get('launching', False)
|
|
}
|
|
else:
|
|
# Process has stopped, remove from tracking
|
|
del self.processes[worker_id]
|
|
|
|
return managed
|
|
|
|
def cleanup_all(self):
|
|
"""Stop all managed workers (called on shutdown)."""
|
|
for worker_id in list(self.processes.keys()):
|
|
try:
|
|
self.stop_worker(worker_id)
|
|
except Exception as e:
|
|
log(f"[Distributed] Error stopping worker {worker_id}: {e}")
|
|
|
|
# Clear all managed processes from config
|
|
config = load_config()
|
|
config['managed_processes'] = {}
|
|
save_config(config)
|
|
|
|
def load_processes(self):
|
|
"""Load persisted process information from config."""
|
|
config = load_config()
|
|
managed_processes = config.get('managed_processes', {})
|
|
|
|
# Verify each saved process is still running
|
|
for worker_id, proc_info in managed_processes.items():
|
|
pid = proc_info.get('pid')
|
|
if pid and self._is_process_running(pid):
|
|
# Reconstruct process info
|
|
self.processes[worker_id] = {
|
|
'pid': pid,
|
|
'process': None, # Can't reconstruct subprocess object
|
|
'started_at': proc_info.get('started_at'),
|
|
'config': proc_info.get('config'),
|
|
'log_file': proc_info.get('log_file')
|
|
}
|
|
debug_log(f"[Distributed] Restored worker {worker_id} (PID: {pid})")
|
|
else:
|
|
if pid:
|
|
debug_log(f"[Distributed] Worker {worker_id} (PID: {pid}) is no longer running")
|
|
|
|
def save_processes(self):
|
|
"""Save process information to config."""
|
|
config = load_config()
|
|
|
|
# Create serializable version of process info
|
|
managed_processes = {}
|
|
for worker_id, proc_info in self.processes.items():
|
|
# Only save if process is running
|
|
is_running, _ = self._check_worker_process(worker_id, proc_info)
|
|
|
|
if is_running:
|
|
managed_processes[worker_id] = {
|
|
'pid': proc_info['pid'],
|
|
'started_at': proc_info['started_at'],
|
|
'config': proc_info['config'],
|
|
'log_file': proc_info.get('log_file'),
|
|
'launching': proc_info.get('launching', False)
|
|
}
|
|
|
|
# Update config with managed processes
|
|
config['managed_processes'] = managed_processes
|
|
save_config(config)
|
|
|
|
def _is_process_running(self, pid):
|
|
"""Check if a process with given PID is running."""
|
|
return is_process_alive(pid)
|
|
|
|
def _check_worker_process(self, worker_id, proc_info):
|
|
"""Check if a worker process is still running and return status.
|
|
|
|
Returns:
|
|
tuple: (is_running, has_subprocess_object)
|
|
"""
|
|
process = proc_info.get('process')
|
|
pid = proc_info.get('pid')
|
|
|
|
if process:
|
|
# Normal case with subprocess object
|
|
return process.poll() is None, True
|
|
elif pid:
|
|
# Restored process without subprocess object
|
|
return self._is_process_running(pid), False
|
|
else:
|
|
# No process or PID
|
|
return False, False
|
|
|
|
def _kill_process_tree(self, pid):
|
|
"""Kill a process and all its children."""
|
|
if PSUTIL_AVAILABLE:
|
|
try:
|
|
parent = psutil.Process(pid)
|
|
children = parent.children(recursive=True)
|
|
|
|
# Log what we're about to kill
|
|
debug_log(f"Killing process tree for PID {pid} ({parent.name()})")
|
|
for child in children:
|
|
debug_log(f" - Child PID {child.pid} ({child.name()})")
|
|
|
|
# Kill children first
|
|
for child in children:
|
|
try:
|
|
debug_log(f"Terminating child {child.pid}")
|
|
child.terminate()
|
|
except psutil.NoSuchProcess:
|
|
pass
|
|
|
|
# Wait a bit for graceful termination
|
|
gone, alive = psutil.wait_procs(children, timeout=PROCESS_WAIT_TIMEOUT)
|
|
|
|
# Force kill any remaining
|
|
for child in alive:
|
|
try:
|
|
debug_log(f"Force killing child {child.pid}")
|
|
child.kill()
|
|
except psutil.NoSuchProcess:
|
|
pass
|
|
|
|
# Finally kill the parent
|
|
try:
|
|
debug_log(f"Terminating parent {pid}")
|
|
parent.terminate()
|
|
parent.wait(timeout=PROCESS_WAIT_TIMEOUT)
|
|
except psutil.TimeoutExpired:
|
|
debug_log(f"Force killing parent {pid}")
|
|
parent.kill()
|
|
except psutil.NoSuchProcess:
|
|
debug_log(f"Parent process {pid} already gone")
|
|
|
|
return True
|
|
|
|
except psutil.NoSuchProcess:
|
|
debug_log(f"Process {pid} does not exist")
|
|
return False
|
|
except Exception as e:
|
|
debug_log(f"Error killing process tree: {e}")
|
|
# Fall through to OS commands
|
|
|
|
# Fallback to OS-specific commands
|
|
debug_log("[Distributed] Using OS commands to kill process tree")
|
|
if platform.system() == "Windows":
|
|
try:
|
|
# Use wmic to find child processes
|
|
result = subprocess.run(['wmic', 'process', 'where', f'ParentProcessId={pid}', 'get', 'ProcessId'],
|
|
capture_output=True, text=True)
|
|
if result.returncode == 0:
|
|
lines = result.stdout.strip().split('\n')[1:] # Skip header
|
|
child_pids = [line.strip() for line in lines if line.strip() and line.strip().isdigit()]
|
|
|
|
debug_log(f"[Distributed] Found child processes: {child_pids}")
|
|
|
|
# Kill each child
|
|
for child_pid in child_pids:
|
|
try:
|
|
subprocess.run(['taskkill', '/F', '/PID', child_pid],
|
|
capture_output=True, check=False)
|
|
except (OSError, FileNotFoundError) as e:
|
|
debug_log(f"[Distributed] Warning: taskkill failed for PID {child_pid}: {e}")
|
|
|
|
# Kill the parent with tree flag
|
|
result = subprocess.run(['taskkill', '/F', '/PID', str(pid), '/T'],
|
|
capture_output=True, text=True)
|
|
debug_log(f"[Distributed] Taskkill result: {result.stdout.strip()}")
|
|
return result.returncode == 0
|
|
except Exception as e:
|
|
log(f"[Distributed] Error with taskkill: {e}")
|
|
return False
|
|
else:
|
|
# Unix: use pkill
|
|
try:
|
|
subprocess.run(['pkill', '-TERM', '-P', str(pid)], check=False)
|
|
time.sleep(WORKER_CHECK_INTERVAL)
|
|
subprocess.run(['pkill', '-KILL', '-P', str(pid)], check=False)
|
|
os.kill(pid, signal.SIGKILL)
|
|
return True
|
|
except Exception as e:
|
|
log(f"[Distributed] Error killing process tree for PID {pid}: {e}")
|
|
return False
|