Recover PR 181 and add pull ready tag (#183)

Author: jyj0w0

jetstream/tests/core/test_orchestrator.py

@@ -42,6 +42,7 @@
 """
 
 import unittest
+from parameterized import parameterized
 from jetstream.core import orchestrator
 from jetstream.core.proto import jetstream_pb2
 from jetstream.core.utils.return_sample import ReturnSample
@@ -50,7 +51,7 @@
 
 class OrchestratorTest(unittest.IsolatedAsyncioTestCase):
 
-  def _setup_driver_interleaved_mode(self):
+  def _setup_driver(self, interleaved_mode: bool = True):
     prefill_engine = mock_engine.TestEngine(
         batch_size=32, cache_length=256, weight=2.0
     )
@@ -64,13 +65,14 @@ def _setup_driver_interleaved_mode(self):
         generate_engines=[generate_engine],
         prefill_params=[prefill_engine.load_params()],
         generate_params=[generate_engine.load_params()],
-        interleaved_mode=True,
+        interleaved_mode=interleaved_mode,
     )
     return driver
 
-  async def test_orchestrator_interleaved_mode(self):
+  @parameterized.expand([True, False])
+  async def test_orchestrator(self, interleaved_mode: bool):
     """Test the multithreaded orchestration."""
-    driver = self._setup_driver_interleaved_mode()
+    driver = self._setup_driver(interleaved_mode)
     client = orchestrator.LLMOrchestrator(driver=driver)
 
     # The string representation of np.array([[65, 66]]), [2] will be prepend
@@ -97,9 +99,10 @@ async def test_orchestrator_interleaved_mode(self):
     driver.stop()
     print("Orchestrator driver stopped.")
 
-  async def test_orchestrator_interleaved_mode_client_tokenization(self):
+  @parameterized.expand([True, False])
+  async def test_orchestrator_client_tokenization(self, interleaved_mode: bool):
     """Test the multithreaded orchestration."""
-    driver = self._setup_driver_interleaved_mode()
+    driver = self._setup_driver(interleaved_mode)
     client = orchestrator.LLMOrchestrator(driver=driver)
 
     # The token ids of  string "AB", [2] will be prepend
@@ -128,8 +131,9 @@ async def test_orchestrator_interleaved_mode_client_tokenization(self):
     driver.stop()
     print("Orchestrator driver stopped.")
 
-  def test_should_buffer_response(self):
-    driver = self._setup_driver_interleaved_mode()
+  @parameterized.expand([True, False])
+  def test_should_buffer_response(self, interleaved_mode: bool):
+    driver = self._setup_driver(interleaved_mode)
     client = orchestrator.LLMOrchestrator(driver=driver)
     self.assertTrue(
         client.should_buffer_response(