robopro-gb10-eval-guide / ROBOPRO_EVAL_CROSSNODE_GB10.md
JackLiu0406's picture
Upload ROBOPRO_EVAL_CROSSNODE_GB10.md with huggingface_hub
30b7f88 verified
|
Raw
History Blame Contribute Delete
5.41 kB

Running RoboPRO eval on GB10 β€” split MODEL and SIM across two nodes

For teams whose model + simulator don't fit on one GB10 node. The RoboPRO eval is a client/server split already β€” you just point the two halves at different nodes.

Architecture (who is server, who is client)

  • Model = TCP server (policy_model_server.py, runs in the xvla env): loads the policy on the GPU once, bind/listens on a port, returns an action chunk per observation.
  • Sim = TCP client (eval_policy_client.py, runs in the robotwin env): runs the SAPIEN simulation + RoboPRO benchmark task in-process, and each control step sends (3-cam RGB + proprio + instruction) to the model server and executes the returned 30-step chunk.

There is no separate sim server β€” the simulator is embedded in the client. The only server is the model. So the sim dials into the model. (If your past setup was sim-as-server, this is the reverse direction β€” same socket, roles swapped.)

The two knobs that make it cross-node

side env var set to default
model node MODEL_SERVER_BIND 0.0.0.0 (listen on all interfaces, not just localhost) localhost
sim node MODEL_SERVER_HOST the model node's hostname (e.g. trt-gb10-7) localhost
both --port same high port, e.g. 29555 random

Single-node (our setup) leaves both at localhost. Cross-node only requires changing these two.

Prereqs on GB10

  • Both nodes need the staged envs + code + assets in node-local /tmp/dxvla_local (NFS mmap is too slow at runtime). Reuse our stager: bash /shared_work/jack/eval_root/stage_k320_node.sh (and stage_giants_node.sh for giant ckpts). Each node self-stages on first run.
  • SAPIEN 3.0.0b1 is required for correct office/kitchen physics (DXVLA_SAPIEN_B1=1); 3.0.3 is a regression. See GB10_EVAL_PIPELINE.md.
  • WORKERS_PER_GPU=1 (MPS is broken on this Blackwell GPU; sharing corrupts the sim).
  • Compute nodes share the internal network and can reach each other on high ports β€” use the node hostname (trt-gb10-N), bind 0.0.0.0. If a port is taken, pick another (29500–29999).

Step-by-step (2 nodes via SLURM)

Allocate two nodes, say node A = model, node B = sim. Pick a port, e.g. PORT=29555.

1. On the MODEL node (A) β€” start the policy server (xvla env), bound to all interfaces:

srun --nodelist=trt-gb10-A --nodes=1 --gres=gpu:1 --cpus-per-task=20 -t 600 bash -c '
  L=/tmp/dxvla_local
  MODEL_SERVER_BIND=0.0.0.0 \
  CKPT_NAME=<your-ckpt> GIANT_CODE_DIR=<your-model-code-dir> \
  DXVLA_SKIP_RELOAD=0 XVLA_POSED_DA3=0 \
  $L/envs/xvla/bin/python \
    /tmp/dxvla_local/code/RoboPRO/customized_robotwin/script/policy_model_server.py --port 29555
'
# wait for: "πŸš€ Model server started on 0.0.0.0:29555"

2. On the SIM node (B) β€” start the eval client (robotwin env), pointed at node A:

srun --nodelist=trt-gb10-B --nodes=1 --gres=gpu:1 --cpus-per-task=18 -t 600 bash -c '
  L=/tmp/dxvla_local
  MODEL_SERVER_HOST=trt-gb10-A \
  DXVLA_SAPIEN_B1=1 WORKERS_PER_GPU=1 DXVLA_DENOISER=optix \
  USE_EVAL_SEEDS=1 BENCH_CONFIG=clean EVAL_TEST_NUM=5 \
  BENCH_EVAL_SEEDS_ROOT=/shared_work/jack/robopro-eval-pipeline/seeds/k320_first10_seeds \
  $L/envs/robotwin/bin/python \
    /tmp/dxvla_local/code/RoboPRO/customized_robotwin/script/eval_policy_client.py \
    --task_name <task> --task_config bench_demo_<scene>_clean --port 29555
'
# expect: "πŸ”— Connected to model server at trt-gb10-A:29555"

The client streams obs β†’ actions to node A and runs the rollout locally on node B. Results land in the run's eval_result/.../_metrics.jsonl + videos.

Easiest path: copy our single-node driver slurm_eval_node.sh (which co-locates both on localhost) and split it β€” launch the server step on node A with MODEL_SERVER_BIND=0.0.0.0, and the client step on node B with MODEL_SERVER_HOST=<nodeA>. Everything else (staging, b1 swap, asset symlink, env exports) is identical.

Plugging in YOUR model

The server loads the policy via get_model() in policy/dxvla/deploy_policy.py. Either:

  • swap your checkpoint via CKPT_NAME + GIANT_CODE_DIR (if it's an X-VLA-family model β€” see the per-ckpt EVAL_INTERFACE.md for the loading contract), or
  • implement your own get_model() returning an object with generate_actions(input_ids, image_input, image_mask, domain_id, proprio, ...) -> [B, num_actions, action_dim]. The server handles the socket; you only provide the policy. The obs the client sends per step: 3-view RGB (countertop = main), proprio, language instruction; it expects a 30-step action chunk back.

GB10 gotchas (don't skip)

  • DXVLA_SAPIEN_B1=1 on the sim node; OptiX denoiser (OIDN is a silent no-op on Blackwell).
  • WORKERS_PER_GPU=1; do not enable MPS.
  • Stage to node-local /tmp/dxvla_local; never run from NFS across the fan-out.
  • Kill background drivers by PID, cancel SLURM by --name/jobid β€” never pkill -f <pat> that matches your own shell (β†’ exit 144).
  • Reserve nodes 1, 2 (Mark). Full issue list: GB10_EVAL_PIPELINE.md.

Validate before a full run

Smoke 1 cell/scene: run one --task_name with EVAL_TEST_NUM=1, confirm the client logs Connected to model server and a _metrics.jsonl line appears with no Traceback/asset errors.