68 lines
2.2 KiB
Python
68 lines
2.2 KiB
Python
#!/usr/bin/env python3
|
|
import asyncio
|
|
import aiohttp
|
|
import base64
|
|
import time
|
|
import json
|
|
from typing import List
|
|
|
|
API_URL = "http://127.0.0.1:8008/api/v1/inpaint"
|
|
|
|
|
|
def img_to_b64(path: str) -> str:
|
|
with open(path, "rb") as f:
|
|
return "data:image/png;base64," + base64.b64encode(f.read()).decode("utf-8")
|
|
|
|
|
|
async def post_inpaint(session: aiohttp.ClientSession, image_b64: str, mask_b64: str):
|
|
payload = {
|
|
"image": image_b64,
|
|
"mask": mask_b64,
|
|
"model_name": "simple-lama",
|
|
}
|
|
t0 = time.time()
|
|
async with session.post(API_URL, json=payload) as resp:
|
|
_ = await resp.read()
|
|
dt = time.time() - t0
|
|
return resp.status, dt
|
|
|
|
|
|
async def run_scenario(concurrency: int, total: int, img_path: str, mask_path: str):
|
|
image_b64 = img_to_b64(img_path)
|
|
mask_b64 = img_to_b64(mask_path)
|
|
async with aiohttp.ClientSession() as session:
|
|
tasks = []
|
|
# 워밍업: 배치 2를 유도하기 위해 빠르게 2개
|
|
for _ in range(2):
|
|
tasks.append(asyncio.create_task(post_inpaint(session, image_b64, mask_b64)))
|
|
results = await asyncio.gather(*tasks)
|
|
print("Warmup:", results)
|
|
|
|
# 본 테스트: 동시성(concurrency)로 total개 요청 보내기
|
|
pending = 0
|
|
idx = 0
|
|
inflight: List[asyncio.Task] = []
|
|
latencies = []
|
|
while idx < total or inflight:
|
|
while idx < total and len(inflight) < concurrency:
|
|
inflight.append(asyncio.create_task(post_inpaint(session, image_b64, mask_b64)))
|
|
idx += 1
|
|
done, inflight = await asyncio.wait(inflight, return_when=asyncio.FIRST_COMPLETED)
|
|
for d in done:
|
|
code, dt = d.result()
|
|
latencies.append(dt)
|
|
print("count=", len(latencies), "avg=", sum(latencies)/len(latencies), "max=", max(latencies))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
import argparse
|
|
p = argparse.ArgumentParser()
|
|
p.add_argument("--img", required=True)
|
|
p.add_argument("--mask", required=True)
|
|
p.add_argument("--concurrency", type=int, default=2)
|
|
p.add_argument("--total", type=int, default=10)
|
|
args = p.parse_args()
|
|
asyncio.run(run_scenario(args.concurrency, args.total, args.img, args.mask))
|
|
|
|
|