Add jetstream_server_startup_latency metric (#118)

* first commit

* no labels on metric

* format

* change measurement

* fmt

* rename metric

* Time -> time

* nits

* fixed args

* int -> float

* int -> float

* move endpoint to server_lib.py

* nit

* missing labels

Author: Bslabe123

jetstream/core/metrics/prometheus.py

@@ -50,6 +50,11 @@ def __new__(cls):
       documentation="The percentage of decode slots currently being used",
       labelnames=["id", "idx"],
   )
+  _server_startup_latency = Gauge(
+      name="jetstream_server_startup_latency",
+      documentation="Total time taken to start the Jetstream server",
+      labelnames=["id"],
+  )
 
   def get_prefill_backlog_metric(self):
     return self._prefill_backlog.labels(id=self._id)
@@ -62,3 +67,6 @@ def get_generate_backlog_metric(self, idx: int):
 
   def get_slots_used_percentage_metric(self, idx: int):
     return self._slots_used_percentage.labels(id=self._id, idx=idx)
+
+  def get_server_startup_latency_metric(self):
+    return self._server_startup_latency.labels(id=self._id)

jetstream/core/server_lib.py

@@ -23,6 +23,7 @@
 import os
 import signal
 import threading
+import time
 import traceback
 from typing import Any, Type
 
@@ -122,6 +123,9 @@ def run(
   Returns:
     JetStreamServer that wraps the grpc server and orchestrator driver.
   """
+
+  server_start_time = time.time()
+
   logging.info("Kicking off gRPC server.")
   engines = config_lib.get_engines(config, devices=devices)
   prefill_params = [pe.load_params() for pe in engines.prefill_engines]
@@ -196,6 +200,11 @@ def run(
 
   jetstream_server.start()
 
+  if metrics_collector:
+    metrics_collector.get_server_startup_latency_metric().set(
+        time.time() - server_start_time
+    )
+
   # Setup Jax Profiler
   if enable_jax_profiler:
     logging.info("Starting JAX profiler server on port %s", jax_profiler_port)