Files
dj/model-bridge/test_gpu_regeneration.py

132 lines
8.3 KiB
Python

"""Isolated HTTP integration: exact cancellation, GPU override, races and history."""
import concurrent.futures, hashlib, json, sqlite3, threading, time, urllib.error, urllib.request, os
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
from pathlib import Path
root = Path(__file__).resolve().parent.parent
data = root / 'integration-work' / os.environ.get('HUOBAO_TEST_FOLDER','hd-test-regenerate')
assert data.name.startswith('hd-test') and data.resolve().is_relative_to(root/'integration-work')
base = 'http://127.0.0.1:'+os.environ.get('HUOBAO_TEST_PORT','18883')+'/api/v1'
dbpath = data / 'huobao.sqlite3'
calls, jobs = [], {}
cancel_mode = 'delay'
def api(path, body=None):
req = urllib.request.Request(base + path, data=json.dumps(body).encode() if body is not None else None, headers={'Content-Type':'application/json'})
try:
with urllib.request.urlopen(req, timeout=65) as r: return r.status, json.load(r)
except urllib.error.HTTPError as e: return e.code, json.load(e)
def sql(query, params=()):
with sqlite3.connect(dbpath) as db:
db.row_factory = sqlite3.Row
return [dict(r) for r in db.execute(query, params).fetchall()]
class Handler(BaseHTTPRequestHandler):
def log_message(self, *args): pass
def reply(self, payload, status=200):
raw=json.dumps(payload).encode(); self.send_response(status); self.send_header('Content-Type','application/json'); self.end_headers(); self.wfile.write(raw)
def do_POST(self):
body=json.loads(self.rfile.read(int(self.headers.get('Content-Length',0))) or '{}')
calls.append((self.path, body))
if self.path.endswith('/cancel'):
jid=self.path.split('/')[-2]
if cancel_mode=='error': return self.reply({'error':'cancel unavailable'},503)
if cancel_mode=='finished': return self.reply({'task':{'id':jid,'status':'succeeded'}})
jobs[jid]['cancelled_at']=time.monotonic()+2
return self.reply({'task':{'id':jid,'status':'running'}})
assert self.path=='/v2/video_generation',self.path
jid='mock-'+str(len(jobs)+1)
jobs[jid]={'status':'running', 'body':body}
gpu=body.get('gpu','gpu2')
routing={'gpu':gpu,'port':{'gpu1':8189,'gpu2':8188,'gpu3':8191}[gpu], 'providerId':{'gpu1':'gpu1_8189','gpu2':'default','gpu3':'gpu3'}[gpu], 'legacy':False}
jobs[jid]['routing']=routing
return self.reply({'task_id':jid,'routing':routing})
def do_GET(self):
if self.path.endswith('.mp4'):
p=root/'integration-work/hd-test-alignment/2K.mp4' if self.path=='/hd.mp4' else data/'static/videos/hd-test-original.mp4'
raw=p.read_bytes(); self.send_response(200); self.send_header('Content-Type','video/mp4'); self.send_header('Content-Length',str(len(raw))); self.end_headers(); self.wfile.write(raw); return
jid=self.path.split('/')[-1]; job=jobs[jid]
if job.get('cancelled_at',float('inf'))<=time.monotonic():job['status']='cancelled'
task={'id':jid,'status':job['status'],'duration':4,'routing':job['routing']}
if job['status']=='succeeded':task['content']={'url':'http://127.0.0.1:18884/'+('hd.mp4' if job['body']['model']=='MiniMax-H3-VDN-H3' else 'video.mp4')}
return self.reply({'task':task})
def row(tid):return sql('SELECT * FROM sys_task WHERE id=?',(tid,))[0]
def wait(tid, predicate):
for _ in range(100):
r=row(tid)
if predicate(r): return r
time.sleep(.2)
raise AssertionError(('timeout',tid,row(tid)['status']))
def payload(**kw):return {'type':'video','storyboard_id':1,'config_id':4,'model':'MiniMax-H3','gpu':'gpu1','prompt':'A quiet lake.','duration':4,'aspect_ratio':'16:9','force_new':True,**kw}
server=ThreadingHTTPServer(('127.0.0.1',18884),Handler)
threading.Thread(target=server.serve_forever,daemon=True).start()
with sqlite3.connect(dbpath) as db:
db.execute("UPDATE ai_service_configs SET base_url='http://127.0.0.1:18884' WHERE service_type IN ('video','image')")
db.execute("UPDATE ai_service_configs SET settings=? WHERE id=4",(json.dumps({'gpu':'gpu2'}),))
db.execute("UPDATE ai_service_configs SET is_active=0 WHERE name='VDN test'")
db.execute('UPDATE dramas SET deleted_at=NULL WHERE id=1')
db.execute('UPDATE episodes SET deleted_at=NULL WHERE id=1')
original=data/'static/videos/hd-test-original.mp4'; digest=hashlib.sha256(original.read_bytes()).hexdigest()
try:
code,result=api('/tasks',payload()); assert code==201,result
old=result['data']['id']; oldrow=wait(old,lambda r:bool(r['task_id'])); jid=oldrow['task_id']
assert jobs[jid]['body']['gpu']=='gpu1'
before=len(calls)
for change in ({'gpu':'gpu0'}, {'model':'unsupported'}, {'config_id':999999}, {'replace_task_id':old+999}):
code,_=api('/tasks',payload(replace_task_id=old,**change)) if 'replace_task_id' not in change else api('/tasks',payload(**change))
assert code==400,(change,code)
assert len(calls)==before and row(old)['status']=='processing'
# Place the real background poll inside the cancellation window.
time.sleep(8)
with concurrent.futures.ThreadPoolExecutor() as pool:
first=pool.submit(api,'/tasks',payload(gpu='gpu3',replace_task_id=old))
time.sleep(.2)
duplicate=api('/tasks',payload(gpu='gpu2',replace_task_id=old))
code,result=first.result()
assert code==201,result
assert duplicate[0]==400,duplicate
new=result['data']['id']; nr=wait(new,lambda r:bool(r['task_id']))
assert jobs[jid]['status']=='cancelled'
assert row(old)['error_code']=='user_cancelled'
assert jobs[nr['task_id']]['body']['gpu']=='gpu3'
assert json.loads(nr['params'])['routing']['port']==8191
assert json.loads(nr['params'])['replacesTaskId']==old
assert json.loads(sql('SELECT settings FROM ai_service_configs WHERE id=4')[0]['settings'])['gpu']=='gpu2'
count=len(calls); assert api('/tasks',payload(replace_task_id=old))[0]==400; assert len(calls)==count
assert api('/tasks/'+str(old)+'/recover',{})[0]==400
jobs[nr['task_id']]['status']='succeeded'; wait(new,lambda r:r['status']=='completed')
assert json.loads(row(new)['params'])['routing']['gpu']=='gpu3'
assert api('/tasks/'+str(new))[1]['data']['routing']['port']==8191
assert row(old)['error_code']=='user_cancelled'
# Cancellation failure and a completed-upstream race must not create replacements.
for mode in ('error','finished'):
code,result=api('/tasks',payload()); assert code==201,result
tid=result['data']['id']; tr=wait(tid,lambda r:bool(r['task_id']))
cancel_mode=mode; count=len(jobs)
code,result=api('/tasks',payload(gpu='gpu2',replace_task_id=tid)); assert code==400,result
assert len(jobs)==count and row(tid)['status']=='processing'
jobs[tr['task_id']]['status']='failed'
wait(tid,lambda r:r['status']=='failed')
# Fixed HD config + per-request override reaches the model, original stays intact.
cancel_mode='delay'
with sqlite3.connect(dbpath) as db:db.execute('UPDATE storyboards SET video_url=? WHERE id=1',('static/videos/hd-test-original.mp4',))
config=next(c for c in api('/ai-configs')[1]['data'] if c['name'].startswith('MiniMax H3-VDN-H3 高清 2K/4K · GPU1'))
code,result=api('/hd/storyboards/1',{'resolution':'2K','config_id':config['id'],'gpu':'gpu3'})
assert code==201,result
tid=result['data']['id']; tr=wait(tid,lambda r:bool(r['task_id']))
assert jobs[tr['task_id']]['body']['gpu']=='gpu3'
jobs[tr['task_id']]['status']='succeeded';wait(tid,lambda r:r['status']=='completed')
hd=api('/hd/episodes/1')[1]['data']; shot=next(s for s in hd['items'] if s['id']==1)
assert shot['current']=='static/videos/hd-test-original.mp4'
assert shot['original']=='static/videos/hd-test-original.mp4'
assert len(hd['configs'])==4 and shot['versions'][-1]['gpu']=='gpu3'
assert hashlib.sha256(original.read_bytes()).hexdigest()==digest
result={'passed':True,'checks':['explicit GPU override','invalid replacement never cancels','targeted cancel confirmed before submission','double click guarded','stale replacement rejected','cancel failure no new task','completion race no new task','cancelled record cannot recover','HD config and override','original media preserved'], 'jobs':len(jobs)}
(root/'integration-work/gpu-regeneration-test.json').write_text(json.dumps(result,indent=2))
print(json.dumps(result))
finally:server.shutdown()