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

Issue:
Return execution time in execution result
Implement https://github.com/rigetti/pyquil/pull/1436 in this SDK. This will be a breaking change because the `ExecuteResult` type returned from execution methods currently only contains a `HashMap` of the register name to its data—we need to start returning structs so we can add additional metadata.


Discussion / hints:

Fix patch:
diff --git a/qcs/examples/parametric_compilation.rs b/qcs/examples/parametric_compilation.rs
index 8abbe0386..e6bef8270 100644
--- a/qcs/examples/parametric_compilation.rs
+++ b/qcs/examples/parametric_compilation.rs
@@ -2,6 +2,7 @@
 //! [pyquil]: https://pyquil-docs.rigetti.com/en/stable/basics.html?highlight=parametric#parametric-compilation
 
 use std::f64::consts::PI;
+use std::time::Duration;
 
 use qcs::Executable;
 
@@ -22,19 +23,24 @@ async fn main() {
     let mut parametric_measurements = Vec::with_capacity(200);
 
     let step = 2.0 * PI / 50.0;
+    let mut total_execution_time = Duration::new(0, 0);
 
     for i in 0..=50 {
         let theta = step * f64::from(i);
-        let result = exe
+        let mut data = exe
             .with_parameter("theta", 0, theta)
             .execute_on_qpu("Aspen-11")
             .await
-            .expect("Failed to execute")
-            .remove("ro")
-            .expect("Missing ro register");
+            .expect("Failed to execute");
+        total_execution_time += data
+            .duration
+            .expect("Aspen-11 should always report duration");
+        let result = data.registers.remove("ro").expect("Missing ro register");
         parametric_measurements.append(&mut result.into_i16().unwrap()[0])
     }
 
+    println!("Total execution time: {:?}", total_execution_time);
+
     for measurement in parametric_measurements {
         if measurement == 1 {
             // We didn't run with all 0 so parametrization worked!
diff --git a/qcs/examples/quil_t.rs b/qcs/examples/quil_t.rs
index b2640d426..32855f6e1 100644
--- a/qcs/examples/quil_t.rs
+++ b/qcs/examples/quil_t.rs
@@ -22,6 +22,7 @@ async fn main() {
         .execute_on_qpu("Aspen-11")
         .await
         .expect("Failed to execute")
+        .registers
         .remove("ro")
         .expect("Missing ro register");
 
diff --git a/qcs/src/executable.rs b/qcs/src/executable.rs
index 460fa9bd3..a66aa4499 100644
--- a/qcs/src/executable.rs
+++ b/qcs/src/executable.rs
@@ -7,14 +7,14 @@ use std::time::Duration;
 
 use crate::configuration::{Configuration, LoadError, RefreshError};
 use crate::qpu::ExecutionError;
-use crate::{qpu, qvm, ExecutionResult};
+use crate::{qpu, qvm, ExecutionData};
 
 /// The builder interface for executing Quil programs on QVMs and QPUs.
 ///
 /// # Example
 ///
 /// ```rust
-/// use qcs::{Executable, ExecutionResult};
+/// use qcs::{Executable, RegisterData};
 ///
 ///
 /// const PROGRAM: &str = r##"
@@ -32,7 +32,7 @@ use crate::{qpu, qvm, ExecutionResult};
 ///     let mut result = Executable::from_quil(PROGRAM).with_shots(4).execute_on_qvm().await.unwrap();
 ///     // We know it's i8 because we declared the memory as `BIT` in Quil.
 ///     // "ro" is the only source read from by default if you don't specify a .read_from()
-///     let data = result.remove("ro").expect("Did not receive ro data").into_i8().unwrap();
+///     let data = result.registers.remove("ro").expect("Did not receive ro data").into_i8().unwrap();
 ///     // In this case, we ran the program for 4 shots, so we know the length is 4.
 ///     assert_eq!(data.len(), 4);
 ///     for shot in data {
@@ -127,11 +127,13 @@ impl<'executable> Executable<'executable, '_> {
     ///         .await
     ///         .unwrap();
     ///     let first = result
+    ///         .registers
     ///         .remove("first")
     ///         .expect("Did not receive first buffer")
     ///         .into_f64()
     ///         .expect("Received incorrect data type for first");
     ///     let second = result
+    ///         .registers
     ///         .remove("second")
     ///         .expect("Did not receive second buffer")
     ///         .into_f64()
@@ -177,7 +179,7 @@ impl<'executable> Executable<'executable, '_> {
     ///             .with_parameter("theta", 0, theta)
     ///             .with_parameter("theta", 1, theta * 2.0)
     ///             .execute_on_qvm().await.unwrap();
-    ///         let data = result.remove("theta").expect("Could not read theta").into_f64().unwrap();
+    ///         let data = result.registers.remove("theta").expect("Could not read theta").into_f64().unwrap();
     ///         assert_eq!(data[0][0], theta);
     ///         assert_eq!(data[0][1], theta * 2.0);
     ///     }
@@ -208,7 +210,7 @@ impl<'executable> Executable<'executable, '_> {
     }
 }
 
-type ExecuteResult = Result<HashMap<Box<str>, ExecutionResult>, Error>;
+type ExecuteResult = Result<ExecutionData, Error>;
 
 impl Executable<'_, '_> {
     /// Specify a number of times to run the program for each execution. Defaults to 1 run or "shot".
@@ -258,7 +260,10 @@ impl Executable<'_, '_> {
             .run(self.shots, self.get_readouts(), &self.params, &config)
             .await;
         self.qvm = Some(qvm);
-        result.map_err(Error::from)
+        result.map_err(Error::from).map(|registers| ExecutionData {
+            registers,
+            duration: None,
+        })
     }
 
     /// Load `self.config` if not yet loaded, then return a reference to it.
diff --git a/qcs/src/execution_data.rs b/qcs/src/execution_data.rs
new file mode 100644
index 000000000..6f034bed3
--- /dev/null
+++ b/qcs/src/execution_data.rs
@@ -0,0 +1,19 @@
+use std::collections::HashMap;
+use std::time::Duration;
+
+use crate::RegisterData;
+
+/// The result of executing an [`Executable`][crate::Executable] via
+/// [`Executable::execute_on_qvm`][crate::Executable::execute_on_qvm] or
+/// [`Executable::execute_on_qpu`][crate::Executable::execute_on_qpu].
+pub struct ExecutionData {
+    /// The data of all registers that were read from
+    /// (via [`Executable::read_from`][crate::Executable::read_from]). Key is the name of the
+    /// register, value is the data of the register after execution.
+    pub registers: HashMap<Box<str>, RegisterData>,
+    /// The time it took to execute the program on the QPU, not including any network or queueing
+    /// time. If paying for on-demand execution, this is the amount you will be billed for.
+    ///
+    /// This will always be `None` for QVM execution.
+    pub duration: Option<Duration>,
+}
diff --git a/qcs/src/lib.rs b/qcs/src/lib.rs
index 9428257b1..b4ba86c4a 100644
--- a/qcs/src/lib.rs
+++ b/qcs/src/lib.rs
@@ -10,10 +10,12 @@
 //! using [`Executable`].
 
 pub use executable::{Error, Executable, Service};
-pub use execution_result::ExecutionResult;
+pub use execution_data::ExecutionData;
+pub use register_data::RegisterData;
 
 pub mod configuration;
 mod executable;
-mod execution_result;
+mod execution_data;
 mod qpu;
 mod qvm;
+mod register_data;
diff --git a/qcs/src/qpu/execution.rs b/qcs/src/qpu/execution.rs
index 6507afbf8..ed14093a8 100644
--- a/qcs/src/qpu/execution.rs
+++ b/qcs/src/qpu/execution.rs
@@ -13,7 +13,7 @@ use qcs_api::models::EngagementWithCredentials;
 
 use crate::configuration::Configuration;
 use crate::executable::Parameters;
-use crate::{qpu, ExecutionResult};
+use crate::{ExecutionData, RegisterData};
 
 use super::quilc::{self, NativeQuil, NativeQuilProgram};
 use super::rewrite_arithmetic::{RewrittenProgram, SUBSTITUTION_NAME};
@@ -110,8 +110,8 @@ pub(crate) enum Unexpected {
     Qcs(String),
 }
 
-impl From<qpu::IsaError> for Error {
-    fn from(source: qpu::IsaError) -> Self {
+impl From<IsaError> for Error {
+    fn from(source: IsaError) -> Self {
         match source {
             IsaError::QpuNotFound => Self::QpuNotFound,
             IsaError::Unauthorized => Self::Unauthorized,
@@ -219,13 +219,13 @@ impl<'a> Execution<'a> {
         params: &Parameters,
         readouts: &[&str],
         config: &Configuration,
-    ) -> Result<HashMap<Box<str>, ExecutionResult>, Error> {
+    ) -> Result<ExecutionData, Error> {
         let qcs = self.refresh_qcs(readouts, config).await?;
         let qcs_for_thread = qcs.clone();
 
         let patch_values = self.get_substitutions(params).map_err(Error::Quil)?;
 
-        let buffers = spawn_blocking(move || {
+        let response = spawn_blocking(move || {
             execute(
                 &qcs_for_thread.executable,
                 &qcs_for_thread.engagement,
@@ -239,9 +239,14 @@ impl<'a> Execution<'a> {
             source,
         })??;
 
-        let registers = process_buffers(buffers, &qcs.buffer_names)?;
-
-        ExecutionResult::try_from_registers(registers, self.shots).map_err(Error::from)
+        let registers = process_buffers(response.buffers, &qcs.buffer_names)?;
+        let register_data = RegisterData::try_from_registers(registers, self.shots)?;
+        Ok(ExecutionData {
+            registers: register_data,
+            duration: response
+                .execution_duration_microseconds
+                .map(Duration::from_micros),
+        })
     }
 
     /// Take or create a [`Qcs`] for this [`Execution`]. This fetches / updates engagements, builds
diff --git a/qcs/src/qpu/runner.rs b/qcs/src/qpu/runner.rs
index 24a0ea72e..c444785cb 100644
--- a/qcs/src/qpu/runner.rs
+++ b/qcs/src/qpu/runner.rs
@@ -1,3 +1,5 @@
+//! Interface to the QPU's API.
+
 use std::collections::HashMap;
 use std::convert::{TryFrom, TryInto};
 use std::fmt::{Display, Formatter};
@@ -19,7 +21,7 @@ pub(crate) fn execute(
     program: &str,
     engagement: &EngagementWithCredentials,
     patch_values: &Parameters,
-) -> Result<HashMap<String, Buffer>, Error> {
+) -> Result<GetExecutionResultsResponse, Error> {
     let params = QPUParams::from_program(program, patch_values);
     let EngagementWithCredentials {
         address,
@@ -44,7 +46,7 @@ pub(crate) fn execute(
     let request = RPCRequest::from(&params);
     let job_id: String = client.run_request(&request)?;
     debug!("Received job ID {} from QPU", &job_id);
-    let get_buffers_request = GetBuffersRequest::new(job_id);
+    let get_buffers_request = GetExecutionResultsRequest::new(job_id);
     client
         .run_request(&RPCRequest::from(&get_buffers_request))
         .map_err(Error::from)
@@ -115,23 +117,32 @@ impl<'request> QPURequest<'request> {
 }
 
 #[derive(Serialize, Debug)]
-struct GetBuffersRequest {
+struct GetExecutionResultsRequest {
     job_id: String,
     wait: bool,
 }
 
-impl GetBuffersRequest {
+impl GetExecutionResultsRequest {
     fn new(job_id: String) -> Self {
         Self { job_id, wait: true }
     }
 }
 
-impl<'request> From<&'request GetBuffersRequest> for RPCRequest<'request, GetBuffersRequest> {
-    fn from(req: &'request GetBuffersRequest) -> Self {
-        RPCRequest::new("get_buffers", req)
+impl<'request> From<&'request GetExecutionResultsRequest>
+    for RPCRequest<'request, GetExecutionResultsRequest>
+{
+    fn from(req: &'request GetExecutionResultsRequest) -> Self {
+        RPCRequest::new("get_execution_results", req)
     }
 }
 
+#[derive(Debug, Deserialize)]
+pub(crate) struct GetExecutionResultsResponse {
+    pub(crate) buffers: HashMap<String, Buffer>,
+    #[serde(default)]
+    pub(crate) execution_duration_microseconds: Option<u64>,
+}
+
 /// The raw form of the data which comes back from an execution.
 ///
 /// Generally this should not be used directly, but converted into an appropriate
diff --git a/qcs/src/qvm/execution.rs b/qcs/src/qvm/execution.rs
index 1d4d9af5f..a16581cf7 100644
--- a/qcs/src/qvm/execution.rs
+++ b/qcs/src/qvm/execution.rs
@@ -8,7 +8,7 @@ use quil_rs::{
 
 use crate::configuration::Configuration;
 use crate::executable::Parameters;
-use crate::ExecutionResult;
+use crate::RegisterData;
 
 use super::{Request, Response};
 
@@ -59,7 +59,7 @@ impl Execution {
         readouts: &[&str],
         params: &Parameters,
         config: &Configuration,
-    ) -> Result<HashMap<Box<str>, ExecutionResult>, Error> {
+    ) -> Result<HashMap<Box<str>, RegisterData>, Error> {
         if shots == 0 {
             return Err(Error::ShotsMustBePositive);
         }
@@ -108,7 +108,7 @@ impl Execution {
         shots: u16,
         readouts: &[&str],
         config: &Configuration,
-    ) -> Result<HashMap<Box<str>, ExecutionResult>, Error> {
+    ) -> Result<HashMap<Box<str>, RegisterData>, Error> {
         let request = Request::new(&self.program.to_string(true), shots, readouts);
 
         let client = reqwest::Client::new();
diff --git a/qcs/src/qvm/mod.rs b/qcs/src/qvm/mod.rs
index f4fa3c798..e43d04e67 100644
--- a/qcs/src/qvm/mod.rs
+++ b/qcs/src/qvm/mod.rs
@@ -7,7 +7,7 @@ use serde::{Deserialize, Serialize};
 
 pub(crate) use execution::{Error, Execution};
 
-use crate::ExecutionResult;
+use crate::RegisterData;
 
 mod execution;
 
@@ -21,7 +21,7 @@ pub(super) enum Response {
 #[derive(Debug, Deserialize)]
 pub(super) struct Success {
     #[serde(flatten)]
-    pub registers: HashMap<String, ExecutionResult>,
+    pub registers: HashMap<String, RegisterData>,
 }
 
 #[derive(Debug, Deserialize)]
diff --git a/qcs/src/execution_result.rs b/qcs/src/register_data.rs
similarity index 85%
rename from qcs/src/execution_result.rs
rename to qcs/src/register_data.rs
index e9e0e806a..c4c3f8ad4 100644
--- a/qcs/src/execution_result.rs
+++ b/qcs/src/register_data.rs
@@ -24,7 +24,7 @@ use crate::qpu::{DecodeError, Register};
 /// convert any variant type to its inner data.
 #[derive(Debug, Deserialize, EnumAsInner, PartialEq, Serialize)]
 #[serde(untagged)]
-pub enum ExecutionResult {
+pub enum RegisterData {
     /// Corresponds to the Quil `BIT` or `OCTET` types.
     I8(Vec<Vec<i8>>),
     /// Corresponds to the Quil `REAL` type.
@@ -36,7 +36,7 @@ pub enum ExecutionResult {
     Complex32(Vec<Vec<Complex32>>),
 }
 
-impl ExecutionResult {
+impl RegisterData {
     pub(crate) fn try_from_registers(
         registers_by_name: HashMap<Box<str>, Vec<Register>>,
         shots: u16,
@@ -135,9 +135,9 @@ mod describe_program_results {
         let registers = hashmap! {
             Box::from(String::from("ro")) => vec![Register::I8(vec![1, 2, 3]), Register::I8(vec![4, 5, 6])]
         };
-        let results = ExecutionResult::try_from_registers(registers, 3).unwrap();
+        let results = RegisterData::try_from_registers(registers, 3).unwrap();
         let expected = hashmap! {
-            Box::from(String::from("ro")) => ExecutionResult::I8(vec![vec![1, 4], vec![2, 5], vec![3, 6]])
+            Box::from(String::from("ro")) => RegisterData::I8(vec![vec![1, 4], vec![2, 5], vec![3, 6]])
         };
         assert_eq!(results, expected);
     }
@@ -147,9 +147,9 @@ mod describe_program_results {
         let registers = hashmap! {
             Box::from(String::from("ro")) => vec![Register::I16(vec![1, 2, 3]), Register::I16(vec![4, 5, 6])]
         };
-        let results = ExecutionResult::try_from_registers(registers, 3).unwrap();
+        let results = RegisterData::try_from_registers(registers, 3).unwrap();
         let expected = hashmap! {
-            Box::from(String::from("ro")) => ExecutionResult::I16(vec![vec![1, 4], vec![2, 5], vec![3, 6]])
+            Box::from(String::from("ro")) => RegisterData::I16(vec![vec![1, 4], vec![2, 5], vec![3, 6]])
         };
         assert_eq!(results, expected);
     }
@@ -162,9 +162,9 @@ mod describe_program_results {
                 Register::F64(vec![4.0, 5.0, 6.0]),
             ]
         };
-        let results = ExecutionResult::try_from_registers(registers, 3).unwrap();
+        let results = RegisterData::try_from_registers(registers, 3).unwrap();
         let expected = hashmap! {
-            Box::from(String::from("ro")) => ExecutionResult::F64(vec![vec![1.0, 4.0], vec![2.0, 5.0], vec![3.0, 6.0]])
+            Box::from(String::from("ro")) => RegisterData::F64(vec![vec![1.0, 4.0], vec![2.0, 5.0], vec![3.0, 6.0]])
         };
         assert_eq!(results, expected);
     }
@@ -185,9 +185,9 @@ mod describe_program_results {
                 ]),
             ]
         };
-        let results = ExecutionResult::try_from_registers(registers, 3).unwrap();
+        let results = RegisterData::try_from_registers(registers, 3).unwrap();
         let expected = hashmap! {
-            Box::from(String::from("ro")) => ExecutionResult::Complex32(vec![
+            Box::from(String::from("ro")) => RegisterData::Complex32(vec![
                 vec![Complex32::from(1.0), Complex32::from(4.0)],
                 vec![Complex32::from(2.0), Complex32::from(5.0)],
                 vec![Complex32::from(3.0), Complex32::from(6.0)],
@@ -201,7 +201,7 @@ mod describe_program_results {
         let registers = hashmap! {
             Box::from(String::from("ro")) => vec![Register::I8(vec![1, 2, 3]), Register::I16(vec![4, 5, 6])]
         };
-        let results = ExecutionResult::try_from_registers(registers, 3);
+        let results = RegisterData::try_from_registers(registers, 3);
         assert!(results.is_err());
     }
 
@@ -211,10 +211,10 @@ mod describe_program_results {
             Box::from(String::from("first")) => vec![Register::I8(vec![1, 2, 3]), Register::I8(vec![4, 5, 6])],
             Box::from(String::from("second")) => vec![Register::I16(vec![1, 2, 3]), Register::I16(vec![4, 5, 6])],
         };
-        let results = ExecutionResult::try_from_registers(registers, 3).unwrap();
+        let results = RegisterData::try_from_registers(registers, 3).unwrap();
         let expected = hashmap! {
-            Box::from(String::from("first")) => ExecutionResult::I8(vec![vec![1, 4], vec![2, 5], vec![3, 6]]),
-            Box::from(String::from("second")) => ExecutionResult::I16(vec![vec![1, 4], vec![2, 5], vec![3, 6]]),
+            Box::from(String::from("first")) => RegisterData::I8(vec![vec![1, 4], vec![2, 5], vec![3, 6]]),
+            Box::from(String::from("second")) => RegisterData::I16(vec![vec![1, 4], vec![2, 5], vec![3, 6]]),
         };
         assert_eq!(results, expected);
     }


Test patch:
diff --git a/qcs/tests/basic_qvm.rs b/qcs/tests/basic_qvm.rs
index 38bc73107..59728a3c1 100644
--- a/qcs/tests/basic_qvm.rs
+++ b/qcs/tests/basic_qvm.rs
@@ -27,11 +27,13 @@ async fn test_bell_state() {
         .expect("Could not run on QVM");
 
     let first: Vec<Vec<i8>> = data
+        .registers
         .remove("first")
         .expect("Missing first buffer")
         .into_i8()
         .expect("Produced wrong data type");
     let second: Vec<Vec<i8>> = data
+        .registers
         .remove("second")
         .expect("Missing second buffer")
         .into_i8()
diff --git a/qcs/tests/mocked_qpu.rs b/qcs/tests/mocked_qpu.rs
index 8c7ea99e4..a3c22b82e 100644
--- a/qcs/tests/mocked_qpu.rs
+++ b/qcs/tests/mocked_qpu.rs
@@ -2,11 +2,12 @@
 //! the `qcs` crate.
 
 use std::thread;
+use std::time::Duration;
 
 use maplit::hashmap;
 
 use qcs::configuration::{SECRETS_PATH_VAR, SETTINGS_PATH_VAR};
-use qcs::{Executable, ExecutionResult};
+use qcs::{Executable, RegisterData};
 
 const BELL_STATE: &str = r#"
 DECLARE ro BIT[2]
@@ -29,9 +30,10 @@ async fn successful_bell_state() {
         .await
         .expect("Failed to run program that should be successful");
     assert_eq!(
-        result,
-        hashmap! {Box::from(String::from("ro")) => ExecutionResult::I8(vec![vec![0, 0], vec![1, 1]])}
+        result.registers,
+        hashmap! {Box::from(String::from("ro")) => RegisterData::I8(vec![vec![0, 0], vec![1, 1]])}
     );
+    assert_eq!(result.duration, Some(Duration::from_micros(8675)));
 }
 
 async fn setup() {
@@ -183,7 +185,7 @@ mod qpu {
     #[serde(untagged)]
     enum Params {
         QPUParams { request: QPURequest, priority: u8 },
-        GetBuffersRequest { job_id: String, wait: bool },
+        GetExecutionResultsRequest { job_id: String, wait: bool },
     }
 
     #[derive(Deserialize, Debug)]
@@ -209,6 +211,12 @@ mod qpu {
         dtype: String,
     }
 
+    #[derive(Serialize, Debug, PartialEq)]
+    pub struct GetExecutionResultsResponse {
+        buffers: HashMap<&'static str, Buffer>,
+        execution_duration_microseconds: u64,
+    }
+
     pub(crate) fn run() {
         let ctx = zmq::Context::new();
 
@@ -238,9 +246,9 @@ mod qpu {
                             id: qpu_request.id,
                             result: "1",
                         }),
-                        Params::GetBuffersRequest { .. } => {
-                            let mut resp = HashMap::new();
-                            resp.insert(
+                        Params::GetExecutionResultsRequest { .. } => {
+                            let mut buffers = HashMap::new();
+                            buffers.insert(
                                 "q0",
                                 Buffer {
                                     shape: (2,),
@@ -248,7 +256,7 @@ mod qpu {
                                     dtype: "int8".to_string(),
                                 },
                             );
-                            resp.insert(
+                            buffers.insert(
                                 "q1",
                                 Buffer {
                                     shape: (2,),
@@ -256,6 +264,10 @@ mod qpu {
                                     dtype: "int8".to_string(),
                                 },
                             );
+                            let resp = GetExecutionResultsResponse {
+                                buffers,
+                                execution_duration_microseconds: 8675,
+                            };
                             Box::new(RPCResponse::RPCReply {
                                 id: qpu_request.id,
                                 result: resp,
diff --git a/qcs/tests/parametric_compilation.rs b/qcs/tests/parametric_compilation.rs
index 50be17893..49685283f 100644
--- a/qcs/tests/parametric_compilation.rs
+++ b/qcs/tests/parametric_compilation.rs
@@ -32,6 +32,7 @@ async fn basic_substitution() {
             .expect("Failed to execute");
         parametric_measurements.append(
             &mut result
+                .registers
                 .remove("ro")
                 .expect("Missing ro register")
                 .into_i8()