File size: 9,500 Bytes
68eed61 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 | """Evaluate Zeva deltas continuously, or run one explicitly selected checkpoint."""
import argparse
import json
import os
import socket
from pathlib import Path
import subprocess
import time
import urllib.request
ROOT=Path(__file__).resolve().parents[1]
# Allow corrected runs to use an isolated artifact directory while preserving
# the historical default for existing watcher processes.
RUN=Path(os.environ.get('WIDOWX_RUN_DIR', str(ROOT/'.artifacts/widowx')))
PYTHON='/home/gpu4/tianyi/cosmos-framework/.venv/bin/python'
SIM_PYTHON='/home/gpu4/miniconda3/envs/simpler_env/bin/python'
PORT=18765
EVAL_GPU=os.environ.get('WIDOWX_EVAL_GPU','3')
RENDER_GPU=os.environ.get('WIDOWX_RENDER_GPU', EVAL_GPU)
def run_one(checkpoint, episodes=10, port=PORT, record_traces=False, max_steps=None, num_steps=4, seed_start=0, exec_horizon=4, guidance=3.0, memory_mode='on'):
# Do not mistake an unrelated local service's health endpoint for ours.
# Fail before spending time/VRAM loading the policy on an occupied port.
with socket.socket() as probe:
probe.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
probe.bind(('127.0.0.1', port))
out=RUN/'eval'; out.mkdir(parents=True,exist_ok=True)
gpu_uuid=subprocess.check_output(['nvidia-smi','-i',EVAL_GPU,'--query-gpu=uuid',
'--format=csv,noheader'],text=True).strip()
if not gpu_uuid.startswith('GPU-') or '\n' in gpu_uuid:
raise ValueError('WIDOWX_EVAL_GPU must select exactly one physical GPU')
render_uuid=subprocess.check_output(['nvidia-smi','-i',RENDER_GPU,'--query-gpu=uuid',
'--format=csv,noheader'],text=True).strip()
if not render_uuid.startswith('GPU-') or '\n' in render_uuid:
raise ValueError('WIDOWX_RENDER_GPU must select exactly one physical GPU')
env=os.environ.copy()
env.update(CUDA_VISIBLE_DEVICES=gpu_uuid,PYTHONPATH=str(ROOT)+':/home/gpu4/tianyi/cosmos-framework',
WIDOWX_POLICY_OUTPUT_DIR=str(out/'server_runtime'),
QWEN_VLM_PATH='/home/gpu4/tianyi/models/Qwen3-VL-8B-Instruct',
WAN_VAE_PATH='/home/gpu4/tianyi/models/wan22_vae/Wan2.2_VAE.pth',
HF_HUB_OFFLINE='1',TOKENIZERS_PARALLELISM='false')
simenv=env.copy()
simenv.update(CUDA_VISIBLE_DEVICES=render_uuid,PYTHONPATH='/opt/benchmarks/SimplerEnv',
VK_ICD_FILENAMES='/etc/vulkan/icd.d/nvidia_icd.json',
MS2_ASSET_DIR='/opt/benchmarks/SimplerEnv/ManiSkill2_real2sim/data')
subprocess.run([SIM_PYTHON,str(ROOT/'scripts/check_sapien_gpu.py'),
'--expected-uuid',render_uuid],env=simenv,check=True)
with (out/(checkpoint.stem+'.server.log')).open('a') as log:
server=subprocess.Popen([PYTHON,str(ROOT/'scripts/widowx_policy_server.py'),
'--delta',str(checkpoint),'--port',str(port),'--num-steps',str(num_steps),
'--guidance',str(guidance),'--memory-mode',memory_mode],
env=env,cwd='/home/gpu4/tianyi/cosmos-framework',stdout=log,stderr=subprocess.STDOUT)
try:
deadline=time.monotonic()+600
while True:
if server.poll() is not None: raise RuntimeError(f'policy server exited {server.returncode}')
try:
with urllib.request.urlopen(f'http://127.0.0.1:{port}/healthz',timeout=3) as r:
if r.status==200: break
except Exception: pass
if time.monotonic()>deadline: raise TimeoutError('policy server startup')
time.sleep(3)
trace_args = ['--trace-dir', str(out/'traces'/checkpoint.stem)] if record_traces else []
limit_args = ['--max-steps', str(max_steps)] if max_steps is not None else []
with (out/(checkpoint.stem+'.episodes.log')).open('a') as eval_log:
subprocess.run([SIM_PYTHON,str(ROOT/'scripts/eval_widowx_simpler.py'),
'--checkpoint',str(checkpoint),'--output',str(out/(checkpoint.stem+'.json')),
'--video-dir',str(out/'videos'/checkpoint.stem),
'--episodes',str(episodes),'--port',str(port),'--seed-start',str(seed_start),
'--exec-horizon',str(exec_horizon),
'--memory-mode',memory_mode,
*trace_args,*limit_args],env=simenv,check=True,
stdout=eval_log,stderr=subprocess.STDOUT)
report=json.loads((out/(checkpoint.stem+'.json')).read_text())
if report.get('status') != 'complete' or len(report.get('episodes', [])) != episodes:
raise RuntimeError('evaluation did not produce the requested number of complete episodes')
report['inference_num_steps'] = num_steps
report['guidance'] = guidance
report['physical_gpu_uuid'] = gpu_uuid
report['render_gpu_uuid'] = render_uuid
report['memory_mode'] = memory_mode
report_path = out/(checkpoint.stem+'.json')
temporary = report_path.with_suffix('.tmp')
temporary.write_text(json.dumps(report, indent=2))
temporary.replace(report_path)
print('EVAL_COMPLETE '+json.dumps(dict(checkpoint=str(checkpoint),
episodes=episodes,successes=sum(bool(x['success']) for x in report['episodes']),
success_rate=report['success_rate'])),flush=True)
finally:
server.terminate()
try: server.wait(timeout=30)
except subprocess.TimeoutExpired: server.kill(); server.wait()
def main():
parser=argparse.ArgumentParser(description=__doc__)
parser.add_argument('--once', action='store_true', help='evaluate one checkpoint and exit; never starts training')
parser.add_argument('--checkpoint', type=Path, help='explicit checkpoint for --once')
parser.add_argument('--episodes', type=int, default=10)
parser.add_argument('--port', type=int, default=PORT)
parser.add_argument('--record-traces', action='store_true')
parser.add_argument('--num-steps', type=int, default=4, help='UniPC denoising steps per prediction')
parser.add_argument('--seed-start', type=int, default=0, help='first environment and inference seed')
parser.add_argument('--exec-horizon', type=int, default=4, choices=range(1,33))
parser.add_argument('--guidance', type=float, default=3.0)
parser.add_argument('--memory-mode',choices=('on','bank_off'),default='on')
parser.add_argument('--max-steps', type=int, default=None,
help='maximum environment control steps per episode')
args=parser.parse_args()
if args.episodes < 1 or not 1 <= args.port <= 65535 or (args.max_steps is not None and args.max_steps < 1):
parser.error('episodes and max-steps must be positive; port must be in [1,65535]')
if args.num_steps < 1 or args.seed_start < 0:
parser.error('num-steps must be positive and seed-start must be nonnegative')
if args.once != (args.checkpoint is not None):
parser.error('--once and --checkpoint must be supplied together')
(RUN/'eval').mkdir(parents=True,exist_ok=True)
if args.once:
checkpoint=args.checkpoint.resolve()
if not checkpoint.is_file():
parser.error(f'checkpoint does not exist: {checkpoint}')
if (RUN/'eval'/(checkpoint.stem+'.json')).exists():
parser.error('a result already exists for this checkpoint; use a separate WIDOWX_RUN_DIR')
print('EVALUATING_ONCE '+json.dumps(dict(checkpoint=str(checkpoint),
episodes=args.episodes,gpu=EVAL_GPU,port=args.port,
inference_num_steps=args.num_steps,seed_start=args.seed_start,max_steps=args.max_steps,
memory_mode=args.memory_mode,render_gpu=RENDER_GPU)),flush=True)
try:
run_one(checkpoint, args.episodes, args.port, args.record_traces, args.max_steps, args.num_steps, args.seed_start, args.exec_horizon, args.guidance, args.memory_mode)
except Exception as e:
(RUN/'eval'/(checkpoint.stem+'.error.json')).write_text(
json.dumps(dict(status='failed',error=repr(e),time=time.time())))
raise
return
while True:
# Evaluate the newest snapshot first so a long backlog cannot delay the
# most useful current-policy result; older 2k snapshots are backfilled.
for checkpoint in sorted((RUN/'checkpoints').glob('step_*.pt'), reverse=True):
if int(checkpoint.stem.split('_')[1]) % 2000:
continue # Startup save/resume probes are not scheduled evaluations.
if (RUN/'eval'/(checkpoint.stem+'.json')).exists(): continue
try:
print('EVALUATING',checkpoint,flush=True)
run_one(checkpoint, args.episodes, args.port, args.record_traces, args.max_steps, args.num_steps, args.seed_start, args.exec_horizon, args.guidance, args.memory_mode)
except Exception as e:
failure=RUN/'eval'/(checkpoint.stem+'.error.json')
failure.write_text(json.dumps(dict(status='failed',error=repr(e),time=time.time())))
print('EVAL_FAILED',repr(e),flush=True)
time.sleep(30)
time.sleep(10)
if __name__=='__main__': main()
|