| 1 | #!/usr/bin/env python3 |
| 2 | # Copyright 2022 Xiaomi Corp. (authors: Fangjun Kuang) |
| 3 | # 2023 Nvidia (authors: Yuekai Zhang) |
| 4 | # 2023 Recurrent.ai (authors: Songtao Shi) |
| 5 | # See LICENSE for clarification regarding multiple authors |
| 6 | # |
| 7 | # Licensed under the Apache License, Version 2.0 (the "License"); |
| 8 | # you may not use this file except in compliance with the License. |
| 9 | # You may obtain a copy of the License at |
| 10 | # |
| 11 | # http://www.apache.org/licenses/LICENSE-2.0 |
| 12 | # |
| 13 | # Unless required by applicable law or agreed to in writing, software |
| 14 | # distributed under the License is distributed on an "AS IS" BASIS, |
| 15 | # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 16 | # See the License for the specific language governing permissions and |
| 17 | # limitations under the License. |
| 18 | """ |
| 19 | This script supports to load dataset from huggingface and sends it to the server |
| 20 | for decoding, in parallel. |
| 21 | |
| 22 | Usage: |
| 23 | num_task=2 |
| 24 | |
| 25 | # For offline F5-TTS |
| 26 | python3 client_grpc.py \ |
| 27 | --server-addr localhost \ |
| 28 | --model-name f5_tts \ |
| 29 | --num-tasks $num_task \ |
| 30 | --huggingface-dataset yuekai/seed_tts \ |
| 31 | --split-name test_zh \ |
| 32 | --log-dir ./log_concurrent_tasks_${num_task} |
| 33 | """ |
| 34 | |
| 35 | import argparse |
| 36 | import asyncio |
| 37 | import json |
| 38 | import os |
| 39 | import time |
| 40 | import types |
| 41 | from pathlib import Path |
| 42 | |
| 43 | import numpy as np |
| 44 | import soundfile as sf |
| 45 | import tritonclient |
| 46 | import tritonclient.grpc.aio as grpcclient |
| 47 | from tritonclient.utils import np_to_triton_dtype |
| 48 | |
| 49 | |
| 50 | def write_triton_stats(stats, summary_file): |
| 51 | with open(summary_file, "w") as summary_f: |
| 52 | model_stats = stats["model_stats"] |
| 53 | # write a note, the log is from triton_client.get_inference_statistics(), to better human readability |
| 54 | summary_f.write( |
| 55 | "The log is parsing from triton_client.get_inference_statistics(), to better human readability. \n" |
| 56 | ) |
| 57 | summary_f.write("To learn more about the log, please refer to: \n") |
| 58 | summary_f.write("1. https://github.com/triton-inference-server/server/blob/main/docs/user_guide/metrics.md \n") |
| 59 | summary_f.write("2. https://github.com/triton-inference-server/server/issues/5374 \n\n") |
| 60 | summary_f.write( |
| 61 | "To better improve throughput, we always would like let requests wait in the queue for a while, and then execute them with a larger batch size. \n" |
| 62 | ) |
| 63 | summary_f.write( |
| 64 | "However, there is a trade-off between the increased queue time and the increased batch size. \n" |
| 65 | ) |
| 66 | summary_f.write( |
| 67 | "You may change 'max_queue_delay_microseconds' and 'preferred_batch_size' in the model configuration file to achieve this. \n" |
| 68 | ) |
| 69 | summary_f.write( |
| 70 | "See https://github.com/triton-inference-server/server/blob/main/docs/user_guide/model_configuration.md#delayed-batching for more details. \n\n" |
| 71 | ) |
| 72 | for model_state in model_stats: |
| 73 | if "last_inference" not in model_state: |
| 74 | continue |
| 75 | summary_f.write(f"model name is {model_state['name']} \n") |
| 76 | model_inference_stats = model_state["inference_stats"] |
| 77 | total_queue_time_s = int(model_inference_stats["queue"]["ns"]) / 1e9 |
| 78 | total_infer_time_s = int(model_inference_stats["compute_infer"]["ns"]) / 1e9 |
| 79 | total_input_time_s = int(model_inference_stats["compute_input"]["ns"]) / 1e9 |
| 80 | total_output_time_s = int(model_inference_stats["compute_output"]["ns"]) / 1e9 |
| 81 | summary_f.write( |
| 82 | f"queue time {total_queue_time_s:<5.2f} s, compute infer time {total_infer_time_s:<5.2f} s, compute input time {total_input_time_s:<5.2f} s, compute output time {total_output_time_s:<5.2f} s \n" # noqa |
| 83 | ) |
| 84 | model_batch_stats = model_state["batch_stats"] |
| 85 | for batch in model_batch_stats: |
| 86 | batch_size = int(batch["batch_size"]) |
| 87 | compute_input = batch["compute_input"] |
| 88 | compute_output = batch["compute_output"] |
| 89 | compute_infer = batch["compute_infer"] |
| 90 | batch_count = int(compute_infer["count"]) |
| 91 | assert compute_infer["count"] == compute_output["count"] == compute_input["count"] |
| 92 | compute_infer_time_ms = int(compute_infer["ns"]) / 1e6 |
| 93 | compute_input_time_ms = int(compute_input["ns"]) / 1e6 |
| 94 | compute_output_time_ms = int(compute_output["ns"]) / 1e6 |
| 95 | summary_f.write( |
| 96 | f"execuate inference with batch_size {batch_size:<2} total {batch_count:<5} times, total_infer_time {compute_infer_time_ms:<9.2f} ms, avg_infer_time {compute_infer_time_ms:<9.2f}/{batch_count:<5}={compute_infer_time_ms / batch_count:.2f} ms, avg_infer_time_per_sample {compute_infer_time_ms:<9.2f}/{batch_count:<5}/{batch_size}={compute_infer_time_ms / batch_count / batch_size:.2f} ms \n" # noqa |
| 97 | ) |
| 98 | summary_f.write( |
| 99 | f"input {compute_input_time_ms:<9.2f} ms, avg {compute_input_time_ms / batch_count:.2f} ms, " # noqa |
| 100 | ) |
| 101 | summary_f.write( |
| 102 | f"output {compute_output_time_ms:<9.2f} ms, avg {compute_output_time_ms / batch_count:.2f} ms \n" # noqa |
| 103 | ) |
| 104 | |
| 105 | |
| 106 | def get_args(): |
| 107 | parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter) |
| 108 | |
| 109 | parser.add_argument( |
| 110 | "--server-addr", |
| 111 | type=str, |
| 112 | default="localhost", |
| 113 | help="Address of the server", |
| 114 | ) |
| 115 | |
| 116 | parser.add_argument( |
| 117 | "--server-port", |
| 118 | type=int, |
| 119 | default=8001, |
| 120 | help="Grpc port of the triton server, default is 8001", |
| 121 | ) |
| 122 | |
| 123 | parser.add_argument( |
| 124 | "--reference-audio", |
| 125 | type=str, |
| 126 | default=None, |
| 127 | help="Path to a single audio file. It can't be specified at the same time with --manifest-dir", |
| 128 | ) |
| 129 | |
| 130 | parser.add_argument( |
| 131 | "--reference-text", |
| 132 | type=str, |
| 133 | default="", |
| 134 | help="", |
| 135 | ) |
| 136 | |
| 137 | parser.add_argument( |
| 138 | "--target-text", |
| 139 | type=str, |
| 140 | default="", |
| 141 | help="", |
| 142 | ) |
| 143 | |
| 144 | parser.add_argument( |
| 145 | "--huggingface-dataset", |
| 146 | type=str, |
| 147 | default="yuekai/seed_tts", |
| 148 | help="dataset name in huggingface dataset hub", |
| 149 | ) |
| 150 | |
| 151 | parser.add_argument( |
| 152 | "--split-name", |
| 153 | type=str, |
| 154 | default="wenetspeech4tts", |
| 155 | choices=["wenetspeech4tts", "test_zh", "test_en", "test_hard"], |
| 156 | help="dataset split name, default is 'test'", |
| 157 | ) |
| 158 | |
| 159 | parser.add_argument( |
| 160 | "--manifest-path", |
| 161 | type=str, |
| 162 | default=None, |
| 163 | help="Path to the manifest dir which includes wav.scp trans.txt files.", |
| 164 | ) |
| 165 | |
| 166 | parser.add_argument( |
| 167 | "--model-name", |
| 168 | type=str, |
| 169 | default="f5_tts", |
| 170 | help="triton model_repo module name to request", |
| 171 | ) |
| 172 | |
| 173 | parser.add_argument( |
| 174 | "--num-tasks", |
| 175 | type=int, |
| 176 | default=1, |
| 177 | help="Number of concurrent tasks for sending", |
| 178 | ) |
| 179 | |
| 180 | parser.add_argument( |
| 181 | "--log-interval", |
| 182 | type=int, |
| 183 | default=5, |
| 184 | help="Controls how frequently we print the log.", |
| 185 | ) |
| 186 | |
| 187 | parser.add_argument( |
| 188 | "--compute-wer", |
| 189 | action="store_true", |
| 190 | default=False, |
| 191 | help="""True to compute WER. |
| 192 | """, |
| 193 | ) |
| 194 | |
| 195 | parser.add_argument( |
| 196 | "--log-dir", |
| 197 | type=str, |
| 198 | required=False, |
| 199 | default="./tests/client_grpc", |
| 200 | help="log directory", |
| 201 | ) |
| 202 | |
| 203 | parser.add_argument( |
| 204 | "--batch-size", |
| 205 | type=int, |
| 206 | default=1, |
| 207 | help="Inference batch_size per request for offline mode.", |
| 208 | ) |
| 209 | |
| 210 | return parser.parse_args() |
| 211 | |
| 212 | |
| 213 | def load_audio(wav_path, target_sample_rate=24000): |
| 214 | assert target_sample_rate == 24000, "hard coding in server" |
| 215 | if isinstance(wav_path, dict): |
| 216 | waveform = wav_path["array"] |
| 217 | sample_rate = wav_path["sampling_rate"] |
| 218 | else: |
| 219 | waveform, sample_rate = sf.read(wav_path) |
| 220 | if sample_rate != target_sample_rate: |
| 221 | from scipy.signal import resample |
| 222 | |
| 223 | waveform = resample(waveform, int(len(waveform) * (target_sample_rate / sample_rate))) |
| 224 | return waveform, target_sample_rate |
| 225 | |
| 226 | |
| 227 | async def send( |
| 228 | manifest_item_list: list, |
| 229 | name: str, |
| 230 | triton_client: tritonclient.grpc.aio.InferenceServerClient, |
| 231 | protocol_client: types.ModuleType, |
| 232 | log_interval: int, |
| 233 | model_name: str, |
| 234 | padding_duration: int = None, |
| 235 | audio_save_dir: str = "./", |
| 236 | save_sample_rate: int = 24000, |
| 237 | ): |
| 238 | total_duration = 0.0 |
| 239 | latency_data = [] |
| 240 | task_id = int(name[5:]) |
| 241 | |
| 242 | print(f"manifest_item_list: {manifest_item_list}") |
| 243 | for i, item in enumerate(manifest_item_list): |
| 244 | if i % log_interval == 0: |
| 245 | print(f"{name}: {i}/{len(manifest_item_list)}") |
| 246 | waveform, sample_rate = load_audio(item["audio_filepath"], target_sample_rate=24000) |
| 247 | duration = len(waveform) / sample_rate |
| 248 | lengths = np.array([[len(waveform)]], dtype=np.int32) |
| 249 | |
| 250 | reference_text, target_text = item["reference_text"], item["target_text"] |
| 251 | |
| 252 | estimated_target_duration = duration / len(reference_text) * len(target_text) |
| 253 | |
| 254 | if padding_duration: |
| 255 | # padding to nearset 10 seconds |
| 256 | samples = np.zeros( |
| 257 | ( |
| 258 | 1, |
| 259 | padding_duration |
| 260 | * sample_rate |
| 261 | * ((int(estimated_target_duration + duration) // padding_duration) + 1), |
| 262 | ), |
| 263 | dtype=np.float32, |
| 264 | ) |
| 265 | |
| 266 | samples[0, : len(waveform)] = waveform |
| 267 | else: |
| 268 | samples = waveform |
| 269 | |
| 270 | samples = samples.reshape(1, -1).astype(np.float32) |
| 271 | |
| 272 | inputs = [ |
| 273 | protocol_client.InferInput("reference_wav", samples.shape, np_to_triton_dtype(samples.dtype)), |
| 274 | protocol_client.InferInput("reference_wav_len", lengths.shape, np_to_triton_dtype(lengths.dtype)), |
| 275 | protocol_client.InferInput("reference_text", [1, 1], "BYTES"), |
| 276 | protocol_client.InferInput("target_text", [1, 1], "BYTES"), |
| 277 | ] |
| 278 | inputs[0].set_data_from_numpy(samples) |
| 279 | inputs[1].set_data_from_numpy(lengths) |
| 280 | |
| 281 | input_data_numpy = np.array([reference_text], dtype=object) |
| 282 | input_data_numpy = input_data_numpy.reshape((1, 1)) |
| 283 | inputs[2].set_data_from_numpy(input_data_numpy) |
| 284 | |
| 285 | input_data_numpy = np.array([target_text], dtype=object) |
| 286 | input_data_numpy = input_data_numpy.reshape((1, 1)) |
| 287 | inputs[3].set_data_from_numpy(input_data_numpy) |
| 288 | |
| 289 | outputs = [protocol_client.InferRequestedOutput("waveform")] |
| 290 | |
| 291 | sequence_id = 100000000 + i + task_id * 10 |
| 292 | start = time.time() |
| 293 | response = await triton_client.infer(model_name, inputs, request_id=str(sequence_id), outputs=outputs) |
| 294 | |
| 295 | audio = response.as_numpy("waveform").reshape(-1) |
| 296 | |
| 297 | end = time.time() - start |
| 298 | |
| 299 | audio_save_path = os.path.join(audio_save_dir, f"{item['target_audio_path']}.wav") |
| 300 | sf.write(audio_save_path, audio, save_sample_rate, "PCM_16") |
| 301 | |
| 302 | actual_duration = len(audio) / save_sample_rate |
| 303 | latency_data.append((end, actual_duration)) |
| 304 | total_duration += actual_duration |
| 305 | |
| 306 | return total_duration, latency_data |
| 307 | |
| 308 | |
| 309 | def load_manifests(manifest_path): |
| 310 | with open(manifest_path, "r") as f: |
| 311 | manifest_list = [] |
| 312 | for line in f: |
| 313 | assert len(line.strip().split("|")) == 4 |
| 314 | utt, prompt_text, prompt_wav, gt_text = line.strip().split("|") |
| 315 | utt = Path(utt).stem |
| 316 | # gt_wav = os.path.join(os.path.dirname(manifest_path), "wavs", utt + ".wav") |
| 317 | if not os.path.isabs(prompt_wav): |
| 318 | prompt_wav = os.path.join(os.path.dirname(manifest_path), prompt_wav) |
| 319 | manifest_list.append( |
| 320 | { |
| 321 | "audio_filepath": prompt_wav, |
| 322 | "reference_text": prompt_text, |
| 323 | "target_text": gt_text, |
| 324 | "target_audio_path": utt, |
| 325 | } |
| 326 | ) |
| 327 | return manifest_list |
| 328 | |
| 329 | |
| 330 | def split_data(data, k): |
| 331 | n = len(data) |
| 332 | if n < k: |
| 333 | print(f"Warning: the length of the input list ({n}) is less than k ({k}). Setting k to {n}.") |
| 334 | k = n |
| 335 | |
| 336 | quotient = n // k |
| 337 | remainder = n % k |
| 338 | |
| 339 | result = [] |
| 340 | start = 0 |
| 341 | for i in range(k): |
| 342 | if i < remainder: |
| 343 | end = start + quotient + 1 |
| 344 | else: |
| 345 | end = start + quotient |
| 346 | |
| 347 | result.append(data[start:end]) |
| 348 | start = end |
| 349 | |
| 350 | return result |
| 351 | |
| 352 | |
| 353 | async def main(): |
| 354 | args = get_args() |
| 355 | url = f"{args.server_addr}:{args.server_port}" |
| 356 | |
| 357 | triton_client = grpcclient.InferenceServerClient(url=url, verbose=False) |
| 358 | protocol_client = grpcclient |
| 359 | |
| 360 | if args.reference_audio: |
| 361 | args.num_tasks = 1 |
| 362 | args.log_interval = 1 |
| 363 | manifest_item_list = [ |
| 364 | { |
| 365 | "reference_text": args.reference_text, |
| 366 | "target_text": args.target_text, |
| 367 | "audio_filepath": args.reference_audio, |
| 368 | "target_audio_path": "test", |
| 369 | } |
| 370 | ] |
| 371 | elif args.huggingface_dataset: |
| 372 | import datasets |
| 373 | |
| 374 | dataset = datasets.load_dataset( |
| 375 | args.huggingface_dataset, |
| 376 | split=args.split_name, |
| 377 | trust_remote_code=True, |
| 378 | ) |
| 379 | manifest_item_list = [] |
| 380 | for i in range(len(dataset)): |
| 381 | manifest_item_list.append( |
| 382 | { |
| 383 | "audio_filepath": dataset[i]["prompt_audio"], |
| 384 | "reference_text": dataset[i]["prompt_text"], |
| 385 | "target_audio_path": dataset[i]["id"], |
| 386 | "target_text": dataset[i]["target_text"], |
| 387 | } |
| 388 | ) |
| 389 | else: |
| 390 | manifest_item_list = load_manifests(args.manifest_path) |
| 391 | |
| 392 | args.num_tasks = min(args.num_tasks, len(manifest_item_list)) |
| 393 | manifest_item_list = split_data(manifest_item_list, args.num_tasks) |
| 394 | |
| 395 | os.makedirs(args.log_dir, exist_ok=True) |
| 396 | tasks = [] |
| 397 | start_time = time.time() |
| 398 | for i in range(args.num_tasks): |
| 399 | task = asyncio.create_task( |
| 400 | send( |
| 401 | manifest_item_list[i], |
| 402 | name=f"task-{i}", |
| 403 | triton_client=triton_client, |
| 404 | protocol_client=protocol_client, |
| 405 | log_interval=args.log_interval, |
| 406 | model_name=args.model_name, |
| 407 | audio_save_dir=args.log_dir, |
| 408 | padding_duration=1, |
| 409 | save_sample_rate=24000, |
| 410 | ) |
| 411 | ) |
| 412 | tasks.append(task) |
| 413 | |
| 414 | ans_list = await asyncio.gather(*tasks) |
| 415 | |
| 416 | end_time = time.time() |
| 417 | elapsed = end_time - start_time |
| 418 | |
| 419 | total_duration = 0.0 |
| 420 | latency_data = [] |
| 421 | for ans in ans_list: |
| 422 | total_duration += ans[0] |
| 423 | latency_data += ans[1] |
| 424 | |
| 425 | rtf = elapsed / total_duration |
| 426 | |
| 427 | s = f"RTF: {rtf:.4f}\n" |
| 428 | s += f"total_duration: {total_duration:.3f} seconds\n" |
| 429 | s += f"({total_duration / 3600:.2f} hours)\n" |
| 430 | s += f"processing time: {elapsed:.3f} seconds ({elapsed / 3600:.2f} hours)\n" |
| 431 | |
| 432 | latency_list = [chunk_end for (chunk_end, chunk_duration) in latency_data] |
| 433 | latency_ms = sum(latency_list) / float(len(latency_list)) * 1000.0 |
| 434 | latency_variance = np.var(latency_list, dtype=np.float64) * 1000.0 |
| 435 | s += f"latency_variance: {latency_variance:.2f}\n" |
| 436 | s += f"latency_50_percentile_ms: {np.percentile(latency_list, 50) * 1000.0:.2f}\n" |
| 437 | s += f"latency_90_percentile_ms: {np.percentile(latency_list, 90) * 1000.0:.2f}\n" |
| 438 | s += f"latency_95_percentile_ms: {np.percentile(latency_list, 95) * 1000.0:.2f}\n" |
| 439 | s += f"latency_99_percentile_ms: {np.percentile(latency_list, 99) * 1000.0:.2f}\n" |
| 440 | s += f"average_latency_ms: {latency_ms:.2f}\n" |
| 441 | |
| 442 | print(s) |
| 443 | if args.manifest_path: |
| 444 | name = Path(args.manifest_path).stem |
| 445 | elif args.split_name: |
| 446 | name = args.split_name |
| 447 | with open(f"{args.log_dir}/rtf-{name}.txt", "w") as f: |
| 448 | f.write(s) |
| 449 | |
| 450 | stats = await triton_client.get_inference_statistics(model_name="", as_json=True) |
| 451 | write_triton_stats(stats, f"{args.log_dir}/stats_summary-{name}.txt") |
| 452 | |
| 453 | metadata = await triton_client.get_model_config(model_name=args.model_name, as_json=True) |
| 454 | with open(f"{args.log_dir}/model_config-{name}.json", "w") as f: |
| 455 | json.dump(metadata, f, indent=4) |
| 456 | |
| 457 | |
| 458 | if __name__ == "__main__": |
| 459 | asyncio.run(main()) |
| 460 |