diff --git a/qcodes/dataset/measurement_extensions.py b/qcodes/dataset/measurement_extensions.py index dbd49582ba4f..53c7e77b2242 100644 --- a/qcodes/dataset/measurement_extensions.py +++ b/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): diff --git a/qcodes/dataset/measurements.py b/qcodes/dataset/measurements.py index 7e147fde9ec3..828357910bc1 100644 --- a/qcodes/dataset/measurements.py +++ b/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, )