Enable tensor disk offload for swarm learning - #4965
Open
YuanTingHsieh wants to merge 11 commits into
Open
Conversation
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #4965 +/- ##
==========================================
+ Coverage 64.24% 64.29% +0.04%
==========================================
Files 1030 1030
Lines 103575 103655 +80
==========================================
+ Hits 66542 66644 +102
+ Misses 37033 37011 -22
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
Comment on lines
+496
to
+497
| if root_dir: | ||
| shutil.rmtree(root_dir, ignore_errors=True) |
Contributor
There was a problem hiding this comment.
When finalization runs while a result callback is downloading tensors or the aggregation worker is materializing a lazy tensor, shutil.rmtree removes the run-owned storage without waiting for those operations, causing FileNotFoundError, dropped aggregation input, or an execution exception during shutdown.
Knowledge Base Used: app_common: Federated Learning Workflows
Contributor
Greptile SummaryAdds opt-in, run-scoped tensor disk offload to PyTorch Swarm learning.
Confidence Score: 3/5The PR should not merge until offload cleanup is synchronized with active result downloads and lazy aggregation consumers. The new finalizer can delete run-owned safetensors storage while independent Cell and aggregation threads are still using it, causing shutdown-time result failures or dropped aggregation work. Files Needing Attention: nvflare/app_common/ccwf/swarm_client_ctl.py Important Files Changed
Sequence DiagramsequenceDiagram
participant T as Trainer
participant CJ as Trainer Client Job
participant ACJ as Aggregation Client Job
participant D as Offload Storage
participant A as Aggregator
T->>CJ: Training result / transport refs
CJ->>ACJ: PASS_THROUGH result
ACJ->>D: Terminal tensor download
D-->>ACJ: Lazy disk-backed tensors
ACJ->>A: Gather contribution
A->>D: Materialize tensors as needed
A-->>ACJ: Aggregated model
Note over ACJ,D: Finalization must wait until downloads and materialization finish
Reviews (1): Last reviewed commit: "Merge main into feat/ccwf-tensor-disk-of..." | Re-trigger Greptile |
Contributor
There was a problem hiding this comment.
Pull request overview
This PR adds fully opt-in support for disk-backed tensor downloads on the receiving aggregation path in Swarm/CCWF when using streamed PyTorch payloads, enabling lower aggregation-client peak RSS while preserving correct ref/transport behavior across mixed execution modes and rotating aggregators.
Changes:
- Wire
server_expected_formatandenable_tensor_disk_offloadthrough Swarm PyTorch recipe → Swarm client config/job →SwarmClientController, including persistor behavior and warnings when misconfigured. - Extend FOBS/ViaDownloader download plumbing to accept call-scoped
fobs_ctx(incl. disk-offload root selection) and pass it through to PyTorch/Numpy decomposers. - Add/expand unit tests and documentation covering execution-mode matrices, PASS_THROUGH boundaries, disk-root ownership/cleanup, and recipe/CLI exposure.
Reviewed changes
Copilot reviewed 24 out of 24 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/unit_test/tool/agent_skill_checks/seed_skills_test.py | Updates recipe catalog assertions to reflect new Swarm PT recipe parameters. |
| tests/unit_test/recipe/swarm_recipe_test.py | Adds recipe-level tests for PyTorch streaming + disk-offload wiring and warning behavior. |
| tests/unit_test/fuel/utils/fobs/test_via_downloader_msg_root.py | Updates downloader stub signature to accept fobs_ctx. |
| tests/unit_test/fuel/utils/fobs/test_via_downloader_download_timeout.py | Updates downloader stub signature to accept fobs_ctx. |
| tests/unit_test/fuel/utils/fobs/decomposers/via_downloader_test.py | Adds test ensuring tensor decomposer honors call-scoped disk-offload settings via fobs_ctx. |
| tests/unit_test/app_common/ccwf/test_swarm_tensor_disk_offload.py | New tests for Swarm controller offload root ownership, cleanup, and secure relay route handling. |
| tests/unit_test/app_common/ccwf/test_swarm_forward_memory_path.py | Refactors/extends tests to validate mode-aware learn-task preparation and PASS_THROUGH handling. |
| tests/unit_test/app_common/ccwf/test_swarm_execution_mode_matrix.py | New 2-site/2-round execution-mode matrix test validating ref ownership and terminal offload policy. |
| tests/unit_test/app_common/ccwf/test_lazy_ref_local_aggr.py | Expands tests for local/remote aggregation boundaries and disk-offload decode context propagation. |
| nvflare/fuel/utils/fobs/decomposers/via_downloader.py | Adds fobs_ctx parameter to download() contract and threads it into remote-cell download path. |
| nvflare/fuel/utils/fobs/init.py | Introduces FOBSContextKey.TENSOR_DISK_OFFLOAD constant for shared context signaling. |
| nvflare/app_opt/pt/tensor_downloader.py | Adds call-scoped root_dir override to disk download helper (fallback to cell context). |
| nvflare/app_opt/pt/recipes/swarm.py | Adds recipe parameters (server_expected_format, enable_tensor_disk_offload), persistor wiring, and warning on misconfiguration. |
| nvflare/app_opt/pt/decomposers.py | Makes tensor decomposer honor call-scoped fobs_ctx for disk-offload enablement and root dir. |
| nvflare/app_common/utils/tensor_disk_offload_context.py | Aligns internal enable flag with FOBSContextKey.TENSOR_DISK_OFFLOAD and documents args. |
| nvflare/app_common/decomposers/numpy_decomposers.py | Updates NumPy decomposer download signature to accept fobs_ctx for interface consistency. |
| nvflare/app_common/ccwf/swarm_client_ctl.py | Implements Swarm-specific transport/offload boundary logic, job-scoped temp root lifecycle, and mode-aware ref termination. |
| nvflare/app_common/ccwf/ccwf_job.py | Wires enable_tensor_disk_offload from SwarmClientConfig into SwarmClientController. |
| docs/user_guide/data_scientist_guide/available_recipes.rst | Documents Swarm + tensor offload usage and adds ExchangeFormat import in examples. |
| docs/user_guide/admin_guide/deployment/overview.rst | Adds client-side deployment note for TMPDIR requirements when Swarm disk offload is enabled. |
| docs/user_guide/admin_guide/deployment/notes_on_large_models.rst | Generalizes guidance to cover Swarm clients as receivers in addition to servers. |
| docs/release_notes/flare_290.rst | Adds release note entry for Swarm tensor streaming + aggregation-client disk offload. |
| docs/programming_guide/controllers/client_controlled_workflows.rst | Updates CCWF/Swarm programming guide with new parameters and streaming/offload guidance. |
| docs/design/tensor_disk_offload.md | Extends design doc scope and details Swarm-specific terminal-download behavior and cleanup. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
934
to
+938
| if self._has_lazy_refs(request): | ||
| request = self._resolve_lazy_refs(request, fl_ctx) | ||
| # PASS_THROUGH is a transport instruction for this one hop, not | ||
| # part of the model result's semantic metadata. | ||
| request.set_header(ReservedHeaderKey.PASS_THROUGH, False) |
Comment on lines
264
to
+266
| _SwarmValidator(initial_ckpt=initial_ckpt) | ||
| warn_on_potential_secrets(command, context="recipe parameter 'command'") | ||
| self.server_expected_format = server_expected_format |
Comment on lines
+445
to
+446
| self.enable_tensor_disk_offload = client_config_args["enable_tensor_disk_offload"] | ||
| if self.enable_tensor_disk_offload and self.server_expected_format != ExchangeFormat.PYTORCH: |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What changed
enable_tensor_disk_offloadwiring to the Swarm client configuration, controller, and PyTorch recipeserver_expected_formatinSwarmLearningRecipeso PyTorch tensors remain on the Cell/FOBS streaming pathSwarmClientControllerdecide where transport refs terminate based on aggregation role and learner execution modeClientAPIExecutorlearner inputs as ordinary in-memory tensorsWhy
The merged
ClientAPIExecutorcan preserve large PyTorch tensor references between an external trainer and the CCWF controller. Swarm needs message-scoped handling because a client can be both trainer and aggregator, sites can use different execution modes, and the aggregation role can move between rounds. Terminal aggregation downloads can therefore use disk-backed lazy refs without exposing process-local disk refs to trainers, persistence, or final-result broadcasts.Usage
This is fully opt-in. Defaults remain NumPy exchange and disk offload disabled. Every client that can become the Swarm aggregator should have
TMPDIRpointed at sufficiently sized disk-backed storage.Disk offload applies on the receiving aggregation path. The trainer's source model and outgoing tensors remain memory-backed while serving the transport ref. An in-process trainer's local contribution is already materialized and remains in memory.
Validation
./runtest.sh -s --skip-installgit diff --check