From 7b5e7c2ab8ae72a0f8c17eabe86771ac363f285d Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Mon, 7 Sep 2026 11:19:21 -0700 Subject: [PATCH 1/3] FEAT: Cache numeric Series apply kernels --- python/cudf/cudf/core/udf/scalar_function.py | 97 +++++++++++++++++++- 1 file changed, 96 insertions(+), 1 deletion(-) diff --git a/python/cudf/cudf/core/udf/scalar_function.py b/python/cudf/cudf/core/udf/scalar_function.py index 8bf39e768025..62f96f32e061 100644 --- a/python/cudf/cudf/core/udf/scalar_function.py +++ b/python/cudf/cudf/core/udf/scalar_function.py @@ -1,6 +1,8 @@ -# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +import warnings +from contextlib import nullcontext from functools import cache from numba import cuda @@ -9,6 +11,7 @@ from cudf.core.dtype.validators import is_dtype_obj_string from cudf.core.udf.api import Masked, pack_return from cudf.core.udf.masked_typing import MaskedType +from cudf.core.udf.nrt_utils import CaptureNRTUsage, nrt_enabled from cudf.core.udf.strings_typing import string_view from cudf.core.udf.templates import ( masked_input_initializer_template, @@ -17,10 +20,74 @@ ) from cudf.core.udf.udf_kernel_base import ApplyKernelBase from cudf.core.udf.utils import ( + DEPRECATED_SM_REGEX, _mask_get, ) +def _compile_unmasked_series_apply_kernel(f_, sig, nrt): + """Compile a source-backed, disk-cacheable Series.apply kernel.""" + ctx = nrt_enabled() if nrt else nullcontext() + with ctx: + with warnings.catch_warnings(): + warnings.simplefilter("default") + warnings.filterwarnings( + "ignore", + message=DEPRECATED_SM_REGEX, + category=UserWarning, + module=r"^numba\.cuda(\.|$)", + ) + + @cuda.jit( + sig, + cache=True, + ) + def _kernel(retval, size, input_col_0, offset_0): + i = cuda.grid(1) + ret_data_arr, ret_mask_arr = retval + + if i < size: + masked_0 = Masked(input_col_0[i + offset_0], True) + ret_masked = pack_return(f_(masked_0)) + ret_data_arr[i] = ret_masked.value + ret_mask_arr[i] = ret_masked.valid + + return _kernel + + +def _compile_masked_series_apply_kernel(f_, sig, nrt): + """Compile a source-backed, disk-cacheable Series.apply kernel.""" + ctx = nrt_enabled() if nrt else nullcontext() + with ctx: + with warnings.catch_warnings(): + warnings.simplefilter("default") + warnings.filterwarnings( + "ignore", + message=DEPRECATED_SM_REGEX, + category=UserWarning, + module=r"^numba\.cuda(\.|$)", + ) + + @cuda.jit( + sig, + cache=True, + ) + def _kernel(retval, size, input_col_0, offset_0): + i = cuda.grid(1) + ret_data_arr, ret_mask_arr = retval + + if i < size: + d_0, m_0 = input_col_0 + masked_0 = Masked( + d_0[i + offset_0], _mask_get(m_0, i + offset_0) + ) + ret_masked = pack_return(f_(masked_0)) + ret_data_arr[i] = ret_masked.value + ret_mask_arr[i] = ret_masked.valid + + return _kernel + + class SeriesApplyKernel(ApplyKernelBase): """ Class representing a kernel that computes the result of @@ -58,6 +125,34 @@ def _get_kernel_string(self): extra_args=extra_args, masked_initializer=masked_initializer ) + def compile_kernel(self): + # The static no-argument kernels below have a real source location, so + # Numba can persist their compiled specializations. Other signatures + # keep the existing generated-kernel path for this prototype. + if self.args or is_dtype_obj_string(self.frame.dtype): + return super().compile_kernel() + + capture_nrt_usage = CaptureNRTUsage() + with capture_nrt_usage: + return_type = self._get_udf_return_type() + + # String allocation uses symbols from UDF_SHIM_FILE, which Numba's + # on-disk cache cannot serialize. Keep the existing linked path for + # those specializations. + if capture_nrt_usage.use_nrt: + return super().compile_kernel() + + self.sig = self._construct_signature(return_type) + kernel_factory = ( + _compile_masked_series_apply_kernel + if self.frame._column.mask is not None + else _compile_unmasked_series_apply_kernel + ) + kernel = kernel_factory( + self.device_func, self.sig, capture_nrt_usage.use_nrt + ) + return kernel, return_type + @cache def _get_kernel_string_exec_context(self): # This is the global execution context that will be used From 10a1e861868995949d1c387c6c01bed50c976b64 Mon Sep 17 00:00:00 2001 From: Vyas Ramasubramani Date: Mon, 7 Sep 2026 11:42:08 -0700 Subject: [PATCH 2/3] FEAT: Cache numeric dataframe apply kernels --- python/cudf/cudf/core/udf/groupby_utils.py | 5 ++ python/cudf/cudf/core/udf/udf_kernel_base.py | 73 ++++++++++++++++++-- 2 files changed, 72 insertions(+), 6 deletions(-) diff --git a/python/cudf/cudf/core/udf/groupby_utils.py b/python/cudf/cudf/core/udf/groupby_utils.py index aa4e498a270a..0fc74255d42d 100644 --- a/python/cudf/cudf/core/udf/groupby_utils.py +++ b/python/cudf/cudf/core/udf/groupby_utils.py @@ -197,6 +197,11 @@ def f(group): def kernel_type(self): return "groupby_apply" + @property + def _requires_linked_udf_shim(self): + # GroupBy UDF reductions call the device functions in UDF_SHIM_FILE. + return True + def _get_frame_type(self): return _get_frame_groupby_type( np.dtype(list(_all_dtypes_from_frame(self.frame).items())), diff --git a/python/cudf/cudf/core/udf/udf_kernel_base.py b/python/cudf/cudf/core/udf/udf_kernel_base.py index 8b54f2678717..579027f6a15d 100644 --- a/python/cudf/cudf/core/udf/udf_kernel_base.py +++ b/python/cudf/cudf/core/udf/udf_kernel_base.py @@ -1,9 +1,12 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 import warnings from abc import ABC, abstractmethod from contextlib import nullcontext +from hashlib import sha256 +from pickle import dumps +from textwrap import indent import numpy as np from numba import cuda, typeof @@ -12,6 +15,7 @@ from numba.types import CPointer, Poison, Tuple, boolean, int64, void from cudf.api.types import is_scalar +from cudf.core.dtype.validators import is_dtype_obj_string from cudf.core.udf.masked_typing import MaskedType from cudf.core.udf.nrt_utils import CaptureNRTUsage, nrt_enabled from cudf.core.udf.strings_typing import str_view_arg_handler @@ -75,6 +79,11 @@ def _get_kernel_string_exec_context(self): string. """ + @property + def _requires_linked_udf_shim(self): + """Whether this API has device-library dependencies beyond NRT.""" + return False + @staticmethod def _format_arg_list(prefix, count): """Build a comma-separated parameter list like 'prefix_0, prefix_1, ...'.""" @@ -156,12 +165,65 @@ def compile_kernel(self): return kernel, return_type + def _disk_cache_identity(self): + """Return a stable identity for globals used by a generated kernel.""" + key = _generate_cache_key( + self.frame, self.func, self.args, suffix=self.kernel_type + ) + return int.from_bytes(sha256(dumps(key)).digest()[:8], "big") + + def _can_use_disk_cache(self, nrt): + return ( + not self._requires_linked_udf_shim + and not nrt + and not any( + is_dtype_obj_string(col.dtype) for col in self.frame._columns + ) + ) + def compile_kernel_string(self, kernel_string, nrt=False): global_exec_context = self._get_kernel_string_exec_context() - global_exec_context["f_"] = self.device_func + use_disk_cache = self._can_use_disk_cache(nrt) + if use_disk_cache: + # The closure carries the UDF and generated-kernel globals into + # Numba's cache key. Compiling with this module's filename gives + # the dispatcher a persistent source locator. + kernel_lines = indent( + kernel_string, " ", lambda _line: True + ).splitlines() + kernel_def_idx = next( + i + for i, line in enumerate(kernel_lines) + if line.startswith(" def _kernel") + ) + kernel_lines[kernel_def_idx + 1 : kernel_def_idx + 1] = [ + " if cache_identity < 0:", + " return", + ] + factory_string = "\n".join( + [ + "def _make_kernel(f_, cache_identity):", + *kernel_lines, + " return _kernel", + ] + ) + exec( + compile(factory_string, __file__, "exec"), global_exec_context + ) + _kernel = global_exec_context["_make_kernel"]( + self.device_func, self._disk_cache_identity() + ) + else: + global_exec_context["f_"] = self.device_func + exec(kernel_string, global_exec_context) + _kernel = global_exec_context["_kernel"] - exec(kernel_string, global_exec_context) - _kernel = global_exec_context["_kernel"] + jit_kwargs = {} + if use_disk_cache: + jit_kwargs["cache"] = True + else: + jit_kwargs["link"] = [UDF_SHIM_FILE] + jit_kwargs["extensions"] = [str_view_arg_handler] ctx = nrt_enabled() if nrt else nullcontext() with ctx: with warnings.catch_warnings(): @@ -174,8 +236,7 @@ def compile_kernel_string(self, kernel_string, nrt=False): ) kernel = cuda.jit( self.sig, - link=[UDF_SHIM_FILE], - extensions=[str_view_arg_handler], + **jit_kwargs, )(_kernel) return kernel From f784670a93b66cdee8eb5d545e41f34fa3dea72d Mon Sep 17 00:00:00 2001 From: anon Date: Mon, 7 Sep 2026 18:54:54 +0000 Subject: [PATCH 3/3] CI: Persist Numba CUDA test cache --- .github/workflows/pr.yaml | 8 ++++++-- .github/workflows/test.yaml | 8 ++++++-- 2 files changed, 12 insertions(+), 4 deletions(-) diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml index 6abfeaad8dc0..b6f8a6864b9b 100644 --- a/.github/workflows/pr.yaml +++ b/.github/workflows/pr.yaml @@ -352,11 +352,13 @@ jobs: packages: read pull-requests: read secrets: inherit # zizmor: ignore[secrets-inherit] - uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@main + # TODO: Return to @main after rapidsai/shared-workflows#629 merges. + uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@e7287514746a7f81aa65e99a4a5783f9882bcc52 if: fromJSON(needs.changed-files.outputs.changed_file_groups).cpp_build_inputs || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_all || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_cudf || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_conda_builds || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_conda_common || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_conda_cudf with: build_type: pull-request script: "ci/test_python_cudf.sh" + numba-cache-key-prefix: cudf-numba-udf-v1 # https://github.com/NVIDIA/cudf/issues/23498 matrix_filter: map(select(.GPU != "gb300" and .GPU != "gh200")) conda-python-other-tests: @@ -549,7 +551,8 @@ jobs: packages: read pull-requests: read secrets: inherit # zizmor: ignore[secrets-inherit] - uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@main + # TODO: Return to @main after rapidsai/shared-workflows#629 merges. + uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@e7287514746a7f81aa65e99a4a5783f9882bcc52 if: fromJSON(needs.changed-files.outputs.changed_file_groups).cpp_build_inputs || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_all || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_cudf || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_cudf_streaming || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_wheels || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_wheel_cudf with: build_type: pull-request @@ -559,6 +562,7 @@ jobs: RUN_CUDF_TESTS=${{ fromJSON(needs.changed-files.outputs.changed_file_groups).cpp_build_inputs || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_all || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_cudf || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_wheels || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_wheel_cudf }} RUN_CUDF_STREAMING_TESTS=${{ fromJSON(needs.changed-files.outputs.changed_file_groups).cpp_build_inputs || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_all || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_cudf_streaming || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_wheels || fromJSON(needs.changed-files.outputs.changed_file_groups).test_python_wheel_cudf }} ci/test_wheel_cudf.sh + numba-cache-key-prefix: cudf-numba-udf-v1 # https://github.com/NVIDIA/cudf/issues/23498 matrix_filter: map(select(.GPU != "gb300" and .GPU != "gh200")) wheel-tests-cudf-polars: diff --git a/.github/workflows/test.yaml b/.github/workflows/test.yaml index 7c73aad711c1..6613a377045f 100644 --- a/.github/workflows/test.yaml +++ b/.github/workflows/test.yaml @@ -116,13 +116,15 @@ jobs: packages: read pull-requests: read secrets: inherit # zizmor: ignore[secrets-inherit] - uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@main + # TODO: Return to @main after rapidsai/shared-workflows#629 merges. + uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@e7287514746a7f81aa65e99a4a5783f9882bcc52 with: build_type: ${{ inputs.build_type }} branch: ${{ inputs.branch }} date: ${{ inputs.date }} sha: ${{ inputs.sha }} script: "ci/test_python_cudf.sh" + numba-cache-key-prefix: cudf-numba-udf-v1 # https://github.com/NVIDIA/cudf/issues/23498 matrix_filter: map(select(.GPU != "gb300" and .GPU != "gh200")) conda-python-other-tests: @@ -200,13 +202,15 @@ jobs: packages: read pull-requests: read secrets: inherit # zizmor: ignore[secrets-inherit] - uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@main + # TODO: Return to @main after rapidsai/shared-workflows#629 merges. + uses: rapidsai/shared-workflows/.github/workflows/wheels-test.yaml@e7287514746a7f81aa65e99a4a5783f9882bcc52 with: build_type: ${{ inputs.build_type }} branch: ${{ inputs.branch }} date: ${{ inputs.date }} sha: ${{ inputs.sha }} script: ci/test_wheel_cudf.sh + numba-cache-key-prefix: cudf-numba-udf-v1 wheel-tests-dask-cudf: permissions: actions: read