| """Run all 60 rollouts exclusively on the eai host; resume completed episodes.""" |
| import datetime |
| import json |
| import os |
| from pathlib import Path |
| import signal |
| import socket |
| import subprocess |
| import time |
|
|
| ROOT=Path(__file__).resolve().parent |
| MODELS=ROOT.parent/'b1k_sft' |
| VARIANTS=['human','codegen_success','mix_h50_c50'] |
| PORT=18080 |
| EXPECTED_HOST='eai-Raider-16-HX-AI-A2XWIG' |
| POLICY_PY=ROOT/'.policy-venv/bin/python' |
| SIM_PY=Path('/home/eai/miniconda3/envs/behavior-392dev/bin/python') |
| B1K=Path('/home/eai/code/b1k') |
| status={'host':socket.gethostname(),'started':datetime.datetime.now().isoformat(),'state':'starting','variants':{}} |
| children=[] |
|
|
| def save(): |
| status['updated']=datetime.datetime.now().isoformat() |
| p=ROOT/'status.tmp';p.write_text(json.dumps(status,indent=2));p.replace(ROOT/'status.json') |
|
|
| def stop(p): |
| if p and p.poll() is None: |
| os.killpg(p.pid,signal.SIGTERM) |
| try:p.wait(timeout=30) |
| except subprocess.TimeoutExpired: |
| os.killpg(p.pid,signal.SIGKILL);p.wait() |
|
|
| def missing(variant): |
| root=ROOT/'results'/variant/'json' |
| ids=[] |
| for i in range(20): |
| p=root/f'turning_on_radio_{301+i}_0.json' |
| if not p.exists():ids.append(i);continue |
| r=json.loads(p.read_text()) |
| assert r['instance_id']==301+i and r['task']=='turning_on_radio' |
| assert 'success' in r and 'final' in r['q_score'] |
| return ids |
|
|
| def main(): |
| assert socket.gethostname()==EXPECTED_HOST,'Evaluation is restricted to eai' |
| policy_env=os.environ.copy() |
| policy_env.update(PYTHONPATH=':'.join(map(str,[ROOT,ROOT/'behavior-1k-solution/src',ROOT/'behavior-1k-solution/openpi/src',ROOT/'behavior-1k-solution/openpi/packages/openpi-client/src'])),CUDA_VISIBLE_DEVICES='0',XLA_PYTHON_CLIENT_PREALLOCATE='false',XLA_PYTHON_CLIENT_ALLOCATOR='platform',XLA_FLAGS='--xla_gpu_enable_triton_gemm=false',JAX_COMPILATION_CACHE_DIR=str(ROOT/'jax-cache'),JAX_PERSISTENT_CACHE_MIN_COMPILE_TIME_SECS='1',TOKENIZERS_PARALLELISM='false',OMP_NUM_THREADS='4') |
| sim_env=os.environ.copy() |
| sim_env.update(PYTHONPATH=':'.join(map(str,[B1K/'OmniGibson',B1K/'bddl3',B1K/'joylo'])),OMNIGIBSON_HEADLESS='1',OMNIGIBSON_DATA_PATH=str(B1K/'datasets'),CUDA_VISIBLE_DEVICES='0',OMP_NUM_THREADS='4',PYTHONUNBUFFERED='1') |
| for key in ['ISAAC_PATH','EXP_PATH','CARB_APP_PATH']: |
| sim_env.pop(key,None) |
| for variant in VARIANTS: |
| directory=ROOT/'results'/variant;directory.mkdir(parents=True,exist_ok=True) |
| remaining=missing(variant) |
| if not remaining: |
| status['variants'][variant]={'state':'complete','completed':20};save();continue |
| ready=directory/'policy-ready.json';ready.unlink(missing_ok=True) |
| status.update(state='loading_policy',active_variant=variant);save() |
| with (directory/'policy.log').open('a') as log: |
| policy=subprocess.Popen([str(POLICY_PY),str(ROOT/'serve_checkpoint.py'),'--variant',str(MODELS/variant),'--port',str(PORT),'--log',str(directory/'policy-events.jsonl'),'--ready',str(ready)],env=policy_env,stdout=log,stderr=subprocess.STDOUT,start_new_session=True) |
| children.append(policy) |
| try: |
| deadline=time.monotonic()+1200 |
| while not ready.exists(): |
| if policy.poll() is not None:raise RuntimeError(f'{variant} policy exited: {policy.returncode}') |
| if time.monotonic()>deadline:raise TimeoutError('Policy startup timeout') |
| time.sleep(2) |
| for attempt in range(1,4): |
| remaining=missing(variant) |
| if not remaining:break |
| status.update(state='evaluating',active_variant=variant,attempt=attempt) |
| status['variants'][variant]={'state':'running','completed':20-len(remaining),'pending_indices':remaining};save() |
| cmd=[str(SIM_PY),str(ROOT/'run_simulator.py'),'--task-name','turning_on_radio','--mode','public_test','--host','127.0.0.1','--port',str(PORT),'--instance-indices',*map(str,remaining),'--output-dir',str(directory)] |
| with (directory/f'simulator-attempt{attempt}.log').open('a') as simlog: |
| sim=subprocess.Popen(cmd,env=sim_env,stdout=simlog,stderr=subprocess.STDOUT,start_new_session=True);children.append(sim) |
| while sim.poll() is None: |
| todo=missing(variant) |
| status['variants'][variant]={'state':'running','completed':20-len(todo),'pending_indices':todo};save() |
| if policy.poll() is not None: |
| stop(sim);raise RuntimeError(f'Policy failed during {variant} evaluation') |
| time.sleep(10) |
| status['last_simulator_exit']=sim.returncode;save() |
| if missing(variant):raise RuntimeError(f'{variant} still missing indices {missing(variant)} after retries') |
| status['variants'][variant]={'state':'complete','completed':20};save() |
| finally:stop(policy) |
| status['state']='complete';save() |
| subprocess.run([str(SIM_PY),str(ROOT/'summarize.py')],check=True) |
|
|
| if __name__=='__main__': |
| try:main() |
| except BaseException as exc: |
| status.update(state='failed',error=repr(exc));save();raise |
| finally: |
| for p in children:stop(p) |
|
|
