[BUG] apply_grouped broken in 23.04

[meprem]

Describe the bug
apply_grouped is broken in 23.04

import numpy as np

import cudf
from cudf.datasets import randomdata
from numba import cuda

df = randomdata(
    nrows=10, dtypes={"a": float, "b": bool, "c": str, "e": float}, seed=12
)

def rolling_avg(e, rolling_avg_e):
    win_size = 3
    for i in range(cuda.threadIdx.x, len(e), cuda.blockDim.x):
        if i < win_size - 1:
            # If there is not enough data to fill the window,
            # take the average to be NaN
            rolling_avg_e[i] = np.nan
        else:
            total = 0
            for j in range(i - win_size + 1, i + 1):
                total += e[j]
            rolling_avg_e[i] = total / win_size

grouped = df.groupby(["b"])

results = grouped.apply_grouped(
    rolling_avg, incols=["e"], outcols=dict(rolling_avg_e=np.float64)
)
results

Steps/Code to reproduce bug
Follow this guide http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports to craft a minimal bug report. This helps us reproduce the issue you're having and resolve the issue more quickly.

I can reproduce the bug with code here: https://docs.rapids.ai/api/cudf/stable/user_guide/guide-to-udfs/

I am attaching conda env file and notebook to reproduce the bug.

apply_grouped_bug.zip

Expected behavior
A clear and concise description of what you expected to happen.

The code from the documentation should run: https://docs.rapids.ai/api/cudf/stable/user_guide/guide-to-udfs/

Please note I use similar code in production that used to work in the previous version.

Environment overview (please complete the following information)

• Environment location: [Bare-metal, Docker, Cloud(specify cloud provider)]
• Method of cuDF install: [conda, Docker, or from source]
  • If method of install is [Docker], provide docker pull & docker run commands used

I am attaching conda env file.

Environment details
Please run and paste the output of the cudf/print_env.sh script here, to gather any other relevant environment details

I am using p3dn.24xlarge AWS machine.

Additional context
Add any other context about the problem here.

[quasiben]

Instead of a zip file can you instead inline the code and data in the comments or host on gist.github.com ?

[meprem]

updated

[shwina]

Thank you for reporting, @meprem - could you please also provide the error and stack trace you see when you run the provided code snippet? When I try to run that code, it works, so I'm not sure where to start debugging.

[meprem]

Env file:

name: rapidsai_ploomber
channels:
  - rapidsai
  - nvidia
  - numba
  - conda-forge
dependencies:
  - rapids=23.04
  - cudatoolkit=11.8
  - nvtabular=23.04.00
  - python=3.8
  - pip
prefix: /home/ubuntu/anaconda3/envs/rapidsai_ploomber

Stacktrace:

Traceback (most recent call last):
  File "1.py", line 26, in <module>
    results = grouped.apply_grouped(
  File "/opt/conda/envs/rapidsai_ploomber/lib/python3.8/site-packages/cudf/core/groupby/groupby.py", line 1458, in apply_grouped
    return grouped_values.apply_chunks(function, **kwargs)
  File "/opt/conda/envs/rapidsai_ploomber/lib/python3.8/site-packages/nvtx/nvtx.py", line 101, in inner
    result = func(*args, **kwargs)
  File "/opt/conda/envs/rapidsai_ploomber/lib/python3.8/site-packages/cudf/core/dataframe.py", line 4582, in apply_chunks
    return applyutils.apply_chunks(
  File "/opt/conda/envs/rapidsai_ploomber/lib/python3.8/site-packages/cudf/utils/applyutils.py", line 105, in apply_chunks
    return applychunks.run(df, chunks=chunks, tpb=tpb)
  File "/opt/conda/envs/rapidsai_ploomber/lib/python3.8/contextlib.py", line 75, in inner
    return func(*args, **kwds)
  File "/opt/conda/envs/rapidsai_ploomber/lib/python3.8/site-packages/cudf/utils/applyutils.py", line 169, in run
    self.launch_kernel(df, bound.args, **launch_params)
  File "/opt/conda/envs/rapidsai_ploomber/lib/python3.8/site-packages/cudf/utils/applyutils.py", line 212, in launch_kernel
    self.kernel.forall(len(df))(len(df), chunks, *args)
  File "/opt/conda/envs/rapidsai_ploomber/lib/python3.8/site-packages/numba/cuda/dispatcher.py", line 489, in __call__
    specialized = self.dispatcher.specialize(*args)
  File "/opt/conda/envs/rapidsai_ploomber/lib/python3.8/site-packages/numba/cuda/dispatcher.py", line 718, in specialize
    specialization.compile(argtypes)
  File "/opt/conda/envs/rapidsai_ploomber/lib/python3.8/site-packages/numba/cuda/dispatcher.py", line 929, in compile
    kernel.bind()
  File "/opt/conda/envs/rapidsai_ploomber/lib/python3.8/site-packages/numba/cuda/dispatcher.py", line 208, in bind
    self._codelibrary.get_cufunc()
  File "/opt/conda/envs/rapidsai_ploomber/lib/python3.8/site-packages/numba/cuda/codegen.py", line 180, in get_cufunc
    module = ctx.create_module_image(cubin)
  File "/opt/conda/envs/rapidsai_ploomber/lib/python3.8/site-packages/numba/cuda/cudadrv/driver.py", line 1450, in create_module_image
    module = load_module_image(self, image)
  File "/opt/conda/envs/rapidsai_ploomber/lib/python3.8/site-packages/numba/cuda/cudadrv/driver.py", line 1551, in load_module_image
    return load_module_image_ctypes(context, image)
  File "/opt/conda/envs/rapidsai_ploomber/lib/python3.8/site-packages/numba/cuda/cudadrv/driver.py", line 1573, in load_module_image_ctypes
    driver.cuModuleLoadDataEx(byref(handle), image, len(options),
  File "/opt/conda/envs/rapidsai_ploomber/lib/python3.8/site-packages/numba/cuda/cudadrv/driver.py", line 339, in safe_cuda_api_call
    retcode = libfn(*args)
ctypes.ArgumentError: argument 2: <class 'TypeError'>: wrong type

numba is causing problem but need it for nvtabular. thanks.

[brandon-b-miller]

What version of numba is being used here?

[gmarkall]

It is probably 0.57.0 - the last line of the traceback:

  File "/opt/conda/envs/rapidsai_ploomber/lib/python3.8/site-packages/numba/cuda/cudadrv/driver.py", line 339, in safe_cuda_api_call
    retcode = libfn(*args)

matches that line in 0.57.0:

https://github.com/numba/numba/blob/4fd4e39c672d119b54a2276d170f270764d2bce7/numba/cuda/cudadrv/driver.py#L339

However with 0.56.0 line 339 is something else:

https://github.com/numba/numba/blob/288a38bbd5a15418a211bf067878dfdf3c139509/numba/cuda/cudadrv/driver.py#L339

@meprem Can you try switching to Numba 0.56.4 please?

[vyasr]

I'm guessing that Graham's diagnosis above was correct. With the latest version of cudf, which supports newer versions of numba, I cannot reproduce this issue:

In [22]: import numpy as np
    ...: 
    ...: import cudf
    ...: from cudf.datasets import randomdata
    ...: from numba import cuda
    ...: 
    ...: df = randomdata(
    ...:     nrows=10, dtypes={"a": float, "b": bool, "c": str, "e": float}, seed=12
    ...: )
    ...: 
    ...: def rolling_avg(e, rolling_avg_e):
    ...:     win_size = 3
    ...:     for i in range(cuda.threadIdx.x, len(e), cuda.blockDim.x):
    ...:         if i < win_size - 1:
    ...:             # If there is not enough data to fill the window,
    ...:             # take the average to be NaN
    ...:             rolling_avg_e[i] = np.nan
    ...:         else:
    ...:             total = 0
    ...:             for j in range(i - win_size + 1, i + 1):
    ...:                 total += e[j]
    ...:             rolling_avg_e[i] = total / win_size
    ...: 
    ...: grouped = df.groupby(["b"])
    ...: 
    ...: results = grouped.apply_grouped(
    ...:     rolling_avg, incols=["e"], outcols=dict(rolling_avg_e=np.float64)
    ...: )
    ...: results
Out[22]: 
          a      b       c         e  rolling_avg_e
1  0.480099  False     Bob -0.729580            NaN
4 -0.970850  False   Sarah  0.342905            NaN
6  0.801430  False   Sarah  0.632337       0.081887
7 -0.933157  False   Quinn -0.420826       0.184805
0 -0.691674   True     Dan -0.958380            NaN
2 -0.473370   True  Xavier -0.767454            NaN
3  0.067479   True   Alice -0.380205      -0.702013
5  0.837494   True   Wendy -0.057540      -0.401733
8  0.913899   True  Ursula  0.466252       0.009502
9 -0.725581   True  George  0.405245       0.271319