instument datasaver_builder and dond_into

qcodes/dataset/measurement_extensions.py

@@ -6,6 +6,8 @@
 from dataclasses import dataclass
 from typing import Any, Callable
 
+from opentelemetry import trace
+
 from qcodes.dataset.dond.do_nd import _Sweeper
 from qcodes.dataset.dond.do_nd_utils import ParamMeasT, catch_interrupts
 from qcodes.dataset.dond.sweeps import AbstractSweep, LinSweep, TogetherSweep
@@ -14,6 +16,8 @@
 from qcodes.dataset.threading import process_params_meas
 from qcodes.parameters.parameter_base import ParameterBase
 
+TRACER = trace.get_tracer(__name__)
+
 
 @dataclass
 class DataSetDefinition:
@@ -85,12 +89,16 @@ def datasaver_builder(
     Yields:
         A list of generated datasavers with parameters registered
     """
+
     measurement_instances = setup_measurement_instances(
         dataset_definitions, override_experiment
     )
-    with catch_interrupts() as _, ExitStack() as stack:
+    with TRACER.start_as_current_span(
+        "qcodes.dataset.datasaver_builder"
+    ), catch_interrupts() as _, ExitStack() as stack:
+        datasaver_builder_span = trace.get_current_span()
         datasavers = [
-            stack.enter_context(measurement.run())
+            stack.enter_context(measurement.run(parent_span=datasaver_builder_span))
             for measurement in measurement_instances
         ]
         yield datasavers
@@ -153,31 +161,40 @@ def dond_into(
         additional_setpoints: A list of setpoint parameters to be registered in the
             measurement but not scanned/swept-over.
     """
-    sweep_instances, params_meas = parse_dond_into_args(*params)
-    sweeper = _Sweeper(sweep_instances, additional_setpoints)
-    for set_events in sweeper:
-        results: dict[ParameterBase, Any] = {}
-        additional_setpoints_data = process_params_meas(additional_setpoints)
-        for set_event in set_events:
-            if set_event.should_set:
-                set_event.parameter(set_event.new_value)
-                for act in set_event.actions:
-                    act()
-                time.sleep(set_event.delay)
-
-            if set_event.get_after_set:
-                results[set_event.parameter] = set_event.parameter()
-            else:
-                results[set_event.parameter] = set_event.new_value
-
-        meas_value_pair = process_params_meas(params_meas)
-        for meas_param, value in meas_value_pair:
-            results[meas_param] = value
-
-        datasaver.add_result(
-            *list(results.items()),
-            *additional_setpoints_data,
-        )
+    # at this stage multiple measurement context managers may be in run state
+    # as datasavers. Here we ensure we bind the parent span to the correct
+    # datasaver.
+    if datasaver._span is not None:
+        context = trace.set_span_in_context(datasaver._span)
+    else:
+        context = None
+
+    with TRACER.start_as_current_span("qcodes.dataset.dond_into", context=context):
+        sweep_instances, params_meas = parse_dond_into_args(*params)
+        sweeper = _Sweeper(sweep_instances, additional_setpoints)
+        for set_events in sweeper:
+            results: dict[ParameterBase, Any] = {}
+            additional_setpoints_data = process_params_meas(additional_setpoints)
+            for set_event in set_events:
+                if set_event.should_set:
+                    set_event.parameter(set_event.new_value)
+                    for act in set_event.actions:
+                        act()
+                    time.sleep(set_event.delay)
+
+                if set_event.get_after_set:
+                    results[set_event.parameter] = set_event.parameter()
+                else:
+                    results[set_event.parameter] = set_event.new_value
+
+            meas_value_pair = process_params_meas(params_meas)
+            for meas_param, value in meas_value_pair:
+                results[meas_param] = value
+
+            datasaver.add_result(
+                *list(results.items()),
+                *additional_setpoints_data,
+            )
 
 
 class LinSweeper(LinSweep):

qcodes/dataset/measurements.py

@@ -84,7 +84,9 @@ def __init__(
         dataset: DataSetProtocol,
         write_period: float,
         interdeps: InterDependencies_,
+        span: trace.Span,
     ) -> None:
+        self._span = span
         self._dataset = dataset
         if (
             DataSaver.default_callback is not None
@@ -504,6 +506,7 @@ def __init__(
         shapes: Shapes | None = None,
         in_memory_cache: bool | None = None,
         dataset_class: DataSetType = DataSetType.DataSet,
+        parent_span: trace.Span | None = None,
     ) -> None:
         if in_memory_cache is None:
             in_memory_cache = qc.config.dataset.in_memory_cache
@@ -529,6 +532,7 @@ def __init__(
         self._extra_log_info = extra_log_info
         self._write_in_background = write_in_background
         self._in_memory_cache = in_memory_cache
+        self._parent_span = parent_span
         self.ds: DataSetProtocol
 
     @staticmethod
@@ -549,10 +553,21 @@ def _calculate_write_period(
         return float(write_period)
 
     def __enter__(self) -> DataSaver:
+        # multiple runners can be active at the same time.
+        # If we just activate them in order the first one
+        # would be the parent of the next one but that is wrong
+        # since they are siblings that should coexist with the
+        # same parent.
+        if self._parent_span is not None:
+            context = trace.set_span_in_context(self._parent_span)
+        else:
+            context = None
         # We want to enter the opentelemetry span here
         # and end it in the `__exit__` of this context manger
         # so here we capture it in a exitstack that we keep around.
-        self._span = TRACER.start_span("qcodes.dataset.Measurement.run")
+        self._span = TRACER.start_span(
+            "qcodes.dataset.Measurement.run", context=context
+        )
         with ExitStack() as stack:
             stack.enter_context(trace.use_span(self._span, end_on_exit=True))
 
@@ -649,6 +664,7 @@ def __enter__(self) -> DataSaver:
             dataset=self.ds,
             write_period=self.write_period,
             interdeps=self._interdependencies,
+            span=self._span,
         )
 
         return self.datasaver
@@ -1248,6 +1264,7 @@ def run(
         write_in_background: bool | None = None,
         in_memory_cache: bool | None = True,
         dataset_class: DataSetType = DataSetType.DataSet,
+        parent_span: trace.Span | None = None,
     ) -> Runner:
         """
         Returns the context manager for the experimental run
@@ -1281,4 +1298,5 @@ def run(
             shapes=self._shapes,
             in_memory_cache=in_memory_cache,
             dataset_class=dataset_class,
+            parent_span=parent_span,
         )