From 2814db709c010aa7f97b9ae3d828ab649d368f59 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Mon, 13 May 2024 22:15:13 +0000
Subject: [PATCH 01/74] initial-commit

---
 jetstream/core/orchestrator.py | 18 ++++++++++--------
 requirements.in                |  1 +
 requirements.txt               |  2 ++
 3 files changed, 13 insertions(+), 8 deletions(-)

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index 34b45d8e..d87253fe 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -96,6 +96,7 @@
 
 import numpy as np
 import prometheus_client
+import shortuuid
 
 root = logging.getLogger()
 root.setLevel(logging.DEBUG)
@@ -210,9 +211,6 @@ class Driver:
   # todo: remove jax_padding after all then engine migrate to np padding
   _jax_padding = True
 
-  # Record metrics for prefill_backlog size
-  _prefill_backlog_size_metric: prometheus_client.Gauge
-
   def __init__(
       self,
       prefill_engines: Optional[list[engine_api.Engine]] = None,
@@ -246,9 +244,10 @@ def __init__(
     # Stage 1
     # At first, a request is placed here in order to get prefilled.
     self._prefill_backlog = queue.Queue()
-    self._prefill_backlog_size_metric = prometheus_client.Gauge(
-        "jetstream_prefill_backlog_size", "Size of prefill queue"
-    )
+    prometheus_client.Gauge(
+        "jetstream_prefill_backlog_size-{shortuuid.uuid()}",
+        "Size of prefill queue",
+    ).set_function(lambda: self._prefill_backlog.qsize())
 
     # Stage 2
     # After prefilling, it is placed here in order to get transferred to
@@ -429,7 +428,6 @@ def place_request_on_prefill_queue(self, request: ActiveRequest):
     """Used to place new requests for prefilling and generation."""
     # Don't block so we can fail and shed load when the queue is full.
     self._prefill_backlog.put(request, block=False)
-    self._prefill_backlog_size_metric.set(self._prefill_backlog.qsize())
 
   def _load_cache_history(self, path: str) -> Union[None, Any]:
     """Loads previous kv cache for a longer conversation."""
@@ -451,7 +449,6 @@ def _prefill_thread(self, idx: int):
       my_transfer_backlog = self._transfer_backlogs[idx]
       # The prefill thread can just sleep until it has work to do.
       request = self._prefill_backlog.get(block=True)
-      self._prefill_backlog_size_metric.set(self._prefill_backlog.qsize())
 
       if request is None:
         break
@@ -532,6 +529,11 @@ def _generate_thread(self, idx: int):
     logging.info("---------Spinning up generate thread %d.---------", idx)
     generate_engine = self._generate_engines[idx]
     my_slots = self._generate_slots[idx]
+    prometheus_client.Gauge(
+        f"jetstream_slots_available_percentage-{shortuuid.uuid()}-generate-{idx}",
+        "The percentage of available slots in decode batch",
+    ).set_function(lambda: my_slots.qsize() / max_concurrent_decodes)
+
     my_generate_backlog = self._generate_backlogs[idx]
     my_detokenize_backlog = self._detokenize_backlogs[idx]
 
diff --git a/requirements.in b/requirements.in
index 6f4ebb60..e2e8e984 100644
--- a/requirements.in
+++ b/requirements.in
@@ -11,3 +11,4 @@ pytest
 seqio
 tiktoken
 blobfile
+shortuuid
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index f2a23c7b..c7bf9e2b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -236,6 +236,8 @@ sentencepiece==0.1.99
     # via seqio
 seqio==0.0.19
     # via -r requirements.in
+shortuuid==1.0.13
+    # via -r requirements.in
 six==1.16.0
     # via
     #   astunparse

From 1fa15d174ba01ad6509c124b089b99aa5403086c Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Mon, 13 May 2024 22:18:48 +0000
Subject: [PATCH 02/74] newline

---
 requirements.in | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/requirements.in b/requirements.in
index e2e8e984..4af40e9a 100644
--- a/requirements.in
+++ b/requirements.in
@@ -11,4 +11,4 @@ pytest
 seqio
 tiktoken
 blobfile
-shortuuid
\ No newline at end of file
+shortuuid

From 2c159725dffdbe21ae24a5f5af05533f6c31c256 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Mon, 13 May 2024 23:01:46 +0000
Subject: [PATCH 03/74] moved gauge

---
 jetstream/core/orchestrator.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index d87253fe..98e2600a 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -529,10 +529,6 @@ def _generate_thread(self, idx: int):
     logging.info("---------Spinning up generate thread %d.---------", idx)
     generate_engine = self._generate_engines[idx]
     my_slots = self._generate_slots[idx]
-    prometheus_client.Gauge(
-        f"jetstream_slots_available_percentage-{shortuuid.uuid()}-generate-{idx}",
-        "The percentage of available slots in decode batch",
-    ).set_function(lambda: my_slots.qsize() / max_concurrent_decodes)
 
     my_generate_backlog = self._generate_backlogs[idx]
     my_detokenize_backlog = self._detokenize_backlogs[idx]
@@ -559,6 +555,11 @@ def _generate_thread(self, idx: int):
 
       max_concurrent_decodes = generate_engine.max_concurrent_decodes
 
+      prometheus_client.Gauge(
+          f"jetstream_slots_available_percentage-{shortuuid.uuid()}-generate-{idx}",
+          "The percentage of available slots in decode batch",
+      ).set_function(lambda: my_slots.qsize() / max_concurrent_decodes)
+
       # Check if there are any free my_slots. We don't want to block here since
       # we can still generate if we can't insert. We do this in a while loop to
       # insert as many sequences as possible.

From dd5932ffd11b726d1f5a7dd0aa04638076cc7d3b Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Tue, 14 May 2024 16:35:13 +0000
Subject: [PATCH 04/74] Added labels for metrics

---
 jetstream/core/orchestrator.py | 20 ++++++++++++++------
 1 file changed, 14 insertions(+), 6 deletions(-)

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index 98e2600a..c0847c37 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -245,9 +245,12 @@ def __init__(
     # At first, a request is placed here in order to get prefilled.
     self._prefill_backlog = queue.Queue()
     prometheus_client.Gauge(
-        "jetstream_prefill_backlog_size-{shortuuid.uuid()}",
+        "jetstream_prefill_backlog_size",
         "Size of prefill queue",
-    ).set_function(lambda: self._prefill_backlog.qsize())
+        labelnames=["uuid"]
+    )
+    self._prefill_backlog.labels("uuid").set(shortuuid.uuid())
+    self._prefill_backlog.set_function(lambda: self._prefill_backlog.qsize())
 
     # Stage 2
     # After prefilling, it is placed here in order to get transferred to
@@ -555,10 +558,15 @@ def _generate_thread(self, idx: int):
 
       max_concurrent_decodes = generate_engine.max_concurrent_decodes
 
-      prometheus_client.Gauge(
-          f"jetstream_slots_available_percentage-{shortuuid.uuid()}-generate-{idx}",
-          "The percentage of available slots in decode batch",
-      ).set_function(lambda: my_slots.qsize() / max_concurrent_decodes)
+      jetstream_slots_available_percentage_metric = prometheus_client.Gauge(
+          name="jetstream_slots_available_percentage",
+          documentation="The percentage of available slots in decode batch",
+          labelnames=["uuid", idx],
+      )
+      jetstream_slots_available_percentage_metric.labels("uuid").set(shortuuid.uuid())
+      jetstream_slots_available_percentage_metric.labels("idx").set(idx)
+      jetstream_slots_available_percentage_metric.set_function(lambda: my_slots.qsize() / max_concurrent_decodes)
+      
 
       # Check if there are any free my_slots. We don't want to block here since
       # we can still generate if we can't insert. We do this in a while loop to

From 471bb24c47d8ac9296a9c18c0e22b5b59d41e8e6 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Tue, 14 May 2024 16:46:31 +0000
Subject: [PATCH 05/74] revert removing driver metric field

---
 jetstream/core/orchestrator.py | 22 +++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index c0847c37..ecc7f388 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -211,6 +211,9 @@ class Driver:
   # todo: remove jax_padding after all then engine migrate to np padding
   _jax_padding = True
 
+  # Record metrics for prefill_backlog size
+  _prefill_backlog_size_metric: prometheus_client.Gauge
+
   def __init__(
       self,
       prefill_engines: Optional[list[engine_api.Engine]] = None,
@@ -244,13 +247,15 @@ def __init__(
     # Stage 1
     # At first, a request is placed here in order to get prefilled.
     self._prefill_backlog = queue.Queue()
-    prometheus_client.Gauge(
+    self._prefill_backlog_metric = prometheus_client.Gauge(
         "jetstream_prefill_backlog_size",
         "Size of prefill queue",
-        labelnames=["uuid"]
+        labelnames=["uuid"],
+    )
+    self._prefill_backlog_metric.labels("uuid").set(shortuuid.uuid())
+    self._prefill_backlog_metric.set_function(
+        lambda: self._prefill_backlog.qsize()
     )
-    self._prefill_backlog.labels("uuid").set(shortuuid.uuid())
-    self._prefill_backlog.set_function(lambda: self._prefill_backlog.qsize())
 
     # Stage 2
     # After prefilling, it is placed here in order to get transferred to
@@ -563,10 +568,13 @@ def _generate_thread(self, idx: int):
           documentation="The percentage of available slots in decode batch",
           labelnames=["uuid", idx],
       )
-      jetstream_slots_available_percentage_metric.labels("uuid").set(shortuuid.uuid())
+      jetstream_slots_available_percentage_metric.labels("uuid").set(
+          shortuuid.uuid()
+      )
       jetstream_slots_available_percentage_metric.labels("idx").set(idx)
-      jetstream_slots_available_percentage_metric.set_function(lambda: my_slots.qsize() / max_concurrent_decodes)
-      
+      jetstream_slots_available_percentage_metric.set_function(
+          lambda: my_slots.qsize() / max_concurrent_decodes
+      )
 
       # Check if there are any free my_slots. We don't want to block here since
       # we can still generate if we can't insert. We do this in a while loop to

From bed3c6ce87c6386f3a330e598f55f94935a8989c Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Tue, 14 May 2024 16:47:54 +0000
Subject: [PATCH 06/74] removed newline

---
 jetstream/core/orchestrator.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index ecc7f388..f579c19b 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -537,7 +537,6 @@ def _generate_thread(self, idx: int):
     logging.info("---------Spinning up generate thread %d.---------", idx)
     generate_engine = self._generate_engines[idx]
     my_slots = self._generate_slots[idx]
-
     my_generate_backlog = self._generate_backlogs[idx]
     my_detokenize_backlog = self._detokenize_backlogs[idx]
 

From 85c8e146798b3897d4cbaeac087527cc79044913 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Tue, 14 May 2024 16:52:41 +0000
Subject: [PATCH 07/74] Removed unneccesary lambda

---
 jetstream/core/orchestrator.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index f579c19b..e622943c 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -253,9 +253,7 @@ def __init__(
         labelnames=["uuid"],
     )
     self._prefill_backlog_metric.labels("uuid").set(shortuuid.uuid())
-    self._prefill_backlog_metric.set_function(
-        lambda: self._prefill_backlog.qsize()
-    )
+    self._prefill_backlog_metric.set_function(self._prefill_backlog.qsize())
 
     # Stage 2
     # After prefilling, it is placed here in order to get transferred to

From 4b40a31efb0db05622950fb0db4f522312baf8dc Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Tue, 14 May 2024 19:50:54 +0000
Subject: [PATCH 08/74] Proper use of labels api

---
 jetstream/core/orchestrator.py | 11 ++---------
 1 file changed, 2 insertions(+), 9 deletions(-)

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index e622943c..948fc9dc 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -252,8 +252,7 @@ def __init__(
         "Size of prefill queue",
         labelnames=["uuid"],
     )
-    self._prefill_backlog_metric.labels("uuid").set(shortuuid.uuid())
-    self._prefill_backlog_metric.set_function(self._prefill_backlog.qsize())
+    self._prefill_backlog_metric.labels(shortuuid.uuid()).set_function(self._prefill_backlog.qsize())
 
     # Stage 2
     # After prefilling, it is placed here in order to get transferred to
@@ -565,13 +564,7 @@ def _generate_thread(self, idx: int):
           documentation="The percentage of available slots in decode batch",
           labelnames=["uuid", idx],
       )
-      jetstream_slots_available_percentage_metric.labels("uuid").set(
-          shortuuid.uuid()
-      )
-      jetstream_slots_available_percentage_metric.labels("idx").set(idx)
-      jetstream_slots_available_percentage_metric.set_function(
-          lambda: my_slots.qsize() / max_concurrent_decodes
-      )
+      jetstream_slots_available_percentage_metric.labels(shortuuid.uuid(), idx).set_function(my_slots.qsize() / max_concurrent_decodes)
 
       # Check if there are any free my_slots. We don't want to block here since
       # we can still generate if we can't insert. We do this in a while loop to

From f1b538f0eb49d991d673bfc5b107020f6c8facc5 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Tue, 14 May 2024 19:56:21 +0000
Subject: [PATCH 09/74] format

---
 jetstream/core/orchestrator.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index 948fc9dc..41d2cc02 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -252,7 +252,9 @@ def __init__(
         "Size of prefill queue",
         labelnames=["uuid"],
     )
-    self._prefill_backlog_metric.labels(shortuuid.uuid()).set_function(self._prefill_backlog.qsize())
+    self._prefill_backlog_metric.labels(shortuuid.uuid()).set_function(
+        self._prefill_backlog.qsize()
+    )
 
     # Stage 2
     # After prefilling, it is placed here in order to get transferred to
@@ -564,7 +566,9 @@ def _generate_thread(self, idx: int):
           documentation="The percentage of available slots in decode batch",
           labelnames=["uuid", idx],
       )
-      jetstream_slots_available_percentage_metric.labels(shortuuid.uuid(), idx).set_function(my_slots.qsize() / max_concurrent_decodes)
+      jetstream_slots_available_percentage_metric.labels(
+          shortuuid.uuid(), idx
+      ).set_function(my_slots.qsize() / max_concurrent_decodes)
 
       # Check if there are any free my_slots. We don't want to block here since
       # we can still generate if we can't insert. We do this in a while loop to

From 7eb34e3dbe86a3f4a644aa53315c2e8db7d9cd68 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Tue, 14 May 2024 20:10:42 +0000
Subject: [PATCH 10/74] missing quotes

---
 jetstream/core/orchestrator.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index 41d2cc02..e3f3f320 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -564,7 +564,7 @@ def _generate_thread(self, idx: int):
       jetstream_slots_available_percentage_metric = prometheus_client.Gauge(
           name="jetstream_slots_available_percentage",
           documentation="The percentage of available slots in decode batch",
-          labelnames=["uuid", idx],
+          labelnames=["uuid", "idx"],
       )
       jetstream_slots_available_percentage_metric.labels(
           shortuuid.uuid(), idx

From 700fc6abf9b9f38507bb57461c347fd31d577efa Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Tue, 14 May 2024 21:07:48 +0000
Subject: [PATCH 11/74] convert result to float, add lambda

---
 jetstream/core/orchestrator.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index e3f3f320..a0e49a6f 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -253,7 +253,7 @@ def __init__(
         labelnames=["uuid"],
     )
     self._prefill_backlog_metric.labels(shortuuid.uuid()).set_function(
-        self._prefill_backlog.qsize()
+        lambda: float(self._prefill_backlog.qsize())
     )
 
     # Stage 2
@@ -568,7 +568,7 @@ def _generate_thread(self, idx: int):
       )
       jetstream_slots_available_percentage_metric.labels(
           shortuuid.uuid(), idx
-      ).set_function(my_slots.qsize() / max_concurrent_decodes)
+      ).set_function(lambda: float(my_slots.qsize() / max_concurrent_decodes))
 
       # Check if there are any free my_slots. We don't want to block here since
       # we can still generate if we can't insert. We do this in a while loop to

From 0cc9bdb8c98928306d29b78bf880d00466509ac1 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Tue, 14 May 2024 21:35:49 +0000
Subject: [PATCH 12/74] Moved registration to driver level

---
 jetstream/core/orchestrator.py | 25 +++++++++++++------------
 1 file changed, 13 insertions(+), 12 deletions(-)

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index a0e49a6f..2ce6f360 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -243,16 +243,23 @@ def __init__(
     self._generate_params = generate_params
     self._interleaved_mode = interleaved_mode
 
-    # Stages 1-4 represent the life cycle of a request.
-    # Stage 1
-    # At first, a request is placed here in order to get prefilled.
-    self._prefill_backlog = queue.Queue()
-    self._prefill_backlog_metric = prometheus_client.Gauge(
+    # Register metric at the driver level to avoid duplicate registration
+    prefill_backlog_metric = prometheus_client.Gauge(
         "jetstream_prefill_backlog_size",
         "Size of prefill queue",
         labelnames=["uuid"],
     )
-    self._prefill_backlog_metric.labels(shortuuid.uuid()).set_function(
+    jetstream_slots_available_percentage_metric = prometheus_client.Gauge(
+        name="jetstream_slots_available_percentage",
+        documentation="The percentage of available slots in decode batch",
+        labelnames=["uuid", "idx"],
+    )
+
+    # Stages 1-4 represent the life cycle of a request.
+    # Stage 1
+    # At first, a request is placed here in order to get prefilled.
+    self._prefill_backlog = queue.Queue()
+    prefill_backlog_metric.labels(shortuuid.uuid()).set_function(
         lambda: float(self._prefill_backlog.qsize())
     )
 
@@ -560,12 +567,6 @@ def _generate_thread(self, idx: int):
         time_of_last_print = time.time()
 
       max_concurrent_decodes = generate_engine.max_concurrent_decodes
-
-      jetstream_slots_available_percentage_metric = prometheus_client.Gauge(
-          name="jetstream_slots_available_percentage",
-          documentation="The percentage of available slots in decode batch",
-          labelnames=["uuid", "idx"],
-      )
       jetstream_slots_available_percentage_metric.labels(
           shortuuid.uuid(), idx
       ).set_function(lambda: float(my_slots.qsize() / max_concurrent_decodes))

From 536a78e194bd28a5d4736ff626303a680598f7c3 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Tue, 14 May 2024 21:41:43 +0000
Subject: [PATCH 13/74] made metric driver property

---
 jetstream/core/orchestrator.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index 2ce6f360..69f1a32f 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -244,12 +244,12 @@ def __init__(
     self._interleaved_mode = interleaved_mode
 
     # Register metric at the driver level to avoid duplicate registration
-    prefill_backlog_metric = prometheus_client.Gauge(
+    self._prefill_backlog_metric = prometheus_client.Gauge(
         "jetstream_prefill_backlog_size",
         "Size of prefill queue",
         labelnames=["uuid"],
     )
-    jetstream_slots_available_percentage_metric = prometheus_client.Gauge(
+    self._jetstream_slots_available_percentage_metric = prometheus_client.Gauge(
         name="jetstream_slots_available_percentage",
         documentation="The percentage of available slots in decode batch",
         labelnames=["uuid", "idx"],
@@ -259,7 +259,7 @@ def __init__(
     # Stage 1
     # At first, a request is placed here in order to get prefilled.
     self._prefill_backlog = queue.Queue()
-    prefill_backlog_metric.labels(shortuuid.uuid()).set_function(
+    self._prefill_backlog_metric.labels(shortuuid.uuid()).set_function(
         lambda: float(self._prefill_backlog.qsize())
     )
 
@@ -567,7 +567,7 @@ def _generate_thread(self, idx: int):
         time_of_last_print = time.time()
 
       max_concurrent_decodes = generate_engine.max_concurrent_decodes
-      jetstream_slots_available_percentage_metric.labels(
+      self._jetstream_slots_available_percentage_metric.labels(
           shortuuid.uuid(), idx
       ).set_function(lambda: float(my_slots.qsize() / max_concurrent_decodes))
 

From 97c27b29ff0587788a404f36ebcdf1a4f18380da Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Tue, 14 May 2024 21:56:24 +0000
Subject: [PATCH 14/74] missing field on class

---
 jetstream/core/orchestrator.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index 69f1a32f..052977d4 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -211,8 +211,9 @@ class Driver:
   # todo: remove jax_padding after all then engine migrate to np padding
   _jax_padding = True
 
-  # Record metrics for prefill_backlog size
+  # Record metrics
   _prefill_backlog_size_metric: prometheus_client.Gauge
+  _jetstream_slots_available_percentage_metric: prometheus_client.Gauge
 
   def __init__(
       self,

From b6cfdf1fcdd03de5b89a35cb67843b26feb7bfe1 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Tue, 14 May 2024 22:24:34 +0000
Subject: [PATCH 15/74] Dont regenerate uuid on each scrape

---
 jetstream/core/orchestrator.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index 052977d4..16704c82 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -212,6 +212,7 @@ class Driver:
   _jax_padding = True
 
   # Record metrics
+  _driver_uuid: str
   _prefill_backlog_size_metric: prometheus_client.Gauge
   _jetstream_slots_available_percentage_metric: prometheus_client.Gauge
 
@@ -260,7 +261,7 @@ def __init__(
     # Stage 1
     # At first, a request is placed here in order to get prefilled.
     self._prefill_backlog = queue.Queue()
-    self._prefill_backlog_metric.labels(shortuuid.uuid()).set_function(
+    self._prefill_backlog_metric.labels(self._driver_uuid).set_function(
         lambda: float(self._prefill_backlog.qsize())
     )
 
@@ -569,7 +570,7 @@ def _generate_thread(self, idx: int):
 
       max_concurrent_decodes = generate_engine.max_concurrent_decodes
       self._jetstream_slots_available_percentage_metric.labels(
-          shortuuid.uuid(), idx
+          self._driver_uuid, idx
       ).set_function(lambda: float(my_slots.qsize() / max_concurrent_decodes))
 
       # Check if there are any free my_slots. We don't want to block here since

From 548cec3827d2d06727fb778a9dde5f1bf1bb5e43 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Tue, 14 May 2024 22:29:07 +0000
Subject: [PATCH 16/74] Init uuid

---
 jetstream/core/orchestrator.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index 16704c82..a2fce737 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -246,6 +246,7 @@ def __init__(
     self._interleaved_mode = interleaved_mode
 
     # Register metric at the driver level to avoid duplicate registration
+    self._driver_uuid = shortuuid.uuid()
     self._prefill_backlog_metric = prometheus_client.Gauge(
         "jetstream_prefill_backlog_size",
         "Size of prefill queue",

From 5554c37883c27567270227e0becfcd981792e9bc Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Wed, 15 May 2024 17:04:31 +0000
Subject: [PATCH 17/74] Moved metrics to separate file

---
 jetstream/core/orchestrator.py | 27 ++++----------------------
 jetstream/metrics/__init__.py  | 35 ++++++++++++++++++++++++++++++++++
 2 files changed, 39 insertions(+), 23 deletions(-)
 create mode 100644 jetstream/metrics/__init__.py

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index 09605bdb..6e752055 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -94,9 +94,8 @@
 from jetstream.core.utils import async_multifuture
 from jetstream.core.utils.return_sample import ReturnSample
 from jetstream.engine import engine_api, tokenizer_api, token_utils
+import jetstream.metrics as jetstream_metrics
 import numpy as np
-import prometheus_client
-import shortuuid
 
 root = logging.getLogger()
 root.setLevel(logging.DEBUG)
@@ -212,11 +211,6 @@ class Driver:
   # todo: remove jax_padding after all then engine migrate to np padding
   _jax_padding = True
 
-  # Record metrics
-  _driver_uuid: str
-  _prefill_backlog_size_metric: prometheus_client.Gauge
-  _jetstream_slots_available_percentage_metric: prometheus_client.Gauge
-
   def __init__(
       self,
       prefill_engines: Optional[list[engine_api.Engine]] = None,
@@ -246,24 +240,11 @@ def __init__(
     self._generate_params = generate_params
     self._interleaved_mode = interleaved_mode
 
-    # Register metric at the driver level to avoid duplicate registration
-    self._driver_uuid = shortuuid.uuid()
-    self._prefill_backlog_metric = prometheus_client.Gauge(
-        "jetstream_prefill_backlog_size",
-        "Size of prefill queue",
-        labelnames=["uuid"],
-    )
-    self._jetstream_slots_available_percentage_metric = prometheus_client.Gauge(
-        name="jetstream_slots_available_percentage",
-        documentation="The percentage of available slots in decode batch",
-        labelnames=["uuid", "idx"],
-    )
-
     # Stages 1-4 represent the life cycle of a request.
     # Stage 1
     # At first, a request is placed here in order to get prefilled.
     self._prefill_backlog = queue.Queue()
-    self._prefill_backlog_metric.labels(self._driver_uuid).set_function(
+    jetstream_metrics.prefill_backlog.labels(jetstream_metrics.instance_uuid).set_function(
         lambda: float(self._prefill_backlog.qsize())
     )
 
@@ -597,8 +578,8 @@ def _generate_thread(self, idx: int):
         time_of_last_print = time.time()
 
       max_concurrent_decodes = generate_engine.max_concurrent_decodes
-      self._jetstream_slots_available_percentage_metric.labels(
-          self._driver_uuid, idx
+      jetstream_metrics.jetstream_slots_available_percentage.labels(
+          jetstream_metrics.instance_uuid, idx
       ).set_function(lambda: float(my_slots.qsize() / max_concurrent_decodes))
 
       # Check if there are any free my_slots. We don't want to block here since
diff --git a/jetstream/metrics/__init__.py b/jetstream/metrics/__init__.py
new file mode 100644
index 00000000..1cdcaa16
--- /dev/null
+++ b/jetstream/metrics/__init__.py
@@ -0,0 +1,35 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import prometheus_client
+import shortuuid
+
+instance_uuid = shortuuid.uuid()
+
+# Metrics we care to observe
+prefill_backlog : prometheus_client.Gague
+jetstream_slots_available_percentage: prometheus_client.Gauge
+
+def register_metrics(): 
+    prefill_backlog_metric = prometheus_client.Gauge(
+        "jetstream_prefill_backlog_size",
+        "Size of prefill queue",
+        labelnames=["uuid"],
+    )
+
+    jetstream_slots_available_percentage_metric = prometheus_client.Gauge(
+        name="jetstream_slots_available_percentage",
+        documentation="The percentage of available slots in decode batch",
+        labelnames=["uuid", "idx"],
+    )
\ No newline at end of file

From e627b128f0af4bec5a3d2a5ba4bc360dc83d5d2d Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Wed, 15 May 2024 17:13:05 +0000
Subject: [PATCH 18/74] Typos

---
 jetstream/core/orchestrator.py |  6 +++---
 jetstream/metrics/__init__.py  | 25 +++++++++++++------------
 2 files changed, 16 insertions(+), 15 deletions(-)

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index 6e752055..9f954c15 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -244,9 +244,9 @@ def __init__(
     # Stage 1
     # At first, a request is placed here in order to get prefilled.
     self._prefill_backlog = queue.Queue()
-    jetstream_metrics.prefill_backlog.labels(jetstream_metrics.instance_uuid).set_function(
-        lambda: float(self._prefill_backlog.qsize())
-    )
+    jetstream_metrics.prefill_backlog.labels(
+        jetstream_metrics.instance_uuid
+    ).set_function(lambda: float(self._prefill_backlog.qsize()))
 
     # Stage 2
     # After prefilling, it is placed here in order to get transferred to
diff --git a/jetstream/metrics/__init__.py b/jetstream/metrics/__init__.py
index 1cdcaa16..292b7fe5 100644
--- a/jetstream/metrics/__init__.py
+++ b/jetstream/metrics/__init__.py
@@ -18,18 +18,19 @@
 instance_uuid = shortuuid.uuid()
 
 # Metrics we care to observe
-prefill_backlog : prometheus_client.Gague
+prefill_backlog: prometheus_client.Gauge
 jetstream_slots_available_percentage: prometheus_client.Gauge
 
-def register_metrics(): 
-    prefill_backlog_metric = prometheus_client.Gauge(
-        "jetstream_prefill_backlog_size",
-        "Size of prefill queue",
-        labelnames=["uuid"],
-    )
 
-    jetstream_slots_available_percentage_metric = prometheus_client.Gauge(
-        name="jetstream_slots_available_percentage",
-        documentation="The percentage of available slots in decode batch",
-        labelnames=["uuid", "idx"],
-    )
\ No newline at end of file
+def register_metrics():
+  prefill_backlog_metric = prometheus_client.Gauge(
+      "jetstream_prefill_backlog_size",
+      "Size of prefill queue",
+      labelnames=["uuid"],
+  )
+
+  jetstream_slots_available_percentage_metric = prometheus_client.Gauge(
+      name="jetstream_slots_available_percentage",
+      documentation="The percentage of available slots in decode batch",
+      labelnames=["uuid", "idx"],
+  )

From ed50c6e15d7fa38ee06351737e5962fd902dd79d Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Wed, 15 May 2024 17:19:12 +0000
Subject: [PATCH 19/74] Initialize metrics

---
 jetstream/core/server_lib.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/jetstream/core/server_lib.py b/jetstream/core/server_lib.py
index cc983535..c272870f 100644
--- a/jetstream/core/server_lib.py
+++ b/jetstream/core/server_lib.py
@@ -31,6 +31,7 @@
 from jetstream.core.proto import jetstream_pb2_grpc
 
 from prometheus_client import start_http_server
+import jetstream.metrics as jetstream_metrics
 
 _HOST = "[::]"
 PROMETHEUS_ENABLED_ON_PORT = (
@@ -144,6 +145,8 @@ def run(
         "Starting Prometheus server on port %d", PROMETHEUS_ENABLED_ON_PORT
     )
     start_http_server(PROMETHEUS_ENABLED_ON_PORT)
+    jetstream_metrics.register_metrics()
+    
   else:
     logging.info(
         "Not starting Prometheus server: PROMETHEUS_ENABLED_ON_PORT not set"

From 95fa94b5e5810cecadbc80ad2da1cd3ee027163e Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Wed, 15 May 2024 17:21:17 +0000
Subject: [PATCH 20/74] reformat

---
 jetstream/core/server_lib.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jetstream/core/server_lib.py b/jetstream/core/server_lib.py
index c272870f..b14a3edd 100644
--- a/jetstream/core/server_lib.py
+++ b/jetstream/core/server_lib.py
@@ -146,7 +146,7 @@ def run(
     )
     start_http_server(PROMETHEUS_ENABLED_ON_PORT)
     jetstream_metrics.register_metrics()
-    
+
   else:
     logging.info(
         "Not starting Prometheus server: PROMETHEUS_ENABLED_ON_PORT not set"

From eef698e91516c637f2985af45ae97ace0db95323 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Wed, 15 May 2024 17:25:21 +0000
Subject: [PATCH 21/74] added 'global'

---
 jetstream/core/orchestrator.py |  2 +-
 jetstream/metrics/__init__.py  | 11 +++++++----
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index 9f954c15..661a8f05 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -578,7 +578,7 @@ def _generate_thread(self, idx: int):
         time_of_last_print = time.time()
 
       max_concurrent_decodes = generate_engine.max_concurrent_decodes
-      jetstream_metrics.jetstream_slots_available_percentage.labels(
+      jetstream_metrics.slots_available_percentage.labels(
           jetstream_metrics.instance_uuid, idx
       ).set_function(lambda: float(my_slots.qsize() / max_concurrent_decodes))
 
diff --git a/jetstream/metrics/__init__.py b/jetstream/metrics/__init__.py
index 292b7fe5..966a37eb 100644
--- a/jetstream/metrics/__init__.py
+++ b/jetstream/metrics/__init__.py
@@ -19,18 +19,21 @@
 
 # Metrics we care to observe
 prefill_backlog: prometheus_client.Gauge
-jetstream_slots_available_percentage: prometheus_client.Gauge
+slots_available_percentage: prometheus_client.Gauge
 
 
 def register_metrics():
-  prefill_backlog_metric = prometheus_client.Gauge(
+
+  global prefill_backlog
+  prefill_backlog = prometheus_client.Gauge(
       "jetstream_prefill_backlog_size",
       "Size of prefill queue",
       labelnames=["uuid"],
   )
 
-  jetstream_slots_available_percentage_metric = prometheus_client.Gauge(
-      name="jetstream_slots_available_percentage",
+  global slots_available_percentage
+  slots_available_percentage = prometheus_client.Gauge(
+      name="slots_available_percentage",
       documentation="The percentage of available slots in decode batch",
       labelnames=["uuid", "idx"],
   )

From cd882800314afc664b0e57d35287947f7c297da7 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Wed, 15 May 2024 17:34:15 +0000
Subject: [PATCH 22/74] Moved registration out of function

---
 jetstream/core/server_lib.py  |  1 -
 jetstream/metrics/__init__.py | 30 ++++++++++++------------------
 2 files changed, 12 insertions(+), 19 deletions(-)

diff --git a/jetstream/core/server_lib.py b/jetstream/core/server_lib.py
index b14a3edd..5c141ab7 100644
--- a/jetstream/core/server_lib.py
+++ b/jetstream/core/server_lib.py
@@ -145,7 +145,6 @@ def run(
         "Starting Prometheus server on port %d", PROMETHEUS_ENABLED_ON_PORT
     )
     start_http_server(PROMETHEUS_ENABLED_ON_PORT)
-    jetstream_metrics.register_metrics()
 
   else:
     logging.info(
diff --git a/jetstream/metrics/__init__.py b/jetstream/metrics/__init__.py
index 966a37eb..6aca93e6 100644
--- a/jetstream/metrics/__init__.py
+++ b/jetstream/metrics/__init__.py
@@ -12,28 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+""" All logic around adding metrics to the Prometheus client registry """
+
 import prometheus_client
 import shortuuid
 
 instance_uuid = shortuuid.uuid()
 
 # Metrics we care to observe
-prefill_backlog: prometheus_client.Gauge
-slots_available_percentage: prometheus_client.Gauge
-
-
-def register_metrics():
-
-  global prefill_backlog
-  prefill_backlog = prometheus_client.Gauge(
-      "jetstream_prefill_backlog_size",
-      "Size of prefill queue",
-      labelnames=["uuid"],
-  )
+prefill_backlog = prometheus_client.Gauge(
+    "jetstream_prefill_backlog_size",
+    "Size of prefill queue",
+    labelnames=["uuid"],
+)
 
-  global slots_available_percentage
-  slots_available_percentage = prometheus_client.Gauge(
-      name="slots_available_percentage",
-      documentation="The percentage of available slots in decode batch",
-      labelnames=["uuid", "idx"],
-  )
+slots_available_percentage = prometheus_client.Gauge(
+    name="slots_available_percentage",
+    documentation="The percentage of available slots in decode batch",
+    labelnames=["uuid", "idx"],
+)

From e46f6b122a49a88153ca8562d8c83db51825e202 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Wed, 15 May 2024 17:38:15 +0000
Subject: [PATCH 23/74] removed unused import

---
 jetstream/core/server_lib.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/jetstream/core/server_lib.py b/jetstream/core/server_lib.py
index 5c141ab7..0f781fca 100644
--- a/jetstream/core/server_lib.py
+++ b/jetstream/core/server_lib.py
@@ -31,7 +31,6 @@
 from jetstream.core.proto import jetstream_pb2_grpc
 
 from prometheus_client import start_http_server
-import jetstream.metrics as jetstream_metrics
 
 _HOST = "[::]"
 PROMETHEUS_ENABLED_ON_PORT = (

From 3eebce074cfc2f46e74a521ef389e0739dc4d014 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <bslabe123@gmail.com>
Date: Wed, 15 May 2024 10:39:05 -0700
Subject: [PATCH 24/74] Update server_lib.py

---
 jetstream/core/server_lib.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/jetstream/core/server_lib.py b/jetstream/core/server_lib.py
index 0f781fca..cc983535 100644
--- a/jetstream/core/server_lib.py
+++ b/jetstream/core/server_lib.py
@@ -144,7 +144,6 @@ def run(
         "Starting Prometheus server on port %d", PROMETHEUS_ENABLED_ON_PORT
     )
     start_http_server(PROMETHEUS_ENABLED_ON_PORT)
-
   else:
     logging.info(
         "Not starting Prometheus server: PROMETHEUS_ENABLED_ON_PORT not set"

From 5348ac8e4b2141bdf729e8a480c6585fc1768d4a Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Wed, 15 May 2024 18:06:12 +0000
Subject: [PATCH 25/74] jetstream_prefill_backlog_size -> prefill_backlog_size

---
 jetstream/metrics/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jetstream/metrics/__init__.py b/jetstream/metrics/__init__.py
index 6aca93e6..007bfa0d 100644
--- a/jetstream/metrics/__init__.py
+++ b/jetstream/metrics/__init__.py
@@ -21,7 +21,7 @@
 
 # Metrics we care to observe
 prefill_backlog = prometheus_client.Gauge(
-    "jetstream_prefill_backlog_size",
+    "prefill_backlog_size",
     "Size of prefill queue",
     labelnames=["uuid"],
 )

From 8f4afedaa6a0b9af0b33b3d552bcaad0ba550c3e Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Wed, 15 May 2024 18:06:40 +0000
Subject: [PATCH 26/74] label fields

---
 jetstream/metrics/__init__.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/jetstream/metrics/__init__.py b/jetstream/metrics/__init__.py
index 007bfa0d..b6937170 100644
--- a/jetstream/metrics/__init__.py
+++ b/jetstream/metrics/__init__.py
@@ -21,8 +21,8 @@
 
 # Metrics we care to observe
 prefill_backlog = prometheus_client.Gauge(
-    "prefill_backlog_size",
-    "Size of prefill queue",
+    name="prefill_backlog_size",
+    documentation="Size of prefill queue",
     labelnames=["uuid"],
 )
 

From 6e9c98d7881f4efe1c81211056b3df043e980b57 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Wed, 15 May 2024 18:09:09 +0000
Subject: [PATCH 27/74] rename metrics

---
 jetstream/metrics/__init__.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/jetstream/metrics/__init__.py b/jetstream/metrics/__init__.py
index b6937170..55a03511 100644
--- a/jetstream/metrics/__init__.py
+++ b/jetstream/metrics/__init__.py
@@ -21,13 +21,13 @@
 
 # Metrics we care to observe
 prefill_backlog = prometheus_client.Gauge(
-    name="prefill_backlog_size",
+    name="jetstream_prefill_backlog_size",
     documentation="Size of prefill queue",
     labelnames=["uuid"],
 )
 
 slots_available_percentage = prometheus_client.Gauge(
-    name="slots_available_percentage",
+    name="jetstream_slots_available_percentage",
     documentation="The percentage of available slots in decode batch",
     labelnames=["uuid", "idx"],
 )

From 3e3d6ad69ebdb3225660be0edff13f559b080ad2 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Wed, 15 May 2024 22:29:25 +0000
Subject: [PATCH 28/74] Moved metrics to singleton class

---
 jetstream/core/metrics/__init__.py   | 13 ++++++++++
 jetstream/core/metrics/prometheus.py | 37 ++++++++++++++++++++++++++++
 jetstream/core/orchestrator.py       | 19 ++++++++------
 jetstream/metrics/__init__.py        | 33 -------------------------
 requirements.in                      |  1 -
 requirements.txt                     |  2 --
 6 files changed, 62 insertions(+), 43 deletions(-)
 create mode 100644 jetstream/core/metrics/__init__.py
 create mode 100644 jetstream/core/metrics/prometheus.py
 delete mode 100644 jetstream/metrics/__init__.py

diff --git a/jetstream/core/metrics/__init__.py b/jetstream/core/metrics/__init__.py
new file mode 100644
index 00000000..2d9f7b38
--- /dev/null
+++ b/jetstream/core/metrics/__init__.py
@@ -0,0 +1,13 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
\ No newline at end of file
diff --git a/jetstream/core/metrics/prometheus.py b/jetstream/core/metrics/prometheus.py
new file mode 100644
index 00000000..38b9159d
--- /dev/null
+++ b/jetstream/core/metrics/prometheus.py
@@ -0,0 +1,37 @@
+import os
+from prometheus_client import Gauge
+from dataclasses import dataclass
+
+# Properties to be set for each metric
+@dataclass
+class _MetricLabels:
+    hostname: str = os.environ['HOSTNAME'];
+    idx: int = None
+
+def __create_labeled_metric(metric, labels: _MetricLabels):
+    return metric.labels(hostname = labels.hostname, idx = labels.idx if labels.idx else None)
+
+# Metric Definitions
+class JetstreamMetricsCollector:
+    def __new__(cls):
+        if not hasattr(cls, 'instance'):
+            cls.instance = super(JetstreamMetricsCollector, cls).__new__(cls)
+        return cls.instance
+    
+    # Metric definitions
+    __prefill_backlog = Gauge(
+            name="jetstream_prefill_backlog_size",
+            documentation="Size of prefill queue",
+            labelnames=["hostname"],
+        )
+    __slots_available_percentage = Gauge(
+        name="jetstream_slots_available_percentage",
+        documentation="The percentage of available slots in decode batch",
+        labelnames=["hostname", "idx"],
+        )
+
+    def get_prefill_backlog_metric(self):
+        return __create_labeled_metric(self.__prefill_backlog, self.MetricLabels(hostname=self.hostname))
+
+    def get_slots_available_percentage_metric(self, idx: int):
+        return __create_labeled_metric(self.__slots_available_percentage, self.MetricLabels(hostname=self.hostname, idx=idx))
\ No newline at end of file
diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index 661a8f05..5deca9f7 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -94,7 +94,7 @@
 from jetstream.core.utils import async_multifuture
 from jetstream.core.utils.return_sample import ReturnSample
 from jetstream.engine import engine_api, tokenizer_api, token_utils
-import jetstream.metrics as jetstream_metrics
+from jetstream.core.metrics.prometheus import JetstreamMetricsCollector
 import numpy as np
 
 root = logging.getLogger()
@@ -211,6 +211,12 @@ class Driver:
   # todo: remove jax_padding after all then engine migrate to np padding
   _jax_padding = True
 
+  # Record prometheus metrics if flag set
+  _metrics_collector: JetstreamMetricsCollector
+  if flag:
+    _metrics_collector = JetstreamMetricsCollector()
+  
+
   def __init__(
       self,
       prefill_engines: Optional[list[engine_api.Engine]] = None,
@@ -244,9 +250,8 @@ def __init__(
     # Stage 1
     # At first, a request is placed here in order to get prefilled.
     self._prefill_backlog = queue.Queue()
-    jetstream_metrics.prefill_backlog.labels(
-        jetstream_metrics.instance_uuid
-    ).set_function(lambda: float(self._prefill_backlog.qsize()))
+    if self._metrics_collector is not None:
+      self._metrics_collector.get_prefill_backlog_hostname_metric().set_function(lambda: float(self._prefill_backlog.qsize()))
 
     # Stage 2
     # After prefilling, it is placed here in order to get transferred to
@@ -578,9 +583,9 @@ def _generate_thread(self, idx: int):
         time_of_last_print = time.time()
 
       max_concurrent_decodes = generate_engine.max_concurrent_decodes
-      jetstream_metrics.slots_available_percentage.labels(
-          jetstream_metrics.instance_uuid, idx
-      ).set_function(lambda: float(my_slots.qsize() / max_concurrent_decodes))
+
+      if self._metrics_collector is not None:
+        self._metrics_collector.get_slots_available_percentage_metric(idx).set_function(lambda: float(my_slots.qsize() / max_concurrent_decodes))
 
       # Check if there are any free my_slots. We don't want to block here since
       # we can still generate if we can't insert. We do this in a while loop to
diff --git a/jetstream/metrics/__init__.py b/jetstream/metrics/__init__.py
deleted file mode 100644
index 55a03511..00000000
--- a/jetstream/metrics/__init__.py
+++ /dev/null
@@ -1,33 +0,0 @@
-# Copyright 2024 Google LLC
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-""" All logic around adding metrics to the Prometheus client registry """
-
-import prometheus_client
-import shortuuid
-
-instance_uuid = shortuuid.uuid()
-
-# Metrics we care to observe
-prefill_backlog = prometheus_client.Gauge(
-    name="jetstream_prefill_backlog_size",
-    documentation="Size of prefill queue",
-    labelnames=["uuid"],
-)
-
-slots_available_percentage = prometheus_client.Gauge(
-    name="jetstream_slots_available_percentage",
-    documentation="The percentage of available slots in decode batch",
-    labelnames=["uuid", "idx"],
-)
diff --git a/requirements.in b/requirements.in
index dfc5c4a4..56b2153c 100644
--- a/requirements.in
+++ b/requirements.in
@@ -12,5 +12,4 @@ seqio
 tiktoken
 blobfile
 parameterized
-shortuuid
 
diff --git a/requirements.txt b/requirements.txt
index ff60cf32..61dc801c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -238,8 +238,6 @@ sentencepiece==0.1.99
     # via seqio
 seqio==0.0.19
     # via -r requirements.in
-shortuuid==1.0.13
-    # via -r requirements.in
 six==1.16.0
     # via
     #   astunparse

From 091ebcc07b1cebc33bc914d863387382279626b0 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Wed, 15 May 2024 22:31:39 +0000
Subject: [PATCH 29/74] Revert requirements change

---
 requirements.in  | 2 +-
 requirements.txt | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/requirements.in b/requirements.in
index 56b2153c..105dbd45 100644
--- a/requirements.in
+++ b/requirements.in
@@ -12,4 +12,4 @@ seqio
 tiktoken
 blobfile
 parameterized
-
+shortuuid
diff --git a/requirements.txt b/requirements.txt
index 61dc801c..ff60cf32 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -238,6 +238,8 @@ sentencepiece==0.1.99
     # via seqio
 seqio==0.0.19
     # via -r requirements.in
+shortuuid==1.0.13
+    # via -r requirements.in
 six==1.16.0
     # via
     #   astunparse

From 8014cfcec790c6c6c575289d3b4c8c7526010512 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Wed, 15 May 2024 22:31:59 +0000
Subject: [PATCH 30/74] revert requirements changes

---
 requirements.in | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/requirements.in b/requirements.in
index 105dbd45..eba423d4 100644
--- a/requirements.in
+++ b/requirements.in
@@ -12,4 +12,4 @@ seqio
 tiktoken
 blobfile
 parameterized
-shortuuid
+shortuuid
\ No newline at end of file

From 717f7424a577d3056121f14f5f32b234426aeeb1 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Thu, 16 May 2024 21:17:16 +0000
Subject: [PATCH 31/74] Plumb metrics config to JetStreamServer run header

---
 jetstream/core/config_lib.py         |  5 +++++
 jetstream/core/metrics/prometheus.py |  1 +
 jetstream/core/server_lib.py         | 17 ++++++++---------
 jetstream/tests/core/test_server.py  |  2 ++
 4 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/jetstream/core/config_lib.py b/jetstream/core/config_lib.py
index fa33e45d..f8697e81 100644
--- a/jetstream/core/config_lib.py
+++ b/jetstream/core/config_lib.py
@@ -17,6 +17,7 @@
 import dataclasses
 import functools
 from typing import Any, Callable, List, Tuple, Type
+from numpy import uint16
 
 from jetstream.engine import engine_api
 from jetstream.engine import mock_engine
@@ -51,6 +52,10 @@ class InstantiatedEngines:
   generate_engines: List[engine_api.Engine]
   interleaved_engines: List[engine_api.Engine]
 
+@dataclasses.dataclass
+class MetricsConfig:
+  port: uint16
+
 
 # ▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼#
 
diff --git a/jetstream/core/metrics/prometheus.py b/jetstream/core/metrics/prometheus.py
index 38b9159d..965be367 100644
--- a/jetstream/core/metrics/prometheus.py
+++ b/jetstream/core/metrics/prometheus.py
@@ -12,6 +12,7 @@ def __create_labeled_metric(metric, labels: _MetricLabels):
     return metric.labels(hostname = labels.hostname, idx = labels.idx if labels.idx else None)
 
 # Metric Definitions
+# Wrapper class should be used to assure all metrics have proper tags
 class JetstreamMetricsCollector:
     def __new__(cls):
         if not hasattr(cls, 'instance'):
diff --git a/jetstream/core/server_lib.py b/jetstream/core/server_lib.py
index cc983535..8faabe6a 100644
--- a/jetstream/core/server_lib.py
+++ b/jetstream/core/server_lib.py
@@ -31,13 +31,9 @@
 from jetstream.core.proto import jetstream_pb2_grpc
 
 from prometheus_client import start_http_server
+from server_lib import MetricsConfig
 
 _HOST = "[::]"
-PROMETHEUS_ENABLED_ON_PORT = (
-    int(os.getenv("PROMETHEUS_ENABLED_ON_PORT"))
-    if os.getenv("PROMETHEUS_ENABLED_ON_PORT")
-    else None
-)
 
 
 class JetStreamServer:
@@ -99,6 +95,9 @@ def run(
     credentials: Any = grpc.insecure_server_credentials(),
     threads: int | None = None,
     jax_padding: bool = True,
+    metricsConfig: MetricsConfig = None
+
+
 ) -> JetStreamServer:
   """Runs a server with a specified config.
 
@@ -139,14 +138,14 @@ def run(
   jetstream_server.start()
 
   # Setup Prometheus server
-  if PROMETHEUS_ENABLED_ON_PORT is not None:
+  if config.prometheus_port is not None:
     logging.info(
-        "Starting Prometheus server on port %d", PROMETHEUS_ENABLED_ON_PORT
+        f"Starting Prometheus server on port {config}", 
     )
-    start_http_server(PROMETHEUS_ENABLED_ON_PORT)
+    start_http_server(config.prometheus_port)
   else:
     logging.info(
-        "Not starting Prometheus server: PROMETHEUS_ENABLED_ON_PORT not set"
+        "Not starting Prometheus server: --prometheus_port flag not set"
     )
   return jetstream_server
 
diff --git a/jetstream/tests/core/test_server.py b/jetstream/tests/core/test_server.py
index 9dda982a..7c7528d4 100644
--- a/jetstream/tests/core/test_server.py
+++ b/jetstream/tests/core/test_server.py
@@ -96,6 +96,8 @@ async def test_server(
         counter += 1
       server.stop()
 
+      # TODO: explore adding prometheus tests
+
 
 if __name__ == "__main__":
   unittest.main()

From ddeb240f5f9b449a4453d186fbe8506a7747b690 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Thu, 16 May 2024 21:44:16 +0000
Subject: [PATCH 32/74] Cleanup of prior commit

---
 jetstream/core/config_lib.py         |  3 +-
 jetstream/core/metrics/__init__.py   |  2 +-
 jetstream/core/metrics/prometheus.py | 61 ++++++++++++++++------------
 jetstream/core/orchestrator.py       | 15 ++++---
 jetstream/core/server_lib.py         | 33 +++++++--------
 5 files changed, 65 insertions(+), 49 deletions(-)

diff --git a/jetstream/core/config_lib.py b/jetstream/core/config_lib.py
index f8697e81..0a83e607 100644
--- a/jetstream/core/config_lib.py
+++ b/jetstream/core/config_lib.py
@@ -52,8 +52,9 @@ class InstantiatedEngines:
   generate_engines: List[engine_api.Engine]
   interleaved_engines: List[engine_api.Engine]
 
+
 @dataclasses.dataclass
-class MetricsConfig:
+class MetricsServerConfig:
   port: uint16
 
 
diff --git a/jetstream/core/metrics/__init__.py b/jetstream/core/metrics/__init__.py
index 2d9f7b38..6d5e14bc 100644
--- a/jetstream/core/metrics/__init__.py
+++ b/jetstream/core/metrics/__init__.py
@@ -10,4 +10,4 @@
 # distributed under the License is distributed on an "AS IS" BASIS,
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
-# limitations under the License.
\ No newline at end of file
+# limitations under the License.
diff --git a/jetstream/core/metrics/prometheus.py b/jetstream/core/metrics/prometheus.py
index 965be367..29d9b3e0 100644
--- a/jetstream/core/metrics/prometheus.py
+++ b/jetstream/core/metrics/prometheus.py
@@ -2,37 +2,48 @@
 from prometheus_client import Gauge
 from dataclasses import dataclass
 
+
 # Properties to be set for each metric
 @dataclass
 class _MetricLabels:
-    hostname: str = os.environ['HOSTNAME'];
-    idx: int = None
+  hostname: str = os.environ["HOSTNAME"]
+  idx: int = None
+
 
 def __create_labeled_metric(metric, labels: _MetricLabels):
-    return metric.labels(hostname = labels.hostname, idx = labels.idx if labels.idx else None)
+  return metric.labels(
+      hostname=labels.hostname, idx=labels.idx if labels.idx else None
+  )
+
 
 # Metric Definitions
 # Wrapper class should be used to assure all metrics have proper tags
 class JetstreamMetricsCollector:
-    def __new__(cls):
-        if not hasattr(cls, 'instance'):
-            cls.instance = super(JetstreamMetricsCollector, cls).__new__(cls)
-        return cls.instance
-    
-    # Metric definitions
-    __prefill_backlog = Gauge(
-            name="jetstream_prefill_backlog_size",
-            documentation="Size of prefill queue",
-            labelnames=["hostname"],
-        )
-    __slots_available_percentage = Gauge(
-        name="jetstream_slots_available_percentage",
-        documentation="The percentage of available slots in decode batch",
-        labelnames=["hostname", "idx"],
-        )
-
-    def get_prefill_backlog_metric(self):
-        return __create_labeled_metric(self.__prefill_backlog, self.MetricLabels(hostname=self.hostname))
-
-    def get_slots_available_percentage_metric(self, idx: int):
-        return __create_labeled_metric(self.__slots_available_percentage, self.MetricLabels(hostname=self.hostname, idx=idx))
\ No newline at end of file
+
+  def __new__(cls):
+    if not hasattr(cls, "instance"):
+      cls.instance = super(JetstreamMetricsCollector, cls).__new__(cls)
+    return cls.instance
+
+  # Metric definitions
+  __prefill_backlog = Gauge(
+      name="jetstream_prefill_backlog_size",
+      documentation="Size of prefill queue",
+      labelnames=["hostname"],
+  )
+  __slots_available_percentage = Gauge(
+      name="jetstream_slots_available_percentage",
+      documentation="The percentage of available slots in decode batch",
+      labelnames=["hostname", "idx"],
+  )
+
+  def get_prefill_backlog_metric(self):
+    return __create_labeled_metric(
+        self.__prefill_backlog, self.MetricLabels(hostname=self.hostname)
+    )
+
+  def get_slots_available_percentage_metric(self, idx: int):
+    return __create_labeled_metric(
+        self.__slots_available_percentage,
+        self.MetricLabels(hostname=self.hostname, idx=idx),
+    )
diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index 5deca9f7..fd9a83a9 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -211,11 +211,8 @@ class Driver:
   # todo: remove jax_padding after all then engine migrate to np padding
   _jax_padding = True
 
-  # Record prometheus metrics if flag set
+  # All metrics we want to monitor should be collected with this
   _metrics_collector: JetstreamMetricsCollector
-  if flag:
-    _metrics_collector = JetstreamMetricsCollector()
-  
 
   def __init__(
       self,
@@ -225,6 +222,7 @@ def __init__(
       generate_params: Optional[list[Any]] = None,
       interleaved_mode: bool = False,
       jax_padding: bool = True,
+      metrics_collector: JetstreamMetricsCollector = None,
   ):
     if prefill_engines is None:
       prefill_engines = []
@@ -245,13 +243,16 @@ def __init__(
     self._prefill_params = prefill_params
     self._generate_params = generate_params
     self._interleaved_mode = interleaved_mode
+    self._metrics_collector = metrics_collector
 
     # Stages 1-4 represent the life cycle of a request.
     # Stage 1
     # At first, a request is placed here in order to get prefilled.
     self._prefill_backlog = queue.Queue()
     if self._metrics_collector is not None:
-      self._metrics_collector.get_prefill_backlog_hostname_metric().set_function(lambda: float(self._prefill_backlog.qsize()))
+      self._metrics_collector.get_prefill_backlog_hostname_metric().set_function(
+          lambda: float(self._prefill_backlog.qsize())
+      )
 
     # Stage 2
     # After prefilling, it is placed here in order to get transferred to
@@ -585,7 +586,9 @@ def _generate_thread(self, idx: int):
       max_concurrent_decodes = generate_engine.max_concurrent_decodes
 
       if self._metrics_collector is not None:
-        self._metrics_collector.get_slots_available_percentage_metric(idx).set_function(lambda: float(my_slots.qsize() / max_concurrent_decodes))
+        self._metrics_collector.get_slots_available_percentage_metric(
+            idx
+        ).set_function(lambda: float(my_slots.qsize() / max_concurrent_decodes))
 
       # Check if there are any free my_slots. We don't want to block here since
       # we can still generate if we can't insert. We do this in a while loop to
diff --git a/jetstream/core/server_lib.py b/jetstream/core/server_lib.py
index 8faabe6a..f10c4d8d 100644
--- a/jetstream/core/server_lib.py
+++ b/jetstream/core/server_lib.py
@@ -28,10 +28,10 @@
 import jax
 from jetstream.core import config_lib
 from jetstream.core import orchestrator
+from jetstream.core.metrics.prometheus import JetstreamMetricsCollector
 from jetstream.core.proto import jetstream_pb2_grpc
 
 from prometheus_client import start_http_server
-from server_lib import MetricsConfig
 
 _HOST = "[::]"
 
@@ -95,9 +95,7 @@ def run(
     credentials: Any = grpc.insecure_server_credentials(),
     threads: int | None = None,
     jax_padding: bool = True,
-    metricsConfig: MetricsConfig = None
-
-
+    metrics_server_config: config_lib.MetricsServerConfig = None,
 ) -> JetStreamServer:
   """Runs a server with a specified config.
 
@@ -121,6 +119,20 @@ def run(
   interleaved_mode = (
       len(config.prefill_slices) + len(config.generate_slices) == 0
   )
+  
+  # Setup Prometheus server
+  metrics_collector: JetstreamMetricsCollector = None
+  if metrics_server_config is not None:
+    logging.info(
+        f"Starting Prometheus server on port {metrics_server_config.port}",
+    )
+    start_http_server(metrics_server_config.port)
+    metrics_collector = JetstreamMetricsCollector()
+  else:
+    logging.info(
+        "Not starting Prometheus server: --prometheus_port flag not set"
+    )
+
   driver = orchestrator.Driver(
       prefill_engines=engines.prefill_engines + engines.interleaved_engines,
       generate_engines=engines.generate_engines + engines.interleaved_engines,
@@ -128,6 +140,7 @@ def run(
       generate_params=generate_params + shared_params,
       interleaved_mode=interleaved_mode,
       jax_padding=jax_padding,
+      metrics_collector=metrics_collector
   )
   # We default threads to the total number of concurrent allowed decodes,
   # to make sure we can fully saturate the model. Set default minimum to 64.
@@ -136,20 +149,8 @@ def run(
   logging.info("Starting server on port %d with %d threads", port, threads)
 
   jetstream_server.start()
-
-  # Setup Prometheus server
-  if config.prometheus_port is not None:
-    logging.info(
-        f"Starting Prometheus server on port {config}", 
-    )
-    start_http_server(config.prometheus_port)
-  else:
-    logging.info(
-        "Not starting Prometheus server: --prometheus_port flag not set"
-    )
   return jetstream_server
 
-
 def get_devices() -> Any:
   """Gets devices."""
   # TODO: Add more logs for the devices.

From ea222f01a41c800eb7688bdcc3c1fbbe716259dd Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Thu, 16 May 2024 22:19:30 +0000
Subject: [PATCH 33/74] Default for hostname metric label

---
 jetstream/core/metrics/prometheus.py | 3 ++-
 jetstream/core/server_lib.py         | 5 +++--
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/jetstream/core/metrics/prometheus.py b/jetstream/core/metrics/prometheus.py
index 29d9b3e0..f8936097 100644
--- a/jetstream/core/metrics/prometheus.py
+++ b/jetstream/core/metrics/prometheus.py
@@ -1,4 +1,5 @@
 import os
+import shortuuid
 from prometheus_client import Gauge
 from dataclasses import dataclass
 
@@ -6,7 +7,7 @@
 # Properties to be set for each metric
 @dataclass
 class _MetricLabels:
-  hostname: str = os.environ["HOSTNAME"]
+  hostname: str = os.getenv("HOSTNAME", shortuuid.uuid())
   idx: int = None
 
 
diff --git a/jetstream/core/server_lib.py b/jetstream/core/server_lib.py
index f10c4d8d..c2c9985d 100644
--- a/jetstream/core/server_lib.py
+++ b/jetstream/core/server_lib.py
@@ -119,7 +119,7 @@ def run(
   interleaved_mode = (
       len(config.prefill_slices) + len(config.generate_slices) == 0
   )
-  
+
   # Setup Prometheus server
   metrics_collector: JetstreamMetricsCollector = None
   if metrics_server_config is not None:
@@ -140,7 +140,7 @@ def run(
       generate_params=generate_params + shared_params,
       interleaved_mode=interleaved_mode,
       jax_padding=jax_padding,
-      metrics_collector=metrics_collector
+      metrics_collector=metrics_collector,
   )
   # We default threads to the total number of concurrent allowed decodes,
   # to make sure we can fully saturate the model. Set default minimum to 64.
@@ -151,6 +151,7 @@ def run(
   jetstream_server.start()
   return jetstream_server
 
+
 def get_devices() -> Any:
   """Gets devices."""
   # TODO: Add more logs for the devices.

From ab0561ba4022180c68e0ac86404cdc67f9986c5b Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Thu, 16 May 2024 22:25:02 +0000
Subject: [PATCH 34/74] __ -> _

---
 jetstream/core/metrics/prometheus.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/jetstream/core/metrics/prometheus.py b/jetstream/core/metrics/prometheus.py
index f8936097..32ff704d 100644
--- a/jetstream/core/metrics/prometheus.py
+++ b/jetstream/core/metrics/prometheus.py
@@ -11,7 +11,7 @@ class _MetricLabels:
   idx: int = None
 
 
-def __create_labeled_metric(metric, labels: _MetricLabels):
+def _create_labeled_metric(metric, labels: _MetricLabels):
   return metric.labels(
       hostname=labels.hostname, idx=labels.idx if labels.idx else None
   )
@@ -39,12 +39,12 @@ def __new__(cls):
   )
 
   def get_prefill_backlog_metric(self):
-    return __create_labeled_metric(
+    return _create_labeled_metric(
         self.__prefill_backlog, self.MetricLabels(hostname=self.hostname)
     )
 
   def get_slots_available_percentage_metric(self, idx: int):
-    return __create_labeled_metric(
+    return _create_labeled_metric(
         self.__slots_available_percentage,
         self.MetricLabels(hostname=self.hostname, idx=idx),
     )

From 2f8b55c988bfc406ed8555e69a11558a788406c2 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Thu, 16 May 2024 22:29:21 +0000
Subject: [PATCH 35/74] Linter error

---
 jetstream/core/metrics/prometheus.py | 6 ++----
 1 file changed, 2 insertions(+), 4 deletions(-)

diff --git a/jetstream/core/metrics/prometheus.py b/jetstream/core/metrics/prometheus.py
index 32ff704d..deec50e0 100644
--- a/jetstream/core/metrics/prometheus.py
+++ b/jetstream/core/metrics/prometheus.py
@@ -39,12 +39,10 @@ def __new__(cls):
   )
 
   def get_prefill_backlog_metric(self):
-    return _create_labeled_metric(
-        self.__prefill_backlog, self.MetricLabels(hostname=self.hostname)
-    )
+    return _create_labeled_metric(self.__prefill_backlog, _MetricLabels())
 
   def get_slots_available_percentage_metric(self, idx: int):
     return _create_labeled_metric(
         self.__slots_available_percentage,
-        self.MetricLabels(hostname=self.hostname, idx=idx),
+        _MetricLabels(idx=idx),
     )

From ea8e2c7015954d03705bdcb6be3fd189988eaa34 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Thu, 16 May 2024 22:35:08 +0000
Subject: [PATCH 36/74] linter error

---
 jetstream/core/metrics/prometheus.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jetstream/core/metrics/prometheus.py b/jetstream/core/metrics/prometheus.py
index deec50e0..6e95e7a0 100644
--- a/jetstream/core/metrics/prometheus.py
+++ b/jetstream/core/metrics/prometheus.py
@@ -8,7 +8,7 @@
 @dataclass
 class _MetricLabels:
   hostname: str = os.getenv("HOSTNAME", shortuuid.uuid())
-  idx: int = None
+  idx: int
 
 
 def _create_labeled_metric(metric, labels: _MetricLabels):

From adf47fd836574127ba2e3be2bb74279496bebc66 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Thu, 16 May 2024 22:40:00 +0000
Subject: [PATCH 37/74] linter error

---
 jetstream/core/metrics/prometheus.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/jetstream/core/metrics/prometheus.py b/jetstream/core/metrics/prometheus.py
index 6e95e7a0..cac7e13e 100644
--- a/jetstream/core/metrics/prometheus.py
+++ b/jetstream/core/metrics/prometheus.py
@@ -12,9 +12,7 @@ class _MetricLabels:
 
 
 def _create_labeled_metric(metric, labels: _MetricLabels):
-  return metric.labels(
-      hostname=labels.hostname, idx=labels.idx if labels.idx else None
-  )
+  return metric.labels(hostname=labels.hostname, idx=labels.idx)
 
 
 # Metric Definitions

From 8694cdda240d2ad90b63ced4dfb5585be4cc153a Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Thu, 16 May 2024 22:43:07 +0000
Subject: [PATCH 38/74] default value for idx

---
 jetstream/core/metrics/prometheus.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jetstream/core/metrics/prometheus.py b/jetstream/core/metrics/prometheus.py
index cac7e13e..9112b7f8 100644
--- a/jetstream/core/metrics/prometheus.py
+++ b/jetstream/core/metrics/prometheus.py
@@ -8,7 +8,7 @@
 @dataclass
 class _MetricLabels:
   hostname: str = os.getenv("HOSTNAME", shortuuid.uuid())
-  idx: int
+  idx: int = 0
 
 
 def _create_labeled_metric(metric, labels: _MetricLabels):

From 859c4ac5522d65929545d39fb5ff915eb59da553 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Thu, 16 May 2024 22:49:16 +0000
Subject: [PATCH 39/74] Linter error

---
 jetstream/core/orchestrator.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index fd9a83a9..8565692a 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -243,14 +243,15 @@ def __init__(
     self._prefill_params = prefill_params
     self._generate_params = generate_params
     self._interleaved_mode = interleaved_mode
-    self._metrics_collector = metrics_collector
+    if metrics_collector is not None:
+      self._metrics_collector = metrics_collector
 
     # Stages 1-4 represent the life cycle of a request.
     # Stage 1
     # At first, a request is placed here in order to get prefilled.
     self._prefill_backlog = queue.Queue()
     if self._metrics_collector is not None:
-      self._metrics_collector.get_prefill_backlog_hostname_metric().set_function(
+      self._metrics_collector.get_prefill_backlog_metric().set_function(
           lambda: float(self._prefill_backlog.qsize())
       )
 

From 8bf88f77745dadb676e7a3889b1ed691d98e2a2f Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Thu, 16 May 2024 22:52:43 +0000
Subject: [PATCH 40/74] Linter error

---
 jetstream/core/orchestrator.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index 8565692a..66854ad1 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -212,7 +212,7 @@ class Driver:
   _jax_padding = True
 
   # All metrics we want to monitor should be collected with this
-  _metrics_collector: JetstreamMetricsCollector
+  _metrics_collector: JetstreamMetricsCollector = None
 
   def __init__(
       self,
@@ -243,8 +243,7 @@ def __init__(
     self._prefill_params = prefill_params
     self._generate_params = generate_params
     self._interleaved_mode = interleaved_mode
-    if metrics_collector is not None:
-      self._metrics_collector = metrics_collector
+    self._metrics_collector = metrics_collector
 
     # Stages 1-4 represent the life cycle of a request.
     # Stage 1

From bb21b7162c8d9363c2e203d5ccfe1b16a470fe7b Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Thu, 16 May 2024 22:56:30 +0000
Subject: [PATCH 41/74] Linter error final fix hopefully

---
 jetstream/core/orchestrator.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index 66854ad1..fa818ea1 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -212,7 +212,7 @@ class Driver:
   _jax_padding = True
 
   # All metrics we want to monitor should be collected with this
-  _metrics_collector: JetstreamMetricsCollector = None
+  _metrics_collector: JetstreamMetricsCollector | None = None
 
   def __init__(
       self,

From da721fba8ac328819449cdec818af1f86a29501f Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Thu, 16 May 2024 22:59:40 +0000
Subject: [PATCH 42/74] Fixed type annotation

---
 jetstream/core/orchestrator.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index fa818ea1..51212ad0 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -222,7 +222,7 @@ def __init__(
       generate_params: Optional[list[Any]] = None,
       interleaved_mode: bool = False,
       jax_padding: bool = True,
-      metrics_collector: JetstreamMetricsCollector = None,
+      metrics_collector: JetstreamMetricsCollector | None = None,
   ):
     if prefill_engines is None:
       prefill_engines = []

From 0b368e3eba629d32e49ad1c2c2d71e652f1186ef Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Thu, 16 May 2024 23:03:47 +0000
Subject: [PATCH 43/74] Type fix in server_lib.py

---
 jetstream/core/server_lib.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jetstream/core/server_lib.py b/jetstream/core/server_lib.py
index c2c9985d..3f51f5ac 100644
--- a/jetstream/core/server_lib.py
+++ b/jetstream/core/server_lib.py
@@ -95,7 +95,7 @@ def run(
     credentials: Any = grpc.insecure_server_credentials(),
     threads: int | None = None,
     jax_padding: bool = True,
-    metrics_server_config: config_lib.MetricsServerConfig = None,
+    metrics_server_config: config_lib.MetricsServerConfig | None = None,
 ) -> JetStreamServer:
   """Runs a server with a specified config.
 

From d767b636735f004d28273e2ea4d1cf86f8b13571 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Thu, 16 May 2024 23:08:06 +0000
Subject: [PATCH 44/74] __ -> _ in JetstreamMetricsCollector property names

---
 jetstream/core/metrics/prometheus.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/jetstream/core/metrics/prometheus.py b/jetstream/core/metrics/prometheus.py
index 9112b7f8..5476fc50 100644
--- a/jetstream/core/metrics/prometheus.py
+++ b/jetstream/core/metrics/prometheus.py
@@ -25,22 +25,22 @@ def __new__(cls):
     return cls.instance
 
   # Metric definitions
-  __prefill_backlog = Gauge(
+  _prefill_backlog = Gauge(
       name="jetstream_prefill_backlog_size",
       documentation="Size of prefill queue",
       labelnames=["hostname"],
   )
-  __slots_available_percentage = Gauge(
+  _slots_available_percentage = Gauge(
       name="jetstream_slots_available_percentage",
       documentation="The percentage of available slots in decode batch",
       labelnames=["hostname", "idx"],
   )
 
   def get_prefill_backlog_metric(self):
-    return _create_labeled_metric(self.__prefill_backlog, _MetricLabels())
+    return _create_labeled_metric(self._prefill_backlog, _MetricLabels())
 
   def get_slots_available_percentage_metric(self, idx: int):
     return _create_labeled_metric(
-        self.__slots_available_percentage,
+        self._slots_available_percentage,
         _MetricLabels(idx=idx),
     )

From 651bbe06b1c2405b50fb5ba64b84e12cf1d7cf71 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Thu, 16 May 2024 23:13:12 +0000
Subject: [PATCH 45/74] added docstrings

---
 jetstream/core/metrics/prometheus.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/jetstream/core/metrics/prometheus.py b/jetstream/core/metrics/prometheus.py
index 5476fc50..134a000d 100644
--- a/jetstream/core/metrics/prometheus.py
+++ b/jetstream/core/metrics/prometheus.py
@@ -4,9 +4,10 @@
 from dataclasses import dataclass
 
 
-# Properties to be set for each metric
 @dataclass
 class _MetricLabels:
+  """All metrics emmitted should be configured with the following labels"""
+
   hostname: str = os.getenv("HOSTNAME", shortuuid.uuid())
   idx: int = 0
 
@@ -15,9 +16,8 @@ def _create_labeled_metric(metric, labels: _MetricLabels):
   return metric.labels(hostname=labels.hostname, idx=labels.idx)
 
 
-# Metric Definitions
-# Wrapper class should be used to assure all metrics have proper tags
 class JetstreamMetricsCollector:
+  """Wrapper class should be used to assure all metrics have proper tags"""
 
   def __new__(cls):
     if not hasattr(cls, "instance"):

From 27b9b43bb82157ce0560a181d36b735326faeede Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Thu, 16 May 2024 23:18:02 +0000
Subject: [PATCH 46/74] added module docstring

---
 jetstream/core/metrics/prometheus.py | 16 ++++++++++++++++
 jetstream/core/server_lib.py         |  1 -
 2 files changed, 16 insertions(+), 1 deletion(-)

diff --git a/jetstream/core/metrics/prometheus.py b/jetstream/core/metrics/prometheus.py
index 134a000d..6502168d 100644
--- a/jetstream/core/metrics/prometheus.py
+++ b/jetstream/core/metrics/prometheus.py
@@ -1,3 +1,19 @@
+# Copyright 2024 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Contains common functions for configuring Jetstream server metrics"""
+
 import os
 import shortuuid
 from prometheus_client import Gauge
diff --git a/jetstream/core/server_lib.py b/jetstream/core/server_lib.py
index 3f51f5ac..6bcd8ea9 100644
--- a/jetstream/core/server_lib.py
+++ b/jetstream/core/server_lib.py
@@ -20,7 +20,6 @@
 import asyncio
 from concurrent import futures
 import logging
-import os
 import threading
 from typing import Any, Type
 

From 78d1825dd3adebb1f96bbf53e921c5576f08951f Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Thu, 16 May 2024 23:22:51 +0000
Subject: [PATCH 47/74] laxy % formatting

---
 jetstream/core/server_lib.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jetstream/core/server_lib.py b/jetstream/core/server_lib.py
index 6bcd8ea9..074ddb3e 100644
--- a/jetstream/core/server_lib.py
+++ b/jetstream/core/server_lib.py
@@ -123,7 +123,7 @@ def run(
   metrics_collector: JetstreamMetricsCollector = None
   if metrics_server_config is not None:
     logging.info(
-        f"Starting Prometheus server on port {metrics_server_config.port}",
+        "Starting Prometheus server on port %d", metrics_server_config.port
     )
     start_http_server(metrics_server_config.port)
     metrics_collector = JetstreamMetricsCollector()

From 3c9b4beb3f9d5d326b5c0712b6558e5e51819225 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 16:56:22 +0000
Subject: [PATCH 48/74] metrics port cannot be 0

---
 jetstream/core/server_lib.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jetstream/core/server_lib.py b/jetstream/core/server_lib.py
index 074ddb3e..b39ff460 100644
--- a/jetstream/core/server_lib.py
+++ b/jetstream/core/server_lib.py
@@ -121,7 +121,7 @@ def run(
 
   # Setup Prometheus server
   metrics_collector: JetstreamMetricsCollector = None
-  if metrics_server_config is not None:
+  if metrics_server_config is not None and metrics_server_config.port is not 0:
     logging.info(
         "Starting Prometheus server on port %d", metrics_server_config.port
     )

From 3d4c5cf9edd4762f6a580210c9b938e2c55df94a Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 17:18:42 +0000
Subject: [PATCH 49/74] idx can be None

---
 jetstream/core/metrics/prometheus.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jetstream/core/metrics/prometheus.py b/jetstream/core/metrics/prometheus.py
index 6502168d..809d4981 100644
--- a/jetstream/core/metrics/prometheus.py
+++ b/jetstream/core/metrics/prometheus.py
@@ -25,7 +25,7 @@ class _MetricLabels:
   """All metrics emmitted should be configured with the following labels"""
 
   hostname: str = os.getenv("HOSTNAME", shortuuid.uuid())
-  idx: int = 0
+  idx: int | None = None
 
 
 def _create_labeled_metric(metric, labels: _MetricLabels):

From 587b3cdfb1f5c94693f3462e75620fd0285c2cee Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 17:37:13 +0000
Subject: [PATCH 50/74] Removed redundant class

---
 jetstream/core/metrics/prometheus.py | 21 +++++----------------
 1 file changed, 5 insertions(+), 16 deletions(-)

diff --git a/jetstream/core/metrics/prometheus.py b/jetstream/core/metrics/prometheus.py
index 809d4981..def4e75f 100644
--- a/jetstream/core/metrics/prometheus.py
+++ b/jetstream/core/metrics/prometheus.py
@@ -20,21 +20,11 @@
 from dataclasses import dataclass
 
 
-@dataclass
-class _MetricLabels:
-  """All metrics emmitted should be configured with the following labels"""
-
-  hostname: str = os.getenv("HOSTNAME", shortuuid.uuid())
-  idx: int | None = None
-
-
-def _create_labeled_metric(metric, labels: _MetricLabels):
-  return metric.labels(hostname=labels.hostname, idx=labels.idx)
-
-
 class JetstreamMetricsCollector:
   """Wrapper class should be used to assure all metrics have proper tags"""
 
+  _hostname: str = os.getenv("HOSTNAME", shortuuid.uuid())
+
   def __new__(cls):
     if not hasattr(cls, "instance"):
       cls.instance = super(JetstreamMetricsCollector, cls).__new__(cls)
@@ -53,10 +43,9 @@ def __new__(cls):
   )
 
   def get_prefill_backlog_metric(self):
-    return _create_labeled_metric(self._prefill_backlog, _MetricLabels())
+    return self._prefill_backlog.labels(hostname=self._hostname)
 
   def get_slots_available_percentage_metric(self, idx: int):
-    return _create_labeled_metric(
-        self._slots_available_percentage,
-        _MetricLabels(idx=idx),
+    return self._slots_available_percentage.labels(
+        hostname=self._hostname, idx=idx
     )

From c597a800f6c07585de9ffb26e87cf9e26e1eac8e Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 17:41:52 +0000
Subject: [PATCH 51/74] removed lingering import

---
 jetstream/core/metrics/prometheus.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/jetstream/core/metrics/prometheus.py b/jetstream/core/metrics/prometheus.py
index def4e75f..35ef3d65 100644
--- a/jetstream/core/metrics/prometheus.py
+++ b/jetstream/core/metrics/prometheus.py
@@ -17,7 +17,6 @@
 import os
 import shortuuid
 from prometheus_client import Gauge
-from dataclasses import dataclass
 
 
 class JetstreamMetricsCollector:

From b0271267733b427cd0c80048e5a9b0e5f51cb863 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <bslabe123@gmail.com>
Date: Fri, 17 May 2024 10:49:57 -0700
Subject: [PATCH 52/74] Update test_server.py

---
 jetstream/tests/core/test_server.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/jetstream/tests/core/test_server.py b/jetstream/tests/core/test_server.py
index 7c7528d4..9dda982a 100644
--- a/jetstream/tests/core/test_server.py
+++ b/jetstream/tests/core/test_server.py
@@ -96,8 +96,6 @@ async def test_server(
         counter += 1
       server.stop()
 
-      # TODO: explore adding prometheus tests
-
 
 if __name__ == "__main__":
   unittest.main()

From c93362087aafb263dbe69f1ed241468161a19ac7 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 18:09:31 +0000
Subject: [PATCH 53/74] hostname -> id

---
 jetstream/core/metrics/prometheus.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/jetstream/core/metrics/prometheus.py b/jetstream/core/metrics/prometheus.py
index 35ef3d65..0347e49a 100644
--- a/jetstream/core/metrics/prometheus.py
+++ b/jetstream/core/metrics/prometheus.py
@@ -22,7 +22,7 @@
 class JetstreamMetricsCollector:
   """Wrapper class should be used to assure all metrics have proper tags"""
 
-  _hostname: str = os.getenv("HOSTNAME", shortuuid.uuid())
+  _id: str = os.getenv("HOSTNAME", shortuuid.uuid())
 
   def __new__(cls):
     if not hasattr(cls, "instance"):
@@ -33,18 +33,18 @@ def __new__(cls):
   _prefill_backlog = Gauge(
       name="jetstream_prefill_backlog_size",
       documentation="Size of prefill queue",
-      labelnames=["hostname"],
+      labelnames=["id"],
   )
   _slots_available_percentage = Gauge(
       name="jetstream_slots_available_percentage",
       documentation="The percentage of available slots in decode batch",
-      labelnames=["hostname", "idx"],
+      labelnames=["id", "idx"],
   )
 
   def get_prefill_backlog_metric(self):
-    return self._prefill_backlog.labels(hostname=self._hostname)
+    return self._prefill_backlog.labels(id=self._id)
 
   def get_slots_available_percentage_metric(self, idx: int):
     return self._slots_available_percentage.labels(
-        hostname=self._hostname, idx=idx
+        id=self._id, idx=idx
     )

From 2f7e8dd6582edb264d7f62d1b5f05bd5bc68d913 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 18:10:45 +0000
Subject: [PATCH 54/74] linter

---
 jetstream/core/metrics/prometheus.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/jetstream/core/metrics/prometheus.py b/jetstream/core/metrics/prometheus.py
index 0347e49a..de0be2c2 100644
--- a/jetstream/core/metrics/prometheus.py
+++ b/jetstream/core/metrics/prometheus.py
@@ -45,6 +45,4 @@ def get_prefill_backlog_metric(self):
     return self._prefill_backlog.labels(id=self._id)
 
   def get_slots_available_percentage_metric(self, idx: int):
-    return self._slots_available_percentage.labels(
-        id=self._id, idx=idx
-    )
+    return self._slots_available_percentage.labels(id=self._id, idx=idx)

From 6203941955f0c327432473b4bf44b070bedaec65 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 18:31:05 +0000
Subject: [PATCH 55/74] is not -> !=

---
 jetstream/core/server_lib.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jetstream/core/server_lib.py b/jetstream/core/server_lib.py
index b39ff460..9d4fc8dd 100644
--- a/jetstream/core/server_lib.py
+++ b/jetstream/core/server_lib.py
@@ -121,7 +121,7 @@ def run(
 
   # Setup Prometheus server
   metrics_collector: JetstreamMetricsCollector = None
-  if metrics_server_config is not None and metrics_server_config.port is not 0:
+  if metrics_server_config is not None and metrics_server_config.port != 0:
     logging.info(
         "Starting Prometheus server on port %d", metrics_server_config.port
     )

From cc17e4606d198137070ab75458d6ee3b347dd5bf Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 20:05:20 +0000
Subject: [PATCH 56/74] Zijun nits

---
 docs/online-inference-with-maxtext-engine.md | 21 +++++++++++++++++++-
 jetstream/core/orchestrator.py               |  4 ++--
 jetstream/core/server_lib.py                 |  2 +-
 jetstream/tests/core/test_server.py          | 12 ++++++++++-
 4 files changed, 34 insertions(+), 5 deletions(-)

diff --git a/docs/online-inference-with-maxtext-engine.md b/docs/online-inference-with-maxtext-engine.md
index 90044b9e..8e612e40 100644
--- a/docs/online-inference-with-maxtext-engine.md
+++ b/docs/online-inference-with-maxtext-engine.md
@@ -109,6 +109,7 @@ export ICI_TENSOR_PARALLELISM=1
 export SCAN_LAYERS=false
 export WEIGHT_DTYPE=bfloat16
 export PER_DEVICE_BATCH_SIZE=11
+export PROMETHEUS_PORT=9090
 ```
 
 #### Create Llama2-7b environment variables for server flags
@@ -127,6 +128,7 @@ export ICI_TENSOR_PARALLELISM=1
 export SCAN_LAYERS=false
 export WEIGHT_DTYPE=bfloat16
 export PER_DEVICE_BATCH_SIZE=11
+export PROMETHEUS_PORT=9090
 ```
 
 #### Create Llama2-13b environment variables for server flags
@@ -147,6 +149,7 @@ export ICI_TENSOR_PARALLELISM=1
 export SCAN_LAYERS=false
 export WEIGHT_DTYPE=bfloat16
 export PER_DEVICE_BATCH_SIZE=4
+export PROMETHEUS_PORT=9090
 ```
 
 ### Run the following command to start the JetStream MaxText server
@@ -165,7 +168,8 @@ python MaxText/maxengine_server.py \
   ici_tensor_parallelism=${ICI_TENSOR_PARALLELISM} \
   scan_layers=${SCAN_LAYERS} \
   weight_dtype=${WEIGHT_DTYPE} \
-  per_device_batch_size=${PER_DEVICE_BATCH_SIZE}
+  per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
+  prometheus_port=${PROMETHEUS_PORT}
 ```
 
 ### JetStream MaxText Server flag descriptions:
@@ -183,6 +187,7 @@ python MaxText/maxengine_server.py \
 *   ici\_tensor\_parallelism: The number of shards for tensor parallelism
 *   weight\_dtype: Weight data type (e.g. bfloat16)
 *   scan\_layers: Scan layers boolean flag (set to `false` for inference)
+*   prometheus\_port: create a prometheus endpoint emitting metrics on this port
 
 Note: these flags are from [MaxText config](https://github.com/google/maxtext/blob/f9e04cdc1eec74a0e648411857c09403c3358461/MaxText/configs/base.yml)
 
@@ -205,6 +210,19 @@ Prompt: Today is a good day
 Response:  to be a fan
 ```
 
+### (optional) Observe Jetstream metrics
+
+Since we configured `prometheus_port=9090` above, we can observe various Jetstream metrics via a HTTP requests to `0.0.0.0:9000`. Towards the end, the response should have content similar to the following:
+
+```
+# HELP jetstream_prefill_backlog_size Size of prefill queue
+# TYPE jetstream_prefill_backlog_size gauge
+jetstream_prefill_backlog_size{id="maxengine-server-58f8786f4c-4587n"} 0.0
+# HELP jetstream_slots_available_percentage The percentage of available slots in decode batch
+# TYPE jetstream_slots_available_percentage gauge
+jetstream_slots_available_percentage{id="maxengine-server-58f8786f4c-4587n",idx="0"} 0.96875
+```
+
 ## Step 6: Run benchmarks with JetStream MaxText server
 
 Note: The JetStream MaxText Server is not running with quantization optimization in Step 3. To get best benchmark results, we need to enable quantization (Please use AQT trained or fine tuned checkpoints to ensure accuracy) for both weights and KV cache, please add the quantization flags and restart the server as following:
@@ -231,6 +249,7 @@ ici_tensor_parallelism=${ICI_TENSOR_PARALLELISM} \
 scan_layers=${SCAN_LAYERS} \
 weight_dtype=${WEIGHT_DTYPE} \
 per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
+prometheus_port=${PROMETHEUS_PORT}
 quantization=${QUANTIZATION} \
 quantize_kvcache=${QUANTIZE_KVCACHE}
 ```
diff --git a/jetstream/core/orchestrator.py b/jetstream/core/orchestrator.py
index 51212ad0..c5af65b7 100644
--- a/jetstream/core/orchestrator.py
+++ b/jetstream/core/orchestrator.py
@@ -249,7 +249,7 @@ def __init__(
     # Stage 1
     # At first, a request is placed here in order to get prefilled.
     self._prefill_backlog = queue.Queue()
-    if self._metrics_collector is not None:
+    if self._metrics_collector:
       self._metrics_collector.get_prefill_backlog_metric().set_function(
           lambda: float(self._prefill_backlog.qsize())
       )
@@ -585,7 +585,7 @@ def _generate_thread(self, idx: int):
 
       max_concurrent_decodes = generate_engine.max_concurrent_decodes
 
-      if self._metrics_collector is not None:
+      if self._metrics_collector:
         self._metrics_collector.get_slots_available_percentage_metric(
             idx
         ).set_function(lambda: float(my_slots.qsize() / max_concurrent_decodes))
diff --git a/jetstream/core/server_lib.py b/jetstream/core/server_lib.py
index 9d4fc8dd..d50f3cc7 100644
--- a/jetstream/core/server_lib.py
+++ b/jetstream/core/server_lib.py
@@ -121,7 +121,7 @@ def run(
 
   # Setup Prometheus server
   metrics_collector: JetstreamMetricsCollector = None
-  if metrics_server_config is not None and metrics_server_config.port != 0:
+  if metrics_server_config and metrics_server_config.port:
     logging.info(
         "Starting Prometheus server on port %d", metrics_server_config.port
     )
diff --git a/jetstream/tests/core/test_server.py b/jetstream/tests/core/test_server.py
index 9dda982a..266fe7ab 100644
--- a/jetstream/tests/core/test_server.py
+++ b/jetstream/tests/core/test_server.py
@@ -18,9 +18,12 @@
 response.
 """
 
+import http
 from typing import Any, Type
 import unittest
 
+import requests
+
 from parameterized import parameterized
 import grpc
 from jetstream.core import config_lib
@@ -60,7 +63,10 @@ async def test_server(
     """Sets up a server and requests token responses."""
     ######################### Server side ######################################
     port = portpicker.pick_unused_port()
+    metrics_port = portpicker.pick_unused_port()
+
     print("port: " + str(port))
+    print("metrics port: " + str(metrics_port))
     credentials = grpc.local_server_credentials()
 
     server = server_lib.run(
@@ -68,6 +74,7 @@ async def test_server(
         config=config,
         devices=devices,
         credentials=credentials,
+        metrics_server_config=config_lib.MetricsServerConfig(port=metrics_port),
     )
     ###################### Requester side ######################################
     async with grpc.aio.secure_channel(
@@ -75,7 +82,7 @@ async def test_server(
     ) as channel:
       stub = jetstream_pb2_grpc.OrchestratorStub(channel)
 
-      # The string representation of np.array([[65, 66]]), [2] will be prependd
+      # The string representation of np.array([[65, 66]]), [2] will be prepended
       # as BOS
       text = "AB"
       request = jetstream_pb2.DecodeRequest(
@@ -96,6 +103,9 @@ async def test_server(
         counter += 1
       server.stop()
 
+    async with requests.get(f"localhost:{metrics_port}") as response:
+      assert response.status_code == requests.status_codes.codes["ok"]
+
 
 if __name__ == "__main__":
   unittest.main()

From 60482a28be8431378a652fdb8e8a137e1c48579e Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 20:06:15 +0000
Subject: [PATCH 57/74] reformat

---
 jetstream/tests/core/test_server.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/jetstream/tests/core/test_server.py b/jetstream/tests/core/test_server.py
index b3262140..56393634 100644
--- a/jetstream/tests/core/test_server.py
+++ b/jetstream/tests/core/test_server.py
@@ -103,10 +103,8 @@ async def test_server(
         counter += 1
       server.stop()
 
-
     async with requests.get(f"localhost:{metrics_port}") as response:
       assert response.status_code == requests.status_codes.codes["ok"]
 
-
   def test_get_devices(self):
     assert len(server_lib.get_devices()) == 1

From 5d1a04a3515bc58ceae24167472599a5d1d87794 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 20:27:26 +0000
Subject: [PATCH 58/74] missing protocol scheme in url

---
 jetstream/tests/core/test_server.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jetstream/tests/core/test_server.py b/jetstream/tests/core/test_server.py
index 56393634..23342ad3 100644
--- a/jetstream/tests/core/test_server.py
+++ b/jetstream/tests/core/test_server.py
@@ -103,7 +103,7 @@ async def test_server(
         counter += 1
       server.stop()
 
-    async with requests.get(f"localhost:{metrics_port}") as response:
+    async with requests.get(f"http://localhost:{metrics_port}") as response:
       assert response.status_code == requests.status_codes.codes["ok"]
 
   def test_get_devices(self):

From b904b32b8e3628cd3d1f091d8a1480196c2dc719 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 20:53:50 +0000
Subject: [PATCH 59/74] requests -> aiohttp

---
 jetstream/tests/core/test_server.py |   8 +-
 requirements.txt                    | 174 ++++++++++++++--------------
 2 files changed, 90 insertions(+), 92 deletions(-)

diff --git a/jetstream/tests/core/test_server.py b/jetstream/tests/core/test_server.py
index 23342ad3..9c8a7936 100644
--- a/jetstream/tests/core/test_server.py
+++ b/jetstream/tests/core/test_server.py
@@ -22,8 +22,7 @@
 from typing import Any, Type
 import unittest
 
-import requests
-
+import aiohttp
 from parameterized import parameterized
 import grpc
 from jetstream.core import config_lib
@@ -103,8 +102,9 @@ async def test_server(
         counter += 1
       server.stop()
 
-    async with requests.get(f"http://localhost:{metrics_port}") as response:
-      assert response.status_code == requests.status_codes.codes["ok"]
+    async with aiohttp.ClientSession() as session:
+        async with session.get(f"http://localhost:{metrics_port}") as response:
+            assert response.status == 200
 
   def test_get_devices(self):
     assert len(server_lib.get_devices()) == 1
diff --git a/requirements.txt b/requirements.txt
index ff60cf32..d8fddb5e 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,13 +4,14 @@
 #
 #    pip-compile requirements.in
 #
-absl-py==1.4.0
+absl-py==2.1.0
     # via
     #   -r requirements.in
     #   array-record
     #   chex
     #   clu
     #   etils
+    #   keras
     #   ml-collections
     #   optax
     #   orbax-checkpoint
@@ -19,80 +20,84 @@ absl-py==1.4.0
     #   tensorflow
     #   tensorflow-metadata
     #   tfds-nightly
-array-record==0.5.0
+aiohttp==3.9.5
+    # via -r requirements.in
+aiosignal==1.3.1
+    # via aiohttp
+array-record==0.5.1
     # via tfds-nightly
 astunparse==1.6.3
     # via tensorflow
+async-timeout==4.0.3
+    # via aiohttp
+attrs==23.2.0
+    # via aiohttp
 blobfile==2.1.1
     # via -r requirements.in
-cachetools==5.3.2
-    # via google-auth
 certifi==2024.2.2
     # via requests
 charset-normalizer==3.3.2
     # via requests
-chex==0.1.7
+chex==0.1.86
     # via optax
 click==8.1.7
     # via tfds-nightly
-clu==0.0.10
+clu==0.0.12
     # via seqio
 contextlib2==21.6.0
     # via ml-collections
-coverage==7.4.4
+coverage==7.5.1
     # via -r requirements.in
 dm-tree==0.1.8
-    # via
-    #   chex
-    #   tfds-nightly
-docstring-parser==0.15
+    # via tfds-nightly
+docstring-parser==0.16
     # via pyglove
-editdistance==0.6.2
+editdistance==0.8.1
     # via seqio
-etils[array-types,enp,epath,epy,etqdm,etree]==1.6.0
+etils[array-types,enp,epath,epy,etqdm,etree]==1.7.0
     # via
     #   array-record
     #   clu
     #   orbax-checkpoint
     #   tfds-nightly
-exceptiongroup==1.2.0
+exceptiongroup==1.2.1
     # via pytest
 filelock==3.14.0
     # via blobfile
-flatbuffers==23.5.26
+flatbuffers==24.3.25
     # via tensorflow
-flax==0.8.0
+flax==0.8.3
     # via
     #   -r requirements.in
     #   clu
-fsspec==2023.12.2
+frozenlist==1.4.1
+    # via
+    #   aiohttp
+    #   aiosignal
+fsspec==2024.5.0
     # via etils
-gast==0.4.0
+gast==0.5.4
     # via tensorflow
-google-auth==2.27.0
-    # via
-    #   google-auth-oauthlib
-    #   tensorboard
-google-auth-oauthlib==1.0.0
-    # via tensorboard
 google-pasta==0.2.0
     # via tensorflow
-googleapis-common-protos==1.62.0
-    # via tensorflow-metadata
-grpcio==1.60.1
+grpcio==1.63.0
     # via
     #   -r requirements.in
     #   tensorboard
     #   tensorflow
-h5py==3.10.0
-    # via tensorflow
+h5py==3.11.0
+    # via
+    #   keras
+    #   tensorflow
 idna==3.7
-    # via requests
-importlib-resources==6.1.1
+    # via
+    #   requests
+    #   yarl
+importlib-resources==6.4.0
     # via etils
 iniconfig==2.0.0
     # via pytest
-jax==0.4.23
+jax==0.4.28
     # via
     #   -r requirements.in
     #   chex
@@ -101,7 +106,7 @@ jax==0.4.23
     #   optax
     #   orbax-checkpoint
     #   seqio
-jaxlib==0.4.23
+jaxlib==0.4.28
     # via
     #   -r requirements.in
     #   chex
@@ -109,13 +114,13 @@ jaxlib==0.4.23
     #   optax
     #   orbax-checkpoint
     #   seqio
-keras==2.13.1
+keras==3.3.3
     # via tensorflow
-libclang==16.0.6
+libclang==18.1.1
     # via tensorflow
 lxml==4.9.4
     # via blobfile
-markdown==3.5.2
+markdown==3.6
     # via tensorboard
 markdown-it-py==3.0.0
     # via rich
@@ -129,14 +134,22 @@ ml-dtypes==0.3.2
     # via
     #   jax
     #   jaxlib
+    #   keras
+    #   tensorflow
     #   tensorstore
-msgpack==1.0.7
+msgpack==1.0.8
     # via
     #   flax
     #   orbax-checkpoint
+multidict==6.0.5
+    # via
+    #   aiohttp
+    #   yarl
+namex==0.0.8
+    # via keras
 nest-asyncio==1.6.0
     # via orbax-checkpoint
-numpy==1.23.1
+numpy==1.26.4
     # via
     #   -r requirements.in
     #   chex
@@ -146,6 +159,7 @@ numpy==1.23.1
     #   h5py
     #   jax
     #   jaxlib
+    #   keras
     #   ml-dtypes
     #   opt-einsum
     #   optax
@@ -154,20 +168,19 @@ numpy==1.23.1
     #   seqio
     #   tensorboard
     #   tensorflow
-    #   tensorflow-hub
     #   tensorstore
     #   tfds-nightly
-oauthlib==3.2.2
-    # via requests-oauthlib
 opt-einsum==3.3.0
     # via
     #   jax
     #   tensorflow
-optax==0.1.8
+optax==0.2.2
     # via flax
-orbax-checkpoint==0.5.2
+optree==0.11.0
+    # via keras
+orbax-checkpoint==0.5.11
     # via flax
-packaging==23.2
+packaging==24.0
     # via
     #   clu
     #   pytest
@@ -175,7 +188,7 @@ packaging==23.2
     #   tensorflow
 parameterized==0.9.0
     # via -r requirements.in
-pluggy==1.4.0
+pluggy==1.5.0
     # via pytest
 portpicker==1.6.0
     # via -r requirements.in
@@ -185,56 +198,45 @@ promise==2.3
     # via tfds-nightly
 protobuf==3.20.3
     # via
-    #   googleapis-common-protos
     #   orbax-checkpoint
     #   seqio
     #   tensorboard
     #   tensorflow
-    #   tensorflow-hub
     #   tensorflow-metadata
     #   tfds-nightly
 psutil==5.9.8
     # via
     #   portpicker
     #   tfds-nightly
-pyasn1==0.5.1
-    # via
-    #   pyasn1-modules
-    #   rsa
-pyasn1-modules==0.3.0
-    # via google-auth
 pycryptodomex==3.20.0
     # via blobfile
 pyglove==0.4.4
     # via seqio
-pygments==2.17.2
+pygments==2.18.0
     # via rich
-pytest==8.1.1
+pytest==8.2.0
     # via -r requirements.in
 pyyaml==6.0.1
     # via
     #   flax
     #   ml-collections
     #   orbax-checkpoint
-regex==2024.4.28
+regex==2024.5.15
     # via tiktoken
 requests==2.31.0
     # via
-    #   requests-oauthlib
-    #   tensorboard
+    #   tensorflow
     #   tfds-nightly
     #   tiktoken
-requests-oauthlib==1.3.1
-    # via google-auth-oauthlib
-rich==13.7.0
-    # via flax
-rsa==4.9
-    # via google-auth
-scipy==1.12.0
+rich==13.7.1
+    # via
+    #   flax
+    #   keras
+scipy==1.13.0
     # via
     #   jax
     #   jaxlib
-sentencepiece==0.1.99
+sentencepiece==0.2.0
     # via seqio
 seqio==0.0.19
     # via -r requirements.in
@@ -246,24 +248,21 @@ six==1.16.0
     #   google-pasta
     #   ml-collections
     #   promise
+    #   tensorboard
     #   tensorflow
-tensorboard==2.13.0
+tensorboard==2.16.2
     # via tensorflow
 tensorboard-data-server==0.7.2
     # via tensorboard
-tensorflow==2.13.1
-    # via tensorflow-text
-tensorflow-estimator==2.13.0
-    # via tensorflow
-tensorflow-hub==0.16.1
+tensorflow==2.16.1
     # via tensorflow-text
-tensorflow-io-gcs-filesystem==0.35.0
+tensorflow-io-gcs-filesystem==0.37.0
     # via tensorflow
-tensorflow-metadata==1.14.0
+tensorflow-metadata==1.15.0
     # via tfds-nightly
-tensorflow-text==2.13.0
+tensorflow-text==2.16.1
     # via seqio
-tensorstore==0.1.52
+tensorstore==0.1.59
     # via
     #   flax
     #   orbax-checkpoint
@@ -271,11 +270,9 @@ termcolor==2.4.0
     # via
     #   tensorflow
     #   tfds-nightly
-tf-keras==2.15.0
-    # via tensorflow-hub
 tfds-nightly==4.9.2.dev202308090034
     # via seqio
-tiktoken==0.6.0
+tiktoken==0.7.0
     # via -r requirements.in
 toml==0.10.2
     # via tfds-nightly
@@ -283,34 +280,35 @@ tomli==2.0.1
     # via pytest
 toolz==0.12.1
     # via chex
-tqdm==4.66.3
+tqdm==4.66.4
     # via
     #   etils
     #   tfds-nightly
-typing-extensions==4.5.0
+typing-extensions==4.11.0
     # via
     #   chex
     #   clu
     #   etils
     #   flax
+    #   optree
     #   orbax-checkpoint
     #   tensorflow
-urllib3==2.2.0
+urllib3==2.2.1
     # via
     #   blobfile
     #   requests
-werkzeug==3.0.1
+werkzeug==3.0.3
     # via tensorboard
-wheel==0.42.0
-    # via
-    #   astunparse
-    #   tensorboard
+wheel==0.43.0
+    # via astunparse
 wrapt==1.16.0
     # via
     #   clu
     #   tensorflow
     #   tfds-nightly
-zipp==3.17.0
+yarl==1.9.4
+    # via aiohttp
+zipp==3.18.2
     # via etils
 
 # The following packages are considered to be unsafe in a requirements file:

From ae37a0c3780ac29742640db90413397705b1e94e Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 20:57:58 +0000
Subject: [PATCH 60/74] linter

---
 jetstream/tests/core/test_server.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/jetstream/tests/core/test_server.py b/jetstream/tests/core/test_server.py
index 9c8a7936..ddf1e927 100644
--- a/jetstream/tests/core/test_server.py
+++ b/jetstream/tests/core/test_server.py
@@ -18,7 +18,6 @@
 response.
 """
 
-import http
 from typing import Any, Type
 import unittest
 
@@ -103,8 +102,8 @@ async def test_server(
       server.stop()
 
     async with aiohttp.ClientSession() as session:
-        async with session.get(f"http://localhost:{metrics_port}") as response:
-            assert response.status == 200
+      async with session.get(f"http://localhost:{metrics_port}") as response:
+        assert response.status == 200
 
   def test_get_devices(self):
     assert len(server_lib.get_devices()) == 1

From a679d455e3521bd46ff346d391aa8e14c172d602 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 21:03:54 +0000
Subject: [PATCH 61/74] Cleaned readme

---
 docs/online-inference-with-maxtext-engine.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/online-inference-with-maxtext-engine.md b/docs/online-inference-with-maxtext-engine.md
index 8e612e40..c5bdce94 100644
--- a/docs/online-inference-with-maxtext-engine.md
+++ b/docs/online-inference-with-maxtext-engine.md
@@ -217,10 +217,10 @@ Since we configured `prometheus_port=9090` above, we can observe various Jetstre
 ```
 # HELP jetstream_prefill_backlog_size Size of prefill queue
 # TYPE jetstream_prefill_backlog_size gauge
-jetstream_prefill_backlog_size{id="maxengine-server-58f8786f4c-4587n"} 0.0
+jetstream_prefill_backlog_size{id="SOME-HOSTNAME-HERE>"} 0.0
 # HELP jetstream_slots_available_percentage The percentage of available slots in decode batch
 # TYPE jetstream_slots_available_percentage gauge
-jetstream_slots_available_percentage{id="maxengine-server-58f8786f4c-4587n",idx="0"} 0.96875
+jetstream_slots_available_percentage{id="<SOME-HOSTNAME-HERE>",idx="0"} 0.96875
 ```
 
 ## Step 6: Run benchmarks with JetStream MaxText server

From 0b5c058953172217e5a25121ca384530ff8b2a7e Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 22:22:43 +0000
Subject: [PATCH 62/74] Changes to tests, readme

---
 docs/online-inference-with-maxtext-engine.md | 32 +++++++++++++++-----
 jetstream/tests/core/test_server.py          | 32 ++++++++++++++++++--
 2 files changed, 53 insertions(+), 11 deletions(-)

diff --git a/docs/online-inference-with-maxtext-engine.md b/docs/online-inference-with-maxtext-engine.md
index c5bdce94..13b3e702 100644
--- a/docs/online-inference-with-maxtext-engine.md
+++ b/docs/online-inference-with-maxtext-engine.md
@@ -109,7 +109,6 @@ export ICI_TENSOR_PARALLELISM=1
 export SCAN_LAYERS=false
 export WEIGHT_DTYPE=bfloat16
 export PER_DEVICE_BATCH_SIZE=11
-export PROMETHEUS_PORT=9090
 ```
 
 #### Create Llama2-7b environment variables for server flags
@@ -128,7 +127,6 @@ export ICI_TENSOR_PARALLELISM=1
 export SCAN_LAYERS=false
 export WEIGHT_DTYPE=bfloat16
 export PER_DEVICE_BATCH_SIZE=11
-export PROMETHEUS_PORT=9090
 ```
 
 #### Create Llama2-13b environment variables for server flags
@@ -149,7 +147,6 @@ export ICI_TENSOR_PARALLELISM=1
 export SCAN_LAYERS=false
 export WEIGHT_DTYPE=bfloat16
 export PER_DEVICE_BATCH_SIZE=4
-export PROMETHEUS_PORT=9090
 ```
 
 ### Run the following command to start the JetStream MaxText server
@@ -168,8 +165,7 @@ python MaxText/maxengine_server.py \
   ici_tensor_parallelism=${ICI_TENSOR_PARALLELISM} \
   scan_layers=${SCAN_LAYERS} \
   weight_dtype=${WEIGHT_DTYPE} \
-  per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
-  prometheus_port=${PROMETHEUS_PORT}
+  per_device_batch_size=${PER_DEVICE_BATCH_SIZE}
 ```
 
 ### JetStream MaxText Server flag descriptions:
@@ -187,7 +183,6 @@ python MaxText/maxengine_server.py \
 *   ici\_tensor\_parallelism: The number of shards for tensor parallelism
 *   weight\_dtype: Weight data type (e.g. bfloat16)
 *   scan\_layers: Scan layers boolean flag (set to `false` for inference)
-*   prometheus\_port: create a prometheus endpoint emitting metrics on this port
 
 Note: these flags are from [MaxText config](https://github.com/google/maxtext/blob/f9e04cdc1eec74a0e648411857c09403c3358461/MaxText/configs/base.yml)
 
@@ -212,7 +207,29 @@ Response:  to be a fan
 
 ### (optional) Observe Jetstream metrics
 
-Since we configured `prometheus_port=9090` above, we can observe various Jetstream metrics via a HTTP requests to `0.0.0.0:9000`. Towards the end, the response should have content similar to the following:
+Metrics are not exported by default, to configure Jetstream to emit metrics start this guide again from step four and replace the `Run the following command to start the JetStream MaxText server` step with the following:
+
+```bash
+export PROMETHEUS_PORT=9090
+
+cd ~/maxtext
+python MaxText/maxengine_server.py \
+  MaxText/configs/base.yml \
+  tokenizer_path=${TOKENIZER_PATH} \
+  load_parameters_path=${LOAD_PARAMETERS_PATH} \
+  max_prefill_predict_length=${MAX_PREFILL_PREDICT_LENGTH} \
+  max_target_length=${MAX_TARGET_LENGTH} \
+  model_name=${MODEL_NAME} \
+  ici_fsdp_parallelism=${ICI_FSDP_PARALLELISM} \
+  ici_autoregressive_parallelism=${ICI_AUTOREGRESSIVE_PARALLELISM} \
+  ici_tensor_parallelism=${ICI_TENSOR_PARALLELISM} \
+  scan_layers=${SCAN_LAYERS} \
+  weight_dtype=${WEIGHT_DTYPE} \
+  per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
+  prometheus_port=${PROMETHEUS_PORT}
+```
+
+Now that we configured `prometheus_port=9090` above, we can observe various Jetstream metrics via a HTTP requests to `0.0.0.0:9000`. Towards the end, the response should have content similar to the following:
 
 ```
 # HELP jetstream_prefill_backlog_size Size of prefill queue
@@ -249,7 +266,6 @@ ici_tensor_parallelism=${ICI_TENSOR_PARALLELISM} \
 scan_layers=${SCAN_LAYERS} \
 weight_dtype=${WEIGHT_DTYPE} \
 per_device_batch_size=${PER_DEVICE_BATCH_SIZE} \
-prometheus_port=${PROMETHEUS_PORT}
 quantization=${QUANTIZATION} \
 quantize_kvcache=${QUANTIZE_KVCACHE}
 ```
diff --git a/jetstream/tests/core/test_server.py b/jetstream/tests/core/test_server.py
index ddf1e927..37a86fba 100644
--- a/jetstream/tests/core/test_server.py
+++ b/jetstream/tests/core/test_server.py
@@ -61,10 +61,8 @@ async def test_server(
     """Sets up a server and requests token responses."""
     ######################### Server side ######################################
     port = portpicker.pick_unused_port()
-    metrics_port = portpicker.pick_unused_port()
 
     print("port: " + str(port))
-    print("metrics port: " + str(metrics_port))
     credentials = grpc.local_server_credentials()
 
     server = server_lib.run(
@@ -72,7 +70,6 @@ async def test_server(
         config=config,
         devices=devices,
         credentials=credentials,
-        metrics_server_config=config_lib.MetricsServerConfig(port=metrics_port),
     )
     ###################### Requester side ######################################
     async with grpc.aio.secure_channel(
@@ -101,6 +98,35 @@ async def test_server(
         counter += 1
       server.stop()
 
+    # prometheus not configured, assert no metrics collector on Driver
+    assert server._driver._metrics_collector is not None
+
+  async def test_prometheus_server(
+      self,
+      config: Type[config_lib.ServerConfig],
+      expected_text: list[str],
+      expected_token_ids: list[int | None],
+      devices: list[Any],
+  ):
+    """Sets up a server and requests token responses."""
+    ######################### Server side ######################################
+    port = portpicker.pick_unused_port()
+    metrics_port = portpicker.pick_unused_port()
+
+    print("port: " + str(port))
+    print("metrics port: " + str(metrics_port))
+    credentials = grpc.local_server_credentials()
+
+    server = server_lib.run(
+        port=port,
+        config=config,
+        devices=devices,
+        credentials=credentials,
+        metrics_server_config=config_lib.MetricsServerConfig(port=metrics_port),
+    )
+    ###################### Requester side ######################################
+    # assert prometheus server is running and responding
+    assert server._driver._metrics_collector is not None
     async with aiohttp.ClientSession() as session:
       async with session.get(f"http://localhost:{metrics_port}") as response:
         assert response.status == 200

From 37c8fcb6e62f0d4a99a2843c207ef96862cd13b2 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 22:25:46 +0000
Subject: [PATCH 63/74] fixed assertion

---
 jetstream/tests/core/test_server.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jetstream/tests/core/test_server.py b/jetstream/tests/core/test_server.py
index 37a86fba..7f4986c7 100644
--- a/jetstream/tests/core/test_server.py
+++ b/jetstream/tests/core/test_server.py
@@ -99,7 +99,7 @@ async def test_server(
       server.stop()
 
     # prometheus not configured, assert no metrics collector on Driver
-    assert server._driver._metrics_collector is not None
+    assert server._driver._metrics_collector is None
 
   async def test_prometheus_server(
       self,

From e2166a9091bed29e816216f48bc3244a87aa4b82 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 22:27:21 +0000
Subject: [PATCH 64/74] better description

---
 jetstream/tests/core/test_server.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/jetstream/tests/core/test_server.py b/jetstream/tests/core/test_server.py
index 7f4986c7..cd9400ed 100644
--- a/jetstream/tests/core/test_server.py
+++ b/jetstream/tests/core/test_server.py
@@ -108,7 +108,7 @@ async def test_prometheus_server(
       expected_token_ids: list[int | None],
       devices: list[Any],
   ):
-    """Sets up a server and requests token responses."""
+    """Same as prior test_server but with metrics server config"""
     ######################### Server side ######################################
     port = portpicker.pick_unused_port()
     metrics_port = portpicker.pick_unused_port()

From abd55550250b606ccfba1eec0d7eb3db0927efb0 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 22:38:06 +0000
Subject: [PATCH 65/74] parameterized test setup

---
 jetstream/tests/core/test_server.py | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/jetstream/tests/core/test_server.py b/jetstream/tests/core/test_server.py
index cd9400ed..d5252a41 100644
--- a/jetstream/tests/core/test_server.py
+++ b/jetstream/tests/core/test_server.py
@@ -101,6 +101,16 @@ async def test_server(
     # prometheus not configured, assert no metrics collector on Driver
     assert server._driver._metrics_collector is None
 
+  @parameterized.expand(
+      [
+          (
+              config_lib.InterleavedCPUTestServer,
+              [],
+              [],
+              [None],
+          ),
+      ]
+  )
   async def test_prometheus_server(
       self,
       config: Type[config_lib.ServerConfig],

From 8170a982d1014eeedc833bffc5e5954a5f98b22b Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 22:42:43 +0000
Subject: [PATCH 66/74] linter

---
 jetstream/tests/core/test_server.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/jetstream/tests/core/test_server.py b/jetstream/tests/core/test_server.py
index d5252a41..3b77b086 100644
--- a/jetstream/tests/core/test_server.py
+++ b/jetstream/tests/core/test_server.py
@@ -105,8 +105,6 @@ async def test_server(
       [
           (
               config_lib.InterleavedCPUTestServer,
-              [],
-              [],
               [None],
           ),
       ]
@@ -114,8 +112,6 @@ async def test_server(
   async def test_prometheus_server(
       self,
       config: Type[config_lib.ServerConfig],
-      expected_text: list[str],
-      expected_token_ids: list[int | None],
       devices: list[Any],
   ):
     """Same as prior test_server but with metrics server config"""

From b3aff259b61779bad7c2dbecc6a4dcbd1050c223 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <bslabe123@gmail.com>
Date: Fri, 17 May 2024 15:44:49 -0700
Subject: [PATCH 67/74] Update online-inference-with-maxtext-engine.md

---
 docs/online-inference-with-maxtext-engine.md | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/docs/online-inference-with-maxtext-engine.md b/docs/online-inference-with-maxtext-engine.md
index 13b3e702..5c6aef5d 100644
--- a/docs/online-inference-with-maxtext-engine.md
+++ b/docs/online-inference-with-maxtext-engine.md
@@ -229,7 +229,7 @@ python MaxText/maxengine_server.py \
   prometheus_port=${PROMETHEUS_PORT}
 ```
 
-Now that we configured `prometheus_port=9090` above, we can observe various Jetstream metrics via a HTTP requests to `0.0.0.0:9000`. Towards the end, the response should have content similar to the following:
+Now that we configured `prometheus_port=9090` above, we can observe various Jetstream metrics via HTTP requests to `0.0.0.0:9000`. Towards the end, the response should have content similar to the following:
 
 ```
 # HELP jetstream_prefill_backlog_size Size of prefill queue
@@ -324,4 +324,4 @@ rm -rf maxtext
 rm -rf JetStream
 # Clean up python virtual environment
 rm -rf .env
-```
\ No newline at end of file
+```

From 35e5a9a84578d92a79949870ffad451c8b1e82ee Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 22:56:07 +0000
Subject: [PATCH 68/74] disable protected-access

---
 jetstream/tests/core/test_server.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/jetstream/tests/core/test_server.py b/jetstream/tests/core/test_server.py
index 3b77b086..978daf3b 100644
--- a/jetstream/tests/core/test_server.py
+++ b/jetstream/tests/core/test_server.py
@@ -99,7 +99,7 @@ async def test_server(
       server.stop()
 
     # prometheus not configured, assert no metrics collector on Driver
-    assert server._driver._metrics_collector is None
+    assert server._driver._metrics_collector is None  # pylint: disable=protected-access
 
   @parameterized.expand(
       [
@@ -132,7 +132,7 @@ async def test_prometheus_server(
     )
     ###################### Requester side ######################################
     # assert prometheus server is running and responding
-    assert server._driver._metrics_collector is not None
+    assert server._driver._metrics_collector is not None  # pylint: disable=protected-access
     async with aiohttp.ClientSession() as session:
       async with session.get(f"http://localhost:{metrics_port}") as response:
         assert response.status == 200

From 332924107600eb25e9291ff2d40b2a40f0a38c95 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 23:43:55 +0000
Subject: [PATCH 69/74] make prometheus test not async, revert requirements.txt
 changes

---
 jetstream/tests/core/test_server.py |   9 +-
 requirements.txt                    | 180 ++++++++++++++--------------
 2 files changed, 96 insertions(+), 93 deletions(-)

diff --git a/jetstream/tests/core/test_server.py b/jetstream/tests/core/test_server.py
index 978daf3b..c39fb02e 100644
--- a/jetstream/tests/core/test_server.py
+++ b/jetstream/tests/core/test_server.py
@@ -21,6 +21,8 @@
 from typing import Any, Type
 import unittest
 
+
+import requests
 import aiohttp
 from parameterized import parameterized
 import grpc
@@ -133,9 +135,10 @@ async def test_prometheus_server(
     ###################### Requester side ######################################
     # assert prometheus server is running and responding
     assert server._driver._metrics_collector is not None  # pylint: disable=protected-access
-    async with aiohttp.ClientSession() as session:
-      async with session.get(f"http://localhost:{metrics_port}") as response:
-        assert response.status == 200
+    assert (
+        requests.get(f"http://localhost:{metrics_port}").status_code
+        == requests.status_codes.codes["ok"]
+    )
 
   def test_get_devices(self):
     assert len(server_lib.get_devices()) == 1
diff --git a/requirements.txt b/requirements.txt
index d8fddb5e..483edea0 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -4,14 +4,13 @@
 #
 #    pip-compile requirements.in
 #
-absl-py==2.1.0
+absl-py==1.4.0
     # via
     #   -r requirements.in
     #   array-record
     #   chex
     #   clu
     #   etils
-    #   keras
     #   ml-collections
     #   optax
     #   orbax-checkpoint
@@ -20,84 +19,80 @@ absl-py==2.1.0
     #   tensorflow
     #   tensorflow-metadata
     #   tfds-nightly
-aiohttp==3.9.5
-    # via -r requirements.in
-aiosignal==1.3.1
-    # via aiohttp
-array-record==0.5.1
+array-record==0.5.0
     # via tfds-nightly
 astunparse==1.6.3
     # via tensorflow
-async-timeout==4.0.3
-    # via aiohttp
-attrs==23.2.0
-    # via aiohttp
 blobfile==2.1.1
     # via -r requirements.in
+cachetools==5.3.2
+    # via google-auth
 certifi==2024.2.2
     # via requests
 charset-normalizer==3.3.2
     # via requests
-chex==0.1.86
+chex==0.1.7
     # via optax
 click==8.1.7
     # via tfds-nightly
-clu==0.0.12
+clu==0.0.10
     # via seqio
 contextlib2==21.6.0
     # via ml-collections
-coverage==7.5.1
+coverage==7.4.4
     # via -r requirements.in
 dm-tree==0.1.8
-    # via tfds-nightly
-docstring-parser==0.16
+    # via
+    #   chex
+    #   tfds-nightly
+docstring-parser==0.15
     # via pyglove
-editdistance==0.8.1
+editdistance==0.6.2
     # via seqio
-etils[array-types,enp,epath,epy,etqdm,etree]==1.7.0
+etils[array-types,enp,epath,epy,etqdm,etree]==1.6.0
     # via
     #   array-record
     #   clu
     #   orbax-checkpoint
     #   tfds-nightly
-exceptiongroup==1.2.1
+exceptiongroup==1.2.0
     # via pytest
 filelock==3.14.0
     # via blobfile
-flatbuffers==24.3.25
+flatbuffers==23.5.26
     # via tensorflow
-flax==0.8.3
+flax==0.8.0
     # via
     #   -r requirements.in
     #   clu
-frozenlist==1.4.1
-    # via
-    #   aiohttp
-    #   aiosignal
-fsspec==2024.5.0
+fsspec==2023.12.2
     # via etils
-gast==0.5.4
+gast==0.4.0
     # via tensorflow
+google-auth==2.27.0
+    # via
+    #   google-auth-oauthlib
+    #   tensorboard
+google-auth-oauthlib==1.0.0
+    # via tensorboard
 google-pasta==0.2.0
     # via tensorflow
-grpcio==1.63.0
+googleapis-common-protos==1.62.0
+    # via tensorflow-metadata
+grpcio==1.60.1
     # via
     #   -r requirements.in
     #   tensorboard
     #   tensorflow
-h5py==3.11.0
-    # via
-    #   keras
-    #   tensorflow
+h5py==3.10.0
+    # via tensorflow
 idna==3.7
-    # via
-    #   requests
-    #   yarl
-importlib-resources==6.4.0
+    # via requests
+importlib-resources==6.1.1
     # via etils
 iniconfig==2.0.0
     # via pytest
-jax==0.4.28
+jax==0.4.23
     # via
     #   -r requirements.in
     #   chex
@@ -106,7 +101,7 @@ jax==0.4.28
     #   optax
     #   orbax-checkpoint
     #   seqio
-jaxlib==0.4.28
+jaxlib==0.4.23
     # via
     #   -r requirements.in
     #   chex
@@ -114,13 +109,13 @@ jaxlib==0.4.28
     #   optax
     #   orbax-checkpoint
     #   seqio
-keras==3.3.3
+keras==2.13.1
     # via tensorflow
-libclang==18.1.1
+libclang==16.0.6
     # via tensorflow
 lxml==4.9.4
     # via blobfile
-markdown==3.6
+markdown==3.5.2
     # via tensorboard
 markdown-it-py==3.0.0
     # via rich
@@ -134,22 +129,14 @@ ml-dtypes==0.3.2
     # via
     #   jax
     #   jaxlib
-    #   keras
-    #   tensorflow
     #   tensorstore
-msgpack==1.0.8
+msgpack==1.0.7
     # via
     #   flax
     #   orbax-checkpoint
-multidict==6.0.5
-    # via
-    #   aiohttp
-    #   yarl
-namex==0.0.8
-    # via keras
 nest-asyncio==1.6.0
     # via orbax-checkpoint
-numpy==1.26.4
+numpy==1.23.1
     # via
     #   -r requirements.in
     #   chex
@@ -159,7 +146,6 @@ numpy==1.26.4
     #   h5py
     #   jax
     #   jaxlib
-    #   keras
     #   ml-dtypes
     #   opt-einsum
     #   optax
@@ -168,19 +154,20 @@ numpy==1.26.4
     #   seqio
     #   tensorboard
     #   tensorflow
+    #   tensorflow-hub
     #   tensorstore
     #   tfds-nightly
+oauthlib==3.2.2
+    # via requests-oauthlib
 opt-einsum==3.3.0
     # via
     #   jax
     #   tensorflow
-optax==0.2.2
+optax==0.1.8
     # via flax
-optree==0.11.0
-    # via keras
-orbax-checkpoint==0.5.11
+orbax-checkpoint==0.5.2
     # via flax
-packaging==24.0
+packaging==23.2
     # via
     #   clu
     #   pytest
@@ -188,55 +175,64 @@ packaging==24.0
     #   tensorflow
 parameterized==0.9.0
     # via -r requirements.in
-pluggy==1.5.0
+pluggy==1.4.0
     # via pytest
 portpicker==1.6.0
     # via -r requirements.in
-prometheus-client==0.20.0
-    # via -r requirements.in
-promise==2.3
+	@@ -185,56 +198,45 @@ promise==2.3
     # via tfds-nightly
 protobuf==3.20.3
     # via
+    #   googleapis-common-protos
     #   orbax-checkpoint
     #   seqio
     #   tensorboard
     #   tensorflow
+    #   tensorflow-hub
     #   tensorflow-metadata
     #   tfds-nightly
 psutil==5.9.8
     # via
     #   portpicker
     #   tfds-nightly
+pyasn1==0.5.1
+    # via
+    #   pyasn1-modules
+    #   rsa
+pyasn1-modules==0.3.0
+    # via google-auth
 pycryptodomex==3.20.0
     # via blobfile
 pyglove==0.4.4
     # via seqio
-pygments==2.18.0
+pygments==2.17.2
     # via rich
-pytest==8.2.0
+pytest==8.1.1
     # via -r requirements.in
 pyyaml==6.0.1
     # via
     #   flax
     #   ml-collections
     #   orbax-checkpoint
-regex==2024.5.15
+regex==2024.4.28
     # via tiktoken
 requests==2.31.0
     # via
-    #   tensorflow
+    #   requests-oauthlib
+    #   tensorboard
     #   tfds-nightly
     #   tiktoken
-rich==13.7.1
-    # via
-    #   flax
-    #   keras
-scipy==1.13.0
+requests-oauthlib==1.3.1
+    # via google-auth-oauthlib
+rich==13.7.0
+    # via flax
+rsa==4.9
+    # via google-auth
+scipy==1.12.0
     # via
     #   jax
     #   jaxlib
-sentencepiece==0.2.0
+sentencepiece==0.1.99
     # via seqio
 seqio==0.0.19
     # via -r requirements.in
@@ -248,21 +244,24 @@ six==1.16.0
     #   google-pasta
     #   ml-collections
     #   promise
-    #   tensorboard
     #   tensorflow
-tensorboard==2.16.2
+tensorboard==2.13.0
     # via tensorflow
 tensorboard-data-server==0.7.2
     # via tensorboard
-tensorflow==2.16.1
+tensorflow==2.13.1
+    # via tensorflow-text
+tensorflow-estimator==2.13.0
+    # via tensorflow
+tensorflow-hub==0.16.1
     # via tensorflow-text
-tensorflow-io-gcs-filesystem==0.37.0
+tensorflow-io-gcs-filesystem==0.35.0
     # via tensorflow
-tensorflow-metadata==1.15.0
+tensorflow-metadata==1.14.0
     # via tfds-nightly
-tensorflow-text==2.16.1
+tensorflow-text==2.13.0
     # via seqio
-tensorstore==0.1.59
+tensorstore==0.1.52
     # via
     #   flax
     #   orbax-checkpoint
@@ -270,9 +269,11 @@ termcolor==2.4.0
     # via
     #   tensorflow
     #   tfds-nightly
+tf-keras==2.15.0
+    # via tensorflow-hub
 tfds-nightly==4.9.2.dev202308090034
     # via seqio
-tiktoken==0.7.0
+tiktoken==0.6.0
     # via -r requirements.in
 toml==0.10.2
     # via tfds-nightly
@@ -280,36 +281,35 @@ tomli==2.0.1
     # via pytest
 toolz==0.12.1
     # via chex
-tqdm==4.66.4
+tqdm==4.66.3
     # via
     #   etils
     #   tfds-nightly
-typing-extensions==4.11.0
+typing-extensions==4.5.0
     # via
     #   chex
     #   clu
     #   etils
     #   flax
-    #   optree
     #   orbax-checkpoint
     #   tensorflow
-urllib3==2.2.1
+urllib3==2.2.0
     # via
     #   blobfile
     #   requests
-werkzeug==3.0.3
+werkzeug==3.0.1
     # via tensorboard
-wheel==0.43.0
-    # via astunparse
+wheel==0.42.0
+    # via
+    #   astunparse
+    #   tensorboard
 wrapt==1.16.0
     # via
     #   clu
     #   tensorflow
     #   tfds-nightly
-yarl==1.9.4
-    # via aiohttp
-zipp==3.18.2
+zipp==3.17.0
     # via etils
 
 # The following packages are considered to be unsafe in a requirements file:
-# setuptools
+# setuptools
\ No newline at end of file

From d02ec13c7b5652cf24f7a51258808359d2ac2c0a Mon Sep 17 00:00:00 2001
From: Brendan Slabe <bslabe123@gmail.com>
Date: Fri, 17 May 2024 16:44:32 -0700
Subject: [PATCH 70/74] Update requirements.txt

---
 requirements.txt | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/requirements.txt b/requirements.txt
index 483edea0..ff60cf32 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -179,7 +179,9 @@ pluggy==1.4.0
     # via pytest
 portpicker==1.6.0
     # via -r requirements.in
-	@@ -185,56 +198,45 @@ promise==2.3
+prometheus-client==0.20.0
+    # via -r requirements.in
+promise==2.3
     # via tfds-nightly
 protobuf==3.20.3
     # via
@@ -312,4 +314,4 @@ zipp==3.17.0
     # via etils
 
 # The following packages are considered to be unsafe in a requirements file:
-# setuptools
\ No newline at end of file
+# setuptools

From aeaca8581282f9c7e34261a47e924a7d0a730127 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 23:46:40 +0000
Subject: [PATCH 71/74] remove aiohttp

---
 jetstream/tests/core/test_server.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/jetstream/tests/core/test_server.py b/jetstream/tests/core/test_server.py
index c39fb02e..5e6b9e7a 100644
--- a/jetstream/tests/core/test_server.py
+++ b/jetstream/tests/core/test_server.py
@@ -23,7 +23,6 @@
 
 
 import requests
-import aiohttp
 from parameterized import parameterized
 import grpc
 from jetstream.core import config_lib

From 532efe257968565e2865f7f690e3346605bcda97 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 23:55:31 +0000
Subject: [PATCH 72/74] moved prometheus test to test_server function

---
 jetstream/tests/core/test_server.py | 59 ++++++++++-------------------
 1 file changed, 21 insertions(+), 38 deletions(-)

diff --git a/jetstream/tests/core/test_server.py b/jetstream/tests/core/test_server.py
index 5e6b9e7a..d8972a7e 100644
--- a/jetstream/tests/core/test_server.py
+++ b/jetstream/tests/core/test_server.py
@@ -62,6 +62,7 @@ async def test_server(
     """Sets up a server and requests token responses."""
     ######################### Server side ######################################
     port = portpicker.pick_unused_port()
+    metrics_port = portpicker.pick_unused_port()
 
     print("port: " + str(port))
     credentials = grpc.local_server_credentials()
@@ -73,6 +74,10 @@ async def test_server(
         credentials=credentials,
     )
     ###################### Requester side ######################################
+
+    # prometheus not configured, assert no metrics collector on Driver
+    assert server._driver._metrics_collector is None  # pylint: disable=protected-access
+
     async with grpc.aio.secure_channel(
         f"localhost:{port}", grpc.local_channel_credentials()
     ) as channel:
@@ -99,45 +104,23 @@ async def test_server(
         counter += 1
       server.stop()
 
-    # prometheus not configured, assert no metrics collector on Driver
-    assert server._driver._metrics_collector is None  # pylint: disable=protected-access
-
-  @parameterized.expand(
-      [
-          (
-              config_lib.InterleavedCPUTestServer,
-              [None],
+      # Now test server with prometheus config
+      server = server_lib.run(
+          port=port,
+          config=config,
+          devices=devices,
+          credentials=credentials,
+          metrics_server_config=config_lib.MetricsServerConfig(
+              port=metrics_port
           ),
-      ]
-  )
-  async def test_prometheus_server(
-      self,
-      config: Type[config_lib.ServerConfig],
-      devices: list[Any],
-  ):
-    """Same as prior test_server but with metrics server config"""
-    ######################### Server side ######################################
-    port = portpicker.pick_unused_port()
-    metrics_port = portpicker.pick_unused_port()
-
-    print("port: " + str(port))
-    print("metrics port: " + str(metrics_port))
-    credentials = grpc.local_server_credentials()
-
-    server = server_lib.run(
-        port=port,
-        config=config,
-        devices=devices,
-        credentials=credentials,
-        metrics_server_config=config_lib.MetricsServerConfig(port=metrics_port),
-    )
-    ###################### Requester side ######################################
-    # assert prometheus server is running and responding
-    assert server._driver._metrics_collector is not None  # pylint: disable=protected-access
-    assert (
-        requests.get(f"http://localhost:{metrics_port}").status_code
-        == requests.status_codes.codes["ok"]
-    )
+      )
+      # assert prometheus server is running and responding
+      assert server._driver._metrics_collector is not None  # pylint: disable=protected-access
+      assert (
+          requests.get(f"http://localhost:{metrics_port}").status_code
+          == requests.status_codes.codes["ok"]
+      )
+      server.stop()
 
   def test_get_devices(self):
     assert len(server_lib.get_devices()) == 1

From 3aba8ed1929332b9dae715474b828e26fc43b666 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Fri, 17 May 2024 23:58:51 +0000
Subject: [PATCH 73/74] log line

---
 jetstream/tests/core/test_server.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/jetstream/tests/core/test_server.py b/jetstream/tests/core/test_server.py
index d8972a7e..f4706b0e 100644
--- a/jetstream/tests/core/test_server.py
+++ b/jetstream/tests/core/test_server.py
@@ -65,6 +65,7 @@ async def test_server(
     metrics_port = portpicker.pick_unused_port()
 
     print("port: " + str(port))
+    print("metrics port: " + str(metrics_port))
     credentials = grpc.local_server_credentials()
 
     server = server_lib.run(

From a5c7539a25806924379240904e7c9686e833c490 Mon Sep 17 00:00:00 2001
From: Brendan Slabe <slabe@google.com>
Date: Sat, 18 May 2024 00:00:56 +0000
Subject: [PATCH 74/74] timeout

---
 jetstream/tests/core/test_server.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/jetstream/tests/core/test_server.py b/jetstream/tests/core/test_server.py
index f4706b0e..68faef15 100644
--- a/jetstream/tests/core/test_server.py
+++ b/jetstream/tests/core/test_server.py
@@ -118,7 +118,9 @@ async def test_server(
       # assert prometheus server is running and responding
       assert server._driver._metrics_collector is not None  # pylint: disable=protected-access
       assert (
-          requests.get(f"http://localhost:{metrics_port}").status_code
+          requests.get(
+              f"http://localhost:{metrics_port}", timeout=5
+          ).status_code
           == requests.status_codes.codes["ok"]
       )
       server.stop()
