API for setting limits for use by Python thread pools

[itamarst]

See #208 for discussion of why the current API is insufficient.

Here is a demonstration of the problem as far as the current API goes, will eventually become a unit test for new API:

from sys import thread_info
from threading import Thread, current_thread

import sklearn
from threadpoolctl import threadpool_info, threadpool_limits

def assert_limit_is(limit):
    print("Checking", current_thread())
    for i in threadpool_info():
        assert i["num_threads"] == limit
    print("OK")


threadpool_limits(limits=2)
assert_limit_is(2)

t = Thread(target=lambda: assert_limit_is(2))
t.start()
t.join()

When run:

Checking <_MainThread(MainThread, started 128633514360960)>
OK
Checking <Thread(Thread-1, started 128632355026624)>
Exception in thread Thread-1:
Traceback (most recent call last):
  File "/usr/lib/python3.9/threading.py", line 980, in _bootstrap_inner
    self.run()
  File "/usr/lib/python3.9/threading.py", line 917, in run
    self._target(*self._args, **self._kwargs)
  File "/home/itamarst/devel/threadpoolctl/theproblem.py", line 17, in <lambda>
    t = Thread(target=lambda: assert_limit_is(2))
  File "/home/itamarst/devel/threadpoolctl/theproblem.py", line 10, in assert_limit_is
    assert i["num_threads"] == limit
AssertionError

Thinking through the design space

Limiting APIs might affect only the current thread, or all threads in the process. If the former, we need to call the API in every thread in the thread pool; if the latter, we need to call it before starting the thread pool in the thread that starts the pool.

If every library has a single exposed API, there's really no difference between the logic in the current API threadpool_limits() and what you'd do in a thread pool thread. You don't have a choice, you just have the one API!

However, sometimes there is a choice. For MKL there will be a choice between "apply to all threads" vs "apply just to this thread". If there is a choice for a library:

• In the main thread, you don't really need to call anything for the library, if the presumed usage pattern is that all work is done in the new thread pool. In fact you don't want to call process-wide APIs because of nested thread pools.
• In the threads in the thread pool, you want to call the thread-local API.

Proposed solutions

Noted that for all options, nested Python thread pools are gonna be half-broken given some limiting APIs are process-wide. But will kinda-sorta work if all threads take the same amount of time to run... I have an idea for a completely different approach to deal with that though.

Option 1: Single function custom API

Sketch of proposed API, updated version of #208:

from multiprocessing.pool import ThreadPool
from os import cpu_count

from threadpoolctl import per_thread_limit


def run(n_threads: int):
    cores_per_thread = max(1, cpu_count()) // n_threads

    with per_thread_limit(cores_per_thread):
        with ThreadPool(n_threads, initializer=lambda: per_thread_limit(cores_per_thread)) as pool:
            # ... do business logic here ...

Option 2: Two-phase custom API

from multiprocessing.pool import ThreadPool
from os import cpu_count

from threadpoolctl import per_thread_limit


def run(n_threads: int):
    with per_thread_limit(max(1, cpu_count()) // n_threads) as limiter:
        with ThreadPool(n_threads, initializer=limiter.limit_in_current_thread) as pool:
            # ... do business logic here ...

Option 3: Reuse current API

The positive: no new APIs, it's just documentation.

The negative:

• If we get better affordances in the future, this is bad because it's ending up using process wide APIs when it shouldn't. E.g. mkl has both process-wide and current-thread get/set-num-threads APIs. threadpoolctl could choose to expose both, and then use both as needed, but this API option would preclude taking advantage of that. And that is actually viable (see Expose thread-local APIs when possible #216) so going to say this is a bad idea.

Option 4: Current API at top-level, new API for threads in thread pool

Rejected because with nested thread pools, in cases where you can choose between process-wide and thread-local APIs, this will unnecessary call process-wide API and impact other thread pools.

Conclusion

I like the two-phase API because:

• You can have different logic for setup thread and in-thread-pool, and that may well be useful.
• You can detect user coding errors (not calling the per-thread API) and warn about it, so it's less error prone.