perf(lmdb): parallelize batch decoding and support epoch-based schedules - #5914
Conversation
Decode same-nloc batches directly into preallocated arrays and use a bounded, rank-shared process pool so both PT and pt_expt can overlap deterministic LMDB reads with GPU work without duplicating dataset metadata. Centralize worker budgeting in DP_LMDB_NUM_WORKERS, preserve synchronous handling for small batches, and harden schema consistency, dtype promotion, path canonicalization, requirement freezing, and LMDB resource lifecycles across normal and exceptional exits. Add regression coverage for serial/parallel parity, finite epochs, tail batches, shared pools, concurrent readers, merge failures, multi-task training, and cleanup behavior, and document the immutable-dataset and aggregate-worker contracts.
pt_expt accepted only an explicit numb_steps, so a run had to be sized in steps even when the goal was a number of passes over the data. The epoch-to-step conversion now lives in a backend-independent resolver that pt uses as well, so both backends share the option validation, the multi-task step split and the cross-rank pinning that keeps num_steps in lockstep. A backend supplies only the epoch length of a task: pt reads its sharded per-rank loader, while pt_expt divides the dataset-wide batch count by the world size because every rank holds the complete data system, so an epoch denotes one pass over the dataset in either backend.
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThe PR adds shared epoch-based training schedule resolution and integrates it into both PyTorch trainers. It also introduces concurrent LMDB decoding, loader adapters, worker-count configuration, resource cleanup, LMDB merging changes, documentation, and regression tests. ChangesTraining scheduling and LMDB data pipelines
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Trainer
participant LmdbBatchDataLoader
participant LmdbBatchIterator
participant LmdbDataReader
Trainer->>LmdbBatchDataLoader: create training loader
LmdbBatchDataLoader->>LmdbBatchIterator: initialize sampler and workers
LmdbBatchIterator->>LmdbDataReader: decode requested frames
LmdbDataReader-->>LmdbBatchIterator: decoded batch arrays
LmdbBatchIterator-->>LmdbBatchDataLoader: return merged batch
Trainer->>LmdbBatchDataLoader: close in finally block
LmdbBatchDataLoader->>LmdbBatchIterator: release executor
LmdbBatchDataLoader->>LmdbDataReader: close reader
Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #5914 +/- ##
==========================================
- Coverage 79.21% 79.05% -0.16%
==========================================
Files 1069 1070 +1
Lines 124070 124537 +467
Branches 4522 4528 +6
==========================================
+ Hits 98278 98456 +178
- Misses 24171 24463 +292
+ Partials 1621 1618 -3 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (2)
deepmd/dpmodel/train/schedule.py (1)
121-123: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick winOver-specification is rejected for multi-task but silently resolved for single-task. The docstring promises
ValueErrorwhen the run length is over-specified, yet a single-task config with bothnumb_stepsandnumb_epochreturns early onnumb_stepswithout a word. Either mirror the mutual-exclusivity check here or soften the docstring to state thatnumb_stepstakes precedence.Also applies to: 142-147
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@deepmd/dpmodel/train/schedule.py` around lines 121 - 123, Update the schedule-building logic around the single-task early return and the corresponding multi-task branch to enforce mutual exclusivity when both num_steps and num_epoch are provided, raising ValueError as documented; do not silently prioritize num_steps, and preserve normal behavior when only one option is set.source/tests/pt/test_lmdb_dataloader.py (1)
1068-1071: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winThis shared-pool assertion can pass vacuously.
_executorstaysNoneuntil a parallel-eligible batch is submitted (num_workers > 1 and len(indices) >= num_workers). If the multitask fixture's batch size is below 2, both sides areNoneand the identity check succeeds without proving pool sharing. Assert the executor is notNonefirst.♻️ Suggested strengthening
+ assert trainer.training_dataloader["model_1"]._batch_iterator.started assert ( trainer.training_dataloader["model_1"]._batch_iterator._executor is trainer.training_dataloader["model_2"]._batch_iterator._executor )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/pt/test_lmdb_dataloader.py` around lines 1068 - 1071, Strengthen the shared-pool assertion in the multitask dataloader test by first asserting that the relevant batch iterator executor is not None, then retain the identity assertion comparing the executors for model_1 and model_2. Use the existing training_dataloader and _batch_iterator symbols without changing the sharing behavior being tested.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@deepmd/dpmodel/train/schedule.py`:
- Around line 121-130: Update the ValueError messages in the non-multi-task
schedule validation around StepSchedule to reference the actual numb_epoch
configuration key instead of training.num_epoch, while preserving the existing
numb_steps wording. Adjust the related test regex expectations in
test_train_schedule.py to accept the corrected key using the requested
num[b]?_epoch-compatible pattern.
In `@deepmd/dpmodel/utils/lmdb_data.py`:
- Around line 334-355: Update the min_pair_dist branch around
compute_min_pair_dist_single so it only computes when both “coord” and “atype”
are present in frame; otherwise populate min_pair_dist using the requirement’s
default value, preserving the existing dtype handling and find_min_pair_dist
marker.
In `@source/tests/pt_expt/test_training_ddp.py`:
- Around line 682-689: Update test_ranks_share_an_epoch to compute the expected
per-rank step count using the production ceil-based rounding rule, such as
math.ceil(nframes / 2), rather than floor division. Add the required math import
and keep both rank assertions aligned with this expected value.
In `@source/tests/pt_expt/utils/test_env.py`:
- Around line 15-29: Update test_lmdb_num_workers_partitions_node_budget to
monkeypatch os.sched_getaffinity with raising=False so it runs on platforms
where the attribute is absent. Refactor
test_lmdb_num_workers_rejects_invalid_override to use pytest.raises(ValueError),
adding the pytest import required by the module’s existing test style.
---
Nitpick comments:
In `@deepmd/dpmodel/train/schedule.py`:
- Around line 121-123: Update the schedule-building logic around the single-task
early return and the corresponding multi-task branch to enforce mutual
exclusivity when both num_steps and num_epoch are provided, raising ValueError
as documented; do not silently prioritize num_steps, and preserve normal
behavior when only one option is set.
In `@source/tests/pt/test_lmdb_dataloader.py`:
- Around line 1068-1071: Strengthen the shared-pool assertion in the multitask
dataloader test by first asserting that the relevant batch iterator executor is
not None, then retain the identity assertion comparing the executors for model_1
and model_2. Use the existing training_dataloader and _batch_iterator symbols
without changing the sharing behavior being tested.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 7ca39f2e-aa59-43fb-a3f7-3a2af68674c9
📒 Files selected for processing (17)
deepmd/dpmodel/train/__init__.pydeepmd/dpmodel/train/schedule.pydeepmd/dpmodel/utils/lmdb_data.pydeepmd/env.pydeepmd/pt/train/training.pydeepmd/pt/utils/lmdb_dataset.pydeepmd/pt_expt/train/training.pydeepmd/pt_expt/utils/lmdb_dataset.pydoc/env.mdsource/tests/common/dpmodel/test_lmdb_data.pysource/tests/common/dpmodel/test_train_schedule.pysource/tests/pt/test_lmdb_dataloader.pysource/tests/pt_expt/test_lmdb_training.pysource/tests/pt_expt/test_multitask.pysource/tests/pt_expt/test_training.pysource/tests/pt_expt/test_training_ddp.pysource/tests/pt_expt/utils/test_env.py
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Requesting changes for the inline LMDB epoch-order regression.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Reviewed the parallel-decode path, the new shared step schedule, and the worker sizing.
One thing worth recording up front, since it was the property I most expected to break: batch contents and ordering are independent of worker count. _submit slices into contiguous chunks in worker order, [f.result() for f in ...] preserves submission order, and _merge_lmdb_chunks concatenates in list order -- so DP_LMDB_NUM_WORKERS is a pure performance knob and does not perturb training numerics. Three separate passes over the code agreed on that. The merge_lmdb refcount fix (try/finally: _close_lmdb(src_path) replacing a direct src_env.close() on a refcounted handle) and the pin_memory=DEVICE.type != "cpu" correction are both genuine fixes in passing.
Four findings below plus two coverage asks. The two correctness ones (the epoch counter and the pt_expt epoch length) are both confined to DDP, which is also why no existing test catches them -- the added distributed tests use mocked epoch_length callables and never cross an epoch boundary at world_size > 1.
None of this is blocking from my side; the rebase question is the one I would settle first.
A run whose LMDB decoders were killed by the hangup of the session that started it stopped after three epochs and sat there for twenty hours. The training process was healthy and its main thread was waiting on a decode that could never arrive: the pool's own manager thread was itself blocked reading the partial result the dying decoder had left in the queue, so the pool never declared itself broken and nothing ever raised. Decoders no longer die of that hangup. A decoder is a background helper that owns no terminal, so the signal carries no meaning for it while the default disposition makes it fatal; it is ignored. The pool is built on `spawn` rather than `forkserver` for the same reason, a fork server being an unshielded intermediary whose own death is reported as the death of every decoder it started, which breaks the pool however well the decoders themselves are protected. Measured over two and sixteen workers: with a fork server a full hangup leaves no decoder alive and the pool unusable, while with `spawn` every decoder survives and the pool keeps serving. Losing a decoder anyway -- to the memory killer, or to a signal no process can ignore -- now costs throughput rather than the run. It reaches the iterator at three points, and all three fall back to decoding in the training process: the pool may refuse the submission, it may fail the futures it holds, or it may go quiet, which a liveness check between waits turns into the same outcome. The batch is re-decoded from the indices the pending batch carries, which is exact, because decoding only reads frames by key. Pool health lives on the shared pool entry, so a rank discovers the loss once however many data tasks share it. A pool stuck on a partial result also has to be dismantled rather than waited on. The interpreter joins every pool manager thread before it exits, so a run would otherwise complete its training, write its checkpoints and then never terminate, which a scheduler records as a timeout. Closing the write end the training process still holds delivers the end of file that read is waiting for. Measured: thirty seconds and counting before, two seconds to a clean exit after. `signal.SIGHUP` is absent on Windows, where the shielding is skipped rather than failing every decoder at startup.
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@source/tests/common/dpmodel/test_lmdb_data.py`:
- Around line 1710-1744: Guard both signal-based tests,
test_decoders_survive_the_hangup_of_their_launching_session and
test_killing_the_real_decoders_does_not_stop_the_run, so they are skipped on
non-POSIX platforms before referencing signal.SIGHUP or signal.SIGKILL. Use the
suite’s existing unittest/platform conventions where available and leave their
POSIX behavior unchanged.
- Around line 1689-1701: Update the subprocess test setup around
_acquire_lmdb_executor and _release_lmdb_executor so it does not unconditionally
access the private entry.executor._result_queue._writer; guard the partial-frame
write with an appropriate hasattr check (or replace it with a supported
mechanism), while preserving the malformed partial-frame scenario when the
writer is available.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: 48c5b852-5841-44d5-8226-bde4dab1f7fe
📒 Files selected for processing (6)
deepmd/dpmodel/utils/lmdb_data.pydeepmd/pt/utils/lmdb_dataset.pydeepmd/pt_expt/utils/lmdb_dataset.pysource/tests/common/dpmodel/test_lmdb_data.pysource/tests/pt/test_lmdb_dataloader.pysource/tests/pt_expt/test_lmdb_training.py
🚧 Files skipped from review as they are similar to previous changes (5)
- deepmd/pt/utils/lmdb_dataset.py
- source/tests/pt_expt/test_lmdb_training.py
- source/tests/pt/test_lmdb_dataloader.py
- deepmd/pt_expt/utils/lmdb_dataset.py
- deepmd/dpmodel/utils/lmdb_data.py
… decode Three defects in the parallel batch decoder surface once frames are grouped by label availability. The batch sequence is partitioned by that availability, so registering data requirements must not follow a partition built without them. The pt_expt data system tried to enforce this by rebuilding an iterator it no longer owns, leaving a stale attribute behind and rebuilding nothing. The reader already refuses a registration that follows any decode, and the sampler derives the partition only when the first batch is drawn, so no batch state can precede the call: the rebuild is unnecessary and the dead assignment goes. A test pins the reader's refusal, which nothing covered. Collation and preallocated decoding must agree on the batch they produce. Collation emitted the availability flags ahead of the payload while the preallocated path kept the frame's own key order, so the same batch was two different mappings depending on which path built it. Availability is now validated before collation begins and the batch is assembled in frame order, which leaves the validation intact and the two paths interchangeable. Frames were also grouped by an availability rule of their own, which could disagree with the one collation enforces and turn a grouping miss into a raise. Both now settle availability through the same predicate.
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
source/tests/pt_expt/test_lmdb_training.py (1)
195-214: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winClose every
LmdbDataSystemcreated by these tests.
serialand the streaming reader remain open, whiletest_data_requirements_freeze_after_first_readstarts a two-worker decoder pipeline without closing it. Usetry/finallyto close each system so test processes and LMDB handles cannot leak into later tests.Also applies to: 216-247, 249-263
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@source/tests/pt_expt/test_lmdb_training.py` around lines 195 - 214, The tests creating LmdbDataSystem instances must close every system, including serial, streaming, and the two-worker decoder used by test_data_requirements_freeze_after_first_read. Wrap each system’s usage in try/finally and call its close method in finally, covering test_streaming_batch_matches_frame_collation and the additional affected test blocks.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@source/tests/pt_expt/test_lmdb_training.py`:
- Around line 303-310: Update the batch assertions in the LMDB training test to
flatten all frame IDs across every rank and batch, then use assertCountEqual
against range(8). Remove the first-two-batches-only overlap check or make the
aggregate assertion the sole validation so duplicates anywhere fail while full
coverage remains required.
---
Outside diff comments:
In `@source/tests/pt_expt/test_lmdb_training.py`:
- Around line 195-214: The tests creating LmdbDataSystem instances must close
every system, including serial, streaming, and the two-worker decoder used by
test_data_requirements_freeze_after_first_read. Wrap each system’s usage in
try/finally and call its close method in finally, covering
test_streaming_batch_matches_frame_collation and the additional affected test
blocks.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro Plus
Run ID: bd40e97c-2bbf-42cd-8610-61504c8053d1
📒 Files selected for processing (8)
deepmd/dpmodel/utils/lmdb_data.pydeepmd/pt_expt/entrypoints/main.pydeepmd/pt_expt/train/training.pydeepmd/pt_expt/utils/lmdb_dataset.pysource/tests/common/dpmodel/test_lmdb_data.pysource/tests/pt/test_lmdb_dataloader.pysource/tests/pt_expt/test_lmdb_training.pysource/tests/pt_expt/utils/test_env.py
🚧 Files skipped from review as they are similar to previous changes (5)
- deepmd/pt_expt/utils/lmdb_dataset.py
- deepmd/pt_expt/train/training.py
- source/tests/common/dpmodel/test_lmdb_data.py
- source/tests/pt/test_lmdb_dataloader.py
- deepmd/dpmodel/utils/lmdb_data.py
wanghan-iapcm
left a comment
wanghan-iapcm
left a comment
There was a problem hiding this comment.
Verified all six threads against the current head rather than the replies.
The rebase reconciled cleanly: #5839 is now in the history, its structural-key hoisting and availability-signature grouping survive, and add_data_requirement still invalidates the signature cache. Handling refresh_batch_count() on the distributed sampler was a good catch that I had not asked for -- the availability partitioning changes the batch count, so the sharding fix and the rebase genuinely interact there.
The two distributed fixes are the substantive ones and both are right. DistributedSameNlocBatchSampler.__len__ is now rank-independent with deterministic padding in _partition_batches, so ranks cross the epoch boundary together instead of drifting apart when the batch count is not divisible by the world size. And plumbing rank/world_size through to LmdbDataSystem fixes the pt_expt side properly -- keeping nbatches as the global count so ceil(total / world_size) still matches the local sampler length is exactly the invariant that keeps the epoch arithmetic honest.
test_distributed_batches_are_sharded_with_equal_epoch_lengths is a real regression test, not a placeholder: three global batches on two ranks is the non-divisible case, and asserting len(sampler) == [2, 2] fails on the old rank-dependent length, while the disjointness assertion fails on the old shared-seed sampler where both ranks drew identical batches. One test that fails for both reasons before the fix.
The four merge guards and the five worker-count branches are now exercised.
On the parallel-decode threshold: your reasoning is fine and I am not pressing it. One leftover, non-blocking -- doc/env.md still says "The next batch is prefetched while the current batch is consumed" one sentence before explaining that small batches decode synchronously, and the deferred path stores indices rather than a future, so there is no prefetch there either. Now that the synchronous fallback is a deliberate choice rather than an oversight, that sentence would be worth qualifying whenever you next touch the file.
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Requesting changes for two reproducible LMDB correctness and lifecycle issues. Details are inline.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
njzjz-bot
left a comment
njzjz-bot
left a comment
There was a problem hiding this comment.
Reviewed current head 0d188183fb5ddd27dbe057fd4efe7fa42011b741. No blocking findings.
Validation:
ruff check .andruff format --check .- 15 shared schedule tests
- 79 shared LMDB tests, including decoder-loss and shutdown paths
- 62 PyTorch LMDB adapter/training tests
- 23 PyTorch Exportable LMDB and worker-configuration tests
- 4 epoch/multitask/distributed schedule tests
- CLI and PyTorch import smoke checks
The failed macOS arm64 wheel job is an infrastructure download failure (No buffer space available while fetching lmdb) after the wheel built, not a failure caused by this change.
Coding agent: Codex
Codex version: codex-cli 0.144.6
Model: gpt-5.6-sol
Reasoning effort: xhigh
2e39220
Summary
Checks
No test suite was run per request; the changes were tested separately on another machine. Repository pre-commit style hooks passed while creating the cherry-pick commit.
Summary by CodeRabbit
DP_LMDB_NUM_WORKERSoverride and auto-selection behavior.