fix(gh-230): top-level progress bar #237
Conversation
WalkthroughReplaces string-based strategy identifiers with (strategy_name, phase) tuples across Sampler and resource bundles, adds a Phase enum, tags bundle step sequences with phases, and wraps Sampler’s strategy loop with a tqdm progress bar. Integration tests updated to pass tuple-based strategy_order. Changes
Sequence Diagram(s)sequenceDiagram
autonumber
actor User
participant Sampler
participant Progress as tqdm
participant Strategy
User->>Sampler: sample(..., strategy_order=[(name, phase), ...])
Sampler->>Progress: create(total=len(strategy_order))
loop for each (strategy, phase) in strategy_order
Sampler->>Progress: set_description(phase)
Sampler->>Strategy: run(strategy, rng_key, resources, last_step, data)
Strategy-->>Sampler: (rng_key, resources, last_step)
Sampler->>Progress: update(1)
end
Sampler-->>User: results
sequenceDiagram
autonumber
participant Bundle as RQSpline_*_Bundle
participant Seq as strategy_order
Note over Bundle,Seq: Build phase-tagged execution sequence
Bundle->>Seq: append (initialize_model, INITIALIZATION)
Bundle->>Seq: append (…training steps…, TRAINING)
Bundle->>Seq: append (reset_steppers/update_state, INTERMEDIATE)
Bundle->>Seq: append (…production steps…, PRODUCTION)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~55 minutes Possibly related PRs
Poem
Tip 🔌 Remote MCP (Model Context Protocol) integration is now available!Pro plan users can now connect to remote MCP servers from the Integrations page. Connect with popular remote MCPs such as Notion and Linear to add more context to your reviews and chats. 📜 Recent review detailsConfiguration used: CodeRabbit UI Review profile: CHILL Plan: Pro 💡 Knowledge Base configuration:
You can enable these sources in your CodeRabbit configuration. 📒 Files selected for processing (1)
🚧 Files skipped from review as they are similar to previous changes (1)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
✨ Finishing Touches
🧪 Generate unit tests
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. 🪧 TipsChatThere are 3 ways to chat with CodeRabbit:
SupportNeed help? Create a ticket on our support page for assistance with any issues or questions. CodeRabbit Commands (Invoked using PR/Issue comments)Type Other keywords and placeholders
CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (8)
src/flowMC/resource_strategy_bundle/base.py (1)
27-27: Type-safety: let strategy_order accept Phase (not just str) to avoid .value plumbingDownstream code currently stores Phase.value strings. Allowing the enum directly improves discoverability and static checking while remaining backward compatible.
Apply this diff:
- strategy_order: list[tuple[str, str]] + strategy_order: list[tuple[str, str | Phase]]If you adopt this, you can pass Phase members directly from bundles and let the Sampler render them (see my Sampler.py comment for a small helper that handles either type).
src/flowMC/Sampler.py (3)
106-119: Progress bar should honor verbose and handle Phase enums transparentlyAlways printing a tqdm bar can clutter logs/tests. Also, accept either str or Phase for the label.
Apply this diff:
- with tqdm.tqdm( - self.strategy_order, - total=len(self.strategy_order), - ) as strategy_order_bar: - for strategy, phase in strategy_order_bar: - strategy_order_bar.set_description(phase) + with tqdm.tqdm( + self.strategy_order, + total=len(self.strategy_order), + disable=not self.verbose, + ) as strategy_order_bar: + for strategy, phase in strategy_order_bar: + phase_label = getattr(phase, "value", phase) # allow Phase or str + strategy_order_bar.set_description(str(phase_label)) ( rng_key, self.resources, last_step, ) = self.strategies[ strategy ](rng_key, self.resources, last_step, data)
119-121: Persist RNG key after samplingWithout storing the advanced rng_key back, repeated calls may reuse the original key and repeat sequences.
Apply this diff:
- + # Persist the advanced RNG key for deterministic continuation across subsequent calls. + self.rng_key = rng_key
34-34: Optional typing: accept Phase in type hints to align with ResourceStrategyBundleThis keeps the public API consistent if bundles start emitting Phase rather than strings.
Apply this diff:
- strategy_order: Optional[list[tuple[str, str]]] + strategy_order: Optional[list[tuple[str, str | "Phase"]]] @@ - strategy_order: None | list[tuple[str, str]] = None, + strategy_order: None | list[tuple[str, str | "Phase"]] = None,Note: Quoted "Phase" keeps runtime import optional; if you prefer, import Phase from flowMC.resource_strategy_bundle.base and drop the quotes.
Also applies to: 48-48
src/flowMC/resource_strategy_bundle/RQSpline_MALA_PT.py (1)
318-326: Prefer storing Phase enums in strategy_order; stringify at render timeYou’re storing Phase.value strings. Keeping the enum in strategy_order improves type safety and avoids stringly-typed APIs. Sampler already handles either type in my suggested change.
Apply this diff:
- ("parallel_tempering", Phase.TRAINING.value), - ("local_stepper", Phase.TRAINING.value), - ("update_global_step", Phase.TRAINING.value), - ("model_trainer", Phase.TRAINING.value), - ("update_model", Phase.TRAINING.value), - ("global_stepper", Phase.TRAINING.value), - ("update_local_step", Phase.TRAINING.value), + ("parallel_tempering", Phase.TRAINING), + ("local_stepper", Phase.TRAINING), + ("update_global_step", Phase.TRAINING), + ("model_trainer", Phase.TRAINING), + ("update_model", Phase.TRAINING), + ("global_stepper", Phase.TRAINING), + ("update_local_step", Phase.TRAINING), @@ - ("parallel_tempering", Phase.PRODUCTION.value), - ("local_stepper", Phase.PRODUCTION.value), - ("update_global_step", Phase.PRODUCTION.value), - ("global_stepper", Phase.PRODUCTION.value), - ("update_local_step", Phase.PRODUCTION.value), + ("parallel_tempering", Phase.PRODUCTION), + ("local_stepper", Phase.PRODUCTION), + ("update_global_step", Phase.PRODUCTION), + ("global_stepper", Phase.PRODUCTION), + ("update_local_step", Phase.PRODUCTION), @@ - strategy_order = [("initialize_tempered_positions", Phase.INITIALIZATION.value)] + strategy_order = [("initialize_tempered_positions", Phase.INITIALIZATION)]Also applies to: 328-333, 335-335
test/integration/test_HMC.py (1)
85-85: Phase label is free-form; consider using the enum for consistency"testing_phase" works since the phase is only used for the progress bar label. If you standardize on Phase, tests can reflect that.
Optionally:
+from flowMC.resource_strategy_bundle.base import Phase @@ - strategy_order=[("take_steps", "testing_phase")], + strategy_order=[("take_steps", Phase.PRODUCTION.value)],test/integration/test_MALA.py (1)
77-77: Same note on phase label consistencyUsing a Phase enum value improves consistency across tests (purely optional).
Example:
+from flowMC.resource_strategy_bundle.base import Phase @@ - strategy_order=[("take_steps", "testing_phase")], + strategy_order=[("take_steps", Phase.PRODUCTION.value)],test/integration/test_RWMCMC.py (1)
84-85: Prefer using the Phase enum over a free-form string for the phase label.Using a literal like "testing_phase" won’t break anything today (Sampler uses it only for tqdm labels), but adopting the enum everywhere reduces typo risk and will survive future validation that may enforce allowed phases.
Apply this diff:
+from flowMC.resource_strategy_bundle.base import Phase @@ - strategy_order=[("take_steps", "testing_phase")], + strategy_order=[("take_steps", Phase.TRAINING.value)],
📜 Review details
Configuration used: CodeRabbit UI
Review profile: CHILL
Plan: Pro
💡 Knowledge Base configuration:
- MCP integration is disabled by default for public repositories
- Jira integration is disabled by default for public repositories
- Linear integration is disabled by default for public repositories
You can enable these sources in your CodeRabbit configuration.
📒 Files selected for processing (7)
src/flowMC/Sampler.py(4 hunks)src/flowMC/resource_strategy_bundle/RQSpline_MALA.py(2 hunks)src/flowMC/resource_strategy_bundle/RQSpline_MALA_PT.py(2 hunks)src/flowMC/resource_strategy_bundle/base.py(2 hunks)test/integration/test_HMC.py(2 hunks)test/integration/test_MALA.py(2 hunks)test/integration/test_RWMCMC.py(2 hunks)
🧰 Additional context used
🧬 Code graph analysis (7)
test/integration/test_MALA.py (2)
src/flowMC/Sampler.py (1)
Sampler(12-128)src/flowMC/strategy/take_steps.py (2)
TakeSerialSteps(147-180)TakeSteps(14-144)
src/flowMC/resource_strategy_bundle/base.py (2)
src/flowMC/resource/base.py (1)
Resource(6-38)src/flowMC/strategy/base.py (1)
Strategy(8-31)
test/integration/test_HMC.py (3)
src/flowMC/resource/states.py (1)
State(8-63)src/flowMC/Sampler.py (1)
Sampler(12-128)src/flowMC/strategy/take_steps.py (4)
TakeSerialSteps(147-180)TakeSteps(14-144)__call__(60-144)body(156-162)
src/flowMC/resource_strategy_bundle/RQSpline_MALA.py (1)
src/flowMC/resource_strategy_bundle/base.py (2)
Phase(8-14)ResourceStrategyBundle(17-27)
src/flowMC/resource_strategy_bundle/RQSpline_MALA_PT.py (1)
src/flowMC/resource_strategy_bundle/base.py (2)
Phase(8-14)ResourceStrategyBundle(17-27)
test/integration/test_RWMCMC.py (3)
src/flowMC/Sampler.py (1)
Sampler(12-128)src/flowMC/strategy/take_steps.py (3)
TakeSteps(14-144)__call__(60-144)sample(164-180)src/flowMC/strategy/base.py (1)
Strategy(8-31)
src/flowMC/Sampler.py (2)
src/flowMC/resource_strategy_bundle/base.py (1)
ResourceStrategyBundle(17-27)src/flowMC/strategy/base.py (1)
Strategy(8-31)
⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (2)
- GitHub Check: build (3.12)
- GitHub Check: build (3.11)
🔇 Additional comments (9)
src/flowMC/resource_strategy_bundle/base.py (1)
8-14: Enum introduction looks goodAdding Phase centralizes legal phase names and prevents typos downstream. Good call.
src/flowMC/resource_strategy_bundle/RQSpline_MALA_PT.py (1)
234-243: No Issue: Lambda Strategies Already Return the Correct TupleThe
Lambdastrategy’s__call__method invokes the provided function for its side effect and then unconditionally returns(rng_key, resources, initial_position)which satisfies the
Strategyprotocol of returning a(rng_key, resources, last_step)triple. No change is needed.• File:
src/flowMC/strategy/lambda_function.pylines 25–30
• The lambda passed in may returnNone, butLambda.__call__ignores that return value and returns the required triple.Likely an incorrect or invalid review comment.
test/integration/test_HMC.py (1)
7-7: Imports reorg is fineThe added imports align with the updated API and local test setup.
Also applies to: 9-9, 11-11
test/integration/test_MALA.py (1)
7-7: Imports reorg is fineMatches the new API usage in the test body.
Also applies to: 9-9, 11-11
test/integration/test_RWMCMC.py (1)
7-11: Imports look correct and consistent with the refactor.Importing GaussianRandomWalk, State, and TakeSerialSteps aligns with the new strategy scheduling and kernel layout.
src/flowMC/resource_strategy_bundle/RQSpline_MALA.py (4)
3-17: Import updates are appropriate for the new functionality.eqx is needed for tree_at, State is used for sampler_state, and Phase/ResourceStrategyBundle/Takes* align with the phase-aware scheduling. No issues spotted.
260-264: Production phase tagging looks consistent.The ordering mirrors training and uses Phase.PRODUCTION.value uniformly. No correctness concerns.
271-273: Nice insertion of INTERMEDIATE steps between training and production.reset_steppers followed by update_state is a sensible boundary to prevent buffer position drift and to flip the training flag. This sequencing matches expectations.
251-258: Allstrategy_orderusages have been migrated to tuple-based style
- Ran a repo-wide regex search; no old-style plain-list definitions remain.
- Every
strategy_orderreference—including in test/integration/*.py—is already using the new(strategy, phase)tuples.No further changes needed here.
Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com>
1 participant
This PR contains the implementation to add top level progress bar.
I have slightly refactored the main sampling loop to check for any invalid strategy outside the main loop.
Summary by CodeRabbit
New Features
Refactor
Tests