87 lines
4.1 KiB
Python
87 lines
4.1 KiB
Python
"""Worker drain and verified history restore for rolling updates.
|
|
|
|
Control files are operator-owned, read-only mounts. They cannot be supplied in a
|
|
generation prompt. Existing history and output reads are never gated by drain.
|
|
"""
|
|
import hashlib
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
import secrets
|
|
import sys
|
|
import uuid
|
|
|
|
|
|
def maintenance_enabled(path):
|
|
try:
|
|
data=Path(path).read_bytes()
|
|
if len(data)>4096:return True
|
|
value=json.loads(data)
|
|
if value.get('schema')!=1 or not isinstance(value.get('maintenance'),bool):return True
|
|
return value['maintenance']
|
|
except (OSError,ValueError,TypeError):return True
|
|
|
|
|
|
def make_middleware(path,canary_token):
|
|
from aiohttp import web
|
|
@web.middleware
|
|
async def gate(request,handler):
|
|
controlled=(request.method=='GET' and request.path in ['/queue','/api/queue']) or (
|
|
request.method=='POST' and request.path in ['/prompt','/api/prompt'])
|
|
supplied=request.headers.get('X-H3-Canary-Token','')
|
|
bypass=bool(canary_token and secrets.compare_digest(supplied,canary_token))
|
|
if controlled and not bypass and maintenance_enabled(path):
|
|
return web.json_response({'error':{'type':'h3_worker_draining',
|
|
'message':'Worker maintenance: existing jobs/results are retained; new jobs use the other worker.'}},status=503,
|
|
headers={'Retry-After':'10'})
|
|
# Authentication is still handled by the original middleware/handler.
|
|
# A canary token only bypasses drain, never normal API authorization.
|
|
return await handler(request)
|
|
return gate
|
|
|
|
|
|
def restore_history(queue,path,expected_sha256,expected_source):
|
|
p=Path(path)
|
|
if p.stat().st_size>256*1024**2:raise ValueError('History snapshot exceeds the audited size bound')
|
|
raw=p.read_bytes()
|
|
if not expected_sha256 or hashlib.sha256(raw).hexdigest()!=expected_sha256:
|
|
raise ValueError('History snapshot SHA256 mismatch')
|
|
snapshot=json.loads(raw)
|
|
if snapshot.get('schema')!=1 or snapshot.get('queue_empty') is not True or snapshot.get('source_container_id')!=expected_source:
|
|
raise ValueError('History snapshot identity or drain state mismatch')
|
|
history=snapshot.get('history')
|
|
if not isinstance(history,dict):raise ValueError('History snapshot must contain a dictionary')
|
|
for key,value in history.items():
|
|
uuid.UUID(key)
|
|
if not isinstance(value,dict) or value.get('status',{}).get('status_str') not in ['success','error']:
|
|
raise ValueError('A nonterminal history entry cannot be imported')
|
|
with queue.mutex:
|
|
if queue.queue or queue.currently_running:
|
|
raise ValueError('Never restore a snapshot into a worker already executing jobs')
|
|
overlap=set(history)&set(queue.history)
|
|
if any(history[k]!=queue.history[k] for k in overlap):raise ValueError('History identifier collision')
|
|
queue.history.update(history)
|
|
return len(history)
|
|
|
|
|
|
def register_worker_state():
|
|
server=sys.modules.get('server')
|
|
instance=getattr(getattr(server,'PromptServer',None),'instance',None)
|
|
if instance is None or getattr(instance,'_h3_worker_state_registered',False):return
|
|
control=os.environ.get('H3_WORKER_CONTROL_FILE','')
|
|
if control:
|
|
token=os.environ.get('H3_CANARY_TOKEN','')
|
|
token_file=os.environ.get('H3_CANARY_TOKEN_FILE','')
|
|
if token_file:
|
|
token=Path(token_file).read_text().strip()
|
|
if len(token)<16:raise ValueError('Canary token file is missing a valid operator token')
|
|
instance.app.middlewares.append(make_middleware(control,token))
|
|
snapshot=os.environ.get('H3_HISTORY_SNAPSHOT','')
|
|
if snapshot:
|
|
async def restore_before_listen(app):
|
|
count=restore_history(instance.prompt_queue,snapshot,os.environ.get('H3_HISTORY_SNAPSHOT_SHA256',''),
|
|
os.environ.get('H3_HISTORY_SOURCE_CONTAINER_ID',''))
|
|
print(f'H3 verified history restored: {count} terminal jobs; no prompt resubmission',flush=True)
|
|
instance.app.on_startup.append(restore_before_listen)
|
|
instance._h3_worker_state_registered=True
|