Add lock in prefill and generate to prevent starvation (#126)

add lock for prefill and generate to prevent starvation

Author: FanhaiLu1

jetstream_pt/ray_engine.py

@@ -1,4 +1,5 @@
 from collections import defaultdict
+import threading
 from typing import Any, Iterable, Optional, Union, Tuple, List
 
 import numpy as np
@@ -38,6 +39,8 @@ def __init__(
     self.batch_size = batch_size
     self.is_disaggregated = is_disaggregated
     self.pod_slice_name = pod_slice_name
+    if not self.is_disaggregated:
+      self._lock = threading.Lock()
 
   # pylint: disable-next=all
   def load_params(self) -> Params:
@@ -66,6 +69,31 @@ def prefill(
       existing_prefix: Optional[Prefix] = None,
       padded_tokens: np.ndarray,  # PrefillInputs[np.ndarray],
       true_length: int,
+  ) -> Prefix:
+    if self.is_disaggregated:
+      return self.prefill_impl(
+          params=params,
+          existing_prefix=existing_prefix,
+          padded_tokens=padded_tokens,
+          true_length=true_length,
+      )
+
+    with self._lock:
+      return self.prefill_impl(
+          params=params,
+          existing_prefix=existing_prefix,
+          padded_tokens=padded_tokens,
+          true_length=true_length,
+      )
+
+  # pylint: disable-next=all
+  def prefill_impl(
+      self,
+      *,
+      params: Any,  # Weights
+      existing_prefix: Optional[Prefix] = None,
+      padded_tokens: np.ndarray,  # PrefillInputs[np.ndarray],
+      true_length: int,
   ) -> Prefix:
     all_outputs = []
     for worker in self.engine_workers:
@@ -116,6 +144,15 @@ def insert(
 
   def generate(
       self, params: Any, decode_state: DecodeState
+  ) -> tuple[None, engine_api.ResultTokens]:
+    if self.is_disaggregated:
+      return self.generate_impl(params=params, decode_state=decode_state)
+    with self._lock:
+      return self.generate_impl(params=params, decode_state=decode_state)
+
+  # pylint: disable-next=all
+  def generate_impl(
+      self, params: Any, decode_state: DecodeState
   ) -> tuple[None, engine_api.ResultTokens]:
     all_outputs = []
     for worker in self.engine_workers: