3
\$\begingroup\$

One function of a program I'm working with benefits from using all available openMP threads. The rest of the program has issues with the env variable OMP_NUM_THREADS being unequal to 1.

I'm currently using the following solution to that, and was wondering whether there were any caveats or room for improvement here. I haven't worked much with context managers yet, and greatly appreciate any help/pointers. Thanks!

import os
import multiprocessing as mp

class MaxOMP:
    def __enter__(self):
        os.environ["OMP_NUM_THREADS"] = str(mp.cpu_count())

    def __exit__(self, exc_type, exc_val, exc_tb):
        os.environ["OMP_NUM_THREADS"] = "1"


def main():
    with MaxOMP():
        run_omp_intensive_job()
\$\endgroup\$
4
\$\begingroup\$

Instead of MaxOMP as a name, I like MaxOmp - in general still CamelCase your names even if they include an acronym.

Additionally, there isn't any reason why this needs to be restricted to any particular environment variable. As such, I would probably make this:

class EnvironmentVariableContextManager:

    _base_value = None
    _variable = None

    def __init__(self, variable_name, value):
        self.variable = variable_name
        self.target_value = value

    @property
    def variable(self):
        return self._variable

    @variable.setter
    def variable(self, name):
        self._base_value = os.environ.get(name, None)

        if self._base_value is None:
            raise Exception(f"{name} is not a valid environment variable")

        self._variable = name

    @property
    def original_value(self):
        return self._base_value

    def __enter__(self):
        os.environ[self.variable] = self.target_value

    def __exit__(self, *):
        os.environ[self.variable] = self.original_value

This has a few benefits:

  • We use properties to properly encapsulate access and add some extra validation (if you don't care about invalid environment variables feel free to remove that)
  • We restore it to the original value, not some hardcoded value. You could also provide a way to specify an "after" value if that was desirable.

As suggested in the comments, we may have a valid use-case for an environment variable that doesn't currently have a value. In that case, we would remove the exception, and instead do this:

@variable.setter
def variable(self, name):
    self._base_value = os.environ.get(name, None)
    self._variable = name

def __exit__(self, *):
    if self.original_value is None:
        os.environ.pop(self.variable)
    else:
        os.environ[self.variable] = self.original_value

Once we have that, then we just inherit.

class MaxOmpThreads(EnvironmentVariableContextManager):

    def __init__(self, max_threads=None):
        max_threads = max_threads or mp.cpu_count()
        super().__init__("OMP_NUM_THREADS", max_threads)

Another thing to note, is that os.environ only reflects the environment variables at the time of Python startup. If this is not what you want then you'll likely need a different solution.

This also may not be thread-safe; if you might have multiple processes working on something at once, and some need one setting and some need another, then you're likely to run into issues. If your program is overall serial, and the pieces that use OMP are well defined and contained, you're probably fine.

\$\endgroup\$
  • 1
    \$\begingroup\$ Instead of raising an exception, I would find it more intuitive if the context manager would set it on entry and delete it on exit (instead of resetting) in case the variable doesn't exist prior to entry. \$\endgroup\$ – 409_Conflict Aug 28 at 10:52
2
\$\begingroup\$

You don't need multiprocessing as os already have what you need : os.cpu_count

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.