from concurrent.futures import ThreadPoolExecutor, ProcessPoolExecutor
from tensorflow_serving.apis import predict_pb2
from tensorflow_serving.apis import prediction_service_pb2_grpc
from statistics import mean
def debug_print(message, debug=False):
def generate_single_request(batch_size=1):
request = predict_pb2.PredictRequest()
request.model_spec.name = 'wdl'
request.model_spec.signature_name = "serving_default"
request.inputs[f"I{i}"].CopyFrom(
tf.make_tensor_proto(np.random.uniform(0,1, [batch_size,1]), dtype=tf.float32)
categories = [f"cat_{i}_{j}" for j in range(5)]
batch_values = [np.random.choice(categories) for _ in range(batch_size)]
request.inputs[f"C{i}"].CopyFrom(
tf.make_tensor_proto(batch_values, dtype=tf.string)
def generate_all_requests(num_requests, batch_size=1, num_workers=None, debug=False):
num_workers = min(multiprocessing.cpu_count(), num_requests)
debug_print(f"{num_workers} processes are used to generate request data in parallel...", debug)
requests_per_worker = [num_requests // num_workers] * num_workers
for i in range(num_requests % num_workers):
requests_per_worker[i] +=1
with ProcessPoolExecutor(max_workers=num_workers) as executor:
for count in requests_per_worker:
seed = np.random.randint(0, 1000000)
future = executor.submit(
_worker_generate_requests, count, batch_size, seed
requests.extend(future.result())
def _worker_generate_requests(num_requests, batch_size, seed):
return [generate_single_request(batch_size) for _ in range(num_requests)]
async def get_request(requests, request_rate, burstiness, max_batches, batch_counter, counter_lock, start_time, duration, debug=False):
assert burstiness >0, "Burstiness factor must be positive."
theta = 1.0 / (request_rate * burstiness)
total_requests = len(requests)
current_time = time.time()
elapsed_time = current_time - start_time
if elapsed_time > duration:
debug_print(f"[get_request] Time limit triggered: {elapsed_time:.4 f} s > {duration} s", debug)
async with counter_lock:
if request_idx >= total_requests or batch_counter[0] >= max_batches:
reason = "Request index exhausted" if request_idx >= total_requests else "Batch count upper limit reached"
debug_print(f"[get_request] Counter limit triggered: {reason}, Batch count: {batch_counter[0]}/{max_batches}", debug)
current_request = requests[request_idx]
if request_rate != float("inf"):
interval = np.random.gamma(shape=burstiness, scale=theta)
await asyncio.sleep(interval)
def call_grpc_stub(stub, request):
start_time = time.time()
result = stub.Predict(request, 10.0)
return end_time - start_time
class ConcurrencyMonitor:
def __init__(self, max_concurrency):
self.max_concurrency = max_concurrency
self.lock = asyncio.Lock()
async def increment(self):
if self.current > self.highest:
self.highest = self.current
async def decrement(self):
async def get_current(self):
async def get_highest(self):
async def concurrency_monitor(monitor, interval=3, debug=False):
"""Periodically print concurrent usage information."""
current = await monitor.get_current()
highest = await monitor.get_highest()
print(f"[monitoring] Current concurrency: {current}/{monitor.max_concurrency}, historical maximum concurrency: {highest}")
await asyncio.sleep(interval)
except asyncio.CancelledError:
debug_print("[Monitoring] The monitoring task has stopped", debug)
debug_print(f"[Monitoring] Monitoring task exception: {e}", debug)
def __init__(self, max_workers, stub, request_generator, semaphore, latencies, monitor, loop, start_time, duration, debug=False):
self.max_workers = max_workers
self.request_generator = request_generator
self.semaphore = semaphore
self.latencies = latencies
self.start_time = start_time
self.duration = duration
self.queue = asyncio.Queue()
for i in range(min(max_workers,5)):
def _create_worker(self):
worker_id = len(self.workers) +1
task = asyncio.create_task(self._worker_loop(worker_id))
self.workers.append(task)
debug_print(f"[WorkerPool] Create worker coroutines. {worker_id}, total:{len(self.workers)}", self.debug)
async def _worker_loop(self, worker_id):
debug_print(f"[worker-{worker_id}] start", self.debug)
request = await asyncio.wait_for(self.queue.get(), timeout=1.0)
current_time = time.time()
if current_time - self.start_time > self.duration:
debug_print(f"[worker-{worker_id}] Exit: Time limit", self.debug)
await self.monitor.increment()
async with self.semaphore:
remaining_time = self.duration - (time.time() - self.start_time)
debug_print(f"[worker-{worker_id}] Insufficient remaining time for the request: {remaining_time:.4 f} s", self.debug)
latency = await asyncio.wait_for(
self.loop.run_in_executor(ThreadPoolExecutor(), call_grpc_stub, self.stub, request),
self.latencies.append(latency)
debug_print(f"[worker-{worker_id}] Request completed, latency: {latency:.4 f} s", self.debug)
except asyncio.TimeoutError:
debug_print(f"[worker-{worker_id}] Request timeout: remaining time {remaining_time:.4 f} s", self.debug)
debug_print(f"[worker-{worker_id}] Request exception: {e}", self.debug)
await self.monitor.decrement()
except asyncio.TimeoutError:
async def submit(self, request):
if self.queue.qsize() >= len(self.workers) and len(self.workers) < self.max_workers:
await self.queue.put(request)
for _ in range(len(self.workers)):
await self.queue.put(None)
await asyncio.gather(*self.workers, return_exceptions=True)
debug_print(f"[WorkerPool] All worker coroutines have stopped, total:{len(self.workers)}", self.debug)
async def async_call_grpc(stub, requests, request_rate, burstiness, max_concurrency, duration, max_batches, debug=False):
start_time = time.time()
semaphore = asyncio.Semaphore(max_concurrency)
loop = asyncio.get_running_loop()
counter_lock = asyncio.Lock()
stop_event = asyncio.Event()
request_generator = get_request(requests, request_rate, burstiness, max_batches, batch_counter, counter_lock, start_time, duration, debug)
concurrency_monitor_obj = ConcurrencyMonitor(max_concurrency)
monitor_task = asyncio.create_task(concurrency_monitor(concurrency_monitor_obj,3, debug))
worker_pool = WorkerPool(
max_workers=max_concurrency,
request_generator=request_generator,
monitor=concurrency_monitor_obj,
async for request in request_generator:
if time.time() - start_time > duration or stop_event.is_set():
await worker_pool.submit(request)
debug_print(f"[feeder] exception: {e}", debug)
await worker_pool.close()
feeder_task = asyncio.create_task(feeder())
await asyncio.gather(feeder_task, return_exceptions=True)
highest_concurrency = await concurrency_monitor_obj.get_highest()
debug_print(f"[async_call_grpc] Highest concurrency: {highest_concurrency}/{max_concurrency}", debug)
valid_latencies = [t for t in latencies if t <= duration]
return batch_counter[0], valid_latencies, highest_concurrency
parser = argparse.ArgumentParser(description='TensorFlow Serving stress test tool')
parser.add_argument('--request-rate', type=float, default=10.0, help='Target QPS')
parser.add_argument('--burstiness', type=float, default=1.0, help='Precipitating factor')
parser.add_argument('--duration', type=float, default=60.0, help='Stress test duration (seconds)')
parser.add_argument('--server', type=str, default='localhost:8500', help='Service address')
parser.add_argument('--num-requests', type=int, default=1000, help='Total request batch count')
parser.add_argument('--max-concurrency', type=int, default=50, help='Maximum concurrency')
parser.add_argument('--batch-size', type=int, default=1, help='Sample batch size for each request')
parser.add_argument('--processes', type=int, default=None, help='Number of request-generating processes')
parser.add_argument('--debug', action='store_true', help='Enable debugging information and monitoring')
args = parser.parse_args()
print(f"Stress test configuration:")
print(f" Request rate: {args.request_rate} QPS")
print(f" Stress test duration: {args.duration}second")
print(f" Maximum number of requests: {args.num_requests} batches")
print(f" Maximum concurrency: {args.max_concurrency}")
print(f" Batch size: {args.batch_size}")
print(f" Target server: {args.server}")
print(f" Debugging mode:{'Enable' if args.debug else 'Disable'}")
generate_count = max(args.num_requests, args.max_concurrency *2)
print(f"Generating {generate_count} request data items, batch size: {args.batch_size}...")
requests = generate_all_requests(generate_count, args.batch_size, args.processes, args.debug)
print(f"Request data generation completed, duration: {time.time()-start_gen:.2 f} seconds")
channel = grpc.insecure_channel(args.server)
stub = prediction_service_pb2_grpc.PredictionServiceStub(channel)
print(f"Start the stress test. Target QPS={args.request_rate}, concurrency={args.max_concurrency}, batch size={args.batch_size}")
start_time = time.time()
completed_batches, latencies, highest_concurrency = await async_call_grpc(
stub, requests, args.request_rate, args.burstiness,
args.max_concurrency, args.duration, args.num_requests,
total_time = max(time.time() - start_time, 0.001)
throughput = completed_batches / total_time
p99_latency = np.percentile(latencies, 99) if latencies else 0
avg_latency = np.mean(latencies) if latencies else 0
total_samples = completed_batches * args.batch_size
samples_per_second = total_samples / total_time
print("\n============ Stress test results ============")
print(f"Completed requests (batches): {completed_batches} / {args.num_requests}")
print(f"Number of completed samples: {total_samples}")
print(f"Total duration: {total_time:.2 f} seconds")
print(f"Request throughput (QPS): {throughput:.2 f}")
print(f"Sample throughput (samples/second): {samples_per_second:.2 f}")
print(f"Average latency: {avg_latency:.4 f} seconds")
print(f"P99 latency: {p99_latency:.4 f} seconds")
print(f"Maximum concurrency: {args.max_concurrency}")
print(f"Actual maximum concurrency: {highest_concurrency}")
print(f"Batch size: {args.batch_size}")
if __name__ == "__main__":
if hasattr(multiprocessing, 'freeze_support'):
multiprocessing.freeze_support()