Skip to content

Fix Engine atexit cleanup lifetime - #3943

Open
galitheia wants to merge 1 commit into
mosaicml:mainfrom
galitheia:fix/engine-atexit-lifetime-upstream
Open

Fix Engine atexit cleanup lifetime#3943
galitheia wants to merge 1 commit into
mosaicml:mainfrom
galitheia:fix/engine-atexit-lifetime-upstream

Conversation

@galitheia

@galitheia galitheia commented Sep 2, 2026

Copy link
Copy Markdown

Motivation

Composer currently registers the same static _close callback for every
Engine and passes each Engine's state and logger into atexit. Manually
closing an Engine does not unregister that entry, retaining training objects
and allowing cleanup to run again during interpreter shutdown.

Additionally, Engine.__del__() unregisters the shared callback globally,
which can remove handlers belonging to other live Engines. This prevents
bounded repeated in-process Trainer lifetimes.

This change gives every Engine a unique weak atexit callback, makes close()
idempotent, and unregisters only the callback owned by that Engine.

Summary

  • Register a unique weak atexit callback for each Engine.
  • Make Engine.close() idempotent and unregister only that Engine's callback.
  • Preserve the interpreter-shutdown sentinel and cleanup ordering.
  • Delegate del cleanup through close().
  • Add coverage for independent Engines, idempotent close, and exactly-once
    interpreter-exit cleanup.

Testing

  • Direct Engine lifecycle and interpreter-exit probes passed.
  • Ruff lint passed.
  • YAPF formatting check passed.
  • git diff --check passed.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant