Repository: rigetti/qcs-sdk-rust (language: python)

Issue:
Make ExecutionData pickleable
It can be convenient to serialize/pickle `ExecutionData` to provide some means of persisting job results. 


Discussion / hints:

Fix patch:
diff --git a/Cargo.lock b/Cargo.lock
index f9cc2c8d0..5040e4151 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -1523,9 +1523,9 @@ dependencies = [
 
 [[package]]
 name = "mio"
-version = "0.8.9"
+version = "0.8.11"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "3dce281c5e46beae905d4de1870d8b1509a9142b62eedf18b443b011ca8343d0"
+checksum = "a4a650543ca06a924e8b371db273b2756685faae30f8487da1b56505a8f78b0c"
 dependencies = [
  "libc",
  "wasi",
@@ -1957,9 +1957,9 @@ dependencies = [
 
 [[package]]
 name = "proc-macro2"
-version = "1.0.70"
+version = "1.0.78"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "39278fbbf5fb4f646ce651690877f89d1c5811a3d4acb27700c1cb3cdb78fd3b"
+checksum = "e2422ad645d89c99f8f3e6b88a9fdeca7fabeac836b1002371c4367c8f984aae"
 dependencies = [
  "unicode-ident",
 ]
@@ -2293,9 +2293,9 @@ dependencies = [
 
 [[package]]
 name = "quote"
-version = "1.0.33"
+version = "1.0.35"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "5267fca4496028628a95160fc423a33e8b2e6af8a5302579e322e4b520293cae"
+checksum = "291ec9ab5efd934aaf503a6466c5d5251535d108ee747472c3977cc5acc868ef"
 dependencies = [
  "proc-macro2",
 ]
@@ -3167,9 +3167,9 @@ checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20"
 
 [[package]]
 name = "tokio"
-version = "1.34.0"
+version = "1.36.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
-checksum = "d0c014766411e834f7af5b8f4cf46257aab4036ca95e9d2c144a10f59ad6f5b9"
+checksum = "61285f6515fa018fb2d1e46eb21223fff441ee8db5d0f1435e8ab4f5cdb80931"
 dependencies = [
  "backtrace",
  "bytes",
diff --git a/Cargo.toml b/Cargo.toml
index d80df0cde..d647abb55 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -9,7 +9,7 @@ qcs-api-client-grpc = "0.7.13"
 qcs-api-client-openapi = "0.8.11"
 serde_json = "1.0.86"
 thiserror = "1.0.37"
-tokio = "1.34.0"
+tokio = "1.36.0"
 # We specify quil-rs as a git and versioned dependency so that we can keep the version of
 # quil-rs used in both the Rust and Python packages in sync. The tag used should always
 # be a `quil-py` tag and should match the version used in `crates/python/pyproject.toml`
diff --git a/crates/lib/src/execution_data.rs b/crates/lib/src/execution_data.rs
index ce7d1713f..a2480a931 100644
--- a/crates/lib/src/execution_data.rs
+++ b/crates/lib/src/execution_data.rs
@@ -48,7 +48,7 @@ use crate::{
 /// knowing the intent of the program that was run. Instead, it's recommended to build the
 /// [`RegisterMatrix`] you need from the inner [`QpuResultData`] data using the knowledge of your
 /// program to choose the correct readout values for each shot.
-#[derive(Debug, Clone, PartialEq, EnumAsInner)]
+#[derive(Debug, Clone, PartialEq, EnumAsInner, Deserialize, Serialize)]
 pub enum ResultData {
     /// Data returned from the QVM, stored as [`QvmResultData`]
     Qvm(QvmResultData),
@@ -57,7 +57,7 @@ pub enum ResultData {
 }
 
 /// The result of executing an [`Executable`](crate::Executable)
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
 pub struct ExecutionData {
     /// The [`ResultData`] that was read from the [`Executable`](crate::Executable).
     pub result_data: ResultData,
diff --git a/crates/lib/src/qpu/result_data.rs b/crates/lib/src/qpu/result_data.rs
index b1806af3c..980077621 100644
--- a/crates/lib/src/qpu/result_data.rs
+++ b/crates/lib/src/qpu/result_data.rs
@@ -3,6 +3,7 @@
 use enum_as_inner::EnumAsInner;
 use num::complex::Complex64;
 use quil_rs::instruction::MemoryReference;
+use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 
 use qcs_api_client_grpc::models::controller::{
@@ -12,7 +13,7 @@ use qcs_api_client_grpc::models::controller::{
 
 /// A row of readout values from the QPU. Each row contains all the values emitted to a
 /// memory reference across all shots.
-#[derive(Debug, Clone, EnumAsInner, PartialEq)]
+#[derive(Debug, Clone, EnumAsInner, PartialEq, Deserialize, Serialize)]
 pub enum ReadoutValues {
     /// Integer readout values
     Integer(Vec<i64>),
@@ -23,7 +24,7 @@ pub enum ReadoutValues {
 }
 
 /// A row of data containing the contents of each memory region at the end of a job.
-#[derive(Debug, Clone, EnumAsInner, PartialEq)]
+#[derive(Debug, Clone, EnumAsInner, PartialEq, Deserialize, Serialize)]
 pub enum MemoryValues {
     /// Values that correspond to a memory region declared with the BIT or OCTET data type.
     Binary(Vec<u8>),
@@ -35,7 +36,7 @@ pub enum MemoryValues {
 
 /// This struct encapsulates data returned from the QPU after executing a job.
 #[allow(clippy::module_name_repetitions)]
-#[derive(Debug, Clone, PartialEq)]
+#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
 pub struct QpuResultData {
     pub(crate) mappings: HashMap<String, String>,
     pub(crate) readout_values: HashMap<String, ReadoutValues>,
diff --git a/crates/lib/src/qvm/mod.rs b/crates/lib/src/qvm/mod.rs
index 042d2a373..c01a2ca0c 100644
--- a/crates/lib/src/qvm/mod.rs
+++ b/crates/lib/src/qvm/mod.rs
@@ -9,7 +9,7 @@ use quil_rs::{
     quil::{Quil, ToQuilError},
     Program,
 };
-use serde::Deserialize;
+use serde::{Deserialize, Serialize};
 
 pub(crate) use execution::Execution;
 
@@ -105,7 +105,7 @@ impl<T: Client + Sync + Send> Client for Arc<T> {
 
 /// Encapsulates data returned after running a program on the QVM
 #[allow(clippy::module_name_repetitions)]
-#[derive(Debug, Deserialize, Clone, PartialEq)]
+#[derive(Debug, Deserialize, Serialize, Clone, PartialEq)]
 pub struct QvmResultData {
     pub(crate) memory: HashMap<String, RegisterData>,
 }
diff --git a/crates/python/src/compiler/quilc.rs b/crates/python/src/compiler/quilc.rs
index 995dd9777..8c4013b37 100644
--- a/crates/python/src/compiler/quilc.rs
+++ b/crates/python/src/compiler/quilc.rs
@@ -255,6 +255,7 @@ impl PyNativeQuilMetadata {
             qpu_runtime_estimation,
         })
     }
+
     pub fn __getstate__<'a>(&self, py: Python<'a>) -> PyResult<&'a PyBytes> {
         Ok(PyBytes::new(
             py,
diff --git a/crates/python/src/execution_data.rs b/crates/python/src/execution_data.rs
index 3277ec411..f095f8219 100644
--- a/crates/python/src/execution_data.rs
+++ b/crates/python/src/execution_data.rs
@@ -1,11 +1,16 @@
+use std::collections::HashMap;
 use std::time::Duration;
 
 use numpy::{Complex64, PyArray2};
-use pyo3::exceptions::PyKeyError;
+use pyo3::exceptions::{PyKeyError, PyRuntimeError};
+use pyo3::pyclass::CompareOp;
 use pyo3::{
-    exceptions::PyValueError, pyclass, pymethods, types::PyDelta, IntoPy, Py, PyObject, PyRef,
-    PyRefMut, PyResult, Python, ToPyObject,
+    exceptions::PyValueError,
+    pyclass, pymethods,
+    types::{PyBytes, PyDelta},
+    IntoPy, Py, PyObject, PyRef, PyRefMut, PyResult, Python, ToPyObject,
 };
+use qcs::qvm::QvmResultData;
 use qcs::{ExecutionData, RegisterMap, RegisterMatrix, ResultData};
 use rigetti_pyo3::{
     impl_repr, py_wrap_data_struct, py_wrap_error, py_wrap_type, py_wrap_union_enum, wrap_error,
@@ -16,6 +21,7 @@ use crate::qpu::PyQpuResultData;
 use crate::qvm::PyQvmResultData;
 
 py_wrap_union_enum! {
+    #[pyo3(module = "qcs_sdk")]
     PyResultData(ResultData) as "ResultData" {
         qpu: Qpu => PyQpuResultData,
         qvm: Qvm => PyQvmResultData
@@ -23,6 +29,14 @@ py_wrap_union_enum! {
 }
 impl_repr!(PyResultData);
 
+impl Default for PyResultData {
+    fn default() -> Self {
+        PyResultData(ResultData::Qvm(QvmResultData::from_memory_map(
+            HashMap::new(),
+        )))
+    }
+}
+
 wrap_error!(RustRegisterMatrixConversionError(
     qcs::RegisterMatrixConversionError
 ));
@@ -56,6 +70,8 @@ impl PyResultData {
 }
 
 py_wrap_data_struct! {
+    #[pyo3(module = "qcs_sdk")]
+    #[derive(Debug, PartialEq)]
     PyExecutionData(ExecutionData) as "ExecutionData" {
         result_data: ResultData => PyResultData,
         duration: Option<Duration> => Option<Py<PyDelta>>
@@ -68,11 +84,13 @@ impl PyExecutionData {
     #[new]
     fn __new__(
         py: Python<'_>,
-        result_data: PyResultData,
+        // result_data is optional here because pickling an object requires calling __new__ without
+        // any arguments.
+        result_data: Option<PyResultData>,
         duration: Option<Py<PyDelta>>,
     ) -> PyResult<Self> {
         Ok(Self(ExecutionData {
-            result_data: ResultData::py_try_from(py, &result_data)?,
+            result_data: ResultData::py_try_from(py, &result_data.unwrap_or_default())?,
             duration: duration
                 .map(|delta| {
                     delta
@@ -84,6 +102,29 @@ impl PyExecutionData {
                 .transpose()?,
         }))
     }
+
+    pub fn __getstate__<'a>(&self, py: Python<'a>) -> PyResult<&'a PyBytes> {
+        Ok(PyBytes::new(
+            py,
+            &serde_json::to_vec(self.as_inner())
+                .map_err(|e| PyRuntimeError::new_err(format!("failed to serialize: {e}")))?,
+        ))
+    }
+
+    pub fn __setstate__<'a>(&mut self, state: &PyBytes) -> PyResult<()> {
+        let execution_data: ExecutionData = serde_json::from_slice(state.as_bytes())
+            .map_err(|e| PyRuntimeError::new_err(format!("failed to deserialize: {e}")))?;
+        *self = PyExecutionData(execution_data);
+        Ok(())
+    }
+
+    fn __richcmp__(&self, other: &Self, op: CompareOp, py: Python<'_>) -> PyObject {
+        match op {
+            CompareOp::Eq => (self == other).into_py(py),
+            CompareOp::Ne => (self != other).into_py(py),
+            _ => py.NotImplemented(),
+        }
+    }
 }
 
 py_wrap_type! {
diff --git a/deny.toml b/deny.toml
index 20d2a2355..a4332dd9e 100644
--- a/deny.toml
+++ b/deny.toml
@@ -28,8 +28,8 @@ ignore = [
     "RUSTSEC-2023-0052",  # Introduced by transitive dependency `webpki`.
                           # `hyper-proxy`, then `qcs-api-client-rust` need to update in order to remove
                           # `webpki`.
+    "RUSTSEC-2024-0006",   # Introduced by bindgen/shlex with no valid upgrade path.
     "RUSTSEC-2023-0055",  # there isn't a newer version of lexical with a fix
-    "RUSTSEC-2024-0006",
 ]
 
 # This section is considered when running `cargo deny check licenses`
@@ -101,7 +101,7 @@ skip-tree = [
   { name = "toml", version = "*", depth = 20 },
   { name = "warp", version = "*", depth = 20 },              # Only used for development
   { name = "hermit-abi", version = "*", depth = 20 },        # Only used for development
-  { name = "hyper-rustls", version = "<=0.23", depth = 20 }, # `hyper-proxy` relies on an older version than `reqwest`
+  { name = "hyper-rustls", version = "<=0.23", depth = 20 }, # `hyper-proxy` relies on an older version than `rqwest`
   { name = "itertools", version = "*", depth = 20 },         # `quil-rs` and various transitive dependencies depend on different versions
 ]
 


Test patch:
diff --git a/crates/python/tests/execution_data/test_execution_data.py b/crates/python/tests/execution_data/test_execution_data.py
index 5e8426717..466aca21b 100644
--- a/crates/python/tests/execution_data/test_execution_data.py
+++ b/crates/python/tests/execution_data/test_execution_data.py
@@ -1,9 +1,13 @@
-from qcs_sdk.qpu import ReadoutValues, MemoryValues, QPUResultData
-from qcs_sdk.qvm import QVMResultData
-from qcs_sdk import ResultData, RegisterData, RegisterMatrix
+from datetime import timedelta
+import pickle
+
+import pytest
 import numpy as np
 from numpy.testing import assert_array_equal
-import pytest
+
+from qcs_sdk.qpu import ReadoutValues, MemoryValues, QPUResultData
+from qcs_sdk.qvm import QVMResultData
+from qcs_sdk import ResultData, RegisterData, RegisterMatrix, ExecutionData
 
 
 class TestResultData:
@@ -108,3 +112,23 @@ def test_iter(self):
             assert np.all(matrix.to_ndarray() == np.matrix([[0, 1, 2], [1, 2, 3]]))
 
         assert expected_keys == actual_keys == set(register_map.keys())
+
+
+class TestExecutionData:
+    def test_pickle(self):
+        """Test that ExecutionData can be pickled and unpickled."""
+        qpu_result_data = QPUResultData(
+            mappings={"q0": "ro"},
+            readout_values={"q0": ReadoutValues([0, 1])},
+            memory_values={"m": MemoryValues([0, 1])},
+        )
+        qvm_result_data = QVMResultData.from_memory_map({"ro": RegisterData([[0, 1]])})
+
+        for execution_data in [
+            ExecutionData(result_data=ResultData(qpu_result_data), duration=timedelta(seconds=1)),
+            ExecutionData(result_data=ResultData(qvm_result_data), duration=None),
+        ]:
+            print(execution_data.result_data)
+            pickled = pickle.dumps(execution_data)
+            unpickled = pickle.loads(pickled)
+            assert execution_data == unpickled