[Core] Make multiprocessing pool raise TypeError when provided a non-iterable in imap or imap_unordered

[cadedaniel]

TL;DR

This PR modifies ray.util.multiprocessing.Pool so that pool.imap and pool.imap_unordered throw TypeError when given non-iterable inputs.

More details

We want the API of ray.util.multiprocessing.Pool to closely resemble that of multiprocessing.Pool. One Ray user pointed out in #24237 that our imap and imap_unordered implementations differ slightly from multiprocessing.Pool in what arguments they accept. See the following usage:

#!/usr/bin/env python3

def main():
    import multiprocessing as mp
    with mp.Pool() as pool:
        assert_throws_on_noniterable(pool.imap)
        assert_throws_on_noniterable(pool.imap_unordered)

    import ray.util.multiprocessing as raymp
    with raymp.Pool() as pool:
        assert_throws_on_noniterable(pool.imap)
        assert_throws_on_noniterable(pool.imap_unordered)

def assert_throws_on_noniterable(method):
    iterable_input = [1, 2, 3]
    noniterable_input = 3

    assert len([_ for _ in method(f, iterable_input)]) == len(iterable_input)

    try:
        [_ for _ in method(f, noniterable_input)]
        assert False, f"expected TypeError for {method}"
    except TypeError:
        pass

def f(_):
    pass

if __name__ == '__main__':
    main()

With this PR, this code produces no exception. Without this PR, this fails:

Traceback (most recent call last):
  File "/data/ray/./cade.py", line 30, in <module>
    main()
  File "/data/ray/./cade.py", line 11, in main
    assert_throws_on_noniterable(pool.imap)
  File "/data/ray/./cade.py", line 22, in assert_throws_on_noniterable
    assert False, f"expected TypeError for {method}"
AssertionError: expected TypeError for <bound method Pool.imap of <ray.util.multiprocessing.pool.Pool object at 0x7f109b9ff640>>

Closes #24237

[cadedaniel]

We will wait until Ray 2.4 to merge this. In Ray 2.3 we will raise a deprecation warning, see PR #31845

[edoakes]

where are these error messages coming from, the iter call?

Not sure what it looks like, but should we instead wrap it and write a more user-friendly error message?

[cadedaniel]

We can; I opted not to because this is the same error that Pool from Python multiprocessing raises.

[cadedaniel]

review or merge pls cc @scv119 @jjyao

[scv119]

nice!