JIT compilation in cuDF on Snowflake Notebooks is broken

Snowflake is still on 26.2 but I just discovered that the udf path that triggers numba is broken. (This was supposed to be one of the tests they run before pushing but it seems that is not being tested)

It seems to need nvjitlink to be updated, but it is unclear why this is happening.

This is the reproducer

import cudf

s = cudf.Series(['a', 'aa', 'b'])
s.apply(lambda x: len(x))

ValueError: user defined function compilation failed.
---------------------------------------------------------------------------
nvJitLinkError                            Traceback (most recent call last)
File /opt/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/site-packages/cudf/core/indexed_frame.py:3586, in IndexedFrame._apply(self, func, kernel_class, *args, **kwargs)
   3585     kr = kernel_class(self, func, args)
-> 3586     kernel, retty = kr.get_kernel()
   3587 except Exception as e:

File /opt/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/site-packages/cudf/core/udf/udf_kernel_base.py:178, in ApplyKernelBase.get_kernel(self)
    177 def get_kernel(self):
--> 178     return self._compile_or_get_kernel()

File /opt/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/site-packages/cudf/core/udf/udf_kernel_base.py:195, in ApplyKernelBase._compile_or_get_kernel(self)
    193     return kernel, masked_or_scalar
--> 195 kernel, scalar_return_type = self.compile_kernel()
    197 np_return_type = (
    198     numpy_support.as_dtype(scalar_return_type)
    199     if scalar_return_type.is_internal
    200     else scalar_return_type.np_dtype
    201 )

File /opt/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/site-packages/cudf/core/udf/udf_kernel_base.py:148, in ApplyKernelBase.compile_kernel(self)
    147 kernel_string = self._get_kernel_string()
--> 148 kernel = self.compile_kernel_string(
    149     kernel_string, nrt=capture_nrt_usage.use_nrt
    150 )
    152 return kernel, return_type

File /opt/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/site-packages/cudf/core/udf/udf_kernel_base.py:170, in ApplyKernelBase.compile_kernel_string(self, kernel_string, nrt)
    164         warnings.filterwarnings(
    165             "ignore",
    166             message=DEPRECATED_SM_REGEX,
    167             category=UserWarning,
    168             module=r"^numba\.cuda(\.|$)",
    169         )
--> 170         kernel = cuda.jit(
    171             self.sig,
    172             link=[UDF_SHIM_FILE],
    173             extensions=[str_view_arg_handler],
    174         )(_kernel)
    175 return kernel

File /opt/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/site-packages/numba_cuda/numba/cuda/decorators.py:207, in jit.<locals>._jit(func)
    206     else:
--> 207         disp.compile(argtypes)
    209 disp._specialized = specialized

File /opt/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/site-packages/numba_cuda/numba/cuda/core/compiler_lock.py:74, in _DualCompilerLock.__call__.<locals>._acquire_compile_lock(*args, **kwargs)
     73 with self:
---> 74     return func(*args, **kwargs)

File /opt/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/site-packages/numba_cuda/numba/cuda/dispatcher.py:1892, in CUDADispatcher.compile(self, sig)
   1891 # We call bind to force codegen, so that there is a cubin to cache
-> 1892 kernel.bind()
   1893 self._cache.save_overload(sig, kernel)

File /opt/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/site-packages/numba_cuda/numba/cuda/dispatcher.py:337, in _Kernel.bind(self)
    334 """
    335 Force binding to current CUDA context
    336 """
--> 337 cufunc = self._codelibrary.get_cufunc()
    339 self.initialize_once(cufunc.module)

File /opt/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/site-packages/numba_cuda/numba/cuda/codegen.py:336, in CUDACodeLibrary.get_cufunc(self)
    335     return cufunc
--> 336 cubin = self.get_cubin(cc=device.compute_capability)
    337 module = ctx.create_module_image(
    338     cubin, self._setup_functions, self._teardown_functions
    339 )

File /opt/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/site-packages/numba_cuda/numba/cuda/codegen.py:315, in CUDACodeLibrary.get_cubin(self, cc)
    314 self._link_all(linker, cc, ignore_nonlto=False)
--> 315 cubin = linker.complete()
    317 self._cubin_cache[cc] = cubin

File /opt/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/site-packages/numba_cuda/numba/cuda/cudadrv/driver.py:2776, in _Linker.complete(self)
   2775 def complete(self):
-> 2776     self.linker = Linker(*self._object_codes, options=self.options)
   2777     result = self.linker.link("cubin")

File cuda/core/_linker.pyx:71, in cuda.core._linker.Linker.__init__()
---> 71 'Could not get source, probably due dynamically evaluated source code.'

File cuda/core/_linker.pyx:494, in cuda.core._linker.Linker_init()
--> 494 'Could not get source, probably due dynamically evaluated source code.'

File cuda/core/_linker.pyx:525, in cuda.core._linker.Linker_add_code_object()
--> 525 'Could not get source, probably due dynamically evaluated source code.'

File cuda/core/_utils/cuda_utils.pyx:126, in cuda.core._utils.cuda_utils.HANDLE_RETURN_NVJITLINK()
--> 126 'Could not get source, probably due dynamically evaluated source code.'

File cuda/core/_utils/cuda_utils.pyx:145, in cuda.core._utils.cuda_utils._raise_nvjitlink_error()
--> 145 'Could not get source, probably due dynamically evaluated source code.'

nvJitLinkError: ERROR_INTERNAL (6)
nvJitLink error log: ERROR 4 in nvvmAddNVVMContainerToProgram, may need newer version of nvJitLink library


The above exception was the direct cause of the following exception:

ValueError                                Traceback (most recent call last)
Cell In[2], line 2
      1 s = cudf.Series(['a', 'aa', 'b'])
----> 2 s.apply(lambda x: len(x))

File /opt/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/site-packages/cudf/core/series.py:2608, in Series.apply(self, func, convert_dtype, args, by_row, **kwargs)
   2605 elif by_row != "compat":
   2606     raise NotImplementedError("by_row is currently not supported.")
-> 2608 result = self._apply(func, SeriesApplyKernel, *args, **kwargs)
   2609 result.name = self.name
   2610 return result

File /opt/python/cpython-3.12.12-linux-x86_64-gnu/lib/python3.12/site-packages/cudf/core/indexed_frame.py:3588, in IndexedFrame._apply(self, func, kernel_class, *args, **kwargs)
   3586     kernel, retty = kr.get_kernel()
   3587 except Exception as e:
-> 3588     raise ValueError(
   3589         "user defined function compilation failed."
   3590     ) from e
   3592 # Mask and data column preallocated
   3593 ans_col = _return_arr_from_dtype(retty, len(self))

ValueError: user defined function compilation failed.

Start with the provided Snowflake Notebook reproducer and trace the failure from cudf.Series.apply through IndexedFrame._apply and the numba CUDA compilation path. Check the nvJitLink version and dependency environment against the ERROR_INTERNAL traceback; done means the reproducer completes successfully in Snowflake Notebooks.