fix: ensure otel context propagated across ffi boundary (#544)

* test: support shared global rust tracing in test_tracing_subscriber.py
* test: add quilc and qvm unit tests in test_tracing_subscriber.py
* feat: update pyo3 opentelemetry dependencies

---------

Co-authored-by: Graham Smith <[email protected]>

Author: erichulburd

Cargo.lock

@@ -3,9 +3,9 @@ members = ["crates/*"]
 resolver = "2"
 
 [workspace.dependencies]
-qcs-api-client-common = "0.11.9"
-qcs-api-client-grpc = "0.11.9"
-qcs-api-client-openapi = "0.12.9"
+qcs-api-client-common = "0.12.0"
+qcs-api-client-grpc = "0.12.0"
+qcs-api-client-openapi = "0.13.0"
 serde_json = "1.0.86"
 thiserror = "1.0.57"
 tokio = "1.36.0"
@@ -24,11 +24,13 @@ quil-rs = { git = "https://github.com/rigetti/quil-rs", tag = "quil-py/v0.15.2"
 # and need to be updated together.
 ndarray = { version = "0.15.6", features = ["serde"] }
 numpy = "0.20.0"
+opentelemetry = { version = "0.27.1" }
+opentelemetry_sdk = { version = "0.27.1" }
 pyo3 = { version = "0.20.0", default-features = false}
 pyo3-asyncio = { version = "0.20", features = ["tokio-runtime"] }
 pyo3-log = { version = "0.8.2" }
-pyo3-opentelemetry = { version = "0.3" }
-pyo3-tracing-subscriber = { version = "0.1.4", default-features = false }
+pyo3-opentelemetry = "0.4.0" 
+pyo3-tracing-subscriber = { version = "0.2.0", default-features = false }
 
 pyo3-build-config = "0.20.0"
 rigetti-pyo3 = { version = "0.4.1", default-features = false, features = ["complex", "time"] }

Cargo.toml

@@ -13,7 +13,6 @@ readme = "./README.md"
 manual-tests = []
 tracing = ["dep:tracing", "qcs-api-client-common/tracing", "qcs-api-client-grpc/tracing", "qcs-api-client-openapi/tracing"]
 tracing-config = ["tracing", "qcs-api-client-common/tracing-config", "qcs-api-client-grpc/tracing-config", "qcs-api-client-openapi/tracing-config"]
-otel-tracing = ["tracing-config", "qcs-api-client-grpc/otel-tracing", "qcs-api-client-openapi/otel-tracing"]
 libquil = ["dep:libquil-sys"]
 grpc-web = ["qcs-api-client-grpc/grpc-web"]
 tracing-opentelemetry = ["tracing-config", "qcs-api-client-grpc/tracing-opentelemetry", "qcs-api-client-openapi/tracing-opentelemetry"]
@@ -27,9 +26,10 @@ indexmap = "2.2.6"
 lazy_static = "1.4.0"
 ndarray.workspace = true
 num = { version = "0.4.0", features = ["serde"] }
-opentelemetry = { version = "0.23.0" }
-opentelemetry_sdk = { version = "0.23.0" }
+opentelemetry = { workspace = true }
+opentelemetry_sdk = { workspace = true }
 pbjson-types = "0.7.0"
+pyo3 = { workspace = true }
 qcs-api-client-common.workspace = true
 qcs-api-client-openapi.workspace = true
 qcs-api-client-grpc.workspace = true

crates/lib/Cargo.toml

@@ -1,6 +1,6 @@
 [env]
 RUST_BACKTRACE = 0
-CARGO_MAKE_CARGO_BUILD_TEST_FLAGS = "--features otel-tracing,libquil"  # Disable --all-features to avoid manual tests in CI for now
+CARGO_MAKE_CARGO_BUILD_TEST_FLAGS = "--features tracing-opentelemetry,libquil"  # Disable --all-features to avoid manual tests in CI for now
 
 [tasks.pre-test]
 command = "docker"

crates/lib/Makefile.toml

@@ -8,7 +8,7 @@ RX(pi) 0
 MEASURE 0 ro[0]
 "#;
 
-const QUANTUM_PROCESSOR_ID: &str = "Aspen-M-3";
+const QUANTUM_PROCESSOR_ID: &str = "Ankaa-3";
 
 async fn quilc_client() -> rpcq::Client {
     let qcs = Qcs::load();

crates/lib/examples/delayed_job_retrieval.rs

@@ -5,6 +5,10 @@
 use std::time::Duration;
 
 use qcs_api_client_common::configuration::{ClientConfiguration, TokenError};
+#[cfg(feature = "tracing")]
+use qcs_api_client_grpc::tonic::wrap_channel_with_tracing;
+#[cfg(feature = "tracing")]
+use qcs_api_client_grpc::tonic::CustomTraceService;
 #[cfg(feature = "grpc-web")]
 use qcs_api_client_grpc::tonic::{wrap_channel_with_grpc_web, GrpcWebWrapperLayerService};
 use qcs_api_client_grpc::{
@@ -15,6 +19,7 @@ use qcs_api_client_grpc::{
     },
 };
 use qcs_api_client_openapi::apis::configuration::Configuration as OpenApiConfiguration;
+#[cfg(not(any(feature = "grpc-web", feature = "tracing")))]
 use tonic::transport::Channel;
 use tonic::Status;
 
@@ -29,17 +34,27 @@ const DEFAULT_MAX_MESSAGE_DECODING_SIZE: usize = 50 * 1024 * 1024;
 /// It is public so that users can create gRPC clients with different APIs using a "raw" connection
 /// initialized by this library. This ensures that the exact Tonic version used for such clients
 /// matches what this library uses.
-#[cfg(not(feature = "grpc-web"))]
+#[cfg(not(any(feature = "grpc-web", feature = "tracing")))]
 pub type GrpcConnection = RetryService<RefreshService<Channel, ClientConfiguration>>;
 
 /// A type alias for the underlying gRPC connection used by all gRPC clients within this library.
 /// It is public so that users can create gRPC clients with different APIs using a "raw" connection
 /// initialized by this library. This ensures that the exact Tonic version used for such clients
 /// matches what this library uses.
-#[cfg(feature = "grpc-web")]
+#[cfg(all(feature = "grpc-web", not(feature = "tracing")))]
 pub type GrpcConnection =
     GrpcWebWrapperLayerService<RetryService<RefreshService<Channel, ClientConfiguration>>>;
 
+/// A type alias for the underlying gRPC connection used by all gRPC clients within this library.
+/// It is public so that users can create gRPC clients with different APIs using a "raw" connection
+/// initialized by this library. This ensures that the exact Tonic version used for such clients
+/// matches what this library uses.
+///
+/// The underlying [`tonic::transport::Channel`] is wrapped by [`CustomTraceService`],
+/// [`RefreshService`], and [`RetryService`] to provide tracing, token refresh, and retry.
+#[cfg(all(not(feature = "grpc-web"), feature = "tracing"))]
+pub type GrpcConnection = RetryService<RefreshService<CustomTraceService, ClientConfiguration>>;
+
 /// TODO: make configurable at the client level.
 /// <https://github.com/rigetti/qcs-sdk-rust/issues/239>
 pub(crate) static DEFAULT_HTTP_API_TIMEOUT: Duration = Duration::from_secs(10);
@@ -103,11 +118,25 @@ impl Qcs {
     ) -> Result<TranslationClient<GrpcConnection>, GrpcError<TokenError>> {
         let uri = parse_uri(translation_grpc_endpoint)?;
         let channel = get_channel(uri)?;
-        let service =
-            wrap_channel_with_retry(wrap_channel_with(channel, self.get_config().clone()));
+
+        // First add tracing if enabled
+        #[cfg(feature = "tracing")]
+        let channel = wrap_channel_with_tracing(
+            channel,
+            translation_grpc_endpoint.to_string(),
+            self.get_config()
+                .tracing_configuration()
+                .cloned()
+                .unwrap_or_default(),
+        );
+
+        // Then wrap with refresh and retry
+        let channel = wrap_channel_with(channel, self.get_config().clone());
+        let channel = wrap_channel_with_retry(channel);
+
         #[cfg(feature = "grpc-web")]
-        let service = wrap_channel_with_grpc_web(service);
-        Ok(TranslationClient::new(service)
+        let channel = wrap_channel_with_grpc_web(service);
+        Ok(TranslationClient::new(channel)
             .max_encoding_message_size(DEFAULT_MAX_MESSAGE_ENCODING_SIZE)
             .max_decoding_message_size(DEFAULT_MAX_MESSAGE_DECODING_SIZE))
     }

crates/lib/src/client.rs

@@ -11,6 +11,8 @@ use derive_builder::Builder;
 use qcs_api_client_common::configuration::TokenError;
 #[cfg(feature = "grpc-web")]
 use qcs_api_client_grpc::tonic::wrap_channel_with_grpc_web;
+#[cfg(feature = "tracing")]
+use qcs_api_client_grpc::tonic::wrap_channel_with_tracing;
 pub use qcs_api_client_grpc::tonic::Error as GrpcError;
 use qcs_api_client_grpc::{
     get_channel_with_timeout,
@@ -581,10 +583,27 @@ pub trait ExecutionTarget<'a> {
         let uri = parse_uri(address).map_err(QpuApiError::GrpcError)?;
         let channel = get_channel_with_timeout(uri, self.timeout())
             .map_err(|err| QpuApiError::GrpcError(err.into()))?;
-        let channel =
-            wrap_channel_with_retry(wrap_channel_with(channel, client.get_config().clone()));
+
+        // First add tracing if enabled
+        #[cfg(feature = "tracing")]
+        let channel = wrap_channel_with_tracing(
+            channel,
+            address.to_string(),
+            client
+                .get_config()
+                .tracing_configuration()
+                .cloned()
+                .unwrap_or_default(),
+        );
+
+        // Then wrap with refresh and retry
+        let channel = wrap_channel_with(channel, client.get_config().clone());
+        let channel = wrap_channel_with_retry(channel);
+
+        // Add grpc-web if enabled
         #[cfg(feature = "grpc-web")]
         let channel = wrap_channel_with_grpc_web(channel);
+
         Ok(channel)
     }
 

crates/lib/src/qpu/api.rs

@@ -3,6 +3,7 @@
 
 use std::{collections::HashMap, time::Duration};
 
+use pyo3::{exceptions::PyValueError, PyErr};
 use qcs_api_client_grpc::{
     models::controller::EncryptedControllerJob,
     services::translation::{
@@ -29,6 +30,16 @@ pub enum Error {
     ClientTimeout(#[from] Elapsed),
 }
 
+impl From<Error> for PyErr {
+    fn from(value: Error) -> Self {
+        let message = value.to_string();
+        match value {
+            Error::Grpc(_) => PyValueError::new_err(message),
+            Error::ClientTimeout(_) => PyValueError::new_err(message),
+        }
+    }
+}
+
 /// An encrypted and translated program, along with `readout_map`
 /// to map job `readout_data` back to program-declared variables.
 #[derive(Debug)]
@@ -44,7 +55,10 @@ pub struct EncryptedTranslationResult {
 }
 
 /// Translate a program, returning an encrypted and translated program.
-#[cfg_attr(feature = "tracing", instrument(skip_all))]
+#[cfg_attr(
+    feature = "tracing",
+    instrument(skip(quil_program, client, translation_options))
+)]
 pub async fn translate<TO>(
     quantum_processor_id: &str,
     quil_program: &str,

crates/lib/src/qpu/translation.rs

@@ -20,6 +20,7 @@ crate-type = ["cdylib", "rlib"]
 
 [dependencies]
 async-trait = "0.1.73"
+futures-util = "0.3.24"
 qcs = { path = "../lib", features = ["tracing-opentelemetry", "experimental"] }
 qcs-api-client-common = { workspace = true, features = ["python"] }
 qcs-api-client-grpc.workspace = true
@@ -37,8 +38,8 @@ rigetti-pyo3.workspace = true
 paste = "1.0.11"
 pyo3-log.workspace = true
 once_cell = "1.18.0"
-opentelemetry = { version = "0.23.0" }
-opentelemetry_sdk = { version = "0.23.0" }
+opentelemetry = { workspace = true }
+opentelemetry_sdk = { workspace = true }
 tracing = { version = "0.1.37" }
 prost = "0.13.3"
 

crates/python/Cargo.toml

@@ -87,8 +87,8 @@ class MultishotRequest:
 
     def __new__(
         cls,
-        compiled_quil: str,
-        trials: int,
+        program: str,
+        shots: int,
         addresses: Mapping[str, AddressRequest],
         measurement_noise: Optional[Tuple[float, float, float]],
         gate_noise: Optional[Tuple[float, float, float]],
@@ -142,7 +142,7 @@ def run(
     """
     ...
 
-def run_async(
+async def run_async(
     request: MultishotRequest,
     client: QVMClient,
     options: Optional[QVMOptions] = None,