nvflare.app_common.utils.tensor_disk_offload_context module

class TensorDiskOffloadContext(previous_value: Any = None, previous_root_dir: str | None = None, root_dir: str | None = None, applied: bool = False)

Bases: object

applied: bool = False

previous_root_dir: str | None = None

previous_value: Any = None

root_dir: str | None = None

cleanup_tensor_disk_offload(engine, context: TensorDiskOffloadContext) → None

Restore the prior FOBS context values and remove any temporary offload root.

setup_tensor_disk_offload(engine, enabled: bool, job_id: str = 'job', root_dir: str | None = None) → TensorDiskOffloadContext

Enable tensor disk offload in the active cell FOBS context.

Parameters:

• engine – engine that owns the active Cell.

• enabled – whether to prepare disk-backed tensor downloads.

• job_id – identifier used to name the temporary offload root.

• root_dir – optional existing parent directory for the temporary root.

Returns:

Context needed to restore the prior setting and cleanup temporary files.