fix(retrieval): do not lose metric records on preemption or error

[sahel-sh]

What does this PR do ?

The bi-encoder recipe only closed its metric loggers on the success path, and records are buffered until close(), so a run that was preempted or raised lost every metric it had produced.

Changelog

• close() both loggers in the finally, so an exception cannot skip them
• poll StepScheduler.sigterm_received per step; both of its iterators already stop on the flag, but nothing polled it inside the step loop, so it was only checked once per epoch
• size the training buffer from ckpt_every_steps, capped at DEFAULT_BUFFER_SIZE
• flush both loggers: draining the buffer only reaches the file object, whose own buffer dies with the process

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?

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Related to # (issue)

[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.

[akoumpa]

/ok to test bcbd19a

[oliverholworthy]

@yuhezhang-ai I mentioned earlier about this property performing a side-effect, but upon closer inspection it only triggers the side-effectful signals_received check at most once per step since #3007

This additional check here is necessary because the preceding is_ckpt_step property can short-circuit before evaluating sigterm_received when the step is already a periodic, epoch-boundary, or final checkpoint. In those cases, the checkpoint is saved but the signal flag would not be updated, so training could continue.

This reveals a potential StepScheduler design improvement: signal polling could happen explicitly once at a defined step boundary, while sigterm_received and is_ckpt_step become side-effect-free cached-state reads. That would make the behavior less dependent on property access order and boolean short-circuiting. Not blocking for this PR, but maybe worth looking into whether this could be clearer.

[yuhezhang-ai]

Thanks, this makes sense. The per-step cache means line 451 does not add a second collective on ordinary steps, while the explicit access covers checkpoint steps where is_ckpt_step short-circuits before polling the signal.

I agree this is not blocking for this PR. We should have a follow-up to make signal polling an explicit StepScheduler operation with state-only properties.

[yuhezhang-ai]

Thanks for fixing metric durability on exceptional and preempted exits. @sahel-sh A few follow-ups:

1. Please run Ruff formatting. The current lint job reports that nemo_automodel/recipes/retrieval/train_bi_encoder.py would be reformatted: https://github.com/NVIDIA-NeMo/Automodel/actions/runs/32207675344/job/95933955404

2. Please add focused CPU regression coverage for both critical paths introduced here:

   • a signal arriving on a scheduled checkpoint step, proving the explicit post-checkpoint poll detects it and stops the loop; and
   • buffered train/validation metrics being persisted when the loop raises and cleanup runs from finally.

3. Please revisit the checkpoint-alignment claim around buffer_size=min(DEFAULT_BUFFER_SIZE, ckpt_every_steps). Record-count buffering is not necessarily checkpoint-aligned after resuming at a nonzero step or when checkpoints occur at epoch boundaries, so a checkpoint may be persisted while its metrics remain buffered. Either flush explicitly after successful checkpoint boundaries, or narrow/correct the guarantee and test the intended behavior.

The broader StepScheduler explicit-poll/state-query redesign can remain a follow-up; I am not asking for that refactor in this PR.

[akoumpa]

/ok to test 634afe5