Add DistributedJoin and dynamic list-queue splitter mode

tests/api/test_usdu_routes.py

@@ -103,6 +103,21 @@ def _load_usdu_routes_module():
     job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
     job_store_module.MAX_PAYLOAD_SIZE = 1024
     job_store_module.ensure_tile_jobs_initialized = lambda: prompt_server_holder["value"]
+
+    async def _init_dynamic_job(multi_job_id, batch_size, enabled_workers, all_indices=None):
+        prompt_server = prompt_server_holder["value"]
+        async with prompt_server.distributed_tile_jobs_lock:
+            if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+                return
+            job_data = ImageJobState(multi_job_id=multi_job_id)
+            job_data.worker_status = {str(worker_id): 0 for worker_id in (enabled_workers or [])}
+            job_data.assigned_to_workers = {str(worker_id): [] for worker_id in (enabled_workers or [])}
+            indices = all_indices if all_indices is not None else list(range(int(batch_size or 0)))
+            for idx in indices:
+                await job_data.pending_images.put(int(idx))
+            prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data
+
+    job_store_module.init_dynamic_job = _init_dynamic_job
     sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
 
     job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
@@ -242,6 +257,34 @@ class USDURoutesTests(unittest.IsolatedAsyncioTestCase):
         self.assertTrue(response.payload.get("batched_static"))
         self.assertEqual(job_data.assigned_to_workers["worker-a"], [4])
 
+    async def test_init_list_queue_initializes_pending_items(self):
+        request = _FakeRequest(
+            json_payload={
+                "multi_job_id": "list-job-1",
+                "list_size": 3,
+                "enabled_workers": ["worker-a", "worker-b"],
+            }
+        )
+        response = await usdu_routes.init_list_queue_endpoint(request)
+
+        self.assertEqual(response.status, 200)
+        self.assertEqual(response.payload.get("status"), "success")
+        self.assertEqual(response.payload.get("remaining"), 3)
+
+    async def test_request_list_item_assigns_next_index(self):
+        prompt_server = usdu_routes._prompt_server_holder["value"]
+        job_data = usdu_routes._ImageJobState("list-job-2")
+        await job_data.pending_images.put(11)
+        prompt_server.distributed_pending_tile_jobs["list-job-2"] = job_data
+
+        request = _FakeRequest(json_payload={"worker_id": "worker-a", "multi_job_id": "list-job-2"})
+        response = await usdu_routes.request_list_item_endpoint(request)
+
+        self.assertEqual(response.status, 200)
+        self.assertEqual(response.payload.get("item_idx"), 11)
+        self.assertEqual(response.payload.get("estimated_remaining"), 0)
+        self.assertEqual(job_data.assigned_to_workers["worker-a"], [11])
+
     async def test_submit_tiles_completion_signal_enqueues_last_marker(self):
         prompt_server = usdu_routes._prompt_server_holder["value"]
         job_data = usdu_routes._TileJobState("job-4")

tests/test_join.py

@@ -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()

tests/test_list_splitter.py

@@ -1,7 +1,9 @@
+import asyncio
 import importlib.util
 import sys
 import types
 import unittest
+from dataclasses import dataclass, field
 from pathlib import Path
 
 import torch
@@ -24,11 +26,77 @@ def _bootstrap_package(package_name):
     utils_pkg.__path__ = []
     sys.modules[f"{package_name}.utils"] = utils_pkg
 
+    upscale_pkg = types.ModuleType(f"{package_name}.upscale")
+    upscale_pkg.__path__ = []
+    sys.modules[f"{package_name}.upscale"] = upscale_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
 
+    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
+
+    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
+
+    job_models_module = types.ModuleType(f"{package_name}.upscale.job_models")
+
+    class BaseJobState:
+        pass
+
+    @dataclass
+    class ImageJobState(BaseJobState):
+        multi_job_id: str
+        mode: str = field(default="dynamic", init=False)
+        queue: asyncio.Queue = field(default_factory=asyncio.Queue)
+        pending_images: asyncio.Queue = field(default_factory=asyncio.Queue)
+        completed_images: dict = field(default_factory=dict)
+        worker_status: dict = field(default_factory=dict)
+        assigned_to_workers: dict = field(default_factory=dict)
+
+        @property
+        def pending_tasks(self):
+            return self.pending_images
+
+    job_models_module.BaseJobState = BaseJobState
+    job_models_module.ImageJobState = ImageJobState
+    sys.modules[f"{package_name}.upscale.job_models"] = job_models_module
+
+    prompt_server_holder = {
+        "value": types.SimpleNamespace(
+            distributed_tile_jobs_lock=asyncio.Lock(),
+            distributed_pending_tile_jobs={},
+        )
+    }
+
+    job_store_module = types.ModuleType(f"{package_name}.upscale.job_store")
+
+    def ensure_tile_jobs_initialized():
+        return prompt_server_holder["value"]
+
+    async def init_dynamic_job(multi_job_id, batch_size, enabled_workers, all_indices=None):
+        prompt_server = prompt_server_holder["value"]
+        async with prompt_server.distributed_tile_jobs_lock:
+            if multi_job_id in prompt_server.distributed_pending_tile_jobs:
+                return
+            job_data = ImageJobState(multi_job_id=multi_job_id)
+            for idx in (all_indices if all_indices is not None else list(range(batch_size))):
+                await job_data.pending_images.put(idx)
+            prompt_server.distributed_pending_tile_jobs[multi_job_id] = job_data
+
+    job_store_module.ensure_tile_jobs_initialized = ensure_tile_jobs_initialized
+    job_store_module.init_dynamic_job = init_dynamic_job
+    job_store_module._prompt_server_holder = prompt_server_holder
+    sys.modules[f"{package_name}.upscale.job_store"] = job_store_module
+
 
 def _load_module(package_name, module_rel_path, module_name):
     module_path = Path(__file__).resolve().parents[1] / module_rel_path
@@ -47,29 +115,60 @@ def _load_list_splitter_module():
     package_name = "dist_list_splitter_testpkg"
     _bootstrap_package(package_name)
     _load_module(package_name, "nodes/utilities.py", "nodes.utilities")
-    return _load_module(package_name, "nodes/list_splitter.py", "nodes.list_splitter")
+    module = _load_module(package_name, "nodes/list_splitter.py", "nodes.list_splitter")
+    return module
 
 
 splitter_module = _load_list_splitter_module()
 
 
+class _FakeResponse:
+    def __init__(self, payload, status=200):
+        self.payload = payload
+        self.status = status
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        return False
+
+    async def json(self):
+        return self.payload
+
+
+class _FakeSession:
+    def __init__(self, responses):
+        self.responses = list(responses)
+        self.calls = []
+
+    def post(self, url, json=None, timeout=None):
+        self.calls.append({"url": url, "json": json})
+        if not self.responses:
+            return _FakeResponse({"item_idx": None}, status=200)
+        payload = self.responses.pop(0)
+        return _FakeResponse(payload, status=200)
+
+
 class DistributedListSplitterTests(unittest.TestCase):
     def setUp(self):
         self.node = splitter_module.DistributedListSplitter()
+        prompt_server = splitter_module.ensure_tile_jobs_initialized()
+        prompt_server.distributed_pending_tile_jobs = {}
 
     def _images(self, count):
         return [torch.full((1, 4, 4, 3), float(idx), dtype=torch.float32) for idx in range(count)]
 
     def test_single_participant_returns_full_list(self):
         images = self._images(5)
-        out = self.node.split(images, participant_index=0, total_participants=1)[0]
+        out = self.node.split(images, mode="static", participant_index=0, total_participants=1)[0]
         self.assertEqual(len(out), 5)
         self.assertEqual([int(t[0, 0, 0, 0].item()) for t in out], [0, 1, 2, 3, 4])
 
     def test_two_participants_split_six_images_evenly(self):
         images = self._images(6)
-        p0 = self.node.split(images, participant_index=0, total_participants=2)[0]
-        p1 = self.node.split(images, participant_index=1, total_participants=2)[0]
+        p0 = self.node.split(images, mode="static", participant_index=0, total_participants=2)[0]
+        p1 = self.node.split(images, mode="static", participant_index=1, total_participants=2)[0]
         self.assertEqual(len(p0), 3)
         self.assertEqual(len(p1), 3)
         self.assertEqual([int(t[0, 0, 0, 0].item()) for t in p0], [0, 1, 2])
@@ -77,17 +176,14 @@ class DistributedListSplitterTests(unittest.TestCase):
 
     def test_three_participants_split_seven_images(self):
         images = self._images(7)
-        p0 = self.node.split(images, participant_index=0, total_participants=3)[0]
-        p1 = self.node.split(images, participant_index=1, total_participants=3)[0]
-        p2 = self.node.split(images, participant_index=2, total_participants=3)[0]
+        p0 = self.node.split(images, mode="static", participant_index=0, total_participants=3)[0]
+        p1 = self.node.split(images, mode="static", participant_index=1, total_participants=3)[0]
+        p2 = self.node.split(images, mode="static", participant_index=2, total_participants=3)[0]
         self.assertEqual([len(p0), len(p1), len(p2)], [3, 2, 2])
-        self.assertEqual([int(t[0, 0, 0, 0].item()) for t in p0], [0, 1, 2])
-        self.assertEqual([int(t[0, 0, 0, 0].item()) for t in p1], [3, 4])
-        self.assertEqual([int(t[0, 0, 0, 0].item()) for t in p2], [5, 6])
 
     def test_more_participants_than_images_yields_empty_slices(self):
         images = self._images(2)
-        outputs = [self.node.split(images, participant_index=i, total_participants=5)[0] for i in range(5)]
+        outputs = [self.node.split(images, mode="static", participant_index=i, total_participants=5)[0] for i in range(5)]
         self.assertEqual([len(chunk) for chunk in outputs], [1, 1, 0, 0, 0])
 
     def test_mixed_dimensions_are_preserved(self):
@@ -96,13 +192,70 @@ class DistributedListSplitterTests(unittest.TestCase):
             torch.zeros((1, 8, 6, 3), dtype=torch.float32),
             torch.zeros((1, 3, 9, 3), dtype=torch.float32),
         ]
-
-        p0 = self.node.split(images, participant_index=0, total_participants=2)[0]
-        p1 = self.node.split(images, participant_index=1, total_participants=2)[0]
-
+        p0 = self.node.split(images, mode="static", participant_index=0, total_participants=2)[0]
+        p1 = self.node.split(images, mode="static", participant_index=1, total_participants=2)[0]
         self.assertEqual([tuple(t.shape) for t in p0], [(1, 4, 4, 3), (1, 8, 6, 3)])
         self.assertEqual([tuple(t.shape) for t in p1], [(1, 3, 9, 3)])
 
+    def test_dynamic_mode_master_pulls_from_local_queue(self):
+        images = self._images(4)
+
+        pulled = self.node.split(
+            images,
+            mode="dynamic",
+            multi_job_id="job-dynamic",
+            is_worker=False,
+            worker_id="master",
+        )[0]
+
+        self.assertEqual([int(t[0, 0, 0, 0].item()) for t in pulled], [0, 1, 2, 3])
+
+    def test_dynamic_mode_worker_pulls_via_http(self):
+        images = self._images(5)
+        fake_session = _FakeSession([
+            {"item_idx": 1, "estimated_remaining": 3},
+            {"item_idx": 3, "estimated_remaining": 2},
+            {"item_idx": None},
+        ])
+
+        async def _get_session():
+            return fake_session
+
+        splitter_module.get_client_session = _get_session
+
+        pulled = self.node.split(
+            images,
+            mode="dynamic",
+            multi_job_id="job-dynamic",
+            is_worker=True,
+            master_url="http://master.local:8188",
+            worker_id="worker-a",
+        )[0]
+
+        self.assertEqual([int(t[0, 0, 0, 0].item()) for t in pulled], [1, 3])
+        self.assertGreaterEqual(len(fake_session.calls), 3)
+
+    def test_dynamic_mode_exhausted_queue_returns_empty(self):
+        images = self._images(3)
+
+        first_pull = self.node.split(
+            images,
+            mode="dynamic",
+            multi_job_id="job-dynamic",
+            is_worker=False,
+            worker_id="master",
+        )[0]
+        second_pull = self.node.split(
+            images,
+            mode="dynamic",
+            multi_job_id="job-dynamic",
+            is_worker=False,
+            worker_id="master",
+        )[0]
+
+        self.assertEqual(len(first_pull), 3)
+        self.assertEqual(second_pull, [])
+
 
 if __name__ == "__main__":
     unittest.main()

tests/test_prompt_transform.py

@@ -103,6 +103,18 @@ def _branch_prompt():
     }
 
 
+def _join_prompt():
+    """1 → 2(DistributedBranch) -> 3(branch0) and 4(branch1) -> 5(DistributedJoin) -> 6(SaveImage)."""
+    return {
+        "1": {"class_type": "KSampler", "inputs": {}},
+        "2": {"class_type": "DistributedBranch", "inputs": {"input": ["1", 0], "num_branches": 2}},
+        "3": {"class_type": "Blur", "inputs": {"image": ["2", 0]}},
+        "4": {"class_type": "Sharpen", "inputs": {"image": ["2", 1]}},
+        "5": {"class_type": "DistributedJoin", "inputs": {"input": ["3", 0], "num_branches": 2}},
+        "6": {"class_type": "SaveImage", "inputs": {"images": ["5", 0]}},
+    }
+
+
 def _apply(prompt, participant_id, enabled_worker_ids=None, delegate_master=False):
     if enabled_worker_ids is None:
         enabled_worker_ids = ["worker-a", "worker-b"]
@@ -290,6 +302,12 @@ class PrunePromptForWorkerTests(unittest.TestCase):
         self.assertIn("5", result)
         self.assertIn("6", result)
 
+    def test_join_anchor_keeps_join_and_downstream(self):
+        prompt = _join_prompt()
+        result = pt.prune_prompt_for_worker(prompt)
+        self.assertIn("5", result)
+        self.assertIn("6", result)
+
 
 # ---------------------------------------------------------------------------
 # prepare_delegate_master_prompt
@@ -375,6 +393,22 @@ class GenerateJobIdMapTests(unittest.TestCase):
         job_map = pt.generate_job_id_map(idx, "run")
         self.assertEqual(job_map["8"], "run_8")
 
+    def test_maps_list_splitter_nodes(self):
+        prompt = {
+            "7": {"class_type": "DistributedListSplitter", "inputs": {}},
+        }
+        idx = pt.PromptIndex(prompt)
+        job_map = pt.generate_job_id_map(idx, "run")
+        self.assertEqual(job_map["7"], "run_7")
+
+    def test_maps_join_nodes(self):
+        prompt = {
+            "10": {"class_type": "DistributedJoin", "inputs": {}},
+        }
+        idx = pt.PromptIndex(prompt)
+        job_map = pt.generate_job_id_map(idx, "run")
+        self.assertEqual(job_map["10"], "run_10")
+
     def test_maps_branch_nodes(self):
         prompt = {
             "9": {"class_type": "DistributedBranch", "inputs": {"num_branches": 2}},
@@ -553,11 +587,18 @@ class ApplyOverridesListSplitterTests(unittest.TestCase):
         result = _apply(self._splitter_prompt(), "master", enabled_worker_ids=["worker-a", "worker-b"])
         self.assertEqual(result["1"]["inputs"]["participant_index"], 0)
         self.assertEqual(result["1"]["inputs"]["total_participants"], 3)
+        self.assertFalse(result["1"]["inputs"]["is_worker"])
+        self.assertNotIn("master_url", result["1"]["inputs"])
+        self.assertEqual(result["1"]["inputs"]["worker_id"], "master")
+        self.assertEqual(result["1"]["inputs"]["multi_job_id"], "run_1")
 
     def test_worker_gets_incremented_participant_index(self):
         result = _apply(self._splitter_prompt(), "worker-a", enabled_worker_ids=["worker-a", "worker-b"])
         self.assertEqual(result["1"]["inputs"]["participant_index"], 1)
         self.assertEqual(result["1"]["inputs"]["total_participants"], 3)
+        self.assertTrue(result["1"]["inputs"]["is_worker"])
+        self.assertEqual(result["1"]["inputs"]["worker_id"], "worker-a")
+        self.assertEqual(result["1"]["inputs"]["master_url"], "http://master.example.com")
 
     def test_delegate_mode_shifts_worker_to_index_zero(self):
         result = _apply(
@@ -648,5 +689,33 @@ class ApplyOverridesBranchTests(unittest.TestCase):
         self.assertNotIn("5", worker_b_result)
         self.assertNotIn("6", worker_b_result)
 
+
+class ApplyOverridesJoinTests(unittest.TestCase):
+    def test_join_inherits_assigned_branch_from_upstream_branch_node(self):
+        master_prompt = {
+            "1": {"class_type": "KSampler", "inputs": {}},
+            "2": {"class_type": "DistributedBranch", "inputs": {"input": ["1", 0], "num_branches": 2}},
+            "3": {"class_type": "DistributedJoin", "inputs": {"input": ["2", 0], "num_branches": 2}},
+        }
+        worker_prompt = {
+            "1": {"class_type": "KSampler", "inputs": {}},
+            "2": {"class_type": "DistributedBranch", "inputs": {"input": ["1", 0], "num_branches": 2}},
+            "3": {"class_type": "DistributedJoin", "inputs": {"input": ["2", 1], "num_branches": 2}},
+        }
+
+        master_result = _apply(master_prompt, "master", enabled_worker_ids=["worker-a"])
+        worker_result = _apply(worker_prompt, "worker-a", enabled_worker_ids=["worker-a"])
+
+        self.assertEqual(master_result["3"]["inputs"]["assigned_branch"], 0)
+        self.assertEqual(worker_result["3"]["inputs"]["assigned_branch"], 1)
+
+    def test_join_sets_multi_job_id(self):
+        prompt = {
+            "1": {"class_type": "DistributedJoin", "inputs": {"input": ["2", 0]}},
+            "2": {"class_type": "KSampler", "inputs": {}},
+        }
+        result = _apply(prompt, "master", enabled_worker_ids=["worker-a"])
+        self.assertEqual(result["1"]["inputs"]["multi_job_id"], "run_1")
+
 if __name__ == "__main__":
     unittest.main()