Limiting the number of cores/threads used in PyMC5.6+

1

There seem to be a number of queries posted here about changing the number of cores used, many of them unanswered or unresolved. And most are old. So I’m starting a new question, but correct me if this is wrong practice.

Is there documentation / examples of how to limit the number of cores used in sampling/estimation? Or how to choose a computation backend, if that’s a thing?

This is a simple four-chain estimate. It uses 128 threads!! :

pymc-sampling-htop
pymc-sampling-htop3840×404 120 KB

Crazy! And I’m not sure how efficient this is? Aren’t all those kernel processes (ie the red color in the load bars) bad news?

I’m not specifying the cores= parameter in pm.sample, which others say does not help this problem.

I’m interested in what is efficient, and also in limiting the number of cores used so that I can estimate more than one model at once, and also so that my server can do other things!!
My computation processes are niced.

Thanks!
c

2

You can try setting the environment variable OMP_NUM_THREADS=1.

3

That did not make any difference!

!echo $OMP_NUM_THREADS
40

But the estimate still used 128 threads.

4

You can try MKL_NUM_THREADS instead

5

Same thing using MKL_NUM_THREADS as well: all 128 threads fully used.

6

Maybe worth trying to reproduce on a more conventional machine and see if the problem also crops up there?

7

And to be sure did you tey setting them to 1, not 40? Are those real cpu cores or virtual ones?

9

Okay, I tried on my laptop with the max set to 2 and then with it set to 1. In both cases, all 16 threads of my laptop are used at 100%.

Just to be clear, this does not happen when my data size is small (1000) but does when it is larger (10000). With smaller samples, four threads are used (one for each chain).

10

How and when do you set the env variables?

11

From within python, before sampling the model.

     
        if max_processor_threads is not None:
            os.environ["OMP_NUM_THREADS"] = str(max_processor_threads)
            print(f"Set OMP_NUM_THREADS to {max_processor_threads}")
            os.environ["MKL_NUM_THREADS"] = str(max_processor_threads)
            print(f"Set MKL_NUM_THREADS to {max_processor_threads}")
        trace_filename = f"{self.basename}.nc"
        
        print(f"Building model {modelclass} for {self.basename}")
        model = self.build_model(df, modelclass, **kwargs)
        with model:
            trace = pm.sample() #return_inferencedata=True)
        
    ```
12

Try to do it before any other imports

1 Like