38 lines
2.3 KiB
Python
38 lines
2.3 KiB
Python
import base64,json,sqlite3,struct,sys,urllib.request
|
|
from pathlib import Path
|
|
root=Path(__file__).resolve().parent.parent
|
|
cfg=json.loads((root/'model-bridge/client.json').read_text())
|
|
base=cfg['url'].rstrip('/')
|
|
def request(path,body=None):
|
|
req=urllib.request.Request(base+path,data=json.dumps(body).encode() if body else None,headers={'Content-Type':'application/json','Authorization':'Bearer '+cfg['token']})
|
|
with urllib.request.urlopen(req,timeout=90) as response:return json.load(response)
|
|
report=root/'integration-work/image-pool-smoke.json'
|
|
if '--submit' in sys.argv:
|
|
if report.exists():raise RuntimeError('Smoke test already submitted; poll the existing job')
|
|
pool=request('/v1/images/pool')
|
|
assert pool['auto_balance'] and len(pool['configured_workers'])>=2,pool
|
|
result=request('/v1/images/generations',{'model':'qwen_image_2_1','prompt':'A photorealistic ceramic teacup on a simple wooden table, soft window light, no text, no watermark.','size':'512x512','n':1,'response_format':'url'})
|
|
data={'pool':pool,'task_id':result['task_id'],'assigned_worker':result.get('provider_id')}
|
|
report.write_text(json.dumps(data,indent=2))
|
|
else:
|
|
data=json.loads(report.read_text())
|
|
result=request('/v1/images/task/'+data['task_id'])
|
|
data.update(status=result['status'],worker=result.get('provider_id'))
|
|
if result['status']=='completed':
|
|
with urllib.request.urlopen(result['image_url'],timeout=60) as response:content=response.read()
|
|
target=root/'integration-work/image-pool-smoke.png';target.write_bytes(content)
|
|
assert content.startswith(b'\x89PNG\r\n\x1a\n')
|
|
data['dimensions']=list(struct.unpack('>II',content[16:24])); data['file']=str(target)
|
|
if result.get('error'):data['error']=result['error']
|
|
report.write_text(json.dumps(data,indent=2))
|
|
print(json.dumps(data))
|
|
with sqlite3.connect((root/'huobao-data/huobao.sqlite3').as_uri()+'?mode=ro',uri=True) as db:
|
|
rows=db.execute("SELECT id,task_id,status FROM sys_task WHERE type='image' AND model='qwen_image_2_1' AND status='completed'").fetchall()
|
|
providers={}
|
|
for tid,jid,status in rows:
|
|
try:
|
|
worker=base64.urlsafe_b64decode(jid+'='*(-len(jid)%4)).decode().split(':',1)[0]
|
|
providers[worker]=providers.get(worker,0)+1
|
|
except (ValueError,TypeError):pass
|
|
print('Previously completed Qwen images by worker:',providers)
|