调试 MLflow Tracing 中真实的并发死锁

Press enter or click to view image in full size Image by author Data Science Open Source Mlflow AI Python Debugging a Real Concurrency Deadlock in MLflow Tracing How a late enqueue could hang flush(), why the first explanation was wrong, and what reviewer-driven debugging taught us about concurrent queue design Mukesh Bhandarkar 11 min read · Just now -- Listen Share TL;DR — The deadlock came from a lifecycle race in MLflow’s AsyncTraceExportQueue. flush() could finish with the consumer thread, and then a concurrent put() could still add one more task to the queue. That extra put() incremented unfinished_tasks, but there was no consumer left to call task_done(), so queue.join() had no way to finish. I fixed the lifecycle handoff so put() and flush() make that decision under the same lock, late tasks fall back to synchronous handling, and concurrent flush() calls are serialized with a separate _flush_lock. I reproduced the race with both terminate=False and terminate=True. The bug looked like a blocked producer. It wasn’t. What made this bug confusing at first was that the visible symptom and the actual failure point were in different threads. The scenario I was chasing was simple: flush() could stay alive forever when another thread called put() during shutdown. My first instinct was to look inside put(), because it contained a threading.Event.wait() call that looked like it could block the producer. That turned out to be the wrong place to focus. An Event that is already set does not wait; wait() returns immediately. The producer was getting past that line. The hang happened later, when flush() reached queue.join() with Queue.unfinished_tasks still above zero and no consumer left that could bring the counter back down. The rest of this article follows how I went from that wrong first diagnosis to a deterministic reproduction, the first locking fix, the regression review that exposed it, and the two-lock design that was eventually merged. Figure 1 is a quick primer on the few Python and MLflow pieces involved. Press enter or click to view image in full size Figure 1. Tools at a glance: MLflow tracing, Python’s queue.Queue, threading.Event/RLock, and pytest’s parametrize — the four pieces this bug sits between. A quick mental model of MLflow’s async tracing queue The queue has two stages of asynchronous work. Application code calls AsyncTraceExportQueue.put() instead of exporting a trace inline. A dedicated consumer thread removes those tasks from a Python queue. Queue and hands the actual export work to a thread pool. That keeps normal tracing calls from waiting on every export operation. Press enter or click to view image in full size Figure 2. Introductory architecture: the normal path, the race before the fix, and the corrected lifecycle. flush() is the other side of that lifecycle. It stops and drains the async machinery and then waits for the queue’s accounting to settle. queue. Queue tracks that accounting with unfinished_tasks: put() increments it, and task_done() decrements it after queued work has been taken care of. queue.join() does not return until that counter reaches zero. The failure window The failure needed one unlucky ordering. flush() could finish with the consumer thread, while put() still had enough time to treat the queue as available. That gap was the whole bug: the producer could hand the queue one more task after the only thread responsible for consuming queued work had already gone away. 1. flush() initiates its action and requests the consumer to exit. 2. The consumer drains the queue, sees it empty, and exits. 3. flush() successfully joins the consumer thread. 4. However, a concurrent put() still sees the queue as active and enqueues a new task. 5. Queue.put() increments unfinished_tasks. 6. There is no consumer anymore who processes the task and calls task_done(). 7. flush() now enters queue.join() and waits forever. What makes this hang easy to misread is that queue.join() is not asking whether the queue happens to look empty at that instant. It is waiting for bookkeeping to balance. Every successful put() increments unfinished_tasks, and some worker path must eventually answer with task_done(). In this interleaving, the late put() increments the counter after the consumer has exited. Nothing remains that can perform the matching task_done(), so the condition queue.join() is waiting for has become impossible to satisfy. The exact pre-fix put() path, and what it actually does: # mlflow/tracing/export/async_export_queue.py, before the fix def put(self, task): if self._stop_event.is_set(): self._stop_event.wait() # queue.unfinished_tasks == 1 # flush() returns: False <- deadlocked, both terminate=True/False Why the Event.wait() explanation was misleading The prefix code checked whether the stop event was set and then called threading.Event.wait(). That looked like a plausible place for put() to block, but UgaTheDev tested the path independently during review. With the event already set, wait() returned in about 3 microseconds, consistent with Python’s documented Event semantics. That ruled out the producer-side wait and shifted the diagnosis to queue.join() and its unfinished-task accounting threading.Event.wait() I updated the PR description after that review to describe the failure we could actually reproduce: a late enqueue could leave unfinished_tasks non-zero after the consumer had exited. Getting that explanation right mattered beyond the patch itself; a misleading causal story makes future changes to concurrent code harder to reason about. Making the race deterministic I did not want this test to depend on getting lucky with thread scheduling. The useful place to stop flush() was after the consumer thread had already been joined but before queue.join() began waiting. I held flush() there and let another thread call put(). That forced the exact ordering I was trying to investigate instead of running the same test repeatedly and hoping the race would eventually appear. On the old implementation, the result was repeatable: put() came back in roughly 0.000 seconds, unfinished_tasks moved to 1, and nothing was left to process that task. flush() then stayed blocked in queue.join(). I reproduced the same failure with terminate=False and terminate=True. The terminate=True path was worth checking separately because MLflow’s atexit handler uses it during interpreter shutdown. Designing the fix around an invariant The most elegant approach to solving the problem was to state what should always be the case. The producer should be guaranteed to end up on only one side of the flush lifecycle transition point. If put() takes place before the transition point, the task may get enqueued, since a live consumer will have to deal with the queued work. If put() happens after the transition, it may not enter a queue whose consumer is going away. Instead, it will be processed synchronously. The existing reentrant lifecycle lock was leveraged to provide mutual exclusion between the producer enqueuing process and the transition that makes the queue inactive and causes the stop event. This matters because the lock must be released before joining threads or the worker pool. Otherwise, another type of deadlock will arise. The final structure grew out of the review discussion. UgaTheDev proposed the _flush_lock change after finding the concurrent-flush regression, and TomeHirata suggested simplifying put() by removing the temporary handle_synchronously flag and returning early while holding the lifecycle lock. I incorporated both ideas into the PR. The excerpt below reflects that final structure (mlflow/tracing/export/async_export_queue.py): def put(self, task: Task): """ Put a new task to the queue for processing.""" with self._lock: if not self.is_active(): if not self._stop_event.is_set(): self.activate() if self.is_active(): try: self._queue.put(task, block=False) except queue_Full: ... # warn on a full queue, unrelated to this bug return # Queue is flushing or was terminated; execute synchronously. task.handle() # in __init__: self._flush_lock = threading.Lock() def flush(self, terminate: bool = False): with self._flush_lock: # NEW: serializes whole flush() with self._lock: if not self.is_active(): return self._is_active = False self._stop_event.set() # _lock released here -- safe to block on the joins below self._consumer_thread.join() self._queue.join() self._worker_threadpool.shutdown(wait=True) if not terminate: self._stop_event.clear() self.activate() Then the first fix created a second concurrency bug My first fix closed the late-enqueue window by setting _is_active = False as soon as flush() entered its lifecycle transition. For producers, that was exactly what I needed: once flushing had started, a new put() could no longer assume that a consumer would still be there to drain the queue. Get Mukesh Bhandarkar’s stories in your inbox Join Medium for free to get updates from this writer. Subscribe Subscribe Remember me for faster sign in But that state flag was doing more than one job. flush() itself already had an early return for an inactive queue. With the flag now flipping earlier, a second flush() arriving while the first one was still draining could see _is_active == False and return immediately. The original deadlock was gone, but the second caller could report that flushing was finished while the first flush was still in progress. That was the regression UgaTheDev caught during review. Why the final design uses two locks The fix was to stop asking one lock and one state flag to represent two different kinds of coordination. self._lock, the existing RLock, stays around the short lifecycle decision shared by put() and flush(): can this task still enter the async queue, or has the transition already started? That section has to stay small because producers need it too. The new _flush_lock has a different lifetime. It is held for the whole flush() operation, including the blocking joins. If another thread calls flush() during that time, it waits for the first call rather than inspecting a temporary inactive state and returning early. In other words, _lock protects a state transition; _flush_lock serializes an operation. Treating those as separate responsibilities resolved the second race without reopening the first. Press enter or click to view image in full size Figure 3. Two locks, two jobs: what each one protects, and how long it’s held. Testing the fix — and testing the tests I added tests for both failure modes instead of treating the original race as the only thing worth protecting. test_concurrent_put_during_flush_does_not_deadlock runs with terminate=False and terminate=True, so the late-put case is covered for a normal flush and for the termination path. test_flush_waits_for_in_progress_flush covers the regression review found by putting two flush callers in competition and checking that the second one does not finish ahead of the first. I also wanted evidence that those tests were actually guarding the behavior they claimed to guard. Removing _flush_lock makes the concurrent-flush test fail, and running the new race coverage against the old implementation reproduces the original hang. That gave me more confidence than a green run on the final code alone, because it showed that the tests could distinguish the broken implementations from the fixed one. Reconstructed to match the real test names and file — tests/tracing/export/test_async_export_queue.py: @pytest.mark.parametrize("terminate", [False, True]) def test_concurrent_put_during_flush_does_not_deadlock(terminate): queue = AsyncTraceExportQueue() pause_flush_before_queue_join(queue) # widen the window flush_thread = Thread(target=queue.flush, kwargs={'terminate': terminate}) flush_thread.start() wait_for_consumer_thread_to_exit(queue) queue.put(make_task()) # the late enqueue flush_thread.join(timeout=5) assert not flush_thread.is_alive() # hangs pre-fix, both terminate values def test_flush_waits_for_in_progress_flush(): queue = AsyncTraceExportQueue() slow_first_flush = Thread(target=queue.flush) slow_first_flush.start() returned_early = run_second_flush_and_check_timing(queue) assert not returned_early # needs the _flush_lock Targeted validation for the final patch included: 8 focused tracing export tests passing Ruff lint passing Ruff format check passing MLflow Clint checks passing git diff — check clean Review was part of the debugging process Review changed both the diagnosis and the final shape of the patch. UgaTheDev first reproduced the original hang independently and showed that Event.wait() was not the blocking point. They traced the deadlock to queue.join() and unfinished_tasks, which led me to correct the PR description. They then tried to break the proposed locking model in two concrete ways: checking whether _lock remained held across the consumer-thread join, and checking whether releasing it before task.handle() could reorder synchronous and asynchronous completion. The first concern did not occur; the second ordering can occur, but the queue did not guarantee that ordering in the first place. The same review uncovered a separate regression in concurrent flush calls. UgaTheDev backed that finding with before-and-after timing and proposed a ready-to-apply _flush_lock diff. I added the flush lock and updated the patch. TomeHirata later proposed the cleaner early-return form of put(); I adopted that structure, rebased the PR, and pushed the final commits. TomeHirata then reviewed, approved, and squash-merged the change. The sequence matters: reviewers diagnosed problems and proposed improvements, and I incorporated the fixes and shipped them through my PR. That, honestly, is one of the best parts of working in the open: review isn’t just an approval gate. When reviewers reproduce behavior and challenge invariants, they turn into another debugging instrument. Press enter or click to view image in full size Figure 4. End-to-end contribution workflow, from issue discovery through reproduction, design, review, validation, and merge. Five lessons that generalize beyond MLflow 1. Follow ownership, not just lock acquisition. The important question in this bug was who still owned responsibility for completing a task. The queue accepted one more task after the only consumer that could finish that ownership chain had gone away. 2. Be exact about the primitives. My first diagnosis leaned too heavily on what Event.wait() looked like it might do. Its actual semantics ruled that theory out quickly once we checked them. 3. Audit every reader of shared lifecycle state. Moving the _is_active transition earlier solved the producer race, but another flush caller was reading that same flag for a different reason. Changing when a flag flips can change the meaning of every branch that reads it. 4. Match lock scope to the responsibility. The RLock only needs to protect a short atomic lifecycle decision. Serializing a complete flush, including joins, is a separate problem and belongs to a separate lock. 5. Make regression tests prove something. I found it useful to run the new tests against the old behavior and to remove the new synchronization selectively. A test is much more convincing when you have seen it fail for the exact regression it is meant to prevent. The result TomeHirata merged the finished patch into mlflow: master via PR #25264, closing issue #25237; the squash-merged commit is 89a0b7e. It did not change the public API. What changed was the shutdown behavior under concurrency: producers can no longer leave unfinishable queue bookkeeping behind during flush(), and overlapping flush callers now wait for one another correctly. For me, the useful part of the contribution was the path to that result. I had to force the interleaving instead of guessing at it, correct my own explanation when the evidence disagreed, update the locking model after review exposed another contract, and then make the tests demonstrate both failures. The patch itself is small compared with the amount of reasoning needed to be confident in it. Closing thought The question that finally made this bug easier for me was not ‘which thread is stuck?’ but ‘what can no longer happen?’ queue.join() was waiting for unfinished_tasks to reach zero, yet the component that could make that happen had already exited. Once I looked at the race as a broken ownership invariant instead of a mysterious blocked thread, both the failure and the shape of the fix became much easier to see. About the author Mukesh Bhandarkar is a software and ML systems practitioner interested in ML infrastructure, MLOps, distributed/concurrent systems, and open-source engineering. References [1] MLflow contributors, mlflow/mlflow (2026), GitHub [2] M. Bhandarkar, fix: prevent async trace queue deadlock during flush (#25264) (2026), mlflow/mlflow, GitHub [3] MLflow contributors, async_export_queue.py (commit 89a0b7e, 2026), mlflow/mlflow, GitHub [4] Python Software Foundation, queue — A synchronized queue class (2026), Python 3 Documentation [5] Python Software Foundation, threading — Thread-based parallelism (2026), Python 3 Documentation