[Data] [1/n] Async iter_batches: Add Threadpool util

[amogkam]

Part 1 of async iter_batches support. Converts the existing _make_async_gen util to be able to use a threadpool instead of a single thread.

Why are these changes needed?

Related issue number

Checks

• I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
  • I've added any new APIs to the API Reference. For example, if I added a
    method in Tune, I've added it in doc/source/tune/api/ under the
    corresponding .rst file.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[ollie-iterators]

It looks like there is an issue on line 169 of the file "python/ray/data/_internal/block_batching/block_batching.py". The text "prefetch_buffer_size" should be removed as it is not a part of the _make_async_gen function anymore.

[ollie-iterators]

It should cause 10 of the failures to be fixed.

[clarkzinzow]

Mostly nits/questions; my biggest questions are:

1. Whether making the threads daemonic will lead to improperly released resources, which might lead to some bad bugs.
2. Whether output_queue.join() should be called from the thread that calls .task_done(), which I think might be nonsensical.

[clarkzinzow]

Why do we want these to be daemon threads? This is slightly concerning since resources might not be released properly when daemonic threads are terminated (i.e. no __exit__s run for context managers or finally blocks for try-finally patterns).

I'm guessing you went this route because we don't want to add the signaling mechanism necessary for graceful worker thread termination, but it doesn't seem like it would be too difficult to capture a threading.Event in the execute_computation closure and to check if it's set every 100ms when blocked while putting in the queue.

    shutdown = threading.Event()

    def try_put(item: Any) -> bool:
        while not shutdown.is_set():
            try:
                output_queue.put(item, block=True, timeout=0.1)
            except queue.Full:
                pass
        else:
            return True
        return False

    def execute_computation(thread_index: int):
        try:
            for item in fn(thread_safe_generator):
                if try_put(item):
                    return
            try_put(Sentinel(thread_index))
        except Exception as e:
            try_put(e)

    while True:
        next_item = output_queue.get(block=True)
        if isinstance(next_item, Exception):
            output_queue.task_done()
            shutdown.set()
            raise next_item

[amogkam]

In this case, the event is only set if one of the child threads raises an exception. But what happens if the main thread is interrupted? In that case, the event will not be set and the child threads will still be alive right?

[ericl]

@clarkzinzow regarding daemon threads, the main thing is not leaking drivers / worker processes when things aren't properly shut down. For non-daemon threads, those hold the entire process open until you force exit.

The streaming executor is run on a daemon thread, for example, due to this issue. On process shutdown you don't care about finally blocks, anyways.

[clarkzinzow]

@ericl Ah, so we're making the worker threads daemonic to prevent spurious blocking within fn (which would circumvent our shutdown event checking) from keeping the process from exiting, and we're only terminating these threads asynchronously when the process is exiting so we don't care about properly releasing resources. Got it, sounds good!