40 lines
1.8 KiB
Python
40 lines
1.8 KiB
Python
import ast,asyncio
|
|
from dataclasses import dataclass,replace
|
|
from pathlib import Path
|
|
@dataclass(frozen=True)
|
|
class Provider:
|
|
id:str
|
|
base_url:str
|
|
class HTTPException(Exception):
|
|
def __init__(self,status_code,detail):self.status_code=status_code;self.detail=detail
|
|
pool={k:Provider(k,'http://original-'+k+':8188') for k in ('default','gpu1','gpu1_8189','gpu3')}
|
|
providers=lambda:pool
|
|
admission_ids=lambda _:set(pool)
|
|
auto_balance_enabled=lambda:True
|
|
provider_reservations={}
|
|
depths={'default':8,'gpu1':0,'gpu1_8189':1,'gpu3':4}
|
|
async def provider_queue_depth(p):return depths[p.id]
|
|
class Env:
|
|
@staticmethod
|
|
def getenv(*a):return 'configured'
|
|
os=Env()
|
|
for variant in ('general','vdn'):
|
|
source=Path(__file__).with_name(variant+'-normalizer-panels.py').read_text(encoding='utf-8')
|
|
node=next(n for n in ast.parse(source).body if isinstance(n,ast.AsyncFunctionDef) and n.name=='select_provider')
|
|
exec(compile(ast.Module(body=[node],type_ignores=[]),variant,'exec'),globals())
|
|
async def test():
|
|
selected=await select_provider('default',False,True)
|
|
assert selected.id=='gpu1_8189' and selected.base_url=='http://192.168.110.111:8189'
|
|
assert (await select_provider('default',False,False)).id=='gpu1'
|
|
for worker,port in [('default',8188),('gpu1_8189',8189),('gpu3',8191)]:
|
|
chosen=await select_provider(worker,True,True)
|
|
assert chosen.id==worker and chosen.base_url.endswith(':'+str(port))
|
|
depths['gpu1_8189']=None
|
|
assert (await select_provider('default',False,True)).id=='gpu3'
|
|
try:await select_provider('gpu1_8189',True,True)
|
|
except HTTPException as e:assert e.status_code==503
|
|
else:raise AssertionError('Pinned unavailable worker fell back')
|
|
depths['gpu1_8189']=1
|
|
asyncio.run(test())
|
|
print('PASS',variant,': actual 8188/8189/8191 targets, strict pin, auto pool excludes legacy, unrelated clients retain original pool')
|