Setting thread count in standard Python libraries
Jump to navigation
Jump to search
All, the issue I had came from under-the-hood python packages automatically starting several processes, even though I was only requesting a single core in my job submission script. I’ve attached a function that will force the number of threads to a specific value, as it is important that the number of threads your program is trying to use aligns with the requested resources. On local machines, python should only use what is available and does pretty well playing nicely with other programs, but this is not the case with how torque distributes jobs.
These should work for most python applications, as I’m using many of the standard python packages (numpy, spicy, pandas, etc.).
# Function to set threading limits
def set_threading_limits(num_threads):
os.environ['OMP_NUM_THREADS'] = str(num_threads)
os.environ['OPENBLAS_NUM_THREADS'] = str(num_threads)
os.environ['MKL_NUM_THREADS'] = str(num_threads)
os.environ['VECLIB_MAXIMUM_THREADS'] = str(num_threads)
os.environ['NUMEXPR_NUM_THREADS'] = str(num_threads)
print(f"Threading limits set to {num_threads}”)
# Function to restore defaults
def unset_threading_limits():
os.environ.pop('OMP_NUM_THREADS', None)
os.environ.pop('OPENBLAS_NUM_THREADS', None)
os.environ.pop('MKL_NUM_THREADS', None)
os.environ.pop('VECLIB_MAXIMUM_THREADS', None)
os.environ.pop('NUMEXPR_NUM_THREADS', None)
print("Threading limits unset")