From 8484247809d53d8e3386e34fe8c8de77142891bd Mon Sep 17 00:00:00 2001 From: "Jens H. Nielsen" Date: Wed, 11 Nov 2026 19:47:12 +0100 Subject: [PATCH 1/5] Fix subscriber races and run all tests in parallel The tests marked serial were flaky in parallel because of two race conditions in _Subscriber, and one relied on wall-clock timing. - Use the queue size instead of a separate counter that was reset after the callback, which held back rows inserted during a callback until the next insert. - Hold a lock while draining the queue and calling the callback so the completion callback cannot run concurrently with one on the subscriber thread, which delivered results out of order. - Drive the DataSaver clock by hand in test_datasaver_scalars instead of sleeping. With those fixed, drop the serial marker and the separate serial CI step, and drop the flaky marker together with pytest-rerunfailures. --- .github/workflows/pytest.yaml | 14 +++------- .../+dataset-subscriber-races.improved | 5 ++++ pyproject.toml | 2 -- requirements.txt | 4 --- src/qcodes/dataset/subscriber.py | 21 ++++++++++----- .../test_measurement_context_manager.py | 27 +++++++++---------- tests/dataset/test_nested_measurements.py | 2 -- tests/dataset/test_subscribing.py | 2 -- 8 files changed, 35 insertions(+), 42 deletions(-) create mode 100644 docs/changes/newsfragments/+dataset-subscriber-races.improved diff --git a/.github/workflows/pytest.yaml b/.github/workflows/pytest.yaml index 54aae58d02f8..84e0f8e9f840 100644 --- a/.github/workflows/pytest.yaml +++ b/.github/workflows/pytest.yaml @@ -120,21 +120,13 @@ jobs: if: ${{ runner.os == 'Windows' }} run: echo "PYTEST_BASETEMP=--basetemp=D:\\tmp" >> $GITHUB_ENV shell: bash - - name: Run parallel tests + - name: Run tests if: ${{ always() }} shell: bash run: | set -x - pytest -m "not serial" --cov --cov-report xml --hypothesis-profile ci --durations=20 $PYTEST_OPT $PYTEST_BASETEMP tests - id: parallel_tests -# a subset of the tests fails when run in parallel on Windows so run those in serial here - - name: Run serial tests - if: ${{ always() }} - shell: bash - run: | - set -x - pytest -m "serial" -n 0 --dist no --cov --cov-report xml --cov-append --hypothesis-profile ci $PYTEST_OPT $PYTEST_BASETEMP tests - id: serial_tests + pytest --cov --cov-report xml --hypothesis-profile ci --durations=20 $PYTEST_OPT $PYTEST_BASETEMP tests + id: tests - name: Upload coverage to Codecov uses: codecov/codecov-action@303a32d7a59b442fa8d48b6a1cc6825c09c847a5 # v7.1.1 with: diff --git a/docs/changes/newsfragments/+dataset-subscriber-races.improved b/docs/changes/newsfragments/+dataset-subscriber-races.improved new file mode 100644 index 000000000000..a069836c2091 --- /dev/null +++ b/docs/changes/newsfragments/+dataset-subscriber-races.improved @@ -0,0 +1,5 @@ +Fixed two race conditions in ``DataSet`` subscribers. A result inserted while a subscriber's +callback was running could be held back until the next insert, and when a dataset was +completed, its final callback could run concurrently with one already in progress on the +subscriber thread, delivering results out of order. Subscriber callbacks now never overlap +and always receive results in the order they were inserted. diff --git a/pyproject.toml b/pyproject.toml index 700f2bde59db..173b964e4761 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -76,7 +76,6 @@ test = [ "pytest-asyncio>=0.24.0", "pytest-cov>=5.0.0", "pytest-mock>=3.14.0", - "pytest-rerunfailures>=14.0", "pytest-xdist>=3.6.1", "PyVisa-sim>=0.6.0", "sphinx>=4.5.0", # sphinx extension tests @@ -215,7 +214,6 @@ minversion = "7.2" testpaths = "tests" addopts = "-n auto --dist=loadfile --cov-config=pyproject.toml" asyncio_default_fixture_loop_scope = "function" -markers = "serial" # we ignore warnings # triggered by third party packages diff --git a/requirements.txt b/requirements.txt index dfc2cd655eae..66a5201cdceb 100644 --- a/requirements.txt +++ b/requirements.txt @@ -243,7 +243,6 @@ packaging==26.3 # matplotlib # nbconvert # pytest - # pytest-rerunfailures # qcodes # sphinx # versioningit @@ -306,7 +305,6 @@ pytest==9.2.0 # pytest-asyncio # pytest-cov # pytest-mock - # pytest-rerunfailures # pytest-xdist pytest-asyncio==1.4.1.post1 # via qcodes (pyproject.toml) @@ -314,8 +312,6 @@ pytest-cov==7.1.0 # via qcodes (pyproject.toml) pytest-mock==3.17.0 # via qcodes (pyproject.toml) -pytest-rerunfailures==16.7 - # via qcodes (pyproject.toml) pytest-xdist==3.8.0 # via qcodes (pyproject.toml) python-dateutil==4.0.0 diff --git a/src/qcodes/dataset/subscriber.py b/src/qcodes/dataset/subscriber.py index 5c87049c3f65..e62294f325d4 100644 --- a/src/qcodes/dataset/subscriber.py +++ b/src/qcodes/dataset/subscriber.py @@ -4,7 +4,7 @@ import logging import time from queue import Empty, Queue -from threading import Thread +from threading import Lock, Thread from typing import TYPE_CHECKING, Any from qcodes.dataset.sqlite.connection import atomic_transaction @@ -50,7 +50,12 @@ def __init__( self.state = state self.data_queue: Queue[Any] = Queue() - self._queue_length: int = 0 + # The callback is called both from this thread and, on completion, + # from the thread that completes the dataset. Holding this lock while + # draining the queue and calling the callback keeps those calls from + # overlapping, so results are delivered in insertion order and + # done_callback only returns once every queued result is delivered. + self._callback_lock = Lock() self._stop_signal: bool = False # convert milliseconds to seconds self._loop_sleep_time = loop_sleep_time / 1000 @@ -83,7 +88,6 @@ def __init__( def _cache_data_to_queue(self, *args: Any) -> None: self.data_queue.put(args) self._data_set_len += 1 - self._queue_length += 1 def run(self) -> None: self.log.debug("Starting subscriber") @@ -100,8 +104,9 @@ def _exhaust_queue(queue: Queue[Any]) -> list[Any]: return result_list def _call_callback_on_queue_data(self) -> None: - result_list = self._exhaust_queue(self.data_queue) - self.callback(result_list, self._data_set_len, self.state) + with self._callback_lock: + result_list = self._exhaust_queue(self.data_queue) + self.callback(result_list, self._data_set_len, self.state) def _loop(self) -> None: while True: @@ -109,9 +114,11 @@ def _loop(self) -> None: self._clean_up() break - if self._queue_length >= self.min_queue_length: + # Use the queue's own size rather than a separate counter: a + # counter reset after the callback would hold back results + # inserted while the callback was running until the next insert. + if self.data_queue.qsize() >= self.min_queue_length: self._call_callback_on_queue_data() - self._queue_length = 0 time.sleep(self._loop_sleep_time) diff --git a/tests/dataset/measurement/test_measurement_context_manager.py b/tests/dataset/measurement/test_measurement_context_manager.py index d07287fb03cb..91c15ac36148 100644 --- a/tests/dataset/measurement/test_measurement_context_manager.py +++ b/tests/dataset/measurement/test_measurement_context_manager.py @@ -6,8 +6,8 @@ import re import traceback from functools import reduce -from time import sleep from typing import TYPE_CHECKING, Any, Literal +from unittest.mock import patch import hypothesis.strategies as hst import numpy as np @@ -674,8 +674,6 @@ def collect_x_vals( assert collected_x_vals == given_x_vals -@pytest.mark.serial -@pytest.mark.flaky(reruns=5) def test_subscribers_called_for_all_data_points( experiment: Experiment, DAC: DummyInstrument, DMM: DummyInstrument ) -> None: @@ -717,9 +715,6 @@ def sub_get_y_vals( assert yvals == list(given_yvals) -# There is no way around it: this test is slow. We test that write_period -# works and hence we must wait for some time to elapse. Sorry. -@pytest.mark.serial @settings( max_examples=5, deadline=None, @@ -752,21 +747,25 @@ def test_datasaver_scalars( meas.register_parameter(DAC.ch1) meas.register_parameter(DMM.v1, setpoints=(DAC.ch1,)) - with meas.run() as datasaver: + # Drive the clock that the DataSaver uses for write_period by hand, so + # that the outcome does not depend on how fast the machine runs the test. + now = 0.0 + + def fake_perf_counter() -> float: + return now + + with ( + patch("qcodes.dataset.measurements.perf_counter", fake_perf_counter), + meas.run() as datasaver, + ): for set_v, get_v in zip(set_values[:breakpoint], get_values[:breakpoint]): datasaver.add_result((DAC.ch1, set_v), (DMM.v1, get_v)) assert datasaver._dataset.number_of_results == 0 - sleep(write_period * 1.1) + now += write_period * 1.1 datasaver.add_result( (DAC.ch1, set_values[breakpoint]), (DMM.v1, get_values[breakpoint]) ) - # test that work with time intervals are often flaky - # so we add a bit more wait time here if the expected number - # of points have not been written - for _ in range(10): - if datasaver.points_written != breakpoint + 1: - sleep(write_period * 1.1) assert datasaver.points_written == breakpoint + 1 assert datasaver.run_id == no_of_runs + 1 diff --git a/tests/dataset/test_nested_measurements.py b/tests/dataset/test_nested_measurements.py index cd151160f092..73172ebb669e 100644 --- a/tests/dataset/test_nested_measurements.py +++ b/tests/dataset/test_nested_measurements.py @@ -189,8 +189,6 @@ def subscriber( return subscriber -@pytest.mark.flaky(reruns=5) -@pytest.mark.serial def test_subscription_on_dual_datasets( experiment: "Experiment", basic_subscriber: "Callable[[list[tuple[VALUE]], int, dict | None], None]", diff --git a/tests/dataset/test_subscribing.py b/tests/dataset/test_subscribing.py index ad5b84fb5794..ccd2dd7a1399 100644 --- a/tests/dataset/test_subscribing.py +++ b/tests/dataset/test_subscribing.py @@ -134,8 +134,6 @@ def _make_broken_subscriber_config(tmp_path: "Path") -> "Generator[None, None, N yield -@pytest.mark.flaky(reruns=5) -@pytest.mark.serial def test_basic_subscription( dataset: "DataSet", basic_subscriber: "Callable[[list[tuple[VALUE]], int, dict | None], None]", From c2c9f585992f53bfc361599e033cb958e82251ec Mon Sep 17 00:00:00 2001 From: "Jens H. Nielsen" Date: Wed, 11 Nov 2026 20:15:22 +0100 Subject: [PATCH 2/5] Raise NotImplementedError for subscribers with background writing A subscriber's trigger calls an SQL function that is only registered on the DataSet's own connection. The background writer inserts through its own connection, so the first insert failed with "no such function", the writer thread died and the measurement hung waiting for it. This combination has never worked, so reject it explicitly for now. - DataSet.subscribe raises while the database is written in the background, and mark_started(start_bg_writer=True) raises if the dataset already has subscribers, before any start action is taken. - Runner.__enter__ raises before the run is created if subscribers are set up via add_subscriber, the default_subscribers config option or DataSaver.default_callback, so no run is left incomplete. --- .../+subscribers-background-writing.improved | 5 +++ src/qcodes/dataset/data_set.py | 13 +++++++ src/qcodes/dataset/measurements.py | 30 ++++++++++++++- .../test_measurement_context_manager.py | 18 +++++++++ tests/dataset/test_datasaver.py | 38 +++++++++++++++++-- tests/dataset/test_subscribing.py | 14 +++++++ 6 files changed, 113 insertions(+), 5 deletions(-) create mode 100644 docs/changes/newsfragments/+subscribers-background-writing.improved diff --git a/docs/changes/newsfragments/+subscribers-background-writing.improved b/docs/changes/newsfragments/+subscribers-background-writing.improved new file mode 100644 index 000000000000..541a3c8b1086 --- /dev/null +++ b/docs/changes/newsfragments/+subscribers-background-writing.improved @@ -0,0 +1,5 @@ +Subscribing to a ``DataSet`` that is written in the background (``write_in_background=True``) +now raises a ``NotImplementedError``. This combination never worked: the background writer +failed to insert the first results and the measurement then hung. ``Measurement.run`` raises +the error before the run is created if subscribers are set up via ``Measurement.add_subscriber``, +the ``subscription.default_subscribers`` config option or ``DataSaver.default_callback``. diff --git a/src/qcodes/dataset/data_set.py b/src/qcodes/dataset/data_set.py index 08536a0855d2..b700e6904d8b 100644 --- a/src/qcodes/dataset/data_set.py +++ b/src/qcodes/dataset/data_set.py @@ -207,6 +207,15 @@ class _WriterStatus: _WRITERS: dict[str, _WriterStatus] = {} +# The SQL function that a subscriber's trigger calls only exists on the +# DataSet's own connection, so inserts made by the background writer, which +# uses its own connection, would fail. +_SUBSCRIBERS_WITH_BACKGROUND_WRITING_ERROR = ( + "Subscribers are not supported for a DataSet that is written in the " + "background. Write the DataSet in the main thread " + "(write_in_background=False) to subscribe to it." +) + class DataSet(BaseDataSet): # the "persistent traits" are the attributes/properties of the DataSet @@ -770,6 +779,8 @@ def _perform_start_actions( "be written either in the background or in the " "main thread. You cannot mix." ) + if start_bg_writer and (self.subscribers or self._pending_subscribers): + raise NotImplementedError(_SUBSCRIBERS_WITH_BACKGROUND_WRITING_ERROR) with atomic(self.conn) as conn: for spec in paramspecs: @@ -1248,6 +1259,8 @@ def subscribe( state: Any | None = None, callback_kwargs: Mapping[str, Any] | None = None, ) -> str: + if self._writer_status.write_in_background: + raise NotImplementedError(_SUBSCRIBERS_WITH_BACKGROUND_WRITING_ERROR) subscriber_id = uuid.uuid4().hex if not self._results_table_exists: # The results table does not exist yet - this happens with split diff --git a/src/qcodes/dataset/measurements.py b/src/qcodes/dataset/measurements.py index 53e687107b54..e996c68c722e 100644 --- a/src/qcodes/dataset/measurements.py +++ b/src/qcodes/dataset/measurements.py @@ -26,7 +26,11 @@ import qcodes as qc import qcodes.validators as vals -from qcodes.dataset.data_set import DataSet, load_by_guid +from qcodes.dataset.data_set import ( + _SUBSCRIBERS_WITH_BACKGROUND_WRITING_ERROR, + DataSet, + load_by_guid, +) from qcodes.dataset.data_set_in_memory import DataSetInMem from qcodes.dataset.data_set_protocol import ( DataSetProtocol, @@ -602,6 +606,22 @@ def _calculate_write_period( return float(write_period) def __enter__(self) -> DataSaver: + # Fail before creating the run: subscribing to a DataSet that is + # written in the background is not supported, and raising only once + # the run is started would leave it incomplete. The DataSaver also + # adds the default callback and the default subscribers from the config. + if ( + self._write_in_background + and self._dataset_class is DataSetType.DataSet + and ( + self.subscribers + or qc.config.subscription.default_subscribers + or "run_tables_subscription_callback" + in (DataSaver.default_callback or {}) + ) + ): + raise NotImplementedError(_SUBSCRIBERS_WITH_BACKGROUND_WRITING_ERROR) + # 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 @@ -1480,6 +1500,9 @@ def add_subscriber( """ Add a subscriber to the dataset of the measurement. + Subscribers are not supported when the dataset is written in the + background, see the ``write_in_background`` argument of :meth:`run`. + Args: func: A function taking three positional arguments: a list of tuples of parameter values, an integer, a mutable variable @@ -1519,7 +1542,10 @@ def run( will be stored in background, without blocking the main thread that is executing the context manager. By default the setting for write in background will be - read from the ``qcodesrc.json`` config file. + read from the ``qcodesrc.json`` config file. Subscribers + are not supported when writing in the background for the + default ``DataSet`` class, so entering the context manager + raises a ``NotImplementedError`` if any are set up. in_memory_cache: Should measured data be keep in memory and available as part of the `dataset.cache` object. dataset_class: Enum representing the Class used to store data diff --git a/tests/dataset/measurement/test_measurement_context_manager.py b/tests/dataset/measurement/test_measurement_context_manager.py index 91c15ac36148..e009b4c5e2be 100644 --- a/tests/dataset/measurement/test_measurement_context_manager.py +++ b/tests/dataset/measurement/test_measurement_context_manager.py @@ -715,6 +715,24 @@ def sub_get_y_vals( assert yvals == list(given_yvals) +def test_subscribers_with_background_writing_not_supported( + experiment: Experiment, DAC: DummyInstrument +) -> None: + meas = Measurement(exp=experiment) + meas.register_parameter(DAC.ch1) + meas.add_subscriber(lambda results, length, state: None, state=[]) + no_of_runs = len(experiment) + + with ( + pytest.raises(NotImplementedError, match="written in the background"), + meas.run(write_in_background=True), + ): + pass + + # the error is raised before the run is created + assert len(experiment) == no_of_runs + + @settings( max_examples=5, deadline=None, diff --git a/tests/dataset/test_datasaver.py b/tests/dataset/test_datasaver.py index a552b1b7c23b..7ae5ae3b08a4 100644 --- a/tests/dataset/test_datasaver.py +++ b/tests/dataset/test_datasaver.py @@ -42,8 +42,7 @@ def reset_callback_globals() -> "None": @pytest.mark.usefixtures("experiment") -@pytest.mark.parametrize("bg_writing", [True, False]) -def test_default_callback(bg_writing: "bool") -> None: +def test_default_callback() -> None: """ The Web UI needs to know the results of an experiment with the metadata. So a default_callback class variable is set by the Web UI with a callback @@ -67,7 +66,7 @@ def test_default_callback(bg_writing: "bool") -> None: interdeps=InterDependencies_(), registered_parameters=[], ) - test_set.mark_started(start_bg_writer=bg_writing) + test_set.mark_started() test_set.mark_completed() assert CALLBACK_SNAPSHOT == "reasonable_snapshot" assert CALLBACK_RUN_ID is not None @@ -79,6 +78,39 @@ def test_default_callback(bg_writing: "bool") -> None: test_set.conn.close() +@pytest.mark.usefixtures("experiment") +def test_default_callback_with_background_writing_not_supported() -> None: + """ + Subscribers, such as the default callback, are not supported for a + dataset that is written in the background, so starting it raises. + """ + test_set = None + try: + DataSaver.default_callback = { + "run_tables_subscription_callback": callback, + "run_tables_subscription_min_wait": 1, + "run_tables_subscription_min_count": 2, + } + + test_set = new_data_set("test-dataset") + test_set.add_metadata("snapshot", "reasonable_snapshot") + DataSaver( + dataset=test_set, + write_period=0, + interdeps=InterDependencies_(), + registered_parameters=[], + ) + with pytest.raises(NotImplementedError, match="written in the background"): + test_set.mark_started(start_bg_writer=True) + assert test_set.pristine + finally: + DataSaver.default_callback = None + if test_set is not None: + # The dataset is never completed, so stop the subscriber explicitly. + test_set.unsubscribe_all() + test_set.conn.close() + + @pytest.mark.usefixtures("experiment") @pytest.mark.parametrize("bg_writing", [True, False]) def test_numpy_types(bg_writing: "bool") -> None: diff --git a/tests/dataset/test_subscribing.py b/tests/dataset/test_subscribing.py index ccd2dd7a1399..f2c5c0e32f46 100644 --- a/tests/dataset/test_subscribing.py +++ b/tests/dataset/test_subscribing.py @@ -177,6 +177,20 @@ def assert_expected_state() -> "None": assert len(triggers) == 0 +def test_subscribe_with_background_writing_not_supported(dataset: "DataSet") -> None: + xparam = ParamSpecBase(name="x", paramtype="numeric") + yparam = ParamSpecBase(name="y", paramtype="numeric") + dataset.set_interdependencies(InterDependencies_(dependencies={yparam: (xparam,)})) + dataset.mark_started(start_bg_writer=True) + try: + with pytest.raises(NotImplementedError, match="written in the background"): + dataset.subscribe(lambda results, length, state: None) + assert dataset.subscribers == {} + finally: + # completing the dataset shuts down the background writer + dataset.mark_completed() + + @pytest.mark.usefixtures("working_subscriber_config") def test_subscription_from_config( dataset: "DataSet", From 703b2d444c189a303622c0072bcd9ff6222b20aa Mon Sep 17 00:00:00 2001 From: "Jens H. Nielsen" Date: Wed, 11 Nov 2026 20:19:01 +0100 Subject: [PATCH 3/5] Rename newsfragments to PR number 8758 --- .../{+dataset-subscriber-races.improved => 8758.improved} | 0 .../{+subscribers-background-writing.improved => 8758.improved.1} | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename docs/changes/newsfragments/{+dataset-subscriber-races.improved => 8758.improved} (100%) rename docs/changes/newsfragments/{+subscribers-background-writing.improved => 8758.improved.1} (100%) diff --git a/docs/changes/newsfragments/+dataset-subscriber-races.improved b/docs/changes/newsfragments/8758.improved similarity index 100% rename from docs/changes/newsfragments/+dataset-subscriber-races.improved rename to docs/changes/newsfragments/8758.improved diff --git a/docs/changes/newsfragments/+subscribers-background-writing.improved b/docs/changes/newsfragments/8758.improved.1 similarity index 100% rename from docs/changes/newsfragments/+subscribers-background-writing.improved rename to docs/changes/newsfragments/8758.improved.1 From cc3dcbebd956014cac7f89cc9effec76d5a4772a Mon Sep 17 00:00:00 2001 From: "Jens H. Nielsen" Date: Wed, 11 Nov 2026 20:21:12 +0100 Subject: [PATCH 4/5] Use pytest-mock to patch the DataSaver clock in test_datasaver_scalars --- .../measurement/test_measurement_context_manager.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/tests/dataset/measurement/test_measurement_context_manager.py b/tests/dataset/measurement/test_measurement_context_manager.py index e009b4c5e2be..9ef5b848b061 100644 --- a/tests/dataset/measurement/test_measurement_context_manager.py +++ b/tests/dataset/measurement/test_measurement_context_manager.py @@ -7,7 +7,6 @@ import traceback from functools import reduce from typing import TYPE_CHECKING, Any, Literal -from unittest.mock import patch import hypothesis.strategies as hst import numpy as np @@ -748,6 +747,7 @@ def test_datasaver_scalars( experiment: Experiment, DAC: DummyInstrument, DMM: DummyInstrument, + mocker: MockerFixture, set_values: list[float], get_values: list[float], breakpoint: int, @@ -772,10 +772,9 @@ def test_datasaver_scalars( def fake_perf_counter() -> float: return now - with ( - patch("qcodes.dataset.measurements.perf_counter", fake_perf_counter), - meas.run() as datasaver, - ): + mocker.patch("qcodes.dataset.measurements.perf_counter", fake_perf_counter) + + with meas.run() as datasaver: for set_v, get_v in zip(set_values[:breakpoint], get_values[:breakpoint]): datasaver.add_result((DAC.ch1, set_v), (DMM.v1, get_v)) From 9fafbdca6defed59f125bdb60abd55f2e10111c8 Mon Sep 17 00:00:00 2001 From: "Jens H. Nielsen" Date: Thu, 12 Nov 2026 20:01:29 +0100 Subject: [PATCH 5/5] Address review comments on subscriber length and background check - Only update the length passed to subscriber callbacks while draining the queue under the callback lock. Since the subscriber thread polls the queue size, a row could be delivered before the inserting thread had counted it, so the callback got a stale length. - Only reject subscribing if this run is the one being written in the background. The writer status is shared by every run in a database, so subscribing to a dataset that has not been started was rejected whenever another run was written in the background. --- src/qcodes/dataset/data_set.py | 9 ++++- src/qcodes/dataset/subscriber.py | 5 ++- tests/dataset/test_subscribing.py | 66 +++++++++++++++++++++++++++++++ 3 files changed, 78 insertions(+), 2 deletions(-) diff --git a/src/qcodes/dataset/data_set.py b/src/qcodes/dataset/data_set.py index b700e6904d8b..e9888fc7e689 100644 --- a/src/qcodes/dataset/data_set.py +++ b/src/qcodes/dataset/data_set.py @@ -1259,7 +1259,14 @@ def subscribe( state: Any | None = None, callback_kwargs: Mapping[str, Any] | None = None, ) -> str: - if self._writer_status.write_in_background: + writer_status = self._writer_status + # Background writing is tracked per database, so also check that it is + # this run that is being written. A dataset that has not been started + # can be subscribed to; starting it in the background then raises. + if ( + writer_status.write_in_background + and self.run_id in writer_status.active_datasets + ): raise NotImplementedError(_SUBSCRIBERS_WITH_BACKGROUND_WRITING_ERROR) subscriber_id = uuid.uuid4().hex if not self._results_table_exists: diff --git a/src/qcodes/dataset/subscriber.py b/src/qcodes/dataset/subscriber.py index e62294f325d4..506cc134632a 100644 --- a/src/qcodes/dataset/subscriber.py +++ b/src/qcodes/dataset/subscriber.py @@ -45,6 +45,9 @@ def __init__( self.dataSet = dataSet self.table_name = dataSet.table_name + # Number of results in the dataset up to and including the ones last + # passed to the callback. Only updated while draining the queue, so it + # always matches the results that the callback has been given. self._data_set_len = len(dataSet) self.state = state @@ -87,7 +90,6 @@ def __init__( def _cache_data_to_queue(self, *args: Any) -> None: self.data_queue.put(args) - self._data_set_len += 1 def run(self) -> None: self.log.debug("Starting subscriber") @@ -106,6 +108,7 @@ def _exhaust_queue(queue: Queue[Any]) -> list[Any]: def _call_callback_on_queue_data(self) -> None: with self._callback_lock: result_list = self._exhaust_queue(self.data_queue) + self._data_set_len += len(result_list) self.callback(result_list, self._data_set_len, self.state) def _loop(self) -> None: diff --git a/tests/dataset/test_subscribing.py b/tests/dataset/test_subscribing.py index f2c5c0e32f46..fde8edc47de1 100644 --- a/tests/dataset/test_subscribing.py +++ b/tests/dataset/test_subscribing.py @@ -7,6 +7,7 @@ from numpy import ndarray import qcodes +from qcodes.dataset import new_data_set from qcodes.dataset.descriptions.dependencies import InterDependencies_ from qcodes.dataset.sqlite.connection import atomic_transaction from qcodes.parameters import ParamSpecBase @@ -191,6 +192,71 @@ def test_subscribe_with_background_writing_not_supported(dataset: "DataSet") -> dataset.mark_completed() +def test_subscribe_while_another_run_writes_in_background( + dataset: "DataSet", +) -> None: + """ + Only the run that is written in the background can not be subscribed to. + Another dataset in the same database can, and can then be written in the + main thread once the background run has completed. + """ + xparam = ParamSpecBase(name="x", paramtype="numeric") + yparam = ParamSpecBase(name="y", paramtype="numeric") + idps = InterDependencies_(dependencies={yparam: (xparam,)}) + + background_run = new_data_set("background-run") + background_run.set_interdependencies(idps) + background_run.mark_started(start_bg_writer=True) + + dataset.set_interdependencies(idps) + received: list[tuple[VALUE, ...]] = [] + try: + dataset.subscribe( + lambda results, length, state: received.extend(results), + min_wait=0, + min_count=1, + ) + finally: + background_run.mark_completed() + background_run.conn.close() + + dataset.mark_started() + dataset.add_results([{"x": 1, "y": 2}]) + dataset.mark_completed() + assert received == [(1, 2)] + + +def test_subscriber_length_matches_delivered_results(dataset: "DataSet") -> None: + """ + The length passed to the callback is the number of results in the dataset + up to and including the ones passed along with it, however the inserts + and the callbacks interleave. + """ + xparam = ParamSpecBase(name="x", paramtype="numeric") + yparam = ParamSpecBase(name="y", paramtype="numeric") + dataset.set_interdependencies(InterDependencies_(dependencies={yparam: (xparam,)})) + dataset.mark_started() + + calls: list[tuple[int, int]] = [] + + def subscriber( + results: "list[tuple[VALUE, ...]]", length: int, state: None + ) -> None: + calls.append((length, len(results))) + + dataset.subscribe(subscriber, min_wait=0, min_count=1) + n_results = 1000 + for x in range(n_results): + dataset.add_results([{"x": x, "y": x}]) + dataset.mark_completed() + + delivered = 0 + for length, n_delivered in calls: + delivered += n_delivered + assert length == delivered + assert delivered == n_results + + @pytest.mark.usefixtures("working_subscriber_config") def test_subscription_from_config( dataset: "DataSet",