Advanced API¶

Lower-level components for custom workflows. In most cases, use the Simplified API or OffloadManager instead.

TensorManager provides direct control over discovery, profiling, and inference phases. See Internal Phases for details on how these components work together.

TensorManager ¶

TensorManager(device_gpu: str | device, tensor_manager_load_strategy: Strategy, pinned_memory: bool = True, pinned_memory_mode: PinnedMemoryMode = 'torch', loader_type: str = 'allocation_block_transfer', remove_layers_operations: list[dict[str, Any]] | None = None, blocks: int = 4, move_top_level_buffers_to_gpu: bool = True, use_shm: bool = False, include_patterns: list[str] | None = None, exclude_patterns: list[str] | None = None, enable_diagnostics: bool = False, max_gpu_mem_fraction: float | None = None, profile_mode: ProfileMode = 'view', *, _direct_mode: bool = True, _use_trace_tensor: bool = False, _rearrange_transfers: bool = False, _compute_transfer_gap: int = 1, _enable_untraced_tensor_discovery: bool = True, _offload_timing: OffloadTimingMode = 'off', _piecewise_prefetch: PiecewisePrefetchMode = 'warn', _offload_timing_log_every: int = OFFLOAD_TIMING_LOG_EVERY, _offload_timing_measure_max_passes: int = OFFLOAD_TIMING_MEASURE_MAX_PASSES)

Initialize TensorManager with configurable tensor loading strategy.

PARAMETER	DESCRIPTION
`device_gpu`	GPU device to use TYPE: `str \| device`
`tensor_manager_load_strategy`	Strategy for loading tensors TYPE: `Strategy`
`pinned_memory`	Whether to use pinned (page-locked) memory for CPU tensors. Required for non-blocking transfers on a separate CUDA stream so that weight copies can overlap with GPU computation. Increases host memory pressure since pinned pages cannot be swapped. (default: True) TYPE: `bool` DEFAULT: `True`
`pinned_memory_mode`	How to pin CPU tensors when `pinned_memory=True`. `"torch"` uses :meth:`torch.Tensor.pin_memory` (fresh pinned copy); `"host_register"` uses `cudaHostRegister` to pin in place (no copy, lower peak host RAM). On a CUDA host whose cudart binding is broken or missing, the requested mode falls back to `"torch"` with a WARNING log naming the cause. On a CPU-only host, `pinned_memory=True` raises `RuntimeError` at construction — set `pinned_memory=False` for CPU-only deployments. SHM segments built with `pinned_memory=True` register in place regardless of this setting; SHM segments built with `pinned_memory=False` are not registered. (default: "torch") TYPE: `PinnedMemoryMode` DEFAULT: `'torch'`
`loader_type`	Type of tensor loader to use. Options: - "strategy": Uses TensorStrategyLoader - "raw_block_transfer": Uses PreallocatedBatchTransferTensorLoader with RawBlockController - "allocation_block_transfer": Uses PreallocatedBatchTransferTensorLoader with AllocationBlockController TYPE: `str` DEFAULT: `'allocation_block_transfer'`
`remove_layers_operations`	List of operations to remove layers from the strategy map TYPE: `list[dict[str, Any]] \| None` DEFAULT: `None`
`blocks`	Number of blocks to use for the block transfer loaders (default: 4) TYPE: `int` DEFAULT: `4`
`move_top_level_buffers_to_gpu`	Whether to move top-level model buffers to GPU during discovery initialization (default: True) TYPE: `bool` DEFAULT: `True`
`include_patterns`	List of patterns to include for offloading. Parameters not matching are kept on GPU permanently (default: None → [""]). TYPE:* `list[str] \| None` DEFAULT: `None`
`exclude_patterns`	List of patterns to exclude from offloading. Passed through to tensor discovery for parameter-level filtering (default: None) TYPE: `list[str] \| None` DEFAULT: `None`
`enable_diagnostics`	Whether to log diagnostic information (layer duration statistics, block assignment table) after strategy computation (default: False) TYPE: `bool` DEFAULT: `False`
`max_gpu_mem_fraction`	Fraction of total GPU memory to use as budget, in (0.0, 1.0]. Resolved to bytes at compute time via :func:`flextensor.gpu_budget.resolve_gpu_budget`. If None, no memory constraint is applied (latency mode). (default: None) TYPE: `float \| None` DEFAULT: `None`
`profile_mode`	Profile-phase mechanism. `"view"`/`"getter"` are profiling variants of the model-patching runtime and affect only the profile phase; `"torch_function"` selects the indirect runtime, which also changes warmup and inference. One of: `"view"` — default. Patches the profile model with views into a rotating GPU block plus a shared prefix for tensors reused across labels. Pre-allocates substantial GPU memory during profile — pick `"getter"` if you are tight on memory. Incompatible with `_use_trace_tensor=True`. See `profile_mode` in `docs/explanation/configuration.md` for the timing model and memory formula. `"getter"` — profile model uses property getters that route every parameter access through `TensorLayerLoader`. Lower profile-phase GPU footprint than `"view"` at the cost of attribute-getter overhead in per-trap durations. `"torch_function"` — fallback using `TorchFunctionMode` traps that rewrite tensor arguments per call, without patching the model. Significant overhead, not torch.compile-compatible; only valid with `loader_type='strategy'`. TYPE: `ProfileMode` DEFAULT: `'view'`
`_direct_mode`	Internal — runtime family. `True` (default) uses the model-patching runtime across warmup/profile/inference; `False` uses the indirect `TorchFunctionMode` runtime (requires `loader_type='strategy'`). `profile_mode='torch_function'` forces this `False`. Not part of the public API. TYPE: `bool` DEFAULT: `True`
`_use_trace_tensor`	Internal debug flag — enables TraceTensor + BenchmarkReplace for per-tensor offload timing. Not part of the public API. TYPE: `bool` DEFAULT: `False`
`_rearrange_transfers`	Internal — whether to enable transfer rearrangement. Auto-enabled when permanent gap layers are detected. (default: False) TYPE: `bool` DEFAULT: `False`
`_compute_transfer_gap`	Internal — minimum layer gap between a block's compute and its next transfer during rearrangement. (default: 1) TYPE: `int` DEFAULT: `1`
`_enable_untraced_tensor_discovery`	Internal — whether to discover untraced tensors (e.g., FP8 weights passed to Triton kernels). (default: True) TYPE: `bool` DEFAULT: `True`
`_offload_timing`	Offload-timing mode forwarded from :attr:`~flextensor.config.OffloadConfig.offload_timing` (`"off"` / `"eager"` / `"cuda_graph"`). (default: `"off"`) TYPE: `OffloadTimingMode` DEFAULT: `'off'`
`_piecewise_prefetch`	Piecewise prefetch policy forwarded from :attr:`~flextensor.config.OffloadConfig.piecewise_prefetch` (`"off"` / `"warn"` / `"error"`). (default: `"warn"`) TYPE: `PiecewisePrefetchMode` DEFAULT: `'warn'`
`_offload_timing_log_every`	Internal — log rolling timing table every N passes; `0` disables. (default: :data:`~flextensor.offload_timing.OFFLOAD_TIMING_LOG_EVERY`) TYPE: `int` DEFAULT: `OFFLOAD_TIMING_LOG_EVERY`
`_offload_timing_measure_max_passes`	Internal — max finalized passes in the durable measure store (ring buffer). (default: :data:`~flextensor.offload_timing.OFFLOAD_TIMING_MEASURE_MAX_PASSES`) TYPE: `int` DEFAULT: `OFFLOAD_TIMING_MEASURE_MAX_PASSES`

set_model ¶

set_model(model: Module | dict[str, object] | ModelDict) -> None

initialize_warmup ¶

initialize_warmup() -> Any

initialize_profile ¶

initialize_profile() -> Any

initialize_inference ¶

initialize_inference() -> Any

prepare_warmup_mode ¶

prepare_warmup_mode() -> None

prepare_profile_mode ¶

prepare_profile_mode() -> None

Wire the profile-phase trap for profile_mode='torch_function'.

Single-call setup (the torch_function path doesn't patch the model).

prepare_infer_mode ¶

prepare_infer_mode() -> None

Prepare inference mode from fresh profiling data.

This method is called after profiling to compute the load strategy and create the appropriate loader for inference.

trap ¶

trap(name: str) -> Any

release_memory ¶

release_memory() -> None

get_gpu_memory_usage ¶

get_gpu_memory_usage() -> GPUMemoryUsage

Get GPU memory usage by FlexTensor in inference mode.

Returns the memory used by GPU transfer blocks and unmapped tensors that were moved to GPU. This method should be called after the manager has transitioned to inference mode (after discovery and profiling phases).

RETURNS	DESCRIPTION
`GPUMemoryUsage`	Memory breakdown with per-component bytes and MB values. See `GPUMemoryUsage` for field details. TYPE: `GPUMemoryUsage`

RAISES	DESCRIPTION
`RuntimeError`	If called before inference mode is initialized (tensor_layer_loader is None)

Examples:

>>> tensor_manager.initialize_inference()
>>> usage = tensor_manager.get_gpu_memory_usage()
>>> print(f"GPU blocks: {usage.blocks_mb:.2f} MB")
>>> print(f"Unmapped tensors: {usage.unmapped_tensors_mb:.2f} MB")
>>> print(f"Total: {usage.total_mb:.2f} MB")

get_memory_transfer_stats ¶

get_memory_transfer_stats() -> dict[int, float] | None

Get the memory transfer statistics computed during profiling.

RETURNS	DESCRIPTION
`dict[int, float] \| None`	Dictionary mapping tensor size (bytes) to transfer time (ms),
`dict[int, float] \| None`	or None if profiling has not yet completed.

benchmark_context ¶

benchmark_context(iterations: int = 10) -> Any

Create a benchmarking context that automatically integrates results into TensorManager.

This method provides a context manager that creates a benchmark instance using the configured benchmark_cls, and automatically copies the benchmark results to the TensorManager's stats attributes upon context exit.

PARAMETER	DESCRIPTION
`iterations`	Number of benchmark iterations to perform TYPE: `int` DEFAULT: `10`

RETURNS	DESCRIPTION
`Any`	Context manager that yields the benchmark instance (TensorBenchmarkMode).

Examples:

>>> with tensor_manager.benchmark_context() as benchmark:
...     model = {name: layer.to(torch.bfloat16) for name, layer in model.items()}
>>> # Stats are automatically integrated into tensor_manager

save_profile ¶

save_profile(profile_directory: str) -> None

Save current profile to directory.

PARAMETER	DESCRIPTION
`profile_directory`	Directory to save profile to TYPE: `str`

load_profile ¶

load_profile(profile_directory: str, model: Module) -> None

Load profile from directory and restore state to model.

This method loads a saved TensorManagerState and restores it to the manager, but does NOT configure loaders or finalize the manager for inference. It only sets the internal tensor_manager_state and associates the model.

After calling this method, you MUST call the following methods in order:

initialize_warmup() - Configures tensor layer loaders via prepare_infer_load_mode() and patches the model with inference traps via prepare_final_model().
initialize_profile() - Prepares the model for profiling mode. When a saved state exists, this is a no-op and returns the model as-is. However, if there are changes that require re-profiling, this method handles the profile mode setup.
initialize_inference() - Finalizes the manager for inference mode and releases the internal model reference (sets self.model = None) to free memory.

Example

tm = TensorManager(device_gpu="cuda:0", ...) tm.load_profile("/path/to/profile", model) model = tm.initialize_warmup() # Required: configures loaders and patches model model = tm.initialize_profile() # Prepares profiling (no-op if state exists) model = tm.initialize_inference() # Releases model reference

Now the model is ready for inference¶

PARAMETER	DESCRIPTION
`profile_directory`	Directory containing the profile. TYPE: `str`
`model`	Model to restore state to. The model should match the architecture used when the profile was originally saved. TYPE: `Module`

Note

Future versions may add an optional parameter to automatically run finalization (initialize_warmup, initialize_profile, initialize_inference) after loading the profile.

load_state ¶

load_state(profile_directory: str) -> TensorManagerState

Load profile state from directory without restoring.

Use this when you need access to the state object before restoring, e.g., to read state.gpu_tensors_names or state.view_tensors_names.

PARAMETER	DESCRIPTION
`profile_directory`	Directory containing the profile TYPE: `str`

RETURNS	DESCRIPTION
`TensorManagerState`	The loaded state object

restore_state ¶

restore_state(model: Module | dict[str, object] | ModelDict, state: TensorManagerState, *, load_from_shm: bool = False) -> None

Delegate to TensorManagerStateHandler for state restoration.

export_state ¶

export_state() -> TensorManagerState

Export a deep-isolated snapshot rebuilt from the live block runtime.

RETURNS	DESCRIPTION
`TensorManagerState`	A state containing current tensor IDs, view IDs, statistic names,
`TensorManagerState`	and SHM block names.

RAISES	DESCRIPTION
`RuntimeError`	If no restored state or live allocation-block controller is available.

plan_state_adoption ¶

plan_state_adoption(model: Module, state: TensorManagerState, *, host_reserve_bytes: int = _STATE_ADOPTION_HOST_RESERVE_BYTES, gpu_reserve_bytes: int = _STATE_ADOPTION_GPU_RESERVE_BYTES) -> StateTransitionPlan

Plan a read-only transition from model to a saved placement.

PARAMETER	DESCRIPTION
`model`	Model whose current tensor placement is captured. TYPE: `Module`
`state`	Saved profile describing the target placement and loader strategy. TYPE: `TensorManagerState`
`host_reserve_bytes`	Available host memory excluded from planning. The 4-GiB default leaves room for Python, framework, and loader setup; tune it for the deployment rather than treating it as an invariant. TYPE: `int` DEFAULT: `_STATE_ADOPTION_HOST_RESERVE_BYTES`
`gpu_reserve_bytes`	Available GPU memory excluded from planning. The 1-GiB default leaves room for CUDA/framework runtime growth; tune it for the deployment rather than treating it as an invariant. TYPE: `int` DEFAULT: `_STATE_ADOPTION_GPU_RESERVE_BYTES`

RETURNS	DESCRIPTION
`StateTransitionPlan`	A transition plan. This method does not move or pin tensors, allocate
`StateTransitionPlan`	loaders, patch `model`, or activate the manager.

RAISES	DESCRIPTION
`ValueError`	If the reserves are invalid, the saved profile is incompatible with this manager or model, or its metadata is invalid.
`RuntimeError`	If the planned transition exceeds available capacity.

execute_state_adoption ¶

execute_state_adoption(model: Module, plan: StateTransitionPlan) -> None

Execute a plan returned by :meth:plan_state_adoption.

This operation is not transactional. If a migration or pinning step fails, the exception reports completed and current groups, but completed groups are not rolled back. Run it before loader construction, torch.compile, or CUDA-graph capture because it replaces live tensor storages and invalidates earlier references or captures. After success, call :meth:restore_adopted_state before constructing the saved-state loader.

PARAMETER	DESCRIPTION
`model`	Model captured when the plan was created. TYPE: `Module`
`plan`	Validated storage migration and pinning plan. TYPE: `StateTransitionPlan`

restore_adopted_state ¶

restore_adopted_state(model: Module, state: TensorManagerState) -> None

Restore saved loader metadata after successful state adoption.

Rebinds saved tensor IDs, strategies, statistics, and model ownership without repeating placement preprocessing established by :meth:execute_state_adoption.

PARAMETER	DESCRIPTION
`model`	Model whose state-adoption plan was executed. TYPE: `Module`
`state`	Saved profile used to create that plan. TYPE: `TensorManagerState`

Benchmark Modes¶

These are used internally by OffloadManager and exposed for users who need custom benchmark instrumentation.

TensorBenchmarkMode ¶

TensorBenchmarkMode(device_gpu: device, pinned_memory: bool = True, iterations: int = 10, pinned_memory_limit_mb: float | None = None, host_pinner: HostPinner | None = None)

Bases: TorchFunctionMode, ABC

Abstract base class for tensor benchmarking modes.

This class defines the common interface that all tensor benchmark implementations must follow to be used interchangeably as benchmark_cls parameter in TensorManager.

Initialize tensor benchmarking mode.

PARAMETER	DESCRIPTION
`device_gpu`	Target GPU device for benchmarking TYPE: `device`
`pinned_memory`	Whether to use pinned memory for tensors TYPE: `bool` DEFAULT: `True`
`iterations`	Number of benchmark iterations TYPE: `int` DEFAULT: `10`
`pinned_memory_limit_mb`	Maximum size of tensors that will be copied to pinned_memory TYPE: `float \| None` DEFAULT: `None`
`host_pinner`	:class:`HostPinner` controlling how CPU tensors are pinned. Defaults to torch mode (PyTorch's caching pinned allocator) when `None`. TYPE: `HostPinner \| None` DEFAULT: `None`

get_results `abstractmethod` ¶

get_results() -> dict[str, Any]

Get benchmark results.

RETURNS	DESCRIPTION
`dict[str, Any]`	Dictionary containing benchmark results with keys:
`dict[str, Any]`	tensor_statistics_map: Mapping of tensor IDs to statistics
`dict[str, Any]`	tensors_map: Mapping of tensor IDs to tensors

BenchmarkReplace ¶

BenchmarkReplace(device_gpu: device, pinned_memory: bool = True, iterations: int = 10, pinned_memory_limit_mb: float | None = None, *, host_pinner: HostPinner)

Bases: TensorBenchmarkMode

Benchmark mode that measures tensor transfer times and collects statistics.

This implementation intercepts tensor operations, measures transfer times from CPU to GPU, and collects comprehensive statistics about tensor usage.

Initialize BenchmarkReplace for tensor benchmarking.

PARAMETER	DESCRIPTION
`device_gpu`	Target GPU device for benchmarking TYPE: `device`
`pinned_memory`	Whether to use pinned memory for tensors TYPE: `bool` DEFAULT: `True`
`iterations`	Number of benchmark iterations TYPE: `int` DEFAULT: `10`
`pinned_memory_limit_mb`	Maximum size of tensors that will be copied to pinned_memory TYPE: `float \| None` DEFAULT: `None`
`host_pinner`	:class:`HostPinner` controlling how CPU tensors are pinned. Required (keyword-only) so profiling stats always reflect the same pinning primitive that production transfers will use — a silent fallback to `HostPinner()` would let `pinned_memory_mode="host_register"` report `"torch"`-mode latencies. TYPE: `HostPinner`

get_results ¶

get_results() -> dict[str, Any]

Get benchmark results.

RETURNS	DESCRIPTION
`dict[str, Any]`	Dictionary containing benchmark results with keys:
`dict[str, Any]`	tensor_statistics_map: Mapping of tensor IDs to statistics
`dict[str, Any]`	tensors_map: Mapping of tensor IDs to tensors

PreloadToDevice ¶

PreloadToDevice(device_gpu: device, pinned_memory: bool = True, iterations: int = 10, pinned_memory_limit_mb: float | None = None, host_pinner: HostPinner | None = None)

Bases: TensorBenchmarkMode

Benchmark mode that immediately transfers tensors to GPU without timing measurements.

This implementation provides a simple preloading strategy without collecting detailed benchmark statistics, suitable for scenarios where you just want tensors moved to GPU without measurement overhead.

Initialize PreloadToDevice for simple tensor preloading.

PARAMETER	DESCRIPTION
`device_gpu`	Target GPU device for preloading TYPE: `device`
`pinned_memory`	Whether to use pinned memory for tensors (ignored in this implementation) TYPE: `bool` DEFAULT: `True`
`iterations`	Number of benchmark iterations (ignored in this implementation) TYPE: `int` DEFAULT: `10`
`pinned_memory_limit_mb`	Maximum size of tensors that will be copied to pinned_memory TYPE: `float \| None` DEFAULT: `None`
`host_pinner`	Stored for interface compatibility but not used; kept so callers like `TensorManager.benchmark_context` can pass it uniformly. TYPE: `HostPinner \| None` DEFAULT: `None`

get_results ¶

get_results() -> dict[str, Any]

Get benchmark results.

RETURNS	DESCRIPTION
`dict[str, Any]`	Dictionary containing benchmark results with keys:
`dict[str, Any]`	tensor_statistics_map: Mapping of tensor IDs to statistics (empty for this implementation)
`dict[str, Any]`	tensors_map: Mapping of tensor IDs to tensors

Host Pinning¶

Lower-level helpers for pinned_memory_mode dispatch and cudaHostRegister registrations. See Configuration for user-facing semantics and Troubleshooting for when to switch modes.

host_pinning ¶

Host memory pinning helpers: :class:HostPinner over PyTorch's pinned allocator and cudaHostRegister in-place registration. Construct via :func:make_host_pinner; see docs/explanation/configuration.md and docs/how-to/troubleshooting.md for user-facing semantics.

PinnedMemoryMode `module-attribute` ¶

PinnedMemoryMode = Literal['torch', 'host_register']

The two supported pinning strategies for pinned_memory_mode.

HostPinner ¶

HostPinner(registry: HostPinRegistry | None = None)

Unified entry point for pinning a CPU tensor for fast GPU transfer.

Encapsulates the pinned_memory_mode dispatch so call sites don't have to branch on "torch" vs "host_register". Always usable: an instance constructed without a registry maps :meth:pin to :meth:torch.Tensor.pin_memory so callers can hold a single object whether or not manual pinning is in effect.

Construction¶

HostPinner() — torch mode (no copy avoidance).
HostPinner(HostPinRegistry()) — host_register mode; the registry's release_all must be invoked on shutdown.

See :func:is_available for whether host_register mode is supported on the current platform.

registry `property` ¶

registry: HostPinRegistry | None

The bound :class:HostPinRegistry, or None in torch mode.

pin ¶

pin(tensor: Tensor) -> Tensor

Pin tensor for fast non-blocking GPU transfer.

Strict in both modes — failures raise RuntimeError:

host_register mode (registry attached): the same tensor with its existing storage pinned in place via cudaHostRegister.
torch mode (no registry): a new tensor backed by a fresh pinned allocation from PyTorch's caching pinned allocator. A tensor.pin_memory() failure (e.g. RLIMIT_MEMLOCK exhaustion) raises.

Requires a CUDA host; constructing a :class:HostPinner on a CPU-only host is a misuse. :func:make_host_pinner enforces this at config time.

Callers should always reassign: t = pinner.pin(t).

No-op (returns the tensor unchanged) for non-CPU tensors, meta tensors, empty tensors, and tensors that are already pinned.

is_pinned ¶

is_pinned(tensor: Tensor) -> bool

Whether tensor is currently pinned by this pinner.

Detects both PyTorch's pinned allocator (torch mode) and this pinner's :class:HostPinRegistry (host_register mode). Returns False for tensors with no usable storage pointer (e.g. meta, lazy, or fake tensors whose untyped_storage() raises); the underlying error is logged at DEBUG.

release_all ¶

release_all() -> None

Release every registration owned by the bound registry.

No-op in torch mode.

HostPinRegistry ¶

HostPinRegistry()

Track cudaHostRegister registrations so they can be released together.

One registry instance is typically owned by a :class:~flextensor.tensor_manager.TensorManager and shared with the processors, loaders, and allocation blocks that it creates. Call :meth:release_all from the manager's shutdown path.

All public methods are thread-safe; the lock is held across each cudart call so the same pointer can't be double-(un)registered.

The registry retains a strong reference to each backing tensor, so its storage stays alive until :meth:release / :meth:release_all; callers don't need a separate keep-alive list.

cudaHostRegister state is per-process. Don't share a registry across processes — children must build their own via :func:make_host_pinner after CUDA is initialised, and must not release pins inherited from the parent. Prefer multiprocessing start method "spawn" over fork.

pin_in_place ¶

pin_in_place(tensor: Tensor) -> Tensor

Pin tensor's storage in place via cudaHostRegister.

Entries are keyed by storage address, so a tensor and a view of the same tensor register only once. The first pin's size and tensor reference are the ones kept — re-pinning a view does not shrink the registered byte range.

No-op for non-CPU tensors, meta tensors, empty tensors, tensors with a null storage pointer, tensors already pinned (by PyTorch's allocator or otherwise), and pointers already registered here. Returns the same tensor object for chaining.

RAISES	DESCRIPTION
`RuntimeError`	If the CUDA runtime cannot be reached via :func:`torch.cuda.cudart` or the registration call itself returns a non-zero error code.

release ¶

release(ptr: int) -> bool

Unregister a single pointer.

Returns True only when cudaHostUnregister ran and returned a success code. Returns False if the pointer wasn't tracked, the CUDA runtime is no longer reachable, or the call returned a non-zero error code. Use :meth:is_registered to query "was this pointer tracked" without attempting to unregister.

The cudart-unavailable and non-zero-rc branches log at WARNING with ptr, rc (where applicable), and the call name so operators running at the default log level can see when registry and kernel state have drifted apart. The "pointer wasn't tracked" branch returns False without a WARNING (DEBUG-level only) — callers can pre-check with :meth:is_registered if they need to distinguish that case.

On failure the entry stays tracked so the tensor is not freed and :meth:release can be called again to retry.

release_all ¶

release_all() -> None

Unregister every pointer the registry is tracking.

Logs at WARNING when the CUDA runtime is no longer reachable (so the unregister calls are skipped) and per cudaHostUnregister failure, so operators running at the default log level see when registry and kernel state have drifted apart at shutdown.

Successfully released entries are removed; failed ones stay tracked so their tensors are not freed and :meth:release_all can be called again to retry.

NoOpHostPinner ¶

NoOpHostPinner()

Bases: HostPinner

Pinner that performs no pinning — every pin call returns the tensor unchanged. Returned by :func:make_host_pinner when pinned_memory=False so call sites can hold a typed :class:HostPinner without each having to branch on the public flag.

registry `property` ¶

registry: HostPinRegistry | None

The bound :class:HostPinRegistry, or None in torch mode.

is_pinned ¶

is_pinned(tensor: Tensor) -> bool

Whether tensor is currently pinned by this pinner.

Detects both PyTorch's pinned allocator (torch mode) and this pinner's :class:HostPinRegistry (host_register mode). Returns False for tensors with no usable storage pointer (e.g. meta, lazy, or fake tensors whose untyped_storage() raises); the underlying error is logged at DEBUG.

release_all ¶

release_all() -> None

Release every registration owned by the bound registry.

No-op in torch mode.

is_available ¶

is_available() -> bool

Return True if manual host pinning via cudaHostRegister is usable.

make_host_pinner ¶

make_host_pinner(pinned_memory: bool, mode: PinnedMemoryMode) -> HostPinner

Build a :class:HostPinner from the public (pinned_memory, mode) config.

Returns :class:NoOpHostPinner when pinned_memory=False (the legitimate opt-out). For pinned_memory=True, requires a CUDA host: offloading without a GPU has no purpose, and silently degrading to pageable transfers would mask the misconfiguration as a perf regression.

For mode="host_register" returns a :class:HostPinner backed by a fresh :class:HostPinRegistry when cudart is usable; on a CUDA host where cudart is broken or missing, falls back to torch mode and surfaces a WARNING log naming the cause.

PARAMETER	DESCRIPTION
`pinned_memory`	Whether pinning is requested at all. `False` short-circuits to a :class:`NoOpHostPinner` regardless of `mode`. TYPE: `bool`
`mode`	Pinning strategy when `pinned_memory=True` — :class:`PinnedMemoryMode` (`"torch"` or `"host_register"`). TYPE: `PinnedMemoryMode`

RETURNS	DESCRIPTION
`A`	class:`HostPinner` (or :class:`NoOpHostPinner` when pinning is TYPE: `HostPinner`
`HostPinner`	disabled) wired to the appropriate backend. Callers always get a
`usable`	class:`HostPinner`; only the kind of pinning differs. TYPE: `HostPinner`

RAISES	DESCRIPTION
`RuntimeError`	If `pinned_memory=True` but `torch.cuda.is_available()` is `False`. Set `pinned_memory=False` on intentional CPU-only deployments.

Note

Logs at WARNING level when mode="host_register" cannot use torch.cuda.cudart() and falls back to torch mode, or when a non-torch mode is ignored because pinned_memory=False.

Shared Memory¶

SharedOffloadSession provides coordinated loading from profiles with application-supplied creator and follower model factories. See Shared Memory for session setup, restoration, and cleanup.

SharedOffloadSession ¶

SharedOffloadSession(namespace: str, *, wait_timeout: Any = 300.0)

Own multiple shared offload managers for one worker process.

Construct one session independently in every worker after process spawning. Workers that use the same base namespace elect a creator independently for each manager name passed to :meth:offload_from_profile.

The session is process- and thread-affine: construct it after spawning and perform every load and close operation from that same thread.

The creator factory is invoked only in the elected creator. Followers wait for publication and then invoke the follower factory, which should normally construct the same model topology on the meta device.

PARAMETER	DESCRIPTION
`namespace`	Run-scoped namespace shared by all participating workers. TYPE: `str`
`wait_timeout`	Timeout in seconds for followers waiting for readiness. Defaults to 300 seconds. Zero means no overall limit; creator heartbeat checks remain active. TYPE: `Any` DEFAULT: `300.0`

validate_options `staticmethod` ¶

validate_options(namespace: object, *, wait_timeout: Any = 300.0) -> None

Validate raw constructor options without creating a session or SHM resources.

PARAMETER	DESCRIPTION
`namespace`	Candidate run-scoped namespace shared by all workers. TYPE: `object`
`wait_timeout`	Candidate follower-readiness timeout, in seconds. Zero means no overall limit. TYPE: `Any` DEFAULT: `300.0`

RAISES	DESCRIPTION
`ValueError`	If the namespace is not POSIX-safe or the timeout is not a non-negative finite number.

namespace_for ¶

namespace_for(name: str) -> str

Return a compact independent SHM namespace for a named manager.

offload_from_profile ¶

offload_from_profile(profile_directory: str | PathLike[str], *, config: OffloadConfig, name: str = DEFAULT_MANAGER_NAME, creator_model_factory: Callable[[], Module], follower_model_factory: Callable[[], Module], materialize_device: device | str | None = None, compile_fn: Callable[[Module], Module] | None = None) -> Module

Elect a role and restore a model directly into its saved placement.

Both factories must construct matching registered parameter and buffer names. The creator factory may load checkpoint data; the follower factory must not depend on checkpoint tensor payload.

PARAMETER	DESCRIPTION
`profile_directory`	Directory containing the saved FlexTensor profile. Only the elected creator reads it. TYPE: `str \| PathLike[str]`
`config`	Enabled allocation-block transfer configuration. The session copies it and supplies its own SHM namespace and wait timeout without mutating the caller's object. TYPE: `OffloadConfig`
`name`	Named offload manager to initialize. The name also selects an independent child namespace and creator election. TYPE: `str` DEFAULT: `DEFAULT_MANAGER_NAME`
`creator_model_factory`	Zero-argument factory that constructs the checkpoint-backed model. It runs only in the elected creator. TYPE: `Callable[[], Module]`
`follower_model_factory`	Zero-argument factory that constructs the same registered model topology without checkpoint payload. It runs only in followers and normally builds on `meta`. TYPE: `Callable[[], Module]`
`materialize_device`	Device for follower-resident parameters and buffers. Defaults to the CUDA device selected by config. TYPE: `device \| str \| None` DEFAULT: `None`
`compile_fn`	Optional per-unit compiler applied during restored-state initialization. TYPE: `Callable[[Module], Module] \| None` DEFAULT: `None`

RETURNS	DESCRIPTION
`Module`	The normal FlexTensor model proxy in inference phase.

RAISES	DESCRIPTION
`ValueError`	If name or config is invalid for shared offload.
`RuntimeError`	If the session or manager lifecycle is incompatible with the requested load, or shared state is invalid.
`TimeoutError`	If a follower exceeds the coordination timeout.
`ShmCreatorFailedError`	If the creator publishes an initialization failure.
`ShmCreatorDiedError`	If the creator heartbeat expires while it is still initializing.
`ShmCoordinatorError`	If the existing generation has no live heartbeat lease. Stop workers and use :meth:`cleanup_namespace` before reusing a dead namespace.

Note

Exceptions raised by either model factory or by compile_fn propagate unchanged after transactional cleanup.

close ¶

close(*, preserve_models: bool = False) -> None

Release loaded managers in reverse order.

Closing is idempotent. Failed teardown retains the remaining handles so another call can retry it.

PARAMETER	DESCRIPTION
`preserve_models`	Clone loader-managed values into private storage before teardown. Keep the default `False` for disposable serving workers; their returned proxies become unusable. TYPE: `bool` DEFAULT: `False`

Coordination, shared ownership, and resident tensor storage for cross-process weight sharing. See Low-level SHM protocol and ownership for publication, failure, and lifetime semantics.

ShmCoordinator ¶

ShmCoordinator(namespace: str, keep_alive_seconds: int | float = DEFAULT_KEEP_ALIVE_SECONDS)

Orchestrates creator/follower SHM sharing.

On init, tries to connect to an existing coordination block: - If it exists: validates version and an existing live participant under the setup lock before becoming a follower. - If it doesn't exist: becomes creator, writes version header.

PARAMETER	DESCRIPTION
`namespace`	SHM namespace shared by the intended creator and followers. TYPE: `str`
`keep_alive_seconds`	Timeout for participant-heartbeat liveness detection. TYPE: `int \| float` DEFAULT: `DEFAULT_KEEP_ALIVE_SECONDS`

memory_owner `property` ¶

memory_owner: FlexibleSharedMemory

The shared lifetime for this manager's profile and tensor blocks.

wait_for_ready ¶

wait_for_ready(timeout: float | None = None) -> None

Block until the creator publishes READY or a terminal failure.

PARAMETER	DESCRIPTION
`timeout`	Maximum seconds to wait, or `None` to wait indefinitely. TYPE: `float \| None` DEFAULT: `None`

RAISES	DESCRIPTION
`TimeoutError`	If the namespace is not ready before timeout expires.
`ShmCreatorFailedError`	If the creator publishes FAILED.
`ShmCreatorDiedError`	If its heartbeat expires while INITIALIZING.

notify_ready ¶

notify_ready() -> None

Publish READY. Followers observe the authoritative header by polling.

mark_failed ¶

mark_failed(error: BaseException) -> None

Publish a creator initialization failure.

PARAMETER	DESCRIPTION
`error`	Original creator-side failure. Only its type and sanitized, bounded string representation are published. TYPE: `BaseException`

close ¶

close() -> None

Release profile data before the control plane, retaining failed handles.

ShmCoordinatorError ¶

Bases: RuntimeError

Base class for coordinator-level shared initialization failures.

ShmCreatorFailedError ¶

ShmCreatorFailedError(namespace: str, failure_message: str)

Bases: ShmCoordinatorError

Raised in followers when the creator explicitly publishes a failure.

ShmCreatorDiedError ¶

ShmCreatorDiedError(namespace: str, creator_identity: int | None)

Bases: ShmCoordinatorError

Raised in followers when the initializing creator heartbeat expires.

creator_identity is the creator participant ID, not an OS PID. The constructor keyword and attribute name are retained for compatibility.

AllocationManager ¶

AllocationManager(shm_block_name_prefix: str | None = None, load_from_shm: bool = False, pinned_memory: bool = False, *, host_pinner: HostPinner, shm_owner: FlexibleSharedMemory | None = None, memory_alignment: int = _DEFAULT_PACKED_TENSOR_ALIGNMENT_BYTES, lock_class: type[BaseLock] | None = None, release_tensor_memory: bool = False)

A manager for allocation blocks.

SharedTensorStore ¶

SharedTensorStore(shm: FlexibleSharedMemory, descriptors: tuple[_TensorDescriptor, ...], data_offset: int)

Share resident model tensors between serving processes.

Instances own one :class:FlexibleSharedMemory handle. Materialization always copies out of that mapping, so closing the store cannot invalidate tensors already installed in a model.

create_from_model `classmethod` ¶

create_from_model(namespace: str, model: Module, *, parameter_names: Iterable[str], shm_owner: FlexibleSharedMemory | None = None) -> SharedTensorStore

Create and fully publish a resident tensor store from model.

attach `classmethod` ¶

attach(namespace: str, *, shm_owner: FlexibleSharedMemory | None = None) -> SharedTensorStore

Attach to and validate an existing resident tensor store.

materialize_into ¶

materialize_into(model: Module, *, device: device | str) -> None

Validate, copy, and rebind stored tensors into model.

All descriptors and all target registrations are validated before the first model attribute changes.

close ¶

close() -> None

Close the FlexibleSharedMemory handle.

Advanced API¶

TensorManager ¶

set_model ¶

initialize_warmup ¶

initialize_profile ¶

initialize_inference ¶

prepare_warmup_mode ¶

prepare_profile_mode ¶

prepare_infer_mode ¶

trap ¶

release_memory ¶

get_gpu_memory_usage ¶

get_memory_transfer_stats ¶

benchmark_context ¶

save_profile ¶

load_profile ¶

Now the model is ready for inference¶

load_state ¶

restore_state ¶

export_state ¶

plan_state_adoption ¶

execute_state_adoption ¶

restore_adopted_state ¶

Benchmark Modes¶

TensorBenchmarkMode ¶

get_results abstractmethod ¶

BenchmarkReplace ¶

get_results ¶

PreloadToDevice ¶

get_results ¶

Host Pinning¶

host_pinning ¶

PinnedMemoryMode module-attribute ¶

HostPinner ¶

Construction¶

registry property ¶

pin ¶

is_pinned ¶

release_all ¶

HostPinRegistry ¶

pin_in_place ¶

release ¶

release_all ¶

NoOpHostPinner ¶

registry property ¶

is_pinned ¶

release_all ¶

is_available ¶

make_host_pinner ¶

Shared Memory¶

SharedOffloadSession ¶

validate_options staticmethod ¶

namespace_for ¶

offload_from_profile ¶

close ¶

ShmCoordinator ¶

memory_owner property ¶

wait_for_ready ¶

notify_ready ¶

mark_failed ¶

close ¶

ShmCoordinatorError ¶

ShmCreatorFailedError ¶

ShmCreatorDiedError ¶

AllocationManager ¶

SharedTensorStore ¶

create_from_model classmethod ¶

attach classmethod ¶

materialize_into ¶

close ¶

get_results `abstractmethod` ¶

PinnedMemoryMode `module-attribute` ¶

registry `property` ¶

registry `property` ¶

validate_options `staticmethod` ¶

memory_owner `property` ¶

create_from_model `classmethod` ¶

attach `classmethod` ¶