Allow setting malloc heap size in string udfs #12094

brandon-b-miller · 2022-11-08T18:25:05Z

Adds a mechanism for setting the default cuda malloc heap size for string UDFs, with 2gb default.

codecov · 2022-11-08T20:58:24Z

Codecov Report

Base: 87.47% // Head: 88.22% // Increases project coverage by +0.75% 🎉

Coverage data is based on head (ba1d690) compared to base (f817d96).
Patch has no changes to coverable lines.

❗ Current head ba1d690 differs from pull request most recent head b400143. Consider uploading reports for the commit b400143 to get more accurate results

Additional details and impacted files

@@               Coverage Diff                @@
##           branch-22.12   #12094      +/-   ##
================================================
+ Coverage         87.47%   88.22%   +0.75%     
================================================
  Files               133      137       +4     
  Lines             21826    22571     +745     
================================================
+ Hits              19093    19914     +821     
+ Misses             2733     2657      -76

Impacted Files	Coverage Δ
python/cudf/cudf/core/column/interval.py	`85.45% <0.00%> (-9.10%)`	⬇️
python/strings_udf/strings_udf/__init__.py	`75.80% <0.00%> (-8.51%)`	⬇️
python/cudf/cudf/io/text.py	`91.66% <0.00%> (-8.34%)`	⬇️
python/cudf/cudf/core/_base_index.py	`81.28% <0.00%> (-4.27%)`	⬇️
python/cudf/cudf/io/json.py	`92.06% <0.00%> (-2.68%)`	⬇️
python/cudf/cudf/utils/utils.py	`89.91% <0.00%> (-0.69%)`	⬇️
python/cudf/cudf/core/column/timedelta.py	`90.17% <0.00%> (-0.58%)`	⬇️
python/cudf/cudf/core/column/datetime.py	`89.21% <0.00%> (-0.51%)`	⬇️
python/cudf/cudf/core/column/column.py	`87.96% <0.00%> (-0.46%)`	⬇️
python/dask_cudf/dask_cudf/core.py	`73.72% <0.00%> (-0.41%)`	⬇️
... and 46 more

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

☔ View full report at Codecov.
📢 Do you have feedback about the report comment? Let us know in this issue.

davidwendt · 2022-11-09T13:26:12Z

python/strings_udf/strings_udf/__init__.py

@@ -87,6 +88,26 @@ def _get_ptx_file():
        return regular_result[1]


+default_heap_size = int(2e6)


I feel like this needs a big comment on why this is the default. This number seems a bit high considering 10M rows of 100 bytes each would only be 1GB. I suppose it makes sense since it is the largest strings column in bytes possible.

Would it be too strange to make this configurable through an environment variable perhaps?
At least for this release until we can come up with a good metric for setting it.

I introduced an environment variable to control this. I'm still open to discussion here as to what to set the value as. Indeed I did base it on the size of a strings column.

python/strings_udf/strings_udf/__init__.py

Co-authored-by: Bradley Dice <[email protected]>

… into fea-stringudf-heap

bdice · 2022-11-16T19:26:28Z

python/strings_udf/strings_udf/__init__.py

@@ -87,6 +88,29 @@ def _get_ptx_file():
        return regular_result[1]


+# Maximum size of a string column is 2gb
+STRINGS_UDF_DEFAULT_HEAP_SIZE = 2**31


Should these names be module-scope public or private?

bdice · 2022-11-16T19:28:30Z

python/strings_udf/strings_udf/__init__.py

+heap_size = 0
+
+
+def set_malloc_heap_size(size=STRINGS_UDF_HEAP_SIZE):


Why should this have a default, rather than reading the env var when it is lazily allocated?

I’m not convinced STRINGS_UDF_HEAP_SIZE needs to be defined.

The pool is very small if unset - something like 8mb I believe. Perhaps it'd be better if it said something like _STRINGS_UDF_INITIAL_HEAP_SIZE?

I think @bdice is just asking why there's a default function parameter at all if this is set by an environment variable. This sort of relates to the other question about why we call this lazily rather than at initialization. If it needs to be done lazily for a good reason, then perhaps a better alternative would be to have the default size=None and read the environment variable inside the function if size==None instead of storing it in a module-scope function. That shouldn't lead to repeated reads since after the first invocation with the current state of the code heap_size will always be equal to size so it won't go into the conditional.

The lazy initialization is kind of a placeholder in case we have an algorithm that can help guess the size based on the input size for example. I cannot speak to the env var access since I'm not familiar with the scoping here as it relates to udf kernel launches. Ideally if it could be set before each kernel launch? Though I understand the env var may not have changed.
Also, it may not be possible to make CUDA API calls during initialization (before the CUDA context is created) as we have found before with other CUDA calls.

+1 on what @davidwendt mentioned about initialization. For the logic discussion, what do you think of the changes in af7cd9f ?

bdice · 2022-11-16T19:30:55Z

python/cudf/cudf/core/udf/__init__.py

@@ -47,7 +47,14 @@
        utils.JIT_SUPPORTED_TYPES |= STRING_TYPES
        _supported_masked_types |= {string_view, udf_string}

-        utils.launch_arg_getters[cudf_str_dtype] = column_to_string_view_array
+        def column_to_string_view_array_init_heap(col):
+            # lazily allocate heap only when a string needs to be returned


The structure of this code makes it look like idempotency is its main goal, not laziness.

strings_udf.set_malloc_heap_size() should only ever be called once, I think?

Maybe this function or module in cudf (and not strings_udf) should own the “is it allocated already?” logic via a cache or global of some kind, and then the strings_udf logic doesn’t need to read env vars for the allocation size, etc.

This is not an allocation. There is currently a malloc-heap-size by default (set to about 8MB on the GPUs I've tested on). This updates the heap size to a new value so "already allocated" is not the goal.

Is set_malloc_heap_size intended to be a public API that can be called by a user, or is it an internal function? It's not entirely clear what the public API is here because this file doesn't define an __all__. If the function is just for internal use, why does it need to be called every time someone accesses a string column rather than just calling cudaDeviceSetLimit once when the package is imported (or after we can guarantee that the CUDA context/etc is initialized if that is the requirement that we're currently addressing by doing this "lazily")?

I think what we want to happen is:

The first time a user runs a string_udf, set the heap size to either default or what is retrieved from their environment.

The next time a user runs a string_udf, do not set the heap size. Or at least early return.

Reset the heap the next time the user runs a string_udf iff the user sets the value of STRING_UDF_HEAP_SIZE to something else during the session.
Does this seems reasonable?

Seems we have 2 comment chains asking about the same thing.
#12094 (comment)

The first time a user runs a string_udf, set the heap size to either default or what is retrieved from their environment.

Yes.

The next time a user runs a string_udf, do not set the heap size. Or at least early return.

The "set" function should not be called. Avoid designs that try to set but exit early.

Reset the heap the next time the user runs a string_udf iff the user sets the value of STRING_UDF_HEAP_SIZE to something else during the session.

No. The STRING_UDF_HEAP_SIZE should not be modifiable (I argued separately that that variable should not exist at all). If the user needs to be able to set a specific size, then the "set" function should be public and the user should be able to call it explicitly with immediate effect (not waiting until the next UDF execution, because memory demands can change in the meantime).

With the changes in af7cd9f I still have heap_size available as a module level attribute of strings_udf, what do you think about checking if heap_size == 0 within the getter to determine if the function need be called or not?

You could use a cached function for this instead of a global that tracks whether the function has already been run. Following ideas from https://mail.python.org/archives/list/[email protected]/thread/5OR3LJO7LOL6SC4OOGKFIVNNH4KADBPG/

from functools import lru_cache # Only executes the body of the function on the first call once = lru_cache(maxsize=None) @once def set_initial_malloc_heap_size(): strings_udf.set_malloc_heap_size() def column_to_string_view_array_init_heap(col): set_initial_malloc_heap_size() return column_to_string_view_array(col)

thanks this worked perfectly!

python/strings_udf/strings_udf/__init__.py

vyasr · 2022-11-16T19:51:25Z

I realized halfway through my review that my screen hadn't refreshed to include @bdice's latest comments, which largely overlapped with mine. I'll resolve my comments for now so that we just focus on his, I can post again if I have questions about the resolution.

python/strings_udf/strings_udf/__init__.py

bdice · 2022-11-16T21:32:17Z

python/strings_udf/strings_udf/__init__.py

+    """
+    global heap_size
+    if size == None:
+        size = os.environ.get(


Generally environment variables take effect at load time. I don't think we should allow/encourage the user to modify the environment of the process while it is running in order to change this.

I think you want to define it on import rather than in the body of this function:

_STRINGS_UDF_DEFAULT_HEAP_SIZE = os.environ.get("STRINGS_UDF_HEAP_SIZE", 2**31)

bdice · 2022-11-16T21:38:32Z

python/cudf/cudf/core/udf/__init__.py

@@ -47,7 +47,14 @@
        utils.JIT_SUPPORTED_TYPES |= STRING_TYPES
        _supported_masked_types |= {string_view, udf_string}

-        utils.launch_arg_getters[cudf_str_dtype] = column_to_string_view_array
+        def column_to_string_view_array_init_heap(col):
+            # lazily allocate heap only when a string needs to be returned


You could use a cached function for this instead of a global that tracks whether the function has already been run. Following ideas from https://mail.python.org/archives/list/[email protected]/thread/5OR3LJO7LOL6SC4OOGKFIVNNH4KADBPG/

from functools import lru_cache # Only executes the body of the function on the first call once = lru_cache(maxsize=None) @once def set_initial_malloc_heap_size(): strings_udf.set_malloc_heap_size() def column_to_string_view_array_init_heap(col): set_initial_malloc_heap_size() return column_to_string_view_array(col)

Co-authored-by: Bradley Dice <[email protected]>

brandon-b-miller · 2022-11-16T23:06:07Z

rerun tests

bdice

Looks like a function is defined twice. Otherwise LGTM.

python/cudf/cudf/core/udf/__init__.py

python/strings_udf/strings_udf/__init__.py

Co-authored-by: Bradley Dice <[email protected]>

brandon-b-miller · 2022-11-17T20:37:50Z

@gpucibot merge

add lazy mechanism for setting the heap size

22bfac9

brandon-b-miller added feature request New feature or request 2 - In Progress Currently a work in progress numba Numba issue Python Affects Python cuDF API. non-breaking Non-breaking change labels Nov 8, 2022

brandon-b-miller self-assigned this Nov 8, 2022

brandon-b-miller requested a review from a team as a code owner November 8, 2022 18:25

brandon-b-miller requested review from mroeschke and skirui-source November 8, 2022 18:25

davidwendt reviewed Nov 9, 2022

View reviewed changes

mroeschke reviewed Nov 10, 2022

View reviewed changes

python/strings_udf/strings_udf/__init__.py Outdated Show resolved Hide resolved

brandon-b-miller added 2 commits November 14, 2022 07:07

Merge branch 'branch-22.12' into fea-stringudf-heap

21369cc

introduce env var to control heap, cleanup

e7cc530

davidwendt reviewed Nov 14, 2022

View reviewed changes

python/strings_udf/strings_udf/__init__.py Outdated Show resolved Hide resolved

fix logic error

80459c5

davidwendt approved these changes Nov 16, 2022

View reviewed changes

bdice requested changes Nov 16, 2022

View reviewed changes

python/strings_udf/strings_udf/__init__.py Outdated Show resolved Hide resolved

python/strings_udf/strings_udf/__init__.py Outdated Show resolved Hide resolved

python/strings_udf/strings_udf/__init__.py Outdated Show resolved Hide resolved

brandon-b-miller and others added 4 commits November 16, 2022 06:40

Merge branch 'branch-22.12' into fea-stringudf-heap

ddab30c

Apply suggestions from code review

efff6c7

Co-authored-by: Bradley Dice <[email protected]>

Merge branch 'fea-stringudf-heap' of github.com:brandon-b-miller/cudf…

3fb8800

… into fea-stringudf-heap

style fix

8949fcb

brandon-b-miller requested a review from bdice November 16, 2022 18:08

bdice reviewed Nov 16, 2022

View reviewed changes

Merge branch 'branch-22.12' into fea-stringudf-heap

68cbaf7

vyasr reviewed Nov 16, 2022

View reviewed changes

python/strings_udf/strings_udf/__init__.py Outdated Show resolved Hide resolved

python/strings_udf/strings_udf/__init__.py Outdated Show resolved Hide resolved

adjust logic

af7cd9f

bdice reviewed Nov 16, 2022

View reviewed changes

brandon-b-miller and others added 4 commits November 16, 2022 15:39

Update python/strings_udf/strings_udf/__init__.py

889aba9

Co-authored-by: Bradley Dice <[email protected]>

define _STRINGS_UDF_DEFAULT_HEAP_SIZE at init

95e330b

use an lru cache

9bda9d6

Merge branch 'branch-22.12' into fea-stringudf-heap

ef99102

brandon-b-miller requested a review from bdice November 17, 2022 13:57

bdice approved these changes Nov 17, 2022

View reviewed changes

python/cudf/cudf/core/udf/__init__.py Outdated Show resolved Hide resolved

python/strings_udf/strings_udf/__init__.py Outdated Show resolved Hide resolved

brandon-b-miller and others added 2 commits November 17, 2022 09:00

Update python/strings_udf/strings_udf/__init__.py

5941048

Co-authored-by: Bradley Dice <[email protected]>

cleanup

b400143

rapids-bot bot merged commit 2f2685f into rapidsai:branch-22.12 Nov 17, 2022

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Allow setting malloc heap size in string udfs #12094

Allow setting malloc heap size in string udfs #12094

brandon-b-miller commented Nov 8, 2022

codecov bot commented Nov 8, 2022 •

edited

Loading

davidwendt Nov 9, 2022 •

edited

Loading

brandon-b-miller Nov 14, 2022

bdice Nov 16, 2022

bdice Nov 16, 2022

brandon-b-miller Nov 16, 2022

vyasr Nov 16, 2022

davidwendt Nov 16, 2022 •

edited

Loading

brandon-b-miller Nov 16, 2022

bdice Nov 16, 2022

davidwendt Nov 16, 2022

vyasr Nov 16, 2022

brandon-b-miller Nov 16, 2022 •

edited

Loading

davidwendt Nov 16, 2022

bdice Nov 16, 2022

brandon-b-miller Nov 16, 2022

bdice Nov 16, 2022 •

edited

Loading

brandon-b-miller Nov 16, 2022

vyasr commented Nov 16, 2022

bdice Nov 16, 2022 •

edited

Loading

bdice Nov 16, 2022 •

edited

Loading

brandon-b-miller commented Nov 16, 2022

bdice left a comment

brandon-b-miller commented Nov 17, 2022

		@@ -87,6 +88,26 @@ def _get_ptx_file():
		return regular_result[1]


		default_heap_size = int(2e6)

		heap_size = 0


		def set_malloc_heap_size(size=STRINGS_UDF_HEAP_SIZE):

Allow setting malloc heap size in string udfs #12094

Allow setting malloc heap size in string udfs #12094

Conversation

brandon-b-miller commented Nov 8, 2022

codecov bot commented Nov 8, 2022 • edited Loading

Codecov Report

davidwendt Nov 9, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

davidwendt Nov 16, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

brandon-b-miller Nov 16, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

bdice Nov 16, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

vyasr commented Nov 16, 2022

bdice Nov 16, 2022 • edited Loading

Choose a reason for hiding this comment

bdice Nov 16, 2022 • edited Loading

Choose a reason for hiding this comment

brandon-b-miller commented Nov 16, 2022

bdice left a comment

Choose a reason for hiding this comment

brandon-b-miller commented Nov 17, 2022

codecov bot commented Nov 8, 2022 •

edited

Loading

davidwendt Nov 9, 2022 •

edited

Loading

davidwendt Nov 16, 2022 •

edited

Loading

brandon-b-miller Nov 16, 2022 •

edited

Loading

bdice Nov 16, 2022 •

edited

Loading

bdice Nov 16, 2022 •

edited

Loading

bdice Nov 16, 2022 •

edited

Loading