Stuck on an issue?

Lightrun Answers was designed to reduce the constant googling that comes with debugging 3rd party libraries. It collects links to all the places you might be looking at while hunting down a tough bug.

And, if you’re still stuck at the end, we’re happy to hop on a call to see how we can help out.

Limit number of threads when running on CPU

See original GitHub issue

I am assembling a single-core CPU performance benchmark for various HPC libraries from the modern Python ecosystem. I would like to include JAX, and first results seem very promising, but I’m failing to restrict it to a single thread. As far as I can tell, there is no corresponding setting in jax.config, and it doesn’t listen to any of the popular flags (e.g. OMP_NUM_THREADS).

I installed JAX and jaxlib from PyPI on OSX.

Is there any way to pull this off?

Issue Analytics

State:
Created 4 years ago
Comments:24 (12 by maintainers)

Top GitHub Comments

4reactions

hawkinspcommented, Apr 15, 2020

A quick note: setting the task affinity map is the correct way to limit JAX’s CPU usage at the moment. JAX sizes its main threadpool using this logic: https://github.com/tensorflow/tensorflow/blob/4b2cb67756009dda843c6b56a8b320c8a54373e0/tensorflow/core/platform/default/port.cc#L67

If launching from mpirun, mpirun knows how to set task affinities correctly.

2reactions

samuelacommented, Oct 21, 2019

Hey @dionhaefner ! I had a similar issue: https://github.com/google/jax/issues/743 HTH

Top Results From Across the Web

Maximum number of threads CPU - Stack Overflow

It really depends. As a rule of thumb, limit the number of threads to something close to the number of cores (otherwise you...

How can I tell the maximum threads my server can run?

So, you can have a total maximum of 20 threads executing in parallel, one thread per CPU/core. That can mean 20 single-threaded jobs, ......

What is the theoretical and practical limit on the numbers of ...

If you have a quad-core processor, you can have four hardware threads, or eight hardware threads if your processor is hyper-threaded. What this...

Maximum Number of Threads Per Process in Linux - Baeldung

The threads-max kernel parameter can be set to ensure that the number of threads per process is always less than or equal to...

Specify Maximum Number of Threads to Run Parallel for ...

Specify Maximum Number of Threads to Run Parallel for -Loops in the Generated Code ; 2. Configuration property (default value = 0):. NumberOfCpuThreads....