40 lines
2.4 KiB
Python
40 lines
2.4 KiB
Python
import ast,asyncio,json,pathlib,types
|
|
root=pathlib.Path(json.loads(pathlib.Path('/home/ps/huobao-gpu-selection/idle-baseline-fix-current.json').read_text())['directory'])
|
|
class TransportError(Exception): pass
|
|
class Client:
|
|
def __init__(self,**kwargs):pass
|
|
async def __aenter__(self):return self
|
|
async def __aexit__(self,*args):pass
|
|
async def get(self,url,**kwargs):
|
|
if fail:raise TransportError()
|
|
if url=='http://legacy/queue':value=old_queue
|
|
elif url.endswith('/system_stats'):value=stats
|
|
elif url.endswith('/object_info/ApplyVDNH3'):value={'ApplyVDNH3':{}}
|
|
else:value={'queue_running':[],'queue_pending':[[1,'new-job']]}
|
|
return types.SimpleNamespace(raise_for_status=lambda:None,json=lambda:value)
|
|
async def post(self,url,**kwargs):
|
|
assert url=='http://legacy/free';assert kwargs['json']=={'unload_models':True,'free_memory':True}
|
|
calls.append(url);return types.SimpleNamespace(raise_for_status=lambda:None)
|
|
httpx=types.SimpleNamespace(AsyncClient=Client,HTTPError=TransportError)
|
|
Provider=object
|
|
providers=lambda:{'gpu1':types.SimpleNamespace(base_url='http://legacy')}
|
|
provider_headers=lambda p:{}
|
|
target=types.SimpleNamespace(id='gpu1_8189',base_url='http://new')
|
|
for variant in ('general','vdn'):
|
|
source=(root/(variant+'.after.py')).read_text()
|
|
node=next(n for n in ast.parse(source).body if isinstance(n,ast.AsyncFunctionDef) and n.name=='provider_queue_depth')
|
|
exec(compile(ast.Module(body=[node],type_ignores=[]),variant,'exec'),globals())
|
|
async def check():
|
|
global old_queue,stats,calls,fail
|
|
fail=False;calls=[];old_queue={};stats={'devices':[{'torch_vram_total':2348810240}]}
|
|
assert await provider_queue_depth(target)==1 and not calls
|
|
old_queue={'queue_running':[[1,'old-job']]};assert await provider_queue_depth(target) is None and not calls
|
|
old_queue={'queue_pending':[[1,'old-job']]};assert await provider_queue_depth(target) is None and not calls
|
|
old_queue={};stats={'devices':[{'torch_vram_total':10*1024**3}]}
|
|
assert await provider_queue_depth(target) is None and calls==['http://legacy/free']
|
|
for devices in (None,[],[{}],[{'torch_vram_total':-1}],[{'torch_vram_total':'unknown'}],[{'torch_vram_total':float('nan')}]):
|
|
stats={'devices':devices};assert await provider_queue_depth(target) is None
|
|
fail=True;assert await provider_queue_depth(target) is None
|
|
asyncio.run(check())
|
|
print('PASS',variant,'idle baseline allowed; active/pending blocked; large cache freed; malformed stats and network failure blocked')
|