Pass metrics client config through to Jetstream (#99)

* first commit

* reformat

Author: Bslabe123

run_server.py

@@ -30,6 +30,7 @@
     "InterleavedCPUTestServer",
     "available servers",
 )
+flags.DEFINE_integer("prometheus_port", 0, "")
 
 
 # pylint: disable-next=all
@@ -48,13 +49,20 @@ def main(argv: Sequence[str]):
   )
   print(f"server_config: {server_config}")
 
+  metrics_server_config: server_lib.MetricsServerConfig | None = None
+  if FLAGS.prometheus_port != 0:
+    metrics_server_config = server_lib.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.
   jetstream_server = server_lib.run(
       threads=FLAGS.threads,
       port=FLAGS.port,
       config=server_config,
       devices=devices,
+      metrics_server_config=metrics_server_config,
   )
   print("Started jetstream_server....")
   jetstream_server.wait_for_termination()