[REVIEW] Copy on write implementation

[galipremsagar]

This PR introduces copy-on-write. As the name suggests when copy-on-write is enabled, when there is a shallow copy of a column made, both the columns share the same memory and only when there is a write operation being performed on either the parent or any of it's copies a true copy will be triggered. Copy-on-write(c-o-w) can be enabled in two ways:

1. Setting CUDF_COPY_ON_WRITE environment variable to 1 will enable c-o-w, unsetting will disable c-o-w.
2. Setting copy_on_write option in cudf options by doing cudf.set_option("copy_on_write", True) to enable it and cudf.set_option("copy_on_write", False) to disable it.

Note: Copy-on-write is not being enabled by default, it is being introduced as an opt-in.

A valid performance comparison can be done only with copy_on_write=OFF + .copy(deep=True) vs copy_on_write=ON + .copy(deep=False):

In [1]: import cudf

In [2]: s = cudf.Series(range(0, 100000000))

# branch-23.02 : 1209MiB
# This-PR : 1209MiB

In [3]: s_copy = s.copy(deep=True) #branch-23.02
In [3]: s_copy = s.copy(deep=False) #This-PR

# branch-23.02 : 1973MiB
# This-PR : 1209MiB

In [4]: s
Out[4]: 
0                  0
1                  1
2                  2
3                  3
4                  4
              ...   
99999995    99999995
99999996    99999996
99999997    99999997
99999998    99999998
99999999    99999999
Length: 100000000, dtype: int64

In [5]: s_copy
Out[5]: 
0                  0
1                  1
2                  2
3                  3
4                  4
              ...   
99999995    99999995
99999996    99999996
99999997    99999997
99999998    99999998
99999999    99999999
Length: 100000000, dtype: int64

In [6]: s[2] = 10001

# branch-23.02 : 3121MiB
# This-PR : 3121MiB

In [7]: s
Out[7]: 
0                  0
1                  1
2              10001
3                  3
4                  4
              ...   
99999995    99999995
99999996    99999996
99999997    99999997
99999998    99999998
99999999    99999999
Length: 100000000, dtype: int64

In [8]: s_copy
Out[8]: 
0                  0
1                  1
2                  2
3                  3
4                  4
              ...   
99999995    99999995
99999996    99999996
99999997    99999997
99999998    99999998
99999999    99999999
Length: 100000000, dtype: int64

Stats around the performance and memory gains :

• Memory usage of new copies will be 0 GPU memory additional overhead i.e., users will save 2x, 5x, 10x,...20x memory usage for making 2x, 5x, 10x,...20x deep copies respectively. So, The more you copy the more you save 😉(as long as you don't write on all of them)
• copying times are now cut by 99% for all dtypes when copy-on-write is enabled(copy_on_write=OFF + .copy(deep=True) vs copy_on_write=ON + .copy(deep=False)).

In [1]: import cudf

In [2]: df = cudf.DataFrame({'a': range(0, 1000000)})

In [3]: df = cudf.DataFrame({'a': range(0, 100000000)})

In [4]: df['b'] = df.a.astype('str')

# GPU memory usage
# branch-23.02 : 2345MiB
# This-PR : 2345MiB

In [5]: df
Out[5]: 
                 a         b
0                0         0
1                1         1
2                2         2
3                3         3
4                4         4
...            ...       ...
99999995  99999995  99999995
99999996  99999996  99999996
99999997  99999997  99999997
99999998  99999998  99999998
99999999  99999999  99999999

[100000000 rows x 2 columns]

In [6]: def make_two_copies(df, deep):
   ...:     return df.copy(deep=deep), df.copy(deep=deep)
   ...: 

In [7]: x, y = make_two_copies(df, deep=True) # branch-23.02
In [7]: x, y = make_two_copies(df, deep=False) # This PR

# GPU memory usage
# branch-23.02 : 6147MiB
# This-PR : 2345MiB

In [8]: %timeit make_two_copies(df, deep=True) # branch-23.02
In [8]: %timeit make_two_copies(df, deep=False) # This PR

# Execution times
# branch-23.02 : 135 ms ± 4.49 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
# This-PR : 100 µs ± 879 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

• Even when copy-on-write is disabled, string, list & struct columns deep copies are now 99% faster

In [1]: import cudf

In [2]: s = cudf.Series(range(0, 100000000), dtype='str')

In [3]: %timeit s.copy(deep=True)


# branch-23.02 : 28.3 ms ± 5.32 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
# This PR : 19.9 µs ± 93.5 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)


In [9]: s = cudf.Series([[1, 2], [2, 3], [3, 4], [4, 5], [6, 7]]* 10000000)

In [10]: %timeit s.copy(deep=True)
# branch-23.02 : 25.7 ms ± 5.99 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
# This-PR : 44.2 µs ± 1.61 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)


In [4]: df = cudf.DataFrame({'a': range(0, 100000000), 'b': range(0, 100000000)[::-1]})

In [5]: s = df.to_struct()

In [6]: %timeit s.copy(deep=True)

# branch-23.02 : 42.5 ms ± 3.5 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)
# This-PR : 89.7 µs ± 1.68 µs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)

• Add pytests
• Docs page explaining copy on write and how to enable/disable it.

[vyasr]

Overall I'm pretty happy with this PR. The level of complexity is significantly lower than I had originally feared it would be, especially after looking through spilling. My main outstanding concerns are around 1) the different new properties introduced for the data an mask, and 2) the potential confusion around the different methods for accessing pointers, especially as it pertains to their different behaviors in base/COW/spilling. We also need to deal with the issue Lawrence raised about slices, but hopefully my suggestion there will be sufficient. I don't want the PR's scope to sprawl, but I also don't want to merge anything that adds even more complexity in those areas, especially if it does so in ways that make the code harder to understand even when COW or spilling are disabled. We may want to meet and just discuss minimal solutions for those problems.

Aside from the API-level implications of COW, I would also like for us to come out of this project with a consistent story on column mutability (or at least a plan for having one). @shwina and I have previously discussed that we would love for columns to be immutable, and that has been about 90% true for a long time. COW gets us closer, but the addition of the mutable_ptr attribute for buffers is a clear indicator that we have one important exception (__setitem__ calls) that may always have to exist. I think finding a clear way not only to document this, but to design the API around this fact, should help reduce some of issue (2) above regarding the different accessors.

Regarding the docs guides, I'm happy to make PRs into yours with proposed changes to the docs so that you don't have to implement all of my suggestions (they're definitely quite a bit of work). I didn't make any new comments there in this review, but there are still some outstanding major requests from the first round (like structuring the developer docs around having already read the user guide) that I am happy to draft for review to help move things forward.

[vyasr]

I agree, my first review was me getting a lot of ideas out without really thinking through scoping, but I'm certainly in favor of limiting the scope of this PR and not trying to rewrite existing internals more than absolutely necessary. We can do that in follow-ups.

[vyasr]

I want to limit the scope of this PR, but TBH I also don't want to merge it until we at least have a plan for how to handle this. The spilling code already makes buffer internals too hard to follow without some serious care. I think we need to put a hard cap on any changes that increase the complexity further without having a plan for how to fix it. Otherwise we will quickly find ourselves in a state where nobody will be able to debug this code without spending way too much time relearning each time (I still feel like I have to relearn the spilling internals every time I review a spilling PR, which I find quite concerning and I think we need to prioritize reducing that complexity and mitigating the fallout).

[vyasr]

IIUC the option being registered below should indeed guarantee that we won't get a KeyError

[vyasr]

Why was this change necessary when it wasn't before? What difference does COW make here?

[galipremsagar]

This isn't COW related. This was necessary to specifically test the special handling in CategoricalColumn.copy, where we will need to deep copy the categories especially. See test_index_copy_category below that uses this now.

[vyasr]

Ah got it that makes sense.

[vyasr]

I think it would be worth us meeting (with other interested parties like @wence- and @shwina) to discuss the minimal changes needed to address the issues with confusing APIs introduced in this PR.

[madsbk]

What about something like this?

Suggested change

	# Because hasattr will trigger invocation of
	# `__cuda_array_interface__` which could
	# be expensive in CopyOnWriteBuffer case.
	value_cai = inspect.getattr_static(
	value,
	"__cuda_array_interface__",
	None
	)
	has_cai = isinstance(value, Buffer) or hasattr(value, "__cuda_array_interface__")

Or even better, we could create a small PR that implement an util function to access __cuda_array_interface__, which we could use exclusively. This would also simplify SpillableBuffer.

[galipremsagar]

I'm planning to add a utility to remove this code duplication across the repo.

[galipremsagar]

@shwina @vyasr and I synced up offline and created and came up with the following todo items:

• @shwina created a feature branch copy-on-write, we will merge this PR into it.
• Prem to open a follow-up PR to copy-on-write making buffer.get_ptr(mode=r/w) , Column.data_array_view(mode=r/w) APIs. No more Buffer.ptr or .mutable_ptr should exist. If things break we can fix it in other repos but we expect not much disruption outside of rapids. This also means getting rid of the spill_lock argument to SpillableBuffer.get_ptr(). In its place, we want get_ptr() to just always be able to call get_spill_lock(), which means only using it in an acquire_spill_lock context. [REVIEW] Change ways to access ptr in Buffer  #12587
• Add stronger reasons to the Advantages section in user docs for COW.
• Needs some final passes through docs.
• Once all of the above are completed merge back copy-on-write to branch-23.02

Far in the future items(next releases):

• To think of how to avoid copying of an entire Buffer when we are modifying a slice of it.
• Try to get rid of cuda_array_interface at Buffer level and see its feasibility.
• Ensuring that dask-cuda and ucx-py aren't interacting with cudf at the buffer level.
• Devising a single long-term API to expose cudf's data pointers externally.

[shwina]

I'm good with merging this to the copy-on-write branch.