From f42f70e347d6097f010aca30fd75688dbd9a5458 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Sat, 31 Oct 2026 12:10:50 -0700 Subject: [PATCH 1/6] Make SyclDevice and SyclQueue picklable Root devices pickle by filter string and name; unpickling checks the name. Queues pickle by device and properties and unpickle into the device's default context; queues on non-default contexts and on sub-devices raise PicklingError. copy and deepcopy return the same object. Closes #1716 --- CHANGELOG.md | 2 + dpctl/_sycl_device.pyx | 34 +++++++++++++++ dpctl/_sycl_queue.pyx | 47 +++++++++++++++++++++ dpctl/tests/helper/__init__.py | 2 + dpctl/tests/helper/_helper.py | 10 +++++ dpctl/tests/test_sycl_device.py | 40 +++++++++++++++++- dpctl/tests/test_sycl_queue.py | 73 ++++++++++++++++++++++++++++++++- 7 files changed, 206 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 0aa6be7eab..e652bf82b3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added +* `dpctl.SyclDevice` and `dpctl.SyclQueue` can be pickled; root devices are restored by filter string, queues by device and properties into the device's default context. `copy.copy` and `copy.deepcopy` return the same object + ### Changed ### Deprecated diff --git a/dpctl/_sycl_device.pyx b/dpctl/_sycl_device.pyx index 624575a5d8..41e491d3aa 100644 --- a/dpctl/_sycl_device.pyx +++ b/dpctl/_sycl_device.pyx @@ -155,6 +155,7 @@ from ._sycl_platform cimport SyclPlatform import collections import functools +import pickle import warnings __all__ = [ @@ -371,6 +372,21 @@ def _cached_filter_string(d : SyclDevice): return ":".join((br_str, dt_str, str(relId))) +def _rebuild_device(str filter_string, str name): + """Unpickles a :class:`dpctl.SyclDevice`. + + Raises: + pickle.UnpicklingError: if the device selected by ``filter_string`` + is not named ``name``. + """ + dev = SyclDevice(filter_string) + if dev.name != name: + raise pickle.UnpicklingError( + f"Device '{filter_string}' is '{dev.name}', expected '{name}'" + ) + return dev + + cdef class SyclDevice(_SyclDevice): """ SyclDevice(arg=None) A Python wrapper for the ``sycl::device`` C++ class. @@ -1652,6 +1668,24 @@ cdef class SyclDevice(_SyclDevice): """ return DPCTLDevice_Hash(self._device_ref) + def __reduce__(self): + """Pickles a root device by its filter string and name. + + Raises: + pickle.PicklingError: if the device is not a root device. + """ + try: + filter_string = self.filter_string + except ValueError as e: + raise pickle.PicklingError(f"Cannot pickle {self!r}: {e}") from e + return _rebuild_device, (filter_string, self.name) + + def __copy__(self): + return self + + def __deepcopy__(self, memo): + return self + cdef tuple create_sub_devices_equally(self, size_t count): """ Returns a tuple of sub-devices partitioned from this SYCL device based on the ``count`` parameter. diff --git a/dpctl/_sycl_queue.pyx b/dpctl/_sycl_queue.pyx index 3fbf9c7bf6..ce5cadaf0a 100644 --- a/dpctl/_sycl_queue.pyx +++ b/dpctl/_sycl_queue.pyx @@ -23,6 +23,7 @@ """ from ._backend cimport ( # noqa: E211 + DPCTLContext_AreEq, DPCTLContext_Create, DPCTLContext_Delete, DPCTLDefaultSelector_Create, @@ -103,6 +104,7 @@ from libc.stdlib cimport free, malloc import collections.abc import logging +import pickle import struct import sys @@ -775,6 +777,13 @@ cdef DPCTLSyclEventRef _memset_impl( return ERef +def _rebuild_queue(SyclDevice dev, tuple props): + """Unpickles a :class:`dpctl.SyclQueue` into the default context of + ``dev``. + """ + return SyclQueue(dev, property=props) + + cdef class _SyclQueue: """ Barebone data owner class used by SyclQueue. """ @@ -2190,6 +2199,44 @@ cdef class SyclQueue(_SyclQueue): """ return DPCTLQueue_Hash(self._queue_ref) + def __reduce__(self): + """Pickles the queue's device and properties. The unpickled queue + uses the default context of the device. + + Raises: + pickle.PicklingError: if the queue's context is not the default + context of its device, or the device is not a root device. + """ + cdef DPCTLSyclContextRef CRef = NULL + cdef bint is_default_ctx = False + # sub-devices have no cached context, SyclDevice.__reduce__ rejects them + if self._device.parent_device is None: + CRef = DPCTLDeviceMgr_GetCachedContext( + self._device.get_device_ref() + ) + if CRef is not NULL: + is_default_ctx = DPCTLContext_AreEq( + CRef, self._context.get_context_ref() + ) + DPCTLContext_Delete(CRef) + if not is_default_ctx: + raise pickle.PicklingError( + f"Cannot pickle {self!r}: its context is not the default " + "context of its device" + ) + props = [] + if DPCTLQueue_IsInOrder(self._queue_ref): + props.append("in_order") + if DPCTLQueue_HasEnableProfiling(self._queue_ref): + props.append("enable_profiling") + return _rebuild_queue, (self._device, tuple(props)) + + def __copy__(self): + return self + + def __deepcopy__(self, memo): + return self + def _get_capsule(self): cdef DPCTLSyclQueueRef QRef = NULL QRef = DPCTLQueue_Copy(self._queue_ref) diff --git a/dpctl/tests/helper/__init__.py b/dpctl/tests/helper/__init__.py index be5f8d2e33..bb80f1d7bb 100644 --- a/dpctl/tests/helper/__init__.py +++ b/dpctl/tests/helper/__init__.py @@ -19,6 +19,7 @@ from ._helper import ( create_invalid_capsule, get_queue_or_skip, + get_sub_device_or_skip, has_cpu, has_gpu, has_sycl_platforms, @@ -30,4 +31,5 @@ "has_gpu", "has_sycl_platforms", "get_queue_or_skip", + "get_sub_device_or_skip", ] diff --git a/dpctl/tests/helper/_helper.py b/dpctl/tests/helper/_helper.py index 5857b23195..23630f0d56 100644 --- a/dpctl/tests/helper/_helper.py +++ b/dpctl/tests/helper/_helper.py @@ -49,3 +49,13 @@ def get_queue_or_skip(args=()): except dpctl.SyclQueueCreationError: pytest.skip(f"Queue could not be created from {args}") return q + + +def get_sub_device_or_skip(): + for dev in dpctl.get_devices(): + for partition in ("next_partitionable", [1, 1]): + try: + return dev.create_sub_devices(partition=partition)[0] + except dpctl.SyclSubDeviceCreationError: + pass + pytest.skip("No device can be partitioned") diff --git a/dpctl/tests/test_sycl_device.py b/dpctl/tests/test_sycl_device.py index dd1a579272..809291e9df 100644 --- a/dpctl/tests/test_sycl_device.py +++ b/dpctl/tests/test_sycl_device.py @@ -16,12 +16,15 @@ """Defines unit test cases for the SyclDevice class.""" +import copy +import pickle + import pytest import dpctl from dpctl import SyclDeviceCreationError -from .helper import get_queue_or_skip +from .helper import get_queue_or_skip, get_sub_device_or_skip def test_standard_selectors(device_selector, check): @@ -455,3 +458,38 @@ def test_peer_access_value_keyword_validation(): bad_value = "wrong" with pytest.raises(ValueError): dev0.can_access_peer(dev1, value=bad_value) + + +def test_device_pickling(): + devs = dpctl.get_devices() + if not devs: + pytest.skip("No devices available") + for dev in devs: + dev2 = pickle.loads(pickle.dumps(dev)) + assert dev2 == dev + assert dev2.filter_string == dev.filter_string + + +def test_device_pickling_sub_device(): + sub_dev = get_sub_device_or_skip() + with pytest.raises(pickle.PicklingError): + pickle.dumps(sub_dev) + + +def test_device_unpickling_name_mismatch(): + from dpctl._sycl_device import _rebuild_device + + try: + dev = dpctl.SyclDevice() + except SyclDeviceCreationError: + pytest.skip("No default device available") + with pytest.raises(pickle.UnpicklingError): + _rebuild_device(dev.filter_string, dev.name + " (other)") + + +def test_device_copy(): + sub_dev = get_sub_device_or_skip() + root_dev = sub_dev.parent_device + for dev in (root_dev, sub_dev): + assert copy.copy(dev) is dev + assert copy.deepcopy(dev) is dev diff --git a/dpctl/tests/test_sycl_queue.py b/dpctl/tests/test_sycl_queue.py index af0fe71dcd..c39865dc70 100644 --- a/dpctl/tests/test_sycl_queue.py +++ b/dpctl/tests/test_sycl_queue.py @@ -16,14 +16,21 @@ """Defines unit test cases for the SyclQueue class.""" +import copy import ctypes +import pickle +import subprocess import sys import pytest import dpctl -from .helper import create_invalid_capsule +from .helper import ( + create_invalid_capsule, + get_queue_or_skip, + get_sub_device_or_skip, +) def test_standard_selectors(device_selector, check): @@ -402,3 +409,67 @@ def test_cython_api(dpctl_cython_extension): except dpctl.SyclDeviceCreationError: pytest.skip("Default-construction of SyclDevice failed") assert q.sycl_device == d + + +@pytest.mark.parametrize( + "props", + [ + (), + ("in_order",), + ("enable_profiling",), + ("in_order", "enable_profiling"), + ], +) +def test_queue_pickling(props): + try: + q = dpctl.SyclQueue(property=props or "default") + except dpctl.SyclQueueCreationError: + pytest.skip("Default-constructed queue could not be created") + q2 = pickle.loads(pickle.dumps(q)) + assert isinstance(q2, dpctl.SyclQueue) + assert q2.sycl_device == q.sycl_device + assert q2.sycl_context == q.sycl_context + assert q2.is_in_order == ("in_order" in props) + assert q2.has_enable_profiling == ("enable_profiling" in props) + + +def test_queue_pickling_non_default_context(): + q = get_queue_or_skip() + dev = q.sycl_device + ctx = dpctl.SyclContext([dev]) + if ctx == q.sycl_context: + pytest.skip("Context created from device list is the default context") + with pytest.raises(pickle.PicklingError): + pickle.dumps(dpctl.SyclQueue(ctx, dev)) + + +def test_queue_pickling_sub_device(): + sub_dev = get_sub_device_or_skip() + q = dpctl.SyclQueue(sub_dev) + with pytest.raises(pickle.PicklingError): + pickle.dumps(q) + + +def test_queue_unpickling_in_subprocess(): + q = get_queue_or_skip() + q = dpctl.SyclQueue(q.sycl_device, property="in_order") + script = ( + "import pickle, sys\n" + "q = pickle.loads(sys.stdin.buffer.read())\n" + "print(q.sycl_device.filter_string, q.is_in_order)\n" + ) + res = subprocess.run( + [sys.executable, "-c", script], + input=pickle.dumps(q), + capture_output=True, + check=True, + ) + assert res.stdout.decode().split() == [q.sycl_device.filter_string, "True"] + + +def test_queue_copy(): + q = get_queue_or_skip() + assert copy.copy(q) is q + assert copy.deepcopy(q) is q + d = copy.deepcopy({"queue": q}) + assert d["queue"] is q From eab831973dbcd4d2af6243ea45267300db218534 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Mon, 2 Nov 2026 09:47:06 -0800 Subject: [PATCH 2/6] Keep the queue when pickling USM memory MemoryUSM* objects now pickle their queue and unpickle onto it instead of the default device. Memory whose queue can not be pickled raises PicklingError. Pickles without a queue still load on the default device. --- dpctl/memory/_memory.pyx | 14 ++++++---- dpctl/tests/test_sycl_usm.py | 53 ++++++++++++++++++++++++++++++------ 2 files changed, 53 insertions(+), 14 deletions(-) diff --git a/dpctl/memory/_memory.pyx b/dpctl/memory/_memory.pyx index 6b2050bb9e..dd85c392c1 100644 --- a/dpctl/memory/_memory.pyx +++ b/dpctl/memory/_memory.pyx @@ -143,18 +143,18 @@ cdef void copy_via_host(void *dest_ptr, SyclQueue dest_queue, DPCTLEvent_Delete(E2Ref) -def _to_memory(unsigned char[::1] b, str usm_kind): +def _to_memory(unsigned char[::1] b, str usm_kind, SyclQueue queue=None): """ - Constructs Memory of the same size as the argument + Constructs Memory of the same size as the argument on ``queue`` and copies data into it""" cdef _Memory res if (usm_kind == "shared"): - res = MemoryUSMShared(len(b)) + res = MemoryUSMShared(len(b), queue=queue) elif (usm_kind == "device"): - res = MemoryUSMDevice(len(b)) + res = MemoryUSMDevice(len(b), queue=queue) elif (usm_kind == "host"): - res = MemoryUSMHost(len(b)) + res = MemoryUSMHost(len(b), queue=queue) else: raise ValueError( f"Unrecognized usm_kind={usm_kind} stored in the pickle" @@ -373,7 +373,9 @@ cdef class _Memory: return self.tobytes() def __reduce__(self): - return _to_memory, (self.copy_to_host(), self.get_usm_type()) + return ( + _to_memory, (self.copy_to_host(), self.get_usm_type(), self.queue) + ) @property def __sycl_usm_array_interface__(self): diff --git a/dpctl/tests/test_sycl_usm.py b/dpctl/tests/test_sycl_usm.py index 398a07b216..81b72a62f6 100644 --- a/dpctl/tests/test_sycl_usm.py +++ b/dpctl/tests/test_sycl_usm.py @@ -179,19 +179,56 @@ def test_pickling(memory_ctor): ), "Pickling/unpickling should be changing pointer" -def test_pickling_reconstructor_invalid_type(memory_ctor): +def test_pickling_reconstructor_invalid_type(): + from dpctl.memory._memory import _to_memory + + with pytest.raises(ValueError): + _to_memory(bytearray(8), "unknown") + + +def test_pickling_preserves_queue(memory_ctor): + import pickle + + devs = dpctl.get_devices() + if not devs: + pytest.skip("No SYCL devices available") + for dev in devs: + q = dpctl.SyclQueue(dev, property="in_order") + mobj = memory_ctor(64, queue=q) + mobj2 = pickle.loads(pickle.dumps(mobj)) + q2 = mobj2.sycl_queue + assert q2.sycl_device == dev + assert q2.sycl_context == q.sycl_context + assert q2.is_in_order + + +def test_pickling_non_default_context(memory_ctor): import pickle try: - mobj = memory_ctor(1024, alignment=64) + dev = dpctl.SyclDevice() except dpctl.SyclDeviceCreationError: pytest.skip("No SYCL devices available") - good_pickle_bytes = pickle.dumps(mobj) - usm_types = expected_usm_type_str(memory_ctor).encode("utf-8") - i = good_pickle_bytes.rfind(usm_types) - bad_pickle_bytes = good_pickle_bytes[:i] + b"u" + good_pickle_bytes[i + 1 :] - with pytest.raises(ValueError): - pickle.loads(bad_pickle_bytes) + ctx = dpctl.SyclContext([dev]) + if ctx == dpctl.SyclContext(dev): + pytest.skip("Context created from device list is the default context") + mobj = memory_ctor(64, queue=dpctl.SyclQueue(ctx, dev)) + with pytest.raises(pickle.PicklingError): + pickle.dumps(mobj) + + +def test_pickling_reconstructor_without_queue(memory_ctor): + from dpctl.memory._memory import _to_memory + + try: + dev = dpctl.SyclDevice() + except dpctl.SyclDeviceCreationError: + pytest.skip("No SYCL devices available") + usm_type = expected_usm_type_str(memory_ctor) + mobj = _to_memory(bytearray(b"abcdefgh"), usm_type) + assert type(mobj) is memory_ctor + assert mobj.sycl_device == dev + assert mobj.tobytes() == b"abcdefgh" @pytest.fixture(params=[MemoryUSMShared, MemoryUSMDevice, MemoryUSMHost]) From 071cc0d5f2bdb3e6044a1fb6c6c6afd5afb4ba0d Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Mon, 2 Nov 2026 09:48:25 -0800 Subject: [PATCH 3/6] Make SyclContext picklable This makes SyclQueue pickling also pickle the context. For non-picklable (i.e., non-default) contexts, the queue will throw a PicklingError --- CHANGELOG.md | 4 ++- dpctl/_sycl_context.pyx | 37 +++++++++++++++++++++++++++ dpctl/_sycl_queue.pyx | 32 ++++------------------- dpctl/tests/test_sycl_context.py | 44 +++++++++++++++++++++++++++++++- 4 files changed, 88 insertions(+), 29 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index e652bf82b3..baab3ca634 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,10 +8,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -* `dpctl.SyclDevice` and `dpctl.SyclQueue` can be pickled; root devices are restored by filter string, queues by device and properties into the device's default context. `copy.copy` and `copy.deepcopy` return the same object +* `dpctl.SyclDevice`, `dpctl.SyclContext`, and `dpctl.SyclQueue` can be pickled; root devices are restored by filter string, default contexts by one of their devices, and queues by context, device, and properties. `copy.copy` and `copy.deepcopy` return the same object ### Changed +* Pickled `dpctl.memory.MemoryUSMShared`, `MemoryUSMDevice`, and `MemoryUSMHost` objects keep their queue instead of being restored on the default device; pickling memory whose queue can not be pickled raises `pickle.PicklingError` + ### Deprecated ### Removed diff --git a/dpctl/_sycl_context.pyx b/dpctl/_sycl_context.pyx index f48b6d34c3..2160202150 100644 --- a/dpctl/_sycl_context.pyx +++ b/dpctl/_sycl_context.pyx @@ -23,6 +23,7 @@ """ import logging +import pickle from cpython cimport pycapsule from cpython.mem cimport PyMem_Free, PyMem_Malloc @@ -129,6 +130,13 @@ cdef class _SyclContext: DPCTLContext_Delete(self._ctxt_ref) +def _rebuild_context(SyclDevice dev): + """Unpickles a :class:`dpctl.SyclContext` as the default context of + ``dev``. + """ + return SyclContext(dev) + + cdef class SyclContext(_SyclContext): """ SyclContext(arg=None) @@ -408,6 +416,35 @@ cdef class SyclContext(_SyclContext): """ return DPCTLContext_Hash(self._ctxt_ref) + def __reduce__(self): + """Pickles the default context of a root device by that device. + + Raises: + pickle.PicklingError: if the context is not the default context + of its devices. + """ + cdef DPCTLSyclContextRef CRef = NULL + cdef bint is_default_ctx = False + cdef SyclDevice dev = self.get_devices()[0] + # sub-devices have no cached context + if dev.parent_device is None: + CRef = DPCTLDeviceMgr_GetCachedContext(dev.get_device_ref()) + if CRef is not NULL: + is_default_ctx = DPCTLContext_AreEq(CRef, self._ctxt_ref) + DPCTLContext_Delete(CRef) + if not is_default_ctx: + raise pickle.PicklingError( + f"Cannot pickle {self!r}: it is not the default context of " + "its devices" + ) + return _rebuild_context, (dev,) + + def __copy__(self): + return self + + def __deepcopy__(self, memo): + return self + cdef DPCTLSyclContextRef get_context_ref(self): return self._ctxt_ref diff --git a/dpctl/_sycl_queue.pyx b/dpctl/_sycl_queue.pyx index ce5cadaf0a..00dc741bdf 100644 --- a/dpctl/_sycl_queue.pyx +++ b/dpctl/_sycl_queue.pyx @@ -23,7 +23,6 @@ """ from ._backend cimport ( # noqa: E211 - DPCTLContext_AreEq, DPCTLContext_Create, DPCTLContext_Delete, DPCTLDefaultSelector_Create, @@ -104,7 +103,6 @@ from libc.stdlib cimport free, malloc import collections.abc import logging -import pickle import struct import sys @@ -777,11 +775,9 @@ cdef DPCTLSyclEventRef _memset_impl( return ERef -def _rebuild_queue(SyclDevice dev, tuple props): - """Unpickles a :class:`dpctl.SyclQueue` into the default context of - ``dev``. - """ - return SyclQueue(dev, property=props) +def _rebuild_queue(SyclContext ctx, SyclDevice dev, tuple props): + """Unpickles a :class:`dpctl.SyclQueue`.""" + return SyclQueue(ctx, dev, property=props) cdef class _SyclQueue: @@ -2200,36 +2196,18 @@ cdef class SyclQueue(_SyclQueue): return DPCTLQueue_Hash(self._queue_ref) def __reduce__(self): - """Pickles the queue's device and properties. The unpickled queue - uses the default context of the device. + """Pickles the queue's context, device, and properties. Raises: pickle.PicklingError: if the queue's context is not the default context of its device, or the device is not a root device. """ - cdef DPCTLSyclContextRef CRef = NULL - cdef bint is_default_ctx = False - # sub-devices have no cached context, SyclDevice.__reduce__ rejects them - if self._device.parent_device is None: - CRef = DPCTLDeviceMgr_GetCachedContext( - self._device.get_device_ref() - ) - if CRef is not NULL: - is_default_ctx = DPCTLContext_AreEq( - CRef, self._context.get_context_ref() - ) - DPCTLContext_Delete(CRef) - if not is_default_ctx: - raise pickle.PicklingError( - f"Cannot pickle {self!r}: its context is not the default " - "context of its device" - ) props = [] if DPCTLQueue_IsInOrder(self._queue_ref): props.append("in_order") if DPCTLQueue_HasEnableProfiling(self._queue_ref): props.append("enable_profiling") - return _rebuild_queue, (self._device, tuple(props)) + return _rebuild_queue, (self._context, self._device, tuple(props)) def __copy__(self): return self diff --git a/dpctl/tests/test_sycl_context.py b/dpctl/tests/test_sycl_context.py index 058bb7ad7f..ddf272ebec 100644 --- a/dpctl/tests/test_sycl_context.py +++ b/dpctl/tests/test_sycl_context.py @@ -16,11 +16,14 @@ """Defines unit test cases for the :class:`dpctl.SyclContext` class.""" +import copy +import pickle + import pytest import dpctl -from .helper import create_invalid_capsule +from .helper import create_invalid_capsule, get_sub_device_or_skip list_of_valid_filter_selectors = [ "opencl", @@ -341,3 +344,42 @@ def test_context_atomic_fence_scope_capabilities(valid_filter): assert all(isinstance(m, dpctl.memory_scope) for m in caps) for d in ctx.get_devices(): assert set(caps).issubset(set(d.atomic_fence_scope_capabilities)) + + +def test_context_pickling(): + devs = dpctl.get_devices() + if not devs: + pytest.skip("No devices available") + for dev in devs: + ctx = dpctl.SyclContext(dev) + ctx2 = pickle.loads(pickle.dumps(ctx)) + assert ctx2 == ctx + assert dev in ctx2.get_devices() + + +def test_context_pickling_non_default(): + try: + dev = dpctl.SyclDevice() + except dpctl.SyclDeviceCreationError: + pytest.skip("No default device available") + ctx = dpctl.SyclContext([dev]) + if ctx == dpctl.SyclContext(dev): + pytest.skip("Context created from device list is the default context") + with pytest.raises(pickle.PicklingError): + pickle.dumps(ctx) + + +def test_context_pickling_sub_device(): + sub_dev = get_sub_device_or_skip() + ctx = dpctl.SyclContext(sub_dev) + with pytest.raises(pickle.PicklingError): + pickle.dumps(ctx) + + +def test_context_copy(): + try: + ctx = dpctl.SyclContext() + except dpctl.SyclContextCreationError: + pytest.skip("Default context could not be created") + assert copy.copy(ctx) is ctx + assert copy.deepcopy(ctx) is ctx From bb1e5e1f566c97d21d15599537dad262c7945258 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Mon, 2 Nov 2026 09:56:25 -0800 Subject: [PATCH 4/6] Unpickle default queues onto the device's cached queue SyclQueue objects with default properties now unpickle as the cached queue of their device, so unpickled USM memory lands on one queue across separate loads. Queues with properties still unpickle as new queues. Document the pickling limitations on SyclQueue. --- CHANGELOG.md | 4 ++-- dpctl/_sycl_queue.pyx | 24 +++++++++++++++++++++++- dpctl/memory/_memory.pyx | 3 +++ dpctl/tests/test_sycl_queue.py | 13 +++++++++++++ dpctl/tests/test_sycl_usm.py | 16 ++++++++++++++++ 5 files changed, 57 insertions(+), 3 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index baab3ca634..2f1f261fcb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,11 +8,11 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added -* `dpctl.SyclDevice`, `dpctl.SyclContext`, and `dpctl.SyclQueue` can be pickled; root devices are restored by filter string, default contexts by one of their devices, and queues by context, device, and properties. `copy.copy` and `copy.deepcopy` return the same object +* `dpctl.SyclDevice`, `dpctl.SyclContext`, and `dpctl.SyclQueue` can be pickled. Root devices are restored by filter string, default contexts by one of their devices, and queues by context, device, and properties, with default-property queues restored as the cached queue of their device. `copy.copy` and `copy.deepcopy` return the same object ### Changed -* Pickled `dpctl.memory.MemoryUSMShared`, `MemoryUSMDevice`, and `MemoryUSMHost` objects keep their queue instead of being restored on the default device; pickling memory whose queue can not be pickled raises `pickle.PicklingError` +* Pickled `dpctl.memory.MemoryUSMShared`, `MemoryUSMDevice`, and `MemoryUSMHost` objects keep their queue instead of being restored on the default device. Pickling memory whose queue can not be pickled raises `pickle.PicklingError` ### Deprecated diff --git a/dpctl/_sycl_queue.pyx b/dpctl/_sycl_queue.pyx index 00dc741bdf..741d4525d6 100644 --- a/dpctl/_sycl_queue.pyx +++ b/dpctl/_sycl_queue.pyx @@ -776,7 +776,13 @@ cdef DPCTLSyclEventRef _memset_impl( def _rebuild_queue(SyclContext ctx, SyclDevice dev, tuple props): - """Unpickles a :class:`dpctl.SyclQueue`.""" + """Unpickles a :class:`dpctl.SyclQueue`. Queues with default properties + are restored as the cached queue of ``(ctx, dev)``. + """ + if not props: + # local import, _sycl_queue_manager cimports this module + from ._sycl_queue_manager import get_device_cached_queue + return get_device_cached_queue((ctx, dev)) return SyclQueue(ctx, dev, property=props) @@ -860,6 +866,22 @@ cdef class SyclQueue(_SyclQueue): to ensure one-time use. A new named capsule can be constructed by using :func:`dpctl.SyclQueue._get_capsule` method. + A queue can be pickled if its device is a root device and its context is + the default context of that device. Unpickling creates or looks up a + queue rather than restoring the original one: + + - a queue with default properties is restored as the cached queue of its + device, as returned by + :func:`dpctl._sycl_queue_manager.get_device_cached_queue`. Distinct + queues on the same device, and USM memory allocated on them, are all + restored onto that one queue. + - a queue with ``"in_order"`` or ``"enable_profiling"`` is restored as a + new queue on every unpickling. Only objects pickled in one + :func:`pickle.dumps` call that share a :class:`dpctl.SyclQueue` object + are restored onto one queue. + - the restored queue is equal to the original only if the original is the + cached queue. + Args: ctx (:class:`dpctl.SyclContext`, optional): Sycl context to create :class:`dpctl.SyclQueue` from. If not specified, a single-device diff --git a/dpctl/memory/_memory.pyx b/dpctl/memory/_memory.pyx index dd85c392c1..b53586752b 100644 --- a/dpctl/memory/_memory.pyx +++ b/dpctl/memory/_memory.pyx @@ -373,6 +373,9 @@ cdef class _Memory: return self.tobytes() def __reduce__(self): + """Pickles a host copy of the data and the queue. See + :class:`dpctl.SyclQueue` for how the queue is restored. + """ return ( _to_memory, (self.copy_to_host(), self.get_usm_type(), self.queue) ) diff --git a/dpctl/tests/test_sycl_queue.py b/dpctl/tests/test_sycl_queue.py index c39865dc70..6fa1546018 100644 --- a/dpctl/tests/test_sycl_queue.py +++ b/dpctl/tests/test_sycl_queue.py @@ -25,6 +25,7 @@ import pytest import dpctl +from dpctl._sycl_queue_manager import get_device_cached_queue from .helper import ( create_invalid_capsule, @@ -431,6 +432,18 @@ def test_queue_pickling(props): assert q2.sycl_context == q.sycl_context assert q2.is_in_order == ("in_order" in props) assert q2.has_enable_profiling == ("enable_profiling" in props) + assert (q2 == get_device_cached_queue(q.sycl_device)) == (not props) + + +def test_queue_unpickling_cached_queue(): + dev = get_queue_or_skip().sycl_device + cached = get_device_cached_queue(dev) + q0, q1 = dpctl.SyclQueue(dev), dpctl.SyclQueue(dev) + assert q0 != q1 + assert pickle.loads(pickle.dumps(q0)) == cached + assert pickle.loads(pickle.dumps(q1)) == cached + q2 = dpctl.SyclQueue(dev, property="in_order") + assert pickle.loads(pickle.dumps(q2)) != pickle.loads(pickle.dumps(q2)) def test_queue_pickling_non_default_context(): diff --git a/dpctl/tests/test_sycl_usm.py b/dpctl/tests/test_sycl_usm.py index 81b72a62f6..ef6f2a9f49 100644 --- a/dpctl/tests/test_sycl_usm.py +++ b/dpctl/tests/test_sycl_usm.py @@ -202,6 +202,22 @@ def test_pickling_preserves_queue(memory_ctor): assert q2.is_in_order +def test_pickling_restores_onto_cached_queue(memory_ctor): + import pickle + + from dpctl._sycl_queue_manager import get_device_cached_queue + + try: + dev = dpctl.SyclDevice() + except dpctl.SyclDeviceCreationError: + pytest.skip("No SYCL devices available") + m0 = memory_ctor(64, queue=dpctl.SyclQueue(dev)) + m1 = memory_ctor(64, queue=dpctl.SyclQueue(dev)) + m2 = pickle.loads(pickle.dumps(m0)) + m3 = pickle.loads(pickle.dumps(m1)) + assert m2.sycl_queue == m3.sycl_queue == get_device_cached_queue(dev) + + def test_pickling_non_default_context(memory_ctor): import pickle From 48501d791e42bf86013939419cb13cec8b7e3b52 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Tue, 3 Nov 2026 11:39:39 -0800 Subject: [PATCH 5/6] address pickling review feedback Raise UnpicklingError when the pickled device is not available, and name the queue in PicklingError for sub-device and non-default-context queues. Document that filter string ids depend on ONEAPI_DEVICE_SELECTOR. --- dpctl/_sycl_device.pyx | 16 +++++++++++++--- dpctl/_sycl_queue.pyx | 23 +++++++++++++++++++++-- dpctl/tests/test_sycl_device.py | 7 +++++++ dpctl/tests/test_sycl_queue.py | 4 ++-- dpctl/tests/test_sycl_usm.py | 2 +- 5 files changed, 44 insertions(+), 8 deletions(-) diff --git a/dpctl/_sycl_device.pyx b/dpctl/_sycl_device.pyx index 41e491d3aa..6cd14f7f02 100644 --- a/dpctl/_sycl_device.pyx +++ b/dpctl/_sycl_device.pyx @@ -376,10 +376,15 @@ def _rebuild_device(str filter_string, str name): """Unpickles a :class:`dpctl.SyclDevice`. Raises: - pickle.UnpicklingError: if the device selected by ``filter_string`` - is not named ``name``. + pickle.UnpicklingError: if no device matches ``filter_string``, or + the device it selects is not named ``name``. """ - dev = SyclDevice(filter_string) + try: + dev = SyclDevice(filter_string) + except SyclDeviceCreationError as e: + raise pickle.UnpicklingError( + f"Device '{filter_string}' ('{name}') is not available" + ) from e if dev.name != name: raise pickle.UnpicklingError( f"Device '{filter_string}' is '{dev.name}', expected '{name}'" @@ -1671,6 +1676,11 @@ cdef class SyclDevice(_SyclDevice): def __reduce__(self): """Pickles a root device by its filter string and name. + The filter string's device id is relative to the devices visible + under ``ONEAPI_DEVICE_SELECTOR``. If the loading process sees a + different set of devices, the pickle may load onto another device + of the same name. + Raises: pickle.PicklingError: if the device is not a root device. """ diff --git a/dpctl/_sycl_queue.pyx b/dpctl/_sycl_queue.pyx index 741d4525d6..b60b943984 100644 --- a/dpctl/_sycl_queue.pyx +++ b/dpctl/_sycl_queue.pyx @@ -23,6 +23,7 @@ """ from ._backend cimport ( # noqa: E211 + DPCTLContext_AreEq, DPCTLContext_Create, DPCTLContext_Delete, DPCTLDefaultSelector_Create, @@ -103,6 +104,7 @@ from libc.stdlib cimport free, malloc import collections.abc import logging +import pickle import struct import sys @@ -2221,9 +2223,26 @@ cdef class SyclQueue(_SyclQueue): """Pickles the queue's context, device, and properties. Raises: - pickle.PicklingError: if the queue's context is not the default - context of its device, or the device is not a root device. + pickle.PicklingError: if the queue's device is not a root device, + or its context is not the default context of that device. """ + cdef DPCTLSyclContextRef CRef = NULL + cdef bint is_default_ctx = False + if self._device.parent_device is not None: + raise pickle.PicklingError( + f"Cannot pickle {self!r}: its device is not a root device" + ) + CRef = DPCTLDeviceMgr_GetCachedContext(self._device.get_device_ref()) + if CRef is not NULL: + is_default_ctx = DPCTLContext_AreEq( + CRef, self._context.get_context_ref() + ) + DPCTLContext_Delete(CRef) + if not is_default_ctx: + raise pickle.PicklingError( + f"Cannot pickle {self!r}: its context is not the default " + "context of its device" + ) props = [] if DPCTLQueue_IsInOrder(self._queue_ref): props.append("in_order") diff --git a/dpctl/tests/test_sycl_device.py b/dpctl/tests/test_sycl_device.py index 809291e9df..c3441966f7 100644 --- a/dpctl/tests/test_sycl_device.py +++ b/dpctl/tests/test_sycl_device.py @@ -487,6 +487,13 @@ def test_device_unpickling_name_mismatch(): _rebuild_device(dev.filter_string, dev.name + " (other)") +def test_device_unpickling_missing_device(): + from dpctl._sycl_device import _rebuild_device + + with pytest.raises(pickle.UnpicklingError, match="is not available"): + _rebuild_device("opencl:gpu:999", "Missing device") + + def test_device_copy(): sub_dev = get_sub_device_or_skip() root_dev = sub_dev.parent_device diff --git a/dpctl/tests/test_sycl_queue.py b/dpctl/tests/test_sycl_queue.py index 6fa1546018..4d9e2fcabe 100644 --- a/dpctl/tests/test_sycl_queue.py +++ b/dpctl/tests/test_sycl_queue.py @@ -452,14 +452,14 @@ def test_queue_pickling_non_default_context(): ctx = dpctl.SyclContext([dev]) if ctx == q.sycl_context: pytest.skip("Context created from device list is the default context") - with pytest.raises(pickle.PicklingError): + with pytest.raises(pickle.PicklingError, match="SyclQueue.*context"): pickle.dumps(dpctl.SyclQueue(ctx, dev)) def test_queue_pickling_sub_device(): sub_dev = get_sub_device_or_skip() q = dpctl.SyclQueue(sub_dev) - with pytest.raises(pickle.PicklingError): + with pytest.raises(pickle.PicklingError, match="not a root device"): pickle.dumps(q) diff --git a/dpctl/tests/test_sycl_usm.py b/dpctl/tests/test_sycl_usm.py index ef6f2a9f49..68e2966ac5 100644 --- a/dpctl/tests/test_sycl_usm.py +++ b/dpctl/tests/test_sycl_usm.py @@ -229,7 +229,7 @@ def test_pickling_non_default_context(memory_ctor): if ctx == dpctl.SyclContext(dev): pytest.skip("Context created from device list is the default context") mobj = memory_ctor(64, queue=dpctl.SyclQueue(ctx, dev)) - with pytest.raises(pickle.PicklingError): + with pytest.raises(pickle.PicklingError, match="SyclQueue"): pickle.dumps(mobj) From 9acf3468b194c374d1fa59149f8cf8139ac9d5e0 Mon Sep 17 00:00:00 2001 From: Nikita Grigorian Date: Wed, 4 Nov 2026 11:25:26 -0800 Subject: [PATCH 6/6] skip devices without the USM aspect in test_pickling_preserves_queue --- dpctl/tests/test_sycl_usm.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/dpctl/tests/test_sycl_usm.py b/dpctl/tests/test_sycl_usm.py index 68e2966ac5..59417e7baf 100644 --- a/dpctl/tests/test_sycl_usm.py +++ b/dpctl/tests/test_sycl_usm.py @@ -189,9 +189,12 @@ def test_pickling_reconstructor_invalid_type(): def test_pickling_preserves_queue(memory_ctor): import pickle - devs = dpctl.get_devices() + aspect = f"usm_{expected_usm_type_str(memory_ctor)}_allocations" + devs = [ + d for d in dpctl.get_devices() if getattr(d, f"has_aspect_{aspect}") + ] if not devs: - pytest.skip("No SYCL devices available") + pytest.skip(f"No SYCL devices with {aspect} available") for dev in devs: q = dpctl.SyclQueue(dev, property="in_order") mobj = memory_ctor(64, queue=q)