Skip to content

treams_rs.parallel

CPU threads of the native numerical work.

treams-rs runs its parallel kernels and dense linear algebra on one pool of worker threads that it owns. The pool starts on the first parallel call. A forked child process (multiprocessing, ProcessPoolExecutor, PyTorch DataLoader workers) builds its own pool instead of waiting on the parent's threads, which do not exist after fork.

The default budget is the first positive integer among TREAMS_RS_NUM_THREADS, RAYON_NUM_THREADS and OMP_NUM_THREADS (the outermost entry of an OpenMP list), otherwise the CPUs this process may use, which respects CPU affinity and container quotas. The environment is read once, when treams_rs is imported; later changes to it have no effect, so call set_num_threads instead, also in process-pool initializers. Empty and zero values mean "not set". Unparsable values raise ThreadingWarning at import; they, budgets above the available CPUs and an OMP_NUM_THREADS limit are listed in thread_info()["diagnostics"]. joblib's loky backend sets OMP_NUM_THREADS in its workers, so each worker uses its share of the CPUs.

When threadpoolctl is installed, threadpoolctl.threadpool_limits limits treams-rs too: threadpool_info() lists it with user_api "treams".

Results do not depend on the budget: for a given build and CPU, every value and gradient repeats bit for bit at any number of threads.

ThreadingWarning

An ignored or oversubscribing thread setting.

get_num_threads

get_num_threads() -> int

Return the worker-thread budget of the next native parallel call.

set_num_threads

set_num_threads(threads: int | None) -> None

Set the process-wide worker-thread budget; None restores the default.

threads must be a positive integer. A parallel region that has started finishes on its pool; later regions use the new budget. A budget above the available CPUs is honored with a ThreadingWarning, issued before the budget changes, because oversubscription slows dense linear algebra. Forked children inherit the budget and build their own pool of that size.

Examples

>>> import treams_rs as tr
>>> tr.set_num_threads(1)
>>> tr.get_num_threads()
1
>>> tr.set_num_threads(None)

thread_info

thread_info() -> dict[str, Any]

Return the thread budget, where it came from, and the pool state.

Keys:

  • threads: the budget of the next parallel call.
  • source: "set_num_threads", the environment variable that set the budget, or "available_parallelism".
  • available: the CPUs this process may use.
  • pool_threads: the workers of the current pool, None before the first parallel call of this process.
  • forked: whether this process was forked after treams-rs was loaded; it then builds a pool of its own.
  • diagnostics: ignored, oversubscribing or limiting settings.
  • environment: the variables consulted, in precedence order.

Never starts the pool.

Examples

>>> import treams_rs as tr
>>> tr.thread_info()["threads"] >= 1
True

threads

threads(n: int | None) -> Iterator[None]

Set the process-wide worker-thread budget for the with block.

Restores the previous budget on exit, like threadpoolctl.threadpool_limits. The budget is process-wide, not thread-local: other Python threads that call treams-rs inside the block use it too. Alternating between two budgets reuses their pools.

Blocks in different Python threads must not overlap: each restores the budget it saw on entry, in exit order. For a thread pool, call set_num_threads(k) once before starting the workers.

Examples

>>> import treams_rs as tr
>>> with tr.threads(1):
...     tr.get_num_threads()
1