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/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);