Skip to content

benchmark: robotics harness reads NDJSON in 1 MiB chunks and the DROID tasks from the index - #924

Merged
merceod merged 4 commits into
mainfrom
users/atindra/robotics-harness
Oct 28, 2026
Merged

merceod merged 4 commits into
mainfrom
users/atindra/robotics-harness

Conversation

@merceod

@merceod merceod commented Oct 28, 2026 •

Copy link
Copy Markdown
Collaborator

Two fixes in the robotics benchmark harness.

  • run_robotics.py read the server's NDJSON with iter_lines() at its 512-byte default, which is roughly quadratic on multi-MB lines. Reading with a 1 MiB chunk takes a 4-rollout V-JEPA2-AC run from 10.0 s to 0.29 s of client time, latents bit-identical.
  • DROIDDataset read tasks_df["task"], but droid_100 is a lerobot v3 dataset whose tasks.parquet keeps the instructions in the index. The lookup warned could not load tasks.parquet: 'task' and every pi0.5 request got the prompt "manipulate the object" (V-JEPA2-AC sends no text, so it was unaffected). It now reads the index, where lerobot v3 keeps the strings: requests 0 and 1 become "Put the marker in the pot" and "Put the candy bar on the left side of the first shelf". pi0.5 numbers from before this PR used the fallback prompt, so they aren't comparable with new ones.

Tests: new test_droid_dataset_tasks.py (1 passed). The eval and bench test modules show the same two box-specific failures as main when run serially. e2e --tier smoke pi05 passes; run_robotics against servers: vjepa2_ac 4 + 16 + 4 rollouts, pi05 eager 4/4 with the real prompts. ruff clean.

iter_lines' 512 B default rejoins a line chunk by chunk, which is quadratic
in its length. A V-JEPA 2-AC rollout line is several MiB, and the client
spent seconds per request reassembling it.
droid_100 is a lerobot v3.0 export, which keeps each instruction in the
index of meta/tasks.parquet rather than a task column. The lookup failed,
so every DROID request was sent the fallback prompt.
@mstar-review

mstar-review Bot commented Oct 28, 2026

Copy link
Copy Markdown

Automated review

🟡 No findings on b1fa740, light review

Description

  • The 10.0 s → 0.29 s client-time figure doesn't say how it was measured (which timer, what hardware, or the server config and command used for the 4-rollout run), so it can't be reproduced.
  • "The same two box-specific failures as main" doesn't name the two failing tests, so a reviewer can't check that the set of failures matches main.
What the review checked
  • _task_strings uses the task column when it exists and otherwise falls back to the index, so the older lerobot layout gives the same dict as the inline code it replaces.
  • benchmark/dataset.py only imports stdlib, numpy, requests and other benchmark.* modules at top level, and other test/modular modules already import benchmark, so the new test can run in CPU Core; it also skips when pandas is missing via importorskip.
  • Both iter_lines call sites in run_robotics.py (_send_pi05 NDJSON branch, _send_vjepa2_ac) now pass the same chunk size, and the line-splitting behaviour and the JSON-decode skip are unchanged.
  • The change is limited to the benchmark harness and touches no engine, worker or model code, so none of the runtime invariants apply.

Advisory, not a merge gate. Rules: AGENTS.md. Reply to discuss.

dghosef and others added 2 commits October 27, 2026 20:51
No lerobot writer puts a "task" column in meta/tasks.parquet: v2.x kept tasks
in meta/tasks.jsonl, and v3.0 (write_tasks, convert_tasks, aggregate) stores the
strings as the index. Drop the column branch and its test.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
@mstar-review mstar-review Bot added size/XS and removed size/S labels Oct 28, 2026

@dghosef dghosef left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

lgtm

@merceod
merceod merged commit e819c8d into main Oct 28, 2026
7 checks passed
@merceod
merceod deleted the users/atindra/robotics-harness branch October 28, 2026 05:24
dghosef added a commit that referenced this pull request Oct 28, 2026
Main since 73c6c17: #883 mstar_rust in the Linux wheels, #901 output
spans (BatchedModelOutput.output_spans, NodeSubmodule.spanned_outputs, the
worker's crop), #904 DiT seed noise drawn at ingest (on_request_added,
NoiseStager deleted, to_device_nonblocking as one memcpy into pinned memory),
#903 API-server torch pool from the CPU quota and pooled image encodes,
#924 robotics harness, #819 Inductor cache under ~/.cache, #881
CommGroup.agree / JointGroups.agree (agree_across_ranks and
flagged_on_any_rank deleted), #724 CUDA IPC KV in one 2D memcpy a page,
#387 GLM-5.2 (MLA, fp8 MoE with NoAuxTcRouter, the DSA resource), #921
vjepa2 AC crop, #922 Cosmos3 contiguous positions, #923 dense attention
falls back when fa3-fwd has no kernel for the GPU.

Conflicts:
- mstar/utils/h2d.py, to_device_nonblocking: the stack's signature
  (data: Tensor | Sequence, device, dtype=None) and DeviceCapability
  pin_memory gate, with main's tensor path (#904): a CPU tensor is written
  into a fresh pinned buffer with one memcpy rather than pin_memory()'s
  copy_ over the intra-op pool. Every CPU tensor now goes through a fresh
  buffer, so the stack's "an already-pinned source is read in place"
  caveat (R3.3) no longer applies and is dropped from the docstring.
  test_accelerator_graph_backend's pinning test checks the pinned
  torch.empty and its non-blocking copy instead of pin_memory().
- mstar/model/components/diffusion/denoise_loop.py: imports, both
  (main's copy, the stack's functools).

Main code on a superseded path, ported here:
- GLM-5.2 (glm52/submodules.py): its own check_stop / check_stop_batched
  (four arguments, step_walk, its own row-map loop) / _stops ->
  token_stop = TokenStop(decode_loop, new_token, eos_token_ids, sampler)
  with a _token_stops override for its two differences (a stop token at
  prefill ends the reply at the first decode step; max_tokens counts the
  prefill's token), and a check_stop override for that closing step,
  whose token postprocess drops. test_glm52_model's helper and
  test_model_kv_layouts' GLM cap test build the real submodule (over a stub
  language model) instead of object.__new__, which would leave token_stop
  unset.
- PinnedStager.to_device (#387), a second "host values -> new device
  tensor" helper beside to_device_nonblocking: deleted with its test
  (test_pinned_stager_to_device.py); its callers, GLM-5.2's prefill
  seq_lens and the DSA resource's page tables, call
  to_device_nonblocking(list, device, dtype=torch.long).
  test_dsa_resource spies on that instead.

The other files both sides touched (engine.py, accelerator_graph_runner.py,
kv/manager.py, attn/base.py, ragged/base.py, communication.py,
model/base.py, submodule_base.py, entrypoint.py, conductor.py, moe.py,
image_io.py, vjepa2_model.py, docs, tests) auto-merged into disjoint
hunks: per file the merged delta equals main's own.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants