import numpy as np from concurrent.futures import ThreadPoolExecutor import time import json import tritonclient.grpc as grpcclient from tritonclient.utils import * import queue from functools import partial import random run_multiple_tests = False resp_list = [] scode_breakup = {} def np_to_server_dtype(np_dtype): if np_dtype == bool: return "BOOL" elif np_dtype == np.int8: return "INT8" elif np_dtype == np.int16: return "INT16" elif np_dtype == np.int32: return "INT32" elif np_dtype == np.int64: return "INT64" elif np_dtype == np.uint8: return "UINT8" elif np_dtype == np.uint16: return "UINT16" elif np_dtype == np.uint32: return "UINT32" elif np_dtype == np.uint64: return "UINT64" elif np_dtype == np.float16: return "FP16" elif np_dtype == np.float32: return "FP32" elif np_dtype == np.float64: return "FP64" elif np_dtype == np.object_ or np_dtype.type == np.bytes_: return "BYTES" return None class UserData: def __init__(self): self._completed_requests = queue.Queue() def callback(user_data, result, error): if error: user_data._completed_requests.put(error) else: user_data._completed_requests.put(result) def prepare_tensor(name: str, data: np.ndarray): server_input = grpcclient.InferInput(name=name, shape=data.shape, datatype=np_to_server_dtype(data.dtype)) server_input.set_data_from_numpy(data) return server_input def process_and_send_request(sample_request): prompt = sample_request['prompt'] negative_prompt = sample_request['negative_prompt'] if 'negative_prompt' in sample_request else None height = sample_request['height'] if 'height' in sample_request else None width = sample_request['width'] if 'width' in sample_request else None num_images_per_prompt = sample_request['num_images_per_prompt'] if 'num_images_per_prompt' in sample_request else 1 num_inference_steps = sample_request['num_inference_steps'] if 'num_inference_steps' in sample_request else 20 image = sample_request['image'] if 'image' in sample_request else None mask_image = sample_request['mask_image'] if 'mask_image' in sample_request else None control_images = sample_request['control_images'] if 'control_images' in sample_request else None control_weightages = sample_request['control_weightages'] if 'control_weightages' in sample_request else None control_modes = sample_request['control_modes'] if 'control_modes' in sample_request else None seed = sample_request['seed'] if 'seed' in sample_request else -1 guidance_scale = sample_request['guidance_scale'] if 'guidance_scale' in sample_request else 7.5 strength = sample_request['strength'] if 'strength' in sample_request else 1 scheduler = sample_request['scheduler'] if 'scheduler' in sample_request else "EULER-A" model_type = sample_request['model_type'] if 'model_type' in sample_request else None lora_weights = sample_request['lora_weights'] if 'lora_weights' in sample_request else None control_guidance_start = sample_request['control_guidance_start'] if 'control_guidance_start' in sample_request else None control_guidance_end = sample_request['control_guidance_end'] if 'control_guidance_end' in sample_request else None inputs = [] inputs.append(prepare_tensor("prompt", np.array([prompt], dtype=np.object_))) if negative_prompt is not None: inputs.append(prepare_tensor("negative_prompt", np.array([negative_prompt], dtype=np.object_))) if height is not None: inputs.append(prepare_tensor("height", np.array([height], dtype=np.int32))) if width is not None: inputs.append(prepare_tensor("width", np.array([width], dtype=np.int32))) if num_images_per_prompt is not None: inputs.append(prepare_tensor("num_images_per_prompt", np.array([num_images_per_prompt], dtype=np.int32))) if num_inference_steps is not None: inputs.append(prepare_tensor("num_inference_steps", np.array([num_inference_steps], dtype=np.int32))) if image is not None: inputs.append(prepare_tensor("image", np.array([image], dtype=np.object_))) if mask_image is not None: inputs.append(prepare_tensor("mask_image", np.array([mask_image], dtype=np.object_))) if seed is not None: inputs.append(prepare_tensor("seed", np.array([seed], dtype=np.int64))) if guidance_scale is not None: inputs.append(prepare_tensor("guidance_scale", np.array([guidance_scale], dtype=np.float32))) if model_type is not None: inputs.append(prepare_tensor("model_type", np.array([model_type], dtype=np.object_))) if strength is not None: inputs.append(prepare_tensor("strength", np.array([strength], dtype=np.float32))) if scheduler is not None: inputs.append(prepare_tensor("scheduler", np.array([scheduler], dtype=np.object_))) if control_images is not None: inputs.append(prepare_tensor("control_images", np.array([control_images], dtype=np.object_))) if control_weightages is not None: inputs.append(prepare_tensor("control_weightages", np.array([control_weightages], dtype=np.float32))) if control_modes is not None: inputs.append(prepare_tensor("control_modes", np.array([control_modes], dtype=np.int32))) if lora_weights is not None: inputs.append(prepare_tensor("lora_weights", np.array([lora_weights], dtype=np.object_))) if control_guidance_start is not None: inputs.append(prepare_tensor("control_guidance_start", np.array([control_guidance_start], dtype=np.float32))) if control_guidance_end is not None: inputs.append(prepare_tensor("control_guidance_end", np.array([control_guidance_end], dtype=np.float32))) outputs = [ grpcclient.InferRequestedOutput("response_id"), grpcclient.InferRequestedOutput("time_taken"), grpcclient.InferRequestedOutput("load_lora"), grpcclient.InferRequestedOutput("output_image_urls"), grpcclient.InferRequestedOutput("error"), # grpcclient.InferRequestedOutput("mega_pixel") ] user_data = UserData() st = time.time() mega_pixel = 0 url = "localhost:8002" with grpcclient.InferenceServerClient(url=url, ssl=False) as triton_client: triton_client.start_stream(callback=partial(callback, user_data)) triton_client.async_stream_infer( model_name="flux", inputs=inputs, outputs=outputs, ) et = time.time() response = user_data._completed_requests.get() print(response) # Check if response is an error (InferenceServerException) if hasattr(response, 'message'): # This is an error response print(f"Server error: {response}") output_image_urls = [] inference_time = 0 lora_time = 0 response_id = None mega_pixel = 0 error = str(response) sCode = 500 else: # This is a successful response try: inference_time = 0 lora_time = 0 response_id = None inference_time = response.as_numpy("time_taken").item() lora_time = response.as_numpy("load_lora").item() response_id = response.as_numpy("response_id").item().decode() if response.as_numpy("response_id").item() else None output_image_urls = response.as_numpy("output_image_urls").tolist() if response.as_numpy("output_image_urls") is not None else [] mega_pixel = response.as_numpy("mega_pixel").item().decode() if response.as_numpy("mega_pixel") is not None else "0" error_tensor = response.as_numpy("error") error = error_tensor.item().decode() if error_tensor is not None and error_tensor.item() else None sCode = 200 except Exception as e: print(f"Error processing response: {e}") output_image_urls = [] inference_time = 0 lora_time = 0 response_id = None mega_pixel = 0 error = str(e) sCode = 500 results = { "response_id": response_id, "total_time_taken": et-st, "inference_time_taken": inference_time, "loading_lora_time": lora_time, "output_image_urls": output_image_urls, "error": error, "mega_pixel": 0 if mega_pixel is None else mega_pixel } print(results) if output_image_urls == []: print("No images generated") results["error"] = "No images generated" return results def warmup_and_load_lora(warmup_json_path): if warmup_json_path is None: return False with open(warmup_json_path, 'r') as f: warmup_data = json.load(f) st = time.time() for request in warmup_data: process_and_send_request(request) resp_time = time.time()-st print(f"Warmup and load lora done in {resp_time:.3f} seconds") return True def generate_jitter_window(): percent_bifer = random.randint(1,100) if percent_bifer >= 1 and percent_bifer <= 50: jitter_window = [1, 5] elif percent_bifer >= 51 and percent_bifer <= 75: jitter_window = [10, 15] else: jitter_window = [20,30] time.sleep(random.randint(jitter_window[0],jitter_window[1])) return True def predict(requests_data,percent_bifer): random_request = random.choice(requests_data) sample_request = random_request['payload'] generate_jitter_window() return process_and_send_request(sample_request) def run_single_test(requests_data,id = 0): return predict(requests_data,id) number_of_users = 1 #change here for concurrent users duration_minutes = 2 def run_concurrent_tests_cont(number_of_users, duration_minutes): start_time = time.time() end_time = start_time + duration_minutes * 60 results = [] with ThreadPoolExecutor(max_workers=number_of_users) as executor: future_to_start_time = {} while time.time() < end_time: # Submit new tasks continuously percent_bifer = random.randint(1,10) if len(future_to_start_time) < number_of_users: future = executor.submit(run_single_test, requests_data) future_to_start_time[future] = time.time() # Process completed tasks and replace them done_futures = [f for f in future_to_start_time if f.done()] for future in done_futures: response_time = future.result() results.append(response_time) del future_to_start_time[future] # Wait for any remaining tasks to finish for future in future_to_start_time: results.append(future.result()) p25 = np.percentile(results, 25) p50 = np.percentile(results, 50) p90 = np.percentile(results, 90) p99 = np.percentile(results, 99) avg = sum(results) / len(results) with open(f"result_dump_{number_of_users}_{duration_minutes}.json", "w") as f: f.write( json.dumps(resp_list, indent=4) ) return p25, p50, p90, p99, avg if run_multiple_tests: p25_result , p50_results, p90_resutls, p99_results, avg = run_concurrent_tests_cont(number_of_users,duration_minutes) load_lora_time = warmup_and_load_lora(requests_data) print(f"25th Percentile: {p25_result:.3f} seconds") print(f"50th Percentile: {p50_results:.3f} seconds") print(f"90th Percentile: {p90_resutls:.3f} seconds") print(f"99th Percentile: {p99_results:.3f} seconds") print(f"Average Response Time: {avg:.3f} seconds") with open(f"test_results.json", "w") as f: f.write( json.dumps({ "p25": p25_result, "p50": p50_results, "p90": p90_resutls, "p99": p99_results, "avg": avg, "sCode_breakup": scode_breakup }, indent=4) ) else: payload = { "prompt": "A girl in city, 25 years old, cool, futuristic ", "negative_prompt": "blurry, low quality, distorted", "height": 1024, "width": 1024, "num_images_per_prompt": 1, "num_inference_steps": 20, "seed": 42424243, "guidance_scale": 7.0, "model_type": "txt2img" } result = process_and_send_request(payload) print(result)