Merge pull request #681 from google:prometheus-flag-import-fix

PiperOrigin-RevId: 640138468
google · Jun 4, 2024 · b7f5864 · b7f5864
2 parents ca55917 + 83159ab
commit b7f5864
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/MaxText/maxengine_server.py b/MaxText/maxengine_server.py
@@ -20,7 +20,7 @@
 import pyconfig
 
 import maxengine_config
-from jetstream.core import server_lib
+from jetstream.core import server_lib, config_lib
 
 # _PORT = flags.DEFINE_integer('port', 9000, 'port to listen on')
 # _THREADS = flags.DEFINE_integer(
@@ -38,9 +38,9 @@ def main(config):
  devices = server_lib.get_devices()
  server_config = maxengine_config.get_server_config("MaxtextInterleavedServer", config)
 
- metrics_server_config : server_lib.MetricsServerConfig | None = None
+ metrics_server_config : config_lib.MetricsServerConfig | None = None
  if config.prometheus_port != 0:
- metrics_server_config = server_lib.MetricsServerConfig(
+ metrics_server_config = config_lib.MetricsServerConfig(
  port=config.prometheus_port
  )