Nvidia Triton Server吞吐量计算的线程实现优化问题

首先得说，你这个思路方向是对的——直接开5000个线程肯定会把系统线程池榨干，用信号量控制并发数完全符合导师要求的“线程完成再建新线程”的核心逻辑。不过确实可以让代码更简洁，同时还能提升可靠性。

先聊聊你当前代码里的小隐患：如果worker函数里抛出异常，信号量可能没被释放，导致后续线程一直阻塞，最后死锁。所以我们可以先把worker函数的信号量释放逻辑用try-finally包起来，确保无论成功失败都能释放信号量。

然后优化主函数的代码，让逻辑更紧凑，同时保留你需要的吞吐量计算逻辑。这里给你调整后的版本：

import threading
import time
from tqdm import tqdm
import numpy as np

# 注意：success_count要做线程安全处理，不然多线程下计数会不准
success_count = 0
count_lock = threading.Lock()

def infer_request(url, image, model_name):
    # 这里替换成你的实际推理逻辑
    # 模拟推理耗时，实际中替换为Triton的请求代码
    time.sleep(0.001)
    global success_count
    with count_lock:
        success_count += 1

def worker(i, semaphore, url, model_name):
    try:
        image = get_image(i)
        infer_request(url, image, model_name)
    finally:
        # 不管成功失败，都释放信号量，避免死锁
        semaphore.release()

def send_batched_reqs():
    NUM_REQS = 5000
    max_concurrent = 500  
    semaphore = threading.Semaphore(max_concurrent)
    URL = "你的Triton服务地址"
    MODEL_NAME = "你的模型名称"
    
    start_time = time.time()
    threads = []

    for i in tqdm(range(NUM_REQS), desc="Sending requests"):
        semaphore.acquire()
        t = threading.Thread(
            target=worker,
            args=(i, semaphore, URL, MODEL_NAME)
        )
        threads.append(t)
        t.start()

    # 等待所有线程完成
    for t in threads:
        t.join()

    end_time = time.time()
    total_time = end_time - start_time
    throughput = NUM_REQS / total_time
    print(f"Success Count: {success_count}")
    print(f"Total Time: {total_time:.2f} sec, Throughput: {throughput:.2f} req/sec")
    return total_time, throughput

优化的几个关键细节： #

• 线程安全的计数：给success_count加了锁，避免多线程同时修改导致计数错误，这个是你原来代码里容易忽略的坑
• 可靠的信号量释放：用try-finally包裹业务逻辑，确保即使推理过程抛出异常，信号量也能被释放，不会导致后续线程阻塞
• 参数传递更清晰：把URL和MODEL_NAME作为参数传给worker，避免全局变量的滥用，代码更易维护
• 逻辑更紧凑：去掉了不必要的冗余代码，主函数的流程更清晰

另外，如果你想模拟真实的请求间隔，原来的指数分布sleep可以加回来，但要放在semaphore.acquire()之后，这样不会影响并发数的控制。

备注：内容来源于stack exchange，提问作者Tanay Joshi