perf(adapter/pytorch): offload spawn manager creation to background thread (threading.Thread) by kkkapu · Pull Request #73 · google/ml-flashpoint

kkkapu · 2026-03-06T21:19:58Z

Why this change is necessary:

Using 'spawn' context is required to prevent CUDA OOM during NVRX restarts, but it introduces a significant ～15s delay.
The previous implementation created the Manager synchronously, blocking the main trainer setup.

Change-Id: I9cdf24a9da40fa19d5b93d5ef9ad7c6f42641b1b

Fixes #<issue_number_goes_here>

It's a good idea to open an issue first for discussion.

Tests pass
Appropriate changes to documentation are included in the PR

…hread Why this change is necessary: 1. Using 'spawn' context is required to prevent CUDA OOM during NVRX restarts, but it introduces a significant ～15s delay. 2. The previous implementation created the Manager synchronously, blocking the main trainer setup. Change-Id: I9cdf24a9da40fa19d5b93d5ef9ad7c6f42641b1b

gemini-code-assist · 2026-03-06T21:23:17Z

Summary of Changes

Hello, 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 significantly improves the performance of the trainer setup by converting the synchronous creation of the torch.multiprocessing.Manager into an asynchronous operation. By moving this potentially long-running task to a background thread, the main trainer setup process is no longer blocked, thereby eliminating a ~15-second delay. This optimization is particularly important for environments requiring the 'spawn' context to mitigate CUDA OOM issues during restarts, ensuring a smoother and faster initialization experience.

Highlights

Asynchronous Manager Creation: The creation of the torch.multiprocessing.Manager is now offloaded to a background threading.Thread and accessed via a concurrent.futures.Future. This change addresses a significant ~15-second delay during trainer setup when using the 'spawn' context, which is crucial for preventing CUDA Out-Of-Memory errors during NVRX restarts.
MemoryStorageWriter Refactor: The MemoryStorageWriter class has been updated to accept a concurrent.futures.Future for the multiprocessing manager instead of the manager instance directly. This allows for asynchronous initialization and access to the manager's shared data structures.
Lazy Initialization of Shared Data Structures: The _write_events_per_checkpoint_id and _write_results_per_checkpoint_id attributes within MemoryStorageWriter are now initialized lazily within the reset method, ensuring the multiprocessing manager is available before these shared dictionaries are created.
Test Suite Updates: All relevant tests have been updated to reflect the new asynchronous manager handling, including changes to fixtures, constructor calls, and assertions to accommodate the mp_manager_future.

🧠 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/adapter/megatron/save_strategies.py
- Updated MemoryStorageWriter instantiation to use mp_manager_future instead of mp_manager.
src/ml_flashpoint/adapter/nemo/wrapper_util.py
- Imported concurrent.futures and threading.
- Implemented asynchronous creation of torch_mp.Manager using concurrent.futures.Future and threading.Thread.
- Passed mp_manager_future to MemoryStorageWriter instead of the direct Manager instance.
src/ml_flashpoint/adapter/pytorch/memory_storage_writer.py
- Modified __init__ signature to accept mp_manager_future: concurrent.futures.Future instead of mp_manager: torch_mp.Manager.
- Updated internal attribute _main_process_torchmp_manager to _main_process_torchmp_manager_future.
- Modified __getstate__ and __setstate__ to handle the new _main_process_torchmp_manager_future attribute.
- Implemented lazy initialization of _write_events_per_checkpoint_id and _write_results_per_checkpoint_id within the reset method, retrieving the Manager from the future.
- Updated calls to mp_manager.Event() to mp_manager_future.result().Event().
tests/adapter/megatron/test_save_strategies.py
- Updated MemoryStorageWriter instantiations in tests to pass mp_manager_future (mocked).
- Adjusted assertions to check _main_process_torchmp_manager_future.
tests/adapter/nemo/test_wrapper_util.py
- Modified assertion to check kwargs["mp_manager_future"].result() instead of kwargs["mp_manager"].
tests/adapter/pytorch/test_memory_storage_writer.py
- Renamed mp_manager fixture to mp_manager_future and wrapped the Manager in a concurrent.futures.Future.
- Updated MemoryStorageWriter instantiations in tests to use mp_manager_future.
- Adjusted assertions for _main_process_torchmp_manager_future and lazy initialization of _write_events_per_checkpoint_id and _write_results_per_checkpoint_id.
- Added writer.reset(checkpoint_id.data) calls where necessary before accessing manager-dependent attributes.
- Updated calls to mp_manager.Event() to mp_manager_future.result().Event().

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.

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 introduces a performance improvement by offloading the creation of the torch.multiprocessing.Manager to a background thread. This avoids blocking the main trainer setup, which previously caused a significant delay. The implementation uses a concurrent.futures.Future to manage the asynchronous creation of the manager, and the MemoryStorageWriter is updated to lazily initialize its shared state by waiting on this future. The changes are logical and well-implemented. I have a couple of minor suggestions to improve code clarity and formatting.

gemini-code-assist · 2026-03-06T21:25:53Z

+        if self._write_events_per_checkpoint_id is None and self._main_process_torchmp_manager_future is not None:
+            self._write_events_per_checkpoint_id = self._main_process_torchmp_manager_future.result().dict()
+
+        if self._write_results_per_checkpoint_id is None and self._main_process_torchmp_manager_future is not None:
+            self._write_results_per_checkpoint_id = self._main_process_torchmp_manager_future.result().dict()


These two if statements can be combined. Since _write_events_per_checkpoint_id and _write_results_per_checkpoint_id are initialized together in __init__ and this is the only place they are assigned from None, you can check for one and initialize both. This also avoids calling _main_process_torchmp_manager_future.result() twice, making the code cleaner and slightly more efficient.

Suggested change

if self._write_events_per_checkpoint_id is None and self._main_process_torchmp_manager_future is not None:

self._write_events_per_checkpoint_id = self._main_process_torchmp_manager_future.result().dict()

if self._write_results_per_checkpoint_id is None and self._main_process_torchmp_manager_future is not None:

self._write_results_per_checkpoint_id = self._main_process_torchmp_manager_future.result().dict()

if self._write_events_per_checkpoint_id is None and self._main_process_torchmp_manager_future is not None:

mp_manager = self._main_process_torchmp_manager_future.result()

self._write_events_per_checkpoint_id = mp_manager.dict()

self._write_results_per_checkpoint_id = mp_manager.dict()

Change-Id: I10b64985354fb808873fc3a1578e0cba9e069bb8

Change-Id: I178993a77f184a355aeae0df99f306cab18eeb49

g-husam · 2026-03-06T21:40:09Z

+        mp_manager_future.set_result(ctx.Manager())
+
+    thread = threading.Thread(target=start_manager, daemon=True)
+    thread.start()


doesnt this approach cause errors?

On errors on our side

what do you mean?

Change-Id: I6a2b9b4f3730eacb3cf1d75e5dcf3659f5bf0c4e

g-husam · 2026-03-10T15:58:38Z

        state = self.__dict__.copy()
-        if "_main_process_torchmp_manager" in state:
-            del state["_main_process_torchmp_manager"]
+        if "_main_process_torchmp_manager_future" in state:


nit: a one-liner alternative to the if-then-del is

state.pop("_main_process_torchmp_manager_future", None)

g-husam · 2026-03-10T16:00:19Z

        # Mimicking existing StorageWriter impls (e.g. `_FileSystemWriter`) by using a random ID as the save ID.
        self._current_save_id = generate_hfid("memwritersave")

+        if self._write_events_per_checkpoint_id is None and self._main_process_torchmp_manager_future is not None:


make sure there's a test for the diff conditions here: where each condition is T/F (4 scenarios)

g-husam · 2026-03-10T16:03:05Z

            # Manually initialize the event for checkpoint_id1, as prepare_write_data_buckets is not called
-            writer._write_events_per_checkpoint_id[checkpoint_id1] = writer._main_process_torchmp_manager.Event()
+            writer._write_events_per_checkpoint_id[checkpoint_id1] = (
+                writer._main_process_torchmp_manager_future.result().Event()


style nit: a cleaner way to access could be to add a property that automatically waits for the future (keep it private tho):

@property def _main_process_torchmp_manager(self): """Returns the torch multiprocessing Manager from the Future, blocking if necessary.""" return self._main_process_torchmp_manager_future.result()

then you can access it the same way you were before, and wouldn't have to modify these lines : writer._main_process_torchmp_manager.Event()

Will open a new PR for this, as the updates affect multiple areas of the codebase and tests.

Change-Id: I1ed36f191e66ba16107f976ecc3a54ae07e8975e

github-actions · 2026-03-10T19:46:42Z

Python Code Coverage Summary

Package	Line Rate	Branch Rate	Health
src.ml_flashpoint	100%	100%	✔
src.ml_flashpoint.adapter	100%	100%	✔
src.ml_flashpoint.adapter.megatron	98%	95%	✔
src.ml_flashpoint.adapter.nemo	98%	94%	✔
src.ml_flashpoint.adapter.pytorch	99%	92%	✔
src.ml_flashpoint.checkpoint_object_manager	92%	91%	➖
src.ml_flashpoint.core	96%	92%	✔
src.ml_flashpoint.replication	81%	81%	❌
Summary	95% (2096 / 2208)	91% (482 / 528)	➖

Minimum allowed line rate is 90%

github-actions · 2026-03-10T19:46:45Z

C++ Code Coverage Summary

Package	Line Rate	Branch Rate	Health
src.ml_flashpoint.checkpoint_object_manager.buffer_object	93%	54%	✔
src.ml_flashpoint.checkpoint_object_manager.object_manager	70%	37%	❌
src.ml_flashpoint.replication.transfer_service	79%	40%	❌
Summary	81% (916 / 1126)	43% (687 / 1604)	➖

Minimum allowed line rate is 80%

kkkapu and others added 2 commits March 6, 2026 21:17

Merge branch 'main' into async_spawn2

2cd57d8

gemini-code-assist Bot reviewed Mar 6, 2026

View reviewed changes

kkkapu added 2 commits March 6, 2026 21:29

fix format

23e9080

Change-Id: I10b64985354fb808873fc3a1578e0cba9e069bb8

address review comments

eac9648

Change-Id: I178993a77f184a355aeae0df99f306cab18eeb49

g-husam requested changes Mar 6, 2026

View reviewed changes

address review comments

6b9617f

Change-Id: I6a2b9b4f3730eacb3cf1d75e5dcf3659f5bf0c4e

kkkapu requested review from Leahlijuan and g-husam March 9, 2026 19:43

g-husam approved these changes Mar 10, 2026

View reviewed changes

address review comments

91c68fc

Change-Id: I1ed36f191e66ba16107f976ecc3a54ae07e8975e

kkkapu merged commit ae95fb6 into main Mar 10, 2026
5 checks passed

kkkapu deleted the async_spawn2 branch March 10, 2026 19:54

ronaldw07 mentioned this pull request Mar 10, 2026

refactor: rename thread_count to files_per_rank #67

Open

Uh oh!

Conversation

kkkapu commented Mar 6, 2026

Uh oh!

gemini-code-assist Bot commented Mar 6, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

gemini-code-assist Bot Mar 6, 2026

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

github-actions Bot commented Mar 10, 2026

Python Code Coverage Summary

Uh oh!

github-actions Bot commented Mar 10, 2026

C++ Code Coverage Summary

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants