Skip to content

[FEA] Add "fallback" memory resource #2074

@GregoryKimball

Description

@GregoryKimball

Is your feature request related to a problem? Please describe.
We would like to test a "fallback" memory resource, where allocations are attempted on the first MR, and if the allocation is not possible, then allocations are attempted on a second MR.

Describe the solution you'd like
In 25.10, cudf-polars uses this memory resource by default:

new_mr = rmm.mr.PrefetchResourceAdaptor(
                rmm.mr.PoolMemoryResource(
                    rmm.mr.ManagedMemoryResource(),
                    initial_pool_size=free_memory,
                )
            )

and we would like to experiment with an MR like this:

async_mr = rmm.mr.CudaAsyncMemoryResource()
managed_mr = rmm.mr.PrefetchResourceAdaptor(
                rmm.mr.PoolMemoryResource(
                    rmm.mr.ManagedMemoryResource(),
                    initial_pool_size=free_memory,
                )
            )
new_mr = rmm.mr.FallbackResourceAdaptor(
    async_mr, managed_mr)

Describe alternatives you've considered
I don't know if there is a way to compose something like this today. CallbackResourceAdaptor might be able to do a try-catch. Maybe there is a way to mock this up without changing the default MR back and forth.

Additional context
We would like to test this idea in cudf-polars and Velox-cuDF to only start using managed memory when it is needed. We might need to set a maximum pool size on the async first MR to give the second managed MR enough room to still prefetch.

Metadata

Metadata

Assignees

Labels

Type

No type

Projects

Status

To-do

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions