Reserve the memory into_packed_data() allocates

[madsbk]

Depends on #23892.

When a table_chunk is not already packed, into_packed_data() calls cudf::pack() straight into br->device_mr(). The allocation is roughly the size of the table and nothing reserved it, so the memory system under-counts and backpressure cannot see it coming.

It now takes a MemoryReservation instead of a BufferResource, and into_packed_data_cost() reports what it will allocate. That is zero when the chunk is already packed, since the packed data is then moved out rather than serialized, so the sole caller reserves the right amount on every path rather than over-reserving on the spilled one.

Both non-packed constructors already compute cudf::packed_size() into data_alloc_size_, so the cost needs no extra work at call time.

Breaking change

into_packed_data() takes a MemoryReservation rather than a BufferResource, in C++ and in Python. The BufferResource overload is gone rather than kept alongside, since every other allocating method on table_chunk already takes a reservation and leaving it would leave a blocking path that is easier to reach than the correct one.

Callers pair it with into_packed_data_cost():

chunk, extra = await make_table_chunks_available_or_wait(
    context, chunk, reserve_extra=chunk.into_packed_data_cost(), net_memory_delta=0
)
packed = chunk.into_packed_data(extra)

Note

The cudf::chunked_pack() TODO on the same line is still open. A bounce buffer would bound the allocation instead of needing room for a whole second copy.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.