feat(worker): terminate pruned worker prompts with a hidden no-op sink (#99)
Worker prompts drop everything downstream of DistributedCollector / UltimateSDUpscaleDistributed, so the master re-added a PreviewImage (or PreviewAudio) to give ComfyUI an output node. That preview wrote every result frame as a PNG into the worker's temp directory: for a 362-frame 1344x768 video, ~2 GB per run (RAM when temp is on tmpfs) plus CPU time spent PNG-encoding frames nobody looks at. - Add DistributedWorkerSink: an output node that discards its input. It is is_dev_only, so it stays out of node search and menus. - The master probes each worker with GET /object_info/DistributedWorkerSink while preparing its prompt and uses the sink when it is registered. Workers on older versions (or a failed probe) keep the PreviewImage / PreviewAudio terminator, so mixed versions keep working. Verified with a real master plus one updated and one old worker: the updated worker ran EmptyImage -> Collector -> Sink and left its temp directory empty; the old worker still received PreviewImage. Claude-Session: https://claude.ai/code/session_0166t8T3CVaxZkJbP2RtSpo6 Co-authored-by: ymiyad <cab-13-might@icloud.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
co-authored by
ymiyad
Claude Opus 5.5
parent
9345a8e71b
commit
40678a51c4
@@ -245,3 +245,70 @@ class RewritePromptMediaInputsTests(unittest.TestCase):
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# fetch_worker_has_node
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
class _FakeResponse:
|
||||
def __init__(self, status, payload):
|
||||
self.status = status
|
||||
self._payload = payload
|
||||
|
||||
async def json(self):
|
||||
if isinstance(self._payload, Exception):
|
||||
raise self._payload
|
||||
return self._payload
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, *exc):
|
||||
return False
|
||||
|
||||
|
||||
class _FakeSession:
|
||||
def __init__(self, response=None, error=None):
|
||||
self.response = response
|
||||
self.error = error
|
||||
self.urls = []
|
||||
|
||||
def get(self, url, timeout=None):
|
||||
self.urls.append(url)
|
||||
if self.error:
|
||||
raise self.error
|
||||
return self.response
|
||||
|
||||
|
||||
class FetchWorkerHasNodeTests(unittest.TestCase):
|
||||
def _probe(self, session):
|
||||
import asyncio
|
||||
|
||||
async def _get_session():
|
||||
return session
|
||||
|
||||
original = ms.get_client_session
|
||||
ms.get_client_session = _get_session
|
||||
try:
|
||||
return asyncio.run(ms.fetch_worker_has_node({"id": "w1"}, "DistributedWorkerSink"))
|
||||
finally:
|
||||
ms.get_client_session = original
|
||||
|
||||
def test_registered_node_is_reported(self):
|
||||
session = _FakeSession(_FakeResponse(200, {"DistributedWorkerSink": {"name": "DistributedWorkerSink"}}))
|
||||
self.assertTrue(self._probe(session))
|
||||
self.assertEqual(session.urls, ["http://localhost/object_info/DistributedWorkerSink"])
|
||||
|
||||
def test_unknown_node_returns_empty_object(self):
|
||||
# ComfyUI answers 200 {} for classes that are not registered.
|
||||
self.assertFalse(self._probe(_FakeSession(_FakeResponse(200, {}))))
|
||||
|
||||
def test_http_error_counts_as_missing(self):
|
||||
self.assertFalse(self._probe(_FakeSession(_FakeResponse(404, None))))
|
||||
|
||||
def test_bad_json_counts_as_missing(self):
|
||||
self.assertFalse(self._probe(_FakeSession(_FakeResponse(200, ValueError("bad json")))))
|
||||
|
||||
def test_connection_error_counts_as_missing(self):
|
||||
self.assertFalse(self._probe(_FakeSession(error=OSError("unreachable"))))
|
||||
|
||||
Reference in New Issue
Block a user