-
Notifications
You must be signed in to change notification settings - Fork 231
Description
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
Projects
Status