34 lines
2.7 KiB
Python
34 lines
2.7 KiB
Python
"""Patch only the existing private bridges; no real jobs are cancelled."""
|
|
from pathlib import Path
|
|
import datetime,subprocess,json
|
|
base=Path('/home/ps/huobao-progress-update')
|
|
backup=base/('cancel-backup-'+datetime.datetime.now().strftime('%Y%m%d-%H%M%S'))
|
|
backup.mkdir(mode=0o700)
|
|
endpoint=''' @app.post("/v1/images/task/{jid}/cancel", dependencies=[Depends(auth)])
|
|
async def cancel_image(jid: str):
|
|
with closing(sqlite3.connect(db_path)) as db:
|
|
found = db.execute("SELECT 1 FROM image_jobs WHERE id=?", (jid,)).fetchone()
|
|
if not found:
|
|
raise HTTPException(404, "Unknown image task")
|
|
result = await upstream("POST", "/v1/jobs/" + quote(jid, safe="") + "/cancel", {})
|
|
return {"task": {"id": jid, "status": result.get("status", "running"),
|
|
"cancelRequested": result.get("cancelRequested", False)}}
|
|
|
|
'''
|
|
for name,location in [('huobao-model-bridge','/home/ps/huobao-model-bridge/bridge.py'),('huobao-vdn-bridge','/home/ps/huobao-vdn/bridge/bridge.py')]:
|
|
p=Path(location);original=p.read_text();s=original
|
|
if 'async def cancel_image(' not in s:
|
|
target=' @app.get("/v1/images/task/{jid}", dependencies=[Depends(auth)])'
|
|
assert s.count(target)==1;s=s.replace(target,endpoint+target)
|
|
for old,new in [
|
|
('return {"status": "failed", "error": {"message": str(result.get("error") or "Image generation failed")[:600]}}', 'return {"task_id": jid, "status": "failed", "cancelled": state == "cancelled", "error": {"message": str(result.get("error") or "Image generation failed")[:600]}}'),
|
|
('return {"status": "processing", "provider_id": result.get("providerId"), "worker_state": result.get("worker_state"), "queue_position": result.get("queue_position")}', 'return {"task_id": jid, "status": "processing", "provider_id": result.get("providerId"), "worker_state": result.get("worker_state"), "queue_position": result.get("queue_position")}'),
|
|
('return {"status": "failed", "error": {"message": "Task completed without an image output"}}', 'return {"task_id": jid, "status": "failed", "error": {"message": "Task completed without an image output"}}'),
|
|
('return {"status": "completed", "image_url": url, "data": [{"url": url}], "provider_id": result.get("providerId")}', 'return {"task_id": jid, "status": "completed", "image_url": url, "data": [{"url": url}], "provider_id": result.get("providerId")}')]:
|
|
if new not in s:assert s.count(old)==1;s=s.replace(old,new)
|
|
compile(s,location,'exec');(backup/(name+'.py')).write_text(original);p.write_text(s)
|
|
subprocess.run(['docker','restart','--time','15',name],check=True,capture_output=True)
|
|
assert json.loads(subprocess.check_output(['docker','inspect',name]))[0]['State']['Running']
|
|
print(name,'updated')
|
|
print('backup',backup)
|