Opt-in scope-profiler regions, off by default

CHANGELOG.md

@@ -17,6 +17,10 @@ format follows [Keep a Changelog](https://keepachangelog.com/en/1.1.0/).
   `gather_columns`.
 - Tests that launch `mpiexec -n 2` and `-n 4` jobs, a tutorial ("Moving markers between ranks"),
   guides and an example.
+- Opt-in scope-profiler regions (`mpiparticles.migrate`, `mpiparticles.compact`, ...) tagged
+  `"mpiparticles"`: off unless the application enables the tag, for example with
+  `scope_profiler.setup(enable_tags=["mpiparticles"])`. scope-profiler (>= 0.7.5) is a new
+  dependency.
 
 ### Changed
 

docs/src/content/docs/guides/index.md

@@ -10,3 +10,5 @@ sidebar:
 - [Ids](/mpiparticles/guides/ids/): globally unique int64 ids with a rank stride.
 - [Decomposition and migration](/mpiparticles/guides/migration/): the `Decomposition` protocol,
   `BoxDecomposition`, and what `migrate` sends.
+- [Profiling](/mpiparticles/guides/profiling/): the opt-in scope-profiler regions and how to turn
+  them on.

docs/src/content/docs/guides/profiling.md

@@ -0,0 +1,29 @@
+---
+title: Profiling
+description: The scope-profiler regions mpiparticles ships, off unless the application turns them on.
+sidebar:
+  order: 4
+---
+
+## Opt-in profiling regions
+
+mpiparticles times its expensive operations with [scope-profiler](https://github.com/max-models/scope-profiler)
+regions named `mpiparticles.<operation>`: `migrate`, `keep_owned`, `gather_columns`, and the store's
+`reserve` (only when the storage grows), `compact`, `sort_by_key`, `snapshot` and `restore`. They are
+tagged `"mpiparticles"`, an opt-in tag, so they stay off (and record nothing) when an application
+profiles its own code, until it asks for them:
+
+```python
+import scope_profiler as sp
+
+sp.setup(enable_tags=["mpiparticles"])
+```
+
+`session(enable_tags=...)` and `ProfilingOptions(enable_tags=...)` take the same argument.
+
+The same works without touching the code: `enable_tags = ["mpiparticles"]` in the TOML
+configuration, `SCOPE_PROFILER_ENABLE_TAGS=mpiparticles` in the environment, or
+`scope-profiler run --enable-tag mpiparticles script.py`. `"*"` enables every opt-in tag.
+
+A region that is off costs a few tenths of a microsecond per call, so only operations over all
+markers are instrumented, never per-marker work.

pyproject.toml

@@ -27,6 +27,7 @@ dependencies = [
   "maybempi>=0.1.2",
   "mpiarray>=0.2.1",
   "numpy",
+  "scope-profiler>=0.7.5",
 ]
 optional-dependencies.dev = [
   "jupyterlab",
@@ -87,3 +88,4 @@ report.show_missing = true
 cunumpy = { git = "https://github.com/max-models/cunumpy", branch = "devel" }
 maybempi = { git = "https://github.com/max-models/maybempi", branch = "devel" }
 mpiarray = { git = "https://github.com/max-models/mpiarray", branch = "devel" }
+scope-profiler = { git = "https://github.com/max-models/scope-profiler", branch = "devel" }

src/mpiparticles/__init__.py

@@ -6,6 +6,10 @@
 moves the markers that left a rank's subdomain to their owners. Without an MPI
 launcher everything runs serially, with no MPI calls.
 
+The expensive operations are scope-profiler regions named ``mpiparticles.<operation>``,
+tagged ``"mpiparticles"`` and off by default; enable them with
+``scope_profiler.setup(enable_tags=["mpiparticles"])``.
+
 Example
 -------
 >>> import cunumpy as xp
@@ -19,12 +23,18 @@
 
 from importlib.metadata import PackageNotFoundError, version
 
+import scope_profiler as sp
+
+from mpiparticles._profiling import PROFILE_TAGS
 from mpiparticles.collective import gather_columns, global_count
 from mpiparticles.decomposition import BoxDecomposition, Decomposition
 from mpiparticles.exchange import keep_owned, migrate
 from mpiparticles.ids import IdAllocator
 from mpiparticles.store import MarkerView, ParticleStore, StoreSnapshot
 
+# The profiling regions are off unless the application enables the tag.
+sp.register_opt_in_tag(*PROFILE_TAGS)
+
 try:
     __version__ = version("mpiparticles")
 except PackageNotFoundError:  # not installed, e.g. run from a source tree

src/mpiparticles/_profiling.py

@@ -0,0 +1,8 @@
+"""The tag of mpiparticles' own scope-profiler regions.
+
+The regions are opt-in: they record nothing unless the application enables the
+tag, for example with ``scope_profiler.setup(enable_tags=["mpiparticles"])``.
+"""
+
+#: Tags of every mpiparticles region; registered as opt-in in ``__init__``.
+PROFILE_TAGS = ("mpiparticles",)

src/mpiparticles/collective.py

@@ -7,7 +7,9 @@
 import cunumpy as xp
 import maybempi
 import numpy as np
+import scope_profiler as sp
 
+from mpiparticles._profiling import PROFILE_TAGS
 from mpiparticles.store import ParticleStore
 
 
@@ -23,6 +25,7 @@ def global_count(store: ParticleStore, comm: Any = None) -> int:
     return int(comm.allreduce(count, op=maybempi.get_mpi().SUM))
 
 
+@sp.profile("mpiparticles.gather_columns", tags=PROFILE_TAGS)
 def gather_columns(array: Any, comm: Any = None, root: int | None = 0) -> Any:
     """Concatenate component-major ``(ncomp, n)`` or ``(n,)`` arrays over the ranks.
 

src/mpiparticles/exchange.py

@@ -4,7 +4,9 @@
 
 import cunumpy as xp
 import mpiarray as mpa
+import scope_profiler as sp
 
+from mpiparticles._profiling import PROFILE_TAGS
 from mpiparticles._types import Array
 from mpiparticles.decomposition import Decomposition
 from mpiparticles.store import ParticleStore
@@ -19,6 +21,7 @@ def _leavers(store: ParticleStore, decomposition: Decomposition) -> Array:
     return xp.nonzero(store.alive[:end] & ~owned)[0]
 
 
+@sp.profile("mpiparticles.migrate", tags=PROFILE_TAGS)
 def migrate(store: ParticleStore, decomposition: Decomposition) -> tuple[int, int]:
     """Send the markers that left this rank's subdomain to their owners.
 
@@ -89,6 +92,7 @@ def migrate(store: ParticleStore, decomposition: Decomposition) -> tuple[int, in
     return count, received
 
 
+@sp.profile("mpiparticles.keep_owned", tags=PROFILE_TAGS)
 def keep_owned(store: ParticleStore, decomposition: Decomposition) -> int:
     """Drop the markers this rank does not own (not collective).
 

src/mpiparticles/store.py

@@ -6,7 +6,9 @@
 from typing import Any
 
 import cunumpy as xp
+import scope_profiler as sp
 
+from mpiparticles._profiling import PROFILE_TAGS
 from mpiparticles._types import Array
 from mpiparticles.ids import IdAllocator
 
@@ -202,14 +204,27 @@ def reserve(self, n: int) -> None:
         capacity = self.capacity
         if n <= capacity:
             return
-        while capacity < n:
-            capacity *= 2
-        old = (self._positions, self._velocities, self._weights, self._alive, self._ids)
-        self._allocate(capacity)
-        end = self._length
-        new = (self._positions, self._velocities, self._weights, self._alive, self._ids)
-        for target, source in zip(new, old, strict=True):
-            target[..., :end] = source[..., :end]
+        with sp.region("mpiparticles.reserve", tags=PROFILE_TAGS):
+            while capacity < n:
+                capacity *= 2
+            old = (
+                self._positions,
+                self._velocities,
+                self._weights,
+                self._alive,
+                self._ids,
+            )
+            self._allocate(capacity)
+            end = self._length
+            new = (
+                self._positions,
+                self._velocities,
+                self._weights,
+                self._alive,
+                self._ids,
+            )
+            for target, source in zip(new, old, strict=True):
+                target[..., :end] = source[..., :end]
 
     # -- adding and removing --------------------------------------------------
 
@@ -322,6 +337,7 @@ def _keep_first(self, count: int) -> None:
         self._clear_rows(count, self._length)
         self._length = count
 
+    @sp.profile("mpiparticles.compact", tags=PROFILE_TAGS)
     def compact(self) -> None:
         """Move the alive markers to the front, in order, and drop the dead ones."""
         end = self._length
@@ -362,18 +378,19 @@ def sort_by_key(self, keys: Array) -> None:
         keys = xp.asarray(keys)
         if keys.shape != (end,):
             raise ValueError(f"keys must have shape ({end},), got {keys.shape}")
-        rows = xp.nonzero(self._alive[:end])[0]
-        count = int(rows.size)
-        if count:
-            _, order = xp.algorithms.sort_by_key(keys[rows])
-            rows = rows[order]
-            # one contiguous row at a time: faster than a 2-D column gather
-            for array in (self._positions, self._velocities):
-                for row in range(array.shape[0]):
-                    array[row, :count] = array[row][rows]
-            for array in (self._weights, self._ids):
-                array[:count] = array[rows]
-        self._keep_first(count)
+        with sp.region("mpiparticles.sort_by_key", tags=PROFILE_TAGS):
+            rows = xp.nonzero(self._alive[:end])[0]
+            count = int(rows.size)
+            if count:
+                _, order = xp.algorithms.sort_by_key(keys[rows])
+                rows = rows[order]
+                # one contiguous row at a time: faster than a 2-D column gather
+                for array in (self._positions, self._velocities):
+                    for row in range(array.shape[0]):
+                        array[row, :count] = array[row][rows]
+                for array in (self._weights, self._ids):
+                    array[:count] = array[rows]
+            self._keep_first(count)
 
     # -- alive markers --------------------------------------------------------
 
@@ -409,6 +426,7 @@ def set_active_weights(self, weights: Array) -> None:
 
     # -- checkpointing --------------------------------------------------------
 
+    @sp.profile("mpiparticles.snapshot", tags=PROFILE_TAGS)
     def snapshot(self) -> StoreSnapshot:
         """Return a copy of all markers (alive and dead) and the id counter."""
         end = self._length
@@ -433,16 +451,17 @@ def restore(self, snapshot: StoreSnapshot) -> None:
                 f"snapshot is for ndim={snapshot.ndim}, nvel={snapshot.nvel}, "
                 f"this store has ndim={self._ndim}, nvel={self._nvel}"
             )
-        end = snapshot.length
-        self.reserve(end)
-        self._clear_rows(0, self._length)
-        self._positions[:, :end] = snapshot.positions
-        self._velocities[:, :end] = snapshot.velocities
-        self._weights[:end] = snapshot.weights
-        self._alive[:end] = snapshot.alive
-        self._ids[:end] = snapshot.ids
-        self._length = end
-        self._allocator.next_id = snapshot.next_id
+        with sp.region("mpiparticles.restore", tags=PROFILE_TAGS):
+            end = snapshot.length
+            self.reserve(end)
+            self._clear_rows(0, self._length)
+            self._positions[:, :end] = snapshot.positions
+            self._velocities[:, :end] = snapshot.velocities
+            self._weights[:end] = snapshot.weights
+            self._alive[:end] = snapshot.alive
+            self._ids[:end] = snapshot.ids
+            self._length = end
+            self._allocator.next_id = snapshot.next_id
 
     # -- plasmacoll's MarkerSet -----------------------------------------------
 

src/mpiparticles/tests/test_profiling.py

@@ -0,0 +1,42 @@
+"""mpiparticles' scope-profiler regions are opt-in."""
+
+import pytest
+import scope_profiler as sp
+
+from mpiparticles import BoxDecomposition, migrate
+from mpiparticles.tests.conftest import FakeComm, filled_store
+
+
+@pytest.fixture(autouse=True)
+def _fresh_profiler(monkeypatch):
+    """Keep the profiling state of one test out of the others."""
+    monkeypatch.delenv("SCOPE_PROFILER_ENABLE_TAGS", raising=False)
+    sp.reset()
+    yield
+    sp.reset()
+
+
+def _profiled_region_names(**setup_kwargs):
+    """Return the regions recorded while migrating and compacting a store."""
+    sp.setup(output=None, **setup_kwargs)
+    store = filled_store(5, capacity=4)  # grows once
+    migrate(store, BoxDecomposition(FakeComm(0, 1), [0, 0], [1, 1], [1, 1]))
+    store.compact()
+    with sp.region("user.step"):
+        pass
+    results = sp.finalize(verbose=False, return_results=True)
+    return set(results.region_names)
+
+
+def test_regions_are_off_by_default():
+    assert _profiled_region_names() == {"user.step"}
+
+
+def test_enable_tags_turns_the_regions_on():
+    names = _profiled_region_names(enable_tags=["mpiparticles"])
+    assert names == {
+        "mpiparticles.compact",
+        "mpiparticles.migrate",
+        "mpiparticles.reserve",
+        "user.step",
+    }