From e26db97b5cc6e29a1e5ca5082bc6e7c0449df71c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Erik=20Fr=C3=B6jdh?= Date: Tue, 11 Aug 2026 15:28:54 +0200 Subject: [PATCH] Prototype multi threaded file reader (#343) - Reading files from nfs shares tops out at ~1GB/s for single threaded reads. - MultiThreadedFileReader uses our File wrapper to read a generic file in parallel - Placed in aare::experimental to show that it's not production ready --- .gitignore | 4 +- CMakeLists.txt | 7 +- docs/src/MultiThreadedFileReader.rst | 37 ++++ docs/src/index.rst | 2 + .../experimental/MultiThreadedFileReader.rst | 30 +++ docs/src/python/experimental/index.rst | 10 + include/aare/File.hpp | 3 +- include/aare/FileInterface.hpp | 13 +- include/aare/JungfrauDataFile.hpp | 3 +- include/aare/MultiThreadedFileReader.hpp | 104 ++++++++++ include/aare/NumpyFile.hpp | 4 +- include/aare/RawFile.hpp | 3 +- include/aare/utils/math_helpers.hpp | 14 ++ python/CMakeLists.txt | 1 + python/aare/__init__.py | 8 +- python/aare/experimental.py | 6 + python/src/bind_MultiThreadedFileReader.hpp | 175 ++++++++++++++++ python/src/module.cpp | 5 + python/tests/test_MultiThreadedFileReader.py | 147 ++++++++++++++ python/tests/test_PixelHistogram.py | 101 ++++++++++ src/Dtype.cpp | 6 + src/Dtype.test.cpp | 5 +- src/File.cpp | 3 +- src/MultiThreadedFileReader.cpp | 151 ++++++++++++++ src/MultiThreadedFileReader.test.cpp | 189 ++++++++++++++++++ src/NumpyHelpers.test.cpp | 5 +- 26 files changed, 1021 insertions(+), 15 deletions(-) create mode 100644 docs/src/MultiThreadedFileReader.rst create mode 100644 docs/src/python/experimental/MultiThreadedFileReader.rst create mode 100644 docs/src/python/experimental/index.rst create mode 100644 include/aare/MultiThreadedFileReader.hpp create mode 100644 include/aare/utils/math_helpers.hpp create mode 100644 python/aare/experimental.py create mode 100644 python/src/bind_MultiThreadedFileReader.hpp create mode 100644 python/tests/test_MultiThreadedFileReader.py create mode 100644 python/tests/test_PixelHistogram.py create mode 100644 src/MultiThreadedFileReader.cpp create mode 100644 src/MultiThreadedFileReader.test.cpp diff --git a/.gitignore b/.gitignore index fe88c6fe..dc40ceac 100644 --- a/.gitignore +++ b/.gitignore @@ -127,6 +127,8 @@ cover/ *.mo *.pot +.* + # Django stuff: *.log local_settings.py @@ -192,4 +194,4 @@ cython_debug/ .ruff_cache/ # user defined -wheelhouse/ \ No newline at end of file +wheelhouse/ diff --git a/CMakeLists.txt b/CMakeLists.txt index f923dfce..a3d96cc4 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -396,6 +396,7 @@ set(PUBLICHEADERS include/aare/FileInterface.hpp include/aare/FilePtr.hpp include/aare/Frame.hpp + include/aare/MultiThreadedFileReader.hpp include/aare/hist/PixelHistogram.hpp include/aare/hist/PixelHistogramImpl.hpp include/aare/hist/PedestalTrackingPixelHistogram.hpp @@ -414,7 +415,9 @@ set(PUBLICHEADERS include/aare/RawMasterFile.hpp include/aare/RawSubFile.hpp include/aare/VarClusterFinder.hpp - include/aare/utils/task.hpp) + include/aare/utils/task.hpp + include/aare/utils/ifstream_helpers.hpp + include/aare/utils/math_helpers.hpp) set(SourceFiles ${CMAKE_CURRENT_SOURCE_DIR}/src/calibration.cpp @@ -431,6 +434,7 @@ set(SourceFiles ${CMAKE_CURRENT_SOURCE_DIR}/src/Frame.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/Interpolator.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/JungfrauDataFile.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/MultiThreadedFileReader.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/NumpyFile.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/NumpyHelpers.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/PixelMap.cpp @@ -502,6 +506,7 @@ if(AARE_TESTS) ${CMAKE_CURRENT_SOURCE_DIR}/src/hist/PixelHistogramImpl.test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/hist/PixelHistogram.test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/JungfrauDataFile.test.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/src/MultiThreadedFileReader.test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/NumpyFile.test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/NumpyHelpers.test.cpp ${CMAKE_CURRENT_SOURCE_DIR}/src/RawFile.test.cpp diff --git a/docs/src/MultiThreadedFileReader.rst b/docs/src/MultiThreadedFileReader.rst new file mode 100644 index 00000000..a2bdebf8 --- /dev/null +++ b/docs/src/MultiThreadedFileReader.rst @@ -0,0 +1,37 @@ +MultiThreadedFileReader +======================= + +``MultiThreadedFileReader`` reads one chunk per worker on each call. Each +worker owns an independent :cpp:class:`aare::File`, while all workers write +into non-overlapping regions of the destination buffer. The output is ordered +by frame index and successive calls advance through the file. + +.. code-block:: cpp + + #include "aare/MultiThreadedFileReader.hpp" + + // Four workers, chunks of 128 frames, and at most 10,000 frames. + aare::experimental::MultiThreadedFileReader reader(path, 4, 128, 10'000); + while (reader.remaining_frames() != 0) { + // Contains at most 4 * 128 frames. + auto batch = reader.read(); + process(batch); + } + +Omit the final argument to read every frame in the source. An explicit value +of zero requests an empty result. The low-level ``read_into`` overload avoids +an allocation when the caller already owns a buffer of at least +``reader.next_read_bytes()`` bytes. Use ``read_all()`` to read every frame +remaining from the current position, and ``seek()`` to reposition the reader. +Call ``close()`` to release all worker file handles early. + +.. note:: + + Multiple workers do not guarantee faster reads. Performance depends on the + storage device and file format, so the thread count and chunk size should be + benchmark-driven. + +.. doxygenclass:: aare::experimental::MultiThreadedFileReader + :members: + :undoc-members: + :private-members: diff --git a/docs/src/index.rst b/docs/src/index.rst index 8b9af286..d60fc78e 100644 --- a/docs/src/index.rst +++ b/docs/src/index.rst @@ -28,6 +28,7 @@ AARE pycalibration python/cluster/index python/file/index + python/experimental/index python/histogram/index pyFit @@ -42,6 +43,7 @@ AARE NDView Frame File + MultiThreadedFileReader Dtype Cluster ClusterFinder diff --git a/docs/src/python/experimental/MultiThreadedFileReader.rst b/docs/src/python/experimental/MultiThreadedFileReader.rst new file mode 100644 index 00000000..d8553e4f --- /dev/null +++ b/docs/src/python/experimental/MultiThreadedFileReader.rst @@ -0,0 +1,30 @@ +MultiThreadedFileReader +======================= + +.. py:currentmodule:: aare.experimental + +The reader returns a NumPy array with shape ``(frames, rows, cols)`` and +preserves the source pixel dtype. Each iteration reads at most +``n_threads * chunk_size`` frames—one chunk per worker. File I/O runs with the +Python GIL released. + +.. code-block:: python + + from aare.experimental import MultiThreadedFileReader + + with MultiThreadedFileReader( + "frames.npy", n_threads=4, chunk_size=128, total_frames=10_000 + ) as reader: + for frames in reader: + process(frames) + +Call ``read()`` directly for the next batch, or ``read_all()`` for all frames +remaining from the current position. ``tell()`` and ``seek()`` expose the +iteration position. The context manager closes all worker files on exit. +``close()`` is also available for explicit cleanup and may be called +repeatedly. + +.. autoclass:: MultiThreadedFileReader + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/src/python/experimental/index.rst b/docs/src/python/experimental/index.rst new file mode 100644 index 00000000..fc683bf5 --- /dev/null +++ b/docs/src/python/experimental/index.rst @@ -0,0 +1,10 @@ +Experimental +============ + +APIs in this module may change without notice. + +.. toctree:: + :caption: Experimental + :maxdepth: 1 + + MultiThreadedFileReader diff --git a/include/aare/File.hpp b/include/aare/File.hpp index 6e65fa26..a931f511 100644 --- a/include/aare/File.hpp +++ b/include/aare/File.hpp @@ -64,8 +64,9 @@ class File { size_t total_frames() const; size_t rows() const; size_t cols() const; + Dtype dtype() const; DetectorType detector_type() const; }; -} // namespace aare \ No newline at end of file +} // namespace aare diff --git a/include/aare/FileInterface.hpp b/include/aare/FileInterface.hpp index 65ac0c2e..87ca3dbf 100644 --- a/include/aare/FileInterface.hpp +++ b/include/aare/FileInterface.hpp @@ -145,10 +145,15 @@ class FileInterface { */ virtual size_t bitdepth() const = 0; - virtual DetectorType detector_type() const = 0; + /** + * @brief get the data type of the pixels + * @return pixel data type + */ + virtual Dtype dtype() const { + return Dtype::from_bitdepth(static_cast(bitdepth())); + } - // function to query the data type of the file - /*virtual DataType dtype = 0; */ + virtual DetectorType detector_type() const = 0; virtual ~FileInterface() = default; @@ -168,4 +173,4 @@ class FileInterface { // size_t current_frame{}; }; -} // namespace aare \ No newline at end of file +} // namespace aare diff --git a/include/aare/JungfrauDataFile.hpp b/include/aare/JungfrauDataFile.hpp index 0911c284..fe74d5e3 100644 --- a/include/aare/JungfrauDataFile.hpp +++ b/include/aare/JungfrauDataFile.hpp @@ -49,6 +49,7 @@ class JungfrauDataFile : public FileInterface { size_t pixels_per_frame() override; size_t bytes_per_pixel() const; size_t bitdepth() const override; + Dtype dtype() const override { return Dtype::UINT16; } void seek(size_t frame_index) override; //!< seek to the given frame index (note not byte offset) size_t tell() override; //!< get the frame index of the file pointer @@ -113,4 +114,4 @@ class JungfrauDataFile : public FileInterface { std::filesystem::path fpath(size_t frame_index) const; }; -} // namespace aare \ No newline at end of file +} // namespace aare diff --git a/include/aare/MultiThreadedFileReader.hpp b/include/aare/MultiThreadedFileReader.hpp new file mode 100644 index 00000000..d7a930b8 --- /dev/null +++ b/include/aare/MultiThreadedFileReader.hpp @@ -0,0 +1,104 @@ +// SPDX-License-Identifier: MPL-2.0 +#pragma once + +#include "aare/Dtype.hpp" +#include "aare/File.hpp" + +#include +#include +#include +#include + +namespace aare::experimental { + +/** + * @brief Read independent chunks of a file in parallel. + * + * Each worker opens its own File instance, so seeking and reading do not share + * mutable file state. Chunks are written directly to their position in the + * destination buffer and the resulting frame order is the same as in the file. + */ +class MultiThreadedFileReader { + public: + /** + * @param fname path accepted by File + * @param n_threads maximum number of worker threads + * @param chunk_size number of frames claimed by a worker at a time + * @param total_frames number of frames to read, or all frames when omitted + */ + MultiThreadedFileReader(std::filesystem::path fname, size_t n_threads, + size_t chunk_size, + std::optional total_frames = std::nullopt); + + MultiThreadedFileReader(const MultiThreadedFileReader &) = delete; + MultiThreadedFileReader & + operator=(const MultiThreadedFileReader &) = delete; + MultiThreadedFileReader(MultiThreadedFileReader &&) noexcept = default; + MultiThreadedFileReader & + operator=(MultiThreadedFileReader &&) noexcept = default; + + /** + * @brief Read one chunk per active worker into a caller-owned buffer. + * + * The buffer must hold at least next_read_bytes() bytes. The reader's + * position advances by the returned number of frames. At the end of the + * configured range this function returns zero and does not access the + * destination. + */ + size_t read_into(std::byte *destination); + + /** @brief Read the next wave of chunks into an owned byte buffer. */ + std::vector read(); + + /** @brief Read every frame remaining from the current position. */ + std::vector read_all(); + + /** @brief Set the next frame index to read. The end position is valid. */ + void seek(size_t frame_index); + + /** @brief Return the next frame index to read. */ + size_t tell() const noexcept { return m_current_frame; } + + /** @brief Close all worker files. Safe to call more than once. */ + void close() noexcept { m_files.clear(); } + + /** @brief Return whether the worker files are open. */ + bool is_open() const noexcept { return !m_files.empty(); } + + size_t n_threads() const noexcept { return m_n_threads; } + size_t chunk_size() const noexcept { return m_chunk_size; } + size_t total_frames() const noexcept { return m_total_frames; } + size_t source_total_frames() const noexcept { + return m_source_total_frames; + } + size_t rows() const noexcept { return m_rows; } + size_t cols() const noexcept { return m_cols; } + size_t bitdepth() const noexcept { return m_bitdepth; } + Dtype dtype() const noexcept { return m_dtype; } + size_t bytes_per_frame() const noexcept { return m_bytes_per_frame; } + size_t total_bytes() const noexcept { return m_total_bytes; } + size_t remaining_frames() const noexcept; + size_t next_read_frames() const noexcept; + size_t next_read_bytes() const noexcept { + return next_read_frames() * m_bytes_per_frame; + } + + private: + std::filesystem::path m_fname; + size_t m_n_threads; + size_t m_chunk_size; + size_t m_total_frames; + size_t m_source_total_frames; + size_t m_rows; + size_t m_cols; + size_t m_bitdepth; + Dtype m_dtype; + size_t m_bytes_per_frame; + size_t m_total_bytes; + size_t m_current_frame; + std::vector m_files; + + void ensure_open() const; +}; + +} // namespace aare::experimental diff --git a/include/aare/NumpyFile.hpp b/include/aare/NumpyFile.hpp index 76d5321c..caa68531 100644 --- a/include/aare/NumpyFile.hpp +++ b/include/aare/NumpyFile.hpp @@ -61,7 +61,7 @@ class NumpyFile : public FileInterface { * @brief get the data type of the numpy file * @return DType */ - Dtype dtype() const { return m_header.dtype; } + Dtype dtype() const override { return m_header.dtype; } /** * @brief get the shape of the numpy file @@ -129,4 +129,4 @@ class NumpyFile : public FileInterface { void write_impl(void *data, uint64_t size); }; -} // namespace aare \ No newline at end of file +} // namespace aare diff --git a/include/aare/RawFile.hpp b/include/aare/RawFile.hpp index 65758017..999b737f 100644 --- a/include/aare/RawFile.hpp +++ b/include/aare/RawFile.hpp @@ -118,6 +118,7 @@ class RawFile : public FileInterface { */ size_t cols(const size_t roi_index) const; size_t bitdepth() const override; + Dtype dtype() const override { return Dtype::from_bitdepth(bitdepth()); } size_t n_modules() const; /** @@ -170,4 +171,4 @@ class RawFile : public FileInterface { void open_subfiles(const size_t roi_index); }; -} // namespace aare \ No newline at end of file +} // namespace aare diff --git a/include/aare/utils/math_helpers.hpp b/include/aare/utils/math_helpers.hpp new file mode 100644 index 00000000..9481bd8a --- /dev/null +++ b/include/aare/utils/math_helpers.hpp @@ -0,0 +1,14 @@ +#pragma once +#include + +namespace aare { + +/** + * @brief Compute the ceiling of the integer division of n by d. + * @param n The numerator. + * @param d The denominator. + * @return The ceiling of the integer division. + */ +constexpr size_t ceil_div(size_t n, size_t d) { return n / d + (n % d != 0); } + +} // namespace aare \ No newline at end of file diff --git a/python/CMakeLists.txt b/python/CMakeLists.txt index 03775481..5f877773 100644 --- a/python/CMakeLists.txt +++ b/python/CMakeLists.txt @@ -40,6 +40,7 @@ set(PYTHON_FILES aare/ClusterVector.py aare/Cluster.py aare/calibration.py + aare/experimental.py aare/func.py aare/RawFile.py aare/transform.py diff --git a/python/aare/__init__.py b/python/aare/__init__.py index 10872f89..c8db2e4e 100644 --- a/python/aare/__init__.py +++ b/python/aare/__init__.py @@ -3,8 +3,14 @@ from . import _aare from . import transform +from . import experimental -from ._aare import File, RawMasterFile, RawSubFile, JungfrauDataFile +from ._aare import ( + File, + JungfrauDataFile, + RawMasterFile, + RawSubFile, +) from ._aare import Pedestal_d, Pedestal_f, ClusterFinder_Cluster3x3i, VarClusterFinder from ._aare import DetectorType, ReadoutMode from ._aare import hitmap diff --git a/python/aare/experimental.py b/python/aare/experimental.py new file mode 100644 index 00000000..2a5210d9 --- /dev/null +++ b/python/aare/experimental.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: MPL-2.0 +"""Experimental APIs that may change without notice.""" + +from ._aare.experimental import MultiThreadedFileReader + +__all__ = ["MultiThreadedFileReader"] diff --git a/python/src/bind_MultiThreadedFileReader.hpp b/python/src/bind_MultiThreadedFileReader.hpp new file mode 100644 index 00000000..f63f79c8 --- /dev/null +++ b/python/src/bind_MultiThreadedFileReader.hpp @@ -0,0 +1,175 @@ +// SPDX-License-Identifier: MPL-2.0 +#pragma once + +#include "aare/MultiThreadedFileReader.hpp" + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +namespace py = pybind11; + +inline py::dtype multi_threaded_reader_numpy_dtype(const aare::Dtype &dtype) { + using aare::Dtype; + if (dtype == Dtype::INT8) + return py::dtype::of(); + if (dtype == Dtype::UINT8) + return py::dtype::of(); + if (dtype == Dtype::INT16) + return py::dtype::of(); + if (dtype == Dtype::UINT16) + return py::dtype::of(); + if (dtype == Dtype::INT32) + return py::dtype::of(); + if (dtype == Dtype::UINT32) + return py::dtype::of(); + if (dtype == Dtype::INT64) + return py::dtype::of(); + if (dtype == Dtype::UINT64) + return py::dtype::of(); + if (dtype == Dtype::FLOAT) + return py::dtype::of(); + if (dtype == Dtype::DOUBLE) + return py::dtype::of(); + throw std::runtime_error("Unsupported pixel data type"); +} + +inline py::array +multi_threaded_reader_read(aare::experimental::MultiThreadedFileReader &reader, + bool read_all) { + const size_t n_frames = + read_all ? reader.remaining_frames() : reader.next_read_frames(); + const std::vector shape{ + static_cast(n_frames), + static_cast(reader.rows()), + static_cast(reader.cols())}; + + py::array image(multi_threaded_reader_numpy_dtype(reader.dtype()), shape); + auto *destination = reinterpret_cast(image.mutable_data()); + { + py::gil_scoped_release release; + if (read_all) { + size_t offset = 0; + while (reader.remaining_frames() != 0) { + const size_t frames_read = + reader.read_into(destination + offset); + offset += frames_read * reader.bytes_per_frame(); + } + } else { + reader.read_into(destination); + } + } + return image; +} + +inline void define_multi_threaded_file_reader_bindings(py::module_ &m) { + using aare::experimental::MultiThreadedFileReader; + + auto reader = + py::class_(m, "MultiThreadedFileReader"); + reader.attr("__module__") = "aare.experimental"; + reader + .def(py::init>(), + py::arg("fname"), py::arg("n_threads"), py::arg("chunk_size"), + py::arg("total_frames") = py::none(), + R"doc( + Read chunks of detector frames concurrently. + + Each worker opens an independent File. The returned array is + ordered by frame index even though chunks are read in parallel. + + Args: + fname: Path accepted by File. + n_threads: Maximum number of worker threads. + chunk_size: Number of frames read per claimed chunk. + total_frames: Optional frame limit. None reads all frames. + )doc") + .def( + "read", + [](MultiThreadedFileReader &self) { + return multi_threaded_reader_read(self, false); + }, + R"doc( + Read one chunk per active worker into a NumPy array. + + Returns: + An array containing at most n_threads * chunk_size frames. + An empty array is returned at the end of the configured + frame range. The GIL is released while file data is read. + )doc") + .def( + "read_all", + [](MultiThreadedFileReader &self) { + return multi_threaded_reader_read(self, true); + }, + R"doc( + Read all frames remaining from the current position. + )doc") + .def_property_readonly("n_threads", &MultiThreadedFileReader::n_threads) + .def_property_readonly("chunk_size", + &MultiThreadedFileReader::chunk_size) + .def_property_readonly("total_frames", + &MultiThreadedFileReader::total_frames) + .def_property_readonly("source_total_frames", + &MultiThreadedFileReader::source_total_frames) + .def_property_readonly("rows", &MultiThreadedFileReader::rows) + .def_property_readonly("cols", &MultiThreadedFileReader::cols) + .def_property_readonly("bitdepth", &MultiThreadedFileReader::bitdepth) + .def_property_readonly("dtype", + [](const MultiThreadedFileReader &self) { + return multi_threaded_reader_numpy_dtype( + self.dtype()); + }) + .def_property_readonly("bytes_per_frame", + &MultiThreadedFileReader::bytes_per_frame) + .def_property_readonly("total_bytes", + &MultiThreadedFileReader::total_bytes) + .def_property_readonly("remaining_frames", + &MultiThreadedFileReader::remaining_frames) + .def_property_readonly("next_read_frames", + &MultiThreadedFileReader::next_read_frames) + .def_property_readonly("next_read_bytes", + &MultiThreadedFileReader::next_read_bytes) + .def("seek", &MultiThreadedFileReader::seek, py::arg("frame_index")) + .def("tell", &MultiThreadedFileReader::tell) + .def("close", &MultiThreadedFileReader::close, + "Close all worker files. Safe to call more than once.") + .def_property_readonly( + "closed", + [](const MultiThreadedFileReader &self) { return !self.is_open(); }) + .def("__len__", &MultiThreadedFileReader::total_frames) + .def( + "__enter__", + [](MultiThreadedFileReader &self) -> MultiThreadedFileReader * { + if (!self.is_open()) { + throw std::runtime_error( + "Cannot enter a closed MultiThreadedFileReader"); + } + return &self; + }, + py::return_value_policy::reference_internal) + .def("__exit__", + [](MultiThreadedFileReader &self, const py::object &, + const py::object &, const py::object &) { + self.close(); + return false; + }) + .def( + "__iter__", [](MultiThreadedFileReader &self) { return &self; }, + py::return_value_policy::reference_internal) + .def("__next__", [](MultiThreadedFileReader &self) { + if (self.remaining_frames() == 0) { + throw py::stop_iteration(); + } + return multi_threaded_reader_read(self, false); + }); +} diff --git a/python/src/module.cpp b/python/src/module.cpp index 023a19ee..52c830b9 100644 --- a/python/src/module.cpp +++ b/python/src/module.cpp @@ -12,6 +12,7 @@ #include "bind_Defs.hpp" #include "bind_Eta.hpp" #include "bind_Interpolator.hpp" +#include "bind_MultiThreadedFileReader.hpp" #include "bind_PedestalTrackingPixelHistogram.hpp" #include "bind_PixelHistogram.hpp" #include "bind_PixelMap.hpp" @@ -59,7 +60,11 @@ double, 'f' for float) define_ClusterCollector(m, "Cluster" #N "x" #M #TYPE_CODE); PYBIND11_MODULE(_aare, m) { + auto experimental = m.def_submodule( + "experimental", "Experimental APIs that may change without notice"); + define_file_io_bindings(m); + define_multi_threaded_file_reader_bindings(experimental); define_raw_file_io_bindings(m); define_raw_sub_file_io_bindings(m); define_ctb_raw_file_io_bindings(m); diff --git a/python/tests/test_MultiThreadedFileReader.py b/python/tests/test_MultiThreadedFileReader.py new file mode 100644 index 00000000..0dc9c654 --- /dev/null +++ b/python/tests/test_MultiThreadedFileReader.py @@ -0,0 +1,147 @@ +# SPDX-License-Identifier: MPL-2.0 +import numpy as np +import pytest + +import aare +from aare.experimental import MultiThreadedFileReader + + +@pytest.fixture +def frame_file(tmp_path): + data = np.arange(10 * 2 * 3, dtype=np.uint16).reshape(10, 2, 3) + path = tmp_path / "frames.npy" + np.save(path, data) + return path, data + + +def test_experimental_import_path(): + assert aare.experimental.MultiThreadedFileReader is MultiThreadedFileReader + assert MultiThreadedFileReader.__module__ == "aare.experimental" + assert not hasattr(aare, "MultiThreadedFileReader") + + +def test_reads_all_frames_in_order(frame_file): + path, expected = frame_file + reader = MultiThreadedFileReader(path, n_threads=2, chunk_size=3) + + first = reader.read() + second = reader.read() + exhausted = reader.read() + + assert np.array_equal(first, expected[:6]) + assert np.array_equal(second, expected[6:]) + assert exhausted.shape == (0, 2, 3) + assert first.dtype == np.uint16 + assert reader.n_threads == 2 + assert reader.chunk_size == 3 + assert reader.total_frames == 10 + assert reader.source_total_frames == 10 + assert reader.rows == 2 + assert reader.cols == 3 + assert reader.bitdepth == 16 + assert reader.dtype == np.dtype(np.uint16) + assert reader.bytes_per_frame == 12 + assert reader.total_bytes == expected.nbytes + assert len(reader) == 10 + assert reader.tell() == 10 + assert reader.remaining_frames == 0 + assert reader.next_read_frames == 0 + assert reader.next_read_bytes == 0 + + +def test_total_frame_limit(frame_file): + path, expected = frame_file + reader = MultiThreadedFileReader( + path, n_threads=2, chunk_size=2, total_frames=7 + ) + + assert np.array_equal(reader.read(), expected[:4]) + assert np.array_equal(reader.read_all(), expected[4:7]) + + +def test_iteration_yields_one_chunk_per_thread(frame_file): + path, expected = frame_file + reader = MultiThreadedFileReader(path, n_threads=2, chunk_size=2) + + batches = list(reader) + + assert [len(batch) for batch in batches] == [4, 4, 2] + assert np.array_equal(np.concatenate(batches), expected) + + +def test_seek_resets_iteration_position(frame_file): + path, expected = frame_file + reader = MultiThreadedFileReader(path, n_threads=2, chunk_size=2) + + reader.read() + assert reader.tell() == 4 + + reader.seek(1) + assert reader.tell() == 1 + assert np.array_equal(reader.read(), expected[1:5]) + + with pytest.raises(IndexError): + reader.seek(11) + + +def test_context_manager_closes_worker_files(frame_file): + path, expected = frame_file + + with MultiThreadedFileReader(path, n_threads=2, chunk_size=2) as reader: + assert not reader.closed + assert np.array_equal(reader.read(), expected[:4]) + + assert reader.closed + reader.close() + with pytest.raises(RuntimeError): + reader.read() + with pytest.raises(RuntimeError): + reader.seek(0) + with pytest.raises(RuntimeError): + with reader: + pass + + +def test_explicit_zero_frame_limit(frame_file): + path, expected = frame_file + reader = MultiThreadedFileReader( + path, n_threads=8, chunk_size=3, total_frames=0 + ) + + actual = reader.read() + assert actual.shape == (0, *expected.shape[1:]) + assert actual.dtype == expected.dtype + + +@pytest.mark.parametrize( + "dtype", [np.int8, np.int32, np.uint64, np.float32, np.float64] +) +def test_preserves_numpy_dtype(tmp_path, dtype): + expected = np.arange(4 * 2 * 3, dtype=dtype).reshape(4, 2, 3) + path = tmp_path / "typed-frames.npy" + np.save(path, expected) + + reader = MultiThreadedFileReader(path, n_threads=2, chunk_size=3) + actual = reader.read() + + assert actual.dtype == expected.dtype + assert reader.dtype == expected.dtype + assert np.array_equal(actual, expected) + + +@pytest.mark.parametrize( + ("n_threads", "chunk_size", "total_frames"), + [(0, 1, None), (1, 0, None), (1, 1, 11)], +) +def test_invalid_configuration( + frame_file, n_threads, chunk_size, total_frames +): + path, _ = frame_file + + with pytest.raises(ValueError): + MultiThreadedFileReader( + path, + n_threads=n_threads, + chunk_size=chunk_size, + total_frames=total_frames, + ) diff --git a/python/tests/test_PixelHistogram.py b/python/tests/test_PixelHistogram.py new file mode 100644 index 00000000..d8563328 --- /dev/null +++ b/python/tests/test_PixelHistogram.py @@ -0,0 +1,101 @@ +# SPDX-License-Identifier: MPL-2.0 +import numpy as np +import pytest + +from aare import PixelHistogram + + +def _random_frames(rows, cols, n, xmin, xmax, seed=0): + rng = np.random.default_rng(seed) + return [rng.uniform(xmin - 0.25, xmax + 0.25, size=(rows, cols)).astype(np.float64) + for _ in range(n)] + + +def _reference_hdata(frames, rows, cols, n_bins, xmin, xmax): + expected = np.zeros((rows, cols, n_bins), dtype=np.uint16) + inv_range = n_bins / (xmax - xmin) + for img in frames: + for r in range(rows): + for c in range(cols): + v = float(img[r, c]) + if not (xmin <= v < xmax): + continue + b = int((v - xmin) * inv_range) + if b >= n_bins: + b = n_bins - 1 + expected[r, c, b] += 1 + return expected + + +def test_async_fill_matches_reference(): + rows, cols, n_bins = 5, 7, 8 + xmin, xmax = 0.0, 2.0 + frames = _random_frames(rows, cols, n=3, xmin=xmin, xmax=xmax, seed=1) + + hist = PixelHistogram(rows=rows, cols=cols, n_bins=n_bins, xmin=xmin, xmax=xmax) + for img in frames: + hist.fill_async(img) + + np.testing.assert_array_equal( + hist.values(), + _reference_hdata(frames, rows, cols, n_bins, xmin, xmax), + ) + + + + +def test_fill_async_copies_buffer(): + # After fill_async returns, the caller should be free to mutate the + # numpy array without affecting the pending fill. + rows, cols, n_bins = 4, 4, 4 + xmin, xmax = 0.0, 1.0 + hist = PixelHistogram(rows=rows, cols=cols, n_bins=n_bins, xmin=xmin, xmax=xmax, n_threads=1, max_pending=8) + + img = np.full((rows, cols), 0.1, dtype=np.float64) # falls in bin 0 + hist.fill_async(img) + # Mutate the original array immediately; this must not affect the + # value that was already enqueued. + img[:] = 0.9 # would be bin 3 + hist.flush() + + h = hist.values() + assert h.shape == (rows, cols, n_bins) + # Every pixel saw one value in bin 0, none elsewhere. + assert (h[:, :, 0] == 1).all() + assert (h[:, :, 1:] == 0).all() + + + def test_fill_async_rejects_wrong_shape(): + hist = PixelHistogram(8, 8, 4, 0.0, 1.0) + bad = np.zeros((4, 4), dtype=np.float32) + with pytest.raises(ValueError): + hist.fill_async(bad) + + +def test_hdata_flushes_pending(): + # Submit several frames with a tiny queue and read hdata() without an + # explicit flush(); hdata() must drain everything first. + rows, cols, n_bins = 3, 3, 4 + xmin, xmax = 0.0, 1.0 + hist = PixelHistogram(rows=rows, cols=cols, n_bins=n_bins, xmin=xmin, xmax=xmax, + n_threads=1, max_pending=1) + frames = _random_frames(rows, cols, n=8, xmin=xmin, xmax=xmax, seed=3) + for img in frames: + hist.fill_async(img) + + h = hist.values() # no explicit flush() + np.testing.assert_array_equal( + h, _reference_hdata(frames, rows, cols, n_bins, xmin, xmax) + ) + + +def test_bin_centers_and_edges(): + n_bins = 5 + xmin, xmax = 0.0, 1.0 + hist = PixelHistogram(rows=2, cols=2, n_bins=n_bins, xmin=xmin, xmax=xmax) + edges = hist.bin_edges() + centers = hist.bin_centers() + assert edges.shape == (n_bins + 1,) + assert centers.shape == (n_bins,) + np.testing.assert_allclose(edges, np.linspace(xmin, xmax, n_bins + 1), atol=1e-6) + np.testing.assert_allclose(centers, 0.5 * (edges[:-1] + edges[1:]), atol=1e-6) diff --git a/src/Dtype.cpp b/src/Dtype.cpp index ad708ad9..08098874 100644 --- a/src/Dtype.cpp +++ b/src/Dtype.cpp @@ -92,6 +92,12 @@ Dtype::Dtype(Dtype::TypeIndex ti) : m_type(ti) {} */ Dtype::Dtype(std::string_view sv) { + // NumPy uses '|' for data types whose byte order is not applicable, + // notably one-byte integer types. + if (!sv.empty() && sv.front() == '|') { + sv.remove_prefix(1); + } + // Check if the file is using our native endianess if (auto pos = sv.find_first_of("<>"); pos != std::string_view::npos) { const auto endianess = [](const char c) { diff --git a/src/Dtype.test.cpp b/src/Dtype.test.cpp index 760a6c83..f87825e8 100644 --- a/src/Dtype.test.cpp +++ b/src/Dtype.test.cpp @@ -12,6 +12,9 @@ TEST_CASE("Construct from typeid") { } TEST_CASE("Construct from string") { + REQUIRE(Dtype("|i1") == typeid(int8_t)); + REQUIRE(Dtype("|u1") == typeid(uint8_t)); + if (endian::native == endian::little) { REQUIRE(Dtype("tell(); } size_t File::rows() const { return file_impl->rows(); } size_t File::cols() const { return file_impl->cols(); } size_t File::bitdepth() const { return file_impl->bitdepth(); } +Dtype File::dtype() const { return file_impl->dtype(); } size_t File::bytes_per_pixel() const { return file_impl->bitdepth() / bits_per_byte; } DetectorType File::detector_type() const { return file_impl->detector_type(); } -} // namespace aare \ No newline at end of file +} // namespace aare diff --git a/src/MultiThreadedFileReader.cpp b/src/MultiThreadedFileReader.cpp new file mode 100644 index 00000000..6810e9bf --- /dev/null +++ b/src/MultiThreadedFileReader.cpp @@ -0,0 +1,151 @@ +// SPDX-License-Identifier: MPL-2.0 +#include "aare/MultiThreadedFileReader.hpp" +#include "aare/File.hpp" +#include "aare/utils/math_helpers.hpp" + +#include +#include +#include +#include +#include + +namespace aare::experimental { +namespace { + +size_t checked_product(size_t lhs, size_t rhs) { + if (lhs != 0 && rhs > std::numeric_limits::max() / lhs) { + throw std::overflow_error( + "MultiThreadedFileReader buffer size overflow"); + } + return lhs * rhs; +} + +} // namespace + +MultiThreadedFileReader::MultiThreadedFileReader( + std::filesystem::path fname, size_t n_threads, size_t chunk_size, + std::optional total_frames) + : m_fname(std::move(fname)), m_n_threads(n_threads), + m_chunk_size(chunk_size), m_total_frames(0), m_source_total_frames(0), + m_rows(0), m_cols(0), m_bitdepth(0), m_dtype(Dtype::NONE), + m_bytes_per_frame(0), m_total_bytes(0), m_current_frame(0) { + if (m_n_threads == 0) { + throw std::invalid_argument( + "MultiThreadedFileReader requires at least one thread"); + } + if (m_chunk_size == 0) { + throw std::invalid_argument( + "MultiThreadedFileReader chunk size must be greater than zero"); + } + + File file(m_fname); + m_source_total_frames = file.total_frames(); + m_total_frames = total_frames.value_or(m_source_total_frames); + if (m_total_frames > m_source_total_frames) { + throw std::invalid_argument( + "Requested frame count exceeds the number of frames in the file"); + } + + m_rows = file.rows(); + m_cols = file.cols(); + m_bitdepth = file.bitdepth(); + m_dtype = file.dtype(); + m_bytes_per_frame = file.bytes_per_frame(); + m_total_bytes = checked_product(m_total_frames, m_bytes_per_frame); + + m_files.reserve(m_n_threads); + m_files.push_back(std::move(file)); + for (size_t i = 1; i < m_n_threads; ++i) { + m_files.emplace_back(m_fname); + } +} + +size_t MultiThreadedFileReader::remaining_frames() const noexcept { + return m_total_frames - m_current_frame; +} + +size_t MultiThreadedFileReader::next_read_frames() const noexcept { + const size_t remaining = remaining_frames(); + if (remaining == 0) { + return 0; + } + + const size_t chunks_remaining = ceil_div(remaining, m_chunk_size); + if (m_n_threads >= chunks_remaining) { + return remaining; + } + // This multiplication is safe: in this branch n_threads * chunk_size is + // strictly smaller than remaining. + return m_n_threads * m_chunk_size; +} + +void MultiThreadedFileReader::ensure_open() const { + if (!is_open()) { + throw std::runtime_error("MultiThreadedFileReader is closed"); + } +} + +size_t MultiThreadedFileReader::read_into(std::byte *destination) { + ensure_open(); + const size_t frames_to_read = next_read_frames(); + if (frames_to_read == 0) { + return 0; + } + if (destination == nullptr) { + throw std::invalid_argument( + "MultiThreadedFileReader destination must not be null"); + } + + const size_t first_frame = m_current_frame; + const size_t active_threads = ceil_div(frames_to_read, m_chunk_size); + auto worker = [&](size_t worker_index) { + File &file = m_files[worker_index]; + const size_t batch_offset = worker_index * m_chunk_size; + const size_t begin = first_frame + batch_offset; + const size_t count = + std::min(m_chunk_size, frames_to_read - batch_offset); + file.seek(begin); + file.read_into(destination + batch_offset * m_bytes_per_frame, count); + }; + + std::vector> workers; + workers.reserve(active_threads); + for (size_t i = 0; i < active_threads; ++i) { + workers.emplace_back(std::async(std::launch::async, worker, i)); + } + + for (auto &future : workers) { + future.get(); + } + m_current_frame += frames_to_read; + return frames_to_read; +} + +std::vector MultiThreadedFileReader::read() { + std::vector data(next_read_bytes()); + read_into(data.data()); + return data; +} + +std::vector MultiThreadedFileReader::read_all() { + ensure_open(); + std::vector data( + checked_product(remaining_frames(), m_bytes_per_frame)); + size_t offset = 0; + while (remaining_frames() != 0) { + const size_t frames_read = read_into(data.data() + offset); + offset += frames_read * m_bytes_per_frame; + } + return data; +} + +void MultiThreadedFileReader::seek(size_t frame_index) { + ensure_open(); + if (frame_index > m_total_frames) { + throw std::out_of_range( + "MultiThreadedFileReader frame index is out of range"); + } + m_current_frame = frame_index; +} + +} // namespace aare::experimental diff --git a/src/MultiThreadedFileReader.test.cpp b/src/MultiThreadedFileReader.test.cpp new file mode 100644 index 00000000..4896d788 --- /dev/null +++ b/src/MultiThreadedFileReader.test.cpp @@ -0,0 +1,189 @@ +// SPDX-License-Identifier: MPL-2.0 +#include "aare/MultiThreadedFileReader.hpp" + +#include "aare/Dtype.hpp" +#include "aare/File.hpp" +#include "aare/FileInterface.hpp" +#include "aare/Frame.hpp" +#include "aare/NumpyFile.hpp" + +#include "test_config.hpp" +#include + +#include +#include +#include +#include +#include +#include +#include + +using aare::File; +using aare::FileConfig; +using aare::Frame; +using aare::NumpyFile; +using aare::experimental::MultiThreadedFileReader; + +namespace { + +class TemporaryNumpyFile { + public: + TemporaryNumpyFile() { + const auto unique = + std::chrono::steady_clock::now().time_since_epoch().count(); + m_path = std::filesystem::temp_directory_path() / + ("aare-mt-reader-" + std::to_string(unique) + ".npy"); + + FileConfig cfg; + cfg.dtype = aare::Dtype::UINT16; + cfg.rows = 2; + cfg.cols = 3; + NumpyFile file(m_path, "w", cfg); + for (uint16_t frame_index = 0; frame_index < 10; ++frame_index) { + Frame frame(cfg.rows, cfg.cols, cfg.dtype); + auto image = frame.view(); + for (ssize_t row = 0; row < image.shape(0); ++row) { + for (ssize_t col = 0; col < image.shape(1); ++col) { + image(row, col) = static_cast(frame_index * 100 + + row * 10 + col); + } + } + file.write(frame); + } + } + + TemporaryNumpyFile(const TemporaryNumpyFile &) = delete; + TemporaryNumpyFile &operator=(const TemporaryNumpyFile &) = delete; + + ~TemporaryNumpyFile() { std::filesystem::remove(m_path); } + + const std::filesystem::path &path() const { return m_path; } + void truncate() { std::filesystem::resize_file(m_path, 0); } + + private: + std::filesystem::path m_path; +}; + +std::vector read_reference(const std::filesystem::path &fpath, + size_t n_frames) { + File file(fpath); + std::vector data(n_frames * file.bytes_per_frame()); + if (n_frames != 0) { + file.read_into(data.data(), n_frames); + } + return data; +} + +} // namespace + +TEST_CASE("Multi-threaded reader preserves numpy frame order", + "[MultiThreadedFileReader]") { + TemporaryNumpyFile file; + const auto &fpath = file.path(); + + MultiThreadedFileReader reader(fpath, 2, 3); + + CHECK(reader.n_threads() == 2); + CHECK(reader.chunk_size() == 3); + CHECK(reader.total_frames() == 10); + CHECK(reader.source_total_frames() == 10); + CHECK(reader.rows() == 2); + CHECK(reader.cols() == 3); + CHECK(reader.bitdepth() == 16); + CHECK(reader.total_bytes() == + reader.total_frames() * reader.bytes_per_frame()); + + const auto reference = read_reference(fpath, reader.total_frames()); + auto first = reader.read(); + auto second = reader.read(); + + CHECK(first == std::vector(reference.begin(), + reference.begin() + 6 * 12)); + CHECK(second == + std::vector(reference.begin() + 6 * 12, reference.end())); + CHECK(reader.read().empty()); + CHECK(reader.tell() == 10); + CHECK(reader.remaining_frames() == 0); +} + +TEST_CASE("Multi-threaded reader handles uneven raw chunks and frame limits", + "[.with-data][MultiThreadedFileReader]") { + const auto fpath = + test_data_path() / "raw/jungfrau/jungfrau_single_master_0.json"; + REQUIRE(std::filesystem::exists(fpath)); + + MultiThreadedFileReader reader(fpath, 2, 3, 9); + + CHECK(reader.total_frames() == 9); + CHECK(reader.source_total_frames() == 10); + CHECK(reader.read_all() == read_reference(fpath, 9)); +} + +TEST_CASE("Multi-threaded reader can seek and reread", + "[MultiThreadedFileReader]") { + TemporaryNumpyFile file; + + MultiThreadedFileReader reader(file.path(), 2, 1); + const auto first = reader.read(); + CHECK(reader.tell() == 2); + CHECK(reader.next_read_frames() == 2); + + reader.seek(0); + CHECK(reader.tell() == 0); + CHECK(reader.read() == first); + CHECK_THROWS_AS(reader.seek(11), std::out_of_range); +} + +TEST_CASE("Multi-threaded reader validates its configuration", + "[MultiThreadedFileReader]") { + TemporaryNumpyFile file; + const auto &fpath = file.path(); + + CHECK_THROWS_AS(MultiThreadedFileReader(fpath, 0, 1), + std::invalid_argument); + CHECK_THROWS_AS(MultiThreadedFileReader(fpath, 1, 0), + std::invalid_argument); + CHECK_THROWS_AS(MultiThreadedFileReader(fpath, 1, 1, 11), + std::invalid_argument); + + MultiThreadedFileReader reader(fpath, 2, 2); + CHECK_THROWS_AS(reader.read_into(nullptr), std::invalid_argument); +} + +TEST_CASE("An explicit zero frame limit produces an empty read", + "[MultiThreadedFileReader]") { + TemporaryNumpyFile file; + + MultiThreadedFileReader reader(file.path(), 8, 3, 0); + CHECK(reader.total_frames() == 0); + CHECK(reader.total_bytes() == 0); + CHECK(reader.read().empty()); + CHECK_NOTHROW(reader.read_into(nullptr)); +} + +TEST_CASE("read_into reads at most one chunk per worker", + "[MultiThreadedFileReader]") { + TemporaryNumpyFile file; + MultiThreadedFileReader reader(file.path(), 3, 2); + std::vector data(reader.next_read_bytes()); + + CHECK(reader.next_read_frames() == 6); + CHECK(reader.read_into(data.data()) == 6); + CHECK(reader.tell() == 6); + CHECK(reader.next_read_frames() == 4); + CHECK(reader.next_read_bytes() == 4 * reader.bytes_per_frame()); +} + +TEST_CASE("Multi-threaded reader can close its worker files", + "[MultiThreadedFileReader]") { + TemporaryNumpyFile file; + MultiThreadedFileReader reader(file.path(), 2, 2); + + CHECK(reader.is_open()); + reader.close(); + CHECK_FALSE(reader.is_open()); + CHECK_NOTHROW(reader.close()); + CHECK_THROWS_AS(reader.read(), std::runtime_error); + CHECK_THROWS_AS(reader.read_all(), std::runtime_error); + CHECK_THROWS_AS(reader.seek(0), std::runtime_error); +} diff --git a/src/NumpyHelpers.test.cpp b/src/NumpyHelpers.test.cpp index da170831..cf2f9fce 100644 --- a/src/NumpyHelpers.test.cpp +++ b/src/NumpyHelpers.test.cpp @@ -34,6 +34,9 @@ TEST_CASE("trim whitespace") { } TEST_CASE("parse data type descriptions") { + REQUIRE(parse_descr("|i1") == aare::Dtype::INT8); + REQUIRE(parse_descr("|u1") == aare::Dtype::UINT8); + REQUIRE(parse_descr("