# Tensor Disk Offload

## Objective

Reduce aggregation peak memory for large PyTorch model updates by materializing
incoming streamed tensor payloads on disk and resolving them lazily during
aggregation.

## Scope

* Applies to streamed **PyTorch tensor** payloads handled by `TensorDecomposer`.
* Controlled by `enable_tensor_disk_offload` in a supported receiving workflow/controller config:
  FedAvg (including FedProx), FedOpt, SCAFFOLD, or Swarm.
* Default is `False` (legacy in-memory behavior).
* If model updates are converted to NumPy before transport, tensor disk offload is not engaged.

Tensor streaming, ref pass-through, and tensor disk offload are separate
behaviors. The sender’s `TensorDownloadable` remains memory-backed while it
serves a transport ref. An intermediate process may preserve that ref without
downloading it. Disk offload applies when the receiving aggregation workflow
terminates the ref and downloads its tensor chunks. It does not spool the
trainer’s model or training result to disk at the source, and it does not reduce
the memory required to load or train the model in the trainer.

## How To Enable

FedAvg:

* `nvflare/recipe/fedavg.py` -> `FedAvgRecipe(..., enable_tensor_disk_offload=True)`
* `nvflare/app_opt/pt/recipes/fedavg.py` -> PT recipe forwards the same flag
* `nvflare/app_common/workflows/fedavg.py` -> `FedAvg(..., enable_tensor_disk_offload=True)`
* PyTorch FedProx uses the same FedAvg recipe and server aggregation path

Other server-controlled workflows:

* `nvflare/app_opt/pt/recipes/fedopt.py` -> `FedOptRecipe(..., enable_tensor_disk_offload=True)`
* `nvflare/app_opt/pt/recipes/scaffold.py` -> `ScaffoldRecipe(..., enable_tensor_disk_offload=True)`
* both require `server_expected_format=ExchangeFormat.PYTORCH`; custom configurations can enable
  the same setting directly on `ScatterAndGather` or `Scaffold`

Swarm/CCWF:

* `nvflare/app_opt/pt/recipes/swarm.py` -> use
  `SwarmLearningRecipe(aggregation_format=ExchangeFormat.PYTORCH, enable_tensor_disk_offload=True)`
* `nvflare/app_common/ccwf/ccwf_job.py` -> use
  `SwarmClientConfig(..., enable_tensor_disk_offload=True)` for custom Job API configurations
* `nvflare/app_common/ccwf/swarm_client_ctl.py` owns an offload root on each
  eligible client because the aggregation role moves between clients, but
  enables disk materialization only for terminal aggregation downloads

If no active Cell is available, the offload context is not enabled and the runtime falls back to in-memory download.

## Data Flow

```default
TensorDownloadable chunks
        |
        v
TensorDecomposer.download()
  - enable_tensor_disk_offload=False -> deserialize in memory
  - enable_tensor_disk_offload=True  -> write safetensors temp files
        |
        v
LazyTensorDict
        |
        v
ViaDownloaderDecomposer.recompose()
        |
        v
Lazy refs in payload tree
        |
        +--> aggregator consumes lazy refs (materialize on demand)
```

Server-controlled FedAvg/FedProx, FedOpt, and SCAFFOLD workflows install their
disk-offload setting on the receiving server Cell. Swarm does not enable disk
offload globally on its client Cells because the same Cells also receive
learner tasks and final-result broadcasts. Instead, the selected aggregation
controller puts the setting and its job-scoped root in the terminal result
download’s FOBS decode context. The destination therefore does not depend on
mutable Cell-wide state, and non-aggregation deliveries remain ordinary
in-memory tensors.

## Runtime Behavior

### FedAvg

In `nvflare/app_common/workflows/fedavg.py`:

* custom aggregators receive `result.params` as-is
* with `enable_tensor_disk_offload=True`, lazy refs are passed through directly
* built-in weighted aggregation materializes per tensor inside `WeightedAggregationHelper.add()`
  and relies on lazy-ref object lifetime / GC for temp-resource cleanup

The built-in weighted path remains lazy-friendly and memory-efficient.

### Filter Compatibility

Tensor disk offload produces per-tensor lazy values with a `materialize()` method. These values are
different from the transport-level `LazyDownloadRef` objects resolved at the Client API site-filter
boundary. Conventional content filters, including model quantizers and dequantizers, expect concrete
NumPy arrays or PyTorch tensors and are not compatible with tensor disk-offload lazy values.

Do not combine `enable_tensor_disk_offload=True` with conventional content filters on the same
receiving path. Disable tensor disk offload for that path, move the transformation to an explicit
send/receive endpoint, or use a consumer such as the built-in weighted aggregator that explicitly
materializes each tensor. Supporting filters while retaining bounded-memory disk offload requires a
separate lazy-ref-aware or streaming/per-tensor filter contract.

### Swarm/CCWF

`ClientAPIExecutor` preserves the Cell/FOBS large-payload references between an
external trainer and the CCWF controller. When a training result is sent to a
remote aggregation client, `SwarmClientController` requests PASS\_THROUGH on that
message. The aggregation controller therefore receives refs instead of
materializing the tensors inside the Cell receive callback, and explicitly
resolves them with its job-scoped disk root. A result returned by a local
external trainer also crosses a Cell boundary as a ref and is resolved by the
local aggregation controller with the same disk root. In both cases tensor disk
offload yields lazy tensor refs, and the built-in
`InTimeAccumulateWeightedAggregator` materializes one tensor at a time.

The Swarm controller owns the decision to preserve or resolve transport refs.
On a non-aggregation client it keeps refs when the local learn executor is
`ClientAPIExecutor(execution_mode="external_process")`, making the external
trainer their single consumer. If that external-process site is also the selected
aggregator, the controller resolves the task once for its aggregation base model
and passes the same in-memory payload to the trainer. This avoids two consumers
racing a one-receiver download transaction. The controller also resolves tasks
into memory for `in_process`, `attach`, and non-`ClientAPIExecutor` learners.
This conservative fallback supports jobs where sites use different learner
execution modes. Disk-backed aggregation refs remain local to the aggregation
client and are never passed to a learner.

The external trainer process’s Cell is not configured as a disk-offload
receiver. It is the terminal consumer of an incoming learn task, and the Client
API materializes ordinary in-memory tensors for model loading. For an outgoing
training result, `TensorDownloadable` still holds the tensors in trainer memory
while serving a transport ref. `ClientAPIExecutor` and the client job preserve
and route that ref; the selected aggregation controller decides whether its
terminal download is materialized on disk. Source-side tensor spooling is not
part of this feature.

If an in-process learner is also the local aggregation client, its own result is
already an in-memory object with no transport ref to preserve. That one local
contribution remains in memory; remote contributions still use terminal
aggregation-CJ disk offload.

`SwarmLearningRecipe` defaults to NumPy exchange for compatibility. Disk offload
therefore requires `aggregation_format=ExchangeFormat.PYTORCH`; streamed
NumPy arrays are not handled by `TensorDecomposer`.

## Custom Aggregator Contract

When a custom aggregator is used, payload params may contain lazy refs (duck-typed object with `materialize()`).

Custom aggregators are responsible for:

1. materializing refs when tensor math is required
2. releasing lazy-ref object references after use so temp resources can be reclaimed

## Temp File Lifecycle

* Each workflow creates a job-scoped offload root (`nvflare_tensor_offload_<job>_*`),
  with safetensors download directories beneath it.
* Temp dir selection follows Python `tempfile` behavior (`TMPDIR` / OS default, typically `/tmp`).
* In containerized deployments, `/tmp` may be tmpfs (RAM-backed); set `TMPDIR` to a disk-backed mount to realize memory offload benefits.
* `LazyTensorDict` owns a shared `_TempDirRef`; each lazy ref keeps this reference alive.
* Lazy download directories are reclaimed when their refs are released, with GC as
  a fallback.
* FedAvg-style workflows restore the prior FOBS context and remove their
  job-scoped root when the workflow exits.
* Swarm keeps its job-scoped root and tensor-forwarding route through workflow
  finalization. Explicit terminal downloads inherit the run abort signal. At job
  `END_RUN`, Swarm cancels active disk consumers under its root, gives the
  controller-owned learning and aggregation threads a bounded drain window, then
  removes the root and its remaining contents. A receiver that abandons an active
  download also sends a capability-negotiated cancellation to the producer. The
  producer records that receiver as failed across every ref in the transaction and
  can release an accepted external-process result source without waiting for its
  normal transfer timeout. Older producers keep the existing bounded timeout path.
* These workflow-local cleanup paths require the server/client job process to reach
  normal teardown. Reclaiming scratch after an abrupt server/client job process
  death requires launcher-owned storage whose cleanup capability survives process
  and container boundaries; that lifecycle is outside this workflow-level feature.

## Failure Behavior

* Download failures trigger `DiskTensorConsumer.download_failed(...)`, which removes the temp dir.
* If an owned external trainer dies after its lazy result envelope is accepted, its CJ sends a
  bounded acknowledged failure notice for the exact source FQCN/reference IDs to the declared
  terminal receiver. Matching active downloads are interrupted immediately; a bounded tombstone
  also covers failure notices that arrive just before download startup. FedAvg and Swarm then use
  their existing materialization/task-error paths instead of retrying the dead source for the
  normal 600-second streaming timeout.
* Mid-transfer receiver cancellation is best effort. If the producer advertised support, it
  promptly settles the abandoned receiver as failed; otherwise the transaction’s existing
  receiver/transaction timeout remains the cleanup backstop.
* Invalid safetensors payload/header parsing fails fast and bubbles up as a download-consume error.
* Existing in-memory download path remains unchanged when offload is disabled.

## Design-Relevant Files

* `nvflare/app_opt/pt/decomposers.py`
* `nvflare/app_opt/pt/lazy_tensor_dict.py`
* `nvflare/app_opt/pt/tensor_downloader.py`
* `nvflare/fuel/utils/fobs/decomposers/via_downloader.py`
* `nvflare/app_common/workflows/fedavg.py`
* `nvflare/app_common/ccwf/swarm_client_ctl.py`
* `nvflare/app_common/ccwf/ccwf_job.py`
* `nvflare/recipe/fedavg.py`
* `nvflare/app_opt/pt/recipes/swarm.py`

## Test Coverage

* `tests/unit_test/app_common/workflow/fedavg_test.py`
* `tests/unit_test/app_common/ccwf/test_swarm_tensor_disk_offload.py`
* `tests/unit_test/recipe/swarm_recipe_test.py`
* `tests/unit_test/app_opt/pt/test_lazy_tensor_dict.py`
* `tests/unit_test/app_opt/pt/test_disk_tensor_consumer.py`
* `tests/unit_test/app_common/aggregators/weighted_aggregation_helper_test.py`
* `tests/unit_test/private/fed/server/server_runner_test.py`
* `tests/stress_test/fedavg_large_model/fedavg_stress_test.py`