Added pickle support for lightning.gpu device's DevPool attribute

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.37.0-dev46"
+__version__ = "0.37.0-dev47"

pennylane_lightning/core/src/simulators/lightning_gpu/bindings/LGPUBindings.hpp

@@ -384,9 +384,20 @@ void registerBackendSpecificInfo(py::module_ &m) {
         .def("acquireDevice", &DevicePool<int>::acquireDevice)
         .def("releaseDevice", &DevicePool<int>::releaseDevice)
         .def("syncDevice", &DevicePool<int>::syncDevice)
+        .def("refresh", &DevicePool<int>::refresh)
         .def_static("getTotalDevices", &DevicePool<int>::getTotalDevices)
         .def_static("getDeviceUIDs", &DevicePool<int>::getDeviceUIDs)
-        .def_static("setDeviceID", &DevicePool<int>::setDeviceIdx);
+        .def_static("setDeviceID", &DevicePool<int>::setDeviceIdx)
+        .def(py::pickle(
+            [](const DevicePool<int> &self) { // __getstate__
+                return py::make_tuple();
+            },
+            [](py::tuple &t) { // __setstate__
+                if (t.size() != 0) {
+                    throw std::runtime_error("Invalid state!");
+                }
+                return DevicePool<int>{};
+            }));
 
     py::class_<DevTag<int>>(m, "DevTag")
         .def(py::init<>())

pennylane_lightning/core/src/utils/cuda_utils/DevicePool.hpp

@@ -28,7 +28,19 @@ namespace Pennylane::LightningGPU {
  */
 template <typename DeviceIndexType = int> class DevicePool {
   public:
-    DevicePool() {
+    DevicePool() { refresh(); }
+
+    DevicePool(DevicePool &&dp) { refresh(); }
+
+    /**
+     * @brief Dynamically refresh the available pool devices.
+     *
+     */
+    void refresh() {
+        std::lock_guard<std::mutex> lg(m_);
+        available_devices_.clear();
+        active_devices_.clear();
+
         for (std::size_t i = 0; i < DevicePool::getTotalDevices(); i++) {
             available_devices_.push(static_cast<DeviceIndexType>(i));
         }

pennylane_lightning/core/src/utils/cuda_utils/TSQueue.hpp

@@ -92,4 +92,12 @@ template <typename T> class TSQueue {
         std::lock_guard<std::mutex> lk(m);
         return q.empty();
     }
+    /**
+     * @brief Clear the queue of all items. Thread safe.
+     *
+     */
+    void clear() {
+        std::lock_guard<std::mutex> lk(m);
+        q = {};
+    }
 };

tests/test_device.py

@@ -14,6 +14,7 @@
 """
 Unit tests for Lightning devices creation.
 """
+import pickle as pkl
 import sys
 
 import numpy as np
@@ -67,6 +68,27 @@ def test_create_device_with_unsupported_mpi_buf_size():
         pass
 
 
+@pytest.mark.skipif(
+    device_name != "lightning.gpu",
+    reason="Check if the method is pickleable throught the cpp layer",
+)
+def test_devpool_is_pickleable():
+    dev = qml.device(device_name, wires=2)
+    try:
+        pickled_devpool = pkl.dumps(dev._dp)
+        un_pickled_devpool = pkl.loads(pickled_devpool)
+
+        from pennylane_lightning.lightning_gpu_ops import DevPool
+
+        d = DevPool()
+
+        assert isinstance(un_pickled_devpool, DevPool)
+        assert un_pickled_devpool.getTotalDevices() == d.getTotalDevices()
+
+    except Exception:
+        pytest.fail("DevPool should be Pickleable")
+
+
 @pytest.mark.skipif(
     (device_name == "lightning.kokkos" and sys.platform == "win32"),
     reason="lightning.kokkos doesn't support 0 wires on Windows.",