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

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

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

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"]

python/src/bind_MultiThreadedFileReader.hpp

@@ -0,0 +1,175 @@
+// SPDX-License-Identifier: MPL-2.0
+#pragma once
+
+#include "aare/MultiThreadedFileReader.hpp"
+
+#include <cstddef>
+#include <cstdint>
+#include <filesystem>
+#include <optional>
+#include <stdexcept>
+#include <vector>
+
+#include <pybind11/numpy.h>
+#include <pybind11/pybind11.h>
+#include <pybind11/stl.h>
+#include <pybind11/stl/filesystem.h>
+
+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<int8_t>();
+    if (dtype == Dtype::UINT8)
+        return py::dtype::of<uint8_t>();
+    if (dtype == Dtype::INT16)
+        return py::dtype::of<int16_t>();
+    if (dtype == Dtype::UINT16)
+        return py::dtype::of<uint16_t>();
+    if (dtype == Dtype::INT32)
+        return py::dtype::of<int32_t>();
+    if (dtype == Dtype::UINT32)
+        return py::dtype::of<uint32_t>();
+    if (dtype == Dtype::INT64)
+        return py::dtype::of<int64_t>();
+    if (dtype == Dtype::UINT64)
+        return py::dtype::of<uint64_t>();
+    if (dtype == Dtype::FLOAT)
+        return py::dtype::of<float>();
+    if (dtype == Dtype::DOUBLE)
+        return py::dtype::of<double>();
+    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<py::ssize_t> shape{
+        static_cast<py::ssize_t>(n_frames),
+        static_cast<py::ssize_t>(reader.rows()),
+        static_cast<py::ssize_t>(reader.cols())};
+
+    py::array image(multi_threaded_reader_numpy_dtype(reader.dtype()), shape);
+    auto *destination = reinterpret_cast<std::byte *>(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_<MultiThreadedFileReader>(m, "MultiThreadedFileReader");
+    reader.attr("__module__") = "aare.experimental";
+    reader
+        .def(py::init<std::filesystem::path, size_t, size_t,
+                      std::optional<size_t>>(),
+             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);
+        });
+}

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<T, N, M, U>(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);

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,
+        )

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)