Always print server starup messages

Author: vipannalla

jetstream/core/server_lib.py

@@ -23,6 +23,7 @@
 import logging
 import os
 import signal
+import sys
 import threading
 import time
 import traceback
@@ -41,6 +42,20 @@
 
 _HOST = "[::]"
 
+# Create seperate logger to log all INFO message for this module. These show
+# stages of server startup and inform user if server is ready to take requests.
+# The default logger created in orchestrator.py only logs WARNINGs and above
+logger = logging.getLogger(__name__)
+logger.propagate = False
+formatter = logging.Formatter(
+    "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
+)
+
+info_handler = logging.StreamHandler(sys.stdout)
+info_handler.setLevel(logging.INFO)
+info_handler.setFormatter(formatter)
+logger.addHandler(info_handler)
+
 
 class JetStreamServer:
   """JetStream grpc server."""
@@ -120,7 +135,7 @@ def create_driver(
   prefill_params = [pe.load_params() for pe in engines.prefill_engines]
   generate_params = [ge.load_params() for ge in engines.generate_engines]
   shared_params = [ie.load_params() for ie in engines.interleaved_engines]
-  logging.info("Loaded all weights.")
+  logger.info("Loaded all weights.")
   if metrics_collector:
     metrics_collector.get_model_load_time_metric().set(
         time.time() - model_load_start_time
@@ -215,19 +230,19 @@ def run(
   del lora_input_adapters_path
 
   server_start_time = time.time()
-  logging.info("Kicking off gRPC server.")
+  logger.info("Kicking off gRPC server.")
   # Setup Prometheus server
   metrics_collector: JetstreamMetricsCollector = None
   if metrics_server_config and metrics_server_config.port:
-    logging.info(
+    logger.info(
         "Starting Prometheus server on port %d", metrics_server_config.port
     )
     start_http_server(metrics_server_config.port)
     metrics_collector = JetstreamMetricsCollector(
         model_name=metrics_server_config.model_name
     )
   else:
-    logging.info(
+    logger.info(
         "Not starting Prometheus server: --prometheus_port flag not set"
     )
 
@@ -256,7 +271,7 @@ def run(
   gc.set_threshold(allocs, gen1, gen2)
   print("GC tweaked (allocs, gen1, gen2): ", allocs, gen1, gen2)
 
-  logging.info("Starting server on port %d with %d threads", port, threads)
+  logger.info("Starting server on port %d with %d threads", port, threads)
   jetstream_server.start()
 
   if metrics_collector:
@@ -266,10 +281,10 @@ def run(
 
   # Setup Jax Profiler
   if enable_jax_profiler:
-    logging.info("Starting JAX profiler server on port %s", jax_profiler_port)
+    logger.info("Starting JAX profiler server on port %s", jax_profiler_port)
     jax.profiler.start_server(jax_profiler_port)
   else:
-    logging.info("Not starting JAX profiler server: %s", enable_jax_profiler)
+    logger.info("Not starting JAX profiler server: %s", enable_jax_profiler)
 
   # Start profiling server by default for proxy backend.
   if jax.config.jax_platforms and "proxy" in jax.config.jax_platforms:
@@ -287,5 +302,5 @@ def get_devices() -> Any:
   """Gets devices."""
   # TODO: Add more logs for the devices.
   devices = jax.devices()
-  logging.info("Using devices: %d", len(devices))
+  logger.info("Using devices: %d", len(devices))
   return devices