feat: worker-thread setting speeds up the output scan and thumbnail generation
Both jobs processed one file at a time on a single core. Pillow decodes the whole image to reach trailing PNG text chunks and releases the GIL while it does, so a bounded thread pool scales almost linearly: on a 1 GB sample of 400 PNGs the metadata pass dropped from 18 s to 2.5 s at 8 threads, and 120 thumbnails from 4.6 s to 0.6 s. - utils/parallel.py: map_parallel(func, items, workers) keeps order, runs on the caller's thread when workers is 1, and caps the pool at the item count - PromptManagerConfig.WORKER_THREADS defaults to half the detected cores, is clamped to max_worker_threads() when loaded from config.json, and is persisted under "performance" - settings API returns worker_threads and cpu_count and validates saves; the settings modal gets a "Worker threads" slider capped at the core count - the scan reads each metadata batch on the configured threads; thumbnail generation processes batches of files per executor job, still emitting one progress event per file Threads rather than processes: this runs inside ComfyUI's server, where forking a CUDA-initialised process is unsafe and the spawn start method would re-import ComfyUI's main module in every child.
This commit is contained in:
+36
-11
@@ -15,9 +15,11 @@ from aiohttp import web
|
||||
from .scan_job import ScanJob
|
||||
|
||||
try:
|
||||
from ...utils.validators import validate_result_timeout
|
||||
from ...utils.parallel import map_parallel
|
||||
from ...utils.validators import validate_result_timeout, validate_worker_threads
|
||||
except ImportError:
|
||||
from utils.validators import validate_result_timeout
|
||||
from utils.parallel import map_parallel
|
||||
from utils.validators import validate_result_timeout, validate_worker_threads
|
||||
|
||||
try:
|
||||
from ...database.operations import PromptDatabase
|
||||
@@ -601,6 +603,8 @@ class AdminRoutesMixin:
|
||||
"settings": {
|
||||
"result_timeout": PromptManagerConfig.RESULT_TIMEOUT,
|
||||
"webui_display_mode": PromptManagerConfig.WEBUI_DISPLAY_MODE,
|
||||
"worker_threads": PromptManagerConfig.WORKER_THREADS,
|
||||
"cpu_count": PromptManagerConfig.max_worker_threads(),
|
||||
"gallery_root_paths": root_paths,
|
||||
"gallery_root_path": root_paths[0] if root_paths else "",
|
||||
"monitored_directories": [
|
||||
@@ -645,6 +649,15 @@ class AdminRoutesMixin:
|
||||
)
|
||||
if "webui_display_mode" in data:
|
||||
PromptManagerConfig.WEBUI_DISPLAY_MODE = data["webui_display_mode"]
|
||||
if "worker_threads" in data:
|
||||
try:
|
||||
PromptManagerConfig.WORKER_THREADS = validate_worker_threads(
|
||||
data["worker_threads"], PromptManagerConfig.max_worker_threads()
|
||||
)
|
||||
except ValueError as ve:
|
||||
return web.json_response(
|
||||
{"success": False, "error": str(ve)}, status=400
|
||||
)
|
||||
|
||||
new_roots, error = self._parse_gallery_roots(data)
|
||||
if error:
|
||||
@@ -738,7 +751,14 @@ class AdminRoutesMixin:
|
||||
monitoring = dict(gallery.get("monitoring") or {})
|
||||
monitoring["directories"] = list(GalleryConfig.MONITORING_DIRECTORIES)
|
||||
gallery["monitoring"] = monitoring
|
||||
config_data = {**existing, "web_ui": web_ui, "gallery": gallery}
|
||||
performance = dict(existing.get("performance") or {})
|
||||
performance["worker_threads"] = PromptManagerConfig.WORKER_THREADS
|
||||
config_data = {
|
||||
**existing,
|
||||
"web_ui": web_ui,
|
||||
"gallery": gallery,
|
||||
"performance": performance,
|
||||
}
|
||||
try:
|
||||
write_private_json(config_file, config_data)
|
||||
self.logger.info(f"Settings saved to {config_file}")
|
||||
@@ -1255,16 +1275,21 @@ class AdminRoutesMixin:
|
||||
return _collect_media_files(output_dirs, IMAGE_EXTENSIONS + VIDEO_EXTENSIONS)
|
||||
|
||||
def _extract_batch_metadata_sync(self, file_batch):
|
||||
"""Extract metadata from a batch of files in one executor call."""
|
||||
results = []
|
||||
for media_file in file_batch:
|
||||
"""Extract metadata from a batch of files in one executor call.
|
||||
|
||||
Files are read on ``PromptManagerConfig.WORKER_THREADS`` threads:
|
||||
Pillow decodes the whole image to reach trailing text chunks and
|
||||
releases the GIL while doing so.
|
||||
"""
|
||||
from ..config import PromptManagerConfig
|
||||
|
||||
def extract(media_file):
|
||||
try:
|
||||
results.append(
|
||||
(media_file, self._extract_comfyui_metadata(str(media_file)))
|
||||
)
|
||||
return (media_file, self._extract_comfyui_metadata(str(media_file)))
|
||||
except Exception:
|
||||
results.append((media_file, {}))
|
||||
return results
|
||||
return (media_file, {})
|
||||
|
||||
return map_parallel(extract, file_batch, PromptManagerConfig.WORKER_THREADS)
|
||||
|
||||
def _ingest_scan_batch_sync(self, batch_results):
|
||||
"""Store the prompts found in a metadata batch (blocking DB work)."""
|
||||
|
||||
+57
-22
@@ -15,6 +15,11 @@ from PIL import Image, UnidentifiedImageError
|
||||
from .prompts import bad_request, parse_page_params, publish_image_paths
|
||||
from .prompts import safe_error_message as _safe_error
|
||||
|
||||
try:
|
||||
from ...utils.parallel import map_parallel
|
||||
except ImportError:
|
||||
from utils.parallel import map_parallel
|
||||
|
||||
# Only these file types are ever served by the image routes, regardless of
|
||||
# what sits inside an allowed directory (defence in depth for the gallery).
|
||||
IMAGE_EXTENSIONS = frozenset(
|
||||
@@ -91,6 +96,21 @@ MAX_SCAN_DEPTH = 12
|
||||
MAX_THUMBNAIL_FILES = 10_000
|
||||
|
||||
|
||||
def _worker_threads():
|
||||
"""Configured thread count for thumbnail work.
|
||||
|
||||
Imported lazily because py.config imports this package.
|
||||
"""
|
||||
from ..config import PromptManagerConfig
|
||||
|
||||
return PromptManagerConfig.WORKER_THREADS
|
||||
|
||||
|
||||
def _thumbnail_batch_size(workers):
|
||||
"""Files per executor job: enough to keep every worker thread busy."""
|
||||
return max(4, workers * 2)
|
||||
|
||||
|
||||
def _write_image_thumbnail(src, dst, thumbnail_size):
|
||||
"""Resize *src* into *dst* (blocking PIL work)."""
|
||||
with Image.open(src) as img:
|
||||
@@ -754,10 +774,8 @@ class ImageRoutesMixin:
|
||||
start = _time.monotonic()
|
||||
targets = self._thumbnail_targets(output_path, thumbnails_dir)
|
||||
stats = _new_thumbnail_stats()
|
||||
for src, dst, is_video in targets:
|
||||
_record_thumbnail_result(
|
||||
stats, self._generate_one(src, dst, thumbnail_size, is_video)
|
||||
)
|
||||
for result in self._generate_batch(targets, thumbnail_size):
|
||||
_record_thumbnail_result(stats, result)
|
||||
payload = _thumbnail_complete_payload(len(targets), stats, start)
|
||||
payload["success"] = True
|
||||
payload["thumbnails_path"] = _public_path(thumbnails_dir)
|
||||
@@ -786,6 +804,18 @@ class ImageRoutesMixin:
|
||||
self.logger.warning(f"Thumbnail scan capped at {MAX_THUMBNAIL_FILES}")
|
||||
return targets
|
||||
|
||||
def _generate_batch(self, targets, thumbnail_size):
|
||||
"""Create thumbnails for ``targets`` on the configured worker threads.
|
||||
|
||||
Blocking; returns one result dict per target, in order.
|
||||
"""
|
||||
|
||||
def one(target):
|
||||
src, dst, is_video = target
|
||||
return self._generate_one(src, dst, thumbnail_size, is_video)
|
||||
|
||||
return map_parallel(one, targets, _worker_threads())
|
||||
|
||||
def _generate_one(self, src, dst, thumbnail_size, is_video):
|
||||
"""Create one thumbnail (blocking). Never raises; returns a result dict."""
|
||||
result = {
|
||||
@@ -861,7 +891,7 @@ class ImageRoutesMixin:
|
||||
return response
|
||||
|
||||
async def _stream_thumbnails(self, response, targets, thumbnail_size):
|
||||
"""Process *targets* one per executor job, emitting SSE after each."""
|
||||
"""Process *targets* in worker-thread batches, emitting SSE per file."""
|
||||
total = len(targets)
|
||||
video_count = sum(1 for _, _, is_video in targets if is_video)
|
||||
connected = await self._emit_sse(
|
||||
@@ -882,26 +912,31 @@ class ImageRoutesMixin:
|
||||
return
|
||||
stats = _new_thumbnail_stats()
|
||||
start = _time.monotonic()
|
||||
for index, (src, dst, is_video) in enumerate(targets, start=1):
|
||||
result = await self._run_in_executor(
|
||||
self._generate_one, src, dst, thumbnail_size, is_video
|
||||
batch_size = _thumbnail_batch_size(_worker_threads())
|
||||
index = 0
|
||||
for batch_start in range(0, total, batch_size):
|
||||
batch = targets[batch_start : batch_start + batch_size]
|
||||
results = await self._run_in_executor(
|
||||
self._generate_batch, batch, thumbnail_size
|
||||
)
|
||||
_record_thumbnail_result(stats, result)
|
||||
if result["action"] == "error":
|
||||
await self._emit_sse(
|
||||
for result in results:
|
||||
index += 1
|
||||
_record_thumbnail_result(stats, result)
|
||||
if result["action"] == "error":
|
||||
await self._emit_sse(
|
||||
response,
|
||||
"file_error",
|
||||
{"file": result["file"], "error": result["error"]},
|
||||
)
|
||||
connected = await self._emit_sse(
|
||||
response,
|
||||
"file_error",
|
||||
{"file": result["file"], "error": result["error"]},
|
||||
"progress",
|
||||
_thumbnail_progress_payload(index, total, stats, result, start),
|
||||
)
|
||||
connected = await self._emit_sse(
|
||||
response,
|
||||
"progress",
|
||||
_thumbnail_progress_payload(index, total, stats, result, start),
|
||||
)
|
||||
if not connected:
|
||||
self.logger.info("Thumbnail client disconnected; stopping early")
|
||||
return
|
||||
await asyncio.sleep(0) # let other requests run between files
|
||||
if not connected:
|
||||
self.logger.info("Thumbnail client disconnected; stopping early")
|
||||
return
|
||||
await asyncio.sleep(0) # let other requests run between batches
|
||||
payload = _thumbnail_complete_payload(total, stats, start)
|
||||
await self._emit_sse(response, "complete", payload)
|
||||
self.logger.info(payload["message"])
|
||||
|
||||
@@ -504,6 +504,29 @@ class PromptManagerConfig:
|
||||
MAX_SEARCH_RESULTS = 100
|
||||
ENABLE_FUZZY_SEARCH = False # Requires fuzzywuzzy
|
||||
AUTO_BACKUP_INTERVAL = 24 # Hours
|
||||
# Threads for the output scan and thumbnail generation; half the cores by
|
||||
# default so ComfyUI's own generation keeps CPU headroom. 1..max_worker_threads().
|
||||
WORKER_THREADS = max(1, (os.cpu_count() or 1) // 2)
|
||||
|
||||
@classmethod
|
||||
def max_worker_threads(cls) -> int:
|
||||
"""Largest accepted worker-thread count: the detected core count."""
|
||||
return os.cpu_count() or 1
|
||||
|
||||
@classmethod
|
||||
def clamp_worker_threads(cls, value, fallback: int) -> int:
|
||||
"""Coerce a worker count from config.json into ``[1, max_worker_threads()]``.
|
||||
|
||||
Non-numeric values (and booleans) return ``fallback`` so a hand-edited
|
||||
file cannot disable the setting.
|
||||
"""
|
||||
if isinstance(value, bool):
|
||||
return fallback
|
||||
if isinstance(value, str) and value.strip().isdigit():
|
||||
value = int(value.strip())
|
||||
if not isinstance(value, int):
|
||||
return fallback
|
||||
return max(1, min(value, cls.max_worker_threads()))
|
||||
|
||||
@classmethod
|
||||
def get_config(cls) -> Dict[str, Any]:
|
||||
@@ -536,6 +559,7 @@ class PromptManagerConfig:
|
||||
"max_search_results": cls.MAX_SEARCH_RESULTS,
|
||||
"enable_fuzzy_search": cls.ENABLE_FUZZY_SEARCH,
|
||||
"auto_backup_interval": cls.AUTO_BACKUP_INTERVAL,
|
||||
"worker_threads": cls.WORKER_THREADS,
|
||||
},
|
||||
"gallery": GalleryConfig.get_config(),
|
||||
"integrations": IntegrationConfig.get_config(),
|
||||
@@ -664,6 +688,10 @@ class PromptManagerConfig:
|
||||
cls.ENABLE_FUZZY_SEARCH = performance["enable_fuzzy_search"]
|
||||
if "auto_backup_interval" in performance:
|
||||
cls.AUTO_BACKUP_INTERVAL = performance["auto_backup_interval"]
|
||||
if "worker_threads" in performance:
|
||||
cls.WORKER_THREADS = cls.clamp_worker_threads(
|
||||
performance["worker_threads"], cls.WORKER_THREADS
|
||||
)
|
||||
|
||||
# Update gallery config
|
||||
if "gallery" in new_config:
|
||||
|
||||
@@ -1629,5 +1629,79 @@ class TestOutputScanIsABackgroundJob(AdminAPITestCase):
|
||||
self.assertEqual(data["last_event"]["type"], "error")
|
||||
|
||||
|
||||
class TestWorkerThreadsSetting(AdminAPITestCase):
|
||||
"""The worker-thread count is exposed, validated, persisted and used by the scan."""
|
||||
|
||||
async def test_get_settings_reports_worker_threads_and_cpu_count(self):
|
||||
data = await (
|
||||
await self.client.request("GET", "/prompt_manager/settings")
|
||||
).json()
|
||||
settings = data["settings"]
|
||||
self.assertEqual(
|
||||
settings["cpu_count"], PromptManagerConfig.max_worker_threads()
|
||||
)
|
||||
self.assertEqual(settings["worker_threads"], PromptManagerConfig.WORKER_THREADS)
|
||||
self.assertGreaterEqual(settings["worker_threads"], 1)
|
||||
|
||||
async def test_save_persists_worker_threads_under_performance(self):
|
||||
self._write_config({"performance": {"max_search_results": 25}})
|
||||
|
||||
resp = await self.client.request(
|
||||
"POST", "/prompt_manager/settings", json={"worker_threads": 1}
|
||||
)
|
||||
|
||||
self.assertEqual(resp.status, 200)
|
||||
self.assertEqual(PromptManagerConfig.WORKER_THREADS, 1)
|
||||
saved = self._read_config()
|
||||
self.assertEqual(saved["performance"]["worker_threads"], 1)
|
||||
self.assertEqual(saved["performance"]["max_search_results"], 25)
|
||||
|
||||
async def test_save_rejects_more_threads_than_cores(self):
|
||||
before = PromptManagerConfig.WORKER_THREADS
|
||||
too_many = PromptManagerConfig.max_worker_threads() + 1
|
||||
|
||||
resp = await self.client.request(
|
||||
"POST", "/prompt_manager/settings", json={"worker_threads": too_many}
|
||||
)
|
||||
|
||||
self.assertEqual(resp.status, 400)
|
||||
self.assertFalse((await resp.json())["success"])
|
||||
self.assertEqual(PromptManagerConfig.WORKER_THREADS, before)
|
||||
|
||||
async def test_scan_reads_metadata_on_several_threads(self):
|
||||
import threading
|
||||
import time
|
||||
|
||||
for i in range(8):
|
||||
_write_png(self.output_dir / f"gen{i}.png", prompt_text=f"prompt {i}")
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
|
||||
PromptManagerConfig.WORKER_THREADS = 4
|
||||
seen = set()
|
||||
lock = threading.Lock()
|
||||
original = self.api._extract_comfyui_metadata
|
||||
|
||||
def recording(path):
|
||||
with lock:
|
||||
seen.add(threading.current_thread().name)
|
||||
time.sleep(0.05)
|
||||
return original(path)
|
||||
|
||||
self.api._extract_comfyui_metadata = recording
|
||||
|
||||
body = await (await self.client.request("POST", "/prompt_manager/scan")).text()
|
||||
|
||||
self.assertIn('"added": 8', body)
|
||||
self.assertGreaterEqual(len(seen), 2)
|
||||
|
||||
async def test_scan_with_one_worker_still_completes(self):
|
||||
_write_png(self.output_dir / "gen.png", prompt_text="solo")
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [str(self.output_dir)]
|
||||
PromptManagerConfig.WORKER_THREADS = 1
|
||||
|
||||
body = await (await self.client.request("POST", "/prompt_manager/scan")).text()
|
||||
|
||||
self.assertIn('"added": 1', body)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -380,6 +380,64 @@ class ThumbnailTestCase(ImageAPITestCase):
|
||||
class TestThumbnailProgressStream(ThumbnailTestCase):
|
||||
"""POST /prompt_manager/images/generate-thumbnails/progress (SSE)."""
|
||||
|
||||
def _record_threads(self):
|
||||
"""Make every thumbnail write record its thread and take a moment."""
|
||||
import threading
|
||||
import time
|
||||
|
||||
from py.api import images as images_module
|
||||
|
||||
seen = set()
|
||||
lock = threading.Lock()
|
||||
original = images_module._write_image_thumbnail
|
||||
|
||||
def recording(src, dst, size):
|
||||
with lock:
|
||||
seen.add(threading.current_thread().name)
|
||||
time.sleep(0.05)
|
||||
return original(src, dst, size)
|
||||
|
||||
images_module._write_image_thumbnail = recording
|
||||
self.addCleanup(setattr, images_module, "_write_image_thumbnail", original)
|
||||
return seen
|
||||
|
||||
async def test_thumbnails_are_generated_on_several_threads(self):
|
||||
from py.config import PromptManagerConfig
|
||||
|
||||
for i in range(8):
|
||||
make_png(self.output_dir / f"p{i}.png", size=(64, 64))
|
||||
orig = PromptManagerConfig.WORKER_THREADS
|
||||
PromptManagerConfig.WORKER_THREADS = 4
|
||||
self.addCleanup(setattr, PromptManagerConfig, "WORKER_THREADS", orig)
|
||||
seen = self._record_threads()
|
||||
|
||||
resp = await self.client.request("POST", self.PROGRESS_URL)
|
||||
events = parse_sse(await resp.text())
|
||||
|
||||
names = [e for e, _ in events]
|
||||
self.assertEqual(names.count("progress"), 8)
|
||||
self.assertEqual(dict(events)["complete"]["count"], 8)
|
||||
self.assertGreaterEqual(len(seen), 2)
|
||||
for i in range(8):
|
||||
self.assertTrue(self._thumb_path(f"p{i}").is_file())
|
||||
|
||||
async def test_blocking_endpoint_also_uses_the_worker_threads(self):
|
||||
from py.config import PromptManagerConfig
|
||||
|
||||
for i in range(8):
|
||||
make_png(self.output_dir / f"p{i}.png", size=(64, 64))
|
||||
orig = PromptManagerConfig.WORKER_THREADS
|
||||
PromptManagerConfig.WORKER_THREADS = 4
|
||||
self.addCleanup(setattr, PromptManagerConfig, "WORKER_THREADS", orig)
|
||||
seen = self._record_threads()
|
||||
|
||||
resp = await self.client.request(
|
||||
"POST", "/prompt_manager/images/generate-thumbnails", json={}
|
||||
)
|
||||
|
||||
self.assertEqual((await resp.json())["count"], 8)
|
||||
self.assertGreaterEqual(len(seen), 2)
|
||||
|
||||
async def test_get_is_method_not_allowed(self):
|
||||
# The stream writes thumbnails, so a plain GET (prefetch, link
|
||||
# preview, <img src>) must not be able to trigger it.
|
||||
@@ -409,8 +467,9 @@ class TestThumbnailProgressStream(ThumbnailTestCase):
|
||||
self.assertTrue(self._thumb_path(name).is_file(), name)
|
||||
with Image.open(self._thumb_path("a")) as thumb:
|
||||
self.assertLessEqual(max(thumb.size), 150)
|
||||
# PIL work ran in the executor, once per file, plus the scan.
|
||||
self.assertEqual(calls.count("_generate_one"), 3)
|
||||
# PIL work ran in the executor in batches (never on the loop), plus the scan.
|
||||
self.assertGreaterEqual(calls.count("_generate_batch"), 1)
|
||||
self.assertNotIn("_generate_one", calls)
|
||||
self.assertIn("_thumbnail_targets", calls)
|
||||
|
||||
async def test_corrupt_png_is_reported_and_loop_continues(self):
|
||||
|
||||
@@ -136,6 +136,32 @@ class TestPromptManagerConfig(unittest.TestCase):
|
||||
PromptManagerConfig.update_config({"performance": {"max_search_results": 50}})
|
||||
self.assertEqual(PromptManagerConfig.MAX_SEARCH_RESULTS, 50)
|
||||
|
||||
def test_worker_threads_default_is_within_the_machine(self):
|
||||
cores = PromptManagerConfig.max_worker_threads()
|
||||
self.assertGreaterEqual(cores, 1)
|
||||
default = PromptManagerConfig.get_config()["performance"]["worker_threads"]
|
||||
self.assertGreaterEqual(default, 1)
|
||||
self.assertLessEqual(default, cores)
|
||||
self.assertEqual(default, max(1, cores // 2))
|
||||
|
||||
def test_update_worker_threads_is_clamped_to_the_machine(self):
|
||||
cores = PromptManagerConfig.max_worker_threads()
|
||||
PromptManagerConfig.update_config({"performance": {"worker_threads": 1}})
|
||||
self.assertEqual(PromptManagerConfig.WORKER_THREADS, 1)
|
||||
PromptManagerConfig.update_config(
|
||||
{"performance": {"worker_threads": cores + 50}}
|
||||
)
|
||||
self.assertEqual(PromptManagerConfig.WORKER_THREADS, cores)
|
||||
PromptManagerConfig.update_config({"performance": {"worker_threads": 0}})
|
||||
self.assertEqual(PromptManagerConfig.WORKER_THREADS, 1)
|
||||
|
||||
def test_update_worker_threads_ignores_junk_from_a_hand_edited_file(self):
|
||||
PromptManagerConfig.update_config({"performance": {"worker_threads": 1}})
|
||||
PromptManagerConfig.update_config({"performance": {"worker_threads": "lots"}})
|
||||
self.assertEqual(PromptManagerConfig.WORKER_THREADS, 1)
|
||||
PromptManagerConfig.update_config({"performance": {"worker_threads": True}})
|
||||
self.assertEqual(PromptManagerConfig.WORKER_THREADS, 1)
|
||||
|
||||
def test_update_propagates_to_gallery(self):
|
||||
PromptManagerConfig.update_config(
|
||||
{"gallery": {"monitoring": {"enabled": False}}}
|
||||
|
||||
@@ -0,0 +1,61 @@
|
||||
"""Tests for utils/parallel.py: bounded thread-pool map used by long jobs."""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
import time
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from utils.parallel import map_parallel # noqa: E402
|
||||
|
||||
|
||||
class TestMapParallel(unittest.TestCase):
|
||||
|
||||
def test_results_keep_input_order(self):
|
||||
def slow_square(n):
|
||||
time.sleep(0.01 * (5 - n)) # later items finish first
|
||||
return n * n
|
||||
|
||||
self.assertEqual(
|
||||
map_parallel(slow_square, [1, 2, 3, 4], workers=4), [1, 4, 9, 16]
|
||||
)
|
||||
|
||||
def test_one_worker_runs_on_the_calling_thread(self):
|
||||
names = []
|
||||
map_parallel(
|
||||
lambda _: names.append(threading.current_thread().name), [1, 2, 3], 1
|
||||
)
|
||||
self.assertEqual(set(names), {threading.current_thread().name})
|
||||
|
||||
def test_several_workers_use_several_threads(self):
|
||||
names = set()
|
||||
lock = threading.Lock()
|
||||
|
||||
def record(_):
|
||||
with lock:
|
||||
names.add(threading.current_thread().name)
|
||||
time.sleep(0.05)
|
||||
|
||||
map_parallel(record, range(8), workers=4)
|
||||
self.assertGreaterEqual(len(names), 2)
|
||||
self.assertNotIn(threading.current_thread().name, names)
|
||||
|
||||
def test_worker_count_is_capped_by_item_count_and_floor_one(self):
|
||||
self.assertEqual(map_parallel(lambda n: n + 1, [1], workers=16), [2])
|
||||
self.assertEqual(map_parallel(lambda n: n + 1, [1, 2], workers=0), [2, 3])
|
||||
self.assertEqual(map_parallel(lambda n: n, [], workers=4), [])
|
||||
|
||||
def test_exception_propagates_to_the_caller(self):
|
||||
def boom(n):
|
||||
if n == 2:
|
||||
raise ValueError("two")
|
||||
return n
|
||||
|
||||
with self.assertRaises(ValueError):
|
||||
map_parallel(boom, [1, 2, 3], workers=3)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -19,6 +19,7 @@ from utils.validators import (
|
||||
sanitize_input,
|
||||
parse_tags_string,
|
||||
validate_result_timeout,
|
||||
validate_worker_threads,
|
||||
)
|
||||
|
||||
|
||||
@@ -293,5 +294,24 @@ class TestValidateResultTimeout(unittest.TestCase):
|
||||
validate_result_timeout(True)
|
||||
|
||||
|
||||
class TestValidateWorkerThreads(unittest.TestCase):
|
||||
"""validate_worker_threads(value, max_workers) -> int in [1, max_workers]."""
|
||||
|
||||
def test_accepts_ints_and_digit_strings_within_range(self):
|
||||
self.assertEqual(validate_worker_threads(1, 8), 1)
|
||||
self.assertEqual(validate_worker_threads(8, 8), 8)
|
||||
self.assertEqual(validate_worker_threads("4", 8), 4)
|
||||
|
||||
def test_rejects_out_of_range_bool_and_junk(self):
|
||||
for value in (0, -1, 9, True, False, "x", None, 2.5, [2]):
|
||||
with self.assertRaises(ValueError, msg=repr(value)):
|
||||
validate_worker_threads(value, 8)
|
||||
|
||||
def test_max_of_one_only_accepts_one(self):
|
||||
self.assertEqual(validate_worker_threads(1, 1), 1)
|
||||
with self.assertRaises(ValueError):
|
||||
validate_worker_threads(2, 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
"""Bounded thread-pool map for the long jobs (output scan, thumbnails).
|
||||
|
||||
Threads, not processes: this code runs inside ComfyUI's server, where forking
|
||||
a CUDA-initialised, multi-threaded process is unsafe and the spawn start
|
||||
method would re-import ComfyUI's main module in every child. Pillow releases
|
||||
the GIL while decoding, resizing and encoding, so threads still scale close
|
||||
to linearly for both jobs.
|
||||
"""
|
||||
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
|
||||
|
||||
def map_parallel(func, items, workers):
|
||||
"""Apply ``func`` to every item on up to ``workers`` threads, keeping order.
|
||||
|
||||
``workers`` of 1 (or fewer) runs on the calling thread with no pool. An
|
||||
exception in ``func`` propagates to the caller, so callers that must
|
||||
survive bad input wrap ``func`` themselves.
|
||||
"""
|
||||
items = list(items)
|
||||
workers = max(1, int(workers or 1))
|
||||
if workers == 1 or len(items) <= 1:
|
||||
return [func(item) for item in items]
|
||||
max_workers = min(workers, len(items))
|
||||
with ThreadPoolExecutor(max_workers, thread_name_prefix="pm-worker") as pool:
|
||||
return list(pool.map(func, items))
|
||||
@@ -130,6 +130,35 @@ def validate_result_timeout(timeout) -> int:
|
||||
return timeout
|
||||
|
||||
|
||||
def validate_worker_threads(value, max_workers) -> int:
|
||||
"""
|
||||
Validate the worker-thread count for the output scan and thumbnail jobs.
|
||||
|
||||
Args:
|
||||
value: Raw value from the settings request body
|
||||
max_workers: Detected core count; the largest accepted value
|
||||
|
||||
Returns:
|
||||
The thread count as an int in ``[1, max_workers]``
|
||||
|
||||
Raises:
|
||||
ValueError: If the value is not a whole number in that range
|
||||
"""
|
||||
if isinstance(value, bool):
|
||||
raise ValueError("Worker threads must be a whole number")
|
||||
|
||||
if isinstance(value, str) and value.strip().isdigit():
|
||||
value = int(value.strip())
|
||||
|
||||
if not isinstance(value, int):
|
||||
raise ValueError("Worker threads must be a whole number")
|
||||
|
||||
if value < 1 or value > max_workers:
|
||||
raise ValueError(f"Worker threads must be between 1 and {max_workers}")
|
||||
|
||||
return value
|
||||
|
||||
|
||||
def validate_tags(tags: Union[str, List[str], None]) -> bool:
|
||||
"""
|
||||
Validate tags input.
|
||||
|
||||
@@ -411,6 +411,21 @@
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Performance Section -->
|
||||
<div class="border-b border-pm pb-3">
|
||||
<h4 class="text-xs font-semibold text-pm-secondary mb-3 uppercase tracking-wide">Performance</h4>
|
||||
|
||||
<div>
|
||||
<div class="flex items-center justify-between mb-1.5">
|
||||
<label for="workerThreads" class="text-xs font-medium text-pm-secondary">Worker threads</label>
|
||||
<span id="workerThreadsValue" class="text-pm-accent text-xs font-mono">1 / 1 cores</span>
|
||||
</div>
|
||||
<input type="range" id="workerThreads" min="1" max="1" step="1" value="1"
|
||||
class="w-full h-1.5 bg-pm-primary rounded-lg appearance-none cursor-pointer accent-[var(--pm-accent)]">
|
||||
<p class="text-xs text-pm-muted mt-1">Threads used by the image rescan and thumbnail generation. More threads finish sooner; fewer leave CPU for ComfyUI while it generates.</p>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Integrations Section -->
|
||||
<div class="border-b border-pm pb-3">
|
||||
<h4 class="text-xs font-semibold text-pm-secondary mb-3 uppercase tracking-wide">Integrations</h4>
|
||||
|
||||
@@ -6,6 +6,8 @@
|
||||
this.settings = {
|
||||
resultTimeout: 5,
|
||||
webuiDisplayMode: 'popup',
|
||||
workerThreads: 1,
|
||||
cpuCount: 1,
|
||||
};
|
||||
this.categories = [];
|
||||
this.tags = [];
|
||||
@@ -230,6 +232,7 @@
|
||||
document.getElementById('wd14GeneralThreshold').addEventListener('input', (e) => {
|
||||
document.getElementById('wd14GeneralThresholdValue').textContent = parseFloat(e.target.value).toFixed(2);
|
||||
});
|
||||
document.getElementById("workerThreads").addEventListener("input", () => this.renderWorkerThreads());
|
||||
document.getElementById('wd14CharacterThreshold').addEventListener('input', (e) => {
|
||||
document.getElementById('wd14CharacterThresholdValue').textContent = parseFloat(e.target.value).toFixed(2);
|
||||
});
|
||||
@@ -280,6 +283,10 @@
|
||||
if (data.success && data.settings) {
|
||||
this.settings.resultTimeout = data.settings.result_timeout ?? 5;
|
||||
this.settings.webuiDisplayMode = data.settings.webui_display_mode || 'popup';
|
||||
this.settings.cpuCount = Math.max(1, parseInt(data.settings.cpu_count, 10) || 1);
|
||||
this.settings.workerThreads = Math.min(
|
||||
this.settings.cpuCount, Math.max(1, parseInt(data.settings.worker_threads, 10) || 1)
|
||||
);
|
||||
this.settings.monitoredDirectories = data.settings.monitored_directories || [];
|
||||
this.settings.galleryRootPaths = data.settings.gallery_root_paths || [];
|
||||
// Backward compat: if server only returned old field
|
||||
@@ -802,9 +809,19 @@
|
||||
});
|
||||
}
|
||||
|
||||
renderWorkerThreads() {
|
||||
const threads = document.getElementById("workerThreads");
|
||||
document.getElementById("workerThreadsValue").textContent =
|
||||
`${threads.value} / ${this.settings.cpuCount} cores`;
|
||||
}
|
||||
|
||||
showSettingsModal() {
|
||||
document.getElementById("resultTimeout").value = this.settings.resultTimeout;
|
||||
document.getElementById("webuiDisplayMode").value = this.settings.webuiDisplayMode;
|
||||
const threads = document.getElementById("workerThreads");
|
||||
threads.max = this.settings.cpuCount;
|
||||
threads.value = this.settings.workerThreads;
|
||||
this.renderWorkerThreads();
|
||||
this.renderScanPaths();
|
||||
this.updateMonitoringStatus();
|
||||
this.detectLoraManager();
|
||||
@@ -999,10 +1016,12 @@
|
||||
async saveSettings() {
|
||||
const timeout = parseInt(document.getElementById("resultTimeout").value);
|
||||
const displayMode = document.getElementById("webuiDisplayMode").value;
|
||||
const workerThreads = parseInt(document.getElementById("workerThreads").value, 10) || 1;
|
||||
const galleryPaths = this._collectScanPaths().filter(p => p !== '');
|
||||
|
||||
this.settings.resultTimeout = timeout;
|
||||
this.settings.webuiDisplayMode = displayMode;
|
||||
this.settings.workerThreads = workerThreads;
|
||||
this.settings.galleryRootPaths = galleryPaths;
|
||||
|
||||
try {
|
||||
@@ -1012,6 +1031,7 @@
|
||||
body: JSON.stringify({
|
||||
result_timeout: timeout,
|
||||
webui_display_mode: displayMode,
|
||||
worker_threads: workerThreads,
|
||||
gallery_root_paths: galleryPaths
|
||||
}),
|
||||
});
|
||||
|
||||
Reference in New Issue
Block a user