Merge

Author: tdene

tests/unit_tests/inference/test_data_parallel_inference_coordinator.py

@@ -79,6 +79,9 @@ async def async_step(
         to_remove = []
         for request_id, request in self.requests.items():
             if request.status == Status.ACTIVE_AND_GENERATING_TOKENS:
+                request.sampling_params.num_tokens_to_generate -= 1
+                if request.sampling_params.num_tokens_to_generate > 0:
+                    continue
                 request.status = Status.COMPLETED
                 self.context.active_cnt -= 1
                 finished_requests.append(request)
@@ -116,6 +119,7 @@ class CoordinatorTestConfig:
     num_requests: int = 10**1
     min_time_offset: float = 10 ** (-4)
     max_time_offset: float = 10 ** (-3)
+    num_steps_to_finish: int = 1
     num_iterations: int = 1
 
     tensor_model_parallel_size: int = 1
@@ -148,7 +152,10 @@ def _build_requests(cls, test_config: CoordinatorTestConfig) -> List[Tuple]:
 
         for _ in range(test_config.num_requests):
             arrival_delta = random.uniform(test_config.min_time_offset, test_config.max_time_offset)
-            ret.append(("Hello world!", SamplingParams(), arrival_delta))
+            num_tokens = test_config.num_steps_to_finish
+            ret.append(
+                ("Hello world!", SamplingParams(num_tokens_to_generate=num_tokens), arrival_delta)
+            )
         return ret
 
     @classmethod
@@ -159,6 +166,7 @@ def _build_test_env(cls, test_config):
         )
         requests = cls._build_requests(test_config)
         engine = DummyEngine()
+        engine.num_steps_to_finish = test_config.num_steps_to_finish
         return CoordinatorTestEnv(config=test_config, requests=requests, engine=engine)
 
     @classmethod
@@ -174,31 +182,42 @@ async def _run_test(cls, **test_config_kwargs):
             launch_inference_coordinator=test_config.launch_inference_coordinator,
         )
 
-        if dist.get_rank() == 0:
-            client = InferenceClient(test_config.port)
-            await client.start()
-            env.timing_data["init_time"] = time.time()
-
-            all_results = []
-            for _ in range(test_config.num_iterations):
-                futures = []
-                for request in tqdm(env.requests, "add_requests"):
-                    prompt, sampling_params, arrival_delta = request
-                    await asyncio.sleep(arrival_delta)
-                    fut = client.add_request(prompt=prompt, sampling_params=sampling_params)
-                    futures.append(fut)
-                results: List[DynamicInferenceRequest] = await asyncio.gather(*futures)
-                all_results.append(results)
-            env.timing_data["done_time"] = time.time()
-
-            if test_config.stop_engines:
-                client.stop_engines()
-            client.stop()
-
-        if test_config.stop_engines:
-            await env.engine.engine_loop_task
+        results_success = False
+        shutdown_success = False
+        try:
+            if dist.get_rank() == 0:
+                client = InferenceClient(test_config.port)
+                await client.start()
+                env.timing_data["init_time"] = time.time()
+
+                all_results = []
+                for _ in range(test_config.num_iterations):
+                    futures = []
+                    for request in tqdm(env.requests, "add_requests"):
+                        prompt, sampling_params, arrival_delta = request
+                        await asyncio.sleep(arrival_delta)
+                        fut = client.add_request(prompt=prompt, sampling_params=sampling_params)
+                        futures.append(fut)
+                    results: List[DynamicInferenceRequest] = await asyncio.gather(*futures)
+                    all_results.append(results)
+                env.timing_data["done_time"] = time.time()
+            results_success = True
+        finally:
+            try:
+                if dist.get_rank() == 0:
+                    if test_config.stop_engines:
+                        client.stop_engines()
+                    client.stop()
+                if test_config.stop_engines:
+                    await env.engine.engine_loop_task
+                shutdown_success = True
+            except:
+                env.engine.engine_loop_task.cancel()
+
         env.timing_data["stop_time"] = time.time()
 
+        assert results_success, "Did not receive all results successfully."
+        assert shutdown_success, "Did not shutdown successfully."
         if dist.get_rank() == 0:
             env.responses = all_results
             if test_config.verify_results:
@@ -297,11 +316,11 @@ def clamp_to_golden_value(value, golden_value, delta=0.1):
 
 if __name__ == "__main__":
     test = TestCoordinator()
-    test.test_simple()
-    test.test_tp()
-    test.test_pp()
-    test_test.tp_pp()
-    test_test.throughput()
+    asyncio.run(test.test_simple())
+    asyncio.run(test.test_tp())
+    asyncio.run(test.test_pp())
+    asyncio.run(test.test_tp_pp())
+    asyncio.run(test.test_throughput())
     test.teardown_method(None)
     print("~~~")
     print("success.")