diff --git a/CHANGELOG.md b/CHANGELOG.md index 5c3c1a1..3dc0c83 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/docs/source/api_reference.rst b/docs/source/api_reference.rst index 0e1120c..7286192 100644 --- a/docs/source/api_reference.rst +++ b/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. diff --git a/src/arraybridge/__init__.py b/src/arraybridge/__init__.py index 5376e43..3ba33a2 100644 --- a/src/arraybridge/__init__.py +++ b/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", diff --git a/src/arraybridge/conversion.py b/src/arraybridge/conversion.py index f810f32..12ef57c 100644 --- a/src/arraybridge/conversion.py +++ b/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,42 +504,37 @@ 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 @@ -546,12 +542,18 @@ 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) diff --git a/src/arraybridge/libraries.py b/src/arraybridge/libraries.py index 17285db..1b1c351 100644 --- a/src/arraybridge/libraries.py +++ b/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: diff --git a/tests/test_conversion.py b/tests/test_conversion.py index 6c23a5e..7f88164 100644 --- a/tests/test_conversion.py +++ b/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] diff --git a/tests/test_devices.py b/tests/test_devices.py index 622b2a6..84b4e97 100644 --- a/tests/test_devices.py +++ b/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) diff --git a/tests/test_libraries.py b/tests/test_libraries.py index 0b0fb56..40e3b0a 100644 --- a/tests/test_libraries.py +++ b/tests/test_libraries.py @@ -228,9 +228,7 @@ def always_fails(): raise MemoryError("out of memory") with pytest.raises(MemoryError, match="out of memory"): - ArrayLibrary.named("torch").run_with_oom_recovery( - always_fails, max_retries=max_retries - ) + ArrayLibrary.named("torch").run_with_oom_recovery(always_fails, max_retries=max_retries) assert len(calls) == expected_calls def test_other_errors_are_not_retried(self): @@ -517,3 +515,47 @@ def test_torch_wraps_a_reversed_view_as_a_contiguous_copy(self, monkeypatch): np.testing.assert_array_equal(host_array, values) assert torch.device_of(on_host) == HOST assert torch.device_of(on_accelerator) == AcceleratorDevice(0) + + +class TestImportRace: + def test_a_module_still_importing_is_waited_for(self, monkeypatch): + import importlib.abc + import importlib.util + import time + from threading import Barrier, Thread + + class SlowLoader(importlib.abc.Loader): + def create_module(self, spec): + return None + + def exec_module(self, module): + time.sleep(0.01) + module.ready = True + + class SlowFinder(importlib.abc.MetaPathFinder): + def find_spec(self, name, path=None, target=None): + if name == "slow_fake_library": + return importlib.util.spec_from_loader(name, SlowLoader()) + return None + + finder = SlowFinder() + monkeypatch.setattr(sys, "meta_path", [finder, *sys.meta_path]) + library = ArrayLibrary.named("numpy") + monkeypatch.setattr(library, "import_name", "slow_fake_library") + monkeypatch.setattr(library, "prepare_import", lambda: None) + + for _ in range(50): + sys.modules.pop("slow_fake_library", None) + seen = [] + barrier = Barrier(2) + + def use(): + barrier.wait() + seen.append(hasattr(library.module, "ready")) + + threads = [Thread(target=use) for _ in range(2)] + for thread in threads: + thread.start() + for thread in threads: + thread.join() + assert seen == [True, True]