Add DistributedJoin and dynamic list-queue splitter mode

2026-03-01 23:29:50 +00:00
parent e37a001d59
commit 34123a0d30
13 changed files with 1096 additions and 24 deletions
@@ -0,0 +1,204 @@
+import asyncio
+import importlib.util
+import json
+import sys
+import types
+import unittest
+from pathlib import Path
+
+import torch
+from PIL import Image
+
+
+def _bootstrap_package(package_name):
+    for mod_name in list(sys.modules):
+        if mod_name == package_name or mod_name.startswith(f"{package_name}."):
+            del sys.modules[mod_name]
+
+    root_pkg = types.ModuleType(package_name)
+    root_pkg.__path__ = []
+    sys.modules[package_name] = root_pkg
+
+    nodes_pkg = types.ModuleType(f"{package_name}.nodes")
+    nodes_pkg.__path__ = []
+    sys.modules[f"{package_name}.nodes"] = nodes_pkg
+
+    utils_pkg = types.ModuleType(f"{package_name}.utils")
+    utils_pkg.__path__ = []
+    sys.modules[f"{package_name}.utils"] = utils_pkg
+
+    logging_module = types.ModuleType(f"{package_name}.utils.logging")
+    logging_module.debug_log = lambda *_args, **_kwargs: None
+    logging_module.log = lambda *_args, **_kwargs: None
+    sys.modules[f"{package_name}.utils.logging"] = logging_module
+
+    config_module = types.ModuleType(f"{package_name}.utils.config")
+    config_module.get_worker_timeout_seconds = lambda: 0.2
+    config_module.is_master_delegate_only = lambda: False
+    sys.modules[f"{package_name}.utils.config"] = config_module
+
+    constants_module = types.ModuleType(f"{package_name}.utils.constants")
+    constants_module.HEARTBEAT_INTERVAL = 1.0
+    sys.modules[f"{package_name}.utils.constants"] = constants_module
+
+    image_module = types.ModuleType(f"{package_name}.utils.image")
+
+    def ensure_contiguous(tensor):
+        return tensor.contiguous()
+
+    def tensor_to_pil(image_batch, index):
+        tensor = image_batch[index].detach().cpu().clamp(0, 1)
+        array = (tensor.numpy() * 255.0).astype("uint8")
+        return Image.fromarray(array)
+
+    image_module.ensure_contiguous = ensure_contiguous
+    image_module.tensor_to_pil = tensor_to_pil
+    sys.modules[f"{package_name}.utils.image"] = image_module
+
+    network_module = types.ModuleType(f"{package_name}.utils.network")
+
+    async def get_client_session():
+        raise RuntimeError("test should patch get_client_session")
+
+    network_module.get_client_session = get_client_session
+    sys.modules[f"{package_name}.utils.network"] = network_module
+
+    async_helpers_module = types.ModuleType(f"{package_name}.utils.async_helpers")
+    async_helpers_module.run_async_in_server_loop = lambda coro: asyncio.run(coro)
+    sys.modules[f"{package_name}.utils.async_helpers"] = async_helpers_module
+
+
+def _load_module(package_name, module_rel_path, module_name):
+    module_path = Path(__file__).resolve().parents[1] / module_rel_path
+    spec = importlib.util.spec_from_file_location(
+        f"{package_name}.{module_name}",
+        module_path,
+    )
+    module = importlib.util.module_from_spec(spec)
+    assert spec is not None and spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+def _load_join_module():
+    package_name = "dist_join_testpkg"
+    _bootstrap_package(package_name)
+    _load_module(package_name, "nodes/utilities.py", "nodes.utilities")
+    return _load_module(package_name, "nodes/join.py", "nodes.join")
+
+
+join_module = _load_join_module()
+
+
+class _FakeResponseCtx:
+    def __init__(self, recorder, payload, url):
+        self._recorder = recorder
+        self._payload = payload
+        self._url = url
+
+    async def __aenter__(self):
+        self._recorder.append({"url": self._url, "payload": self._payload})
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        return False
+
+    def raise_for_status(self):
+        return None
+
+
+class _FakeSession:
+    def __init__(self):
+        self.calls = []
+
+    def post(self, url, json=None, timeout=None):
+        return _FakeResponseCtx(self.calls, json, url)
+
+
+class _FakePromptServer:
+    def __init__(self):
+        self.distributed_pending_jobs = {}
+        self.distributed_jobs_lock = asyncio.Lock()
+
+
+class DistributedJoinTests(unittest.IsolatedAsyncioTestCase):
+    def setUp(self):
+        self.node = join_module.DistributedJoin()
+
+    async def test_worker_mode_sends_result_with_assigned_branch_index(self):
+        fake_session = _FakeSession()
+
+        async def _get_session():
+            return fake_session
+
+        join_module.get_client_session = _get_session
+
+        tensor = torch.ones((1, 4, 4, 3), dtype=torch.float32)
+        outputs = await self.node.execute(
+            tensor,
+            num_branches=2,
+            multi_job_id="join-job",
+            is_worker=True,
+            master_url="http://master.local:8188",
+            worker_id="worker-a",
+            assigned_branch=1,
+        )
+
+        self.assertEqual(len(fake_session.calls), 1)
+        payload = fake_session.calls[0]["payload"]
+        self.assertEqual(payload["batch_idx"], 1)
+        self.assertTrue(payload["is_last"])
+        self.assertEqual(payload["worker_id"], "worker-a")
+        self.assertTrue(torch.equal(outputs[1], tensor))
+
+    async def test_master_mode_collects_worker_results_into_branch_slots(self):
+        fake_server = _FakePromptServer()
+        join_module._get_prompt_server_instance = lambda: fake_server
+
+        queue = asyncio.Queue()
+        worker_1 = torch.full((1, 4, 4, 3), 0.5, dtype=torch.float32)
+        worker_2 = torch.full((1, 4, 4, 3), 0.8, dtype=torch.float32)
+        await queue.put({"tensor": worker_1, "worker_id": "worker-a", "image_index": 1, "is_last": True})
+        await queue.put({"tensor": worker_2, "worker_id": "worker-b", "image_index": 2, "is_last": True})
+        fake_server.distributed_pending_jobs["join-job"] = queue
+
+        master_tensor = torch.zeros((1, 4, 4, 3), dtype=torch.float32)
+        outputs = await self.node.execute(
+            master_tensor,
+            num_branches=3,
+            multi_job_id="join-job",
+            is_worker=False,
+            enabled_worker_ids=json.dumps(["worker-a", "worker-b"]),
+            assigned_branch=0,
+        )
+
+        self.assertTrue(torch.equal(outputs[0], master_tensor))
+        self.assertTrue(torch.equal(outputs[1], worker_1))
+        self.assertTrue(torch.equal(outputs[2], worker_2))
+
+    async def test_unassigned_slots_remain_none(self):
+        fake_server = _FakePromptServer()
+        join_module._get_prompt_server_instance = lambda: fake_server
+
+        queue = asyncio.Queue()
+        worker_1 = torch.full((1, 4, 4, 3), 0.5, dtype=torch.float32)
+        await queue.put({"tensor": worker_1, "worker_id": "worker-a", "image_index": 1, "is_last": True})
+        fake_server.distributed_pending_jobs["join-job"] = queue
+
+        master_tensor = torch.zeros((1, 4, 4, 3), dtype=torch.float32)
+        outputs = await self.node.execute(
+            master_tensor,
+            num_branches=2,
+            multi_job_id="join-job",
+            is_worker=False,
+            enabled_worker_ids=json.dumps(["worker-a"]),
+            assigned_branch=0,
+        )
+
+        self.assertIsNone(outputs[2])
+        self.assertIsNone(outputs[9])
+
+
+if __name__ == "__main__":
+    unittest.main()