fixes to formatting and benchmark eval (#224)

Author: vivianrwu

benchmarks/eval_accuracy.py

@@ -66,11 +66,8 @@ def eval_accuracy(request_outputs_dict, match_type):
   for output in request_outputs_dict:
     preds.append(output["generated_text"])
     targets.append(output["original_output"])
-  if match_type != "math":
-    preds, targets = postprocess_text(preds, targets)
 
   if match_type == "math":
-
     correct_ans = 0
     wrong_ans = 0
     for p, t in zip(preds, targets):
@@ -87,11 +84,11 @@ def eval_accuracy(request_outputs_dict, match_type):
     result["literal"] = correct_ans / total_ans if total_ans > 0 else 0.0
     result["gen_len"] = total_ans
     result["gen_num"] = total_ans
-  if match_type == "rouge":
+
+  else:
     metric = evaluate.load("rouge")
-    nltk.download("punkt")
-    preds = []
-    targets = []
+    nltk.download("punkt_tab")
+    preds, targets = postprocess_text(preds, targets)
     result = metric.compute(
         predictions=preds,
         references=targets,

jetstream/core/orchestrator.py

@@ -1384,8 +1384,8 @@ async def Decode(  # pylint: disable=invalid-overridden-method
       if ttft == 0:
         ttft = time.perf_counter() - request_start_time
         if ttft > 2.0:
-          logging.info(
-              datetime.now(),
+          logger.info(  # pylint: disable=logging-fstring-interpolation
+              f"{datetime.now()}: "
               f"Slow TTFT: {ttft:.2f}s,"
               f" stats={active_request.metadata.stats()},"
               f" prefill_qsize={self._driver.prefill_backlog_size()}",

jetstream/engine/engine_api.py

@@ -367,12 +367,14 @@ def insert(
       prefix: Prefix,
       decode_state: DecodeState,
       slot: int,
+      request_id: Optional[uuid.UUID] = None,
   ) -> DecodeState:
 
     decode_state = self._downstream_engine.insert(
         prefix=prefix,
         decode_state=decode_state,
         slot=slot,
+        request_id=request_id,
     )
     return decode_state
 
@@ -438,3 +440,17 @@ def mesh(self) -> jax.sharding.Mesh:
   @property
   def colocated_cpus(self) -> Union[list[CpuDevices], None]:
     return self._downstream_engine.colocated_cpus
+
+  @property
+  def use_chunked_prefill(self) -> bool:
+    return self._downstream_engine.use_chunked_prefill
+
+  @property
+  def chunk_size(self) -> bool:
+    """Maximum prefill length."""
+    return self._downstream_engine.prefill_chunk_size
+
+  @property
+  def prefill_chunk_size(self) -> int:
+    """Maximum prefill length."""
+    return self._downstream_engine.prefill_chunk_size