Repository: google/jax

Issue:
LLVM worker thread limit
### Description

Recently we have been seeing a lot of crashed jobs on our cluster, with the following error message:
```
LLVM ERROR: pthread_create failed: Resource temporarily unavailable
```
caused by the jobs hitting the thread limit per user per node (set to 4096).

Further investigation revealed that this is caused by jax spawning hundreds/thousands of llvm-worker threads, and never exiting them.
This happens whenever we run a script which contains many compilations, and a certain subset of jax primitives, e.g. qr decomposition.

Here is a minimal working example to reproduce the aforementioned behaviour:

```python
import jax
import jax.numpy as jnp
from functools import partial

@partial(jax.jit, static_argnums=0)
def expensive_computation(f, x):
    x = f(x)
    q, r = jnp.linalg.qr(x)
    return q.sum()+r.sum()

x = jnp.ones((16,16))

i =1
while True:
    i = i+1
    if i%1000 ==0: print(i)
    f = lambda x: x #trigger recompilation
    u = jax.block_until_ready(expensive_computation(f, x))
```
After a few thousand iterations, ```top -Hp $(pgrep python3)``` shows ~6100 llvm-worker threads which have been spawned:
```
top - 12:04:00 up  3:05,  3 users,  load average: 1.04, 0.52, 0.26
Threads: 6178 total,   1 running, 6177 sleeping,   0 stopped,   0 zombie
%Cpu(s): 33.7 us,  1.0 sy,  0.0 ni, 65.3 id,  0.0 wa,  0.0 hi,  0.0 si,  0.0 st
MiB Mem :  60000.0 total,  56564.2 free,   3173.1 used,    262.8 buff/cache
MiB Swap:      0.0 total,      0.0 free,      0.0 used.  56826.9 avail Mem 

    PID USER      PR  NI    VIRT    RES    SHR S  %CPU  %MEM     TIME+ COMMAND                                                                                            
 507363 clemens   20   0   52.8g   2.9g 103472 R  99.0   5.0   3:05.13 python3                                                                                            
 507364 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.06 python3                                                                                            
 507365 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.07 python3                                                                                            
 507367 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507368 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507369 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507370 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507371 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507372 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507373 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507374 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507375 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507376 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507377 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507378 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507379 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507380 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507381 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507382 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507383 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507384 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507385 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507386 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507387 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507388 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 python3                                                                                            
 507390 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.02 python3                                                                                            
 507391 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.06 python3                                                                                            
 507392 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.07 python3                                                                                            
 521501 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-0                                                                                      
 521502 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-1                                                                                      
 521503 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-2                                                                                      
 521507 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-0                                                                                      
 521508 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-1                                                                                      
 521509 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-2                                                                                      
 521513 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-0                                                                                      
 521514 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-1                                                                                      
 521515 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-2                                                                                      
 521519 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-0                                                                                      
 521520 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-1                                                                                      
 521521 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-2                                                                                      
 521525 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-0                                                                                      
 521526 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-1                                                                                      
 521527 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-2                                                                                      
 521531 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-0                                                                                      
 521532 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-1                                                                                      
 521533 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-2                                                                                      
 521537 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-0                                                                                      
 521538 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-1                                                                                      
 521539 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-2                                                                                      
 521543 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-0                                                                                      
 521544 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-1                                                                                      
 521545 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-2                                                                                      
 521549 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-0                                                                                      
 521550 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-1                                                                                      
 521551 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-2                                                                                      
 521555 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-0                                                                                      
 521556 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-1                                                                                      
 521557 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-2                                                                                      
 521561 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-0                                                                                      
 521562 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-1                                                                                      
 521563 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-2                                                                                      
 521567 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-0                                                                                      
 521568 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-1                                                                                      
 521569 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-2                                                                                      
 521573 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-0                                                                                      
 521574 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-1                                                                                      
 521575 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-2                                                                                      
 521579 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-0                                                                                      
 521580 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-1                                                                                      
 521581 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-2                                                                                      
 521585 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-0                                                                                      
 521586 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-1                                                                                      
 521587 clemens   20   0   52.8g   2.9g 103472 S   0.0   5.0   0:00.00 llvm-worker-2
...                            
```
(this was run locally, in an lxc container limited to 3 cpu cores)

A few more remarks:
- if the `jnp.linalg.qr` decomposition is replaced with some normal jax function, e.g. `jax.lax.sin` the issue does not happen
- as can be seen from the thread names (repeated indices), they are not all spawned by the same `llvm::ThreadPool`, but somewhere else (I am not familiar enough with XLA to find where)

Context:
- possibly related to https://github.com/google/jax/issues/15819 and https://github.com/google/jax/issues/16215
- this issue is different from https://github.com/google/jax/issues/1539 and https://github.com/google/jax/issues/15866  , as we are not worried about the number of threads jax is using for doing the acutal computations (which can be limited via cpu pinning e.g. taskset) but rather about the compilation threads


Let me know if this is the right place or if I should open an issue on https://github.com/openxla/xla instead.



### What jax/jaxlib version are you using?

jax 0.4.11, jaxlib 0.4.11

### Which accelerator(s) are you using?

CPU


Discussion:

Resolving patch:
diff --git a/jax/_src/interpreters/mlir.py b/jax/_src/interpreters/mlir.py
--- a/jax/_src/interpreters/mlir.py
+++ b/jax/_src/interpreters/mlir.py
@@ -360,6 +360,14 @@ def _source_info_to_location(
 def make_ir_context() -> ir.Context:
   """Creates an MLIR context suitable for JAX IR."""
   context = ir.Context()
+
+  # If threading is enabled, each MLIR context will keep alive a thread pool.
+  # Since we cache MLIR modules (and hence contexts), this means we might keep
+  # several threads alive for each cache entry. This is a terrible idea. However
+  # we don't do any heavy computation on MLIR modules from Python anyway, so we
+  # just disable threading.
+  context.enable_multithreading(False)
+
   dialects.mhlo.register_mhlo_dialect(context)
   dialects.chlo.register_dialect(context)
   dialects.stablehlo.register_dialect(context)


Test patch: