diff --git a/nodes/branch_collector.py b/nodes/branch_collector.py index d85c0cc..4cabd91 100644 --- a/nodes/branch_collector.py +++ b/nodes/branch_collector.py @@ -79,6 +79,14 @@ class DistributedBranchCollector: if assigned_value is not None: return assigned_value, assigned_idx + # In participant-pruned prompts, each participant should typically have + # exactly one local branch input. If socket names drifted (e.g. branch_2 + # connected for a participant assigned to slot 0), remap that sole value + # to the assigned slot so the participant still contributes correctly. + non_none_indices = [idx for idx, value in enumerate(branch_values) if value is not None] + if len(non_none_indices) == 1: + return branch_values[non_none_indices[0]], assigned_idx + for idx, value in enumerate(branch_values): if value is not None: return value, idx @@ -256,6 +264,8 @@ class DistributedBranchCollector: outputs[idx] = branch_values[idx] local_value, local_branch_idx = self._resolve_local_branch_value(branch_values, assigned_branch) + if 0 <= local_branch_idx < MAX_BRANCH_OUTPUTS and local_value is not None: + outputs[local_branch_idx] = local_value if is_worker: if 0 <= local_branch_idx < MAX_BRANCH_OUTPUTS and local_value is not None: diff --git a/tests/test_join.py b/tests/test_join.py index b5a9ab3..f537b10 100644 --- a/tests/test_join.py +++ b/tests/test_join.py @@ -223,6 +223,26 @@ class DistributedBranchCollectorTests(unittest.IsolatedAsyncioTestCase): self.assertEqual(float(outputs[1].abs().sum().item()), 0.0) self.assertEqual(float(outputs[2].abs().sum().item()), 0.0) + async def test_assigned_slot_can_be_filled_from_single_mismatched_local_input(self): + fake_server = _FakePromptServer() + branch_collector_module._get_prompt_server_instance = lambda: fake_server + + # No worker results expected for this check. + fake_server.distributed_pending_jobs["join-job"] = asyncio.Queue() + + local_tensor = torch.full((1, 4, 4, 3), 0.33, dtype=torch.float32) + outputs = await self.node.execute( + # Local value arrives on slot 1 even though participant is assigned slot 0. + [None, local_tensor, None, None, None, None, None, None, None, None], + num_branches=3, + multi_job_id="join-job", + is_worker=False, + enabled_worker_ids=json.dumps([]), + assigned_branch=0, + ) + + self.assertTrue(torch.equal(outputs[0], local_tensor)) + if __name__ == "__main__": unittest.main()