Metrics bug: server_lib should be config_lib (#113)

* Metrics bug: server_lib should be config_lib * format
AI-Hypercomputer · Jun 4, 2024 · a32be5d · a32be5d
1 parent 958678c
commit a32be5d
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/run_server.py b/run_server.py
@@ -20,7 +20,7 @@
 import jetstream_pt
 from absl import app, flags
 from jetstream.core import server_lib
-from jetstream.core.config_lib import ServerConfig
+from jetstream.core.config_lib import ServerConfig, MetricsServerConfig
 from jetstream_pt.config import FLAGS, create_engine_from_config_flags
 
 flags.DEFINE_integer("port", 9000, "port to listen on")
@@ -49,11 +49,9 @@ def main(argv: Sequence[str]):
   )
   print(f"server_config: {server_config}")
 
-  metrics_server_config: server_lib.MetricsServerConfig | None = None
+  metrics_server_config: MetricsServerConfig | None = None
   if FLAGS.prometheus_port != 0:
-    metrics_server_config = server_lib.MetricsServerConfig(
-        port=FLAGS.prometheus_port
-    )
+    metrics_server_config = MetricsServerConfig(port=FLAGS.prometheus_port)
 
   # We separate credential from run so that we can unit test it with local credentials.
   # We would like to add grpc credentials for OSS.