[Bug]: Training shutdown decorator leaves signal handlers installed after return or error

[LOGO127]

Bug

At 3419401a380db305ed6493ba4680ad8a3c3e0e45, with_graceful_shutdown() installs SIGINT/SIGTERM handlers but restores them only when it catches TrainingInterruptedError. After a normal return or an unrelated exception, the installed handlers still belong to the completed training call.

This matters when invoking training from a longer-lived Python process: a later interrupt enters the stale training handler, rather than the handler that was active before training. Repeated/nested decorated calls can also capture and restore stale handlers.

Minimal reproduction

From a checkout, this uses only Python's standard library and loads the actual production file. It avoids unrelated package initialization and restores the process's handlers before exiting.

import runpy
import signal

module = runpy.run_path("src/speculators/train/graceful_shutdown.py")
before = {sig: signal.getsignal(sig) for sig in (signal.SIGINT, signal.SIGTERM)}

@module["with_graceful_shutdown"]()
def train(self):
    return "finished"

try:
    print(train(object()))
    print({sig.name: signal.getsignal(sig) == handler
           for sig, handler in before.items()})
finally:
    for sig, handler in before.items():
        signal.signal(sig, handler)

Observed on Windows / Python 3.13:

finished
{'SIGINT': False, 'SIGTERM': False}

Expected: both values are True. Raising an ordinary RuntimeError inside the decorated call also leaves the handlers installed.

Proposed scope and verification

A minimal fix is to add an outer finally: handler.restore() to the existing try/except, preserving the early restore before interrupt-checkpoint saving. A local two-line fix plus six regression cases covers return values, ordinary exceptions, KeyboardInterrupt, checkpoint success/failure, and nested calls.

Isolated tests loading the real stdlib-only module: 4 failed / 2 passed before, 6 passed after. A tiny CUDA linear-model workload also reproduces the leak and verifies restoration after the fix for FP32/BF16 and return/error paths; this is not a full Speculators Trainer or distributed-training test.

The standard pytest entry point is currently blocked on this Windows environment by the package's fcntl import. A follow-up PR still needs the contributor's human review/sign-off and supported-environment checks. This report was prepared with coding-agent assistance; no human-review checklist is being asserted here.