Repository navigation
njit(parallel = True) gives wrong output #9379
Description
njit with parallel = True gives the wrong output.
The function _groupby_apply_nb_parallel_dynamic_size partially replicates a dataframe.groupby().apply(func) using numba for speedup. Full minimal example code available below.
The function _groupby_apply_nb_parallel_dynamic_size returns wrong output when parallel is set to true.
I have also noticed that when parallel = True, the output of _groupby_apply_nb_parallel_dynamic_size (in the minimal example below) returns os.cpu_count() unique values, when it should return 1001 unique values. The function works correctly when setting Parallel = False. I have also managed to get the correct output by doing some modifications, such as:
out_cols_size = 1 # Setting this to 1 statically also seems to solve the issue
Even though I don't know why this should fix the issue.
I can't see any data race conditions, thus why I think it is a Numba bug.
OS: Ubuntu 22.04
Python 3.11.7
Numba version: '0.58.1', installed through pip
Numpy '1.26.2'
Pandas : '2.1.4' (it's only used to generate the expected array, nothing else). The version doesn't matter
The below code has similar behavior as pandas.DataFrame.groupby
import os
import numpy as np
import numba
from numba import njit, prange
from numba.typed import List as NumbaList
import pandas as pd
PARALLEL = True # Set it to false so it works with no issues
@njit(cache=False)
def get_max_first_array_of_list_of_arrays(y_x: NumbaList, return_dtype=np.float64):
"""
This function only uses the first element of y_x as a minimal example
"""
y = y_x[0]
ans = np.array([y.max()], dtype=return_dtype)
return ans
@njit(parallel=PARALLEL, cache=False)
def _groupby_apply_nb_parallel_dynamic_size(
start_indices: np.array,
end_indices: np.array,
# dataargs is a list of numpy 2d arrays
dataargs: numba.typed.List,
func: callable,
funcargs,
return_dtype=np.float64,
):
"""
start_indices : 1d np.array with the start indices of each group
end_indices : 1d np.array with the end indices of each group
dataargs: List of np 2d arrays. These are arguments of func
func: jitted function
funcargs: other arguments to func.
return_dtype : output array dtype
"""
ngroups = len(start_indices)
if ngroups == 0:
raise Exception("empty number of grouped, (start_indices)")
# First call to determinize the size of the output. Assumes the first group will have output size
# equal to the output size of all other groups
res = func(
[data[start_indices[0] : end_indices[0], ...] for data in dataargs], *funcargs
)
out_cols_size = res.shape[0] # Setting this to 1 statically also seems to solve the issue
result = np.empty(shape=(ngroups, out_cols_size), dtype=return_dtype)
for k in prange(ngroups):
start_idx = start_indices[k]
end_idx = end_indices[k]
res = func([data[start_idx:end_idx, ...] for data in dataargs], *funcargs)
result[k, :] = res
return result
if __name__ == "__main__":
np.random.seed(42)
size = 100_000
qty_intervals = 1_000
save_dir = os.path.join("random_experiments", "parallel_py")
samples = np.sort(np.random.choice(range(size), qty_intervals, replace=False))
start_indices = np.append(0, samples)
end_indices = np.append(samples, size)
x = np.zeros(shape = (size)) # This could be any np.array
y = np.exp(np.cumsum(np.random.normal(0, 1e-2, size)))
os.makedirs(save_dir, exist_ok=True)
for i in range(10):
max_vals = _groupby_apply_nb_parallel_dynamic_size(
start_indices = start_indices,
end_indices = end_indices,
dataargs = NumbaList([y, x]),
func = get_max_first_array_of_list_of_arrays,
funcargs=(),
return_dtype=np.float64,
)
save_path = os.path.join(save_dir, f'output_{i}_.npy')
np.save(save_path, max_vals)
groups = np.zeros(size, dtype=np.int64)
groups[start_indices] = 1
groups = np.cumsum(groups)
max_expected = pd.DataFrame(y).groupby(by = groups).max().values
for i in range(10):
save_path = os.path.join(save_dir, f'output_{i}_.npy')
max_vals = np.load(save_path)
equal = np.isclose(max_expected, max_vals, equal_nan = True, rtol = 1e-6).all()
if not equal:
print(f"ERROR - parallel output {i} is wrong")
print(f"Number of unique values : {len(np.unique(max_vals))}")
print(f"Number of unique expected : {len(np.unique(max_expected))}")
Activity
gmarkall commentedon Dec 27, 2023
Many thanks for the bug report and reproducer - I can reproduce the issue.
I can't see any obvious issue in your code, but I'm not an expert in the parallel functionality so I can't tell if there's also anything that's not expected to work when
parallel=True- I'll bring this up in the next triage meeting, but note that it's not until Tuesday January 16th.netomenoci commentedon Dec 27, 2023
Thank you so much @gmarkall
I think I have found the reason behind the issue and managed to make it work, even though I think there is a serious problem that needs to be fixed with the hoisting step.
By calling _groupby_apply_nb_parallel_dynamic_size.parallel_diagnostics(level=4), I was able to check that everything seemed to be fine with the serialization step, but it seems there is something wrong with the hoisting step.
This is what it looks like:
(Copying the relevant part of the code below)
It seems that it has a lot of things being hoisted, in particular:
closure__locals___listcomp__v6__v6build_list_0 = build_list(items=[])Which looks like the list comprehension I have in[data[start_idx:end_idx, ...]If it indeed is, then that should be the issue, as this obviously cannot be hoisted outside of the loop since the indices are defined by the prange variable.
I went further and tried to make the dependence even more explicit to help the compiler. I have then replaced
by the below one liner
This is what the hoisting diagnosis looks like after this change:
Ie, there a single difference: the line below disappeared from the hoisting.
closure__locals___listcomp__v6__v6build_list_0 = build_list(items=[])After that, the code works both with PARALLEL = True and PARALLEL = False
I am now afraid of using the parallel optimization from Numba, as it seems it cannot be trusted to perform the optimizations. I will be happy to hear back from you guys on what could be the cause of the issue, so I can feel safer to continue using it after understand in detail what might be going wrong. I would also be happy to be part of the PR to solve the issue.
Best regards,
Neto
netomenoci commentedon Jan 3, 2024
I have found another similar issue, so I will just add it here rather than creating a separate new one.
Replacing
with
ie, using Ellipsis instead of ":" , results in the same issue of Parallel giving wrong results (not only that, but also a number of unique results equal to the number of cpus in the machine).
It is not possible to reproduce it with the above example (sorry). I was unable to create a minimum example for the Ellipsis issue, as it only happens when I run multiple tests at the same time (using pytest) and the issue doesn't happen when running the function in isolation.
I have captured though the diagnostic on the function, and as you can see, it also seems that the result[k, ...] is being hoisted out of the prange loop, which is wrong.
gmarkall commentedon Jan 16, 2024
@DrTodd13 Do you have any thoughts on this issue?
github-actions commentedon Feb 16, 2024
This issue is marked as stale as it has had no activity in the past 30 days. Please close this issue if no further response or action is needed. Otherwise, please respond with any updates and confirm that this issue still needs to be addressed.
netomenoci commentedon Feb 19, 2024
This issue shouldn't be closed. The pr #9397 is still waiting for review
github-actions commentedon Mar 21, 2024
This issue is marked as stale as it has had no activity in the past 30 days. Please close this issue if no further response or action is needed. Otherwise, please respond with any updates and confirm that this issue still needs to be addressed.
sklam commentedon May 28, 2024
Doing some git archaeology here, this issue is fixed by #9532.