refactor(core/loader): abstract torch.distributed APIs in CheckpointLoader

[ronaldw07]

Summary

Closes #30

• Adds injectable callable parameters (global_rank_getter, local_rank_getter, broadcast_object_list_func, all_gather_object_func, get_world_size_func) to DefaultMLFlashpointCheckpointLoader.__init__
• Replaces all direct torch.distributed calls in method bodies with these injected callables
• All parameters default to the corresponding torch.distributed functions, preserving full backwards compatibility
• Mirrors the same dependency-injection pattern already used in DefaultMLFlashpointCheckpointSaver

Test plan

• Existing tests pass without modification (defaults preserve the original behavior)
• The new parameters allow passing stub/mock callables in tests without patching torch.distributed
• NeMoMLFlashpointCheckpointLoader and other subclasses continue to work unchanged since they call super().__init__ with only the original positional arguments

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[gemini-code-assist]

Summary of Changes

Hello @ronaldw07, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request refactors the DefaultMLFlashpointCheckpointLoader by abstracting its dependencies on torch.distributed APIs. This change enhances testability and flexibility by allowing torch.distributed functions to be injected as parameters, making it easier to mock or stub these dependencies in tests without altering the core torch.distributed module. The refactoring maintains existing behavior through default parameter values.

Highlights

• Dependency Injection: Introduced dependency injection for torch.distributed APIs within DefaultMLFlashpointCheckpointLoader.
• API Abstraction: Replaced all direct calls to torch.distributed functions with injectable callable parameters.
• Backward Compatibility: Ensured full backward compatibility by setting default values for new parameters to their corresponding torch.distributed functions.
• Architectural Consistency: Aligned CheckpointLoader with the existing dependency injection pattern already used in DefaultMLFlashpointCheckpointSaver.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• src/ml_flashpoint/core/checkpoint_loader.py
  • Imported Callable from the typing module.
  • Modified the __init__ method of DefaultMLFlashpointCheckpointLoader to accept five new callable parameters for distributed operations: global_rank_getter, local_rank_getter, broadcast_object_list_func, all_gather_object_func, and get_world_size_func.
  • Added corresponding instance attributes to store these injected callables.
  • Replaced direct calls to dist.get_rank() with self._global_rank_getter().
  • Replaced direct calls to dist.broadcast_object_list() with self._broadcast_object_list_func().
  • Replaced direct calls to dist.get_world_size() with self._get_world_size_func().
  • Replaced direct calls to dist.get_node_local_rank() with self._local_rank_getter().
  • Replaced direct calls to dist.all_gather_object() with self._all_gather_object_func().

Activity

• Closed issue Abstract torch.distributed APIs in CheckpointLoader #30.
• Confirmed existing tests pass without modification, ensuring no regression in original behavior.
• Verified that the new parameters allow for easier testing with stub/mock callables.
• Confirmed that subclasses like NeMoMLFlashpointCheckpointLoader continue to work unchanged.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request refactors DefaultMLFlashpointCheckpointLoader to abstract away direct calls to torch.distributed by using dependency injection for distributed communication functions. This is a solid improvement that will make the code more modular and easier to test. The implementation is clean and follows the stated goal. I have one suggestion to make the new __init__ parameters keyword-only to improve API robustness.

[g-husam]

should this still be assert_called_once?

[g-husam]

Thanks for this! Looks good, just need to double check and ensure assertions are not modified, and address lint/build issues