fix refactor

jetstream/entrypoints/http/api_server.py

@@ -15,16 +15,19 @@
 """JetStream Http API server."""
 
 import json
+import logging
 from typing import Sequence
 from absl import app as abslapp
 from absl import flags
 from fastapi import APIRouter, Response
 import fastapi
 from fastapi.responses import StreamingResponse
+from prometheus_client import start_http_server
 import uvicorn
 from google.protobuf.json_format import Parse
 
 from jetstream.core import config_lib, orchestrator, server_lib
+from jetstream.core.metrics.prometheus import JetstreamMetricsCollector
 from jetstream.core.proto import jetstream_pb2
 from jetstream.entrypoints.config import get_server_config
 from jetstream.entrypoints.http.protocol import DecodeRequest
@@ -93,17 +96,28 @@ def server(argv: Sequence[str]):
   del argv
 
   metrics_server_config: config_lib.MetricsServerConfig | None = None
+  # Setup Prometheus server
+  metrics_collector: JetstreamMetricsCollector = None
   if flags.FLAGS.prometheus_port != 0:
     metrics_server_config = config_lib.MetricsServerConfig(
         port=flags.FLAGS.prometheus_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"
+    )
 
   global llm_orchestrator
   llm_orchestrator = orchestrator.LLMOrchestrator(
       driver=server_lib.create_driver(
           config=server_config,
           devices=devices,
-          metrics_server_config=metrics_server_config,
+          metrics_collector=metrics_collector,
       )
   )