Skip to content

TrackingContext is not thread-safe: broadcast_event iterates _trackers without a lock → RuntimeError: Set changed size during iteration #1319

Description

@ryanhill1

Describe the bug

braket.tracking.tracking_context.TrackingContext stores active trackers in a plain set and iterates it without any synchronization. Because a single module-global instance (_tracking_context) is shared across the whole process, and broadcast_event() is invoked on every AwsSession.get_quantum_task() call, concurrent use from multiple threads races: one thread iterating _trackers in broadcast_event() while another thread registers/deregisters a Tracker (e.g. entering/exiting a Tracker() context) mutates the set mid-iteration.

# braket/tracking/tracking_context.py
class TrackingContext:
    def __init__(self):
        self._trackers = set()

    def register_tracker(self, tracker):
        self._trackers.add(tracker)        # mutates

    def deregister_tracker(self, tracker):
        self._trackers.remove(tracker)     # mutates

    def broadcast_event(self, event):
        for tracker in self._trackers:     # unsynchronized iteration  <-- races
            tracker.receive_event(event)

_tracking_context = TrackingContext()      # process-global singleton

Since AwsSession.get_quantum_task() calls broadcast_event(_TaskStatusEvent(...)), any application that (a) submits/polls Braket tasks concurrently across threads while (b) any Tracker is active will intermittently crash with:

File ".../braket/aws/aws_session.py", line 326, in get_quantum_task
    broadcast_event(_TaskStatusEvent(arn=response["quantumTaskArn"], status=response["status"]))
File ".../braket/tracking/tracking_context.py", line 43, in broadcast_event
    for tracker in self._trackers:
RuntimeError: Set changed size during iteration

To reproduce

Minimal deterministic reproduction of the underlying race (no AWS calls needed):

import threading
from braket.tracking.tracker import Tracker
from braket.tracking.tracking_context import broadcast_event
from braket.tracking.tracking_events import _TaskStatusEvent

stop = False

def churn():
    # Simulates Tracker context managers entering/exiting on worker threads.
    while not stop:
        with Tracker():
            pass

t = threading.Thread(target=churn, daemon=True)
t.start()

with Tracker():  # keep the global set non-empty
    event = _TaskStatusEvent(arn="arn:aws:braket:us-east-1:000:quantum-task/x", status="COMPLETED")
    for _ in range(1_000_000):
        broadcast_event(event)   # iterates _trackers while churn() mutates it

stop = True

This raises RuntimeError: Set changed size during iteration from broadcast_event. In production the mutating thread is any concurrent Tracker enter/exit and the iterating thread is a concurrent get_quantum_task() (task submission or status polling).

Validated on amazon-braket-sdk 1.110.1 / Python 3.12: the snippet above raised on every run, at iterations 13,191 / 157,252 / 7,600 across three runs (the exact point varies with thread timing).

Expected behavior

Cost tracking should be thread-safe. Concurrent get_quantum_task() calls and Tracker lifecycle events across threads should not crash, since neither is documented as single-threaded-only.

Suggested fix

Guard the shared set with a lock and iterate over a snapshot, e.g.:

import threading

class TrackingContext:
    def __init__(self):
        self._trackers = set()
        self._lock = threading.Lock()

    def register_tracker(self, tracker):
        with self._lock:
            self._trackers.add(tracker)

    def deregister_tracker(self, tracker):
        with self._lock:
            self._trackers.discard(tracker)   # discard avoids KeyError on double-deregister

    def broadcast_event(self, event):
        with self._lock:
            trackers = list(self._trackers)   # snapshot; release lock before receive_event
        for tracker in trackers:
            tracker.receive_event(event)

Iterating a snapshot (rather than holding the lock across receive_event) keeps the critical section small and avoids holding a lock across tracker callbacks. Applying this implementation to the reproduction above eliminates the crash (no RuntimeError in 1,000,000 iterations under the same concurrent churn).

Environment

  • amazon-braket-sdk-python: 1.110.1
  • Python: 3.12
  • OS: reproduced on macOS; originally observed in a concurrent thread-pool service on Linux

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions