diff --git a/.github/workflows/build-test.yaml b/.github/workflows/build-test.yaml index 72c08a8..2716bbd 100644 --- a/.github/workflows/build-test.yaml +++ b/.github/workflows/build-test.yaml @@ -271,6 +271,16 @@ jobs: run: | python -m pytest -vv --reruns 3 --reruns-delay 30 --only-rerun "(?i)http|timeout|connection|socket|resolve" + # Docker and many JupyterHubs give containers a 64 MB /dev/shm. Results + # larger than that must fall back to the pipe instead of crashing. + - name: Run shared-memory tests with a 64 MB /dev/shm + if: matrix.os == 'ubuntu-latest' && matrix.python-version == '3.10' + run: | + sudo mount -o remount,size=64m /dev/shm + df -h /dev/shm + python -m pytest tests/test_transfer.py -vv -rs | tee transfer-tests.log + grep -q "test_result_larger_than_dev_shm_falls_back_to_pipe PASSED" transfer-tests.log + - name: Run the ported Geant4 examples if: matrix.python-version == '3.12' && matrix.os == 'ubuntu-latest' run: | diff --git a/PYTHON_DEVELOPER_GUIDE.md b/PYTHON_DEVELOPER_GUIDE.md index 16052c4..c8b58b1 100644 --- a/PYTHON_DEVELOPER_GUIDE.md +++ b/PYTHON_DEVELOPER_GUIDE.md @@ -1985,7 +1985,8 @@ app.command("/gun/energy 1 MeV") -> Message(target="", method="command", args=("/gun/energy 1 MeV",), kwargs={}) -> pickled and sent through a multiprocessing Pipe with a sequence number -> worker: getattr(cpp_app, "command")("/gun/energy 1 MeV") - -> the result (or the exception) is pickled back + -> the result (or the exception) is pickled back; a result holding more than + 1 MB of awkward arrays goes through shared memory instead (below) -> Application raises the exception again in your process, or returns the result ``` @@ -2003,7 +2004,17 @@ Consequences to keep in mind: (section 3). - **Everything crosses the process boundary by pickling.** Arguments and results must be picklable, and large event arrays are copied from the worker to your - process. That copy is part of the cost of a run (section 30). + process. That copy is part of the cost of a run (section 30). To keep it + cheap, `_transfer.py` writes the raw buffers of a large result + (`ak.to_buffers`) into one `multiprocessing.shared_memory` block and pickles + only the forms, lengths and offsets. Your process unlinks the block, copies it + out and rebuilds the arrays with `ak.from_buffers`. Pickling through the pipe + remains the fallback for small results and for blocks that can't be allocated + (for example a full `/dev/shm` on Linux). Your process names each block, so it + unlinks the block itself if the worker dies mid-transfer. + `GEANT4_PYTHON_APPLICATION_SHM_MIN_BYTES` (set before `start()`) changes the 1 + MB threshold: `0` sends every array result through shared memory, a huge value + turns it off. - **Python code cannot run inside transport.** A Python function cannot be sent to the worker and called from a Geant4 stepping action, which is why `run_with_callbacks` calls your callbacks after the run, over the returned diff --git a/pyproject.toml b/pyproject.toml index 3d20533..9fbe267 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -99,6 +99,7 @@ filterwarnings = [ "error" ] log_cli_level = "INFO" +markers = ["slow: takes more than a few seconds"] minversion = "6.0" testpaths = ["tests"] xfail_strict = true diff --git a/src/geant4_python_application/_transfer.py b/src/geant4_python_application/_transfer.py new file mode 100644 index 0000000..b4b56e7 --- /dev/null +++ b/src/geant4_python_application/_transfer.py @@ -0,0 +1,276 @@ +"""Hand results from the Geant4 child process to the parent. + +Pickling a run's awkward arrays through ``multiprocessing.Pipe`` runs at +~60-80 MB/s on macOS, which made it most of ``Application.run``'s wall time. +Large results go through one POSIX shared-memory block instead: the child +writes the arrays' raw buffers (``ak.to_buffers``) into it and sends only the +forms, lengths and buffer offsets through the pipe. The parent unlinks the +block, copies it out and rebuilds the arrays with ``ak.from_buffers``. + +Small results, results without awkward arrays, and results whose block can't +be allocated are pickled through the pipe as before. + +Cleanup: the parent chooses each block's name (``block_name``), so it can +unlink the block itself when the child dies mid-transfer. The child unlinks the +block if anything fails before the descriptor is sent, and the parent unlinks +it as soon as it has attached. Should both processes crash, multiprocessing's +resource tracker, which they share, unlinks it once they have exited. + +Environment variables, read by the child when it packs a result: + +``GEANT4_PYTHON_APPLICATION_SHM_MIN_BYTES`` + Results with fewer bytes of array data go through the pipe (default + 1 MiB). ``0`` sends every array result through shared memory; a huge value + disables shared memory. +``GEANT4_PYTHON_APPLICATION_SHM_TEST_PAUSE`` + Testing only: seconds to sleep after writing a block, before sending it. +""" + +from __future__ import annotations + +import os +import sys +import time +from multiprocessing import shared_memory + +import awkward as ak +import numpy as np + +# Below this many bytes of array data the pipe is as fast and needs no cleanup. +MIN_SHARED_BYTES = 1 << 20 +MIN_BYTES_VARIABLE = "GEANT4_PYTHON_APPLICATION_SHM_MIN_BYTES" +TEST_PAUSE_VARIABLE = "GEANT4_PYTHON_APPLICATION_SHM_TEST_PAUSE" +_ALIGNMENT = 64 +_BUFFER_KEY = "{form_key}-{attribute}" +_BYTEORDER = "<" if sys.byteorder == "little" else ">" + +# Parent side: how many results holding awkward arrays arrived by each path. +stats = {"shared_memory": 0, "pipe": 0} + + +class SharedArray: + """Stands in for an ``ak.Array`` whose buffers are in a ``SharedResult`` block.""" + + __slots__ = ("attrs", "behavior", "buffers", "form", "length") + + def __init__(self, form, length, buffers, behavior, attrs): + self.form = form + self.length = length + self.buffers = buffers # buffer key -> (offset, nbytes) + self.behavior = behavior + self.attrs = attrs + + def __getstate__(self): + return {name: getattr(self, name) for name in self.__slots__} + + def __setstate__(self, state): + for name, value in state.items(): + setattr(self, name, value) + + +class SharedResult: + """What goes through the pipe instead of a large result.""" + + __slots__ = ("name", "payload", "size") + + def __init__(self, name: str, size: int, payload): + self.name = name + self.size = size + # The result with every ak.Array replaced by a SharedArray. + self.payload = payload + + def __getstate__(self): + return {name: getattr(self, name) for name in self.__slots__} + + def __setstate__(self, state): + for name, value in state.items(): + setattr(self, name, value) + + +def block_name(prefix: str, counter: int) -> str: + """Name of the block for message ``counter`` (macOS allows 30 characters).""" + return f"{prefix}_{counter:x}" + + +def _map_arrays(obj, function): + """Copy of nested dicts/lists/tuples with each ak.Array replaced by function(array).""" + if isinstance(obj, ak.Array): + return function(obj) + if type(obj) is dict: + return {key: _map_arrays(value, function) for key, value in obj.items()} + if type(obj) in (list, tuple): + return type(obj)(_map_arrays(value, function) for value in obj) + return obj + + +def _arrays(obj): + found = [] + _map_arrays(obj, found.append) + return found + + +def _min_bytes() -> int: + try: + return int(os.environ[MIN_BYTES_VARIABLE]) + except (KeyError, ValueError): + return MIN_SHARED_BYTES + + +def _allocate(size: int, name: str | None = None) -> shared_memory.SharedMemory: + block = shared_memory.SharedMemory(name=name, create=True, size=size) + if sys.platform.startswith("linux"): + # /dev/shm is a tmpfs that can be smaller than the result (64 MB in + # Docker by default). Writing past its limit raises SIGBUS and kills + # the process, so reserve the pages now and fall back if they don't fit. + try: + os.posix_fallocate(block._fd, 0, size) + except BaseException: + block.close() + block.unlink() + raise + return block + + +def destroy(name: str) -> None: + """Unlink the block called name, if it exists.""" + try: + block = shared_memory.SharedMemory(name=name) + except FileNotFoundError: + return + block.close() + block.unlink() + + +def pack(result, name: str | None = None, min_bytes: int | None = None): + """Child side: the object to send through the pipe in place of result.""" + if min_bytes is None: + min_bytes = _min_bytes() + arrays = _arrays(result) + if not arrays or sum(array.layout.nbytes for array in arrays) < min_bytes: + return result + + try: + flattened = {} + offset = 0 + for array in arrays: + if id(array) in flattened: + continue + form, length, container = ak.to_buffers( + ak.to_packed(array, highlevel=False), + buffer_key=_BUFFER_KEY, + form_key="node{id}", + byteorder=_BYTEORDER, + ) + buffers = {} + for key, values in container.items(): + raw = np.ascontiguousarray(values).reshape(-1).view(np.uint8) + buffers[key] = (offset, raw) + offset += -(-raw.nbytes // _ALIGNMENT) * _ALIGNMENT + flattened[id(array)] = (array, form, length, buffers) + size = max(offset, 1) + + block = _allocate(size, name) + except Exception: + return result + + try: + _write(block, size, flattened.values()) + + def to_shared(array): + _, form, length, buffers = flattened[id(array)] + attrs = getattr(array, "attrs", None) + if attrs is not None: + attrs = {k: v for k, v in attrs.items() if not k.startswith("@")} + return SharedArray( + form, + length, + { + key: (start, buffer.nbytes) + for key, (start, buffer) in buffers.items() + }, + None if array.behavior is ak.behavior else array.behavior, + attrs, + ) + + packed = SharedResult(block.name, size, _map_arrays(result, to_shared)) + except BaseException as error: + block.close() + block.unlink() + if isinstance(error, Exception): + return result + raise + block.close() + + pause = os.environ.get(TEST_PAUSE_VARIABLE) + if pause: + time.sleep(float(pause)) + return packed + + +def _write(block, size, flattened) -> None: + data = np.ndarray((size,), dtype=np.uint8, buffer=block.buf) + try: + for _, _, _, buffers in flattened: + for start, buffer in buffers.values(): + data[start : start + buffer.nbytes] = buffer + finally: + # block.close() refuses while a view of its buffer is alive. + del data + + +def discard(message) -> None: + """Free the block of a packed result that won't be unpacked.""" + if isinstance(message, SharedResult): + destroy(message.name) + + +def unpack(message): + """Parent side: the result that pack(result) stood in for.""" + if not isinstance(message, SharedResult): + if _arrays(message): + stats["pipe"] += 1 + return message + + try: + block = shared_memory.SharedMemory(name=message.name) + except FileNotFoundError: + msg = f"Shared memory block {message.name} holding the result is gone" + raise RuntimeError(msg) from None + try: + # Unlinking first means nothing is left behind if the copy is interrupted; + # the mapping stays valid until close(). + block.unlink() + data = np.empty(message.size, dtype=np.uint8) + data[:] = np.ndarray((message.size,), dtype=np.uint8, buffer=block.buf) + finally: + block.close() + + def from_shared(shared): + container = { + key: data[start : start + nbytes] + for key, (start, nbytes) in shared.buffers.items() + } + kwargs = {} if shared.attrs is None else {"attrs": shared.attrs} + return ak.from_buffers( + shared.form, + shared.length, + container, + buffer_key=_BUFFER_KEY, + byteorder=_BYTEORDER, + behavior=shared.behavior, + **kwargs, + ) + + result = _replace_shared(message.payload, from_shared) + stats["shared_memory"] += 1 + return result + + +def _replace_shared(obj, function): + if isinstance(obj, SharedArray): + return function(obj) + if type(obj) is dict: + return {key: _replace_shared(value, function) for key, value in obj.items()} + if type(obj) in (list, tuple): + return type(obj)(_replace_shared(value, function) for value in obj) + return obj diff --git a/src/geant4_python_application/application.py b/src/geant4_python_application/application.py index efe5de6..17515b0 100644 --- a/src/geant4_python_application/application.py +++ b/src/geant4_python_application/application.py @@ -3,6 +3,7 @@ import importlib.util import multiprocessing import os +import secrets import threading from collections import namedtuple from pathlib import Path @@ -11,6 +12,7 @@ import geant4_python_application import geant4_python_application.events +from geant4_python_application import _transfer from geant4_python_application._geant4_application import ( Application as Geant4Application, ) @@ -75,7 +77,7 @@ def _visualize(app, commands: list[str]) -> None: _geant4_vis.start_visualization(commands) -def _start_application(pipe: multiprocessing.Pipe): +def _start_application(pipe: multiprocessing.Pipe, block_prefix: str): app = Geant4Application() app.set_event_fields(_default_event_fields) while True: @@ -106,7 +108,13 @@ def _start_application(pipe: multiprocessing.Pipe): raise ValueError(f"Unknown method: {method}") result = getattr(target, method)(*message.args, **message.kwargs) - pipe.send((counter, result)) + # Large awkward results go through shared memory, not the pipe. + result = _transfer.pack(result, _transfer.block_name(block_prefix, counter)) + try: + pipe.send((counter, result)) + except BaseException: + _transfer.discard(result) + raise except KeyboardInterrupt: pass except EOFError or BrokenPipeError: @@ -129,8 +137,13 @@ def __init__( geant4_python_application.install_datasets(show_progress=True) self._pipe, child_pipe = multiprocessing.Pipe() + # Names this application's shared-memory blocks, so a block left by a + # child that died mid-transfer can be found and unlinked. + self._block_prefix = "g4" + secrets.token_hex(5) self._process = multiprocessing.Process( - target=_start_application, args=(child_pipe,), daemon=True + target=_start_application, + args=(child_pipe, self._block_prefix), + daemon=True, ) self._message_counter = 0 self._lock = threading.Lock() @@ -190,14 +203,21 @@ def _send_and_recv(self, message: Message): if self._process is None or not self._process.is_alive(): raise RuntimeError("Application is not running") with self._lock: + block = _transfer.block_name(self._block_prefix, self._message_counter) try: self._send(self._message_counter, message) counter, response = self._recv() if counter != self._message_counter: + _transfer.discard(response) raise RuntimeError("Message counter mismatch") self._message_counter += 1 + response = _transfer.unpack(response) except (EOFError, BrokenPipeError): raise RuntimeError("Application process died. Recreate the application") + finally: + # unpack has normally unlinked it already. This covers a child + # that died mid-transfer and an interrupted receive. + _transfer.destroy(block) if isinstance(response, Exception): raise response return response diff --git a/tests/test_transfer.py b/tests/test_transfer.py new file mode 100644 index 0000000..906e901 --- /dev/null +++ b/tests/test_transfer.py @@ -0,0 +1,256 @@ +from __future__ import annotations + +import os +import pickle +import signal +import sys +import threading +import time +from multiprocessing import shared_memory +from pathlib import Path + +import awkward as ak +import numpy as np +import pytest + +from geant4_python_application import Application, _transfer, basic_gdml + + +def _block_exists(name: str) -> bool: + try: + block = shared_memory.SharedMemory(name=name) + except FileNotFoundError: + return False + block.close() + return True + + +def _thread_result(n: int) -> dict: + return { + "id": ak.Array(np.arange(n, dtype=np.int64)), + "track_particle": ak.Array([["e-", "gamma"][: i % 3] for i in range(n)]), + "step_position": ak.Array( + [[[{"x": 1.0 * i, "y": 2.0, "z": 3.0}] * (i % 4)] for i in range(n)] + ), + "maybe": ak.Array([None if i % 2 else i for i in range(n)]), + "point": ak.Array([{"x": 1.0 * i} for i in range(n)], with_name="Point"), + } + + +def test_large_result_goes_through_shared_memory(): + result = [_thread_result(50_000), _thread_result(7)] + packed = _transfer.pack(result, min_bytes=0) + assert isinstance(packed, _transfer.SharedResult) + assert len(pickle.dumps(packed)) < 10_000 + assert _block_exists(packed.name) + + unpacked = _transfer.unpack(pickle.loads(pickle.dumps(packed))) + assert not _block_exists(packed.name) + assert isinstance(unpacked, list) + for original, copy in zip(result, unpacked): + assert original.keys() == copy.keys() + for key in original: + assert copy[key].type == original[key].type + assert ak.array_equal(copy[key], original[key]) + + +def test_small_and_non_array_results_are_pickled(): + small = {"id": ak.Array([1, 2, 3])} + assert _transfer.pack(small) is small + for value in (None, 3, "text", {"a"}, [1, 2], ValueError("x")): + assert _transfer.pack(value, min_bytes=0) is value + assert _transfer.unpack(value) is value + + +def test_discard_unlinks_block(): + packed = _transfer.pack(_thread_result(100), min_bytes=0) + assert _block_exists(packed.name) + _transfer.discard(packed) + assert not _block_exists(packed.name) + with pytest.raises(RuntimeError, match="gone"): + _transfer.unpack(packed) + + +def test_allocation_failure_falls_back_to_pickle(monkeypatch): + def fail(*_): + msg = "no space" + raise OSError(msg) + + monkeypatch.setattr(_transfer, "_allocate", fail) + result = {"id": ak.Array(np.arange(1000))} + assert _transfer.pack(result, min_bytes=0) is result + + +def test_unpack_failure_still_unlinks_block(): + packed = _transfer.pack(_thread_result(100), min_bytes=0) + for shared in packed.payload.values(): + shared.length += 10**6 # buffers too short: from_buffers raises + with pytest.raises(Exception): # noqa: B017, PT011 + _transfer.unpack(packed) + assert not _block_exists(packed.name) + + +# -- through the Geant4 child process ------------------------------------------ + +# 1 GeV e- into ~20 X0 of PbWO4: ~4 MB of default event fields per event. +PBWO4_GDML = """ + + + + + + + + + + + + + + + + + + + + + + + +""" +ELECTRON_COMMANDS = ["/gun/particle e-", "/gun/energy 1 GeV"] +CALO_COMMANDS = [*ELECTRON_COMMANDS, "/gun/position 0 0 -10 mm", "/gun/direction 0 0 1"] + +DEV_SHM = Path("/dev/shm") + + +def _dev_shm_blocks() -> set[str] | None: + """Blocks made by this package or by SharedMemory's defaults (Linux only).""" + if not sys.platform.startswith("linux") or not DEV_SHM.is_dir(): + return None + return {p.name for p in DEV_SHM.iterdir() if p.name.startswith(("psm_", "g4"))} + + +def _next_block(app: Application) -> str: + return _transfer.block_name(app._block_prefix, app._message_counter) + + +@pytest.fixture +def no_leaks(): + """Fails the test if it leaves blocks in /dev/shm (checked on Linux).""" + before = _dev_shm_blocks() + yield + after = _dev_shm_blocks() + if before is not None: + assert after - before == set() + + +def test_large_result_uses_shared_memory_small_uses_pipe(no_leaks): + with Application(gdml=basic_gdml, seed=1) as app: + app.commands(ELECTRON_COMMANDS) + before = dict(_transfer.stats) + name = _next_block(app) + events = app.run(100) + assert events.layout.nbytes > _transfer.MIN_SHARED_BYTES + assert _transfer.stats["shared_memory"] == before["shared_memory"] + 1 + assert not _block_exists(name) + + before = dict(_transfer.stats) + events = app.run(1) + assert events.layout.nbytes < _transfer.MIN_SHARED_BYTES + assert _transfer.stats["pipe"] == before["pipe"] + 1 + assert _transfer.stats["shared_memory"] == before["shared_memory"] + + +def _seeded_events(monkeypatch, min_bytes: int): + monkeypatch.setenv(_transfer.MIN_BYTES_VARIABLE, str(min_bytes)) + with Application(gdml=basic_gdml, seed=1234) as app: + app.commands(ELECTRON_COMMANDS) + before = dict(_transfer.stats) + events = app.run(30) + used = {k: _transfer.stats[k] - before[k] for k in before} + return events, used + + +def test_shared_memory_and_pipe_give_identical_arrays(monkeypatch, no_leaks): + shared, used_shared = _seeded_events(monkeypatch, 0) + piped, used_pipe = _seeded_events(monkeypatch, 10**15) + assert used_shared == {"shared_memory": 1, "pipe": 0} + assert used_pipe == {"shared_memory": 0, "pipe": 1} + + assert shared.fields == piped.fields + assert shared.track.fields == piped.track.fields + assert shared.track.step.fields == piped.track.step.fields + assert str(shared.type) == str(piped.type) # nesting, dtypes and record names + assert ak.array_equal(shared, piped, dtype_exact=True, check_parameters=True) + assert ak.sum(ak.num(shared.track.step.energy, axis=2)) > 0 + + +def test_no_blocks_left_after_a_run_that_raises(no_leaks): + with Application(gdml=basic_gdml) as app: + app.commands(ELECTRON_COMMANDS) + app.run(100) # through shared memory + app.command("/gun/position 0 0 1 m") + name = _next_block(app) + with pytest.raises(RuntimeError, match="outside the world"): + app.run(100) + assert not _block_exists(name) + app.command("/gun/position 0 0 0 cm") + assert len(app.run(100)) == 100 + + +def test_no_blocks_left_when_child_is_killed_mid_transfer(monkeypatch, no_leaks): + # The child sleeps after writing the block, before sending its descriptor. + monkeypatch.setenv(_transfer.TEST_PAUSE_VARIABLE, "60") + app = Application(gdml=basic_gdml) + app.start() + app.commands(ELECTRON_COMMANDS) + name = _next_block(app) + errors = [] + + def run(): + try: + app.run(100) + except RuntimeError as error: + errors.append(error) + + thread = threading.Thread(target=run) + thread.start() + deadline = time.monotonic() + 120 + while not _block_exists(name): + assert time.monotonic() < deadline, "block was never written" + assert thread.is_alive(), "run returned before writing a block" + time.sleep(0.05) + + os.kill(app._process.pid, signal.SIGKILL) + thread.join(30) + assert not thread.is_alive() + assert len(errors) == 1 + assert "died" in str(errors[0]) + assert not _block_exists(name) + + +@pytest.mark.slow +@pytest.mark.skipif( + not sys.platform.startswith("linux") + or not DEV_SHM.is_dir() + or os.statvfs(DEV_SHM).f_blocks * os.statvfs(DEV_SHM).f_frsize > 128 * 2**20, + reason="needs a Linux /dev/shm of at most 128 MB (e.g. docker --shm-size=64m)", +) +def test_result_larger_than_dev_shm_falls_back_to_pipe(no_leaks): + shm_bytes = os.statvfs(DEV_SHM).f_blocks * os.statvfs(DEV_SHM).f_frsize + with Application(gdml=PBWO4_GDML, n_threads=8, seed=1) as app: + app.commands(CALO_COMMANDS) + before = dict(_transfer.stats) + events = app.run(48) + assert events.layout.nbytes > 1.5 * shm_bytes + assert _transfer.stats["pipe"] == before["pipe"] + 1 + assert _transfer.stats["shared_memory"] == before["shared_memory"] + assert len(events) == 48 + assert ak.all(events.id == np.arange(48)) + assert ak.all(ak.num(events.track.id, axis=1) > 100) + # Smaller results still fit and use shared memory. + events = app.run(2) + assert events.layout.nbytes < shm_bytes / 4 + assert _transfer.stats["shared_memory"] == before["shared_memory"] + 1 diff --git a/tools/profile_breakdown.py b/tools/profile_breakdown.py new file mode 100644 index 0000000..b462a30 --- /dev/null +++ b/tools/profile_breakdown.py @@ -0,0 +1,546 @@ +"""Wall-time breakdown of geant4_python_application runs, to judge GPU offload. + +Question: would offloading e-/e+/gamma transport to a GPU (Celeritas/AdePT) +or optical photons to Opticks speed this package up, or is the time spent +recording steps into awkward builders and handing the result to Python? + +Every workload/thread-count pair is measured in two fresh processes: + +1. In-process: the compiled ``_geant4_application.Application`` (what the + public ``Application`` drives in its child process), timing ``run()``. + That is BeamOn plus the builder -> awkward snapshot in EndOfRunAction, + with no pipe. Event-field variants: + + * ``minimal`` -- fields {"id"}: Geant4 transport with step recording + reduced as far as the API allows. SteppingAction still calls + data::InsertStep for every step; with no step field selected it only + does ~10 hash-set lookups and copies the process name. + * ``default`` -- the package's default event fields. + * ``no-optical`` -- optical workload only: minimal with Scintillation + and Cerenkov inactivated, so minimal - no-optical is the cost of + generating and tracking optical photons. + +2. Public API: ``geant4_python_application.Application.run`` with the + default fields, timed in this process, with the time spent waiting for + the child process (``_send_and_recv``) recorded separately. + +Then, for the default fields, + + transport = in-process minimal + recording = in-process default - in-process minimal - snapshot estimate + snapshot = ak.to_buffers + ak.from_buffers over the returned arrays + (a proxy for the C++ to_buffers + from_buffers in + EndOfRunAction, which cannot be timed from outside) + ipc = public child wait - in-process default (moving the + awkward arrays to the parent: one shared-memory block for + results over 1 MB, else pickle + multiprocessing.Pipe; see + geant4_python_application._transfer) + python = public total - public child wait (concatenate, build the + Geant4Event record, sort by id) + +Usage:: + + python tools/profile_breakdown.py --threads 0 8 --out tmp/profile.json + +``--threads 0`` is the serial run manager. Prints one markdown table per +workload; ``--out`` also saves the raw numbers as JSON. Multithreaded runs +use events-per-thread times the thread count. +""" + +from __future__ import annotations + +import argparse +import json +import multiprocessing +import re +import statistics +import sys +import time +from pathlib import Path + +import awkward as ak + +import geant4_python_application as g4 +from geant4_python_application import _transfer +from geant4_python_application.application import Message, _default_event_fields + +# ~20 X0 of PbWO4 (X0 = 0.89 cm) along z, 30 cm x 30 cm transverse +# (Moliere radius 2.0 cm), in vacuum. +PBWO4_GDML = R""" + + + + + + + + + + + + + + + + + + + + + +""" + +# 1 m water cube; a 10 MeV gamma has a ~50 cm mean free path in water. +WATER_GDML = R""" + + + + + + + + + + + + + + + + + + + + + +""" + +WORKLOADS = { + "em_calo": { + "title": "1 GeV e- into 20 X0 PbWO4", + "gdml": PBWO4_GDML, + "optical": False, + "commands": [ + "/gun/particle e-", + "/gun/energy 1 GeV", + "/gun/position 0 0 -10 mm", + "/gun/direction 0 0 1", + ], + "events": 10, + }, + "gamma_water": { + "title": "10 MeV gamma into water", + "gdml": WATER_GDML, + "optical": False, + "commands": [ + "/gun/particle gamma", + "/gun/energy 10 MeV", + "/gun/position 0 0 -499 mm", + "/gun/direction 0 0 1", + ], + "events": 2000, + }, + "optical": { + "title": "10 MeV e- in optical water scintillator (optical.optical_water_gdml)", + "gdml": g4.optical.optical_water_gdml, + "optical": True, + # As in examples/geant4-examples/05_optical_scintillation.py + "commands": g4.optical.scintillation_commands( + stack_photons=True, max_cerenkov=100 + ) + + [ + "/gun/particle e-", + "/gun/energy 10 MeV", + "/gun/position 0 0 0 mm", + "/gun/direction 1 0 0", + ], + "events": 50, + }, +} + +MINIMAL_FIELDS = {"id"} +DEFAULT_FIELDS = set(_default_event_fields) +OPTICAL_PROCESSES = ("Scintillation", "Cerenkov") + +CATEGORIES = ["e-", "e+", "gamma", "opticalphoton", "hadrons", "other"] +_ION = re.compile(r"^[A-Z][a-z]?\d+(\[.*\])?$") +HADRONS = { + "proton", + "neutron", + "pi+", + "pi-", + "pi0", + "kaon+", + "kaon-", + "kaon0L", + "kaon0S", +} + + +def categorize(name: str) -> str: + if name in ("e-", "e+", "gamma", "opticalphoton"): + return name + if ( + name in HADRONS + or name in ("deuteron", "triton", "alpha", "He3") + or _ION.match(name) + ): + # nuclear recoils and fragments ("O16", "Pb207[569.7]") count as hadrons + return "hadrons" + return "other" + + +def count_particles(events) -> dict: + names = ak.to_list(ak.flatten(events.track.particle, axis=None)) + steps = ak.to_numpy(ak.flatten(ak.num(events.track.step.energy, axis=2), axis=None)) + tracks_by = dict.fromkeys(CATEGORIES, 0) + steps_by = dict.fromkeys(CATEGORIES, 0) + for name, n in zip(names, steps): + cat = categorize(name) + tracks_by[cat] += 1 + steps_by[cat] += int(n) + return {"tracks": tracks_by, "steps": steps_by, "events": len(events)} + + +def snapshot_estimate(raw) -> float: + """Seconds to flatten and rebuild every per-thread array (EndOfRunAction proxy).""" + t0 = time.perf_counter() + for per_thread in raw: + for array in per_thread.values(): + form, length, buffers = ak.to_buffers(array) + ak.from_buffers(form, length, buffers) + return time.perf_counter() - t0 + + +def median_and_samples(samples: list[float]) -> dict: + return {"median": statistics.median(samples), "samples": samples} + + +# -- 1. in-process ------------------------------------------------------------ + + +def _inprocess( + workload: str, n_threads: int, n_events: int, repeats: int, seed: int, queue +) -> None: + try: + from geant4_python_application._geant4_application import Application + + g4.install_datasets(show_progress=False) + spec = WORKLOADS[workload] + app = Application() + app.set_seed(seed) + app.setup_manager(n_threads) + app.setup_physics("custom", spec["optical"]) + app.setup_detector(spec["gdml"]) + app.setup_action() + for command in spec["commands"]: + app.command(command) + app.initialize() + + warm = min(n_events, 2 * max(n_threads, 1)) + + def timed(fields: set[str]) -> list[float]: + app.set_event_fields(fields) + app.run(warm) + out = [] + for _ in range(repeats): + t0 = time.perf_counter() + app.run(n_events) + out.append(time.perf_counter() - t0) + return out + + result = {"minimal": timed(MINIMAL_FIELDS), "default": timed(DEFAULT_FIELDS)} + if spec["optical"]: + # (In)activating a process marks the physics modified, so the next + # BeamOn rebuilds tables on every thread; timed() warms up first. + for process in OPTICAL_PROCESSES: + app.command(f"/process/deactivate {process}") + result["no-optical"] = timed(MINIMAL_FIELDS) + app.set_event_fields({"id", "track_particle"}) + check = app.run(warm) + result["no_optical_photons_seen"] = int( + sum(ak.sum(d["track_particle"] == "opticalphoton") for d in check) + ) + queue.put(result) + except BaseException as error: # report instead of hanging the parent + queue.put({"error": repr(error)}) + raise + + +def run_inprocess( + workload: str, n_threads: int, n_events: int, repeats: int, seed: int +) -> dict: + context = multiprocessing.get_context("spawn") + queue = context.Queue() + process = context.Process( + target=_inprocess, args=(workload, n_threads, n_events, repeats, seed, queue) + ) + process.start() + result = queue.get() + process.join() + if "error" in result: + raise RuntimeError(f"in-process run failed: {result['error']}") + return { + key: median_and_samples(value) if isinstance(value, list) else value + for key, value in result.items() + } + + +# -- 2. public API ------------------------------------------------------------ + + +class TimedApplication(g4.Application): + """Records how long Application.run waits for the Geant4 child process.""" + + last_child_seconds: float = 0.0 + last_raw = None + + def _send_and_recv(self, message: Message): + if message.method != "run": + return super()._send_and_recv(message) + t0 = time.perf_counter() + result = super()._send_and_recv(message) + self.last_child_seconds = time.perf_counter() - t0 + self.last_raw = result + return result + + +def run_public( + workload: str, n_threads: int, n_events: int, repeats: int, seed: int +) -> dict: + spec = WORKLOADS[workload] + app = TimedApplication( + n_threads=n_threads, gdml=spec["gdml"], optical=spec["optical"], seed=seed + ) + totals, child_seconds, snapshots, sizes = [], [], [], [] + counts = None + stats_before = dict(_transfer.stats) + with app: + app.commands(spec["commands"]) + app.initialize() + app.set_event_fields({"id"}) + app.run(min(n_events, 2 * max(n_threads, 1))) + app.set_event_fields(DEFAULT_FIELDS) + for _ in range(repeats): + t0 = time.perf_counter() + events = app.run(n_events) + totals.append(time.perf_counter() - t0) + child_seconds.append(app.last_child_seconds) + snapshots.append(snapshot_estimate(app.last_raw)) + sizes.append(sum(a.layout.nbytes for d in app.last_raw for a in d.values())) + app.last_raw = None + if counts is None: + counts = count_particles(events) + del events + return { + "transfer": {k: _transfer.stats[k] - stats_before[k] for k in stats_before}, + "total": median_and_samples(totals), + "child": median_and_samples(child_seconds), + "snapshot_est": median_and_samples(snapshots), + "bytes": statistics.median(sizes), + "counts": counts, + } + + +# -- report ------------------------------------------------------------------- + + +def breakdown(res: dict) -> dict: + inproc, public = res["inprocess"], res["public"] + minimal = inproc["minimal"]["median"] + default = inproc["default"]["median"] + snapshot = public["snapshot_est"]["median"] + total = public["total"]["median"] + child = public["child"]["median"] + out = { + "total": total, + "transport": minimal, + "recording": default - minimal - snapshot, + "snapshot": snapshot, + "ipc": child - default, + "python": total - child, + } + if "no-optical" in inproc: + out["optical_photons"] = minimal - inproc["no-optical"]["median"] + return out + + +def pct(x: float, total: float) -> str: + return f"{x:.2f} ({100 * x / total:.0f}%)" + + +def markdown(all_results: list[dict]) -> str: + lines = [] + by_workload: dict = {} + for r in all_results: + by_workload.setdefault(r["workload"], []).append(r) + for workload, results in by_workload.items(): + lines.append(f"### {WORKLOADS[workload]['title']}\n") + bds = [breakdown(r) for r in results] + header = [ + "seconds (share of public-API wall)", + *[f"{r['threads']} thr, {r['events']} ev" for r in results], + ] + rows = [ + [ + "**wall, `Application.run`, default fields**", + *[f"**{b['total']:.2f}**" for b in bds], + ], + [ + "G4 transport (fields {id})", + *[pct(b["transport"], b["total"]) for b in bds], + ], + ] + if "optical_photons" in bds[0]: + rows.append( + [ + " of which optical photons", + *[pct(b["optical_photons"], b["total"]) for b in bds], + ] + ) + rows += [ + [ + "step/track recording (InsertStep etc.)", + *[pct(b["recording"], b["total"]) for b in bds], + ], + [ + "builder -> awkward snapshot (est.)", + *[pct(b["snapshot"], b["total"]) for b in bds], + ], + [ + "result transfer to parent process", + *[pct(b["ipc"], b["total"]) for b in bds], + ], + [ + " runs via shm / pipe (incl. warm-up)", + *[ + "{} / {}".format( + *( + r["public"].get("transfer", {}).get(k, "?") + for k in ("shared_memory", "pipe") + ) + ) + for r in results + ], + ], + [ + "Python assembly in Application.run", + *[pct(b["python"], b["total"]) for b in bds], + ], + [ + "in-process wall, fields {id} / default", + *[ + f"{r['inprocess']['minimal']['median']:.2f} / {r['inprocess']['default']['median']:.2f}" + for r in results + ], + ], + [ + "events/s: public API / in-process {id}", + *[ + f"{r['events'] / b['total']:.0f} / {r['events'] / b['transport']:.0f}" + for r, b in zip(results, bds) + ], + ], + [ + "output size (MB, awkward buffers)", + *[f"{r['public']['bytes'] / 1e6:.0f}" for r in results], + ], + ] + lines.append("| " + " | ".join(header) + " |") + lines.append("|" + "---|" * len(header)) + lines += ["| " + " | ".join(row) + " |" for row in rows] + lines.append("") + + c = results[0]["public"]["counts"] + total_steps = sum(c["steps"].values()) or 1 + total_tracks = sum(c["tracks"].values()) or 1 + lines.append(f"Tracks and steps by particle ({c['events']} events):\n") + lines.append("| particle | tracks/event | steps/event | step share |") + lines.append("|---|---|---|---|") + for cat in CATEGORIES: + t, s = c["tracks"][cat], c["steps"][cat] + if t or s: + lines.append( + f"| {cat} | {t / c['events']:.1f} | {s / c['events']:.1f} | {100 * s / total_steps:.1f}% |" + ) + em = sum(c["steps"][k] for k in ("e-", "e+", "gamma")) + lines.append( + f"| **all** | {total_tracks / c['events']:.1f} | {total_steps / c['events']:.1f} | " + f"e±/γ {100 * em / total_steps:.1f}%, optical {100 * c['steps']['opticalphoton'] / total_steps:.1f}% |" + ) + seen = results[0]["inprocess"].get("no_optical_photons_seen") + if seen is not None: + lines.append( + f"\n(no-optical check: {seen} optical photons with Scintillation and Cerenkov inactivated)" + ) + lines.append("") + return "\n".join(lines) + + +def main(argv=None) -> int: + parser = argparse.ArgumentParser( + description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter + ) + parser.add_argument( + "--workloads", nargs="+", default=list(WORKLOADS), choices=list(WORKLOADS) + ) + parser.add_argument( + "--threads", nargs="+", type=int, default=[0, 8], help="0 = serial run manager" + ) + parser.add_argument( + "--events", type=float, default=1.0, help="scale the per-workload event counts" + ) + parser.add_argument( + "--no-scale", + action="store_true", + help="do not multiply events by the thread count", + ) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument( + "--out", type=Path, default=None, help="write raw results as JSON" + ) + args = parser.parse_args(argv) + + print( + f"geant4_python_application {g4.__version__}, Geant4 {g4.geant4_version}, awkward {ak.__version__}", + file=sys.stderr, + ) + all_results = [] + for workload in args.workloads: + for n_threads in args.threads: + n = int(round(WORKLOADS[workload]["events"] * args.events)) + if not args.no_scale and n_threads > 1: + n *= n_threads + print( + f"[{workload}] threads={n_threads} events={n} ...", + file=sys.stderr, + flush=True, + ) + t0 = time.perf_counter() + res = { + "workload": workload, + "threads": max(n_threads, 1), + "events": n, + "inprocess": run_inprocess( + workload, n_threads, n, args.repeats, args.seed + ), + "public": run_public(workload, n_threads, n, args.repeats, args.seed), + } + print( + f" done in {time.perf_counter() - t0:.1f}s: {json.dumps(breakdown(res))}", + file=sys.stderr, + flush=True, + ) + all_results.append(res) + if args.out is not None: + args.out.parent.mkdir(parents=True, exist_ok=True) + args.out.write_text(json.dumps(all_results, indent=1)) + print(markdown(all_results)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main())