0.5.2: counters keep declined routes, import race, lists of arrays

CHANGELOG.md

@@ -1,5 +1,29 @@
 # Changelog
 
+## 0.5.2
+
+### Counters keep declined routes
+
+`ConversionCounters` keyed its totals by source, target and mechanism, so two
+conversions that declined different routes were merged and `declined_mechanisms`
+was lost in `rows`, `merge`, `to_json`/`from_json` and `format_table`. The key is
+now derived in one place from every `ConversionRecord` field except `nbytes` and
+`seconds`, so a field added to the record cannot fall out of the counters. Rows
+gain `declined_mechanisms` (a list); `format_table` shows it in a last column.
+
+### A library is not loaded while another thread is importing it
+
+`ArrayLibrary.loaded_module()` answers `None` for a module whose
+`__spec__._initializing` is true, so `module` goes through
+`importlib.import_module`, which waits on the module lock. Two threads first
+using a library together no longer see `partially initialized module`.
+
+### A list of accelerator arrays gets its stream and OOM retry
+
+The device of a call is read from direct array arguments and from the items of
+list and tuple arguments (one level, no deeper). Arrays of one library on
+different devices in one call now raise `ValueError` naming the devices.
+
 ## 0.5.1
 
 ### No library inherits an operation it cannot perform

docs/source/api_reference.rst

@@ -204,7 +204,7 @@ plan, says what was done to that array.
 
 ``ConversionCounters``
   An observer that keeps count, bytes and seconds per source library and
-  device, target library and device, and mechanism. ``rows()`` (also named
+  device, target library and device, mechanism and declined mechanisms. ``rows()`` (also named
   ``to_json()``) returns one row per kind, largest byte total first;
   ``from_json(rows)`` rebuilds counters; ``merge(other)`` adds another
   process's totals; ``format_table()`` returns aligned text.

src/arraybridge/__init__.py

@@ -10,7 +10,7 @@
 NumPy/numcodecs import cost.
 """
 
-__version__ = "0.5.1"
+__version__ = "0.5.2"
 
 _LAZY_EXPORTS: dict[str, str] = {
     "MemoryType": ".types",

src/arraybridge/conversion.py

@@ -22,7 +22,7 @@
 from collections.abc import Iterator, Mapping
 from contextlib import contextmanager
 from contextvars import ContextVar
-from dataclasses import asdict, dataclass, field
+from dataclasses import asdict, dataclass, field, fields
 from types import NotImplementedType
 from typing import TYPE_CHECKING, ClassVar, TypeGuard, TypeVar
 
@@ -454,7 +454,17 @@ def record(self, record: ConversionRecord) -> None:
     "arraybridge conversion observer", default=IgnoreConversions()
 )
 
-ConversionKey = tuple[str, str, str, str, str]
+# A conversion kind is identified by every field of the record except its measurements.
+_MEASUREMENTS = frozenset({"nbytes", "seconds"})
+_IDENTITY_FIELDS = tuple(
+    each.name for each in fields(ConversionRecord) if each.name not in _MEASUREMENTS
+)
+ConversionKey = tuple[str | tuple[str, ...], ...]
+JsonRow = dict[str, str | int | float | list[str]]
+
+
+def _key_of(record: ConversionRecord) -> ConversionKey:
+    return tuple(getattr(record, name) for name in _IDENTITY_FIELDS)
 
 
 @dataclass(slots=True)
@@ -468,27 +478,18 @@ class ConversionTotal:
 
 @dataclass(slots=True)
 class ConversionCounters(ConversionObserver):
-    """Per-process totals keyed by source, target and mechanism.
+    """Per-process totals keyed by what identifies a conversion.
 
-    Counters from worker processes are plain data: pickle them or send
-    ``to_json()`` back, then ``merge`` them into the parent's.
+    The key is every field of ``ConversionRecord`` but ``nbytes`` and
+    ``seconds``: the libraries and devices, the mechanism and the declined
+    mechanisms. Counters from worker processes are plain data: pickle them or
+    send ``to_json()`` back, then ``merge`` them into the parent's.
     """
 
     totals: dict[ConversionKey, ConversionTotal] = field(default_factory=dict)
 
     def record(self, record: ConversionRecord) -> None:
-        self._add(
-            (
-                record.source_library,
-                record.source_device,
-                record.target_library,
-                record.target_device,
-                record.mechanism,
-            ),
-            1,
-            record.nbytes,
-            record.seconds,
-        )
+        self._add(_key_of(record), 1, record.nbytes, record.seconds)
 
     def _add(self, key: ConversionKey, count: int, nbytes: int, seconds: float) -> None:
         total = self.totals.setdefault(key, ConversionTotal())
@@ -503,55 +504,56 @@ def merge(self, other: ConversionCounters) -> ConversionCounters:
             self._add(key, total.count, total.nbytes, total.seconds)
         return self
 
-    def rows(self) -> list[dict[str, str | int | float]]:
+    def _largest_first(self) -> list[tuple[ConversionKey, ConversionTotal]]:
+        return sorted(self.totals.items(), key=lambda item: (-item[1].nbytes, item[0]))
+
+    def rows(self) -> list[JsonRow]:
         """One JSON-ready row per kind of conversion, largest byte total first."""
 
         return [
             {
-                "source_library": key[0],
-                "source_device": key[1],
-                "target_library": key[2],
-                "target_device": key[3],
-                "mechanism": key[4],
+                **{
+                    name: list(part) if isinstance(part, tuple) else part
+                    for name, part in zip(_IDENTITY_FIELDS, key, strict=True)
+                },
                 "count": total.count,
                 "nbytes": total.nbytes,
                 "seconds": total.seconds,
             }
-            for key, total in sorted(
-                self.totals.items(), key=lambda item: (-item[1].nbytes, item[0])
-            )
+            for key, total in self._largest_first()
         ]
 
     to_json = rows
 
     @classmethod
-    def from_json(cls, rows: list[dict[str, str | int | float]]) -> ConversionCounters:
+    def from_json(cls, rows: list[JsonRow]) -> ConversionCounters:
         counters = cls()
         for row in rows:
+            key = tuple(
+                tuple(map(str, part)) if isinstance(part, list) else str(part)
+                for part in (row[name] for name in _IDENTITY_FIELDS)
+            )
             counters._add(
-                (
-                    str(row["source_library"]),
-                    str(row["source_device"]),
-                    str(row["target_library"]),
-                    str(row["target_device"]),
-                    str(row["mechanism"]),
-                ),
-                int(row["count"]),
-                int(row["nbytes"]),
-                float(row["seconds"]),
+                key, int(str(row["count"])), int(str(row["nbytes"])), float(str(row["seconds"]))
             )
         return counters
 
     def format_table(self) -> str:
         """The totals as aligned text, one line per kind of conversion."""
 
         lines = [
-            f"{'source':<28} {'target':<28} {'mechanism':<12} "
-            f"{'count':>7} {'MB':>10} {'seconds':>9}"
+            f"{'source':<28} {'target':<28} {'mechanism':<20} "
+            f"{'count':>7} {'MB':>10} {'seconds':>9}  declined"
         ]
-        for key, total in sorted(self.totals.items(), key=lambda item: (-item[1].nbytes, item[0])):
+        for key, total in self._largest_first():
+            row = dict(zip(_IDENTITY_FIELDS, key, strict=True))
+            declined = row["declined_mechanisms"]
+            assert isinstance(declined, tuple)
             lines.append(
-                f"{key[0] + ' ' + key[1]:<28} {key[2] + ' ' + key[3]:<28} {key[4]:<12} "
-                f"{total.count:>7} {total.nbytes / 1e6:>10.1f} {total.seconds:>9.3f}"
+                f"{row['source_library']} {row['source_device']}".ljust(29)
+                + f"{row['target_library']} {row['target_device']}".ljust(29)
+                + f"{row['mechanism']!s:<20} "
+                + f"{total.count:>7} {total.nbytes / 1e6:>10.1f} {total.seconds:>9.3f}  "
+                + ", ".join(declined)
             )
         return "\n".join(lines)

src/arraybridge/libraries.py

@@ -370,9 +370,17 @@ def prepare_import(self) -> None:
             os.environ.setdefault(name, value)
 
     def loaded_module(self) -> ModuleType | None:
-        """The library's module if something already imported it."""
+        """The library's module if something has finished importing it.
 
-        return sys.modules.get(self.import_name)
+        A module another thread is still executing is not loaded yet: it is
+        answered ``None``, and ``import_module`` waits for it.
+        """
+
+        module = sys.modules.get(self.import_name)
+        spec = getattr(module, "__spec__", None)
+        if spec is not None and getattr(spec, "_initializing", False):
+            return None
+        return module
 
     def import_module(self) -> ModuleType:
         """Import the library's module."""
@@ -895,15 +903,28 @@ def _upload(self, data: np.ndarray, device: AcceleratorDevice) -> NativeT:
         """Create a native array on ``device``; called inside that device's scope."""
 
     def _argument_device(self, arguments: Sequence[object]) -> Device | None:
-        """The device of the first array of this library among a call's arguments."""
+        """The one device holding the arrays of this library among a call's arguments.
+
+        An array is an argument itself or an item of an argument that is a list
+        or a tuple; nothing deeper is searched. Arrays on different devices in
+        one call raise ``ValueError`` naming the devices.
+        """
 
         if self.loaded_module() is None:
             return None
-        for value in arguments:
-            device = self._device_if_accepted(value)
-            if device is not None:
-                return device
-        return None
+        devices: dict[str, Device] = {}
+        for argument in arguments:
+            items = argument if isinstance(argument, (list, tuple)) else (argument,)
+            for value in items:
+                device = self._device_if_accepted(value)
+                if device is not None:
+                    devices.setdefault(device.label, device)
+        if len(devices) > 1:
+            raise ValueError(
+                f"The arrays of one {self.display_name} call are on different devices: "
+                + ", ".join(sorted(devices))
+            )
+        return next(iter(devices.values()), None)
 
     @abstractmethod
     def _device_if_accepted(self, value: object) -> Device | None:

tests/test_conversion.py

@@ -757,6 +757,7 @@ def test_totals_are_kept_per_kind_of_conversion(self):
                 "target_library": "torch",
                 "target_device": "accelerator:0",
                 "mechanism": "host copy",
+                "declined_mechanisms": [],
                 "count": 2,
                 "nbytes": 40,
                 "seconds": 0.5,
@@ -767,6 +768,7 @@ def test_totals_are_kept_per_kind_of_conversion(self):
                 "target_library": "torch",
                 "target_device": "accelerator:0",
                 "mechanism": "dlpack",
+                "declined_mechanisms": [],
                 "count": 1,
                 "nbytes": 10,
                 "seconds": 0.25,
@@ -809,7 +811,15 @@ def test_table_has_one_line_per_row(self):
 
         header, *lines = counters.format_table().splitlines()
 
-        assert header.split() == ["source", "target", "mechanism", "count", "MB", "seconds"]
+        assert header.split() == [
+            "source",
+            "target",
+            "mechanism",
+            "count",
+            "MB",
+            "seconds",
+            "declined",
+        ]
         assert len(lines) == len(counters.rows())
         for line, row in zip(lines, counters.rows(), strict=True):
             for field in ("source_library", "source_device", "target_library", "target_device"):
@@ -818,3 +828,40 @@ def test_table_has_one_line_per_row(self):
             assert f"{row['count']:>7}" in line
             assert f"{row['nbytes'] / 1e6:.1f}" in line
             assert f"{row['seconds']:.3f}" in line
+
+
+class TestCounterKey:
+    def test_every_field_but_the_measurements_identifies_a_kind(self):
+        import dataclasses
+
+        base = ConversionRecord("numpy", "host", "torch", "accelerator:0", "host copy", 1, 1.0)
+        values = {
+            "source_library": "cupy",
+            "source_device": "accelerator:1",
+            "target_library": "jax",
+            "target_device": "host",
+            "mechanism": "dlpack",
+            "declined_mechanisms": ("dlpack",),
+        }
+        for each in dataclasses.fields(ConversionRecord):
+            if each.name in ("nbytes", "seconds"):
+                continue
+            other = dataclasses.replace(base, **{each.name: values[each.name]})
+            counters = ConversionCounters()
+            counters.record(base)
+            counters.record(other)
+            assert len(counters.totals) == 2, each.name
+            assert ConversionCounters.from_json(counters.to_json()) == counters, each.name
+            assert pickle.loads(pickle.dumps(counters)) == counters, each.name
+            assert ConversionCounters().merge(counters) == counters, each.name
+
+    def test_declined_routes_are_shown_and_survive_json(self):
+        counters = ConversionCounters()
+        counters.record(
+            ConversionRecord("numpy", "host", "torch", "host", "host copy", 5, 1.0, ("dlpack",))
+        )
+        rows = json.loads(json.dumps(counters.to_json()))
+
+        assert rows[0]["declined_mechanisms"] == ["dlpack"]
+        assert ConversionCounters.from_json(rows) == counters
+        assert "dlpack" in counters.format_table().splitlines()[1]

tests/test_devices.py

@@ -13,6 +13,7 @@
 from arraybridge.libraries import ArrayLibrary, HostDevices
 from tests.fakes import (
     FakeCupy,
+    FakeCupyArray,
     FakePyclesperanto,
     FakePyclesperantoArray,
     FakeTorch,
@@ -262,3 +263,39 @@ def test_current_device_is_the_one_the_library_has_selected(self, monkeypatch):
 
         assert ArrayLibrary.named("cupy").current_device() == AcceleratorDevice(1)
         assert TORCH.current_device() is None
+
+
+class TestArraysInLists:
+    @staticmethod
+    def seen_device(monkeypatch, *arguments):
+        module = FakeCupy()
+        monkeypatch.setitem(sys.modules, "cupy", module)
+        cupy = ArrayLibrary.named("cupy")
+        wrapped = cupy.execution_wrapper(lambda *_: None, oom_recovery=True)
+        wrapped(*[argument(module) for argument in arguments])
+        streams = [stream for event, stream in module.events if event == "stream enter"]
+        return [stream.device_index for stream in streams]
+
+    def test_a_list_of_arrays_runs_on_their_device(self, monkeypatch):
+        def arrays(module):
+            return [FakeCupyArray(np.zeros(2), 1, module), FakeCupyArray(np.zeros(2), 1, module)]
+
+        assert self.seen_device(monkeypatch, arrays) == [1]
+
+    def test_a_tuple_of_arrays_runs_on_their_device(self, monkeypatch):
+        assert self.seen_device(monkeypatch, lambda m: (FakeCupyArray(np.zeros(2), 1, m),)) == [1]
+
+    def test_only_one_level_is_searched(self, monkeypatch):
+        module = FakeCupy()
+        monkeypatch.setitem(sys.modules, "cupy", module)
+        cupy = ArrayLibrary.named("cupy")
+        nested = [[FakeCupyArray(np.zeros(2), 1, module)]]
+
+        assert cupy._argument_device((nested,)) is None
+
+    def test_arrays_on_different_devices_name_them(self, monkeypatch):
+        def arrays(module):
+            return [FakeCupyArray(np.zeros(2), 0, module), FakeCupyArray(np.zeros(2), 1, module)]
+
+        with pytest.raises(ValueError, match="accelerator:0.*accelerator:1"):
+            self.seen_device(monkeypatch, arrays)