Ray lint

Author: qihqi

jetstream_pt/engine.py

@@ -614,7 +614,7 @@ def false_comp(b, i, bk, start, end):
     return b_next, i_next
 
   def generate(
-      self, params: Any, decode_state: DecodeState, sampler = None
+      self, params: Any, decode_state: DecodeState, sampler=None
   ) -> tuple[DecodeState, engine_api.ResultTokens]:
     # seq_len = padded_tokens.shape[0]
     pos = decode_state.current_position

jetstream_pt/ray_engine.py

@@ -70,6 +70,7 @@ def prefill(
       existing_prefix: Optional[Prefix] = None,
       padded_tokens: np.ndarray,  # PrefillInputs[np.ndarray],
       true_length: int,
+      sampler=None,
   ) -> Tuple[Prefix, engine_api.ResultTokens]:
     if self.is_disaggregated:
       return self.prefill_impl(
@@ -144,7 +145,10 @@ def insert(
     return None
 
   def generate(
-      self, params: Any, decode_state: DecodeState
+      self,
+      params: Any,
+      decode_state: DecodeState,
+      sampler=None,
   ) -> tuple[None, engine_api.ResultTokens]:
     if self.is_disaggregated:
       return self.generate_impl(params=params, decode_state=decode_state)