Copy for mutable frames can introduce a slowdown

[jakirkham]

Recently we made a fix to ensure all frames are mutable in PR ( #3967 ). This solves some issues where immutable frames were used to back Python objects.

One downside of that change is it effectively requires a copy of the frames be taken for all objects. This can be unnecessary in cases where the data may already get copied as is the case with Python builtin objects that like to own their memory or GPU objects that require data be moved from host to device (effectively a copy). As a result this copy can introduce a slowdown.

Included below is an example with and without that change using cuDF:

Without mutable frames:

In [1]: import cupy 
   ...: import cudf 
   ...: import rmm 
   ...:  
   ...: from distributed.protocol import serialize_bytes, deserialize_bytes     

In [2]: rmm.reinitialize(pool_allocator=True, 
   ...:                  initial_pool_size=int(30 * 2**30))                     

In [3]: df = cudf.DataFrame({ 
   ...:     k: cupy.random.random(1_000_000) 
   ...:     for i, k in enumerate(map(chr, range(ord("A"), ord("K")))) 
   ...: })                                                                      

In [4]: b = serialize_bytes(df)                                                 

In [5]: %timeit deserialize_bytes(b)                                            
10.1 ms ± 61.6 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

With mutable frames:

In [1]: import cupy 
   ...: import cudf 
   ...: import rmm
   ...: from distributed.protocol import serialize_bytes, deserialize_bytes

In [2]: rmm.reinitialize(pool_allocator=True,
   ...:                  initial_pool_size=int(30 * 2**30))

In [3]: df = cudf.DataFrame({
   ...:     k: cupy.random.random(1_000_000)
   ...:     for i, k in enumerate(map(chr, range(ord("A"), ord("K"))))
   ...: })

In [4]: b = serialize_bytes(df)

In [5]: %timeit deserialize_bytes(b)
45.4 ms ± 65.3 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

[jsignell]

That is a pretty sizable performance hit. Do you have a proposed solution? I suppose the user could pick whether frames are mutable or not. Is there a smart way for us to determine that for them so that we can provide the best default experience?

[jakirkham]

Yeah I have a WIP branch where I tried to push this more into specific object or protocol serialization. It has some test failures though, which suggest it's not quite sufficient.

The idea of picking (or perhaps even identifying) which frames are mutable is an interesting one. Maybe that's a better way to go. Thanks for the suggestion 🙂

[jakirkham]

That wound up being much easier to do as well. Submitted as PR ( #4004 ) 😄

[jakirkham]

This is in 2.23.0+.