39 lines
2.6 KiB
Python
39 lines
2.6 KiB
Python
"""Read-only live image workflow and scheduler checks; never submits a prompt."""
|
|
import asyncio,json
|
|
from unittest.mock import patch
|
|
import httpx,app
|
|
from workflow import compile_workflow,load_registry
|
|
async def main():
|
|
spec=load_registry()['qwen_image_2_1']
|
|
report={'auto_balance':app.auto_balance_enabled(),'admitted':sorted(app.admission_ids(app.providers())),'workers':[]}
|
|
async with httpx.AsyncClient(timeout=30) as client:
|
|
for p in app.providers().values():
|
|
r=await client.get(p.base_url+'/object_info',headers=app.provider_headers(p));r.raise_for_status();info=r.json()
|
|
modes={}
|
|
for refs in (0,1):
|
|
graph=compile_workflow(spec,{'prompt':'A ceramic teacup on a wooden table.','width':512,'height':512,'batch_size':1},['preflight.png']*refs,'image-pool-preflight')
|
|
missing=sorted({n['class_type'] for n in graph.values()}-set(info))
|
|
choices=[]
|
|
for node in graph.values():
|
|
ni=info.get(node['class_type'],{}).get('input',{})
|
|
fields={**ni.get('required',{}),**ni.get('optional',{})}
|
|
for name,value in node.get('inputs',{}).items():
|
|
field=fields.get(name)
|
|
if not isinstance(value,str) or not field or name=='image':continue
|
|
options=field[0] if isinstance(field[0],list) else field[1].get('options') if len(field)>1 and isinstance(field[1],dict) else None
|
|
if options: options=[o.get('key') if isinstance(o,dict) else o for o in options]
|
|
if options and value not in options: choices.append([node['class_type'],name,value])
|
|
modes['reference' if refs else 'text']={'missing_nodes':missing,'missing_models_or_options':choices}
|
|
report['workers'].append({'id':p.id,'queue_depth':await app.provider_queue_depth(p),'modes':modes})
|
|
async def depth(p):return depths[p.id]
|
|
with patch.object(app,'provider_queue_depth',depth):
|
|
report['scheduler']={}
|
|
for label,depths in [('gpu2_busy',{'default':2,'gpu3':0}),('gpu3_busy',{'default':0,'gpu3':2}),('gpu2_unavailable',{'default':None,'gpu3':0})]:
|
|
report['scheduler'][label]=(await app.select_provider('default')).id
|
|
depths={'default':0,'gpu3':0}
|
|
first=await app.reserve_provider('default');second=await app.reserve_provider('default')
|
|
report['scheduler']['parallel_submissions']=[first.id,second.id]
|
|
await app.release_provider(first);await app.release_provider(second)
|
|
print(json.dumps(report))
|
|
asyncio.run(main())
|