Handle missing join worker results with safe fallback outputs

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)

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