Expose Python sampler controls - #165
Draft
hartikainen wants to merge 8 commits into
Draft
hartikainen wants to merge 8 commits into
hartikainen wants to merge 8 commits into
Conversation
Currently, Reverb passes `py::call_guard` directly to the `episode_steps` property declaration. As explained in [pybind11's upstream change](pybind/pybind11#5533), property declarations silently ignored these guards, and the added compile-time check rejects that usage, preventing the binding from compiling. This PR wraps the `episode_steps` getter in `py::cpp_function` and applies the call guard to that function so that reading the property releases the GIL as intended.
Currently, an op wrapper with no `ops_lib` drops its kernel dependency because operator precedence makes the conditional apply to the entire dependency list. This PR adds parentheses around the optional `ops_lib` dependency so that the generated shared library always links its kernel.
The build helpers currently have two issues: they rely on native rules and pass `depset` values directly to attributes that expect lists of labels. The rule references fail when their built-in or automatically loaded definitions are unavailable, while the dependency arguments cause attribute type errors. This PR addresses both issues by explicitly loading `cc_library`, using the already-loaded `py_library`, and converting the dependency sets to lists.
Currently, downstream Bazel builds do not inherit Reverb's `.bazelrc` flags, which can cause unresolved-symbol linker errors when building macOS shared libraries that depend on symbols supplied at runtime. This PR moves the dynamic symbol lookup setting from the global linker flags to the shared-library targets so that downstream builds receive the required linker configuration.
This PR adds Bzlmod support to the build system. Depends on google-deepmind#162.
As a follow-up to google-deepmind#163, this PR replaces `WORKSPACE`-based builds with Bzlmod. Depends on google-deepmind#163.
With builds using Bzlmod as of google-deepmind#163 and google-deepmind#164, this change upgrades Bazel from `7.7.0` to `9.2.0` and updates the dependency rules and configuration for that version. The Bazel rules use Protobuf `35.1`, while a separate native Protobuf dependency stays at `31.1` for TensorFlow compatibility. Keeping these dependencies separate allows the build rules to advance without changing the native runtime linked with TensorFlow. The dependency patches use the corresponding rule providers and native libraries.
Expose cancellation, worker counts, and sampling deadlines through the Python bindings. Deadline errors preserve their exception type, and invalid timeout values are rejected. Include the binding declarations and direct regression tests.
hartikainen
force-pushed
the
kristian/sampler-python-controls
branch
from
September 23, 2026 19:26
7d280e6 to
f932644
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
I was experimenting with Reverb and wanted to cancel Python sampler reads and configure sampling workers and deadlines without using
tf.data. This PR exposesSampler.Close()and optionalnum_workersandrate_limiter_timeout_msarguments onClient.NewSampler(), preserving existing positional calls and default behavior.GetNextTrajectory()now reports a sampling deadline asreverb.DeadlineExceededError, and timeout values below-1raiseValueError. The PR also adds regression tests for cancellation, worker configuration, timeout handling, and compatibility with existing calls.