Decode Batch Percentage Metrics/Improved Scraping (#82)

* initial-commit

* newline

* moved gauge

* Added labels for metrics

* revert removing driver metric field

* removed newline

* Removed unneccesary lambda

* Proper use of labels api

* format

* missing quotes

* convert result to float, add lambda

* Moved registration to driver level

* made metric driver property

* missing field on class

* Dont regenerate uuid on each scrape

* Init uuid

* Moved metrics to separate file

* Typos

* Initialize metrics

* reformat

* added 'global'

* Moved registration out of function

* removed unused import

* Update server_lib.py

* jetstream_prefill_backlog_size -> prefill_backlog_size

* label fields

* rename metrics

* Moved metrics to singleton class

* Revert requirements change

* revert requirements changes

* Plumb metrics config to JetStreamServer run header

* Cleanup of prior commit

* Default for hostname metric label

* __ -> _

* Linter error

* linter error

* linter error

* default value for idx

* Linter error

* Linter error

* Linter error final fix hopefully

* Fixed type annotation

* Type fix in server_lib.py

* __ -> _ in JetstreamMetricsCollector property names

* added docstrings

* added module docstring

* laxy % formatting

* metrics port cannot be 0

* idx can be None

* Removed redundant class

* removed lingering import

* Update test_server.py

* hostname -> id

* linter

* is not -> !=

* Zijun nits

* reformat

* missing protocol scheme in url

* requests -> aiohttp

* linter

* Cleaned readme

* Changes to tests, readme

* fixed assertion

* better description

* parameterized test setup

* linter

* Update online-inference-with-maxtext-engine.md

* disable protected-access

* make prometheus test not async, revert requirements.txt changes

* Update requirements.txt

* remove aiohttp

* moved prometheus test to test_server function

* log line

* timeout

Author: Bslabe123

docs/online-inference-with-maxtext-engine.md

@@ -205,6 +205,41 @@ Prompt: Today is a good day
 Response:  to be a fan
 ```
 
+### (optional) Observe Jetstream metrics
+
+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 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="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="<SOME-HOSTNAME-HERE>",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:
@@ -289,4 +324,4 @@ rm -rf maxtext
 rm -rf JetStream
 # Clean up python virtual environment
 rm -rf .env
-```
+```

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
@@ -46,6 +47,11 @@ class InstantiatedEngines:
   interleaved_engines: List[engine_api.Engine]
 
 
+@dataclasses.dataclass
+class MetricsServerConfig:
+  port: uint16
+
+
 # ▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼▼#
 
 

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.

jetstream/core/metrics/prometheus.py

@@ -0,0 +1,48 @@
+# 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
+
+
+class JetstreamMetricsCollector:
+  """Wrapper class should be used to assure all metrics have proper tags"""
+
+  _id: str = os.getenv("HOSTNAME", shortuuid.uuid())
+
+  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=["id"],
+  )
+  _slots_available_percentage = Gauge(
+      name="jetstream_slots_available_percentage",
+      documentation="The percentage of available slots in decode batch",
+      labelnames=["id", "idx"],
+  )
+
+  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)

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
+from jetstream.core.metrics.prometheus import JetstreamMetricsCollector
 import numpy as np
-import prometheus_client
-import shortuuid
 
 root = logging.getLogger()
 root.setLevel(logging.DEBUG)
@@ -212,8 +211,8 @@ 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
+  # All metrics we want to monitor should be collected with this
+  _metrics_collector: JetstreamMetricsCollector | None = None
 
   def __init__(
       self,
@@ -223,6 +222,7 @@ def __init__(
       generate_params: Optional[list[Any]] = None,
       interleaved_mode: bool = False,
       jax_padding: bool = True,
+      metrics_collector: JetstreamMetricsCollector | None = None,
   ):
     if prefill_engines is None:
       prefill_engines = []
@@ -243,15 +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()
-    self._prefill_backlog_size_metric = prometheus_client.Gauge(
-        f"jetstream_prefill_backlog_size_{shortuuid.uuid()}",
-        "Size of prefill queue",
-    )
+    if self._metrics_collector:
+      self._metrics_collector.get_prefill_backlog_metric().set_function(
+          lambda: float(self._prefill_backlog.qsize())
+      )
 
     # Stage 2
     # After prefilling, it is placed here in order to get transferred to
@@ -432,7 +433,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 _process_prefill_content(
       self,
@@ -474,7 +474,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
@@ -579,6 +578,11 @@ def _generate_thread(self, idx: int):
 
       max_concurrent_decodes = generate_engine.max_concurrent_decodes
 
+      if self._metrics_collector:
+        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
       # insert as many sequences as possible.

jetstream/core/server_lib.py

@@ -20,24 +20,19 @@
 import asyncio
 from concurrent import futures
 import logging
-import os
 import threading
 from typing import Any, Type
 
 import grpc
 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
 
 _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 +94,7 @@ def run(
     credentials: Any = grpc.insecure_server_credentials(),
     threads: int | None = None,
     jax_padding: bool = True,
+    metrics_server_config: config_lib.MetricsServerConfig | None = None,
 ) -> JetStreamServer:
   """Runs a server with a specified config.
 
@@ -122,13 +118,28 @@ def run(
   interleaved_mode = (
       len(config.prefill_slices) + len(config.generate_slices) == 0
   )
+
+  # Setup Prometheus server
+  metrics_collector: JetstreamMetricsCollector = None
+  if metrics_server_config and metrics_server_config.port:
+    logging.info(
+        "Starting Prometheus server on port %d", 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,
       prefill_params=prefill_params + shared_params,
       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.
@@ -137,17 +148,6 @@ def run(
   logging.info("Starting server on port %d with %d threads", port, threads)
 
   jetstream_server.start()
-
-  # Setup Prometheus server
-  if PROMETHEUS_ENABLED_ON_PORT is not None:
-    logging.info(
-        "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"
-    )
   return jetstream_server
 
 

jetstream/tests/core/test_server.py

@@ -21,6 +21,8 @@
 from typing import Any, Type
 import unittest
 
+
+import requests
 from parameterized import parameterized
 import grpc
 from jetstream.core import config_lib
@@ -60,7 +62,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(
@@ -70,12 +75,16 @@ 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:
       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,5 +105,25 @@ async def test_server(
         counter += 1
       server.stop()
 
+      # 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
+          ),
+      )
+      # 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}", timeout=5
+          ).status_code
+          == requests.status_codes.codes["ok"]
+      )
+      server.stop()
+
   def test_get_devices(self):
     assert len(server_lib.get_devices()) == 1