File size: 5,408 Bytes
30b7f88
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
# 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/listen`s 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:
```bash
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:
```bash
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.
```