sinter 1.15.0__py3-none-any.whl

sinter/_collection/_collection_worker_state.py

@@ -0,0 +1,259 @@
+import queue
+import time
+from typing import Any
+from typing import Optional
+from typing import TYPE_CHECKING
+
+import stim
+
+from sinter._data import AnonTaskStats
+from sinter._data import CollectionOptions
+from sinter._data import Task
+from sinter._decoding import CompiledSampler
+from sinter._decoding import Sampler
+
+if TYPE_CHECKING:
+    import multiprocessing
+
+
+def _fill_in_task(task: Task) -> Task:
+    changed = False
+    circuit = task.circuit
+    if circuit is None:
+        circuit = stim.Circuit.from_file(task.circuit_path)
+        changed = True
+    dem = task.detector_error_model
+    if dem is None:
+        try:
+            dem = circuit.detector_error_model(decompose_errors=True, approximate_disjoint_errors=True)
+        except ValueError:
+            try:
+                dem = circuit.detector_error_model(approximate_disjoint_errors=True)
+            except ValueError:
+                dem = circuit.detector_error_model(approximate_disjoint_errors=True, flatten_loops=True)
+        changed = True
+    if not changed:
+        return task
+    return Task(
+        circuit=circuit,
+        decoder=task.decoder,
+        detector_error_model=dem,
+        postselection_mask=task.postselection_mask,
+        postselected_observables_mask=task.postselected_observables_mask,
+        json_metadata=task.json_metadata,
+        collection_options=task.collection_options,
+    )
+
+
+class CollectionWorkerState:
+    def __init__(
+            self,
+            *,
+            flush_period: float,
+            worker_id: int,
+            inp: 'multiprocessing.Queue',
+            out: 'multiprocessing.Queue',
+            sampler: Sampler,
+            custom_error_count_key: Optional[str],
+    ):
+        assert isinstance(flush_period, (int, float))
+        assert isinstance(sampler, Sampler)
+        self.max_flush_period = flush_period
+        self.cur_flush_period = 0.01
+        self.inp = inp
+        self.out = out
+        self.sampler = sampler
+        self.compiled_sampler: CompiledSampler | None = None
+        self.worker_id = worker_id
+
+        self.current_task: Task | None = None
+        self.current_error_cutoff: int | None = None
+        self.custom_error_count_key = custom_error_count_key
+        self.current_task_shots_left: int = 0
+        self.unflushed_results: AnonTaskStats = AnonTaskStats()
+        self.last_flush_message_time = time.monotonic()
+        self.soft_error_flush_threshold: int = 1
+
+    def _send_message_to_manager(self, message: Any):
+        self.out.put(message)
+
+    def state_summary(self) -> str:
+        lines = [
+            f'Worker(id={self.worker_id}) [',
+            f'    max_flush_period={self.max_flush_period}',
+            f'    cur_flush_period={self.cur_flush_period}',
+            f'    sampler={self.sampler}',
+            f'    compiled_sampler={self.compiled_sampler}',
+            f'    current_task={self.current_task}',
+            f'    current_error_cutoff={self.current_error_cutoff}',
+            f'    custom_error_count_key={self.custom_error_count_key}',
+            f'    current_task_shots_left={self.current_task_shots_left}',
+            f'    unflushed_results={self.unflushed_results}',
+            f'    last_flush_message_time={self.last_flush_message_time}',
+            f'    soft_error_flush_threshold={self.soft_error_flush_threshold}',
+            f']',
+        ]
+        return '\n' + '\n'.join(lines) + '\n'
+
+    def flush_results(self):
+        if self.unflushed_results.shots > 0:
+            self.last_flush_message_time = time.monotonic()
+            self.cur_flush_period = min(self.cur_flush_period * 1.4, self.max_flush_period)
+            self._send_message_to_manager((
+                'flushed_results',
+                self.worker_id,
+                (self.current_task.strong_id(), self.unflushed_results),
+            ))
+            self.unflushed_results = AnonTaskStats()
+            return True
+        return False
+
+    def accept_shots(self, *, shots_delta: int):
+        assert shots_delta >= 0
+        self.current_task_shots_left += shots_delta
+        self._send_message_to_manager((
+            'accepted_shots',
+            self.worker_id,
+            (self.current_task.strong_id(), shots_delta),
+        ))
+
+    def return_shots(self, *, requested_shots: int):
+        assert requested_shots >= 0
+        returned_shots = max(0, min(requested_shots, self.current_task_shots_left))
+        self.current_task_shots_left -= returned_shots
+        if self.current_task_shots_left <= 0:
+            self.flush_results()
+        self._send_message_to_manager((
+            'returned_shots',
+            self.worker_id,
+            (self.current_task.strong_id(), returned_shots),
+        ))
+
+    def compute_strong_id(self, *, new_task: Task):
+        strong_id = _fill_in_task(new_task).strong_id()
+        self._send_message_to_manager((
+            'computed_strong_id',
+            self.worker_id,
+            strong_id,
+        ))
+
+    def change_job(self, *, new_task: Task, new_collection_options: CollectionOptions):
+        self.flush_results()
+
+        self.current_task = _fill_in_task(new_task)
+        self.current_error_cutoff = new_collection_options.max_errors
+        self.compiled_sampler = self.sampler.compiled_sampler_for_task(self.current_task)
+        assert self.current_task.strong_id() is not None
+        self.current_task_shots_left = 0
+        self.last_flush_message_time = time.monotonic()
+
+        self._send_message_to_manager((
+            'changed_job',
+            self.worker_id,
+            (self.current_task.strong_id(),),
+        ))
+
+    def process_messages(self) -> int:
+        num_processed = 0
+        while True:
+            try:
+                message = self.inp.get_nowait()
+            except queue.Empty:
+                return num_processed
+
+            num_processed += 1
+            message_type, message_body = message
+
+            if message_type == 'stop':
+                return -1
+
+            elif message_type == 'flush_results':
+                self.flush_results()
+
+            elif message_type == 'compute_strong_id':
+                assert isinstance(message_body, Task)
+                self.compute_strong_id(new_task=message_body)
+
+            elif message_type == 'change_job':
+                new_task, new_collection_options, soft_error_flush_threshold = message_body
+                self.cur_flush_period = 0.01
+                self.soft_error_flush_threshold = soft_error_flush_threshold
+                assert isinstance(new_task, Task)
+                self.change_job(new_task=new_task, new_collection_options=new_collection_options)
+
+            elif message_type == 'set_soft_error_flush_threshold':
+                soft_error_flush_threshold = message_body
+                self.soft_error_flush_threshold = soft_error_flush_threshold
+
+            elif message_type == 'accept_shots':
+                job_key, shots_delta = message_body
+                assert isinstance(shots_delta, int)
+                assert job_key == self.current_task.strong_id()
+                self.accept_shots(shots_delta=shots_delta)
+
+            elif message_type == 'return_shots':
+                job_key, requested_shots = message_body
+                assert isinstance(requested_shots, int)
+                assert job_key == self.current_task.strong_id()
+                self.return_shots(requested_shots=requested_shots)
+
+            else:
+                raise NotImplementedError(f'{message_type=}')
+
+    def num_unflushed_errors(self) -> int:
+        if self.custom_error_count_key is not None:
+            return self.unflushed_results.custom_counts[self.custom_error_count_key]
+        return self.unflushed_results.errors
+
+    def do_some_work(self) -> bool:
+        did_some_work = False
+
+        # Sample some stats.
+        if self.current_task_shots_left > 0:
+            # Don't keep sampling if we've exceeded the number of errors needed.
+            if self.current_error_cutoff is not None and self.current_error_cutoff <= 0:
+                return self.flush_results()
+
+            some_work_done = self.compiled_sampler.sample(self.current_task_shots_left)
+            if some_work_done.shots < 1:
+                raise ValueError(f"Sampler didn't do any work. It returned statistics with shots == 0: {some_work_done}.")
+            assert isinstance(some_work_done, AnonTaskStats)
+            self.current_task_shots_left -= some_work_done.shots
+            if self.current_error_cutoff is not None:
+                errors_done = some_work_done.custom_counts[self.custom_error_count_key] if self.custom_error_count_key is not None else some_work_done.errors
+                self.current_error_cutoff -= errors_done
+            self.unflushed_results += some_work_done
+            did_some_work = True
+
+        # Report them periodically.
+        should_flush = False
+        if self.num_unflushed_errors() >= self.soft_error_flush_threshold:
+            should_flush = True
+        if self.unflushed_results.shots > 0:
+            if self.current_task_shots_left <= 0 or self.last_flush_message_time + self.cur_flush_period < time.monotonic():
+                should_flush = True
+        if should_flush:
+            did_some_work |= self.flush_results()
+
+        return did_some_work
+
+    def run_message_loop(self):
+        try:
+            while True:
+                num_messages_processed = self.process_messages()
+                if num_messages_processed == -1:
+                    break
+                did_some_work = self.do_some_work()
+                if not did_some_work and num_messages_processed == 0:
+                    time.sleep(0.01)
+
+        except KeyboardInterrupt:
+            pass
+
+        except BaseException as ex:
+            import traceback
+            self._send_message_to_manager((
+                'stopped_due_to_exception',
+                self.worker_id,
+                (None if self.current_task is None else self.current_task.strong_id(), self.current_task_shots_left, self.unflushed_results, traceback.format_exc(), ex),
+            ))

sinter/_collection/_collection_worker_test.py

@@ -0,0 +1,222 @@
+import collections
+import multiprocessing
+import time
+from typing import Any, List
+
+import sinter
+import stim
+
+from sinter._collection._collection_worker_state import CollectionWorkerState
+
+
+class MockWorkHandler(sinter.Sampler, sinter.CompiledSampler):
+    def __init__(self):
+        self.expected_task = None
+        self.expected = collections.deque()
+
+    def compiled_sampler_for_task(self, task: sinter.Task) -> sinter.CompiledSampler:
+        assert task == self.expected_task
+        return self
+
+    def handles_throttling(self) -> bool:
+        return True
+
+    def sample(self, shots: int) -> sinter.AnonTaskStats:
+        assert self.expected
+        expected_shots, response = self.expected.popleft()
+        assert shots == expected_shots
+        return response
+
+
+def _assert_drain_queue(q: multiprocessing.Queue, expected_contents: List[Any]):
+    for v in expected_contents:
+        assert q.get(timeout=0.1) == v
+    assert q.empty()
+
+
+def _put_wait_not_empty(q: multiprocessing.Queue, item: Any):
+    q.put(item)
+    while q.empty():
+        time.sleep(0.0001)
+
+
+def test_worker_stop():
+    handler = MockWorkHandler()
+
+    inp = multiprocessing.Queue()
+    out = multiprocessing.Queue()
+    inp.cancel_join_thread()
+    out.cancel_join_thread()
+
+    worker = CollectionWorkerState(
+        flush_period=-1,
+        worker_id=5,
+        sampler=handler,
+        inp=inp,
+        out=out,
+        custom_error_count_key=None,
+    )
+
+    assert worker.process_messages() == 0
+    _assert_drain_queue(out, [])
+
+    t0 = sinter.Task(
+        circuit=stim.Circuit('H 0'),
+        detector_error_model=stim.DetectorErrorModel(),
+        decoder='mock',
+        collection_options=sinter.CollectionOptions(max_shots=100_000_000),
+        json_metadata={'a': 3},
+    )
+    handler.expected_task = t0
+
+    _put_wait_not_empty(inp, ('change_job', (t0, sinter.CollectionOptions(max_errors=100_000_000), 100_000_000)))
+    assert worker.process_messages() == 1
+    _assert_drain_queue(out, [('changed_job', 5, (t0.strong_id(),))])
+
+    _put_wait_not_empty(inp, ('stop', None))
+    assert worker.process_messages() == -1
+
+
+def test_worker_skip_work():
+    handler = MockWorkHandler()
+
+    inp = multiprocessing.Queue()
+    out = multiprocessing.Queue()
+    inp.cancel_join_thread()
+    out.cancel_join_thread()
+
+    worker = CollectionWorkerState(
+        flush_period=-1,
+        worker_id=5,
+        sampler=handler,
+        inp=inp,
+        out=out,
+        custom_error_count_key=None,
+    )
+
+    assert worker.process_messages() == 0
+    _assert_drain_queue(out, [])
+
+    t0 = sinter.Task(
+        circuit=stim.Circuit('H 0'),
+        detector_error_model=stim.DetectorErrorModel(),
+        decoder='mock',
+        collection_options=sinter.CollectionOptions(max_shots=100_000_000),
+        json_metadata={'a': 3},
+    )
+    handler.expected_task = t0
+    _put_wait_not_empty(inp, ('change_job', (t0, sinter.CollectionOptions(max_errors=100_000_000), 100_000_000)))
+    assert worker.process_messages() == 1
+    _assert_drain_queue(out, [('changed_job', 5, (t0.strong_id(),))])
+
+    _put_wait_not_empty(inp, ('accept_shots', (t0.strong_id(), 10000)))
+    assert worker.process_messages() == 1
+    _assert_drain_queue(out, [('accepted_shots', 5, (t0.strong_id(), 10000))])
+
+    assert worker.current_task == t0
+    assert worker.current_task_shots_left == 10000
+    assert worker.process_messages() == 0
+    _assert_drain_queue(out, [])
+
+    _put_wait_not_empty(inp, ('return_shots', (t0.strong_id(), 2000)))
+    assert worker.process_messages() == 1
+    _assert_drain_queue(out, [
+        ('returned_shots', 5, (t0.strong_id(), 2000)),
+    ])
+
+    _put_wait_not_empty(inp, ('return_shots', (t0.strong_id(), 20000000)))
+    assert worker.process_messages() == 1
+    _assert_drain_queue(out, [
+        ('returned_shots', 5, (t0.strong_id(), 8000)),
+    ])
+
+    assert not worker.do_some_work()
+
+
+def test_worker_finish_work():
+    handler = MockWorkHandler()
+
+    inp = multiprocessing.Queue()
+    out = multiprocessing.Queue()
+    inp.cancel_join_thread()
+    out.cancel_join_thread()
+
+    worker = CollectionWorkerState(
+        flush_period=-1,
+        worker_id=5,
+        sampler=handler,
+        inp=inp,
+        out=out,
+        custom_error_count_key=None,
+    )
+
+    assert worker.process_messages() == 0
+    _assert_drain_queue(out, [])
+
+    ta = sinter.Task(
+        circuit=stim.Circuit('H 0'),
+        detector_error_model=stim.DetectorErrorModel(),
+        decoder='mock',
+        collection_options=sinter.CollectionOptions(max_shots=100_000_000),
+        json_metadata={'a': 3},
+    )
+    handler.expected_task = ta
+    _put_wait_not_empty(inp, ('change_job', (ta, sinter.CollectionOptions(max_errors=100_000_000), 100_000_000)))
+    _put_wait_not_empty(inp, ('accept_shots', (ta.strong_id(), 10000)))
+    t0 = time.monotonic()
+    num_processed = 0
+    while True:
+        num_processed += worker.process_messages()
+        if num_processed >= 2:
+            break
+        if time.monotonic() - t0 > 1:
+            raise ValueError("Messages not processed")
+    assert num_processed == 2
+    _assert_drain_queue(out, [
+        ('changed_job', 5, (ta.strong_id(),)),
+        ('accepted_shots', 5, (ta.strong_id(), 10000)),
+    ])
+
+    assert worker.current_task == ta
+    assert worker.current_task_shots_left == 10000
+    assert worker.process_messages() == 0
+    _assert_drain_queue(out, [])
+
+    handler.expected.append((
+        10000,
+        sinter.AnonTaskStats(
+            shots=1000,
+            errors=23,
+            discards=0,
+            seconds=1,
+        ),
+    ))
+
+    assert worker.do_some_work()
+    worker.flush_results()
+    _assert_drain_queue(out, [
+        ('flushed_results', 5, (ta.strong_id(), sinter.AnonTaskStats(shots=1000, errors=23, discards=0, seconds=1)))])
+
+    handler.expected.append((
+        9000,
+        sinter.AnonTaskStats(
+            shots=9000,
+            errors=13,
+            discards=0,
+            seconds=1,
+        ),
+    ))
+
+    assert worker.do_some_work()
+    worker.flush_results()
+    _assert_drain_queue(out, [
+        ('flushed_results', 5, (ta.strong_id(), sinter.AnonTaskStats(
+            shots=9000,
+            errors=13,
+            discards=0,
+            seconds=1,
+        ))),
+    ])
+    assert not worker.do_some_work()
+    worker.flush_results()
+    _assert_drain_queue(out, [])

sinter/_collection/_mux_sampler.py

@@ -0,0 +1,56 @@
+import pathlib
+from typing import Optional
+from typing import Union
+
+from sinter._data import Task
+from sinter._decoding._decoding_all_built_in_decoders import BUILT_IN_SAMPLERS
+from sinter._decoding._decoding_decoder_class import Decoder
+from sinter._decoding._sampler import CompiledSampler
+from sinter._decoding._sampler import Sampler
+from sinter._decoding._stim_then_decode_sampler import StimThenDecodeSampler
+
+
+class MuxSampler(Sampler):
+    """Looks up the sampler to use for a task, by the task's decoder name."""
+
+    def __init__(
+        self,
+        *,
+        custom_decoders: Union[dict[str, Union[Decoder, Sampler]], None],
+        count_observable_error_combos: bool,
+        count_detection_events: bool,
+        tmp_dir: Optional[pathlib.Path],
+    ):
+        self.custom_decoders = custom_decoders
+        self.count_observable_error_combos = count_observable_error_combos
+        self.count_detection_events = count_detection_events
+        self.tmp_dir = tmp_dir
+
+    def compiled_sampler_for_task(self, task: Task) -> CompiledSampler:
+        return self._resolve_sampler(task.decoder).compiled_sampler_for_task(task)
+
+    def _resolve_sampler(self, name: str) -> Sampler:
+        sub_sampler: Union[Decoder, Sampler]
+
+        if name in self.custom_decoders:
+            sub_sampler = self.custom_decoders[name]
+        elif name in BUILT_IN_SAMPLERS:
+            sub_sampler = BUILT_IN_SAMPLERS[name]
+        else:
+            raise NotImplementedError(f'Not a recognized decoder or sampler: {name=}. Did you forget to specify custom_decoders?')
+
+        if isinstance(sub_sampler, Sampler):
+            if self.count_detection_events:
+                raise NotImplementedError("'count_detection_events' not supported when using a custom Sampler (instead of a custom Decoder).")
+            if self.count_observable_error_combos:
+                raise NotImplementedError("'count_observable_error_combos' not supported when using a custom Sampler (instead of a custom Decoder).")
+            return sub_sampler
+        elif isinstance(sub_sampler, Decoder) or hasattr(sub_sampler, 'compile_decoder_for_dem'):
+            return StimThenDecodeSampler(
+                decoder=sub_sampler,
+                count_detection_events=self.count_detection_events,
+                count_observable_error_combos=self.count_observable_error_combos,
+                tmp_dir=self.tmp_dir,
+            )
+        else:
+            raise NotImplementedError(f"Don't know how to turn this into a Sampler: {sub_sampler!r}")

sinter/_collection/_printer.py

@@ -0,0 +1,65 @@
+import threading
+from typing import List, Any
+
+import sys
+import time
+
+
+class ThrottledProgressPrinter:
+    """Handles printing progress updates interspersed amongst output.
+
+    Throttles the progress updates to not flood the screen when 100 show up
+    at the same time, and instead only show the latest one.
+    """
+    def __init__(self, *, outs: List[Any], print_progress: bool, min_progress_delay: float):
+        self.outs = outs
+        self.print_progress = print_progress
+        self.next_can_print_time = time.monotonic()
+        self.latest_msg = ''
+        self.latest_printed_msg = ''
+        self.min_progress_delay = min_progress_delay
+        self.is_worker_running = False
+        self.lock = threading.Lock()
+
+    def print_out(self, msg: str) -> None:
+        with self.lock:
+            for out in self.outs:
+                print(msg, file=out, flush=True)
+
+    def show_latest_progress(self, msg: str) -> None:
+        if not self.print_progress:
+            return
+        with self.lock:
+            if msg == self.latest_msg:
+                return
+            self.latest_msg = msg
+            if not self.is_worker_running:
+                dt = self._try_print_else_delay()
+                if dt > 0:
+                    self.is_worker_running = True
+                    threading.Thread(target=self._print_worker).start()
+
+    def flush(self):
+        with self.lock:
+            if self.latest_msg != "" and self.latest_printed_msg != self.latest_msg:
+                print('\033[31m' + self.latest_msg + '\033[0m', file=sys.stderr, flush=True)
+                self.latest_printed_msg = self.latest_msg
+
+    def _try_print_else_delay(self) -> float:
+        t = time.monotonic()
+        dt = self.next_can_print_time - t
+        if dt <= 0:
+            self.next_can_print_time = t + self.min_progress_delay
+            self.is_worker_running = False
+            if self.latest_msg != "" and self.latest_msg != self.latest_printed_msg:
+                print('\033[31m' + self.latest_msg + '\033[0m', file=sys.stderr, flush=True)
+                self.latest_printed_msg = self.latest_msg
+        return max(dt, 0)
+
+    def _print_worker(self):
+        while True:
+            with self.lock:
+                dt = self._try_print_else_delay()
+            if dt == 0:
+                break
+            time.sleep(dt)

sinter/_collection/_sampler_ramp_throttled.py

@@ -0,0 +1,66 @@
+import time
+
+from sinter._decoding import Sampler, CompiledSampler
+from sinter._data import Task, AnonTaskStats
+
+
+class RampThrottledSampler(Sampler):
+    """Wraps a sampler to adjust requested shots to hit a target time.
+
+    This sampler will initially only take 1 shot per call. If the time taken
+    significantly undershoots the target time, the maximum number of shots per
+    call is increased by a constant factor. If it exceeds the target time, the
+    maximum is reduced by a constant factor. The result is that the sampler
+    "ramps up" how many shots it does per call until it takes roughly the target
+    time, and then dynamically adapts to stay near it.
+    """
+
+    def __init__(self, sub_sampler: Sampler, target_batch_seconds: float, max_batch_shots: int):
+        self.sub_sampler = sub_sampler
+        self.target_batch_seconds = target_batch_seconds
+        self.max_batch_shots = max_batch_shots
+
+    def __str__(self) -> str:
+        return f'CompiledRampThrottledSampler({self.sub_sampler})'
+
+    def compiled_sampler_for_task(self, task: Task) -> CompiledSampler:
+        compiled_sub_sampler = self.sub_sampler.compiled_sampler_for_task(task)
+        if compiled_sub_sampler.handles_throttling():
+            return compiled_sub_sampler
+
+        return CompiledRampThrottledSampler(
+            sub_sampler=compiled_sub_sampler,
+            target_batch_seconds=self.target_batch_seconds,
+            max_batch_shots=self.max_batch_shots,
+        )
+
+
+class CompiledRampThrottledSampler(CompiledSampler):
+    def __init__(self, sub_sampler: CompiledSampler, target_batch_seconds: float, max_batch_shots: int):
+        self.sub_sampler = sub_sampler
+        self.target_batch_seconds = target_batch_seconds
+        self.batch_shots = 1
+        self.max_batch_shots = max_batch_shots
+
+    def __str__(self) -> str:
+        return f'CompiledRampThrottledSampler({self.sub_sampler})'
+
+    def sample(self, max_shots: int) -> AnonTaskStats:
+        t0 = time.monotonic()
+        actual_shots = min(max_shots, self.batch_shots)
+        result = self.sub_sampler.sample(actual_shots)
+        dt = time.monotonic() - t0
+
+        # Rebalance number of shots.
+        if self.batch_shots > 1 and dt > self.target_batch_seconds * 1.3:
+            self.batch_shots //= 2
+        if result.shots * 2 >= actual_shots:
+            for _ in range(4):
+                if self.batch_shots * 2 > self.max_batch_shots:
+                    break
+                if dt > self.target_batch_seconds * 0.3:
+                    break
+                self.batch_shots *= 2
+                dt *= 2
+
+        return result