PyTorch #182948 — Why is_initialized() Was Not Enough to Prevent _lazy_call() Deadlocks

PyTorch Lazy Initialization, Part 2 of 3 — The fully initialized path, the run_now lock race, and the thread-local state between native initialization and Python completion

Part 1 examined the immediate cause of PyTorch’s _lazy_call() deadlock.

The old implementation could execute a callback while holding:

_initialization_lock

That lock was a non-reentrant threading.Lock.

If the callback invoked another operation that eventually called _lazy_call() again, the same thread attempted to acquire the same lock twice.

Outer _lazy_call
        ↓
Acquire initialization lock
        ↓
Execute callback
        ↓
Callback re-enters _lazy_call
        ↓
Attempt to acquire the same lock
        ↓
Wait forever

Moving callback execution outside the lock fixes the obvious case in which CUDA or ROCm is already initialized.

But PR #182948 had to solve a second path that looked similar from the outside and was structurally different underneath.

During _lazy_init(), PyTorch had already completed the native backend initialization:

torch._C._cuda_init()

but had not yet set:

_initialized = True

The initialization thread still held _initialization_lock while draining queued Python callbacks.

A queued callback could therefore re-enter _lazy_call() while:

Native CUDA/ROCm initialization
→ complete

Public Python initialized flag
→ still False

Initialization lock
→ still held by the current thread

A simple check of is_initialized() could not distinguish this transitional state from a backend that had not started initialization at all.

The final fix therefore uses three mechanisms:

is_initialized()
→ stable fully initialized path

run_now
→ initialization completed between the first check and lock acquisition

_tls.is_initializing
→ current thread is inside the transitional queue-draining phase

These are not three versions of the same check.

Each protects a different race.


The fixed _lazy_call() is a small state machine

The current function is effectively:

def _lazy_call(callback, **kwargs):
    if is_initialized() or hasattr(_tls, "is_initializing"):
        callback()
        return

    run_now = False

    with _initialization_lock:
        if is_initialized():
            run_now = True
        else:
            queue(callback)

    if run_now:
        callback()

The decision tree is:

Is the backend fully initialized?
        │
        ├─ Yes
        │   → run callback immediately
        │   → no initialization lock
        │
        └─ No
            ↓
Is this thread currently performing initialization?
        │
        ├─ Yes
        │   → run callback immediately
        │   → do not reacquire the lock this thread already owns
        │
        └─ No
            ↓
Acquire initialization lock
            ↓
Did another thread finish initialization while we were waiting?
        │
        ├─ Yes
        │   → record run_now
        │   → release lock
        │   → run callback
        │
        └─ No
            → queue callback

The implementation performs callback execution outside _lazy_call()’s own critical section, while allowing the current initialization thread to re-enter without acquiring the lock again.


There are more than two relevant initialization states

A Boolean model suggests only:

Not initialized

Initialized

The actual lifecycle contains at least three operational states.

State 1 — Not initialized

Native backend
→ not ready

Queued Python callbacks
→ not yet processed

_initialized
→ False

_tls.is_initializing
→ absent

A new lazy request should be queued.


State 2 — Native initialization is complete, but queued callbacks are still being drained

torch._C._cuda_init()
→ complete

Queued Python callbacks
→ currently executing

_initialized
→ False

_tls.is_initializing
→ present on the initializer thread

_initialization_lock
→ held

A reentrant callback on the initializer thread must be allowed to execute without acquiring the lock again.

Other threads must still remain outside.


State 3 — Fully initialized

Native backend
→ ready

Queued callbacks
→ completed

_initialized
→ True

_tls.is_initializing
→ absent

_initialization_lock
→ released

New callbacks can run immediately without initialization coordination.

The important state was State 2.

is_initialized() alone reports:

False

for both State 1 and State 2.

But the correct _lazy_call() behavior is different in each state.

State 1
→ queue the callback

State 2, initializer thread
→ execute the callback immediately

Path 1: the backend was already initialized before _lazy_call() began

The simplest repaired path begins after initialization has completed.

Consider the first regression script:

import torch

torch.cuda.init()

state = torch.cuda.get_rng_state()

torch.cuda._lazy_call(
    lambda: torch.cuda.set_rng_state(state)
)

print("done")

The outer _lazy_call() receives a callback that calls:

torch.cuda.set_rng_state(state)

set_rng_state() itself creates another callback and submits it through _lazy_call(). The same pattern is used by manual_seed(), which creates a generator-seeding callback and calls _lazy_call(..., seed=True).

Before the fix:

Outer _lazy_call
        ↓
Acquire initialization lock
        ↓
is_initialized() is True
        ↓
Execute outer callback under lock
        ↓
set_rng_state()
        ↓
Inner _lazy_call
        ↓
Attempt same lock
        ↓
Deadlock

After the fix:

Outer _lazy_call
        ↓
is_initialized() is True
        ↓
Execute outer callback without lock
        ↓
set_rng_state()
        ↓
Inner _lazy_call
        ↓
is_initialized() is True
        ↓
Execute nested callback without lock
        ↓
Return

No recursive lock acquisition occurs.


The pre-lock fast path removes the common lock entirely

The first line of the new decision logic is:

if is_initialized() or hasattr(_tls, "is_initializing"):
    callback()
    return

For a stable initialized process, _lazy_call() no longer acquires _initialization_lock simply to learn that initialization has already finished.

Before

Acquire lock
→ inspect state
→ execute callback
→ release lock
After

Inspect stable state
→ execute callback

This solves two problems at once.

Reentrancy

The callback can call _lazy_call() again without recursively acquiring the initialization lock.

Lock duration

Arbitrary callback work no longer extends the initialized-path critical section.

A callback may perform RNG work, allocate resources, invoke other framework APIs, block, or raise an exception.

None of that now happens while _lazy_call() owns the initialization lock.


The first state check cannot be the only state check

A check performed before acquiring a lock is not stable by itself.

Consider two threads.

Thread A enters _lazy_call()

is_initialized()
→ False

Before Thread A acquires the lock, it is descheduled.

Thread B initializes the backend

Acquire initialization lock

Run native initialization

Drain queued callbacks

Set _initialized = True

Release initialization lock

Thread A resumes

Thread A now acquires the lock.

The world has changed since its first check.

Initial check:
not initialized

Inside-lock check:
initialized

The function therefore needs a second is_initialized() check after acquiring the lock.

This is the same double-checking structure used by _lazy_init() itself.


Why the callback cannot run inside the second check

Suppose Thread A performed this:

with _initialization_lock:
    if is_initialized():
        callback()

That would recreate the original deadlock boundary.

The callback could re-enter _lazy_call() while the lock was held.

The corrected code instead records:

run_now = True

inside the lock.

Then:

if run_now:
    callback()

runs only after the with block has ended.

Acquire lock
        ↓
Discover that initialization already completed
        ↓
Set run_now
        ↓
Release lock
        ↓
Execute callback

run_now is therefore not merely a style choice.

It transfers callback execution across the lock boundary.


Why the callback should not be queued after initialization has completed

When Thread A obtains the lock and sees:

is_initialized() == True

the initialization queue has already been drained.

Appending the callback to _queued_calls at that point would not match the current lifecycle.

Initialization completed

Queue already processed

New callback appended afterward

There may be no future initialization pass that naturally drains it.

The correct action is:

Run the callback now

but only after releasing the lock.

run_now preserves both conditions:

Do not lose the callback

Do not execute it under the lock

run_now solves a race between threads

The run_now branch exists for a specific interleaving:

Thread A
→ checks before lock
→ sees not initialized

Thread B
→ completes initialization

Thread A
→ acquires lock
→ sees initialized

It is not the mechanism that solves reentrancy inside the initialization thread.

That second path occurs while _initialized is still false.

For that path, run_now is never reached because attempting to enter the lock would already deadlock.

The thread-local marker must intercept the call before lock acquisition.


Path 2: a queued callback re-entered during _lazy_init()

The second regression script is:

import torch

torch.cuda._lazy_call(
    lambda: torch.cuda.manual_seed(42)
)

torch.cuda.init()

print("done")

The first call occurs before initialization.

The callback is queued.

Backend not initialized
        ↓
Acquire lock
        ↓
Store callback
        ↓
Release lock

Later, torch.cuda.init() enters _lazy_init().

The relevant sequence is:

Acquire _initialization_lock
        ↓
Run torch._C._cuda_init()
        ↓
Set _tls.is_initializing = True
        ↓
Add deferred seed callbacks
        ↓
Execute queued callbacks
        ↓
Delete _tls.is_initializing
        ↓
Set _initialized = True
        ↓
Release lock

The implementation sets the public _initialized flag only after queued callbacks have completed.


The queued callback runs while _lazy_init() still owns the lock

The patch does not move the entire queue-draining loop outside _initialization_lock.

That distinction matters.

During queue draining:

Initializer thread
→ holds _initialization_lock

Queued callback
→ runs under that outer lock

The lock prevents unrelated threads from entering the incomplete initialization state.

The source comment is explicit about the requirement:

Do not let other threads enter while initialization remains incomplete.

The fix therefore cannot simply release the lock for every queued callback without redesigning the broader initialization protocol.

Instead, it allows only the current initializer thread to re-enter _lazy_call() without reacquiring the lock.


The old queued path deadlocked because _initialized was still false

The queued callback invokes:

torch.cuda.manual_seed(42)

manual_seed() builds another callback and calls:

_lazy_call(cb, seed=True)

At that point:

_initialized
→ False

The old _lazy_call() therefore did not execute the nested callback immediately.

It attempted:

with _initialization_lock:

But _lazy_init() already held that lock on the same thread.

_lazy_init owns lock
        ↓
Queued callback runs
        ↓
manual_seed()
        ↓
Nested _lazy_call()
        ↓
Attempts lock owned by current _lazy_init()
        ↓
Self-deadlock

The backend’s native initialization had already completed.

The Python initialization protocol had not.

That was the transitional state missing from the old decision logic.


_tls.is_initializing is a selective reentrancy permit

Immediately after:

torch._C._cuda_init()

the initializer thread sets:

_tls.is_initializing = True

The new _lazy_call() checks for the presence of this marker before attempting the lock.

if hasattr(_tls, "is_initializing"):
    callback()
    return

The nested path becomes:

Queued callback
        ↓
manual_seed()
        ↓
Nested _lazy_call()
        ↓
Current thread has is_initializing marker
        ↓
Run nested callback immediately
        ↓
Do not acquire lock

The outer _lazy_init() continues to own the lock.

The inner _lazy_call() no longer waits for it.

PR #182948 explicitly states that this is safe because torch._C._cuda_init() has already completed before the marker is installed.


The marker does not globally declare initialization complete

_tls is a threading.local() object.

The is_initializing attribute exists only in the initializer thread’s local view.

Initializer thread

hasattr(_tls, "is_initializing")
→ True

It can safely bypass the lock for a reentrant lazy callback.

Another thread

hasattr(_tls, "is_initializing")
→ False

It cannot use the transitional shortcut.

It attempts _initialization_lock and waits until initialization completes.

Initializer thread
→ permitted to re-enter locally

Other threads
→ still excluded globally

This is the key property.

The marker grants reentrancy to the lock-owning thread without publishing a false “fully initialized” state to every thread.


A global initialization_in_progress flag would be too broad

Imagine replacing the thread-local marker with:

initialization_in_progress = True

Every thread could see it.

If _lazy_call() treated that global flag as permission to execute immediately, unrelated threads could enter while:

  • queued callbacks were unfinished,

  • Python-side initialization state was incomplete,

  • _initialized remained false,

  • the initializer still owned the transition.

Thread A
→ performing initialization

Thread B
→ sees global initialization_in_progress
→ bypasses lock
→ executes callback too early

The thread-local marker avoids that.

Only the thread already inside the critical section
receives the reentrancy exception.

Why the attribute is deleted instead of set to False

The code tests:

hasattr(_tls, "is_initializing")

rather than:

_tls.is_initializing is True

The lifecycle is:

_tls.is_initializing = True

followed by:

delattr(_tls, "is_initializing")

The existence of the attribute is the state marker.

Attribute absent
→ this thread is not the initializer

Attribute present
→ this thread is inside initialization

Using presence avoids leaving a persistent thread-local field whose false value must be interpreted separately.


_lazy_init() uses the same marker to avoid recursive initialization

At the top of _lazy_init():

if is_initialized() or hasattr(_tls, "is_initializing"):
    return

This prevents a queued callback from recursively starting the entire initialization procedure again.

The same marker therefore protects two related boundaries.

Recursive _lazy_init()

Already the initializer
→ return
→ do not repeat native initialization

Reentrant _lazy_call()

Already the initializer
→ run callback
→ do not reacquire initialization lock

The meanings are different but coordinated.

_lazy_init
→ do not initialize again

_lazy_call
→ do not queue or lock again

Why _initialized is not set immediately after native initialization

The code could theoretically do this:

Run torch._C._cuda_init()

Set _initialized = True

Drain queued callbacks

Then the reentrant callback would see:

is_initialized() == True

and avoid the deadlock without a thread-local marker.

PR #182948 does not do that.

The implementation keeps _initialized false until every queued callback has completed successfully.

Native backend initialization
        ↓
Queued Python initialization work
        ↓
Only then:
_initialized = True

This preserves a stronger meaning for the public state:

PyTorch’s complete lazy-initialization procedure has finished, not merely the native C++ call.


A queued callback can still fail

During _lazy_init(), each queued callback runs inside:

try:
    queued_call()
except Exception as e:
    raise DeferredCudaCallError(...)

If one callback fails, control does not reach:

_initialized = True

The failure is reported with the traceback captured when the deferred callback was originally registered.

Setting _initialized early would make other code believe initialization had completed even though one of the deferred initialization actions had failed.

The thread-local marker provides the narrower state the implementation actually needs:

Native runtime ready enough for this thread’s reentrant callback

but

full PyTorch lazy initialization not yet successful

“Ready enough” and “fully initialized” are different contracts

During the queue-draining phase:

torch._C._cuda_init()
→ complete

This makes certain backend operations safe for the initializer thread.

But:

_initialized
→ False

because PyTorch’s complete deferred-work contract is unfinished.

The two facts can coexist.

Native readiness
≠
public initialization completion

The marker expresses native readiness locally.

The Boolean _initialized expresses full completion globally.


The final code distinguishes four _lazy_call() outcomes

The behavior can be summarized in a table.

ConditionLock behaviorCallback behavior
Fully initialized before entryNo lockRun immediately
Current thread is _lazy_init() ownerNo nested lockRun immediately
Initialization completes before lock is acquiredAcquire lock, set run_now, releaseRun after lock
Still uninitialized after lock is acquiredAcquire lockQueue callback

This table is the real contract of the fixed function.

A single:

if is_initialized()

could not represent all four outcomes.


Why run_now and _tls.is_initializing are not interchangeable

It may seem that either mechanism could be used everywhere.

They solve different timing relationships.

run_now

Another thread completed initialization.

This thread acquires the lock afterward.

_initialized is now True.

The lock can be acquired safely.

The callback must be moved outside it.

_tls.is_initializing

This same thread is currently inside _lazy_init().

It already owns the lock.

_initialized is still False.

Attempting the lock is itself the bug.

The callback must bypass lock acquisition entirely.

run_now
→ post-wait race between threads

is_initializing
→ same-thread reentrancy during transition

Why an RLock would not express this distinction

A reentrant lock could allow the initialization thread to acquire the lock again.

That would address the immediate self-deadlock.

It would not distinguish:

Fully initialized callbacks
→ should not need the lock at all

from:

Initializer-thread callbacks
→ may re-enter while the outer transition remains locked

It would also leave arbitrary callback execution inside a critical section on the common initialized path.

The landed design retains the ordinary non-reentrant lock and makes the callback boundary explicit.

Stable initialized path
→ callback outside lock

Transitional initializer-thread path
→ callback inside outer initialization phase,
   but without nested acquisition

The source proves the chosen design; the broader comparison with RLock is an architectural interpretation of that design.


The fix did not make every callback universally reentrant

The patch removes one specific lock cycle.

It does not guarantee that callbacks cannot:

  • recursively call themselves forever,

  • acquire another lock in an incompatible order,

  • wait for another thread that is waiting for initialization,

  • block on GPU work,

  • or deadlock through an unrelated resource.

The new guarantee is narrower:

A callback re-entering _lazy_call() no longer deadlocks solely because _initialization_lock is already owned by the same thread.

That is the contract demonstrated by the two regression tests.


The first test protects the stable initialized state

The test:

Initialize CUDA/ROCm

Read RNG state

Call _lazy_call(callback)

Callback calls set_rng_state

set_rng_state re-enters _lazy_call

protects:

is_initialized() fast path

callback outside lock

nested initialized callback

Before the fix, the subprocess timed out.

After the fix, it printed:

done

and exited successfully.


The second test protects the transitional initializer state

The second test:

Queue callback before initialization

Begin _lazy_init()

Complete native backend initialization

Set thread-local initializer marker

Drain queued callback

Callback calls manual_seed

manual_seed re-enters _lazy_call

protects:

_tls.is_initializing path

same-thread lock bypass

nested callback during queue drain

Before the fix, the nested call attempted the same lock held by _lazy_init().

After the fix, it executed immediately through the thread-local branch.


The tests deliberately use different RNG APIs

The initialized-path test uses:

set_rng_state()

The queued-path test uses:

manual_seed()

Both APIs eventually call _lazy_call(), but they exercise different surrounding behavior.

set_rng_state()

Backend already initialized

State object already available

Nested callback submitted immediately

manual_seed()

Initial callback queued before backend initialization

During queue drain, manual_seed creates another deferred-style callback

Thread-local marker converts it into immediate execution

The two tests are not duplicate versions of the same script.

They target different lifecycle positions.


The second callback executes inline during queue draining

When _tls.is_initializing is present, _lazy_call() invokes the nested callback directly.

Outer queued callback begins

Nested lazy operation begins

Nested callback completes

Outer queued callback resumes

The execution is depth-first.

The nested callback is not added to the end of the already-draining queue.

This ensures that the callback requested by the current initialization action completes as part of that action.

The patch does not define a universal ordering model for arbitrary recursive callbacks beyond this inline behavior.


Other threads still see one atomic initialization transition

While the initializer thread drains callbacks:

_initialization_lock
→ held

Another thread entering _lazy_call() sees:

is_initialized()
→ False

hasattr(its own _tls, "is_initializing")
→ False

It attempts the lock and waits.

Only after the initializer:

finishes queued callbacks

deletes its thread-local marker

sets _initialized = True

releases the lock

can the waiting thread proceed.

It then sees the initialized state inside the lock, records run_now, releases the lock, and invokes its callback.

Initializer thread
→ local reentrancy permitted

Other threads
→ global transition remains serialized

The patch preserves queue-or-run atomicity

The lock is still necessary for requests arriving while initialization status can change.

It protects the decision:

This callback belongs to the pre-initialization queue

or

Initialization has completed and this callback must run now

Without the lock, a callback could be:

  • queued after the queue was drained,

  • executed twice,

  • lost between state transitions,

  • or executed while another thread still considered initialization incomplete.

The fix reduces the lock’s scope.

It does not remove its state-transition role.


The central invariant after the patch

The repaired behavior can be stated as:

No callback is executed by _lazy_call() while _lazy_call() itself owns _initialization_lock.

For callbacks drained by _lazy_init(), a different invariant applies:

The initializer may continue to hold the outer lock, but a callback on that same thread may re-enter _lazy_call() without trying to acquire it again.

These two rules are complementary.

Post-initialization callback
→ move execution outside lock

Initialization-phase callback
→ preserve outer lock
→ bypass nested acquisition for owner thread

What the patch directly confirms

The source and tests support the following statements:

  • _initialization_lock remains a non-reentrant threading.Lock.

  • Fully initialized callbacks now execute before lock acquisition.

  • If initialization completes during the lock race, _lazy_call() records run_now and invokes the callback after releasing the lock.

  • _lazy_init() completes torch._C._cuda_init() before setting _tls.is_initializing.

  • _initialized remains false while queued callbacks are drained.

  • The initialization marker is thread-local.

  • Reentrant callbacks on the initializer thread execute immediately.

  • Other threads remain serialized by _initialization_lock.

  • Two separate subprocess tests protect the two paths.


What the patch does not confirm

The public evidence does not establish that:

  • every CUDA or ROCm initialization callback is safe to execute recursively,

  • arbitrary callback recursion terminates,

  • no other lock cycle exists around initialization,

  • setting _initialized later is the only possible correct design,

  • an RLock could never be used in a different redesign,

  • XPU or MTIA has reproduced the same deadlock,

  • the two subprocess tests cover every possible initialization interleaving,

  • callback ordering under all recursive combinations is formally specified.

The confirmed result is narrower:

The final state machine prevents the two documented _lazy_call() self-deadlocks while preserving the initialization lock for cross-thread serialization.


A Boolean was not enough because readiness was local before it was global

The most important distinction in this patch is:

Backend ready for the initializer thread

versus:

Initialization complete for every thread

During queue draining:

Native runtime
→ ready

Current initializer thread
→ allowed to execute nested callback

Other threads
→ still blocked

Public initialized flag
→ still False

That state cannot be represented by one process-wide Boolean without losing either safety or reentrancy.

The thread-local marker provides the missing dimension:

Who is asking?

The final logic separates state, ownership, and timing

The three checks answer three different questions.

is_initialized()

Has the complete initialization lifecycle finished?

_tls.is_initializing

Is this the thread that currently owns the incomplete initialization lifecycle?

run_now

Did initialization complete while this thread was waiting to make the queue-or-run decision?

The patch succeeds because it does not treat these questions as equivalent.


The deeper lesson is that transitional states are real states

Many concurrency bugs arise when code models a lifecycle as:

before

after

while the implementation actually contains:

before

transition in progress

after

The transition may have its own permissions.

Some resources are ready.

Some callbacks are unfinished.

One thread owns the transition.

Other threads must still wait.

Treating that state as ordinary “not initialized” caused the second deadlock.

Treating it globally as “initialized” could have exposed unfinished state to other threads.

The correct answer was a thread-local transitional contract.

A state machine becomes unsafe when an operationally distinct transition is represented as though it were only one side of a Boolean.


Previous article

PyTorch #182948 — Why torch.cuda._lazy_call() Deadlocked When Its Callback Re-entered It


Related material


Patch status: Landed in PyTorch main as commit cec50d89
Stable fast path: is_initialized()
Initializer-thread fast path: hasattr(_tls, "is_initializing")
Lock-race handoff: run_now is set inside the lock and executed afterward
Public initialized boundary: _initialized = True only after queued callbacks finish
Cross-thread behavior: Other threads remain blocked on _initialization_lock during queue draining
Regression scenarios: Already-initialized set_rng_state() reentrancy and queued manual_seed() reentrancy
Validation recorded by the patch: Both scenarios timed out before the fix and completed after it on a ROCm gfx90a host

This is Part 2 of a three-part series on PyTorch’s lazy-initialization reentrancy deadlock.

Part 1 examined how callback execution under a non-reentrant lock caused one thread to wait for itself.

This article separated the fully initialized path, the race in which another thread completes initialization before lock acquisition, and the transitional state in which native initialization is complete but queued callbacks are still being drained.

Part 3 examines why both deadlock scenarios were placed in fresh subprocesses with a 60-second timeout, and how the similar XPU and MTIA source patterns should be distinguished from a backend failure that has actually been reproduced.

#PyTorch #CUDA #ROCm #Python #Concurrency #Deadlock #LazyInitialization #ThreadLocal #ThreadingLock #StateMachine #CodeAnalysis

Popular posts from this blog

AMD Is Trying to Turn Kernel Diversity Into One Software Asset

PyTorch #188031 — Why Did Forward Use 64-Bit Indexing While Backward Still Used uint32_t?

NVIDIA CUTLASS #3017 — Why load() and store() Mapped the Same Tile Coordinate to Different Addresses