Repository: numpy/numpy
File: numpy/_core/tests/test_multithreading.py (Python)

Code before:

        for f in futures:
            f.result()


def test_legacy_usertype_cast_init_thread_safety():
    def closure(b):
        b.wait()
        np.full((10, 10), 1, _rational_tests.rational)

    try:
        run_threaded(closure, 250, pass_barrier=True)
    except RuntimeError:
        # The 32 bit linux runner will trigger this with 250 threads. I can
        # trigger it on my Linux laptop with 500 threads but the CI runner is
        # more resource-constrained.
        # Reducing the number of threads means the test doesn't trigger the
        # bug. Better to skip on some platforms than add a useless test.
        pytest.skip("Couldn't spawn enough threads to run the test")

@pytest.mark.parametrize("dtype", [bool, int, float])
def test_nonzero_bool(dtype):
    # np.nonzero uses np.count_nonzero to determine the size of the output array
    # In a second pass the indices of the non-zero elements are determined, but they can have changed
    x = np.random.randint(4, size=10_000).astype(dtype)

    def func(seed):
        x[::2] = np.random.randint(2)
        try:
            _ = np.nonzero(x)
        except RuntimeError as ex:
            assert 'number of non-zero array elements changed during function execution' in str(ex)            

    run_threaded(func, max_workers=10, pass_count=True, outer_iterations=50)



Review comment (bug):
Maybe add that this test triggers a known data race and the behavior is undefined, we just want to make sure numpy doesn't seg fault.

Code after:

        for f in futures:
            f.result()


def test_legacy_usertype_cast_init_thread_safety():
    def closure(b):
        b.wait()
        np.full((10, 10), 1, _rational_tests.rational)

    try:
        run_threaded(closure, 250, pass_barrier=True)
    except RuntimeError:
        # The 32 bit linux runner will trigger this with 250 threads. I can
        # trigger it on my Linux laptop with 500 threads but the CI runner is
        # more resource-constrained.
        # Reducing the number of threads means the test doesn't trigger the
        # bug. Better to skip on some platforms than add a useless test.
        pytest.skip("Couldn't spawn enough threads to run the test")

@pytest.mark.parametrize("dtype", [bool, int, float])
def test_nonzero(dtype):
    # See: gh-28361
    #
    # np.nonzero uses np.count_nonzero to determine the size of the output array
    # In a second pass the indices of the non-zero elements are determined, but they can have changed
    #
    # This test triggers a data race which is suppressed in the TSAN CI. The test is to ensure
    # np.nonzero does not generate a segmentation fault
    x = np.random.randint(4, size=10_000).astype(dtype)

    def func(index):
        for _ in range(10):
            if index == 0:
                x[::2] = np.random.randint(2)
            else:
                try:
                    _ = np.nonzero(x)
                except RuntimeError as ex:
                    assert 'number of non-zero array elements changed during function execution' in str(ex)

    run_threaded(func, max_workers=10, pass_count=True, outer_iterations=50)