Adding following metrics into JetStream server:

1) kv_cache_utilization: This refers to percentage of memory in the allocated kv-cache on TPU HBM, that is actually used during decode. It is based on the percentage of slots used.
2) num_requests_waiting: Total number of requests which are waiting to be decoded.
3) lora_requests_info: List of LoRA adapters that are loaded into the TPU HBM for serving the requests.

Author: aman2930

jetstream/core/adapter_tensorstore.py

@@ -139,6 +139,17 @@ def convert_if_np(leaf):
     return jax.tree_util.tree_map(convert_if_np, params)
 
 
+  async def get_hbm_loaded_adapters(self):
+    hbm_loaded_adapters = []
+
+    async with self.lock:
+      for adapter_id, metadata in self.adapter_registry.items():
+        if metadata.status == "loaded_hbm":
+          hbm_loaded_adapters.append(adapter_id)
+
+    return ", ".join(hbm_loaded_adapters)
+
+
   async def load_adapter(self, adapter_id: str, adapter_weights = None, to_hbm: bool = True):
     """Loads a LoRA adapter's weights, managing HBM and CPU memory."""
     if adapter_id not in self.adapter_registry:

jetstream/core/metrics/prometheus.py

@@ -214,6 +214,31 @@ def __new__(cls):
       ],
   )
 
+  _num_requests_waiting = Gauge(
+      name="num_requests_waiting",
+      documentation="Number of requests waiting to be processed for inference.",
+      labelnames=["id"],
+      multiprocess_mode="sum",
+  )
+
+  _kv_cache_utilization = Gauge(
+      name="kv_cache_utilization_perc",
+      documentation="Percentage of kv-cache utilized by the requests under processing.",
+      labelnames=["id"],
+      multiprocess_mode="sum",
+  )
+
+  _lora_request_info = Gauge(
+      name="lora_request_info",
+      documentation="Information about LoRA adapters loaded into TPU Memory for serving current requests.",
+      labelnames=[
+        "id",
+        "max_lora",
+        "running_lora_adapters",
+      ],
+      multiprocess_mode="livemostrecent",
+  )
+
   def get_prefill_backlog_metric(self):
     return self._prefill_backlog.labels(id=self._id)
 
@@ -255,3 +280,12 @@ def get_request_output_length(self):
 
   def get_request_success_count_metric(self):
     return self._request_success_count.labels(id=self._id)
+
+  def get_num_requests_waiting_metric(self):
+    return self._num_requests_waiting.labels(id=self._id)
+
+  def get_kv_cache_utilization_metric(self):
+    return self._kv_cache_utilization.labels(id=self._id)
+
+  def get_lora_request_info_metric(self, max_lora: int, loaded_adapters: str):
+    return self._lora_request_info.labels(id=self._id, max_lora=max_lora, running_lora_adapters=loaded_adapters)

jetstream/core/orchestrator.py

@@ -314,6 +314,7 @@ def __init__(
         self._metrics_collector.get_generate_backlog_metric(idx).set_function(
             functools.partial(float, backlog.qsize())
         )
+
     # Stage 4
     # After prefill and generation, ActiveRequests are placed on the
     # detokenization backlog for tokens to be sent into each ActiveRequest's
@@ -433,6 +434,12 @@ def __init__(
     self.live = True
     self._is_ray_backend = is_ray_backend
 
+    if self._metrics_collector:
+      self._metrics_collector.get_num_requests_waiting_metric().set_function(
+          self._get_total_requests_waiting_decode)
+      self._metrics_collector.get_kv_cache_utilization_metric().set_function(
+          self._get_kv_cache_utilization)
+
     # Start all threads
     for t in self._all_threads:
       t.start()
@@ -481,6 +488,28 @@ def stop(self):
     for t in self._all_threads:
       t.join()
 
+  def _get_kv_cache_utilization(self):
+    """Calculated the kv_cache utilization in percentage based on requests being decoded."""
+    total_slots = 0
+    empty_slots = 0
+    for idx, engine in enumerate(self._generate_engines):
+      total_slots += engine.max_concurrent_decodes
+      empty_slots += self._generate_slots[idx].qsize()
+
+    return ((total_slots - empty_slots) * 100 / total_slots)
+
+  def _get_total_requests_waiting_decode(self):
+    """Calculate the total size of all relevant queues."""
+    total_size = self._prefill_backlog.qsize()
+
+    for transfer_queue in self._transfer_backlogs:
+      total_size += transfer_queue.qsize()
+
+    for gen_queue in self._generate_backlogs.values():
+      total_size += gen_queue.qsize()
+
+    return float(total_size)
+
   def get_total_concurrent_requests(self) -> int:
     """Gets the total number of concurrent requests the driver can handle."""
     # We don't support filling all backlogs at once because it can cause GIL
@@ -819,6 +848,14 @@ def _generate_thread(self, idx: int):
 
       start_time = time.perf_counter()
 
+      if self._metrics_collector:
+        adapters_list_str = asyncio.run(self._adapter_tensorstore.get_hbm_loaded_adapters())
+
+        max_loras = max_concurrent_decodes
+
+        self._metrics_collector.get_lora_request_info_metric(max_loras,
+            adapters_list_str).set_to_current_time()
+
       # Now we actually take a generate step on requests in the slots.
       decode_state, sampled_tokens = generate_engine.generate(
           generate_params[adapter_id], decode_state