[rob-luke]

What does this PR do ?

Install the OS-level SIGTERM trap so cfg.train.exit_signal_handler=True catches signals on the main process.

Changelog

• src/megatron/bridge/training/state.py: call .__enter__() on the DistributedSignalHandler constructed in GlobalState._set_signal_handler so signal.signal(SIGTERM, handler) runs. After this, sending SIGTERM to the main process flips signal_handler._signal_received, and the polled exit branch at train.py:~1320 triggers a graceful checkpoint-and-exit.
• tests/unit_tests/training/test_state.py: add two unit tests using the real DistributedSignalHandler — one asserts signal.getsignal(SIGTERM) changes after _set_signal_handler, the other sends a real SIGTERM via os.kill and asserts _signal_received flips. A class-scoped fixture snapshots and restores signal.getsignal(SIGTERM) so the installed trap does not leak across tests.

GitHub Actions CI

See the CI section in the Contributing doc for how to trigger the CI. A Nvidia developer will need to approve and trigger the CI for external contributors.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?
• Does the PR affect components that are optional to install? (Ex: Numba, Pynini, Apex etc)

Additional Information

• Related to # (issue) — none filed

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yaoyu-33]

/claude review

[claude]

Bug: reset_for_restart (line 494) sets self._signal_handler = None without calling self._signal_handler.release() first. Before this PR that was harmless because __enter__() was never called, so no OS trap was installed. Now that the trap IS installed, discarding the handler without release() orphans the OS-level signal handler — the original SIGTERM handler is never restored, and on restart a new handler stacks on top of the leaked one.

Suggested fix — guard the teardown in reset_for_restart:

if self._signal_handler is not None:
    self._signal_handler.release()
self._signal_handler = None

[rob-luke]

Good suggestion. Suggested fix has been applied.

[yaoyu-33]

@rob-luke can you check ai comment?

[claude]

Code Review

Bug: leaked OS signal handler on reset_for_restart

reset_for_restart (state.py:494) sets self._signal_handler = None without calling release() first. Before this PR that was harmless -- enter() was never called, so no OS trap existed. Now that the trap IS installed, discarding the handler orphans the OS-level SIGTERM handler: the original handler is never restored, and on in-process restart a new handler stacks on top of the leaked one.

Suggested fix in reset_for_restart:

if self._signal_handler is not None:
    self._signal_handler.release()
self._signal_handler = None

(Also left an inline comment on the same line.)

Existing mock test does not verify the fix

test_set_signal_handler_with_train_config (test_state.py:631-648) still passes only because MagicMock.enter() returns self by default -- the assertion state._signal_handler == mock_signal_handler is incidentally correct. Reverting the enter() call in state.py would NOT break this test. Consider adding:

mock_signal_handler.__enter__.assert_called_once()

so the mock-based test actually guards the fix.

New tests look good

The two new regression tests (test_set_signal_handler_installs_os_trap, test_sigterm_flips_signal_received_flag) are well-written -- the restore_sigterm_handler fixture properly isolates signal state, the safety-net no-op handler in the SIGTERM test prevents process death on regression, and the end-to-end test drives the public cfg setter path.

---

Suggested test cases

No perf tests impacted.

[rob-luke]

@yaoyu-33 the feedback from ai was good. I have implemented the suggested change. This PR is ready for review again.

[rob-luke]

@kamran-nvidia and @yaoyu-33 this issue is tagged as waiting on customer, but i think I addressed the earlier concerns. Do you have further changes you'd like me to make? Thank you for your help and the great package.

[rob-luke]

Thank you @adityavavreNVDA , @JRD971000 , @yaoyu-33. I think the ci failure was unrelated to the PR and occurred when I merged main to the branch, a skill was missing from the docs index. I’ve resolved the issue by adding the skill to the doc page, but can revert if the fix lands in main.

[adityavavreNVDA]

/ok to test 3af3319

[rob-luke]

Thank you @adityavavreNVDA for running the tests. Please let me know if there are additional changes you'd like to the PR.

[rob-luke]

Thank you @yaoyu-33 & @adityavavreNVDA