Make statistics a mandatory arg for statistics providers by nirandaperera · Pull Request #1017 · rapidsai/rapidsmpf

nirandaperera · 2026-05-07T21:20:18Z

Make Statistics a mandatory constructor argument for providers

Follow-up to #1009. After the initial stats cleanup, BufferResource, ProgressThread, and streaming::Context still defaulted their statistics parameter to Statistics::disabled(), which made it easy to silently drop stats by forgetting to pass an instance and obscured ownership at the call site (e.g. communicator vs. progress thread vs. buffer resource each potentially holding a different Statistics).

In this PR:

BufferResource(...)/BufferResource::from_options(...) now require an explicit Statistics (passed as the first ctor arg); ProgressThread(...) does too. Communicator stores a shared_ptr<ProgressThread> and forwards statistics() to it, so Single/MPI/UCXX drop their own progress-thread members. Tests, examples, and Python bindings are updated; the Python Communicator gains a statistics property.
New MutableStatisticsProvider concept and ProgressThread::set_statistics() (wired through to Python) allow swapping the instance at runtime; the progress loop is paused around the swap.
MemoryRecorder and StreamOrderedTiming now hold Statistics weakly so a swap can destroy the original even with outstanding scopes; the destructors lock and skip publishing when the instance has expired.
Benchmarks default to stats enabled (previously Statistics::disabled() / from_options defaulting to off) and use clear() before the final run to keep per-run reporting.

Depends on #1009

Related to #1008

copy-pr-bot · 2026-05-07T21:20:28Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

nirandaperera · 2026-05-27T00:30:22Z

/ok to test e5214b4

Signed-off-by: niranda perera <niranda.perera@gmail.com>

This reverts commit a6803e6.

Signed-off-by: niranda perera <niranda.perera@gmail.com>

pentschev · 2026-05-28T10:03:30Z

+    auto stats = rapidsmpf::Statistics::from_options(options);
+    auto progress_thread = std::make_shared<rapidsmpf::ProgressThread>(stats);


This changes the existing behavior, I think unintentionally, from statistics being enabled by default to disabled by default. As is, users now require specifying RAPIDSMPF_STATISTICS=1, which I think is undesirable.

I think all benchmarks should have enabled statistics. So yes, I think you should create() the stats. Rather than from_options()ing them

wence- · 2026-05-28T10:51:33Z

+    /**
+     * @brief Resets the current state of this instance based on @p new_options.
+     *
+     * @warning Live `MemoryRecorder` instances (constructed before this call
+     * but not yet destroyed) are *not* tracked. Their destructors will run
+     * after the reset and publish their scope into the freshly cleared
+     * `memory_records_` map, mixing pre-reset measurements into the new
+     * session. Likewise, a `StreamOrderedTiming` whose stop callback is
+     * already mid-execution at reset time may still record one stale stat
+     * (its global-map entry was extracted before `cancel_inflight_timings`
+     * could remove it). Callers that need clean boundaries must quiesce
+     * recorders and synchronise the relevant CUDA streams before calling.
+     *
+     * @param new_options Configuration options whose `"statistics"` field
+     * controls the new enabled state.
+     */
+    void reset_from_options(config::Options new_options);
+


What do you need this for?

@wence- In cudf-polars engines, there is a _reset_worker(...options) method. There, we are only resetting buffer resource, context, etc with a new stats object. Then, after calling _reset_worker, we would have multiple instances of statistics in the engine ([commuicator, progress thread] vs [buffer resource, context])

But why do we need to reset the statistics in this setup? When is _reset_worker called? Is it only for tests? For every query, something else?

What do we actually need to achieve here? It seems like we want to zero the statistics object that we have around no?

I was looking into two options.

Introduce set_stats for communicator and progress thread and reset the new stats object

reset the stats object in-place (clear + enable/disable based on new options)

I think 1. is the safer option.

That is the mechanism for how to reset stats. I am asking a different question: why do we want to reset stats?

I think this stems from StreamingEngine._reset method. https://github.com/rapidsai/cudf/blob/main/python/cudf_polars/cudf_polars/engine/core.py#L289

OK, _reset is only called in the tests (it is a "hack" to workaround that booting a dask or ray cluster is expensive), and I do not think that we have to tie ourselves in knots to publicly support that api also being able to reset statistics. Note that the docstring of _reset already says that it doesn't reset and reinitialize with all the options.

I think we have enough rope with .enable() and .clear() on the stats object, without introducing a new thing.

okay. Let's do that.

I felt clear() was insufficient because it only clears the stats, but not the memory records and report entries. So, we can have a scenario where we create a new Buffer resource with new rmm resource adapter, but stats contains old memory records from a previous buffer resource.

ah, but that problem can be fixed (separately) by clearing memory records

wence- · 2026-05-28T10:56:44Z

+    auto stats = rapidsmpf::Statistics::from_options(options);
+    auto progress_thread = std::make_shared<rapidsmpf::ProgressThread>(stats);


I think all benchmarks should have enabled statistics. So yes, I think you should create() the stats. Rather than from_options()ing them

wence- · 2026-05-28T10:57:26Z

        if (i == total_num_runs - 1) {
-            ctx->statistics()->clear();
+            stats->clear();
        }


Does the ctx not advertise a statistics object any more? That seems wrong.

It does, but its the same as stats in this scope

wence- · 2026-05-28T11:11:38Z

Some broader comments:

I still think that statistics providers should return a shared_ptr<Statistics> const&. That way the caller can either upgrade to an owning shared_ptr or just keep a reference. The converse is not possible: if we return a shared_ptr<Statistics> there is no way to downgrade it to a reference.
Everything that has a Statistics object "somewhere" should be a statistics provider.

pentschev · 2026-05-28T11:34:36Z

I still think that statistics providers should return a shared_ptr<Statistics> const&. That way the caller can either upgrade to an owning shared_ptr or just keep a reference. The converse is not possible: if we return a shared_ptr<Statistics> there is no way to downgrade it to a reference.

What I don't like about returning a const-ref to a shared_ptr is that it's easy to overlook and assume you're getting a copy with potential for use-after-free. One can argue that's true for all const-ref, but a smart pointer falls under a "safe" category in my mental model at least.

wence- · 2026-05-28T11:56:50Z

I still think that statistics providers should return a shared_ptr<Statistics> const&. That way the caller can either upgrade to an owning shared_ptr or just keep a reference. The converse is not possible: if we return a shared_ptr<Statistics> there is no way to downgrade it to a reference.

What I don't like about returning a const-ref to a shared_ptr is that it's easy to overlook and assume you're getting a copy with potential for use-after-free. One can argue that's true for all const-ref, but a smart pointer falls under a "safe" category in my mental model at least.

So there are some usual options, I think:

struct Foo {
    std::shared_ptr<Statistics> const& statistics();
};

struct Owns {
   std::shared_ptr<Statistics> stats_;
};

Owns takes_ref(std::shared_ptr<Statistics> const& stat) {
   return Owns{stat};
}
...

auto stats = std::make_shared<Statistics>();
Foo foo{stats};

foo.statistics()->some_stats_method(); // This can never be use-after-free because foo is alive

auto owning = takes_ref(foo.statistics()); // transparently upgraded in Owns ctor

auto owning_stat = foo.statistics(); // owning

auto &ref_stat = foo.statistics(); // reference, but obvious

pentschev · 2026-05-28T13:22:12Z

So there are some usual options, I think:

struct Foo {
    std::shared_ptr<Statistics> const& statistics();
};

struct Owns {
   std::shared_ptr<Statistics> stats_;
};

Owns takes_ref(std::shared_ptr<Statistics> const& stat) {
   return Owns{stat};
}
...

auto stats = std::make_shared<Statistics>();
Foo foo{stats};

foo.statistics()->some_stats_method(); // This can never be use-after-free because foo is alive

auto owning = takes_ref(foo.statistics()); // transparently upgraded in Owns ctor

auto owning_stat = foo.statistics(); // owning

auto &ref_stat = foo.statistics(); // reference, but obvious

That looks reasonable to me.

madsbk · 2026-05-28T14:21:22Z

But you only get a reference if you explicitly assign to one, right?

auto& ref_stat = foo.statistics();  // reference, but explicit

Please, let's not introduce a special "reference semantics for smart pointers" concept like:

Owns takes_ref(std::shared_ptr<Statistics> const& stat) {
    return Owns{stat};
}

If you assign to a reference, you are explicitly choosing to take on the associated lifetime risk. That is already true for references in general, and arguably even more so for non-const references. So if someone is not confident about the lifetime guarantees, they simply should not bind to a reference.

wence- · 2026-05-28T14:44:51Z

But you only get a reference if you explicitly assign to one, right?
auto& ref_stat = foo.statistics();  // reference, but explicit

Yes. You have to request a reference explicitly to get a reference

wence- · 2026-05-28T14:46:20Z

Note, I too do not advocate the "takes a reference and upgrades it some time later" model.

nirandaperera · 2026-05-28T16:17:34Z

@wence- @pentschev I am marking this as draft because, I want to revisit MutableStatisticsProvider, set_statistics() approach.
a6803e6

I think this is a more safer approach.

2ce411f 891af67 Signed-off-by: niranda perera <niranda.perera@gmail.com>

Signed-off-by: niranda perera <niranda.perera@gmail.com>

…ory-arg-for-providers Signed-off-by: niranda perera <niranda.perera@gmail.com>

Co-authored-by: Peter Andreas Entschev <peter@entschev.com>

Signed-off-by: niranda perera <niranda.perera@gmail.com>

…perera/rapidsmpf into stat-mandatory-arg-for-providers

Signed-off-by: niranda perera <niranda.perera@gmail.com>

pentschev · 2026-05-29T09:41:40Z

+
+        Warnings
+        --------
+        Concurrent calls to ``set_statistics`` will result in undefined behavior.


Why is this warning not documented in the C++ docstring?

It is, isnt it? https://github.com/rapidsai/rapidsmpf/pull/1017/changes/BASE..78d8a04e9a906720832b17d43a8192ea0823c691#diff-f9320c48070fa670c43a05334be536be21c0c779e5ab0a8e44d41f978339b0d3R201

Yeah, sorry about that, don't know how I missed it. 😅

pentschev · 2026-05-29T09:48:30Z

+    /**
+     * @brief Replace the statistics instance held by this progress thread.
+     *
+     * This method pauses the progress loop for the duration of the swap so the loop
+     * cannot dereference `statistics_` while it is being replaced, then restores the
+     * prior running state.
+     *
+     * @param statistics The new statistics instance. Pass `Statistics::disabled()` to opt
+     * out of statistics collection.
+     *
+     * @throws std::invalid_argument If @p statistics is null. Validated
+     * before the pause so a null argument does not leave the progress loop
+     * paused.
+     *
+     * @warning Concurrent calls to `set_statistics()` will result in undefined behavior.
+     */
+    void set_statistics(std::shared_ptr<Statistics> statistics);


This makes ProgressThread a mutable provider, but existing Statistics handles are snapshots. A caller can do auto stats = progress->statistics(), then another path calls set_statistics(new_stats), and the caller continues recording/reporting against the old object while ProgressThread and Communicator report the new one. I think this reintroduces the "which stats object owns this?" ambiguity the PR is trying to remove. If runtime reconfiguration is needed, mutating the existing Statistics instance is safer than swapping identity. If swapping remains, docs/tests should explicitly state that previously returned handles are not retargeted.

Fair point. This is how I looked at this.
Currently, Statistics only gets reset when we reset the engine. While running, we initiate async stat recording, like Memory records, stream ordered timings etc. Recording only happens at a later point, and there could be a scenario where the recording happens after a reset. If stat object is the same, we might end up with some stats from before.
This PR handles it by handing off a weak ptr to async stats recording, and if the ptr has expired, recording is skipped.

Ok, but what will happen to those objects that rely on the Statistics to be still alive? Do they have a mechanism to eventually get a new one, or will they just fall out of the Statistics loop forever?

@pentschev If they rely on Statistics, it could be a problem. If they want to get the most current statistics, they would have to keep a strong reference to long-living stats provider like ProgressThread.

will they just fall out of the Statistics loop forever?

Yes. At least in current usages, IMO it will be safe to miss some recordings on the edge cases.

What are the edge cases?

Currently, Statistics only gets reset when we reset the engine.

Or more generically, when is the engine expected to reset? I assume be "engine reset" you're referring to cudf-polars' engine, is that right? One case where reset may occur is part of the testing framework, is that what you have in mind? What would other cases be?

Yes, edge cases I was referring to was, in-flight stats recordings while engine is being reset.

During tests, I think this is not happening, because we have a long running communicator/ progress thread in the test environment.

Ok, can you please at a minimum document that previously returned handles are not retargeted?

Signed-off-by: niranda perera <niranda.perera@gmail.com>

pentschev · 2026-06-01T15:11:29Z

+ * @warning Concurrent calls to `set_statistics()` are implementation-defined and
+ * likely result in undefined behavior unless the implementation explicitly
+ * documents otherwise. Callers should assume external synchronization is required.


Nit: let's not leave this hanging on implementation details and just say this is not thread-safe/concurrent calls are UB.

pentschev · 2026-06-01T15:16:56Z

+
+        Warnings
+        --------
+        Concurrent calls to ``set_statistics`` will result in undefined behavior.


Yeah, sorry about that, don't know how I missed it. 😅

pentschev · 2026-06-01T16:55:45Z

+    /**
+     * @brief Replace the statistics instance held by this progress thread.
+     *
+     * This method pauses the progress loop for the duration of the swap so the loop
+     * cannot dereference `statistics_` while it is being replaced, then restores the
+     * prior running state.
+     *
+     * @param statistics The new statistics instance. Pass `Statistics::disabled()` to opt
+     * out of statistics collection.
+     *
+     * @throws std::invalid_argument If @p statistics is null. Validated
+     * before the pause so a null argument does not leave the progress loop
+     * paused.
+     *
+     * @warning Concurrent calls to `set_statistics()` will result in undefined behavior.
+     */
+    void set_statistics(std::shared_ptr<Statistics> statistics);


Ok, but what will happen to those objects that rely on the Statistics to be still alive? Do they have a mechanism to eventually get a new one, or will they just fall out of the Statistics loop forever?

jameslamb · 2026-06-01T21:57:57Z

          - '!.github/release.yml'
          - '!.pre-commit-config.yaml'
          - '!.shellcheckrc'
+          - '!SECURITY.md'


It's surprising to see these changes from #1064 showing up in the diff again. Those are already out on main.

Same with the cupy changes throughout packaging files. I suspect something has gone wrong with resolving a merge conflict here.

Oh yes! I think I've messed up bigtime during a merge..

I fixed the merge mess @jameslamb

Ah great, thanks very much. Commit history looks much better now, and you no longer need ci-codeowners / packaging-codeowners so I'll unsubscribe from this.

@ me back if you need my help with anything.

…ory-arg-for-providers Signed-off-by: niranda perera <niranda.perera@gmail.com>

nirandaperera · 2026-06-08T15:14:29Z

closing in favor of #1082

nirandaperera requested review from a team as code owners May 7, 2026 21:20

nirandaperera added breaking Introduces a breaking change improvement Improves an existing functionality labels May 7, 2026

nirandaperera marked this pull request as draft May 7, 2026 21:20

nirandaperera mentioned this pull request May 26, 2026

Initial stats cleanup #1009

Merged

nirandaperera force-pushed the stat-mandatory-arg-for-providers branch 2 times, most recently from afd1218 to 3b2d2da Compare May 26, 2026 22:56

making statistics mandatory arg

9445b50

nirandaperera force-pushed the stat-mandatory-arg-for-providers branch from e5214b4 to 9445b50 Compare May 27, 2026 20:57

nirandaperera marked this pull request as ready for review May 27, 2026 20:58

missing changes

2d93022

Signed-off-by: niranda perera <niranda.perera@gmail.com>

nirandaperera changed the title ~~[WIP] Make statistics a mandatory arg for statistics providers~~ Make statistics a mandatory arg for statistics providers May 27, 2026

nirandaperera added 3 commits May 27, 2026 16:47

set_stats

a6803e6

Signed-off-by: niranda perera <niranda.perera@gmail.com>

Revert "set_stats"

891af67

This reverts commit a6803e6.

reset from options

2ce411f

Signed-off-by: niranda perera <niranda.perera@gmail.com>

nirandaperera mentioned this pull request May 28, 2026

Confusing behavior of Context.from_options and statistics #1065

Open

pentschev reviewed May 28, 2026

View reviewed changes

wence- reviewed May 28, 2026

View reviewed changes

wence- marked this pull request as draft May 28, 2026 17:07

nirandaperera added 4 commits May 28, 2026 10:57

revert

36cd37a

2ce411f 891af67 Signed-off-by: niranda perera <niranda.perera@gmail.com>

set stats and weak ptr

2eca536

Signed-off-by: niranda perera <niranda.perera@gmail.com>

move progress thread to communicator

9ae3d1c

Signed-off-by: niranda perera <niranda.perera@gmail.com>

Merge branch 'main' of github.com:rapidsai/rapidsmpf into stat-mandat…

9f0c309

…ory-arg-for-providers Signed-off-by: niranda perera <niranda.perera@gmail.com>

nirandaperera marked this pull request as ready for review May 28, 2026 22:05

nirandaperera requested a review from wence- May 28, 2026 22:05

nirandaperera and others added 3 commits May 28, 2026 15:17

Apply suggestions from code review

5659dc0

Co-authored-by: Peter Andreas Entschev <peter@entschev.com>

addressing initial comments

4e0d47c

Signed-off-by: niranda perera <niranda.perera@gmail.com>

Merge branch 'stat-mandatory-arg-for-providers' of github.com:niranda…

bdda47a

…perera/rapidsmpf into stat-mandatory-arg-for-providers

nirandaperera requested a review from pentschev May 28, 2026 22:20

addressing comments

37ff748

Signed-off-by: niranda perera <niranda.perera@gmail.com>

nirandaperera mentioned this pull request May 28, 2026

[WIP] Rapidsmpf Stats API changes rapidsai/cudf#22690

Closed

3 tasks

add test

78d8a04

Signed-off-by: niranda perera <niranda.perera@gmail.com>

pentschev reviewed May 29, 2026

View reviewed changes

addressing comments

ab50f50

Signed-off-by: niranda perera <niranda.perera@gmail.com>

pentschev reviewed Jun 1, 2026

View reviewed changes

nirandaperera requested review from a team as code owners June 1, 2026 20:37

nirandaperera requested a review from KyleFromNVIDIA June 1, 2026 20:37

jameslamb reviewed Jun 1, 2026

View reviewed changes

Merge branch 'main' of github.com:rapidsai/rapidsmpf into stat-mandat…

30df058

…ory-arg-for-providers Signed-off-by: niranda perera <niranda.perera@gmail.com>

nirandaperera force-pushed the stat-mandatory-arg-for-providers branch from e8d1b76 to 30df058 Compare June 2, 2026 18:33

jameslamb removed the request for review from KyleFromNVIDIA June 3, 2026 21:34

nirandaperera closed this Jun 8, 2026

		auto stats = rapidsmpf::Statistics::from_options(options);
		auto progress_thread = std::make_shared<rapidsmpf::ProgressThread>(stats);

Uh oh!

Conversation

nirandaperera commented May 7, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

copy-pr-bot Bot commented May 7, 2026

Uh oh!

nirandaperera commented May 27, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nirandaperera May 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nirandaperera May 28, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

wence- commented May 28, 2026

Uh oh!

pentschev commented May 28, 2026

Uh oh!

wence- commented May 28, 2026

Uh oh!

pentschev commented May 28, 2026

Uh oh!

madsbk commented May 28, 2026

Uh oh!

wence- commented May 28, 2026

Uh oh!

wence- commented May 28, 2026

Uh oh!

nirandaperera commented May 28, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

nirandaperera commented May 7, 2026 •

edited

Loading

nirandaperera May 28, 2026 •

edited

Loading

nirandaperera May 28, 2026 •

edited

Loading