Setting thread count in standard Python libraries

From NEClusterWiki
Revision as of 13:05, 25 July 2024 by Ondrejch (talk | contribs) (Created page with "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 sc...")
(diff) ← Older revision | Latest revision (diff) | Newer revision → (diff)
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")