feat: save checkpoint before timeout to avoid 4-hour runtime limit

[wedu-nvidia]

…g time lmit

What does this PR do ?

Since the server automatically stops after 4 hours, it's recommended to save a checkpoint beforehand. For example, set the timeout to 3 hours and 45 minutes to ensure check point saved is saved in time

Issues

List issues that this PR closes (syntax):

Usage

• You can potentially add a usage example below

# Add a code snippet demonstrating how to use this

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 run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• ...

[terrykong]

thanks for contributing. this would definitely be a valuable feature to have. I've left some comments

[wedu-nvidia]

@terrykong I revised based on your suggestions and let me know if have more comments

[terrykong]

@wedu-nvidia could you address the DCO failure and run the pre-commit hooks. See https://github.com/NVIDIA-NeMo/RL/blob/main/CONTRIBUTING.md

[wedu-nvidia]

Hi @terrykong, all DCO and pre-commit issues have been resolved, and the commits are now properly signed.

Please help approve the pending workflows and review the change request when convenient — thanks!

[terrykong]

@wedu-nvidia looks like there are still some failures, this time with pyrefly

ERROR `float` is not assignable to attribute `previous_iteration_time` with type `None` [bad-assignment]
   --> /home/runner/work/RL/RL/nemo_rl/utils/timer.py:311:40
    |
311 |         self.previous_iteration_time = time.time()
    |                                        ^^^^^^^^^^^
    |
ERROR `-` is not supported between `float` and `None` [bad-argument-type]
   --> /home/runner/work/RL/RL/nemo_rl/utils/timer.py:318:24
    |
318 |         elapsed_time = current_time - self.previous_iteration_time
    |                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
    |
  Argument `None` is not assignable to parameter `value` with type `float` in function `float.__sub__`
ERROR `float` is not assignable to attribute `previous_iteration_time` with type `None` [bad-assignment]
   --> /home/runner/work/RL/RL/nemo_rl/utils/timer.py:319:40
    |
319 |         self.previous_iteration_time = current_time
    |                                        ^^^^^^^^^^^^
    |
 INFO errors shown: 3, errors ignored: 24, modules: 82, transitive dependencies: 4,273, lines: 2,048,943, time: 4.08s, peak memory: physical 786.5 MiB

[terrykong]

This one looks to have a legitimate unit test failure:

FAILED unit/algorithms/test_sft.py::test_exit_on_max_steps - KeyError: 'save_...

When you've resolved, we can retry

[wedu-nvidia]

@terrykong I added another parameter, and hope it can pass all.

[wedu-nvidia]

@terrykong The previous error seems solved and I saw another error and I added in
checkpoint_must_save_by: NotRequired[str | None]
in following config

class CheckpointingConfig(TypedDict):
    """Configuration for checkpoint management.

    Attributes:
    enabled (bool): Whether checkpointing is enabled.
    checkpoint_dir (PathLike): Directory where checkpoints will be saved.
    metric_name (str): Name of the metric to use for determining best checkpoints.
    higher_is_better (bool): Whether higher values of the metric indicate better performance.
    keep_top_k (Optional[int]): Number of best checkpoints to keep. If None, all checkpoints are kept.
    """

    enabled: bool
    checkpoint_dir: PathLike
    metric_name: str
    higher_is_better: bool
    save_period: int
    keep_top_k: NotRequired[int]
    checkpoint_must_save_by: NotRequired[str | None]

[wedu-nvidia]

@terrykong Can you help add it the merge queue again? Thanks so much

[wedu-nvidia]

@terrykong can you put it into mergequeue?

[wedu-nvidia]

@terrykong Why I did not see the conflict?