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
Describe the bug
braket.tracking.tracking_context.TrackingContextstores active trackers in a plainsetand iterates it without any synchronization. Because a single module-global instance (_tracking_context) is shared across the whole process, andbroadcast_event()is invoked on everyAwsSession.get_quantum_task()call, concurrent use from multiple threads races: one thread iterating_trackersinbroadcast_event()while another thread registers/deregisters aTracker(e.g. entering/exiting aTracker()context) mutates the set mid-iteration.Since
AwsSession.get_quantum_task()callsbroadcast_event(_TaskStatusEvent(...)), any application that (a) submits/polls Braket tasks concurrently across threads while (b) anyTrackeris active will intermittently crash with:To reproduce
Minimal deterministic reproduction of the underlying race (no AWS calls needed):
This raises
RuntimeError: Set changed size during iterationfrombroadcast_event. In production the mutating thread is any concurrentTrackerenter/exit and the iterating thread is a concurrentget_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 andTrackerlifecycle 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.:
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 (noRuntimeErrorin 1,000,000 iterations under the same concurrent churn).Environment