Support gracefully stopping orchestrator and server (#6)

* Support gracefully stopping orchestrator and server

* Add JetStream server

Author: JoeZijunZhou

.github/workflows/UnitTests.yaml

@@ -69,4 +69,10 @@ jobs:
         python -m jetstream.engine.utils_test
     - name: Test mock JetStream engine implementation
       run: |
-        python -m jetstream.engine.mock_engine_test
+        python -m jetstream.engine.mock_engine_test
+    - name: Test JetStream core orchestrator
+      run: |
+        python -m jetstream.core.orchestrator_test
+    - name: Test JetStream core server library
+      run: |
+        python -m jetstream.core.server_test

benchmarks/benchmark_serving.py

@@ -242,7 +242,7 @@ async def send_request(
     max_tokens: int,
     threads: int,
 ) -> RequestFuncOutput:
-  """Send the request to wiz server."""
+  """Send the request to JetStream server."""
   loop = asyncio.get_running_loop()
   loop.set_default_executor(ThreadPoolExecutor(max_workers=threads))
   request = jetstream_pb2.DecodeRequest(
@@ -406,7 +406,7 @@ def main(args: argparse.Namespace):
     # Save to file
     base_model_id = model_id.split("/")[-1]
     file_name = (
-        f"JetEngine-{args.request_rate}qps-{base_model_id}-{current_dt}.json"
+        f"JetStream-{args.request_rate}qps-{base_model_id}-{current_dt}.json"
     )
     with open(file_name, "w") as outfile:
       json.dump(result_json, outfile)
@@ -433,7 +433,7 @@ def main(args: argparse.Namespace):
       help=(
           "Name of the model. (it's just used to label the benchmark, the model"
           " config is defined in config_lib, and passed as the server config"
-          " flag when we run the wiz-pathways server)"
+          " flag when we run the JetStream server)"
       ),
   )
   parser.add_argument(

jetstream/core/orchestrator.py

@@ -76,6 +76,7 @@
 
 import dataclasses
 import functools
+import itertools
 import logging
 import os
 import queue
@@ -100,15 +101,19 @@
 
 handler = logging.StreamHandler(sys.stdout)
 handler.setLevel(logging.DEBUG)
-formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
+formatter = logging.Formatter(
+    '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
 handler.setFormatter(formatter)
 root.addHandler(handler)
 
+
 def delete_pytree(p):
   def delete_leaf(leaf):
     if isinstance(leaf, jax.Array):
       leaf.delete()
     del leaf
+
   jax.tree_map(delete_leaf, p)
 
 
@@ -185,19 +190,19 @@ class Driver:
   _prefill_params: Optional[dict[int, Any]] = {}
   _generate_params: Optional[dict[int, Any]] = {}
   # Stage 1
-  _prefill_backlog: queue.Queue[ActiveRequest]
+  _prefill_backlog: queue.Queue[ActiveRequest | None]
   # Stage 2
   # We keep this as a dict to avoid a possibly expensive object comparison
   # when logging the index of the generate engine we send a prefill result
   # to, it allows us to natively have the index from the min operation, rather
   # than have to call .index()
-  _generate_backlogs: dict[int, queue.Queue[ActiveRequest]] = {}
+  _generate_backlogs: dict[int, queue.Queue[ActiveRequest | None]] = {}
   # Stage 3
   # This can be a list because we can pass it as an arg to generate and
   # detokenize threads. It is a list of tokens to be detokenized.
   _detokenize_backlogs: list[queue.Queue[engine_api.ResultTokens]] = []
   _generate_slots: list[queue.Queue[int]] = []
-  _active_requests: list[queue.Queue[tuple[int, ActiveRequest]]] = []
+  _active_requests: list[queue.Queue[tuple[int, ActiveRequest | None]]] = []
 
   def __init__(
       self,
@@ -296,11 +301,58 @@ def __init__(
         )
         for idx, engine in enumerate(self._generate_engines)
     ]
+    self._all_threads = list(
+        itertools.chain(
+            self._prefill_threads,
+            self._generate_threads,
+            self.detokenize_threads,
+        )
+    )
     self.live = True
     # Kick off all threads
-    _ = [f.start() for f in self._prefill_threads]
-    _ = [f.start() for f in self._generate_threads]
-    _ = [f.start() for f in self.detokenize_threads]
+    for t in self._all_threads:
+      t.start()
+
+  def stop(self):
+    """Stops the driver and all background threads."""
+    # Signal to all threads that they should stop.
+    self.live = False
+
+    all_backlogs = list(
+        itertools.chain(
+            [self._prefill_backlog],
+            self._generate_backlogs.values(),
+            self._detokenize_backlogs,
+        )
+    )
+
+    while any(t.is_alive() for t in self._all_threads):
+      # Empty all backlogs and mark any remaining requests as cancelled.
+      for q in all_backlogs:
+        while True:
+          try:
+            r = q.get_nowait()
+            if r is None:
+              continue
+            elif isinstance(r, ActiveRequest):
+              r.return_channel = None
+            else:  # detokenize backlog
+              _, r = r
+              if isinstance(r, ActiveRequest):
+                r.return_channel = None
+          except queue.Empty:
+            break
+
+      # Put sentinels to unblock threads.
+      for q in all_backlogs:
+        try:
+          q.put_nowait(None)
+        except queue.Full:
+          pass
+
+    # Wait for all threads to stop.
+    for t in self._all_threads:
+      t.join()
 
   def get_total_concurrent_requests(self) -> int:
     """Returns the total number of concurrent requests the driver can service."""
@@ -338,10 +390,12 @@ def _prefill_thread(
     while self.live:
       # We don't want to keep lots of kv caches live in memory on the prefill
       # slice that aren't about to be sent over to a generation slice.
-      if (self._generate_backlogs[idx].qsize() < generate_backpressure):
+      if self._generate_backlogs[idx].qsize() < generate_backpressure:
         # Check if there is anything on the prefill backlog, pop if so.
         try:
           request = self._prefill_backlog.get(block=True)
+          if request is None:
+            break
           # TODO: Implement hot/cold cache for history.
           history = self._load_cache_history(request.history_path)  # pylint: disable = assignment-from-none
           # Tokenize, and introduce a leading dimension
@@ -372,7 +426,8 @@ def _prefill_thread(
           # Once prefill is complete, place it on the generation queue.
           self._generate_backlogs[idx].put(request)
           logging.info(
-              f'Placed request on the generate queue, {self._generate_backlogs[idx].qsize()=}'
+              'Placed request on the generate queue,'
+              f' {self._generate_backlogs[idx].qsize()=}'
           )
         except queue.Empty:
           # Otherwise, don't do anything!
@@ -410,13 +465,16 @@ def _generate_thread(
     while self.live:
       if (time.time() - time_of_last_print) > 1:
         logging.info(
-           f'Generate thread making a decision with: prefill_backlog={self._prefill_backlog.qsize()} generate_free_slots={my_slots.qsize()}'
+            'Generate thread making a decision with:'
+            f' prefill_backlog={self._prefill_backlog.qsize()} generate_free_slots={my_slots.qsize()}'
         )
         time_of_last_print = time.time()
       # Check if there are any free my_slots.
       if not my_slots.empty() and not self._generate_backlogs[idx].empty():
         # Only get requests from the backlog corresponding to this engine.
         new_request = self._generate_backlogs[idx].get()
+        if new_request is None:
+          break
         slot = my_slots.get()
         logging.info(
             'Generate slice %d slot %d step %d',
@@ -475,6 +533,8 @@ def _detokenize_thread(
     while self.live:
       try:
         data = my_detokenize_backlog.get(block=True)
+        if data is None:
+          break
         start_detokenise_time = time.time()
         if isinstance(data[1], engine_api.ResultTokens):
           # We want to detokenise them.

jetstream/core/orchestrator_test.py

@@ -41,10 +41,10 @@
 tokenizer returns).
 """
 
-from jetstream.engine import mock_engine
+from absl.testing import absltest
 from jetstream.core import orchestrator
 from jetstream.core.proto import jetstream_pb2
-from absl.testing import absltest
+from jetstream.engine import mock_engine
 
 
 class OrchestratorTest(absltest.TestCase):
@@ -87,12 +87,16 @@ def test_orchestrator(self):
     counter = 0
     for token in iterator:
       # Tokens come through as bytes.
-      print('actual output: ' + bytes(token.response[0], encoding='utf-8').decode())
+      print(
+          'actual output: '
+          + bytes(token.response[0], encoding='utf-8').decode()
+      )
       assert (
           bytes(token.response[0], encoding='utf-8').decode()
           == expected_tokens[counter]
       )
       counter += 1
+    driver.stop()
 
 
 if __name__ == '__main__':

jetstream/core/server_lib.py

@@ -31,22 +31,45 @@
 _HOST = '[::]'
 
 
+class JetStreamServer:
+  """JetStream grpc server."""
+
+  def __init__(self, driver: orchestrator.Driver, server: grpc.Server):
+    self._driver = driver
+    self._server = server
+
+  def start(self, port, credentials) -> None:
+    self._server.add_secure_port(f'{_HOST}:{port}', credentials)
+    self._server.start()
+
+  def stop(self) -> None:
+    # Gracefully clean up threads in the orchestrator.
+    self._driver.stop()
+    self._server.stop(0)
+
+  def wait_for_termination(self) -> None:
+    self._server.wait_for_termination()
+
+
 def run(
     port: int,
     config: Type[config_lib.ServerConfig],
     devices: Any,
     credentials: Any = grpc.insecure_server_credentials(),
     threads: int | None = None,
-) -> grpc.Server:
+) -> JetStreamServer:
   """Runs a server with a specified config.
 
   Args:
     port: Port on which the server will be made available.
     config: A ServerConfig to config engine, model, device slices, etc.
-    device: Device objects, will be used to get engine with proper slicing.
+    devices: Device objects, will be used to get engine with proper slicing.
     credentials: Should use grpc credentials by default.
     threads: Number of RPC handlers worker threads. This should be at least
       equal to the decoding batch size to fully saturate the decoding queue.
+
+  Returns:
+    JetStreamServer that wraps the grpc server and orchestrator driver.
   """
   logging.info('Kicking off gRPC server.')
   engines = config_lib.get_engines(config, devices=devices)
@@ -69,9 +92,9 @@ def run(
   )
   logging.info('Starting server on port %d with %d threads', port, threads)
 
-  server.add_secure_port(f'{_HOST}:{port}', credentials)
-  server.start()
-  return server
+  jetstream_server = JetStreamServer(driver, server)
+  jetstream_server.start(port, credentials)
+  return jetstream_server
 
 
 def get_devices() -> Any:

jetstream/core/server_test.py

@@ -20,14 +20,13 @@
 
 from typing import Any, Type
 
+from absl.testing import absltest, parameterized
 import grpc
-import portpicker
-
 from jetstream.core import config_lib
 from jetstream.core import server_lib
 from jetstream.core.proto import jetstream_pb2
 from jetstream.core.proto import jetstream_pb2_grpc
-from absl.testing import absltest, parameterized
+import portpicker
 
 
 class ServerTest(parameterized.TestCase):
@@ -58,7 +57,7 @@ def test_server(
     print('port: ' + str(port))
     credentials = grpc.local_server_credentials()
 
-    _ = server_lib.run(
+    server = server_lib.run(
         port=port,
         config=config,
         devices=devices,
@@ -83,12 +82,16 @@ def test_server(
     counter = 0
     for token in iterator:
       # Tokens come through as bytes
-      print('actual output: ' + bytes(token.response[0], encoding='utf-8').decode())
+      print(
+          'actual output: '
+          + bytes(token.response[0], encoding='utf-8').decode()
+      )
       assert (
           bytes(token.response[0], encoding='utf-8').decode()
           == expected_tokens[counter]
       )
       counter += 1
+    server.stop()
 
 
 if __name__ == '__main__':