Recompute when the deadlock is detected

mii/batching/ragged_batching.py

@@ -133,7 +133,7 @@ def from_msg(msg: Dict[str, int]) -> Self:
 class RaggedRequest:
     uid: int
     input_tokens: torch.Tensor
-    prompt_length: int
+    prompt_tokens: torch.Tensor
     seq_length: int
     max_length: int
     max_new_tokens: int
@@ -147,6 +147,10 @@ class RaggedRequest:
     _generated_tokens: List[torch.Tensor] = field(default_factory=list)
     _finish_reason: GenerationFinishReason = GenerationFinishReason.NONE
 
+    @property
+    def prompt_length(self) -> int:
+        return len(self.prompt_tokens)
+
     @property
     def next_token(self) -> Union[None, torch.Tensor]:
         return self._next_token
@@ -204,6 +208,9 @@ def accumulate_generated_token(self) -> None:
         if not self.is_done:
             self._generated_tokens.append(self.next_token)
 
+    def clear_generated_token(self) -> None:
+        self._generated_tokens.clear()
+
     def set_next_as_input(self) -> None:
         if self.next_token is not None:
             self.input_tokens = self.next_token.unsqueeze(0)
@@ -551,11 +558,30 @@ def schedule_requests(self) -> None:
         self._do_schedule_requests(prompt_reqs)
 
         if len(self.buffer) > 0 and len(self.scheduled_requests) == 0:
-            raise RuntimeError("Deadlock detected: No requests were scheduled.")
+            print(
+                "Deadlock detected. Resetting KV cache and recomputing requests. Consider limiting number of concurrent requests or decreasing max lengths of prompts/generations."
+            )
+            self.scheduled_requests = RaggedRequestBatch([])
+            self.reset_request_status()
+        else:
+            scheduled_requests_ids = set(id(r) for r in self.scheduled_requests)
+            self.buffer = deque(
+                [r for r in self.buffer if id(r) not in scheduled_requests_ids])
+
+    def reset_request_status(self):
+        self.flush([r.uid for r in self.buffer if r.seq_length > 0])
+
+        new_buffer = deque()
+        for r in self.buffer:
+            new_req = copy.copy(r)
+            new_req.prompt_tokens = new_req.input_tokens = torch.concat(
+                [r.prompt_tokens] + [t.unsqueeze(0) for t in r.generated_tokens])
+            new_req.seq_length = 0
+            new_req.max_new_tokens = r.max_new_tokens - len(r.generated_tokens)
+            new_req.clear_generated_token()
+            new_buffer.append(new_req)
 
-        scheduled_requests_ids = set(id(r) for r in self.scheduled_requests)
-        self.buffer = deque(
-            [r for r in self.buffer if id(r) not in scheduled_requests_ids])
+        self.buffer = new_buffer
 
     def make_request(self,
                      uid: int,
@@ -609,7 +635,7 @@ def make_request(self,
             RaggedRequest(
                 uid=uid,
                 input_tokens=input_tokens,
-                prompt_length=prompt_length,
+                prompt_tokens=input_tokens,
                 seq_length=0,
                 max_length=max_length,
                 max_new_tokens=max_new_tokens,
@@ -660,7 +686,7 @@ def __call__(self, inputs: Union[str, List[str]], **kwargs) -> ResponseBatch:
                         RaggedRequest(
                             uid=uid,
                             input_tokens=None,
-                            prompt_length=None,
+                            prompt_tokens=None,
                             seq_length=None,
                             max_length=None,
                             max_new_tokens=None,
@@ -823,7 +849,7 @@ def destroy_session(self,
                     RaggedRequest(
                         uid=uid,
                         input_tokens=None,
-                        prompt_length=None,
+                        prompt_tokens=None,
                         seq_length=None,
                         max_length=None,
                         max_new_tokens=None,