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

Issue:
Stop blocking the Tokio runtime
If connecting to `quilc` or the eventual QPU over RPCQ fails, the SDK currently just spins indefinitely. A timeout should be set (and preferably configurable) so that the SDK gives up on executing and specifies where it failed (was it quilc or QPU connection).


Discussion / hints:
A better approach is to make this non-blocking (i.e., `tokio::task::spawn_blocking`) so that callers can set  their own timeouts by wrapping the future.

Fix patch:
diff --git a/qcs/src/executable.rs b/qcs/src/executable.rs
index 3eae3a223..f525778fc 100644
--- a/qcs/src/executable.rs
+++ b/qcs/src/executable.rs
@@ -2,6 +2,7 @@
 //! users will interact with QCS, quilc, and QVM.
 
 use std::collections::HashMap;
+use std::sync::Arc;
 use std::time::Duration;
 
 use eyre::{eyre, Report, Result, WrapErr};
@@ -53,17 +54,17 @@ use crate::{qpu, qvm, ExecutionResult};
 /// the methods most likely needs to live as long as this struct. Check individual methods for
 /// specifics. If only using `'static` strings then everything should just work.
 pub struct Executable<'executable, 'execution> {
-    quil: &'executable str,
+    quil: Arc<str>,
     shots: u16,
     readout_memory_region_names: Option<Vec<&'executable str>>,
-    params: Parameters<'executable>,
+    params: Parameters,
     compile_with_quilc: bool,
-    config: Option<Configuration>,
+    config: Option<Arc<Configuration>>,
     qpu: Option<qpu::Execution<'execution>>,
     qvm: Option<qvm::Execution>,
 }
 
-pub(crate) type Parameters<'a> = HashMap<&'a str, Vec<f64>>;
+pub(crate) type Parameters = HashMap<Box<str>, Vec<f64>>;
 
 impl<'executable> Executable<'executable, '_> {
     /// Create an [`Executable`] from a string containing a  [quil](https://github.com/quil-lang/quil)
@@ -81,9 +82,9 @@ impl<'executable> Executable<'executable, '_> {
     /// 1. `quil` is a string slice representing the original program to be run. The returned
     ///     [`Executable`] will only live as long as this reference.
     #[must_use]
-    pub fn from_quil(quil: &'executable str) -> Self {
+    pub fn from_quil<Quil: Into<Arc<str>>>(quil: Quil) -> Self {
         Self {
-            quil,
+            quil: quil.into(),
             shots: 1,
             readout_memory_region_names: None,
             params: Parameters::new(),
@@ -185,28 +186,24 @@ impl<'executable> Executable<'executable, '_> {
     /// ```
     ///
     /// [parametric compilation]: https://pyquil-docs.rigetti.com/en/stable/basics.html?highlight=parametric#parametric-compilation
-    pub fn with_parameter(
+    pub fn with_parameter<Param: Into<Box<str>>>(
         &mut self,
-        param_name: &'executable str,
+        param_name: Param,
         index: usize,
         value: f64,
     ) -> &mut Self {
-        let values = if let Some(values) = self.params.get_mut(param_name) {
-            values
-        } else {
-            self.params.insert(param_name, vec![0.0; index]);
-            if let Some(values) = self.params.get_mut(param_name) {
-                values
-            } else {
-                unreachable!("Set in the line above")
-            }
-        };
+        let param_name = param_name.into();
+        let mut values = self
+            .params
+            .remove(&param_name)
+            .unwrap_or_else(|| vec![0.0; index]);
 
         if index + 1 > values.len() {
             values.resize(index + 1, 0.0);
         }
 
         values[index] = value;
+        self.params.insert(param_name, values);
 
         self
     }
@@ -241,8 +238,8 @@ impl Executable<'_, '_> {
     ///
     /// # Warning
     ///
-    /// This function is `async` because of the HTTP client under the hood, but it will block your
-    /// thread waiting on the RPCQ-based functions.
+    /// This function uses [`tokio::task::spawn_blocking`] internally. See the docs for that function
+    /// to avoid blocking shutdown of the runtime.
     ///
     /// # Returns
     ///
@@ -263,31 +260,40 @@ impl Executable<'_, '_> {
     ///
     /// A number of errors could occur if `program` is malformed.
     pub async fn execute_on_qvm(&mut self) -> ExecuteResult {
-        let config = self.take_or_load_config().await;
+        let config = self.get_config().await;
         let mut qvm = if let Some(qvm) = self.qvm.take() {
             qvm
         } else {
-            qvm::Execution::new(self.quil)?
+            qvm::Execution::new(&self.quil)?
         };
         let result = qvm
             .run(self.shots, self.get_readouts(), &self.params, &config)
             .await;
         self.qvm = Some(qvm);
-        self.config = Some(config);
         result.map_err(Error::from)
     }
 
-    /// Remove and return `self.config` if set. Otherwise, load it from disk.
-    async fn take_or_load_config(&mut self) -> Configuration {
-        if let Some(config) = self.config.take() {
-            config
+    /// Load `self.config` if not yet loaded, then return a reference to it.
+    async fn get_config(&mut self) -> Arc<Configuration> {
+        if let Some(config) = &self.config {
+            config.clone()
         } else {
-            Configuration::load().await.unwrap_or_else(|e| {
-                log::error!("Got an error when loading config: {:#?}", e);
+            let config = Arc::new(Configuration::load().await.unwrap_or_else(|e| {
+                log::warn!("Got an error when loading config: {:#?}", e);
                 Configuration::default()
-            })
+            }));
+            self.config = Some(config.clone());
+            config
         }
     }
+
+    /// Refresh `self.config` and return it.
+    async fn refresh_config(&mut self) -> Result<Arc<Configuration>> {
+        let config = self.get_config().await.as_ref().clone();
+        let refreshed = Arc::new(config.refresh().await?);
+        self.config = Some(refreshed.clone());
+        Ok(refreshed)
+    }
 }
 
 impl<'execution> Executable<'_, 'execution> {
@@ -298,23 +304,33 @@ impl<'execution> Executable<'_, 'execution> {
                 return Ok(qpu);
             }
         }
-        let mut config = self.take_or_load_config().await;
-        let result =
-            match qpu::Execution::new(self.quil, self.shots, id, &config, self.compile_with_quilc)
+        let mut config = self.get_config().await;
+        let result = match qpu::Execution::new(
+            self.quil.clone(),
+            self.shots,
+            id,
+            config.clone(),
+            self.compile_with_quilc,
+        )
+        .await
+        {
+            Ok(result) => Ok(result),
+            Err(qpu::Error::Qcs { .. }) => {
+                config = self
+                    .refresh_config()
+                    .await
+                    .wrap_err("When refreshing authentication token")?;
+                qpu::Execution::new(
+                    self.quil.clone(),
+                    self.shots,
+                    id,
+                    config,
+                    self.compile_with_quilc,
+                )
                 .await
-            {
-                Ok(result) => Ok(result),
-                Err(qpu::Error::Qcs { .. }) => {
-                    config = config
-                        .refresh()
-                        .await
-                        .wrap_err("When refreshing authentication token")?;
-                    qpu::Execution::new(self.quil, self.shots, id, &config, self.compile_with_quilc)
-                        .await
-                }
-                err => err,
-            };
-        self.config = Some(config);
+            }
+            err => err,
+        };
         result.wrap_err_with(|| eyre!("When executing on {}", id))
     }
 
@@ -327,8 +343,8 @@ impl<'execution> Executable<'_, 'execution> {
     ///
     /// # Warning
     ///
-    /// This function is `async` because of the HTTP client under the hood, but it will block your
-    /// thread waiting on the RPCQ-based services.
+    /// This function uses [`tokio::task::spawn_blocking`] internally. See the docs for that function
+    /// to avoid blocking shutdown of the runtime.
     ///
     /// # Returns
     ///
@@ -346,21 +362,19 @@ impl<'execution> Executable<'_, 'execution> {
     /// [quilc]: https://github.com/quil-lang/quilc
     pub async fn execute_on_qpu(&mut self, quantum_processor_id: &'execution str) -> ExecuteResult {
         let mut qpu = self.qpu_for_id(quantum_processor_id).await?;
-        let mut config = self.take_or_load_config().await;
-
-        let readouts = self.get_readouts();
+        let mut config = self.get_config().await;
 
-        let response = match qpu.run(&self.params, readouts, &config).await {
+        let response = match qpu.run(&self.params, self.get_readouts(), &config).await {
             Ok(result) => Ok(result),
             Err(qpu::Error::Qcs {
                 retry_after: None, ..
             }) => {
                 // If retry_after is set, don't retry now
-                config = config
-                    .refresh()
+                config = self
+                    .refresh_config()
                     .await
                     .wrap_err("When refreshing authentication token")?;
-                qpu.run(&self.params, readouts, &config).await
+                qpu.run(&self.params, self.get_readouts(), &config).await
             }
             err => err,
         };
diff --git a/qcs/src/qpu/execution.rs b/qcs/src/qpu/execution.rs
index 87bc05516..118991576 100644
--- a/qcs/src/qpu/execution.rs
+++ b/qcs/src/qpu/execution.rs
@@ -2,11 +2,13 @@
 
 use std::collections::HashMap;
 use std::convert::TryFrom;
+use std::sync::Arc;
 use std::time::Duration;
 
 use eyre::{eyre, Report, Result, WrapErr};
 use log::trace;
 use quil_rs::expression::Expression;
+use tokio::task::spawn_blocking;
 
 use qcs_api::models::EngagementWithCredentials;
 
@@ -29,7 +31,7 @@ pub(crate) struct Execution<'a> {
     pub(crate) quantum_processor_id: &'a str,
     pub(crate) shots: u16,
     // All the stuff needed to actually make requests to QCS, lazily initialized
-    qcs: Option<Qcs>,
+    qcs: Option<Arc<Qcs>>,
 }
 
 #[derive(Debug, thiserror::Error)]
@@ -75,13 +77,13 @@ impl<'a> Execution<'a> {
     /// 1. Unable to rewrite the Native Quil for a QPU. This may mean that the `quil` was invalid
     ///     for the QPU or that there is a bug in this library.
     pub(crate) async fn new(
-        quil: &str,
+        quil: Arc<str>,
         shots: u16,
         quantum_processor_id: &'a str,
-        config: &Configuration,
+        config: Arc<Configuration>,
         compile_with_quilc: bool,
     ) -> Result<Execution<'a>, Error> {
-        let isa = get_isa(quantum_processor_id, config)
+        let isa = get_isa(quantum_processor_id, &config)
             .await
             .map_err(|err| Error::Qcs {
                 source: err.wrap_err("When getting ISA"),
@@ -90,11 +92,16 @@ impl<'a> Execution<'a> {
 
         let native_quil = if compile_with_quilc {
             trace!("Converting to Native Quil");
-            quilc::compile_program(quil, &isa, config)
-                .wrap_err("When attempting to compile your program to Native Quil")?
+            let thread_config = config.clone();
+            spawn_blocking(move || {
+                quilc::compile_program(&quil, &isa, &thread_config)
+                    .wrap_err("When attempting to compile your program to Native Quil")
+            })
+            .await
+            .map_err(|_| eyre!("Task running quilc did not complete."))??
         } else {
             trace!("Skipping conversion to Native Quil");
-            NativeQuil::assume_native_quil(String::from(quil))
+            NativeQuil::assume_native_quil(quil.to_string())
         };
 
         let program =
@@ -112,34 +119,41 @@ impl<'a> Execution<'a> {
     /// Run on a real QPU and wait for the results.
     pub(crate) async fn run(
         &mut self,
-        params: &Parameters<'_>,
+        params: &Parameters,
         readouts: &[&str],
         config: &Configuration,
     ) -> Result<HashMap<Box<str>, ExecutionResult>, Error> {
         let qcs = self.refresh_qcs(readouts, config).await?;
+        let qcs_for_thread = qcs.clone();
 
-        let result = self
+        let patch_values = self
             .get_substitutions(params)
-            .map_err(|e| Error::Quil(e.wrap_err("When setting provided parameters")))
-            .and_then(|patch_values| {
-                execute(&qcs.executable, &qcs.engagement, &patch_values).map_err(|e| Error::Qcs {
-                    source: e.wrap_err("When executing program."),
-                    retry_after: None,
-                })
-            })
-            .and_then(|buffers| {
-                process_buffers(buffers, &qcs.buffer_names)
-                    .wrap_err("When processing execution results")
-                    .map_err(Error::from)
+            .map_err(|e| Error::Quil(e.wrap_err("When setting provided parameters")))?;
+
+        let buffers = spawn_blocking(move || {
+            execute(
+                &qcs_for_thread.executable,
+                &qcs_for_thread.engagement,
+                &patch_values,
+            )
+            .map_err(|e| Error::Qcs {
+                source: e.wrap_err("While executing"),
+                retry_after: None,
             })
-            .and_then(|registers| {
-                ExecutionResult::try_from_registers(registers, self.shots)
-                    .wrap_err("When decoding execution results")
-                    .map_err(Error::from)
-            });
-
-        self.qcs = Some(qcs);
-        result
+        })
+        .await
+        .map_err(|_| Error::Qcs {
+            source: eyre!("Execution thread did not complete."),
+            retry_after: None,
+        })??;
+
+        let registers = process_buffers(buffers, &qcs.buffer_names)
+            .wrap_err("When processing execution results")
+            .map_err(Error::from)?;
+
+        ExecutionResult::try_from_registers(registers, self.shots)
+            .wrap_err("When decoding execution results")
+            .map_err(Error::from)
     }
 
     /// Take or create a [`Qcs`] for this [`Execution`]. This fetches / updates engagements, builds
@@ -149,9 +163,9 @@ impl<'a> Execution<'a> {
         &mut self,
         readouts: &[&str],
         config: &Configuration,
-    ) -> Result<Qcs, Error> {
-        if let Some(qcs) = self.qcs.take() {
-            return Ok(qcs);
+    ) -> Result<Arc<Qcs>, Error> {
+        if let Some(qcs) = &self.qcs {
+            return Ok(qcs.clone());
         }
 
         let response = build_executable(
@@ -183,11 +197,13 @@ impl<'a> Execution<'a> {
                     retry_after: None,
                 },
             })?;
-        Ok(Qcs {
+        let qcs = Arc::new(Qcs {
             buffer_names,
             engagement,
             executable: response.program,
-        })
+        });
+        self.qcs = Some(qcs.clone());
+        Ok(qcs)
     }
 
     /// Take the user-provided map of [`Parameters`] and produce the map of substitutions which
@@ -219,18 +235,19 @@ impl<'a> Execution<'a> {
     /// Because QPUs do not evaluate expressions themselves. This function creates the values for
     /// `__SUBST` by calculating the original expressions given the user-provided params (in this
     /// case just `theta`).
-    fn get_substitutions<'params>(
-        &self,
-        params: &'params Parameters,
-    ) -> Result<Parameters<'params>> {
-        let mut patch_values = params.clone();
+    fn get_substitutions(&self, params: &Parameters) -> Result<Parameters> {
+        // Convert into the format that quil-rs expects.
+        let params: HashMap<&str, Vec<f64>> = params
+            .iter()
+            .map(|(key, value)| (key.as_ref(), value.clone()))
+            .collect();
         let values = self
             .program
             .substitutions
             .iter()
             .map(|substitution: &Expression| {
                 substitution
-                    .evaluate(&HashMap::new(), params)
+                    .evaluate(&HashMap::new(), &params)
                     .map_err(|_| eyre!("Could not evaluate expression {}", substitution))
                     .and_then(|complex| {
                         if complex.im == 0.0 {
@@ -243,7 +260,12 @@ impl<'a> Execution<'a> {
                     })
             })
             .collect::<Result<Vec<f64>>>()?;
-        patch_values.insert(SUBSTITUTION_NAME, values);
+        // Convert back to the format that this library expects
+        let mut patch_values: Parameters = params
+            .into_iter()
+            .map(|(key, value)| (key.into(), value))
+            .collect();
+        patch_values.insert(SUBSTITUTION_NAME.into(), values);
         Ok(patch_values)
     }
 }
diff --git a/qcs/src/qpu/quilc/isa/qubit.rs b/qcs/src/qpu/quilc/isa/qubit.rs
index 0b2511912..c5cb1e914 100644
--- a/qcs/src/qpu/quilc/isa/qubit.rs
+++ b/qcs/src/qpu/quilc/isa/qubit.rs
@@ -256,41 +256,12 @@ fn rz_gates(node_id: i32) -> Vec<Operator> {
 
 #[cfg(test)]
 mod describe_rz_gates {
-    use qcs_api::models::OperationSite;
-
-    use super::*;
+    use super::{rz_gates, Arguments, Operator, Parameters};
 
     /// This data is copied from the pyQuil ISA integration test.
     #[test]
     fn it_passes_the_pyquil_aspen_8_test() {
         let node_id = 1;
-        let frb_sim_1q = Operation {
-            characteristics: vec![],
-            name: "randomized_benchmark_simultaneous_1q".to_string(),
-            node_count: Some(30),
-            parameters: vec![],
-            sites: vec![OperationSite {
-                characteristics: vec![
-                    Characteristic {
-                        name: "fRB".to_string(),
-                        value: 0.989821537688075,
-                        error: Some(0.000699235456806402),
-                        node_ids: Some(vec![0]),
-                        parameter_values: None,
-                        timestamp: "1970-01-01T00:00:00+00:00".to_string(),
-                    },
-                    Characteristic {
-                        name: "fRB".to_string(),
-                        value: 0.996832638579018,
-                        error: Some(0.00010089678215399),
-                        node_ids: Some(vec![1]),
-                        timestamp: "1970-01-01T00:00:00+00:00".to_string(),
-                        parameter_values: None,
-                    },
-                ],
-                node_ids: vec![0, 1],
-            }],
-        };
         let gates = rz_gates(node_id);
         let expected = vec![Operator::Gate {
             arguments: Arguments::Int(1),
diff --git a/qcs/src/qpu/runner.rs b/qcs/src/qpu/runner.rs
index 78612d5a9..f5ad00832 100644
--- a/qcs/src/qpu/runner.rs
+++ b/qcs/src/qpu/runner.rs
@@ -60,7 +60,7 @@ struct QPUParams<'request> {
 }
 
 impl<'request> QPUParams<'request> {
-    fn from_program(program: &'request str, patch_values: PatchValues<'request>) -> Self {
+    fn from_program(program: &'request str, patch_values: &'request Parameters) -> Self {
         Self {
             request: QPURequest::from_program(program, patch_values),
             priority: 1,
@@ -74,18 +74,16 @@ impl<'params> From<&'params QPUParams<'params>> for RPCRequest<'params, QPUParam
     }
 }
 
-type PatchValues<'request> = &'request HashMap<&'request str, Vec<f64>>;
-
 #[derive(Serialize, Debug)]
 #[serde(tag = "_type")]
 struct QPURequest<'request> {
     id: String,
     program: &'request str,
-    patch_values: PatchValues<'request>,
+    patch_values: &'request Parameters,
 }
 
 impl<'request> QPURequest<'request> {
-    fn from_program(program: &'request str, patch_values: PatchValues<'request>) -> Self {
+    fn from_program(program: &'request str, patch_values: &'request Parameters) -> Self {
         Self {
             id: uuid::Uuid::new_v4().to_string(),
             program,
diff --git a/qcs/src/qvm/execution.rs b/qcs/src/qvm/execution.rs
index 841fbf0ff..ad77959bc 100644
--- a/qcs/src/qvm/execution.rs
+++ b/qcs/src/qvm/execution.rs
@@ -58,7 +58,7 @@ impl Execution {
         &mut self,
         shots: u16,
         readouts: &[&str],
-        params: &Parameters<'_>,
+        params: &Parameters,
         config: &Configuration,
     ) -> Result<HashMap<Box<str>, ExecutionResult>> {
         if shots == 0 {
@@ -69,7 +69,7 @@ impl Execution {
 
         let mut instruction_count = 0;
         for (name, values) in params {
-            match memory.get(*name) {
+            match memory.get(name.as_ref()) {
                 Some(region) => {
                     if region.size.length != values.len() as u64 {
                         return Err(eyre!(
@@ -89,7 +89,7 @@ impl Execution {
                     0,
                     Instruction::Move(Move {
                         destination: ArithmeticOperand::MemoryReference(MemoryReference {
-                            name: String::from(*name),
+                            name: name.to_string(),
                             index: index as u64,
                         }),
                         source: ArithmeticOperand::LiteralReal(*value),
@@ -139,16 +139,14 @@ impl Execution {
 
 #[cfg(test)]
 mod describe_execution {
-    use std::collections::HashMap;
-
-    use super::*;
+    use super::{Configuration, Execution, Parameters};
 
     #[tokio::test]
     async fn it_errs_on_excess_parameters() {
         let mut exe = Execution::new("DECLARE ro BIT").unwrap();
 
-        let mut params = HashMap::new();
-        params.insert("doesnt_exist", vec![0.0]);
+        let mut params = Parameters::new();
+        params.insert("doesnt_exist".into(), vec![0.0]);
 
         let result = exe.run(1, &[], &params, &Configuration::default()).await;
         if let Err(e) = result {
@@ -162,15 +160,15 @@ mod describe_execution {
     async fn it_errors_when_any_param_is_the_wrong_size() {
         let mut exe = Execution::new("DECLARE ro BIT[2]").unwrap();
 
-        let mut params = HashMap::new();
-        params.insert("ro", vec![0.0]);
+        let mut params = Parameters::new();
+        params.insert("ro".into(), vec![0.0]);
 
         let result = exe.run(1, &[], &params, &Configuration::default()).await;
         if let Err(e) = result {
             let err_string = e.to_string();
             assert!(err_string.contains("ro"));
-            assert!(err_string.contains("1"));
-            assert!(err_string.contains("2"));
+            assert!(err_string.contains('1'));
+            assert!(err_string.contains('2'));
         } else {
             panic!("Expected an error but got none.");
         }


Test patch:
diff --git a/qcs/tests/mocked_qpu.rs b/qcs/tests/mocked_qpu.rs
index 523fd63cc..b3d04af6e 100644
--- a/qcs/tests/mocked_qpu.rs
+++ b/qcs/tests/mocked_qpu.rs
@@ -43,6 +43,7 @@ async fn setup() {
     tokio::spawn(mock_qcs::run());
 }
 
+#[allow(dead_code)]
 mod auth_server {
     use serde::{Deserialize, Serialize};
     use warp::Filter;
@@ -74,6 +75,7 @@ mod auth_server {
     }
 }
 
+#[allow(dead_code)]
 mod mock_qcs {
     use serde::{Deserialize, Serialize};
     use warp::Filter;
@@ -157,6 +159,7 @@ mod mock_qcs {
     }
 }
 
+#[allow(dead_code)]
 mod qpu {
     use std::collections::HashMap;