Enable two-rank CFG-parallel RoboLab serving

[vivekgoe]

Summary

Adds explicit two-rank CFG-parallel support to the RoboLab action-policy server without enabling FSDP or context parallelism.

Changes

• Adds a single --cfg-parallel CLI option.
• Supports:
  • Single process: dp_shard=1, cfgp=1, cp=1
  • Two ranks with --cfg-parallel: dp_shard=1, dp_replicate=2, cfgp=2, cp=1
• Rejects unsupported multi-rank configurations with a clear error.
• Runs the public WebSocket server only on global rank 0.
• Keeps rank 1 in a distributed inference worker loop.
• Broadcasts each observation and seed so both ranks enter CFGP collectives with identical inputs.
• Extends the request lock to cover seed selection, request broadcast, and inference.
• Disables unused OmniInference guardrails, avoiding unnecessary Guardrail1 downloads and model allocation.
• Adds tests for supported and rejected serving topologies.

Runtime behavior

Rank 0 receives the WebSocket request and broadcasts {obs, seed} to rank 1. Both ranks construct the same transformed batch and participate in CFG-parallel generation:

• Rank 0 executes the conditional CFG branch and returns the action to RoboLab.
• Rank 1 executes the unconditional CFG branch and discards its local response.

Both GPUs hold a complete copy of the policy model. Model weights are not FSDP-sharded.

Usage

torchrun --standalone --nproc-per-node=2 \
  -m cosmos_framework.scripts.action_policy_server_robolab \
  --checkpoint-path nvidia/Cosmos3-Edge-Policy-DROID \
  --port 8000 \
  --format-prompt-as-json True \
  --cfg-parallel

Limitations

• CFG-parallel mode currently requires exactly two ranks.
• The complete policy model must fit on each GPU.
• Requests are processed serially to preserve distributed collective ordering.
• FSDP and custom CP/CFGP combinations are intentionally not exposed.

Validation

• Ruff checks pass.
• Python syntax compilation passes.
• git diff --check passes.
• Unit coverage was added for supported and rejected parallelism configurations.
• Tested cfg-parallel on a DualThor set-up. Confirmed that inference latency for Robolab "BananaInBowlTask" reduces by ~40% when using both Thor with cfg-parallel (compared to using single Thor without cfg-parallel).

[vivekgoe]

@pengcuo Please help review this PR. Thanks.

[pengcuo]

Thanks for the PR, Vivek. The CFG-parallel idea is good and the core logic is correct: both ranks run the same collectives in the same order, dp_shard=1 / cfgp=2 / cp=1 builds the right meshes, and the seed is synced by the model. Tests pass.

Two things will take the 2-rank server down in real use, so I'd like them fixed before merge:

1. One bad request kills the server. infer() broadcasts obs before validating it. On rank 0 the error is caught by openpi (connection closed, server keeps running). On rank 1 the same error escapes worker_loop, the process exits, and torchrun kills rank 0. Single-rank mode survives the same request. Fix: validate (_build_sample) on rank 0 before broadcasting, and wrap the worker's per-request body in try/except so it keeps looping.

2. Idle for 30 min = crash. Between requests rank 1 waits inside dist.broadcast_object_list on the NCCL world group. Our PG timeout is 1800 s with TORCH_NCCL_ASYNC_ERROR_HANDLING=1, so if no request arrives for 30 min the NCCL watchdog aborts rank 1. Fix: wait for requests on a gloo control group created with a very large timeout (or a TCPStore), and only issue NCCL collectives while a request is running.

Smaller items, fine as follow-ups:

• Both ranks download the HF checkpoint at the same time. Please wrap _resolve_checkpoint_path in _download_on_rank0 (our loader already says concurrent downloads are unsafe on Lustre/NFS).
• No shutdown signal: if rank 0 exits, rank 1 stays blocked in NCCL holding its GPU. A {"shutdown": True} sentinel in a try/finally around serve_forever() fixes it.
• /healthz still returns OK after rank 1 dies.
• --cfg-parallel --guidance 1.0 never enters the CFG path, so the second GPU does nothing. Worth rejecting that combo.
• Rank 1 also runs the discarded post-processing and VAE decode. Splitting _infer_impl into generate / postprocess would avoid it.
• "guardrails": False also changes single-rank startup and differs from the LIBERO server. Fine by me, but please mention it in the PR description.
• Only dp_shard_size=1 is actually new; cfgp=2 / cp=1 already come from the default latency preset. Consider reusing --dp-shard-size / --cfgp-size instead of a new --cfg-parallel flag, and validating the launch shape before the download.
• Comments mention multi-threading and FSDP; openpi is single-threaded asyncio and FSDP is off here.
• Tests only cover _resolve_parallelism_overrides; a stubbed-broadcast test of infer() -> worker_loop() would be cheap. ruff format flags the new test hunk.

I reproduced item 1 with a 2-process gloo simulation; happy to share the script.

[vivekgoe]

@pengcuo Thanks for detailed review comments. I will address these and push another patch soon.