fix(pt,pt-expt): guard thread setters

Author: Han Wang

deepmd/pt/utils/env.py

@@ -93,9 +93,20 @@
 set_default_nthreads()
 intra_nthreads, inter_nthreads = get_default_nthreads()
 if inter_nthreads > 0:  # the behavior of 0 is not documented
-    torch.set_num_interop_threads(inter_nthreads)
+    # torch.set_num_interop_threads can only be called once per process.
+    # Guard to avoid RuntimeError when multiple backends are imported.
+    try:
+        if torch.get_num_interop_threads() != inter_nthreads:
+            torch.set_num_interop_threads(inter_nthreads)
+    except RuntimeError as err:
+        log.warning(f"Could not set torch interop threads: {err}")
 if intra_nthreads > 0:
-    torch.set_num_threads(intra_nthreads)
+    # torch.set_num_threads can also fail if called after threads are created.
+    try:
+        if torch.get_num_threads() != intra_nthreads:
+            torch.set_num_threads(intra_nthreads)
+    except RuntimeError as err:
+        log.warning(f"Could not set torch intra threads: {err}")
 
 __all__ = [
     "CACHE_PER_SYS",

deepmd/pt_expt/utils/env.py

@@ -93,9 +93,20 @@
 set_default_nthreads()
 intra_nthreads, inter_nthreads = get_default_nthreads()
 if inter_nthreads > 0:  # the behavior of 0 is not documented
-    torch.set_num_interop_threads(inter_nthreads)
+    # torch.set_num_interop_threads can only be called once per process.
+    # Guard to avoid RuntimeError when both pt and pt_expt env modules are imported.
+    try:
+        if torch.get_num_interop_threads() != inter_nthreads:
+            torch.set_num_interop_threads(inter_nthreads)
+    except RuntimeError as err:
+        log.warning(f"Could not set torch interop threads: {err}")
 if intra_nthreads > 0:
-    torch.set_num_threads(intra_nthreads)
+    # torch.set_num_threads can also fail if called after threads are created.
+    try:
+        if torch.get_num_threads() != intra_nthreads:
+            torch.set_num_threads(intra_nthreads)
+    except RuntimeError as err:
+        log.warning(f"Could not set torch intra threads: {err}")
 
 __all__ = [
     "CACHE_PER_SYS",

source/tests/pt/test_env_threads.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import importlib
+import logging
+
+import torch
+
+import deepmd.env as common_env
+
+
+def test_env_threads_guard_handles_runtimeerror(monkeypatch, caplog) -> None:
+    def raise_err(*_args, **_kwargs) -> None:
+        raise RuntimeError("boom")
+
+    monkeypatch.setattr(common_env, "set_default_nthreads", lambda: None)
+    monkeypatch.setattr(common_env, "get_default_nthreads", lambda: (1, 1))
+    monkeypatch.setattr(torch, "get_num_interop_threads", lambda: 2)
+    monkeypatch.setattr(torch, "set_num_interop_threads", raise_err)
+    monkeypatch.setattr(torch, "get_num_threads", lambda: 2)
+    monkeypatch.setattr(torch, "set_num_threads", raise_err)
+
+    caplog.set_level(logging.WARNING, logger="deepmd.pt.utils.env")
+    import deepmd.pt.utils.env as env
+
+    importlib.reload(env)
+
+    messages = [record.getMessage() for record in caplog.records]
+    assert any("Could not set torch interop threads" in msg for msg in messages)
+    assert any("Could not set torch intra threads" in msg for msg in messages)

source/tests/pt_expt/utils/test_env.py

@@ -0,0 +1,28 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import importlib
+import logging
+
+import torch
+
+import deepmd.env as common_env
+
+
+def test_env_threads_guard_handles_runtimeerror(monkeypatch, caplog) -> None:
+    def raise_err(*_args, **_kwargs) -> None:
+        raise RuntimeError("boom")
+
+    monkeypatch.setattr(common_env, "set_default_nthreads", lambda: None)
+    monkeypatch.setattr(common_env, "get_default_nthreads", lambda: (1, 1))
+    monkeypatch.setattr(torch, "get_num_interop_threads", lambda: 2)
+    monkeypatch.setattr(torch, "set_num_interop_threads", raise_err)
+    monkeypatch.setattr(torch, "get_num_threads", lambda: 2)
+    monkeypatch.setattr(torch, "set_num_threads", raise_err)
+
+    caplog.set_level(logging.WARNING, logger="deepmd.pt_expt.utils.env")
+    import deepmd.pt_expt.utils.env as env
+
+    importlib.reload(env)
+
+    messages = [record.getMessage() for record in caplog.records]
+    assert any("Could not set torch interop threads" in msg for msg in messages)
+    assert any("Could not set torch intra threads" in msg for msg in messages)