IWeightsManager

class tensorrt_rtx.IWeightsManager

Manages the saved and loaded weights of a weight-stripped engine.

Create this object with ICudaEngine.create_weights_manager(). TensorRT-RTX keeps the engine alive while its weights manager exists. Methods that change weight memory must not run concurrently with inference, refit, serialization, execution-context creation, weight-streaming changes, or another weights-manager method for the same engine.

__del__(self: tensorrt_rtx.tensorrt_rtx.IWeightsManager) → None
__init__(*args, **kwargs)
get_size(self: tensorrt_rtx.tensorrt_rtx.IWeightsManager) → int

Return the number of bytes written when saving or required when restoring the engine weights.

The size includes any padding required by CUDA virtual-memory allocation granularity.

restore_from_buffer(self: tensorrt_rtx.tensorrt_rtx.IWeightsManager, source: SupportsInt, size: SupportsInt, location: tensorrt_rtx.tensorrt_rtx.TensorLocation, stream_handle: SupportsInt) → bool

Restore the engine weights by copying from an application buffer.

The weights must be unloaded or not yet loaded. The copy is ordered on stream_handle. The application must keep the source memory valid until the operation completes.

Parameters:

  • source – Integer address of the source buffer.

  • size – Source size in bytes; must equal get_size().

  • location – Location of the source buffer.

  • stream_handle – Integer handle of the CUDA stream used to order the operation.

Returns:

True on success, false otherwise.

restore_from_vmm_allocation(self: tensorrt_rtx.tensorrt_rtx.IWeightsManager, allocation_handle: SupportsInt, offset: SupportsInt, size: SupportsInt, stream_handle: SupportsInt) → bool

Restore a range of engine weights by mapping an application CUDA VMM allocation without copying.

The weights must be unloaded or not yet fully initialized before the first call. Multiple calls may map nonoverlapping ranges from different allocations. The engine remains unavailable until the mappings cover the complete range [0, get_size()).

The allocation must have been created by cuMemCreate() on the engine’s CUDA device and be at least size bytes. Its contents beginning at physical allocation offset zero must match the range [offset, offset + size) of the weight data produced by this weights-manager interface. Both offset and size must satisfy CUDA’s minimum virtual-memory allocation granularity. The range [offset, offset + size) must lie within [0, get_size()).

The application owns the allocation handle and must keep it valid until this call returns. After a successful call, the application may release the handle because TensorRT-RTX’s mapping keeps the allocation alive. To reuse the allocation after unload() or engine destruction, the application must retain the handle.

Parameters:

  • allocation_handle – Integer value of the CUmemGenericAllocationHandle.

  • offset – Byte offset of the destination range within the saved weight data.

  • size – Size of the range in bytes.

  • stream_handle – Integer handle of the CUDA stream used to order the operation.

Returns:

True on success, false otherwise.

save_to_buffer(self: tensorrt_rtx.tensorrt_rtx.IWeightsManager, destination: SupportsInt, size: SupportsInt, location: tensorrt_rtx.tensorrt_rtx.TensorLocation, stream_handle: SupportsInt) → bool

Save the loaded engine weights to an application buffer.

The copy is ordered on stream_handle. The application must keep the destination memory valid until the operation completes.

Parameters:

  • destination – Integer address of the destination buffer.

  • size – Destination size in bytes; must equal get_size().

  • location – Location of the destination buffer.

  • stream_handle – Integer handle of the CUDA stream used to order the operation.

Returns:

True on success, false otherwise.

save_to_stream(self: tensorrt_rtx.tensorrt_rtx.IWeightsManager, writer: tensorrt_rtx.v_1_0::IStreamWriter, stream_handle: typing.SupportsInt) → bool

Save the loaded engine weights to an application stream writer.

TensorRT-RTX synchronizes stream_handle before invoking the writer for each copied chunk.

Parameters:

  • writer – Destination IStreamWriter.

  • stream_handle – Integer handle of the CUDA stream containing prior engine work.

Returns:

True on success, false otherwise.

unload(self: tensorrt_rtx.tensorrt_rtx.IWeightsManager, stream_handle: SupportsInt) → bool

Unload the engine weights while preserving their virtual address range.

The application must order all prior engine work on stream_handle. This call synchronizes the stream before unmapping the physical memory. TensorRT-RTX preserves the virtual address range until the associated ICudaEngine is destroyed, allowing the weights to be restored later.

Parameters:

stream_handle – Integer handle of the CUDA stream containing prior engine work.

Returns:

True on success, false otherwise.