There have been several issues in the past about limiting threads, and the recommended way to do it is by setting task affinities (see https://github.com/google/jax/issues/1539#issuecomment-613743209). I'm running under slurm which does set task affinities, nevertheless jax allocates many (~25) threads on first call. I'm testing like this:
So I have two cores available, but 26 threads are spawned. The machine I'm running on has 40 cores total, so it's not just spawning a thread per core, but it's still a lot of threads. I'm also using a GPU so I wasn't expecting any CPU worker threads to get allocated.
What jax/jaxlib version are you using?
jax 0.4.2, jaxlib 0.4.1
Which accelerator(s) are you using?
GPU
Additional system info
No response
NVIDIA GPU info
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 515.65.01 Driver Version: 515.65.01 CUDA Version: 11.7 |
|-------------------------------+----------------------+----------------------+
| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |
| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |
| | | MIG M. |
|===============================+======================+======================|
| 0 Tesla V100-SXM2... On | 00000000:3A:00.0 Off | 0 |
| N/A 28C P0 41W / 300W | 0MiB / 32768MiB | 0% Default |
| | | N/A |
+-------------------------------+----------------------+----------------------+
+-----------------------------------------------------------------------------+
| Processes: |
| GPU GI CI PID Type Process name GPU Memory |
| ID ID Usage |
|=============================================================================|
| No running processes found |
+-----------------------------------------------------------------------------+
Description
There have been several issues in the past about limiting threads, and the recommended way to do it is by setting task affinities (see https://github.com/google/jax/issues/1539#issuecomment-613743209). I'm running under slurm which does set task affinities, nevertheless jax allocates many (~25) threads on first call. I'm testing like this:
and getting outputs like
So I have two cores available, but 26 threads are spawned. The machine I'm running on has 40 cores total, so it's not just spawning a thread per core, but it's still a lot of threads. I'm also using a GPU so I wasn't expecting any CPU worker threads to get allocated.
What jax/jaxlib version are you using?
jax 0.4.2, jaxlib 0.4.1
Which accelerator(s) are you using?
GPU
Additional system info
No response
NVIDIA GPU info