Fixed xtb excesses process spawning with the set_env_variable context…

… manager.

CHANGELOG.md

@@ -9,6 +9,19 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Added
 
 - `xtb` to optional installs and error message for `XtbAdapter` to indicate it can be installed with `pip install qcop[xtb]`.
+- `qcop.adapters.utils.set_env_variable` context manager to address the fact that one must set the env var `OMP_NUM_THREADS=1` _before_ importing any of the `xtb` library or else performance is severely impacted (around 10x). They spawn threads within threads so if you leave the default value (which is the number of cores on your machine) you get way more threads than this! E.g., on my 16 core machine I'll get 47 threads spawned for a single calculation. If I set this to `1` I actually get the correct number of threads, which is `16`. One must set this value _before_ importing any piece of the `xtb` library, so if your script imports anything from `xtb` before calling `qcop.compute("xtb", ...)` wrap this import with the `qcop.adapters.utils.set_env_variable` wrapper like this:
+
+  ```python
+  from qcop.adapters.utils.set_env_variable
+
+  with set_env_variable("OMP_NUM_THREADS", "1"):
+      import xtb
+      import xtb.interface
+      import xtb.libxtb
+      from xtb.utils import Solvent
+
+  ### The rest of your script
+  ```
 
 ## [0.7.1] - 2024-07-10
 

qcop/adapters/utils.py

@@ -265,3 +265,22 @@ def tmpdir(
             if rmdir:  # After exiting context manager
                 shutil.rmtree(temp_dir)
             os.chdir(cwd)
+
+
+@contextmanager
+def set_env_variable(var_name, value):
+    """Context manager to set an environment variable temporarily.
+
+    Args:
+        var_name: The name of the environment variable.
+        value: The value to set the environment variable to.
+    """
+    original_value = os.environ.get(var_name)
+    try:
+        os.environ[var_name] = str(value)
+        yield
+    finally:
+        if original_value:
+            os.environ[var_name] = original_value
+        else:
+            del os.environ[var_name]

qcop/adapters/xtb.py

@@ -18,7 +18,7 @@
 )
 
 from .base import ProgramAdapter
-from .utils import capture_sys_stdout
+from .utils import capture_sys_stdout, set_env_variable
 
 
 class XTBAdapter(ProgramAdapter[ProgramInput, SinglePointResults]):
@@ -58,22 +58,34 @@ def program_version(self, stdout: Optional[str] = None) -> str:
 
     @staticmethod
     def _ensure_xtb():
-        try:
-            xtb = importlib.import_module("xtb")
-            # xtb import structure is screwed up so I have to do this too
-            importlib.import_module("xtb.interface")
-            importlib.import_module("xtb.libxtb")
-
-            return xtb
-        except ModuleNotFoundError:
-            raise ProgramNotFoundError(
-                "xtb",
-                install_msg=(
-                    "Program not found: 'xtb'. To use xtb please install it with pip "
-                    "install qcop[xtb] or add '' if your shell requires it. e.g., "
-                    "pip install 'qcop[xtb]'."
-                ),
-            )
+        # NOTE: xtb-python is 10x slower with the default number of OpenMP threads,
+        # which I think is the number of CPUs on the machine. The xtb library reads in
+        # this env variable upon import so I have to set it before importing xtb. If you
+        # import any other xtb module in your code, e.g.,
+        # `from xtb.utils import Solvent` you have to set this env variable before
+        # importing that module or xtb will use the default number of threads which will
+        # be a 10x slower calculation despite the context manager used here. Strangely,
+        # setting this to one still results in spawning of a thread per logical core on
+        # your machine. And if you set this to the number of cores on your machine (in
+        # my case 16), xtb spawns 47!!! threads. So they are spawning threads within
+        # threads or some other insanity.
+        with set_env_variable("OMP_NUM_THREADS", "1"):
+            try:
+                xtb = importlib.import_module("xtb")
+                # xtb import structure is screwed up so I have to do this too
+                importlib.import_module("xtb.interface")
+                importlib.import_module("xtb.libxtb")
+
+                return xtb
+            except ModuleNotFoundError:
+                raise ProgramNotFoundError(
+                    "xtb",
+                    install_msg=(
+                        "Program not found: 'xtb'. To use xtb please install it with "
+                        "pip install qcop[xtb] or add '' if your shell requires it. "
+                        "e.g., pip install 'qcop[xtb]'."
+                    ),
+                )
 
     def compute_results(
         self,