Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2ccfcafad1 | ||
|
|
3a51ac92e4 | ||
|
|
220251ffdb |
@@ -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
@@ -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):
|
||||
|
||||
@@ -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()
|
||||
@@ -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(
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user