"""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)