diff --git a/nodes/join.py b/nodes/join.py index c162bd3..429b2fd 100644 --- a/nodes/join.py +++ b/nodes/join.py @@ -115,21 +115,38 @@ class DistributedJoin: return workers def _expected_worker_ids_for_branches(self, enabled_workers, delegate_mode, num_branches): - participants = list(enabled_workers) if delegate_mode else ["master"] + list(enabled_workers) - participant_count = len(participants) - if participant_count <= 0: - return [] - + slots_by_participant = self._participant_branch_slots(enabled_workers, delegate_mode, num_branches) expected = [] - for participant_index, participant_id in enumerate(participants): - # A participant receives at least one branch iff it appears in the first `num_branches` slots. - if participant_index >= num_branches: - continue - if participant_id == "master": + for participant_id, slots in slots_by_participant.items(): + if participant_id == "master" or not slots: continue expected.append(str(participant_id)) return expected + def _participant_branch_slots(self, enabled_workers, delegate_mode, num_branches): + participants = list(enabled_workers) if delegate_mode else ["master"] + list(enabled_workers) + participant_count = len(participants) + if participant_count <= 0: + return {} + + slots_by_participant = {} + for branch_slot in range(int(num_branches)): + participant_id = str(participants[branch_slot % participant_count]) + slots_by_participant.setdefault(participant_id, []).append(branch_slot) + return slots_by_participant + + def _fallback_for_missing_branch(self, input_value): + if isinstance(input_value, torch.Tensor): + tensor = input_value + if tensor.ndim == 3: + tensor = tensor.unsqueeze(0) + if tensor.ndim == 4: + fallback = torch.zeros_like(tensor) + if fallback.is_cuda: + fallback = fallback.cpu() + return ensure_contiguous(fallback) + return input_value + async def _send_branch_result_to_master(self, value, branch_idx, multi_job_id, master_url, worker_id): if not isinstance(value, torch.Tensor): raise ValueError("DistributedJoin currently supports torch.Tensor results for worker transfer.") @@ -204,6 +221,7 @@ class DistributedJoin: branch_count = max(2, min(branch_count, MAX_BRANCH_OUTPUTS)) enabled_workers = self._parse_enabled_workers(enabled_worker_ids) + slots_by_participant = self._participant_branch_slots(enabled_workers, delegate_mode, branch_count) expected_worker_ids = self._expected_worker_ids_for_branches(enabled_workers, delegate_mode, branch_count) expected_workers = set(expected_worker_ids) @@ -266,4 +284,12 @@ class DistributedJoin: async with prompt_server.distributed_jobs_lock: prompt_server.distributed_pending_jobs.pop(multi_job_id, None) + missing_workers = sorted(expected_workers - workers_done) + if missing_workers: + fallback = self._fallback_for_missing_branch(input) + for missing_worker_id in missing_workers: + for slot_idx in slots_by_participant.get(str(missing_worker_id), []): + if 0 <= slot_idx < MAX_BRANCH_OUTPUTS and outputs[slot_idx] is None: + outputs[slot_idx] = fallback + return tuple(outputs) diff --git a/tests/test_join.py b/tests/test_join.py index 4595cb4..fbb24df 100644 --- a/tests/test_join.py +++ b/tests/test_join.py @@ -199,6 +199,30 @@ class DistributedJoinTests(unittest.IsolatedAsyncioTestCase): self.assertIsNone(outputs[2]) self.assertIsNone(outputs[9]) + async def test_missing_worker_results_use_safe_fallback_for_expected_slots(self): + fake_server = _FakePromptServer() + join_module._get_prompt_server_instance = lambda: fake_server + + queue = asyncio.Queue() + # No worker payloads enqueued; master should timeout and fill expected worker slots. + fake_server.distributed_pending_jobs["join-job"] = queue + + master_tensor = torch.full((1, 4, 4, 3), 0.7, 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.assertIsInstance(outputs[1], torch.Tensor) + self.assertIsInstance(outputs[2], torch.Tensor) + self.assertEqual(float(outputs[1].abs().sum().item()), 0.0) + self.assertEqual(float(outputs[2].abs().sum().item()), 0.0) + if __name__ == "__main__": unittest.main()