Request input/output size metrics (#123)

* first commit

* remove unused code

* fmt

* changed buckets

* now using DEFAULT_PREFILL_BUCKETS

* missing parenthese

Author: Bslabe123

jetstream/core/metrics/prometheus.py

@@ -16,7 +16,9 @@
 
 import os
 import shortuuid
-from prometheus_client import Counter, Gauge
+from prometheus_client import Counter, Gauge, Histogram
+
+from jetstream.engine.token_utils import DEFAULT_PREFILL_BUCKETS
 
 
 class JetstreamMetricsCollector:
@@ -55,6 +57,39 @@ def __new__(cls):
       documentation="Total time taken to start the Jetstream server",
       labelnames=["id"],
   )
+  _request_input_length = Histogram(
+      name="jetstream_request_input_length",
+      documentation="Number of input tokens per request",
+      labelnames=["id"],
+      buckets=DEFAULT_PREFILL_BUCKETS,
+  )
+  _request_output_length = Histogram(
+      name="jetstream_request_output_length",
+      documentation="Number of output tokens per request",
+      labelnames=["id"],
+      buckets=[
+          1,
+          2,
+          5,
+          10,
+          20,
+          50,
+          100,
+          200,
+          500,
+          1000,
+          2000,
+          5000,
+          10000,
+          20000,
+          50000,
+          100000,
+          200000,
+          500000,
+          1000000,
+          2000000,
+      ],
+  )
   _request_success_count = Counter(
       name="jetstream_request_success_count",
       documentation="Number of requests successfully completed",
@@ -76,5 +111,11 @@ def get_slots_used_percentage_metric(self, idx: int):
   def get_server_startup_latency_metric(self):
     return self._server_startup_latency.labels(id=self._id)
 
+  def get_request_input_length(self):
+    return self._request_input_length.labels(id=self._id)
+
+  def get_request_output_length(self):
+    return self._request_output_length.labels(id=self._id)
+
   def get_request_success_count_metric(self):
     return self._request_success_count.labels(id=self._id)

jetstream/core/orchestrator.py

@@ -109,15 +109,6 @@
 root.addHandler(handler)
 
 
-def delete_pytree(p):
-  def delete_leaf(leaf):
-    if isinstance(leaf, jax.Array):
-      leaf.delete()
-    del leaf
-
-  jax.tree_map(delete_leaf, p)
-
-
 @dataclasses.dataclass
 class ActiveRequest:
   """Current state of the driver."""
@@ -532,6 +523,8 @@ def _prefill_thread(self, idx: int):
           idx,
           my_transfer_backlog.qsize(),
       )
+      if self._metrics_collector:
+        self._metrics_collector.get_request_input_length().observe(true_length)
 
       del prefill_result
       del request
@@ -781,6 +774,9 @@ def _detokenize_thread(self, idx: int):
             request.enqueue_samples(results)
             if request.complete.all():
               if self._metrics_collector:
+                self._metrics_collector.get_request_output_length().observe(
+                    result_tokens.get_result_at_slot(slot).lengths
+                )
                 self._metrics_collector.get_request_success_count_metric().inc()
               request.return_channel.close()
               # Place the slot back on the free queue.