diff --git a/AGENTS.md b/AGENTS.md index 4a587917..0857eb53 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -14,7 +14,7 @@ the same pull request that causes it. | Path | Holds | |---|---| -| `src/` | The pybind11 binding, 42 `.cpp` and 42 `.hpp`, compiled into the `rexlib._binding` extension | +| `src/` | The pybind11 binding, 47 `.cpp` and 47 `.hpp`, compiled into the `rexlib._binding` extension | | `python/rexlib/` | The Python package, 11 `.py`: everything the binding cannot express | | `tests/` | pytest suites, mirroring the binding's module structure | | `tests/assets/` | Two dummy plugins, built by CMake, that the plugin tests discover | @@ -69,8 +69,9 @@ name that names them. See #143. | `_binding.ndarray` | `src/core/ndarray/` | `Array`, `ArrayDescriptor` | | `_binding.hardware` | `src/core/hardware/` | Devices, sessions, queues, events, memory resources | | `_binding.dispatch` | `src/core/dispatch/` | `ExecutionContext`, `Dispatcher`, `ProgramManager` | +| `_binding.concurrency` | `src/core/concurrency/` | `Executor` and its two kinds, `Completion` | | `_binding.functional` | `src/functional/` | The operations, each taking an explicit context | -| `_binding.em.image` | `src/em/image/` | `ImageLocation`, the read and write format managers, `read` and `write` | +| `_binding.em.image` | `src/em/image/` | `ImageLocation`, the format managers and reader providers, the sources, `read`, `write` and the shape queries | A submodule is created by the `main.cpp` above it, which then hands it to the `bind_` function of the directory it stands for: `src/main.cpp` creates `em` @@ -246,6 +247,19 @@ Where it does not, `device_index` being the one left, the hash has to be written, which is rexlib's to do rather than this binding's; `_session_pool.py` keys on a `(backend, id)` tuple in the meantime. +A binding that blocks releases the GIL with +`py::call_guard()`. `Completion.wait`, `Completion.get` +and `ImageBatchSource.read` all do: held through a read, the GIL would freeze +every other Python thread for its duration and there would be nothing +asynchronous left about the interface. pybind11 converts the arguments before +the guard is constructed and the return value after it is destroyed, so +nothing touches Python without it. + +`rexlib::array` is move-only, and anything taking one by value has to be given +`destination.share()` rather than the caller's own. Bound plainly, pybind11 +would move the array out of the Python object and hand the caller back an +empty one. + A type rexlib parses from a string carries a `from_string` static method and no constructor that parses. `DeviceIndex` and `ImageLocation` both do, and `ImageLocation` is why: its constructor already gives a lone string the diff --git a/BACKLOG.md b/BACKLOG.md index 8a7ea396..d5f5916f 100644 --- a/BACKLOG.md +++ b/BACKLOG.md @@ -56,55 +56,12 @@ One trap before `compare` lands: a type that defines `__eq__` and no rexlib, so the comparison operators need an answer about identity before they go in. -## 3. Bind the batch source - -`image_batch_source` is the class meant for pipelines: asynchronous, batched, -running its own threads, and where the performance is. The slow path -(`em.image.read`) stays what it is, for one-off reads and tests. - -The chain: - -``` -ImageReadFormatManager ← bound - └→ DirectImageReaderProvider(formats) │ CachingImageReaderProvider(backing, capacity) - └→ ImageSource(readers, executor) ← ThreadPoolExecutor(n) │ SynchronousExecutor() - └→ ImageBatchSource(source) - └→ read(destination, locations) → Completion{wait, is_ready, get} -``` - -It splits in two, by risk rather than by size. - -**3a — geometry.** Bind `query_extents` and `query_core_extents` over a format -manager. Read-only surface, no threads, no GIL, and useful for sizing any read -rather than only a batch. Their overloads over a reader provider wait for 3b, -which is what binds the provider; those take a non-const reference where the -format manager ones take a const one. - -`query_core_extents` returns the shape of a single image or volume, which is -what a batch destination carries beside its leading extent, so no arithmetic -has to be reproduced here. Whether a `make_batch_descriptor` helper on top is -worth its weight is decided when this is written, with the code in view; it -would only save a caller from getting the leading extent wrong. - -**3b — the pipeline.** Executors, providers, source, batch source, completion, -the Python layer, and the GIL work. - -Settled before starting, recorded in [DECISIONS.md](DECISIONS.md): the -destination's data type is the caller's choice rather than the file's; the -thread pool is per instance with an executor accepted as an override; the -completion is mirrored rather than wrapped in a future; and `wait`, `get` and -`read` release the GIL. - -Two facts to build against, both verified in rexlib's sources: - -- `rexlib::array` is move-only and `read` takes it by value. Bound plainly, - pybind11 would move the array out of the caller's Python object and leave - them holding an empty one. The binding takes `array&` and passes - `destination.share()`. -- Nothing else needs keeping alive. The `locations` span is consumed - synchronously while the plan is built; the destination is a handle whose - storage outlives a dropped Python reference; and `image_source::read` - extracts what each task needs at submit time, so its plan does not dangle. +## 3. Bind the patch source + +`image_patch_source` arrived with the batch source upstream and crops batches +of patches out of an image, clipped at the borders. The batch source is bound; +this one is not, and it stands on the same chain, so it is a smaller piece of +work than the first was. ## 4. Translate the library's exceptions diff --git a/external/rexlib b/external/rexlib index 43bbffec..ce67271d 160000 --- a/external/rexlib +++ b/external/rexlib @@ -1 +1 @@ -Subproject commit 43bbffecbfc7b04daeb032d99391d58453d1b499 +Subproject commit ce67271ddebea6f0ac80a2e9ac08ae3bde1345ec diff --git a/python/rexlib/__init__.py b/python/rexlib/__init__.py index a145a53e..3d79b93c 100644 --- a/python/rexlib/__init__.py +++ b/python/rexlib/__init__.py @@ -25,6 +25,7 @@ from . import _paths as _paths from ._binding import ( + concurrency as concurrency, dispatch as dispatch, hardware as hardware, Plugin as Plugin, @@ -91,6 +92,7 @@ "add", "cast", "cast_copy", + "concurrency", "copy", "device", "dispatch", diff --git a/python/rexlib/em/image/__init__.py b/python/rexlib/em/image/__init__.py index 0d8ece47..915ea37a 100644 --- a/python/rexlib/em/image/__init__.py +++ b/python/rexlib/em/image/__init__.py @@ -21,13 +21,19 @@ from __future__ import annotations from ..._binding.em.image import ( + CachingImageReaderProvider as CachingImageReaderProvider, + DirectImageReaderProvider as DirectImageReaderProvider, + ImageBatchSource as ImageBatchSource, ImageLocation as ImageLocation, + ImageReaderProvider as ImageReaderProvider, + ImageSource as ImageSource, ImageReadFormatManager as ImageReadFormatManager, ImageWriteFormatManager as ImageWriteFormatManager, get_image_read_format_manager as get_image_read_format_manager, get_image_write_format_manager as get_image_write_format_manager, ) from ._functions import ( + batch_source as batch_source, query_core_extents as query_core_extents, query_extents as query_extents, read as read, @@ -35,9 +41,15 @@ ) __all__ = [ + "CachingImageReaderProvider", + "DirectImageReaderProvider", + "ImageBatchSource", "ImageLocation", "ImageReadFormatManager", + "ImageReaderProvider", + "ImageSource", "ImageWriteFormatManager", + "batch_source", "get_image_read_format_manager", "get_image_write_format_manager", "query_core_extents", diff --git a/python/rexlib/em/image/_functions.py b/python/rexlib/em/image/_functions.py index 0ba17fc9..847f4f9d 100644 --- a/python/rexlib/em/image/_functions.py +++ b/python/rexlib/em/image/_functions.py @@ -10,10 +10,17 @@ from __future__ import annotations +import os + +from ..._binding.concurrency import Executor, ThreadPoolExecutor from ..._binding.dispatch import ExecutionContext from ..._binding.em import image as _raw from ..._binding.em.image import ( + CachingImageReaderProvider, + DirectImageReaderProvider, + ImageBatchSource, ImageLocation, + ImageSource, ImageReadFormatManager, ImageWriteFormatManager, get_image_read_format_manager, @@ -24,6 +31,9 @@ from ..._catalog import get_default_catalog from ..._functional import _resolve_context +def _default_worker_count() -> int: + return os.cpu_count() or 1 + def _resolve_read_manager( manager: ImageReadFormatManager | None ) -> ImageReadFormatManager: @@ -38,6 +48,45 @@ def _resolve_write_manager( manager = get_image_write_format_manager(get_default_catalog()) return manager +def batch_source( + workers: int | None = None, + cache: int | None = None, + manager: ImageReadFormatManager | None = None, + executor: Executor | None = None +) -> ImageBatchSource: + """ + Assemble a batch source over a thread pool. + + Wires the four objects a batch source stands on — the formats, a + reader provider, an executor and the source itself — so that a + caller reaches the pipeline without naming them. + + Each source built this way owns its executor. Two of them means two + thread pools, each sized to the machine, which is worth avoiding by + passing one executor to both. + + Args: + workers: How many threads to run reads on. Defaults to what the + machine reports. Ignored when `executor` is given. + cache: How many open readers to keep between reads. Defaults to + opening a file every time it is asked for, which is what + suits reads that do not revisit a file. + manager: The formats to recognize files with. Defaults to the + ones the default catalog holds. + executor: Where reads run. Defaults to a thread pool of its own. + + Returns: + ImageBatchSource: The assembled source. + """ + readers = DirectImageReaderProvider(_resolve_read_manager(manager)) + if cache is not None: + readers = CachingImageReaderProvider(readers, cache) + if executor is None: + executor = ThreadPoolExecutor( + workers if workers is not None else _default_worker_count() + ) + return ImageBatchSource(ImageSource(readers, executor)) + def query_extents( path: str, manager: ImageReadFormatManager | None = None diff --git a/src/core/concurrency/completion.cpp b/src/core/concurrency/completion.cpp new file mode 100644 index 00000000..9c9a991d --- /dev/null +++ b/src/core/concurrency/completion.cpp @@ -0,0 +1,29 @@ +// SPDX-License-Identifier: GPL-3.0-only + +#include "completion.hpp" + +namespace rexlib +{ + +namespace py = pybind11; + +completion_class declare_completion(pybind11::module_ &m) +{ + return completion_class(m, "Completion"); +} + +void define_completion(completion_class &c) +{ + c + .def( + "wait", &completion::wait, + py::call_guard() + ) + .def( + "get", &completion::get, + py::call_guard() + ) + .def_property_readonly("is_ready", &completion::is_ready); +} + +} // namespace rexlib diff --git a/src/core/concurrency/completion.hpp b/src/core/concurrency/completion.hpp new file mode 100644 index 00000000..e5d93354 --- /dev/null +++ b/src/core/concurrency/completion.hpp @@ -0,0 +1,20 @@ +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include + +#include + +#include + +namespace rexlib +{ + +using completion_class = + pybind11::class_>; + +completion_class declare_completion(pybind11::module_ &m); +void define_completion(completion_class &c); + +} // namespace rexlib diff --git a/src/core/concurrency/executor.cpp b/src/core/concurrency/executor.cpp new file mode 100644 index 00000000..a5c394b8 --- /dev/null +++ b/src/core/concurrency/executor.cpp @@ -0,0 +1,37 @@ +// SPDX-License-Identifier: GPL-3.0-only + +#include "executor.hpp" + +#include + +namespace rexlib +{ + +namespace py = pybind11; + +executor_class declare_executor(pybind11::module_ &m) +{ + return executor_class(m, "Executor"); +} + +synchronous_executor_class declare_synchronous_executor(pybind11::module_ &m) +{ + return synchronous_executor_class(m, "SynchronousExecutor"); +} + +thread_pool_executor_class declare_thread_pool_executor(pybind11::module_ &m) +{ + return thread_pool_executor_class(m, "ThreadPoolExecutor"); +} + +void define_synchronous_executor(synchronous_executor_class &c) +{ + c.def(py::init<>()); +} + +void define_thread_pool_executor(thread_pool_executor_class &c) +{ + c.def(py::init(), py::arg("worker_count")); +} + +} // namespace rexlib diff --git a/src/core/concurrency/executor.hpp b/src/core/concurrency/executor.hpp new file mode 100644 index 00000000..00a6095a --- /dev/null +++ b/src/core/concurrency/executor.hpp @@ -0,0 +1,30 @@ +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include + +#include +#include +#include + +#include + +namespace rexlib +{ + +using executor_class = pybind11::class_>; +using synchronous_executor_class = pybind11::class_< + synchronous_executor, executor, std::shared_ptr +>; +using thread_pool_executor_class = pybind11::class_< + thread_pool_executor, executor, std::shared_ptr +>; + +executor_class declare_executor(pybind11::module_ &m); +synchronous_executor_class declare_synchronous_executor(pybind11::module_ &m); +thread_pool_executor_class declare_thread_pool_executor(pybind11::module_ &m); +void define_synchronous_executor(synchronous_executor_class &c); +void define_thread_pool_executor(thread_pool_executor_class &c); + +} // namespace rexlib diff --git a/src/core/concurrency/main.cpp b/src/core/concurrency/main.cpp new file mode 100644 index 00000000..2a5c9259 --- /dev/null +++ b/src/core/concurrency/main.cpp @@ -0,0 +1,23 @@ +// SPDX-License-Identifier: GPL-3.0-only + +#include "main.hpp" + +#include "completion.hpp" +#include "executor.hpp" + +namespace rexlib +{ + +void bind_concurrency(pybind11::module_ &m) +{ + auto completion = declare_completion(m); + declare_executor(m); + auto synchronous_executor = declare_synchronous_executor(m); + auto thread_pool_executor = declare_thread_pool_executor(m); + + define_completion(completion); + define_synchronous_executor(synchronous_executor); + define_thread_pool_executor(thread_pool_executor); +} + +} // namespace rexlib diff --git a/src/core/concurrency/main.hpp b/src/core/concurrency/main.hpp new file mode 100644 index 00000000..5798c388 --- /dev/null +++ b/src/core/concurrency/main.hpp @@ -0,0 +1,12 @@ +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include + +namespace rexlib +{ + +void bind_concurrency(pybind11::module_ &m); + +} // namespace rexlib diff --git a/src/em/image/image_batch_source.cpp b/src/em/image/image_batch_source.cpp new file mode 100644 index 00000000..60cff37c --- /dev/null +++ b/src/em/image/image_batch_source.cpp @@ -0,0 +1,69 @@ +// SPDX-License-Identifier: GPL-3.0-only + +#include "image_batch_source.hpp" + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include + +namespace rexlib +{ + +namespace py = pybind11; + +// array is move-only and read takes one by value, so the caller's array is +// shared rather than moved out of: a moved-from Array would be left empty in +// their hands. +static std::shared_ptr py_read( + const em::image_batch_source &self, + array &destination, + const std::vector &locations +) +{ + return self.read(destination.share(), make_span(locations)); +} + +image_source_class declare_image_source(pybind11::module_ &m) +{ + return image_source_class(m, "ImageSource"); +} + +image_batch_source_class declare_image_batch_source(pybind11::module_ &m) +{ + return image_batch_source_class(m, "ImageBatchSource"); +} + +void define_image_source(image_source_class &c) +{ + c.def( + py::init< + std::shared_ptr, + std::shared_ptr + >(), + py::arg("readers"), py::arg("executor") + ); +} + +void define_image_batch_source(image_batch_source_class &c) +{ + c + .def( + py::init>(), + py::arg("source") + ) + .def( + "read", &py_read, + py::arg("destination"), py::arg("locations"), + py::call_guard() + ); +} + +} // namespace rexlib diff --git a/src/em/image/image_batch_source.hpp b/src/em/image/image_batch_source.hpp new file mode 100644 index 00000000..4518eb2f --- /dev/null +++ b/src/em/image/image_batch_source.hpp @@ -0,0 +1,26 @@ +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include + +#include +#include + +#include + +namespace rexlib +{ + +using image_source_class = + pybind11::class_>; +using image_batch_source_class = pybind11::class_< + em::image_batch_source, std::shared_ptr +>; + +image_source_class declare_image_source(pybind11::module_ &m); +image_batch_source_class declare_image_batch_source(pybind11::module_ &m); +void define_image_source(image_source_class &c); +void define_image_batch_source(image_batch_source_class &c); + +} // namespace rexlib diff --git a/src/em/image/image_reader_provider.cpp b/src/em/image/image_reader_provider.cpp new file mode 100644 index 00000000..0c9fbb46 --- /dev/null +++ b/src/em/image/image_reader_provider.cpp @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: GPL-3.0-only + +#include "image_reader_provider.hpp" + +#include + +#include +#include +#include + +namespace rexlib +{ + +namespace py = pybind11; + +static py::tuple to_shape(const std::vector &extents) +{ + py::tuple shape(extents.size()); + for (std::size_t i = 0; i < extents.size(); ++i) + { + shape[i] = extents[i]; + } + return shape; +} + +static py::tuple query_extents( + em::image_reader_provider &readers, + const std::string &path +) +{ + return to_shape(em::query_extents(readers, path)); +} + +static py::tuple query_core_extents( + em::image_reader_provider &readers, + const std::string &path +) +{ + return to_shape(em::query_core_extents(readers, path)); +} + +image_reader_provider_class +declare_image_reader_provider(pybind11::module_ &m) +{ + return image_reader_provider_class(m, "ImageReaderProvider"); +} + +direct_image_reader_provider_class +declare_direct_image_reader_provider(pybind11::module_ &m) +{ + return direct_image_reader_provider_class(m, "DirectImageReaderProvider"); +} + +caching_image_reader_provider_class +declare_caching_image_reader_provider(pybind11::module_ &m) +{ + return caching_image_reader_provider_class(m, "CachingImageReaderProvider"); +} + +void define_image_reader_provider(pybind11::module_ &m) +{ + m.def( + "query_extents", &query_extents, + py::arg("readers"), py::arg("path") + ); + m.def( + "query_core_extents", &query_core_extents, + py::arg("readers"), py::arg("path") + ); +} + +void define_direct_image_reader_provider( + direct_image_reader_provider_class &c +) +{ + c.def( + py::init>(), + py::arg("formats") + ); +} + +void define_caching_image_reader_provider( + caching_image_reader_provider_class &c +) +{ + c.def( + py::init, std::size_t>(), + py::arg("backing"), py::arg("capacity") + ); +} + +} // namespace rexlib diff --git a/src/em/image/image_reader_provider.hpp b/src/em/image/image_reader_provider.hpp new file mode 100644 index 00000000..629449ac --- /dev/null +++ b/src/em/image/image_reader_provider.hpp @@ -0,0 +1,44 @@ +// SPDX-License-Identifier: GPL-3.0-only + +#pragma once + +#include + +#include +#include +#include + +#include + +namespace rexlib +{ + +using image_reader_provider_class = pybind11::class_< + em::image_reader_provider, std::shared_ptr +>; +using direct_image_reader_provider_class = pybind11::class_< + em::direct_image_reader_provider, + em::image_reader_provider, + std::shared_ptr +>; +using caching_image_reader_provider_class = pybind11::class_< + em::caching_image_reader_provider, + em::image_reader_provider, + std::shared_ptr +>; + +image_reader_provider_class +declare_image_reader_provider(pybind11::module_ &m); +direct_image_reader_provider_class +declare_direct_image_reader_provider(pybind11::module_ &m); +caching_image_reader_provider_class +declare_caching_image_reader_provider(pybind11::module_ &m); +void define_image_reader_provider(pybind11::module_ &m); +void define_direct_image_reader_provider( + direct_image_reader_provider_class &c +); +void define_caching_image_reader_provider( + caching_image_reader_provider_class &c +); + +} // namespace rexlib diff --git a/src/em/image/main.cpp b/src/em/image/main.cpp index fc1a3442..f5f8596f 100644 --- a/src/em/image/main.cpp +++ b/src/em/image/main.cpp @@ -2,9 +2,11 @@ #include "main.hpp" +#include "image_batch_source.hpp" #include "image_location.hpp" #include "image_read.hpp" #include "image_read_format_manager.hpp" +#include "image_reader_provider.hpp" #include "image_write.hpp" #include "image_write_format_manager.hpp" @@ -14,10 +16,20 @@ namespace rexlib void bind_image(pybind11::module_ &m) { auto image_location = declare_image_location(m); + declare_image_reader_provider(m); + auto direct_provider = declare_direct_image_reader_provider(m); + auto caching_provider = declare_caching_image_reader_provider(m); + auto image_source = declare_image_source(m); + auto image_batch_source = declare_image_batch_source(m); declare_image_read_format_manager(m); declare_image_write_format_manager(m); define_image_location(image_location); + define_image_reader_provider(m); + define_direct_image_reader_provider(direct_provider); + define_caching_image_reader_provider(caching_provider); + define_image_source(image_source); + define_image_batch_source(image_batch_source); define_image_read_format_manager(m); define_image_write_format_manager(m); diff --git a/src/main.cpp b/src/main.cpp index 4887d26e..8bc57b6a 100644 --- a/src/main.cpp +++ b/src/main.cpp @@ -2,6 +2,7 @@ #include "core/hardware/main.hpp" #include "core/ndarray/main.hpp" #include "core/numerical/main.hpp" +#include "core/concurrency/main.hpp" #include "core/dispatch/main.hpp" #include "em/main.hpp" #include "functional/main.hpp" @@ -33,6 +34,9 @@ PYBIND11_MODULE(_binding, m) { auto dispatch_module = m.def_submodule("dispatch"); bind_dispatch(dispatch_module); + auto concurrency_module = m.def_submodule("concurrency"); + bind_concurrency(concurrency_module); + auto functional_module = m.def_submodule("functional"); bind_functional(functional_module); diff --git a/tests/em/image/test_batch_source.py b/tests/em/image/test_batch_source.py new file mode 100644 index 00000000..dbf596cd --- /dev/null +++ b/tests/em/image/test_batch_source.py @@ -0,0 +1,127 @@ +# SPDX-License-Identifier: GPL-3.0-only + +import pytest + +import rexlib + +image = rexlib.em.image + +def test_reads_a_batch(__written_files, __setup_context): + source = image.batch_source() + destination = __setup_destination(__written_files, __setup_context) + source.read(destination, __written_files).get() + assert destination.shape[0] == len(__written_files) + +def test_the_destination_survives_the_read(__written_files, __setup_context): + # read takes the array by value and array is move-only, so a binding + # that moved rather than shared would leave this one empty. + source = image.batch_source() + destination = __setup_destination(__written_files, __setup_context) + source.read(destination, __written_files).get() + assert destination.shape == (len(__written_files), 4, 6) + +def test_batches_in_flight_can_be_collected_together( + __written_files, __setup_context +): + source = image.batch_source() + first = __setup_destination(__written_files, __setup_context) + second = __setup_destination(__written_files, __setup_context) + completions = [ + source.read(first, __written_files), + source.read(second, __written_files), + ] + for completion in completions: + completion.get() + assert all(c.is_ready for c in completions) + +def test_a_completion_reports_when_it_is_done(__written_files, __setup_context): + source = image.batch_source() + destination = __setup_destination(__written_files, __setup_context) + completion = source.read(destination, __written_files) + completion.wait() + assert completion.is_ready + +def test_an_empty_batch_is_already_done(__setup_context): + source = image.batch_source() + descriptor = rexlib.make_contiguous_array_descriptor( + (0, 4, 6), rexlib.NumericalType.float32 + ) + destination = rexlib.zeros( + descriptor, rexlib.hardware.MemoryResourceAffinity.host, __setup_context + ) + assert source.read(destination, []).is_ready + +def test_a_destination_of_the_wrong_batch_size_is_refused( + __written_files, __setup_context +): + source = image.batch_source() + descriptor = rexlib.make_contiguous_array_descriptor( + (len(__written_files) + 1, 4, 6), rexlib.NumericalType.float32 + ) + destination = rexlib.zeros( + descriptor, rexlib.hardware.MemoryResourceAffinity.host, __setup_context + ) + with pytest.raises(ValueError): + source.read(destination, __written_files) + +def test_runs_on_a_synchronous_executor_too(__written_files, __setup_context): + source = image.batch_source( + executor=rexlib.concurrency.SynchronousExecutor() + ) + destination = __setup_destination(__written_files, __setup_context) + source.read(destination, __written_files).get() + assert destination.shape[0] == len(__written_files) + +def test_a_caching_provider_serves_repeated_reads( + __written_files, __setup_context +): + source = image.batch_source(cache=4) + destination = __setup_destination(__written_files, __setup_context) + source.read(destination, __written_files).get() + source.read(destination, __written_files).get() + assert destination.shape[0] == len(__written_files) + +def test_assembled_by_hand(__written_files, __setup_context): + catalog = rexlib.ServiceCatalog() + readers = image.DirectImageReaderProvider( + image.get_image_read_format_manager(catalog) + ) + executor = rexlib.concurrency.ThreadPoolExecutor(2) + source = image.ImageBatchSource(image.ImageSource(readers, executor)) + destination = __setup_destination(__written_files, __setup_context) + source.read(destination, __written_files).get() + assert destination.shape[0] == len(__written_files) + +def __setup_destination(locations, context): + descriptor = rexlib.make_contiguous_array_descriptor( + (len(locations), *image.query_core_extents(locations[0].path)), + rexlib.NumericalType.float32 + ) + return rexlib.zeros( + descriptor, rexlib.hardware.MemoryResourceAffinity.host, context + ) + +@pytest.fixture +def __written_files(tmp_path, __setup_context): + descriptor = rexlib.make_contiguous_array_descriptor( + (4, 6), rexlib.NumericalType.float32 + ) + source = rexlib.zeros( + descriptor, rexlib.hardware.MemoryResourceAffinity.host, __setup_context + ) + locations = [] + for index in range(3): + path = str(tmp_path / f'image{index}.mrc') + image.write(source, path) + locations.append(image.ImageLocation(path)) + return locations + +@pytest.fixture +def __setup_context(): + catalog = rexlib.ServiceCatalog() + manager = rexlib.hardware.get_device_manager(catalog) + session = manager.create_device_session(rexlib.hardware.DeviceIndex('cpu', 0)) + device_context = rexlib.hardware.DeviceContext(session) + program_manager = rexlib.dispatch.get_program_manager(catalog) + dispatcher = rexlib.dispatch.make_eager_dispatcher(program_manager) + return rexlib.dispatch.ExecutionContext(device_context, dispatcher)