5 changed files with 336 additions and 40 deletions
+8 -2
View File
@@ -54,10 +54,16 @@ Workers try those paths in order and use the first one that contains `ComfyUI\ma
## Reusing a worker's GUI session
Set these Deadline worker extra-info keys for each ComfyUI worker: `ComfyUIApiUrl` (for example `http://127.0.0.1:8188`) and `ComfyUILaunchGpuUuid` (an NVIDIA `GPU-...` UUID). The configured local endpoint is verified before a job queues its own prompt behind existing GUI work. Verification prefers the plugin's `/deadline/session` identity response and matching ComfyUI root; for an older already-running session, Windows port ownership, executable path, PID, and command line must match the configured installation. An unrelated service that happens to own the port is never reused. Deadline tracks only that prompt and never stops the GUI backend. If the endpoint is absent, fallback starts only on the configured port and resolves the configured NVIDIA UUID. The older `Per-Worker ComfyUI Endpoints` JSON setting remains a fallback for existing configurations.
Set these Deadline worker extra-info keys for each ComfyUI worker: `ComfyUIApiUrl` (for example `http://127.0.0.1:8188`) and `ComfyUILaunchGpuUuid` (an NVIDIA `GPU-...` UUID). The configured local endpoint is verified before a job queues its own prompt behind existing GUI work. Verification prefers the plugin's `/deadline/session` identity response and matching ComfyUI root; for an older already-running session, Windows port ownership, executable path, PID, and command line must match the configured installation. An unrelated service that happens to own the port is never reused. Deadline tracks only that prompt and never stops the GUI backend. If the endpoint is absent, fallback starts on the configured port and resolves the configured NVIDIA UUID. An occupied, unverified endpoint fails the task. The older `Per-Worker ComfyUI Endpoints` JSON setting remains a fallback for existing configurations.
For a verified reused endpoint, task inputs are copied into an isolated `input/deadline/<submission-id>` subtree and recorded outputs are copied to the job output directory.
## Multiple workers on one machine
Normal render jobs use a fixed port per CUDA device: `ComfyUIPort + 100 * GPU index`, with a default base of `8188`. GPU 0 uses `8188`, GPU 1 uses `8288`, and GPU 2 uses `8388`. Configure one active worker/task per GPU and enable each worker's Deadline GPU affinity for its assigned GPU. If a port is occupied, the task fails with an actionable error instead of switching ports. Port checks only bind a temporary socket; they create no lock folder and require no write access to `ProgramData`.
Each new launch receives a unique `--database-url` in Deadline's task temporary directory, avoiding the shared `ComfyUI/user/comfyui.db` lock. Workflow submission starts only after ComfyUI confirms that the selected listening port has bound. After deploying an update, restart workers when idle to clear cached plugin copies.
## Use It
Add `Submit to Deadline` to your workflow, set `output_directory` to a path the farm can see, and run the workflow in ComfyUI.
@@ -99,7 +105,7 @@ The worker writes the normal workflow metadata back into the output image when C
Every variation is checked on the selected Deadline worker after input staging, seed expansion, and fixed-switch resolution. The worker:
- validates every `class_type` against that worker's live `/object_info`;
- validates every `class_type` using that worker's `/object_info/<class_type>` endpoint, requesting only node types present in the workflow;
- rejects missing staged inputs and known browser-interactive nodes;
- resolves fixed selections from literal values or safe primitive value nodes, then removes `ImpactSwitch`, `LatentSwitch`, and `SEGSSwitch` pass-through nodes by reconnecting consumers to the selected upstream input;
- sends the complete submitted `extra_pnginfo.workflow` to metadata-aware nodes;
+62 -37
View File
@@ -20,6 +20,7 @@ import base64
import shutil
import copy
import uuid
from typing import Tuple
"""
@@ -34,7 +35,6 @@ Supports both existing ComfyUI instances and launching new ones.
# Constants
DEFAULT_PORT = 8188
PORT_OFFSET_PER_GPU = 100
MAX_PORT_SEARCH_RANGE = 100
DEFAULT_POLLING_INTERVAL = 10 # seconds
MAX_SEED_VALUE = 2147483647
PROGRESS_LOG_INTERVAL = 10 # Log every 10 polls
@@ -151,10 +151,8 @@ class ComfyUI(DeadlinePlugin):
def _setup_stdout_handlers(self):
"""Setup stdout handlers for ComfyUI output parsing"""
# Server startup handlers
self.AddStdoutHandlerCallback(".*Starting server.*").HandleCallback += self.HandleServerStarted
# Also catch the GUI message as backup
self.AddStdoutHandlerCallback(".*To see the GUI go to.*").HandleCallback += self.HandleServerStarted
# ComfyUI prints this only after the listening socket has bound.
self.AddStdoutHandlerCallback(r".*To see the GUI go to: http://\S+:(\d+).*").HandleCallback += self.HandleServerStarted
self.AddStdoutHandlerCallback(".*Error:.*").HandleCallback += self.HandleStdoutError
self.AddStdoutHandlerCallback(".*Exception:.*").HandleCallback += self.HandleStdoutError
@@ -661,9 +659,8 @@ class ComfyUI(DeadlinePlugin):
self.use_existing_comfyui, self.server_started = True, True
self.LogInfo(f"Reusing verified ComfyUI endpoint {endpoint}; its running session determines GPU and this task queues only its own prompts.")
return self.comfyui_port
if self._is_port_in_use(port):
raise ComfyUIError(f"Configured endpoint port {port} is occupied by an unverified service; refusing another port or replacement.")
self.LogInfo(f"Configured endpoint unavailable; fallback will start only at {endpoint}.")
self._check_launch_port(port)
self.LogInfo(f"Configured endpoint unavailable; starting isolated ComfyUI at {self.comfyui_api_url}.")
return self.comfyui_port
def _determine_final_port(self, base_port: int) -> str:
@@ -673,13 +670,14 @@ class ComfyUI(DeadlinePlugin):
self.use_existing_comfyui = False
if worker_mode or distributed_mode:
worker_port = self._calculate_worker_port(base_port)
self.comfyui_port = self._find_available_port(worker_port)
self.comfyui_port = str(worker_port)
self.LogInfo(f"Worker/distributed mode: starting isolated ComfyUI on port {self.comfyui_port}")
else:
self.comfyui_port = self._find_available_port(base_port)
self.comfyui_port = str(base_port)
self.LogInfo(f"Normal render mode: starting isolated ComfyUI on port {self.comfyui_port}")
self.comfyui_api_url = f"http://127.0.0.1:{self.comfyui_port}"
self._check_launch_port(int(self.comfyui_port))
return self.comfyui_port
@@ -707,6 +705,10 @@ class ComfyUI(DeadlinePlugin):
self.LogInfo("ComfyUI PreRenderTasks started.")
try:
self.server_started = False
self.workflow_submitted = False
self.thread_running = True
self.submission_error = ""
self._setup_batch_processing()
self._load_worker_endpoint_policy()
comfyui_path = self._resolve_comfyui_install_path()
@@ -727,30 +729,30 @@ class ComfyUI(DeadlinePlugin):
self.temp_dir = self.CreateTempDirectory("comfyui_job")
def _is_port_in_use(self, port: int) -> bool:
"""Check if the given port is already in use"""
"""Check whether a new server can bind the port."""
try:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.settimeout(1)
result = s.connect_ex(('127.0.0.1', port))
return result == 0
except:
if os.name == "nt":
s.setsockopt(socket.SOL_SOCKET, socket.SO_EXCLUSIVEADDRUSE, 1)
s.bind(('0.0.0.0', port))
return False
except OSError:
return True
def _find_available_port(self, start_port: int) -> str:
"""Find an available port starting from start_port"""
port = start_port
max_port = start_port + MAX_PORT_SEARCH_RANGE
while port < max_port:
if not self._is_port_in_use(port):
return str(port)
port += 1
return str(start_port)
def _check_launch_port(self, port: int):
if not 1 <= port <= 65535:
raise ComfyUIError(f"ComfyUI port {port} is outside the TCP port range.")
if self._is_port_in_use(port):
raise ComfyUIError(
f"ComfyUI port {port} is already in use on worker {self.GetSlaveName()}. "
"Stop the previous instance or set a different ComfyUIPort base. "
"Each GPU must have only one active Deadline worker/task."
)
def PostRenderTasks(self):
"""Cleanup tasks after rendering"""
self.LogInfo("ComfyUI PostRenderTasks started.")
self.thread_running = False
if not self.task_completed:
message = self.submission_error or (
@@ -835,6 +837,8 @@ class ComfyUI(DeadlinePlugin):
# If using existing ComfyUI instance, return dummy command
if self.use_existing_comfyui:
return self._create_dummy_command()
self._check_launch_port(int(self.comfyui_port))
# Build arguments for new ComfyUI instance
return self._build_comfyui_arguments(comfyui_main_py)
@@ -884,6 +888,9 @@ sys.exit(0 if result.get('success') else 1)
# Add main script and port
args_list.append(f'"{comfyui_main_py}"')
args_list.append(f"--port {port}")
database_path = os.path.abspath(os.path.join(self.temp_dir, f"comfyui-{uuid.uuid4().hex}.db"))
database_url = "sqlite:///" + database_path.replace("\\", "/")
args_list.append(f'--database-url "{database_url}"')
# Add CUDA device argument
cuda_arg = self._get_cuda_device_arg()
@@ -929,6 +936,8 @@ sys.exit(0 if result.get('success') else 1)
def HandleServerStarted(self):
"""Called when the ComfyUI server has started"""
if self.GetRegexMatch(1) != self.comfyui_port or not self.thread_running:
return
# Prevent multiple triggers from stdout handlers
if self.workflow_submitted:
self.LogInfo("ComfyUI server startup detected, but workflow already submitted - ignoring")
@@ -1448,7 +1457,7 @@ sys.exit(0 if result.get('success') else 1)
if not self.task_completed:
self.FailRender(f"ComfyUI error: {error_msg}")
def initialize_api_connection(self) -> bool:
def initialize_api_connection(self, workflow_data: dict) -> bool:
"""Initialize connection to ComfyUI API and get client ID"""
try:
time.sleep(2) # Wait for server to be fully initialized
@@ -1461,15 +1470,31 @@ sys.exit(0 if result.get('success') else 1)
self.client_id = f"deadline-{uuid.uuid4().hex}"
self.LogInfo(f"Got client ID: {self.client_id}")
object_info_response = self.http_request(f"{self.comfyui_api_url}/object_info")
if object_info_response['status_code'] != 200:
raise ComfyUIError(
f"Worker {self.GetSlaveName()} did not return /object_info "
f"(HTTP {object_info_response['status_code']})."
)
self.object_info = object_info_response['json']()
if not isinstance(self.object_info, dict) or not self.object_info:
raise ComfyUIError(f"Worker {self.GetSlaveName()} returned invalid /object_info data.")
node_types = sorted({node["class_type"] for node in workflow_data.values()
if isinstance(node, dict) and isinstance(node.get("class_type"), str) and node["class_type"]})
self.object_info = {}
for class_type in node_types:
url = f"{self.comfyui_api_url}/object_info/{urllib.parse.quote(class_type, safe='')}"
try:
object_info_response = self.http_request(url)
except OSError as e:
raise ComfyUIError(
f"Worker {self.GetSlaveName()} could not return metadata for node type "
f"'{class_type}' from {url}: {e}"
) from e
if object_info_response['status_code'] != 200:
raise ComfyUIError(
f"Worker {self.GetSlaveName()} did not return metadata for node type "
f"'{class_type}' (HTTP {object_info_response['status_code']})."
)
definitions = object_info_response['json']()
if not isinstance(definitions, dict):
raise ComfyUIError(f"Worker {self.GetSlaveName()} returned invalid metadata for '{class_type}'.")
if class_type in definitions:
if not isinstance(definitions[class_type], dict):
raise ComfyUIError(f"Worker {self.GetSlaveName()} returned invalid metadata for '{class_type}'.")
self.object_info[class_type] = definitions[class_type]
self.LogInfo(f"Loaded metadata for {len(self.object_info)} workflow node types on worker {self.GetSlaveName()}")
return True
except Exception as e:
self.LogWarning(f"Error initializing API connection: {e}")
@@ -2230,7 +2255,7 @@ sys.exit(0 if result.get('success') else 1)
self._load_standard_workflow_metadata()
if not self.initialize_api_connection():
if not self.initialize_api_connection(workflow_data):
raise ComfyUIError("Could not initialize the ComfyUI API connection.")
if not self.queue_workflow(workflow_data):
+107
View File
@@ -0,0 +1,107 @@
from __future__ import annotations
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
import json
import threading
import unittest
import urllib.parse
from unittest.mock import MagicMock, patch
from test_reuse_endpoint import ComfyUIError, new_plugin, plugin_module
class NodeMetadataTests(unittest.TestCase):
def test_api_requests_only_used_types_and_encodes_custom_node_names(self):
requests = []
definitions = {
"ImageResize+": {"input": {"required": {}}, "output_node": False},
"easy cleanGpuUsed": {"input": {"hidden": {"prompt": "PROMPT"}}, "output_node": False},
"SaveVideo": {"input": {"required": {}}, "output_node": True},
}
class Handler(BaseHTTPRequestHandler):
def do_GET(self):
requests.append(self.path)
if self.path == "/prompt":
data = {"exec_info": {"queue_remaining": 0}}
elif self.path.startswith("/object_info/"):
node_type = urllib.parse.unquote(self.path[len("/object_info/"):])
data = {node_type: definitions[node_type]} if node_type in definitions else {}
else:
self.send_error(500, "Global catalog must not be requested")
return
body = json.dumps(data).encode("utf-8")
self.send_response(200)
self.send_header("Content-Type", "application/json")
self.send_header("Content-Length", str(len(body)))
self.end_headers()
self.wfile.write(body)
def log_message(self, *_args):
pass
server = ThreadingHTTPServer(("127.0.0.1", 0), Handler)
thread = threading.Thread(target=server.serve_forever, daemon=True)
thread.start()
plugin = new_plugin(comfyui_api_url=f"http://127.0.0.1:{server.server_port}")
prompt = {
"1": {"class_type": "ImageResize+", "inputs": {}},
"2": {"class_type": "ImageResize+", "inputs": {}},
"3": {"class_type": "easy cleanGpuUsed", "inputs": {}},
"4": {"class_type": "SaveVideo", "inputs": {}},
}
try:
with patch.object(plugin_module.time, "sleep"):
self.assertTrue(plugin.initialize_api_connection(prompt))
self.assertEqual(requests, [
"/prompt", "/object_info/ImageResize%2B", "/object_info/SaveVideo",
"/object_info/easy%20cleanGpuUsed",
])
self.assertEqual(plugin.object_info, definitions)
self.assertEqual(plugin._preflight_worker_payload(prompt)[1], ["4"])
finally:
server.shutdown()
server.server_close()
thread.join(timeout=5)
def test_metadata_timeout_identifies_worker_and_requested_node(self):
request = MagicMock(side_effect=[{"status_code": 200}, TimeoutError("timed out")])
plugin = new_plugin(comfyui_api_url="http://127.0.0.1:8188", http_request=request)
prompt = {"1": {"class_type": "SlowNode", "inputs": {}}}
with patch.object(plugin_module.time, "sleep"):
with self.assertRaisesRegex(ComfyUIError, "M21.*'SlowNode'.*timed out"):
plugin.initialize_api_connection(prompt)
self.assertEqual(request.call_count, 2)
def test_missing_node_still_fails_preflight_with_node_context(self):
request = MagicMock(side_effect=[
{"status_code": 200},
{"status_code": 200, "json": lambda: {}},
{"status_code": 200, "json": lambda: {"SaveVideo": {"output_node": True}}},
])
plugin = new_plugin(comfyui_api_url="http://127.0.0.1:8188", http_request=request)
prompt = {
"1": {"class_type": "MissingOnWorker", "inputs": {}},
"2": {"class_type": "SaveVideo", "inputs": {}},
}
with patch.object(plugin_module.time, "sleep"):
self.assertTrue(plugin.initialize_api_connection(prompt))
with self.assertRaisesRegex(ComfyUIError, "node 1.*MissingOnWorker.*M21.*not installed"):
plugin._preflight_worker_payload(prompt)
def test_failed_or_invalid_metadata_is_rejected_with_node_name(self):
for response in (
{"status_code": 500},
{"status_code": 200, "json": lambda: []},
{"status_code": 200, "json": lambda: {"SaveVideo": "invalid"}},
):
with self.subTest(response=response):
request = MagicMock(side_effect=[{"status_code": 200}, response])
plugin = new_plugin(comfyui_api_url="http://127.0.0.1:8188", http_request=request)
with patch.object(plugin_module.time, "sleep"):
with self.assertRaisesRegex(ComfyUIError, "SaveVideo"):
plugin.initialize_api_connection({"1": {"class_type": "SaveVideo", "inputs": {}}})
if __name__ == "__main__":
unittest.main()
+30 -1
View File
@@ -13,6 +13,7 @@ import sys
import tempfile
import types
import unittest
from unittest.mock import patch, MagicMock
from pathlib import Path
@@ -251,8 +252,36 @@ class ReuseEndpointCurrentTests(unittest.TestCase):
http_request=lambda url, **kwargs: {"status_code": 503, "json": lambda: {}},
_is_port_in_use=lambda port: True,
)
with self.assertRaises(ComfyUIError):
with self.assertRaisesRegex(ComfyUIError, "8188 is already in use"):
plugin._configure_policy_endpoint()
self.assertEqual(plugin.comfyui_api_url, "http://127.0.0.1:8188")
self.assertFalse(plugin.use_existing_comfyui)
def test_bound_port_without_listener_is_unavailable(self):
with plugin_module.socket.socket() as sock:
sock.bind(("127.0.0.1", 0))
plugin = new_plugin()
self.assertTrue(plugin._is_port_in_use(sock.getsockname()[1]))
def test_access_denied_port_is_unavailable(self):
sock = MagicMock()
sock.__enter__.return_value = sock
sock.bind.side_effect = PermissionError("port reserved")
with patch.object(plugin_module.socket, "socket", return_value=sock):
self.assertTrue(new_plugin()._is_port_in_use(8188))
def test_render_arguments_recheck_fixed_port_before_launch(self):
plugin = new_plugin(
comfyui_port="8188",
_resolve_comfyui_install_path=lambda: "C:/Comfy",
_is_port_in_use=lambda port: True,
_build_comfyui_arguments=MagicMock(),
)
with patch.object(plugin_module.os.path, "exists", return_value=True):
with self.assertRaisesRegex(ComfyUIError, "8188 is already in use"):
plugin.RenderArgument()
self.assertEqual(plugin.comfyui_port, "8188")
plugin._build_comfyui_arguments.assert_not_called()
def test_fallback_launch_resolves_uuid_and_honors_deadline_affinity(self):
plugin = new_plugin(
+129
View File
@@ -0,0 +1,129 @@
from __future__ import annotations
import os
from pathlib import Path
import re
import socket
import tempfile
import types
import unittest
from unittest.mock import MagicMock, patch
from test_reuse_endpoint import ComfyUIError, new_plugin, plugin_module
class WorkerStartupTests(unittest.TestCase):
def test_worker_gpu_affinity_drives_both_fixed_port_and_launch_device(self):
ports = []
with tempfile.TemporaryDirectory() as temp:
for gpu in (0, 1, 4, 7):
with self.subTest(gpu=gpu):
plugin = new_plugin(
OverrideGpuAffinity=lambda: True,
GpuAffinity=lambda: [gpu],
GetThreadNumber=lambda: 0,
_is_port_in_use=lambda port: False,
temp_dir=temp,
)
port = plugin._calculate_comfyui_port()
self.assertEqual(port, str(8188 + gpu * 100))
self.assertEqual(plugin.comfyui_api_url, "http://127.0.0.1:" + port)
arguments = plugin._build_comfyui_arguments("C:/Comfy/ComfyUI/main.py")
self.assertIn(f"--port {port}", arguments)
self.assertIn(f"--cuda-device {gpu}", arguments)
ports.append(port)
self.assertEqual(len(set(ports)), 4)
def test_occupied_gpu_port_fails_without_switching_ports(self):
with socket.socket() as sock:
sock.bind(("127.0.0.1", 0))
port = sock.getsockname()[1]
plugin = new_plugin(
OverrideGpuAffinity=lambda: True,
GpuAffinity=lambda: [0],
GetThreadNumber=lambda: 0,
GetPluginInfoEntryWithDefault=lambda key, default: str(port) if key == "ComfyUIPort" else default,
)
with self.assertRaisesRegex(ComfyUIError, f"{port} is already in use"):
plugin._calculate_comfyui_port()
self.assertEqual(plugin.comfyui_port, str(port))
def test_fixed_ports_need_no_lock_folder(self):
plugin = new_plugin(_get_cuda_device_arg=lambda: "--cuda-device 1", _is_port_in_use=lambda port: False)
with patch.object(plugin_module.os, "makedirs", side_effect=PermissionError("access denied")) as mkdir, \
patch("builtins.open", side_effect=PermissionError("access denied")) as open_file:
self.assertEqual(plugin._calculate_comfyui_port(), "8288")
mkdir.assert_not_called()
open_file.assert_not_called()
def test_invalid_fixed_port_fails_before_launch(self):
for port in (0, 65536):
with self.subTest(port=port):
with self.assertRaisesRegex(ComfyUIError, "outside the TCP port range"):
new_plugin()._check_launch_port(port)
def test_new_instances_get_separate_task_databases(self):
with tempfile.TemporaryDirectory(prefix="deadline job ") as temp:
plugin = new_plugin(temp_dir=temp, _get_cuda_device_arg=lambda: "--cuda-device 0")
urls = []
for _ in range(2):
arguments = plugin._build_comfyui_arguments("C:/Comfy/ComfyUI/main.py")
match = re.search(r'--database-url "([^"]+)"', arguments)
self.assertIsNotNone(match)
urls.append(match.group(1))
path = urls[-1].removeprefix("sqlite:///")
self.assertEqual(Path(path).parent, Path(temp))
self.assertIn("--cuda-device 0", arguments)
self.assertNotEqual(*urls)
def test_starting_server_and_bind_failure_do_not_trigger_submission(self):
handlers = {}
def add_handler(pattern):
handlers[pattern] = types.SimpleNamespace(HandleCallback=MagicMock())
return handlers[pattern]
plugin = new_plugin(AddStdoutHandlerCallback=add_handler)
plugin._setup_stdout_handlers()
startup_patterns = [pattern for pattern in handlers if "To see the GUI" in pattern]
self.assertEqual(len(startup_patterns), 1)
pattern = startup_patterns[0]
for line in ("Starting server", "Port 8188 is already in use on address 127.0.0.1"):
self.assertIsNone(re.match(pattern, line))
match = re.match(pattern, "[INFO] To see the GUI go to: http://127.0.0.1:8188")
self.assertEqual(match.group(1), "8188")
plugin.comfyui_port = "8188"
plugin.GetRegexMatch = match.group
with patch.object(plugin_module.threading, "Thread") as thread:
plugin.HandleServerStarted()
plugin.HandleServerStarted()
thread.assert_called_once_with(target=plugin.submit_workflow)
thread.return_value.start.assert_called_once()
self.assertTrue(plugin.server_started)
self.assertTrue(plugin.workflow_submitted)
def test_wrong_port_or_stopped_task_does_not_trigger_submission(self):
plugin = new_plugin(comfyui_port="8288", GetRegexMatch=lambda group: "8188")
with patch.object(plugin_module.threading, "Thread") as thread:
plugin.HandleServerStarted()
plugin.GetRegexMatch = lambda group: "8288"
plugin.thread_running = False
plugin.HandleServerStarted()
thread.assert_not_called()
self.assertFalse(plugin.server_started)
self.assertFalse(plugin.workflow_submitted)
def test_reused_instance_launch_does_not_check_its_occupied_port_or_build_arguments(self):
plugin = new_plugin(
use_existing_comfyui=True,
_resolve_comfyui_install_path=lambda: "C:/Comfy",
_create_dummy_command=lambda: "reuse waiter",
_check_launch_port=MagicMock(),
_build_comfyui_arguments=MagicMock(),
)
with patch.object(plugin_module.os.path, "exists", return_value=True):
self.assertEqual(plugin.RenderArgument(), "reuse waiter")
plugin._check_launch_port.assert_not_called()
plugin._build_comfyui_arguments.assert_not_called()
if __name__ == "__main__":
unittest.main()