diff --git a/TESTING.md b/TESTING.md index 6ac18b58..00393f34 100644 --- a/TESTING.md +++ b/TESTING.md @@ -349,9 +349,12 @@ and the time left. The steps, in order: observer join get what is left, which can be 0 s. `stop()` joins watchdog's emitters without a limit, so it runs on a daemon thread and is waited for only until the deadline; a later stop waits for that same call (#561). -3. `close_auto_thumbnails(final=True)` stops the preview worker (30 s) and - kills ffmpeg. -4. `PromptTracker.shutdown` stops the prompt cleanup thread (2 s). +3. `close_auto_thumbnails(final=True)` stops the preview worker (30 s), + kills ffmpeg and shuts the interactive image pool down. +4. `shutdown_job_pool` (`utils/parallel.py`) cancels the long-lived job pool's + queued items and joins its threads (5 s; part `workers`, only when a job + used the pool, #523). +5. `PromptTracker.shutdown` stops the prompt cleanup thread (2 s). A failing step, or a second Ctrl+C, is logged with its traceback and does not skip the later steps. A clean run logs one INFO line to PromptManager's log: @@ -469,6 +472,25 @@ differ are files whose pixel data is undecodable or cut short: legacy found nothing, the walk keeps the text before the pixels. Animated PNGs still use `.text`. +Image work runs on PromptManager's own threads (#523). glibc keeps each thread's +malloc arena at its high-water mark, so image buffers spread over the default +executor's threads stay resident. Interactive image work (view previews behind their +gate, workspace card facts, both inspector routes) runs on the API's fixed +`pm-image` pool (`IMAGE_POOL_SIZE` 3, `_run_image`, `_map_image`; card facts already +cached answer without a trip; while the gate's 3 preview writes hold every thread, a +card or inspector read waits for one of them, by design of the merged pool; after the +package teardown the pool is not recreated and late calls use the default executor). Jobs (scan metadata and file facts, thumbnail and +convert jobs, the legacy gallery generation, capture previews) run through +`map_parallel` on one long-lived `pm-worker` pool of Worker threads, reused between +jobs; a new Worker threads value replaces it (the teardown also joins a replaced pool +still finishing), and a call made from a job thread runs inline. After a scan, thumbnail or convert job, or the legacy gallery generation, +`release_free_memory` (`utils/native_memory.py`) calls glibc's `malloc_trim(0)` once; +it does nothing on Windows, macOS and musl and never raises. No allocator setting +changes. `tests/test_image_pools.py` asserts the thread of each read and write and the +trim; `tests/test_utils_parallel.py` the pool's reuse, resize, nesting and shutdown. +`tools/perf/memory_drive.py` (manual, Linux, never in CI) measures RSS, heap, threads +and arenas under the #523 load against the disposable browser test server. + The `perf budget` job in `.github/workflows/test.yml` is the reference measurement: `ubuntu-latest`, one Python, no coverage, only the wall-clock tests. It runs on every push to `main` and on `workflow_dispatch`, never on pull @@ -3871,8 +3893,8 @@ workspace never asks; previews are a cache the server keeps: - `GET /images/serve/thumbnails/_thumb` for a missing preview is refused as text first (traversal, a backslash, a drive or stream colon, empty or dot segments, so `D:/x`, `\\server\share` and a leading `/` never reach a joined path), then admitted - on the event loop (`_PreviewGate`): 3 previews are written at once on the API's own - `pm-preview` pool (never the default executor), one flight runs per file (the history + on the event loop (`_PreviewGate`): 3 previews are written at once on the API's + interactive `pm-image` pool (never the default executor; #523), one flight runs per file (the history route shares the flight key and the admission), and at most 64 requests wait; one more gets its fallback at once without listing a folder. Inside the admission the preview maps back to its media file under the same root (exact names, then one folder @@ -3905,7 +3927,7 @@ workspace never asks; previews are a cache the server keeps: `exclusive` job runs (`_auto_thumbnail_blocker`), checked when a write starts (after its slot) and again just before its rename. - Shutdown is bounded and logged: the worker stops at a file boundary, ffmpeg is - killed, an idle worker leaves (the next output starts a fresh one), the view pool + killed, an idle worker leaves (the next output starts a fresh one), the image pool shuts down, and after the app's cleanup a late monitor event starts nothing. A package reload stops the workers the previous import left. - Cards and strips load previews only: Images in every layout and size, Library cards, diff --git a/__init__.py b/__init__.py index d2e80c1a..b269c00f 100644 --- a/__init__.py +++ b/__init__.py @@ -280,11 +280,11 @@ def _teardown_logger(): def _stop_background_work(): - """Stop the image monitor, the preview worker and the prompt cleanup - thread against one TEARDOWN_SECONDS deadline, once. + """Stop the image monitor, the preview worker, the job worker threads + and the prompt cleanup thread against one TEARDOWN_SECONDS deadline, once. - Returns ``{"monitor", "thumbnails", "tracker"}`` booleans (True: it - stopped), leaving out parts that never started. Later calls return a + Returns ``{"monitor", "thumbnails", "workers", "tracker"}`` booleans + (True: it stopped), leaving out parts that never started. Later calls return a copy of the first result. A Ctrl+C inside a step is logged and the step is False; one between steps is re-raised after the result (parts not reached are False) is stored and the ``atexit`` entry removed. @@ -298,11 +298,14 @@ def _stop_background_work(): api, _global_prompt_tracker, ) + from .utils import parallel + # A part whose global is set has started; False until its step says # otherwise, so an interrupt between steps leaves an honest result. started = ( ("monitor", monitor), ("thumbnails", thumbnails), + ("workers", parallel.job_pool_started() or None), ("tracker", tracker), ) result = {name: False for name, part in started if part is not None} @@ -345,6 +348,16 @@ def _stop_background_work(): result["thumbnails"] = _teardown_step( "the preview worker", close_previews, logger ) + if "workers" in result: + + def stop_workers(): + from .utils.parallel import JOB_POOL_JOIN_SECONDS + + return parallel.shutdown_job_pool(left(JOB_POOL_JOIN_SECONDS)) + + result["workers"] = _teardown_step( + "the job worker threads", stop_workers, logger + ) if tracker is not None: def stop_tracker(): diff --git a/py/api/__init__.py b/py/api/__init__.py index 9c71f7a7..b312328e 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -85,6 +85,7 @@ try: from ...database.operations import PromptDatabase from ...utils.file_prompt import parse_comfyui_prompt from ...utils.logging_config import get_logger + from ...utils.native_memory import release_free_memory from ...utils.png_text import png_text from ...utils.video_metadata import is_video_path, read_video_metadata except ImportError: @@ -96,6 +97,7 @@ except ImportError: from database.operations import PromptDatabase from utils.file_prompt import parse_comfyui_prompt from utils.logging_config import get_logger + from utils.native_memory import release_free_memory from utils.png_text import png_text from utils.video_metadata import is_video_path, read_video_metadata @@ -528,6 +530,14 @@ class PromptManagerAPI( return await loop.run_in_executor(None, call) return await loop.run_in_executor(None, func, *args) + async def _release_job_memory(self): + """After a scan or thumbnail job: glibc's free heap pages go back to + the OS once (``malloc_trim(0)``; nothing elsewhere, #523).""" + try: + await self._run_in_executor(release_free_memory) + except Exception: # best effort: the job's result already stands + self.logger.debug("Releasing free memory failed", exc_info=True) + def _root_listing(self, root) -> RootListing: """The shared listing record of one output root (created empty).""" return self._root_listings.setdefault(str(root), RootListing()) diff --git a/py/api/admin.py b/py/api/admin.py index dc70b9f0..f0ea05f6 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -1454,6 +1454,7 @@ class AdminRoutesMixin: verb = "cancelled" if state == "cancelled" else "finished" job.log_line(f"Scan {verb} · {_plural(job.processed, 'file')} processed") job.finish(state) + await self._release_job_memory() if state == "finished": # After the scan, so the guarded thumbnail job is admitted (#249). await self._hand_off_scan_thumbnails(job, output_dirs) @@ -1668,7 +1669,8 @@ class AdminRoutesMixin: def _extract_batch_metadata_sync(self, file_batch): """Extract metadata from a batch of files in one executor call. - Files are read on ``PromptManagerConfig.WORKER_THREADS`` threads. + Files are read on the long-lived job pool of + ``PromptManagerConfig.WORKER_THREADS`` threads (#523). PNG text is read chunk by chunk without decoding pixels (#523); file reads and ffprobe release the GIL. """ diff --git a/py/api/history.py b/py/api/history.py index a62c9821..131fca8d 100644 --- a/py/api/history.py +++ b/py/api/history.py @@ -843,7 +843,10 @@ class HistoryRoutesMixin: raise async def execute(): - pending = asyncio.create_task(self._run_in_executor(project)) + # An inspection reads the file's header and text: the + # interactive image pool (#523); media stays on the default. + run = self._run_image if kind == "inspection" else self._run_in_executor + pending = asyncio.create_task(run(project)) try: return await asyncio.shield(pending) except asyncio.CancelledError: diff --git a/py/api/images.py b/py/api/images.py index afa4f319..7d530555 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -2975,6 +2975,7 @@ class ImageRoutesMixin: result = await self._run_in_executor( self._generate_gallery_thumbnails_sync, output_dirs ) + await self._release_job_memory() return web.json_response(result) except Exception as e: @@ -3997,7 +3998,8 @@ class ImageRoutesMixin: **observed, ) - body = await self._run_in_executor(inspect) + # Header and text reads: the interactive image pool (#523). + body = await self._run_image(inspect) if body is None: return web.json_response( { @@ -4210,6 +4212,22 @@ class ImageRoutesMixin: # (which Windows refuses on open files) is not held up by it. return _CardFacts(width, height, *self._card_parse(path, records)) + def _cached_card_facts(self, entry: tuple) -> _CardFacts | None: + path, mtime, root, index = entry[:4] + return self._workspace_card_cache.get((index, str(root)), {}).get((path, mtime)) + + def _card_facts_of(self, entries: list, *, parse: bool) -> list[_CardFacts]: + """Card facts of *entries*, in order; files not cached yet are read on + the interactive image pool (#523), cached ones answer here.""" + found = [self._cached_card_facts(entry) for entry in entries] + missing = [i for i, facts in enumerate(found) if facts is None] + read = self._map_image( + lambda i: self._card_facts(entries[i], parse=parse), missing + ) + for i, facts in zip(missing, read): + found[i] = facts + return found + def _card_facts(self, entry: tuple, *, parse: bool) -> _CardFacts: """One listed image's card facts, read once per ``(path, mtime)``. @@ -4218,11 +4236,11 @@ class ImageRoutesMixin: and drops the rest, and a roots change drops them all. Only files of the current scan are stored, so the cache is bounded by the listing. """ - path, mtime, root, index = entry[:4] - root_key, key = (index, str(root)), (path, mtime) - cached = self._workspace_card_cache.get(root_key, {}).get(key) + cached = self._cached_card_facts(entry) if cached is not None: return cached + path, mtime, root, index = entry[:4] + root_key, key = (index, str(root)), (path, mtime) facts = self._read_card_facts(path, parse=parse) if facts is None: # Not stored: a locked or half-written file may read next time. @@ -4248,15 +4266,17 @@ class ImageRoutesMixin: if linked else {} ) + # Videos: the card reads no video metadata. Library previews show + # no card facts, so they skip the parse (``parse=linked``). + stills = [image for image in images if not image["is_video"]] + read = self._card_facts_of( + [entries[(image["root_index"], image["path"])] for image in stills], + parse=linked, + ) + facts_of = {id(image): facts for image, facts in zip(stills, read)} for image in images: entry = entries[(image["root_index"], image["path"])] - # Videos: the card reads no video metadata. Library previews show - # no card facts, so they skip the parse (``parse=linked``). - facts = ( - _NO_CARD_FACTS - if image["is_video"] - else self._card_facts(entry, parse=linked) - ) + facts = facts_of.get(id(image), _NO_CARD_FACTS) image["width"], image["height"] = facts.width, facts.height if not linked: continue diff --git a/py/api/thumbnail_auto.py b/py/api/thumbnail_auto.py index 5cd662b7..4dbb5954 100644 --- a/py/api/thumbnail_auto.py +++ b/py/api/thumbnail_auto.py @@ -8,7 +8,8 @@ previews behind those URLs are made here, in three ways: - **On first view.** ``GET /prompt_manager/images/serve/thumbnails/_thumb `` for a preview that does not exist yet is admitted on the event loop (``_PreviewGate``): at most ``ON_DEMAND_LIMIT`` previews are written at once, - on this API's own small pool (never the default executor), one flight runs + on this API's interactive image pool (``IMAGE_POOL_SIZE`` threads, shared + with card facts and inspector reads; never the default executor), one flight runs per file, and the requests waiting are capped (``MAX_WAITING``); past the cap a request gets its fallback at once. Inside the admission the preview is mapped back to its media file under the same root (the request text is @@ -24,7 +25,8 @@ previews behind those URLs are made here, in three ways: - **On capture.** The image monitor hands every new output it settles to ``queue_capture_thumbnail``, which queues its preview for one background worker thread (``AutoThumbnailer``). Each wake takes at most Worker threads - (``WORKER_THREADS``) files and writes them on that many threads; enqueuing + (``WORKER_THREADS``) files and writes them on the long-lived job pool + (``utils.parallel``, shared with scans and thumbnail jobs); enqueuing is O(1), so the monitor and ComfyUI's queue never wait. A file already queued is not queued again, the queue is bounded (``MAX_QUEUED``) and a preview left out is made on its first view instead. @@ -49,7 +51,7 @@ write starts (after its slot) and again just before its rename. Shutdown: ``close_auto_thumbnails`` stops the worker at a file boundary, kills a running ffmpeg, waits a bounded time (and logs when the wait ran -out) and shuts the view pool down; after the app's cleanup +out) and shuts the image pool down; after the app's cleanup (``final=True``) nothing starts a worker again. Importing this module again (a package reload) stops the workers the previous import left. """ @@ -60,7 +62,7 @@ import functools import io import os import threading -from concurrent.futures import ThreadPoolExecutor +from concurrent.futures import CancelledError, ThreadPoolExecutor from pathlib import Path from aiohttp import web @@ -104,8 +106,11 @@ PLACEHOLDER_SIZE = (thumbnail_settings.DEFAULT_EDGE, thumbnail_settings.DEFAULT_ PLACEHOLDER_QUALITY = 85 # Source extensions in the order a shared preview name looks for them (R2). _SOURCE_ORDER = thumbnail_settings.IMAGE_SOURCE_ORDER + tuple(sorted(VIDEO_EXTENSIONS)) -# Previews views may write at once (the size of the view pool). -ON_DEMAND_LIMIT = 3 +# Threads of the interactive image pool (#523): view previews, card facts and +# inspector reads. Fixed, so image buffers stay in a few threads' arenas. +IMAGE_POOL_SIZE = 3 +# Previews views may write at once (the whole image pool). +ON_DEMAND_LIMIT = IMAGE_POOL_SIZE # View requests that may wait on the gate; one more gets its fallback at once. MAX_WAITING = 64 # How long a pool thread waits for the capture worker's flight of the same file. @@ -124,7 +129,7 @@ AUTO_REFUSED_KINDS = frozenset( {RESTORE_JOB_KIND, DUPLICATE_CLEANUP_JOB_KIND, SOURCE_CHANGE_JOB_KIND} ) WORKER_NAME = "PromptManagerThumbnailWorker" -POOL_PREFIX = "pm-preview" +POOL_PREFIX = "pm-image" # Identifies this import's workers: a reload stops the ones it did not start. _MODULE_TOKEN = object() @@ -294,7 +299,7 @@ class AutoThumbnailer: ``generate(src, dst, is_video, *, stop, guard, track)`` writes one preview and returns ``_generate_one``'s result dict; ``may_write()`` says whether automatic previews may be written now. Views call - ``ensure`` on the view pool, the monitor calls ``enqueue`` from its + ``ensure`` on the image pool, the monitor calls ``enqueue`` from its thread. """ @@ -491,12 +496,15 @@ class AutoThumbnailer: self._worker_running = False self._changed.notify_all() return - take = min(len(self._queue), max(1, int(self._workers() or 1))) + workers = max(1, int(self._workers() or 1)) + take = min(len(self._queue), workers) batch = [self._queue.popleft() for _ in range(take)] self._batch = len(batch) self._changed.notify_all() try: - map_parallel(self._write_queued, batch, len(batch)) + # The job pool's size is the setting; a smaller batch uses + # fewer of its threads (#523: one long-lived pool). + map_parallel(self._write_queued, batch, workers) finally: with self._changed: self._batch = 0 @@ -567,7 +575,7 @@ class _PreviewGate: """Admission of view requests on the event loop (#249 I2). Waiters are coroutines: at most ``limit`` previews are written at once - on the view pool, one flight runs per key, and past ``MAX_WAITING`` + on the image pool, one flight runs per key, and past ``MAX_WAITING`` waiting requests ``run`` answers ``OVER`` at once. """ @@ -622,13 +630,17 @@ class _PreviewGate: flight.set_result(result) +def _on_image_thread(): + return threading.current_thread().name.startswith(POOL_PREFIX + "_") + + class AutoThumbnailMixin: """Automatic previews on the API: views, the monitor's queue, scans.""" _auto_thumbnailer = None _auto_thumbnails_shut = False _preview_gate_obj = None - _preview_pool_obj = None + _image_pool_obj = None _auto_thumbnailer_lock = threading.Lock() def _auto_thumbnails(self): @@ -651,7 +663,7 @@ class AutoThumbnailMixin: self._auto_thumbnails_shut = True def close_auto_thumbnails(self, final=False, timeout=None): - """Stop the capture worker and the view pool, and forget them. + """Stop the capture worker and the image pool, and forget them. Bounded by ``timeout`` (see AutoThumbnailer.close). ``final`` (the package teardown) also keeps a late monitor event from starting a @@ -662,7 +674,7 @@ class AutoThumbnailMixin: if final: self._auto_thumbnails_shut = True thumbnailer, self._auto_thumbnailer = self._auto_thumbnailer, None - pool, self._preview_pool_obj = self._preview_pool_obj, None + pool, self._image_pool_obj = self._image_pool_obj, None self._preview_gate_obj = None closed = True if thumbnailer is not None: @@ -706,16 +718,67 @@ class AutoThumbnailMixin: gate = self._preview_gate_obj = _PreviewGate(ON_DEMAND_LIMIT, self.logger) return gate - def _preview_pool(self): - with AutoThumbnailMixin._auto_thumbnailer_lock: - if self._preview_pool_obj is None: - self._preview_pool_obj = ThreadPoolExecutor( - ON_DEMAND_LIMIT, thread_name_prefix=POOL_PREFIX - ) - return self._preview_pool_obj + def _image_pool(self): + """The interactive image pool (#523): IMAGE_POOL_SIZE threads for view + previews (behind their gate), card facts and inspector reads, so + image buffers never spread over the default executor's threads. - def _preview_pool_threads(self): - pool = self._preview_pool_obj + None after the package teardown (``close_auto_thumbnails(final=True)``): + a late request then runs on the default executor (``run_in_executor`` + with None) and starts no thread the teardown no longer stops. + """ + with AutoThumbnailMixin._auto_thumbnailer_lock: + if self._image_pool_obj is None and not self._auto_thumbnails_shut: + self._image_pool_obj = ThreadPoolExecutor( + IMAGE_POOL_SIZE, thread_name_prefix=POOL_PREFIX + ) + return self._image_pool_obj + + async def _run_image(self, func, *args, **kwargs): + """``_run_in_executor`` for interactive image work: the image pool. + + After the teardown, or with the pool shut down in between (a test + reset), the call runs on the default executor as before. + """ + call = functools.partial(func, *args, **kwargs) + loop = asyncio.get_running_loop() + try: + future = loop.run_in_executor(self._image_pool(), call) + except RuntimeError: # submit refused: the pool was shut down + future = loop.run_in_executor(None, call) + return await future + + def _map_image(self, func, items): + """``func`` over *items* on the image pool, in order (blocking). + + Called from an image pool thread, after the teardown, or with the + pool shut down in between (its queued items are cancelled), the + items it did not run run on the calling thread. + """ + items = list(items) + pool = None if not items or _on_image_thread() else self._image_pool() + if pool is None: + return [func(item) for item in items] + futures = [] + try: + for item in items: + futures.append(pool.submit(func, item)) + except RuntimeError: + pass # shut down in between: the rest run here + results = [] + for index, item in enumerate(items): + future = futures[index] if index < len(futures) else None + if future is None or future.cancelled(): + results.append(func(item)) + continue + try: + results.append(future.result()) + except CancelledError: # cancelled while this caller waited + results.append(func(item)) + return results + + def _image_pool_threads(self): + pool = self._image_pool_obj return [] if pool is None else list(getattr(pool, "_threads", ())) def _video_placeholder_bytes(self): @@ -730,7 +793,7 @@ class AutoThumbnailMixin: return await self._preview_gate().run( flight_key(Path(dst).parent, (Path(dst).name,)), functools.partial(self._ensure_view, src, dst, is_video), - self._preview_pool(), + self._image_pool(), ) def _ensure_view(self, src, dst, is_video): @@ -774,13 +837,13 @@ class AutoThumbnailMixin: The flight answers only the outcome (a history flight may be the one joined); the file it previews is derived here: the preview path is the request's own, and a fallback looks its source up again on the - view pool (bounded by it, never the default executor). + image pool (bounded by it, never the default executor). """ if _has_traversal(filepath) or not _is_preview_request(filepath): return None parts = tuple(filepath.split("/")) gate = self._preview_gate() - pool = self._preview_pool() + pool = self._image_pool() for root in output_dirs: root = Path(root) outcome = await gate.run( diff --git a/py/api/thumbnail_convert.py b/py/api/thumbnail_convert.py index 3a9db4e3..a12d8be1 100644 --- a/py/api/thumbnail_convert.py +++ b/py/api/thumbnail_convert.py @@ -153,6 +153,7 @@ class ThumbnailConvertMixin: job.log_line("Conversion failed · internal error") state = "failed" job.finish(state) + await self._release_job_memory() async def _convert_job_body(self, job, options): output_dirs = await self._run_in_executor(self._get_all_output_dirs) diff --git a/py/api/thumbnail_jobs.py b/py/api/thumbnail_jobs.py index 6477f7d7..463eb891 100644 --- a/py/api/thumbnail_jobs.py +++ b/py/api/thumbnail_jobs.py @@ -226,6 +226,7 @@ class ThumbnailJobsMixin: ) state = "failed" job.finish(state) + await self._release_job_memory() async def _thumbnail_job_body(self, job, options): """Plan, then generate in worker batches; returns the final state. diff --git a/tests/test_api_workspace_images.py b/tests/test_api_workspace_images.py index 82507e7a..14cf8d11 100644 --- a/tests/test_api_workspace_images.py +++ b/tests/test_api_workspace_images.py @@ -908,20 +908,20 @@ class TestWorkspaceImages(ImageAPITestCase): async def test_one_cards_parse_error_spares_the_page(self): prompt = self._save_prompt("a fox") - for name, mtime in (("bad.png", 200), ("ok.png", 100)): + # Cards are read on the image pool (#523), in no fixed order: the + # parser fails on bad.png's own record (seed 6), not on a call count. + for name, mtime, seed in (("bad.png", 200, 6), ("ok.png", 100, 5)): path = records_png( - self.output_dir / name, seeded_records("a fox", 5), mtime + self.output_dir / name, seeded_records("a fox", seed), mtime ) self.api.db.link_image_to_prompt(prompt, str(path)) real_parse = header_params_of_images() def parsing(info): - if parsing.calls == 0: - parsing.calls += 1 + if '"seed": 6' in info["prompt"]: raise RuntimeError("unexpected parser bug") return real_parse(info) - parsing.calls = 0 with ( patch("py.api.images.header_params", side_effect=parsing), self.assertLogs(self.api.logger, "WARNING") as logged, diff --git a/tests/test_exit_window_threads.py b/tests/test_exit_window_threads.py index 30e39bdd..e5d86117 100644 --- a/tests/test_exit_window_threads.py +++ b/tests/test_exit_window_threads.py @@ -146,6 +146,18 @@ class TestPreviewQueueRefusedAtShutdown(unittest.TestCase): class TestMapParallelRefusedAtShutdown(unittest.TestCase): + def setUp(self): + # The job pool is long-lived (#523): start each test without one, so + # the patched pool class is the one created. + self.assertTrue(parallel.shutdown_job_pool(30)) + self.addCleanup(parallel.shutdown_job_pool, 30) + + def test_the_refused_pool_is_not_reused(self): + with mock.patch.object(parallel, "ThreadPoolExecutor", _ShutPool): + parallel.map_parallel(lambda n: n, [1, 2], 2) + self.assertFalse(parallel.job_pool_started()) + self.assertEqual(parallel.map_parallel(lambda n: n + 1, [1, 2], 2), [2, 3]) + def test_a_shut_pool_runs_the_items_on_the_calling_thread(self): records = _capture(self, "prompt_manager.parallel") with mock.patch.object(parallel, "ThreadPoolExecutor", _ShutPool): diff --git a/tests/test_image_pools.py b/tests/test_image_pools.py new file mode 100644 index 00000000..38384a5e --- /dev/null +++ b/tests/test_image_pools.py @@ -0,0 +1,241 @@ +"""Image work runs on PromptManager's own pools, never the default executor (#523). + +glibc keeps each thread's malloc arena at its high-water mark, so image work +spread over the default executor's threads (min(32, cpus + 4)) kept an +image-sized buffer per thread. Interactive image work (card facts, inspector +reads, view previews) runs on the fixed interactive image pool +(``pm-image``); jobs (scan metadata, thumbnail jobs, capture previews) on the +long-lived Worker threads pool (``pm-worker``). After a scan or thumbnail +job, glibc's free pages are released once (``malloc_trim``). +""" + +import asyncio +import threading +from unittest.mock import patch + +import py.api as api_package +import py.api.images as images_module +import py.api.inspection_service as inspection_service +import py.api.thumbnail_auto as auto +from utils import parallel + +try: + from tests.test_api_images import make_png + from tests.test_thumbnail_auto import WAIT, AutoThumbnailCase, _entry +except ImportError: # discovered with tests/ as the top-level directory + from test_api_images import make_png + from test_thumbnail_auto import WAIT, AutoThumbnailCase, _entry + +IMAGE_THREAD = auto.POOL_PREFIX + "_" +JOB_THREAD = parallel.JOB_THREAD_PREFIX + "_" + + +class ImagePoolCase(AutoThumbnailCase): + async def asyncSetUp(self): + await super().asyncSetUp() + self.assertTrue(parallel.shutdown_job_pool(WAIT)) + self.addCleanup(parallel.shutdown_job_pool, WAIT) + self.threads = [] + self.lock = threading.Lock() + + def spy(self, owner, name): + """Record the thread every call of ``owner.name`` runs on.""" + original = getattr(owner, name) + + def recording(*args, **kwargs): + with self.lock: + self.threads.append(threading.current_thread().name) + return original(*args, **kwargs) + + patcher = patch.object(owner, name, recording) + patcher.start() + self.addCleanup(patcher.stop) + + def assert_threads(self, prefix): + self.assertTrue(self.threads, "the spied call never ran") + self.assertTrue( + all(name.startswith(prefix) for name in self.threads), self.threads + ) + + async def run_job(self, path, body=None): + response = await self.client.post(path, json=body or {}) + self.assertEqual(response.status, 202, await response.text()) + job = self.api._jobs.get((await response.json())["job"]["id"]) + await asyncio.wait_for(asyncio.shield(job.task), WAIT) + return job + + +class TestInteractiveImagePool(ImagePoolCase): + async def test_workspace_card_facts_are_read_on_the_image_pool(self): + for index in range(5): + make_png(self.output_dir / f"c{index}.png", size=(32, 32)) + self.spy(images_module, "_card_header") + response = await self.client.get( + "/prompt_manager/images/output?mode=workspace&limit=100" + ) + self.assertEqual(response.status, 200, await response.text()) + images = (await response.json())["images"] + self.assertEqual([image["width"] for image in images], [32] * 5) + self.assert_threads(IMAGE_THREAD) + self.assertEqual(len(self.threads), 5) + + async def test_cached_card_facts_need_no_image_pool_trip(self): + make_png(self.output_dir / "c.png", size=(32, 32)) + url = "/prompt_manager/images/output?mode=workspace&limit=100" + await (await self.client.get(url)).read() + self.spy(images_module, "_card_header") + response = await self.client.get(url) + self.assertEqual((await response.json())["images"][0]["width"], 32) + self.assertEqual(self.threads, []) + + async def test_the_inspector_reads_on_the_image_pool(self): + make_png(self.output_dir / "i.png", size=(32, 32)) + self.spy(inspection_service, "read_inspection_file") + response = await self.client.get("/prompt_manager/images/inspect/i.png?root=0") + self.assertEqual(response.status, 200, await response.text()) + self.assert_threads(IMAGE_THREAD) + + async def test_view_previews_and_card_reads_share_three_threads(self): + names = [f"v{i}.png" for i in range(auto.ON_DEMAND_LIMIT + 3)] + for name in names: + make_png(self.output_dir / name, size=(64, 64)) + calls = [] + original = self.api._generate_one + + def recording(*args, **kwargs): + calls.append(threading.current_thread().name) + return original(*args, **kwargs) + + self.api._generate_one = recording + for name in names: + url = _entry(self.api, self.output_dir, name)["thumbnail_url"] + await (await self.client.get(url)).read() + self.assertEqual(len(calls), len(names)) + self.assertTrue(all(name.startswith(IMAGE_THREAD) for name in calls), calls) + self.assertLessEqual(len(self.api._image_pool_threads()), auto.IMAGE_POOL_SIZE) + + def test_map_image_keeps_order_and_runs_inline_on_a_pool_thread(self): + self.assertEqual(self.api._map_image(lambda n: n * 2, [3, 1, 2]), [6, 2, 4]) + nested = self.api._image_pool().submit( + self.api._map_image, lambda _: threading.current_thread().name, [1] + ) + (name,) = nested.result(WAIT) + self.assertTrue(name.startswith(IMAGE_THREAD), name) + + def test_map_image_after_a_close_runs_on_the_calling_thread(self): + pool = self.api._image_pool() + pool.shutdown(wait=True) + here = threading.current_thread().name + with patch.object(self.api, "_image_pool", return_value=pool): + names = self.api._map_image( + lambda _: threading.current_thread().name, [1, 2] + ) + self.assertEqual(names, [here, here]) + + async def test_after_the_final_close_no_image_thread_starts(self): + await asyncio.get_running_loop().run_in_executor( + None, lambda: self.api.close_auto_thumbnails(final=True) + ) + self.assertIsNone(self.api._image_pool()) + here = threading.current_thread().name + names = self.api._map_image(lambda _: threading.current_thread().name, [1]) + self.assertEqual(names, [here]) + name = await self.api._run_image(lambda: threading.current_thread().name) + self.assertFalse(name.startswith(IMAGE_THREAD), name) + self.assertEqual(self.api._image_pool_threads(), []) + + def test_map_image_runs_items_a_shutdown_cancelled_here(self): + entered = threading.Barrier(auto.IMAGE_POOL_SIZE + 1) + release = threading.Event() + self.addCleanup(release.set) + pool = self.api._image_pool() + for _ in range(auto.IMAGE_POOL_SIZE): # occupy every pool thread + pool.submit(lambda: (entered.wait(WAIT), release.wait(WAIT))) + entered.wait(WAIT) + here = threading.current_thread().name + results, submitted = {}, threading.Event() + + class Counting: + calls = 0 + + def submit(self, *args): + future = pool.submit(*args) + Counting.calls += 1 + if Counting.calls == 2: + submitted.set() + return future + + def call(): + with patch.object(self.api, "_image_pool", return_value=Counting()): + results["names"] = self.api._map_image( + lambda _: threading.current_thread().name, [1, 2] + ) + + caller = threading.Thread(target=call, name="map-image-caller") + caller.start() + # The caller's two items are queued behind the busy threads; the + # shutdown cancels them and the caller runs them itself. + self.assertTrue(submitted.wait(WAIT)) + pool.shutdown(wait=False, cancel_futures=True) + caller.join(WAIT) + release.set() + self.assertFalse(caller.is_alive()) + self.assertEqual(results["names"], ["map-image-caller"] * 2) + self.assertNotEqual(here, "map-image-caller") + + async def test_run_image_after_a_close_uses_the_default_executor(self): + pool = self.api._image_pool() + pool.shutdown(wait=True) + with patch.object(self.api, "_image_pool", return_value=pool): + name = await self.api._run_image(lambda: threading.current_thread().name) + self.assertFalse(name.startswith(IMAGE_THREAD), name) + + +class TestJobPool(ImagePoolCase): + def test_scan_metadata_is_read_on_the_job_pool(self): + paths = [ + make_png(self.output_dir / f"s{i}.png", size=(16, 16)) for i in range(4) + ] + self.spy(self.api, "_extract_comfyui_metadata") + results = self.api._extract_batch_metadata_sync(paths) + self.assertEqual([path for path, _ in results], paths) + self.assert_threads(JOB_THREAD) + + async def test_thumbnail_job_writes_on_the_job_pool_and_trims_after(self): + for index in range(3): + make_png(self.output_dir / f"t{index}.png", size=(64, 64)) + self.spy(self.api, "_generate_one") + with patch.object(api_package, "release_free_memory") as trim: + job = await self.run_job("/prompt_manager/jobs/thumbnails") + self.assertEqual(job.state, "finished") + self.assert_threads(JOB_THREAD) + trim.assert_called_once_with() + + async def test_scan_job_trims_after_it_finishes(self): + make_png(self.output_dir / "s.png", size=(16, 16)) + states = [] + + def trim(): + states.append(self.api._jobs.active("scan")) + return True + + with patch.object(api_package, "release_free_memory", side_effect=trim): + job = await self.run_job("/prompt_manager/jobs/scan") + self.assertEqual(job.state, "finished") + self.assertGreaterEqual(len(states), 1) + self.assertIsNone(states[0], "trimmed after the scan finished") + + async def test_a_failing_trim_never_fails_the_job(self): + make_png(self.output_dir / "t.png", size=(64, 64)) + with patch.object( + api_package, "release_free_memory", side_effect=RuntimeError("no") + ): + job = await self.run_job("/prompt_manager/jobs/thumbnails") + self.assertEqual(job.state, "finished") + + async def test_capture_previews_write_on_the_job_pool(self): + path = make_png(self.output_dir / "a.png", size=(64, 64)) + self.spy(self.api, "_generate_one") + self.assertEqual(self.api.queue_capture_thumbnail(str(path)), "queued") + await self._idle() + self.assert_threads(JOB_THREAD) diff --git a/tests/test_init_package.py b/tests/test_init_package.py index 6483e809..9742601a 100644 --- a/tests/test_init_package.py +++ b/tests/test_init_package.py @@ -799,6 +799,45 @@ class TestPackageTeardown(PackageTeardownCase): self.assertIsNone(seen["worker"]) self.assertTrue(all(result.values()), result) + def test_teardown_stops_the_job_pool_once_a_job_used_it(self): + """#523: the long-lived job pool is a teardown part of its own.""" + package = self.load() + parallel = sys.modules[package.__name__ + ".utils.parallel"] + self.assertEqual(parallel.map_parallel(lambda n: n * 2, [1, 2], 2), [2, 4]) + # This package's own pool (the suite's top-level one may also run). + threads = list(parallel._pool[1]._threads) + self.assertTrue(threads) + + result = package._stop_background_work() + + self.assertEqual(list(result), ["monitor", "thumbnails", "workers", "tracker"]) + self.assertTrue(all(result.values()), result) + self.assertFalse(parallel.job_pool_started()) + self.assertFalse(any(t.is_alive() for t in threads)) + + def test_teardown_reports_a_job_thread_past_its_bound(self): + package = self.load(require_stopped=False) + parallel = sys.modules[package.__name__ + ".utils.parallel"] + entered, release = threading.Event(), threading.Event() + + def hold(_): + entered.set() + release.wait(30) + + caller = threading.Thread(target=parallel.map_parallel, args=(hold, [1], 1)) + caller.start() + self.addCleanup(caller.join, 30) + self.addCleanup(release.set) + self.assertTrue(entered.wait(30)) + with ( + mock.patch.object(parallel, "JOB_POOL_JOIN_SECONDS", 0.05), + self.assertLogs("prompt_manager.init", "WARNING") as logs, + ): + result = package._stop_background_work() + self.assertFalse(result["workers"]) + self.assertTrue(result["tracker"]) + self.assertRegex("\n".join(logs.output), r"not stopped in time .*workers") + def test_teardown_kills_ffmpeg_and_stops_the_preview_worker(self): package = self.load() held = self.hold_a_preview(package, advance=4.0, kill_releases=True) diff --git a/tests/test_native_memory.py b/tests/test_native_memory.py new file mode 100644 index 00000000..c1892a33 --- /dev/null +++ b/tests/test_native_memory.py @@ -0,0 +1,82 @@ +"""``release_free_memory``: glibc's malloc_trim(0), a no-op elsewhere (#523).""" + +import sys +import unittest +from unittest import mock + +from utils import native_memory + + +class ReleaseFreeMemoryTests(unittest.TestCase): + def setUp(self): + # Each test looks the symbol up again. + patcher = mock.patch.multiple(native_memory, _trim=None, _looked_up=False) + patcher.start() + self.addCleanup(patcher.stop) + + @unittest.skipUnless(native_memory._glibc(), "glibc only") + def test_glibc_trims_once_per_call(self): + self.assertTrue(native_memory.release_free_memory()) + self.assertTrue(native_memory.release_free_memory()) + + def test_other_platforms_do_nothing(self): + for platform in ("win32", "darwin", "cygwin"): + with self.subTest(platform), mock.patch.object(sys, "platform", platform): + native_memory._looked_up = False + with mock.patch.object(native_memory.ctypes, "CDLL") as cdll: + self.assertFalse(native_memory.release_free_memory()) + cdll.assert_not_called() + + def test_musl_does_nothing(self): + with ( + mock.patch.object(sys, "platform", "linux"), + mock.patch.object( + native_memory.os, "confstr", side_effect=ValueError("unknown name") + ), + mock.patch.object(native_memory.ctypes, "CDLL") as cdll, + ): + self.assertFalse(native_memory.release_free_memory()) + cdll.assert_not_called() + + def test_a_missing_symbol_is_false(self): + library = mock.Mock(spec=[]) # no malloc_trim attribute + with ( + mock.patch.object(native_memory, "_glibc", return_value=True), + mock.patch.object(native_memory.ctypes, "CDLL", return_value=library), + ): + self.assertFalse(native_memory.release_free_memory()) + + def test_an_unloadable_libc_is_false(self): + with ( + mock.patch.object(native_memory, "_glibc", return_value=True), + mock.patch.object(native_memory.ctypes, "CDLL", side_effect=OSError), + ): + self.assertFalse(native_memory.release_free_memory()) + + def test_a_failing_trim_never_raises(self): + trim = mock.Mock(side_effect=OSError("boom")) + library = mock.Mock(malloc_trim=trim) + with ( + mock.patch.object(native_memory, "_glibc", return_value=True), + mock.patch.object(native_memory.ctypes, "CDLL", return_value=library), + ): + self.assertFalse(native_memory.release_free_memory()) + trim.assert_called_once_with(0) + + def test_the_symbol_is_looked_up_once(self): + trim = mock.Mock(return_value=1) + library = mock.Mock(malloc_trim=trim) + with ( + mock.patch.object(native_memory, "_glibc", return_value=True), + mock.patch.object( + native_memory.ctypes, "CDLL", return_value=library + ) as cdll, + ): + self.assertTrue(native_memory.release_free_memory()) + self.assertTrue(native_memory.release_free_memory()) + cdll.assert_called_once_with(None) + self.assertEqual(trim.call_count, 2) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_thumbnail_auto.py b/tests/test_thumbnail_auto.py index 18c92db1..4ccf7e52 100644 --- a/tests/test_thumbnail_auto.py +++ b/tests/test_thumbnail_auto.py @@ -291,7 +291,7 @@ class TestThumbnailOnFirstView(AutoThumbnailCase): await self._writers(state, auto.ON_DEMAND_LIMIT) await self._waiting(len(names)) self.assertEqual(default_calls, [], "no view waits on a default thread") - self.assertTrue(all(name.startswith("pm-preview") for name in state["threads"])) + self.assertTrue(all(name.startswith("pm-image") for name in state["threads"])) status = await self.client.get("/prompt_manager/images/thumbnails/status") self.assertEqual(status.status, 200, "another route answers meanwhile") await status.read() @@ -678,7 +678,7 @@ class TestShutdown(AutoThumbnailCase): _entry(self.api, self.output_dir, "viewed.png")["thumbnail_url"] ) await resp.read() - pool_threads = list(self.api._preview_pool_threads()) + pool_threads = list(self.api._image_pool_threads()) self.assertTrue(pool_threads) loop = asyncio.get_running_loop() await loop.run_in_executor( @@ -904,7 +904,8 @@ class TestThumbnailOnCapture(AutoThumbnailCase): self.api.queue_capture_thumbnail(str(path)) release.set() await self._idle() - self.assertEqual(batches, [(1, 1), (2, 2), (2, 2)]) + # The job pool's size is the setting; each batch is at most that (#523). + self.assertEqual(batches, [(1, 2), (2, 2), (2, 2)]) self.assertLessEqual(state["peak"], 2) for i in range(5): self.assertTrue( diff --git a/tests/test_utils_parallel.py b/tests/test_utils_parallel.py index 46bd7e41..21fc717c 100644 --- a/tests/test_utils_parallel.py +++ b/tests/test_utils_parallel.py @@ -1,4 +1,8 @@ -"""Tests for utils/parallel.py: bounded thread-pool map used by long jobs.""" +"""Tests for utils/parallel.py: bounded thread-pool map used by long jobs. + +The jobs share one long-lived pool of Worker threads (#523): its threads (and +their malloc arenas) are reused from job to job instead of a new set per call. +""" import os import sys @@ -8,10 +12,22 @@ import unittest sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) -from utils.parallel import map_parallel # noqa: E402 +from utils import parallel # noqa: E402 +from utils.parallel import JOB_THREAD_PREFIX, map_parallel # noqa: E402 + +WAIT = 30 + + +def _job_threads(): + return { + t for t in threading.enumerate() if t.name.startswith(JOB_THREAD_PREFIX + "_") + } class TestMapParallel(unittest.TestCase): + def setUp(self): + self.assertTrue(parallel.shutdown_job_pool(WAIT)) + self.addCleanup(parallel.shutdown_job_pool, WAIT) def test_results_keep_input_order(self): def slow_square(n): @@ -22,12 +38,102 @@ class TestMapParallel(unittest.TestCase): map_parallel(slow_square, [1, 2, 3, 4], workers=4), [1, 4, 9, 16] ) - def test_one_worker_runs_on_the_calling_thread(self): + def test_one_worker_runs_on_the_one_job_thread(self): names = [] map_parallel( lambda _: names.append(threading.current_thread().name), [1, 2, 3], 1 ) - self.assertEqual(set(names), {threading.current_thread().name}) + self.assertEqual(len(set(names)), 1) + self.assertTrue(names[0].startswith(JOB_THREAD_PREFIX + "_"), names) + + def test_jobs_reuse_the_same_long_lived_threads(self): + """Arenas are reused instead of a new set of threads per job.""" + first, second = set(), set() + map_parallel(lambda _: first.add(threading.current_thread()), range(8), 2) + map_parallel(lambda _: second.add(threading.current_thread()), range(8), 2) + self.assertTrue(second <= first, (first, second)) + self.assertLessEqual(len(_job_threads()), 2) + self.assertTrue(all(t.is_alive() for t in first | second)) + + def test_a_new_worker_count_replaces_the_pool(self): + map_parallel(lambda n: n, range(4), 2) + old = _job_threads() + map_parallel(lambda n: n, range(4), 3) + for thread in old: + thread.join(WAIT) + self.assertFalse(thread.is_alive(), thread.name) + self.assertLessEqual(len(_job_threads()), 3) + + def test_shutdown_also_joins_a_pool_a_resize_retired(self): + entered, release = threading.Event(), threading.Event() + + def hold(_): + entered.set() + release.wait(WAIT) + + caller = threading.Thread(target=map_parallel, args=(hold, [1], 1)) + caller.start() + self.addCleanup(caller.join, WAIT) + self.addCleanup(release.set) + self.assertTrue(entered.wait(WAIT)) + map_parallel(lambda n: n, [1, 2], 2) # a new size retires the busy pool + self.assertFalse(parallel.shutdown_job_pool(0.05)) + release.set() + caller.join(WAIT) + self.assertTrue(parallel.shutdown_job_pool(WAIT)) + self.assertFalse(parallel.job_pool_started()) + + def test_an_item_run_from_a_job_thread_runs_inline(self): + """A nested map never waits for the threads it is running on.""" + + def outer(n): + return map_parallel(lambda m: (m, threading.current_thread()), [n], 1) + + results = map_parallel(outer, [1, 2], 1) + self.assertEqual([r[0][0] for r in results], [1, 2]) + + def test_shutdown_stops_the_threads_and_a_later_job_starts_fresh(self): + map_parallel(lambda n: n, range(4), 2) + threads = _job_threads() + self.assertTrue(threads) + self.assertTrue(parallel.job_pool_started()) + self.assertTrue(parallel.shutdown_job_pool(WAIT)) + self.assertFalse(parallel.job_pool_started()) + self.assertFalse(any(t.is_alive() for t in threads)) + self.assertEqual(map_parallel(lambda n: n + 1, [1], 2), [2]) + + def test_shutdown_reports_a_thread_still_running(self): + entered, release = threading.Event(), threading.Event() + + def hold(_): + entered.set() + release.wait(WAIT) + + caller = threading.Thread(target=map_parallel, args=(hold, [1], 1)) + caller.start() + self.assertTrue(entered.wait(WAIT)) + self.assertFalse(parallel.shutdown_job_pool(0.05)) + release.set() + caller.join(WAIT) + self.assertFalse(caller.is_alive()) + + def test_a_failing_item_waits_for_the_running_ones(self): + """The caller sees the error only after its other items stopped.""" + running = threading.Event() + finished = [] + + def item(n): + if n == 0: + self.assertTrue(running.wait(WAIT)) + raise ValueError("zero") + running.set() + time.sleep(0.05) + finished.append(n) + return n + + with self.assertRaises(ValueError): + map_parallel(item, [0, 1], 2) + self.assertEqual(finished, [1]) def test_several_workers_use_several_threads(self): names = set() diff --git a/tools/perf/memory_drive.py b/tools/perf/memory_drive.py new file mode 100644 index 00000000..70f7d07e --- /dev/null +++ b/tools/perf/memory_drive.py @@ -0,0 +1,395 @@ +"""Measure the server's resident memory under a busy Images/Library session (#523). + +Manual only, Linux only (it reads ``/proc`` and glibc's ``malloc_info``); never +run in CI. It starts the production API inside the disposable browser test +server (``tests/browser/server.py``: its own temporary library, database and +output folder; never a live ComfyUI) and drives it for ``--minutes``: + +- ``--clients`` threads loop over: a workspace grid page of 100 images, 40 of + its thumbnails, 6 inspector reads and 6 full-file reads, the Library list + with 3 entries' history and images, and the recent prompt and image lists; +- every minute, ``--batch`` new PNGs (832 to 1536 px, ComfyUI-shaped + metadata) are copied into the output folder, then a scan job and a + thumbnail job run to completion. + +Every 5 s the server samples its RSS, Python heap (tracemalloc), threads and +glibc arena count. The test server pins glibc's mmap threshold (#516); this +driver turns that off so the production allocator behaviour is measured, +unless ``--mallopt`` asks for the pinned comparison run. + + python tools/perf/memory_drive.py --label after --minutes 12 + python tools/perf/memory_drive.py --repo /path/to/other/checkout --label before + +The checkout under ``--repo`` (default: this one) is the code measured; the +driver itself only needs Pillow. The last line printed is a Markdown table row. +""" + +import argparse +import ctypes +import json +import os +import random +import shutil +import subprocess # nosec B404 - starts this script's own server child +import sys +import tempfile +import threading +import time +import tracemalloc +import urllib.parse +import urllib.request +from concurrent.futures import ThreadPoolExecutor + +HERE = os.path.dirname(os.path.abspath(__file__)) +REPO = os.path.dirname(os.path.dirname(HERE)) +PORT_ANNOUNCEMENT = "PM_TEST_SERVER_PORT=" +OUTPUT_ANNOUNCEMENT = "output=" +SAMPLE_SECONDS = 5 +MEDIA_FILES = 120 +MEDIA_SIDES = (832, 1024, 1216, 1536) +REQUEST_TIMEOUT = 120 +JOB_POLL_LIMIT = 600 +FINAL_STATES = ("finished", "cancelled", "failed") +M_MMAP_THRESHOLD = -3 + + +# ── the server child ──────────────────────────────────────────────────── + + +def _arena_count(libc): + """glibc's arena count (``> 10, + "anon": status["RssAnon"] >> 10, + "heap": tracemalloc.get_traced_memory()[0] >> 20, + "threads": status["Threads"], + "arenas": _arena_count(libc), + } + ) + + "\n" + ) + time.sleep(SAMPLE_SECONDS if announced else 0.5) + + +def serve(repo, samples_path, mallopt): + """Run the browser test server of *repo* in this process, sampled.""" + sys.path.insert(0, repo) + os.chdir(repo) + tracemalloc.start(1) + libc = ctypes.CDLL(None) + libc.open_memstream.restype = ctypes.c_void_p + libc.open_memstream.argtypes = [ + ctypes.POINTER(ctypes.c_char_p), + ctypes.POINTER(ctypes.c_size_t), + ] + libc.fclose.argtypes = [ctypes.c_void_p] + libc.malloc_info.argtypes = [ctypes.c_int, ctypes.c_void_p] + libc.free.argtypes = [ctypes.c_void_p] + if mallopt: + libc.mallopt(M_MMAP_THRESHOLD, 128 * 1024) + threading.Thread( + target=_sample_forever, args=(libc, samples_path), daemon=True + ).start() + from tests.browser import server + + if not mallopt: + server.pin_mmap_threshold = lambda: False # production allocator + server.main(["--port", "0", "--exit-on-stdin-eof"]) + + +# ── the driver ────────────────────────────────────────────────────────── + + +def make_media(folder, count=MEDIA_FILES, seed=523): + """Synthetic ComfyUI-shaped PNGs (smooth noise, prompt and workflow).""" + from PIL import Image, PngImagePlugin + + rng = random.Random(seed) + os.makedirs(folder, exist_ok=True) + for index in range(count): + side = MEDIA_SIDES[index % len(MEDIA_SIDES)] + small = Image.frombytes( + "RGB", (64, 64), bytes(rng.randrange(256) for _ in range(64 * 64 * 3)) + ) + image = small.resize((side, side), Image.BICUBIC) + text = f"synthetic scene {index}, soft light, detailed" + prompt = { + "3": { + "class_type": "KSampler", + "inputs": { + "seed": index, + "steps": 20, + "cfg": 7.0, + "sampler_name": "euler", + "scheduler": "normal", + "denoise": 1.0, + "model": ["4", 0], + "positive": ["6", 0], + "negative": ["7", 0], + "latent_image": ["5", 0], + }, + }, + "4": { + "class_type": "CheckpointLoaderSimple", + "inputs": {"ckpt_name": "model.safetensors"}, + }, + "5": { + "class_type": "EmptyLatentImage", + "inputs": {"width": side, "height": side, "batch_size": 1}, + }, + "6": { + "class_type": "CLIPTextEncode", + "inputs": {"text": text, "clip": ["4", 1]}, + }, + "7": { + "class_type": "CLIPTextEncode", + "inputs": {"text": "blurry", "clip": ["4", 1]}, + }, + "9": { + "class_type": "SaveImage", + "inputs": {"filename_prefix": "img", "images": ["8", 0]}, + }, + } + workflow = { + "nodes": [ + {"id": int(k), "type": v["class_type"]} for k, v in prompt.items() + ] + } + info = PngImagePlugin.PngInfo() + info.add_text("prompt", json.dumps(prompt)) + info.add_text("workflow", json.dumps(workflow)) + image.save( + os.path.join(folder, f"m{index:03d}.png"), pnginfo=info, compress_level=4 + ) + + +class Driver: + def __init__(self, base, output, media, batch): + self.base, self.output, self.media, self.batch = base, output, media, batch + self.errors, self.batches = [], 0 + self.names = sorted(os.listdir(media)) + + def call(self, method, path, body=None, raw=False): + data = None if body is None else json.dumps(body).encode() + request = urllib.request.Request( # nosec B310 - loopback test server + self.base + path, + data=data, + method=method, + headers={"Content-Type": "application/json"}, + ) + try: + with urllib.request.urlopen( # nosec B310 + request, timeout=REQUEST_TIMEOUT + ) as response: + payload = response.read() + except Exception as error: # keep driving; errors are counted + self.errors.append(f"{method} {path.split('?')[0]}: {error}") + return None + return payload if raw else json.loads(payload or b"{}") + + def add_batch(self): + self.batches += 1 + folder = os.path.join(self.output, f"batch{self.batches:03d}") + os.makedirs(folder) + for index in range(self.batch): + name = self.names[(self.batches * self.batch + index) % len(self.names)] + shutil.copyfile( + os.path.join(self.media, name), + os.path.join(folder, f"img_{self.batches:03d}_{index:03d}.png"), + ) + + def wait_job(self, response): + job = ((response or {}).get("job") or {}).get("id") + for _ in range(JOB_POLL_LIMIT if job else 0): + state = self.call("GET", f"/prompt_manager/jobs/{job}") + if not state or (state.get("job") or {}).get("state") in FINAL_STATES: + return + time.sleep(1) + + def scan_and_thumbnails(self): + self.add_batch() + self.wait_job(self.call("POST", "/prompt_manager/jobs/scan", {})) + self.wait_job(self.call("POST", "/prompt_manager/jobs/thumbnails", {})) + + def browse(self, _index=None): + offset = random.randint(0, 3) * 100 # nosec B311 - load shape only + page = self.call( + "GET", + f"/prompt_manager/images/output?mode=workspace&limit=100&offset={offset}", + ) + images = [i for i in (page or {}).get("images") or [] if isinstance(i, dict)] + urls = [i.get("thumbnail_url") or i.get("url") for i in images] + urls = [u for u in urls if u] + for url in random.sample(urls, min(40, len(urls))): + self.call("GET", url, raw=True) + for image in random.sample(images, min(6, len(images))): + self.inspect(image) + if image.get("url"): + self.call("GET", image["url"], raw=True) + library = self.call("GET", "/prompt_manager/library?limit=100") or {} + entries = [e for e in library.get("items") or [] if isinstance(e, dict)] + for entry in random.sample(entries, min(3, len(entries))): + prompt_id = entry.get("id") or entry.get("positive_id") + if prompt_id is not None: + self.call( + "GET", f"/prompt_manager/library/{prompt_id}/history", raw=True + ) + self.call( + "GET", f"/prompt_manager/library/{prompt_id}/images", raw=True + ) + self.call("GET", "/prompt_manager/recent?limit=50", raw=True) + self.call("GET", "/prompt_manager/images/recent?limit=50", raw=True) + + def inspect(self, image): + relative, index = image.get("path"), image.get("root_index") + if not relative or not isinstance(index, int): + return + query = {"root": str(index)} + if image.get("root_scope"): + query["root_scope"] = image["root_scope"] + self.call( + "GET", + "/prompt_manager/images/inspect/" + + urllib.parse.quote(relative) + + "?" + + urllib.parse.urlencode(query), + raw=True, + ) + + +def _wait_for_output(samples_path, process): + while process.poll() is None: + with open(samples_path, encoding="utf-8") as samples: + for line in samples: + if line.startswith(OUTPUT_ANNOUNCEMENT): + return line[len(OUTPUT_ANNOUNCEMENT) :].strip() + time.sleep(0.2) + raise RuntimeError("the server exited before it started") + + +def summarize(samples_path, label, batches, errors): + rows = [] + with open(samples_path, encoding="utf-8") as samples: + for line in samples: + if line.startswith("{"): + rows.append(json.loads(line)) + if not rows: + raise RuntimeError("no samples were taken") + rss = [row["rss"] for row in rows] + late = [row for row in rows if row["t"] >= 120] or rows + end = rows[-1] + return ( + f"| {label} | {rss[0]} → {end['rss']} | {min(r['rss'] for r in late)}–" + f"{max(r['rss'] for r in late)} | {max(rss)} | {rows[0]['heap']} → " + f"{end['heap']} | {max(r['threads'] for r in rows)} | " + f"{max(r['arenas'] for r in rows)} | {batches} | {len(errors)} |" + ) + + +def drive(args): + work = args.out or tempfile.mkdtemp(prefix="pm-memory-") + media = os.path.join(work, "media") + if not os.path.isdir(media): + print(f"Writing {MEDIA_FILES} PNGs to {media}", flush=True) + make_media(media) + samples_path = os.path.join(work, f"{args.label}.samples") + open(samples_path, "w", encoding="utf-8").close() + command = [sys.executable, os.path.abspath(__file__), "--serve", samples_path] + command += ["--repo", os.path.abspath(args.repo)] + command += ["--mallopt"] if args.mallopt else [] + with open(os.path.join(work, f"{args.label}.stderr"), "w") as stderr: + process = subprocess.Popen( # nosec B603 + command, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=stderr + ) + try: + port = None + for line in process.stdout: + text = line.decode(errors="replace") + if text.startswith(PORT_ANNOUNCEMENT): + port = int(text[len(PORT_ANNOUNCEMENT) :]) + break + if port is None: + raise RuntimeError("the server announced no port") + output = _wait_for_output(samples_path, process) + driver = Driver(f"http://127.0.0.1:{port}", output, media, args.batch) + run(driver, args) + finally: + process.stdin.close() + process.wait(60) + print(f"Samples: {samples_path}") + print( + "| run | RSS start → end (MB) | after 2 min | max | heap (MB) | threads " + "| arenas | batches | errors |" + ) + print(summarize(samples_path, args.label, driver.batches, driver.errors)) + + +def run(driver, args): + driver.call("POST", "/__test__/reset", {}) + driver.scan_and_thumbnails() + deadline = time.monotonic() + args.minutes * 60 + next_batch = time.monotonic() + 60 + with ThreadPoolExecutor(args.clients) as pool: + while time.monotonic() < deadline: + list(pool.map(driver.browse, range(args.clients))) + if time.monotonic() > next_batch: + driver.scan_and_thumbnails() + next_batch = time.monotonic() + 60 + time.sleep(SAMPLE_SECONDS + 1) # one more sample after the load stops + + +def main(argv=None): + parser = argparse.ArgumentParser(description=__doc__.split("\n\n")[0]) + parser.add_argument("--repo", default=REPO, help="checkout to measure") + parser.add_argument("--label", default="run") + parser.add_argument("--minutes", type=float, default=12) + parser.add_argument("--clients", type=int, default=6) + parser.add_argument("--batch", type=int, default=60, help="new PNGs a minute") + parser.add_argument("--out", help="work folder (media, samples); default temp") + parser.add_argument( + "--mallopt", action="store_true", help="pin the mmap threshold (comparison)" + ) + parser.add_argument("--serve", metavar="SAMPLES", help=argparse.SUPPRESS) + args = parser.parse_args(argv) + if not sys.platform.startswith("linux"): + parser.error("Linux only: it reads /proc and glibc's malloc_info") + if args.serve: + serve(os.path.abspath(args.repo), args.serve, args.mallopt) + else: + drive(args) + + +if __name__ == "__main__": + main() diff --git a/utils/native_memory.py b/utils/native_memory.py new file mode 100644 index 00000000..c7f7c3b0 --- /dev/null +++ b/utils/native_memory.py @@ -0,0 +1,65 @@ +"""Give glibc's free heap pages back to the OS after a long image job (#523). + +Image work (decode, resize, encode) leaves image-sized free blocks in the +malloc arenas of the threads that did it; glibc keeps those pages resident. +``release_free_memory`` calls ``malloc_trim(0)`` once, which returns the free +pages of every arena and changes no allocator setting (no ``mallopt``, no +``MALLOC_ARENA_MAX``: ComfyUI owns this process's allocator). + +glibc only: on Windows, macOS and musl (Alpine) it does nothing. It never +raises; a missing or failing ``malloc_trim`` is a quiet ``False``. +""" + +import ctypes +import os +import sys +import threading + +_lock = threading.Lock() +_trim = None +_looked_up = False + + +def _glibc() -> bool: + """Whether this process runs on glibc (musl has no CS_GNU_LIBC_VERSION).""" + if not sys.platform.startswith("linux"): + return False + try: + version = os.confstr("CS_GNU_LIBC_VERSION") + except (AttributeError, ValueError, OSError): + return False + return bool(version) and version.startswith("glibc") + + +def _malloc_trim(): + """glibc's ``malloc_trim``, looked up once; None where there is none.""" + global _trim, _looked_up + with _lock: + if not _looked_up: + _looked_up = True + _trim = _lookup() + return _trim + + +def _lookup(): + if not _glibc(): + return None + try: + trim = ctypes.CDLL(None).malloc_trim + except (OSError, AttributeError): + return None + trim.argtypes = [ctypes.c_size_t] + trim.restype = ctypes.c_int + return trim + + +def release_free_memory() -> bool: + """Return free heap pages to the OS once (glibc); True when it ran.""" + trim = _malloc_trim() + if trim is None: + return False + try: + trim(0) + except Exception: # a foreign libc symbol: never fail the caller's job + return False + return True diff --git a/utils/parallel.py b/utils/parallel.py index b790794b..99094224 100644 --- a/utils/parallel.py +++ b/utils/parallel.py @@ -15,18 +15,36 @@ create new thread at interpreter shutdown") only once the interpreter is marked finalizing, after the ``atexit`` callbacks. ``refused_at_shutdown`` recognises both; ``note_refusal`` logs one DEBUG line per place, with no traceback. + +One long-lived pool (#523): every ``map_parallel`` call shares one pool of +``workers`` threads (``JOB_THREAD_PREFIX``), kept between jobs, instead of a +new pool per call. Image work then reuses the same few threads, and their +malloc arenas, rather than leaving each call's image-sized buffers in a new +set of arenas. A different ``workers`` count (the Worker threads setting +changed) replaces the pool; the old one finishes what it holds and exits. +``shutdown_job_pool`` stops it (the package teardown); the next call starts +a fresh one. """ import logging import threading -from concurrent.futures import ThreadPoolExecutor +import time +from concurrent.futures import ThreadPoolExecutor, wait from .exit_signal import checked SHUTDOWN_REFUSAL = "interpreter shutdown" +# ``submit`` on a pool already shut down (this module's teardown or a resize). +SHUT_POOL_REFUSAL = "after shutdown" +JOB_THREAD_PREFIX = "pm-worker" +# The package teardown's bound for the job threads to finish their item. +JOB_POOL_JOIN_SECONDS = 5 _logger = logging.getLogger("prompt_manager.parallel") _noted = set() _noted_lock = threading.Lock() +_pool_lock = threading.Lock() +_pool = None # (workers, ThreadPoolExecutor) +_retired = [] # pools replaced by a new size, finishing what they hold def refused_at_shutdown(error): @@ -45,39 +63,113 @@ def note_refusal(logger, what): logger.debug(f"{what}: refused, the interpreter is exiting") -def map_parallel(func, items, workers): - """Apply ``func`` to every item on up to ``workers`` threads, keeping order. +def job_pool_started(): + """Whether a job pool exists (the teardown stops it).""" + return _pool is not None or bool(_retired) - ``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. A pool refused at interpreter - exit runs the items it did not start on the calling thread (#375). An item + +def _job_pool(workers): + """The shared pool of ``workers`` threads, replacing one of another size.""" + global _pool + with _pool_lock: + if _pool is not None and _pool[0] == workers: + return _pool[1] + retired = _pool + _pool = ( + workers, + ThreadPoolExecutor(workers, thread_name_prefix=JOB_THREAD_PREFIX), + ) + current = _pool[1] + if retired is not None: + retired[1].shutdown(wait=False) # it still runs what it holds + with _pool_lock: + _retired[:] = [p for p in _retired if _alive(p)] + [retired[1]] + return current + + +def shutdown_job_pool(timeout=None): + """Stop the job pool (and any pool a resize retired): queued items are + cancelled (their callers run them on their own thread), running ones are + joined up to ``timeout`` seconds. True when no job thread is left.""" + global _pool + with _pool_lock: + current, _pool = _pool, None + executors = _retired[:] + ([] if current is None else [current[1]]) + _retired.clear() + threads = [] + for executor in executors: + # After shutdown: it takes the lock ``submit`` holds while it starts + # and records a thread, so every thread is listed and none starts. + executor.shutdown(wait=False, cancel_futures=True) + threads.extend(getattr(executor, "_threads", ())) + deadline = None if timeout is None else time.monotonic() + timeout + for thread in threads: + thread.join(None if deadline is None else max(0.0, deadline - time.monotonic())) + return not any(thread.is_alive() for thread in threads) + + +def _alive(executor): + return any(t.is_alive() for t in getattr(executor, "_threads", ())) + + +def _forget(executor): + """Drop *executor* as the shared pool (a later call starts a new one).""" + global _pool + with _pool_lock: + if _pool is not None and _pool[1] is executor: + _pool = None + + +def _on_job_thread(): + return threading.current_thread().name.startswith(JOB_THREAD_PREFIX + "_") + + +def map_parallel(func, items, workers): + """Apply ``func`` to every item on the shared pool of ``workers`` + threads, keeping order. + + ``workers`` below 1 counts as 1. A call made from a job thread runs on + that thread, so nested work never waits for its own pool. An exception + in ``func`` propagates to the caller once the call's other items have + stopped (those not started are dropped), so callers that must survive + bad input wrap ``func`` themselves. A pool refused at interpreter exit + runs the items it did not start on the calling thread (#375). An item that has not started when the interpreter starts exiting raises ExitRequested instead (#376). """ func = checked(func) items = list(items) - workers = max(1, int(workers or 1)) - if workers == 1 or len(items) <= 1: + if not items or _on_job_thread(): return [func(item) for item in items] - max_workers = min(workers, len(items)) - with ThreadPoolExecutor(max_workers, thread_name_prefix="pm-worker") as pool: - futures = [] - try: - for item in items: - futures.append(pool.submit(func, item)) - except RuntimeError as error: - if not refused_at_shutdown(error): - raise + pool = _job_pool(max(1, int(workers or 1))) + futures = [] + try: + for item in items: + futures.append(pool.submit(func, item)) + except RuntimeError as error: + if refused_at_shutdown(error): note_refusal(_logger, "A worker pool") # Drop what was queued but not started; it runs here instead. # (A refused Thread.start leaves its item queued without a future # here: it is cancelled too, unless a busy pool thread took it in # between, the one case where an item can run twice.) - pool.shutdown(wait=True, cancel_futures=True) - else: - return [future.result() for future in futures] - return _run_rest(func, items, futures) + _forget(pool) + pool.shutdown(wait=False, cancel_futures=True) + elif SHUT_POOL_REFUSAL not in str(error): + raise + return _collect(func, items, futures) + + +def _collect(func, items, futures): + """``_run_rest``; on an error, first drop the call's items not started + and wait for its running ones, as the old per-call pool's exit did.""" + try: + return _run_rest(func, items, futures) + except BaseException: + for future in futures: + future.cancel() + wait(futures) + raise def _run_rest(func, items, futures):