(streaming: bool, exclude_input_from_output: bool,
model_files, model_path)
| 322 | @pytest.mark.parametrize("streaming", [False, True]) |
| 323 | @pytest.mark.parametrize("exclude_input_from_output", [False]) |
| 324 | def test_multi_request(streaming: bool, exclude_input_from_output: bool, |
| 325 | model_files, model_path): |
| 326 | output_config = trtllm.OutputConfig() |
| 327 | output_config.exclude_input_from_output = exclude_input_from_output |
| 328 | |
| 329 | # Create executor |
| 330 | beam_width = 1 |
| 331 | executor_config = trtllm.ExecutorConfig( |
| 332 | beam_width, |
| 333 | kv_cache_config=trtllm.KvCacheConfig(free_gpu_memory_fraction=0.5)) |
| 334 | executor = trtllm.Executor(model_path, trtllm.ModelType.DECODER_ONLY, |
| 335 | executor_config) |
| 336 | |
| 337 | num_requests = 20 |
| 338 | max_prompt_len = 20 |
| 339 | max_max_tokens = 20 |
| 340 | end_id = -1 |
| 341 | |
| 342 | # Enqueue the requests |
| 343 | tokens = {} |
| 344 | expected_num_tokens = {} |
| 345 | for i in range(num_requests): |
| 346 | prompt_len = random.randint(1, max_prompt_len) |
| 347 | max_tokens = random.randint(1, max_max_tokens) |
| 348 | input_tokens = [1] * prompt_len |
| 349 | |
| 350 | # Some requests has num_return_sequences > 1. |
| 351 | num_return_sequences = 2 if i % 5 == 1 else 1 |
| 352 | |
| 353 | request = trtllm.Request(input_tokens, |
| 354 | max_tokens=max_tokens, |
| 355 | streaming=streaming, |
| 356 | sampling_config=trtllm.SamplingConfig( |
| 357 | num_return_sequences=num_return_sequences), |
| 358 | output_config=output_config, |
| 359 | end_id=end_id) |
| 360 | request_id = executor.enqueue_request(request) |
| 361 | tokens[request_id] = [ |
| 362 | [] for _ in range(request.sampling_config.num_return_sequences) |
| 363 | ] |
| 364 | expected_num_tokens[request_id] = get_expected_num_tokens( |
| 365 | prompt_len, max_tokens, streaming, exclude_input_from_output) |
| 366 | |
| 367 | # Get the new tokens for each request |
| 368 | num_finished = 0 |
| 369 | i = 0 |
| 370 | num_responses = 0 |
| 371 | max_wait_ms = 10000 |
| 372 | while num_finished < num_requests and i < max_wait_ms: |
| 373 | wait_time = datetime.timedelta(milliseconds=1) |
| 374 | responses = executor.await_responses(wait_time) |
| 375 | for response in responses: |
| 376 | num_responses += 1 |
| 377 | assert not response.has_error( |
| 378 | ), f"Request id {response.request_id} failed with err {response.error_msg}" |
| 379 | result = response.result |
| 380 | num_finished += result.is_final |
| 381 | new_tokens = result.output_token_ids[beam_width - 1] |
nothing calls this directly
no test coverage detected