[Mooncake][Worker] Fail only the request whose async tensor read failed, at TP>1 too

mstar/api_server/data_worker.py

@@ -27,7 +27,11 @@
 )
 from mstar.communication.communicator import BaseCommunicator, CommProtocol, make_communicator
 from mstar.communication.event import EventWakeup
-from mstar.communication.tensors import NameToTensorList, create_tensor_communication_manager
+from mstar.communication.tensors import (
+    FailedRead,
+    NameToTensorList,
+    create_tensor_communication_manager,
+)
 from mstar.engine.resources.kv.config import KVSpec, PagedKVConfig
 from mstar.engine.resources.kv.keys import PageItem, chain, end_key, fingerprint, page_items
 from mstar.engine.resources.spec import apply_yaml_overrides
@@ -886,7 +890,7 @@ def _summarize_inputs(input: PreprocessInput) -> list[InputInfo]:
 
     def _fail_request(
         self, request_id: str, exc: BaseException, stage: str, count: int = 1,
-        sequence: int | None = None, abort: bool = True,
+        sequence: int | None = None, abort: bool = True, traceback: bool = True,
     ):
         """Report a per-request data-worker failure to the API server.
 
@@ -901,8 +905,16 @@ def _fail_request(
 
         ``abort`` (False for a request the conductor has not seen) stops the
         request's walk, since its client already has its answer.
+
+        ``traceback`` False logs one line, for an error already logged where
+        it happened.
         """
-        logger.exception("%s failed for request %s", stage, request_id)
+        if traceback:
+            logger.error("%s failed for request %s", stage, request_id, exc_info=exc)
+        else:
+            logger.error(
+                "%s failed for request %s: %s: %s", stage, request_id, type(exc).__name__, exc,
+            )
         status = 400 if isinstance(exc, (ValueError, TypeError)) else 500
         chunks = [
             ResultChunk(
@@ -1040,9 +1052,29 @@ def _discard_result_tensor(
             graph_edges=[result.graph_edge],
         )
 
+    def _fail_unread_outputs(self, failed: FailedRead) -> None:
+        """Fail the request of a read that will not be delivered: one error
+        chunk in the output slot of each tensor it would have delivered. The
+        manager already acked or kept its buffers."""
+        request_id = failed.request_id
+        state = self.request_output_state.get(request_id)
+        metadata = self.tensor_uuid_to_metadata_per_request.get(request_id, {})
+        for edge in failed.graph_edges:
+            for info in edge.tensor_info:
+                metadata.pop(info.uuid, None)
+                slot = state.order.pop(info.uuid, None) if state is not None else None
+                self._fail_request(
+                    request_id, failed.error, "output tensor read",
+                    sequence=slot[0] if slot is not None else None, traceback=False,
+                )
+
     def _process_read_tensors(self):
         did_work = False
-        for request_id, graph_edges in self.tensor_manager.get_ready_tensors().items():
+        ready = self.tensor_manager.get_ready_tensors()
+        for failed in self.tensor_manager.take_failed_reads():
+            did_work = True
+            self._fail_unread_outputs(failed)
+        for request_id, graph_edges in ready.items():
             did_work = True
             for graph_edge in graph_edges:
                 modality = graph_edge.name.replace("_output", "")
@@ -1179,12 +1211,13 @@ def _send_enable_metrics_if_requested(self) -> None:
     def _finish_reading(self, request_id: str, drained: bool) -> None:
         """The API server is done with this rid's outputs. ``drained`` (every
         chunk delivered) means no read can still be in flight, so ACK straight
-        away — the happy path pays nothing. Otherwise delivery was abandoned
-        (TTL, client gone) with reads possibly still running, so gate the ACK on
+        away, unless a failed read's transfers are still running: its error
+        chunks count as delivered. Otherwise delivery was abandoned (TTL,
+        client gone) with reads possibly still running, so gate the ACK on
         them; ACKing there would let the conductor unlink under a read."""
         if request_id not in self._live_requests:
             return
-        if drained:
+        if drained and not self.tensor_manager.has_inflight_reads(request_id):
             self._send_reads_done(request_id)
         else:
             self._begin_drain(request_id)