Files
dj/integration-work/test_gpu_replacement_references.py

177 lines
13 KiB
Python

"""Exercise GPU replacement preflight/cancellation with local images, no real GPU jobs."""
import base64, json, os, pathlib, re, sqlite3, subprocess, threading, time, urllib.error, urllib.request
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
root = pathlib.Path(__file__).resolve().parent.parent
data = root / 'integration-work' / ('replacement-test-' + time.strftime('%Y%m%d-%H%M%S'))
data.mkdir(); (data / 'static/images').mkdir(parents=True)
png = base64.b64decode('iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO+aN1kAAAAASUVORK5CYII=')
(data / 'static/images/reference.png').write_bytes(png)
stage = root / 'integration-work/stage-4.0.11/app'
source = (root / 'huobao-source/desktop/build/backend.mjs').read_text('utf-8')
source = source.replace('const require = __hbCreateRequire(import.meta.url)', 'const require = __hbCreateRequire(' + json.dumps(str(stage / 'resources/app.asar/build/backend.mjs')) + ')', 1)
source = re.sub(r'^import (\w+) from "(better-sqlite3|sharp)";', r'const \1 = require("\2");', source, flags=re.M)
needle = 'target.hostname === "192.168.110.111" && ["18792", "18794"].includes(target.port)'
assert source.count(needle) == 1
source = source.replace(needle, 'target.hostname === "127.0.0.1" && target.port === "18935"')
source, count = re.subn(r'(url\d*)\.hostname === "192\.168\.110\.111" && \["18792", "18794"\]\.includes\(\1.port\)',
r'\1.hostname === "127.0.0.1" && \1.port === "18935"', source)
assert count == 1
# Fake IDs are private to the loopback fixture.
source, count = re.subn(r'&& !!taskRouting\(\{ provider: config\d*\.provider, taskId \}\)', '&& !!taskId', source)
assert count == 1
(data / 'backend.mjs').write_text(source, 'utf-8')
events = []; reject = False; process = None; cancel_modes = {}; job_states = {}
class Mock(BaseHTTPRequestHandler):
def log_message(self, *args): pass
def reply(self, status, body):
raw = json.dumps(body).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['Content-Length'])))
events.append((self.path, body))
if self.path.endswith('/validate'):
images = [x['image_url']['url'] for x in body['content'] if x['type'] == 'image_url']
if reject or not images or any(not x.startswith('data:image/') for x in images):
self.reply(400, {'error': {'message': 'Local video requires inline image data'}})
else: self.reply(200, {'valid': True})
elif self.path.endswith('/cancel'):
jid=self.path.split('/')[-2];mode=cancel_modes.get(jid,'cancelled')
if mode=='wrong-id':self.reply(200,{'task':{'id':'different-job','status':'cancelled'}})
elif mode=='transport':self.reply(503,{'error':'transient'})
else:
job_states[jid]=mode
self.reply(200, {'task': {'id':jid, 'status':mode}})
else:
self.reply(200, {'task_id': 'new-' + str(len(events)), 'routing': {'gpu': body['gpu'], 'port': 8191, 'providerId': 'gpu3', 'legacy': False}})
def do_GET(self):
if self.path.startswith('/v1/images/task/'):
jid=self.path.split('/')[-1];state=job_states.get(jid,'running')
self.reply(200,{'task_id':jid,'status':'failed' if state=='cancelled' else 'processing',
'cancelled':state=='cancelled','worker_state':'running','error':{'message':'cancelled'}});return
self.reply(200, {'task': {'id': self.path.split('/')[-1], 'status': 'running', 'worker_state': 'running',
'progress': {'phase':'sampling','step':4,'totalSteps':20,'observedAt':round(time.time()*1000),'samplingMinSeconds':120,'samplingMaxSeconds':180}}})
server = ThreadingHTTPServer(('127.0.0.1', 18935), Mock)
threading.Thread(target=server.serve_forever, daemon=True).start()
env = {**os.environ, 'ELECTRON_RUN_AS_NODE': '1', 'PORT': '18934', 'HUOBAO_DATA_DIR': str(data),
'SQLITE_PATH': str(data / 'huobao.sqlite3'), 'FRONTEND_DIST': str(root / 'huobao-source/frontend/.output/public'),
'WORKSPACE_PATH': str(stage / 'resources/workspace-template')}
log = (data / 'backend.log').open('w', encoding='utf-8')
def api(method, path, body=None, expected=200):
request = urllib.request.Request('http://127.0.0.1:18934/api/v1' + path, method=method,
data=None if body is None else json.dumps(body).encode(), headers={'Content-Type': 'application/json'})
try:
with urllib.request.urlopen(request, timeout=40) as response: status = response.status; result = json.load(response)
except urllib.error.HTTPError as error: status = error.code; result = json.load(error)
assert status == expected, (path, status, result)
return result.get('data', result)
def sql(statement, args=()):
with sqlite3.connect(data / 'huobao.sqlite3') as db:
db.row_factory = sqlite3.Row
return [dict(row) for row in db.execute(statement, args)]
def replace(shot, old, **extra):
return api('POST', '/tasks', {'type': 'video', 'storyboard_id': shot, 'config_id': 1,
'model': 'MiniMax-H3', 'gpu': 'gpu3', 'replace_task_id': old, 'prompt': 'reference fixture',
'reference_image_urls': ['static/images/reference.png'], **extra.pop('body', {})}, **extra)
try:
process = subprocess.Popen([str(stage / 'HuobaoDrama.exe'), str(data / 'backend.mjs')], cwd=data,
env=env, stdout=log, stderr=log, creationflags=subprocess.CREATE_NO_WINDOW)
for _ in range(150):
assert process.poll() is None, 'Backend exited'
try: api('GET', '/health'); break
except OSError: time.sleep(.1)
ts = '2026-11-12T00:00:00Z'
sql('INSERT INTO dramas(id,title,created_at,updated_at) VALUES(1,?,?,?)', ('fixture',ts,ts))
sql('INSERT INTO episodes(id,drama_id,episode_number,title,created_at,updated_at) VALUES(1,1,1,?,?,?)', ('fixture',ts,ts))
sql('INSERT INTO ai_service_configs(id,service_type,provider,name,base_url,api_key,model,is_active,created_at,updated_at) VALUES(1,?,?,?,?,?,?,1,?,?)',
('video','minimax','fixture','http://127.0.0.1:18935','test-key',json.dumps(['MiniMax-H3']),ts,ts))
for i in range(1,5):
sql('INSERT INTO storyboards(id,episode_id,storyboard_number,created_at,updated_at) VALUES(?,1,?,?,?)',(i,i,ts,ts))
sql('INSERT INTO sys_task(id,type,storyboard_id,provider,model,status,stage,task_id,config_id,config_base_url,created_at,updated_at) VALUES(?,?,?,?,?,?,?,?,?,?,?,?)',
(i,'video',i,'minimax','MiniMax-H3','processing','running','old-'+str(i),1,'http://127.0.0.1:18935',ts,ts))
# Real local image bytes must reach preflight before cancelling exactly one old job.
result = replace(1,1,expected=201)
for _ in range(100):
if len(events) >= 3: break
time.sleep(.1)
assert [p for p,b in events[:3]] == ['/v2/video_generation/validate','/v2/video_generation/old-1/cancel','/v2/video_generation']
assert events[0][1]['content'] == events[2][1]['content']
assert base64.b64decode(events[0][1]['content'][1]['image_url']['url'].split(',')[1]) == png
assert events[0][1]['gpu'] == events[2][1]['gpu'] == 'gpu3'
assert sql('SELECT error_code FROM sys_task WHERE id=1')[0]['error_code'] == 'user_cancelled'
assert sql('SELECT status FROM sys_task WHERE id=2')[0]['status'] == 'processing'
# Missing reference fails before any upstream call, leaving the old job intact.
before = len(events)
response = replace(2,2,expected=400,body={'reference_image_urls':['static/images/missing.png']})
assert '参考图读取失败' in response['message'] and len(events) == before
assert sql('SELECT status FROM sys_task WHERE id=2')[0]['status'] == 'processing'
# Unready destination also leaves old task intact; no silent duplicate submission.
reject = True; before = len(events)
replace(3,3,expected=400)
assert len(events) == before+1 and events[-1][0].endswith('/validate')
assert sql('SELECT status FROM sys_task WHERE id=3')[0]['status'] == 'processing'
# First-frame replacements use the same conversion and role.
reject = False; before = len(events)
replace(4,4,expected=201,body={'reference_image_urls':[], 'first_frame_url':'/static/images/reference.png'})
for _ in range(100):
if len(events) >= before+3: break
time.sleep(.1)
image = events[before][1]['content'][1]
assert image['role'] == 'first_frame' and image['image_url']['url'].startswith('data:image/')
assert events[before][1]['content'] == events[before+2][1]['content']
for _ in range(140):
row=sql('SELECT params FROM sys_task WHERE id=?',(result['id'],))[0]
progress=json.loads(row['params']).get('progress')
if progress: break
time.sleep(.1)
assert progress['step']==4 and progress['totalSteps']==20
assert progress['samplingMaxSeconds']==180
# Standalone cancellation never creates another generation or touches neighbors.
before=len(events);assert api('POST','/tasks/2/cancel',{})['outcome']=='cancelled'
assert events[-1][0]=='/v2/video_generation/old-2/cancel' and len(events)==before+1
assert sql('SELECT status FROM sys_task WHERE id=3')[0]['status']=='processing'
before=len(events);assert api('POST','/tasks/2/cancel',{})['outcome']=='cancelled';assert len(events)==before
sql('INSERT INTO ai_service_configs(id,service_type,provider,name,base_url,api_key,model,is_active,created_at,updated_at) VALUES(2,?,?,?,?,?,?,1,?,?)',
('image','openai','fixture-image','http://127.0.0.1:18935/v1','test-key',json.dumps(['qwen_image_2_1']),ts,ts))
for i in [20,21,22,23]:
sql('INSERT INTO sys_task(id,type,provider,model,status,stage,task_id,config_id,config_base_url,created_at,updated_at) VALUES(?,?,?,?,?,?,?,?,?,?,?)',
(i,'image','openai','qwen_image_2_1','processing','submitting' if i==23 else 'queued',None if i==23 else 'image-'+str(i),2,'http://127.0.0.1:18935/v1',ts,ts))
sql('INSERT INTO sys_task(id,type,status,stage,local_path,created_at,updated_at) VALUES(19,?,?,?,?,?,?)',('image','completed','completed','static/images/reference.png',ts,ts))
assert api('GET','/tasks/20')['canCancel'] is True
before=len(events);assert api('POST','/tasks/20/cancel',{})['outcome']=='cancelled'
assert events[-1][0]=='/v1/images/task/image-20/cancel' and len(events)==before+1
assert sql('SELECT local_path FROM sys_task WHERE id=19')[0]['local_path']=='static/images/reference.png'
assert (data/'static/images/reference.png').read_bytes()==png
# Wrong receipt / transport failures are unconfirmed, never falsely cancelled.
cancel_modes['image-21']='wrong-id'
assert api('POST','/tasks/21/cancel',{})['outcome']=='pending'
assert sql('SELECT status,stage FROM sys_task WHERE id=21')[0]=={'status':'processing','stage':'cancelling'}
cancel_modes['image-21']='transport'
assert api('POST','/tasks/21/cancel',{})['outcome']=='pending'
# A completion racing with cancellation is preserved for normal download.
cancel_modes['image-22']='succeeded'
assert api('POST','/tasks/22/cancel',{})['outcome']=='completed'
assert sql('SELECT error_code,status FROM sys_task WHERE id=22')[0]=={'error_code':None,'status':'processing'}
before=len(events);api('POST','/tasks/23/cancel',{},expected=400);api('POST','/tasks/19/cancel',{},expected=400)
assert len(events)==before and api('GET','/tasks/23')['canCancel'] is False
sql('INSERT INTO sys_task(id,type,status,stage,created_at,updated_at) VALUES(24,?,?,?,?,?)',('video','processing','waiting_predecessor',ts,ts))
assert api('POST','/tasks/24/cancel',{})['outcome']=='cancelled' and len(events)==before
# The unconfirmed request survives restart; polling confirms the same ID,
# without another cancel or generation request.
job_states['image-21']='cancelled'
process.terminate();process.wait(timeout=10)
process=subprocess.Popen([str(stage/'HuobaoDrama.exe'),str(data/'backend.mjs')],cwd=data,env=env,stdout=log,stderr=log,creationflags=subprocess.CREATE_NO_WINDOW)
for _ in range(180):
row=sql('SELECT status,error_code FROM sys_task WHERE id=21')[0]
if row['error_code']=='user_cancelled':break
time.sleep(.1)
assert row=={'status':'failed','error_code':'user_cancelled'} and len(events)==before
print(json.dumps({'passed':True,'inline_reference_preflight':True,'first_frame':True,'cancel_only_selected_shot':True,
'preserve_old_on_missing_image':True,'preserve_old_on_validation_failure':True,'image_cancel':True,'cancel_idempotent':True,
'completion_race_preserved':True,'bad_receipt_not_cancelled':True,'cancel_survives_restart':True,'fixture':str(data)}))
finally:
if process and process.poll() is None: process.terminate(); process.wait(timeout=10)
server.shutdown(); log.close()