jobs/v7_cache_l4.py

from __future__ import annotations
import json,os,subprocess,sys
from pathlib import Path
import torch
from huggingface_hub import HfApi,hf_hub_download

RUNTIME_REPO='patdev/Companion-Forge-L4-ONNX'

def main(request_path:str):
    reqp=Path(request_path);run_dir=reqp.parent;req=json.loads(reqp.read_text());token=os.environ['HF_TOKEN']
    inp=run_dir/req['input_file'];assert inp.exists(),inp
    cache_repo=req['cache_repo'];HfApi(token=token).create_repo(cache_repo,repo_type='dataset',private=True,exist_ok=True)
    worker=hf_hub_download(RUNTIME_REPO,'v7/jobs/build_teacher_cache_l4.py',repo_type='model',token=token,local_dir='/tmp/cf-v7-cache-worker',force_download=True)
    out=run_dir/'teacher-cache.pt';key=req.get('key') or f"reference-{req.get('seed',42)}"
    cmd=[sys.executable,worker,str(inp),'--out',str(out),'--seed',str(int(req.get('seed',42))),'--ss-steps',str(int(req.get('ss_steps',10))),'--upload-repo',cache_repo,'--key',key]
    env=os.environ.copy();env['CF_CACHE_HARD_EXIT']='1'
    proc=subprocess.run(cmd,env=env,check=False)
    if proc.returncode!=0:raise RuntimeError(f'v7 cache child failed with exit code {proc.returncode}')
    record=torch.load(out,map_location='cpu',weights_only=False);meta=record.get('meta',{});size=out.stat().st_size;out.unlink(missing_ok=True)
    result={'ok':True,'kind':'v7-teacher-cache','cache_repo':cache_repo,'key':key,'path':f'cache/{key}.pt','bytes':size,'coords':meta.get('coords'),'coords_skl':meta.get('coords_skl'),'conditioning_s':meta.get('conditioning_s'),'ss_steps':meta.get('ss_steps')}
    (run_dir/'result.json').write_text(json.dumps(result,indent=2));print('RESULT',json.dumps(result),flush=True)
if __name__=='__main__':main(sys.argv[1])