Create main.py

main.py

@@ -0,0 +1,330 @@
+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 <lora:https://huggingface.co/XLabs-AI/flux-lora-collection/resolve/main/art_lora.safetensors:0.5>",
+          "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)