11# SPDX-License-Identifier: LGPL-3.0-or-later
2+ import os
23import platform
34from ctypes import (
45 CDLL ,
78from importlib import (
89 metadata ,
910)
11+ from pathlib import (
12+ Path ,
13+ )
1014
1115import torch
1216from packaging .version import (
1923)
2024
2125
26+ _OP_LIB_OVERRIDE_ENV = "DEEPMD_OP_PT_LIB"
27+ _ALLOW_PRELOADED_ENV = "DEEPMD_OP_PT_ALLOW_PRELOADED"
28+
29+
30+ def _is_truthy_env (name : str ) -> bool :
31+ value = os .environ .get (name , "" )
32+ return value .strip ().lower () in {"1" , "true" , "yes" , "on" }
33+
34+
35+ def _resolve_module_file (
36+ module_name : str ,
37+ prefix : str ,
38+ ext : str ,
39+ ) -> tuple [Path , bool ]:
40+ override_path = os .environ .get (_OP_LIB_OVERRIDE_ENV )
41+ if override_path :
42+ return Path (override_path ).expanduser ().resolve (), True
43+ return (SHARED_LIB_DIR / (prefix + module_name )).with_suffix (ext ).resolve (), False
44+
45+
46+ def _loaded_library_hints () -> list [str ]:
47+ loaded_libraries = sorted (torch .ops .loaded_libraries )
48+ deepmd_libraries = [
49+ lib
50+ for lib in loaded_libraries
51+ if "deepmd_op_pt" in lib or "libdeepmd_op_pt" in lib
52+ ]
53+ return deepmd_libraries if deepmd_libraries else loaded_libraries
54+
55+
2256def load_library (module_name : str ) -> bool :
2357 """Load OP library.
2458
@@ -39,69 +73,85 @@ def load_library(module_name: str) -> bool:
3973 ext = ".so"
4074 prefix = "lib"
4175
42- module_file = ( SHARED_LIB_DIR / ( prefix + module_name )). with_suffix ( ext ). resolve ( )
76+ module_file , from_env_override = _resolve_module_file ( module_name , prefix , ext )
4377
44- if module_file .is_file ():
45- # Skip if this library was already loaded by torch.ops.load_library.
46- if str (module_file ) in torch .ops .loaded_libraries :
47- return True
48- # Skip if ops were already registered via C++ shared-library linkage
49- # (e.g. LAMMPS plugin links libdeepmd_op_pt.so at the C++ level).
50- # TORCH_LIBRARY(deepmd, m) in print_summary.cc registers "enable_mpi"
51- # as the first op; if it's accessible, the library is already loaded.
52- # Calling torch.ops.load_library again would abort() the process.
53- if hasattr (torch .ops , "deepmd" ) and hasattr (torch .ops .deepmd , "enable_mpi" ):
78+ if not module_file .is_file ():
79+ if from_env_override :
80+ raise RuntimeError (
81+ f"Environment variable { _OP_LIB_OVERRIDE_ENV } { module_file }
82+ )
83+ return False
84+
85+ # Skip if this exact library path was already loaded by torch.ops.load_library.
86+ if str (module_file ) in torch .ops .loaded_libraries :
87+ return True
88+
89+ # If deepmd ops are already registered before this call, abort by default
90+ # to avoid silently using an unexpected preloaded library.
91+ if hasattr (torch .ops , "deepmd" ) and hasattr (torch .ops .deepmd , "enable_mpi" ):
92+ if _is_truthy_env (_ALLOW_PRELOADED_ENV ):
5493 return True
55- try :
56- torch .ops .load_library (module_file )
57- except OSError as e :
58- # check: CXX11_ABI_FLAG; version
59- # from our op
60- PT_VERSION = GLOBAL_CONFIG ["pt_version" ]
61- PT_CXX11_ABI_FLAG = int (GLOBAL_CONFIG ["pt_cxx11_abi_flag" ])
62- # from torch
63- # strip the local version
64- pt_py_version = Version (torch .__version__ ).public
65- pt_cxx11_abi_flag = int (torch .compiled_with_cxx11_abi ())
66-
67- if PT_CXX11_ABI_FLAG != pt_cxx11_abi_flag :
68- raise RuntimeError (
69- "This deepmd-kit package was compiled with "
70- f"CXX11_ABI_FLAG={ PT_CXX11_ABI_FLAG }
71- f"with CXX11_ABI_FLAG={ pt_cxx11_abi_flag }
72- f"incompatible and thus an error is raised when loading { module_name }
73- "You need to rebuild deepmd-kit against this PyTorch "
74- "runtime."
75- ) from e
76-
77- # different versions may cause incompatibility, see TF
78- if PT_VERSION != pt_py_version :
79- raise RuntimeError (
80- "The version of PyTorch used to compile this "
81- f"deepmd-kit package is { PT_VERSION }
82- f"runtime you are using is { pt_py_version }
83- f"incompatible and thus an error is raised when loading { module_name }
84- f"You need to install PyTorch { PT_VERSION }
85- f"against PyTorch { pt_py_version } \n If you are using a wheel from "
86- "PyPI, you may consider to install deepmd-kit execuating "
87- "`DP_ENABLE_PYTORCH=1 pip install deepmd-kit --no-binary deepmd-kit` "
88- "instead."
89- ) from e
90- error_message = (
91- "This deepmd-kit package is inconsistent with PyTorch "
92- f"Runtime, thus an error is raised when loading { module_name }
94+ loaded_hints = _loaded_library_hints ()
95+ hint_text = "\n " .join (loaded_hints ) if loaded_hints else "(none reported by torch.ops.loaded_libraries)"
96+ raise RuntimeError (
97+ "DeepMD custom ops are already registered before deepmd.pt.cxx_op.load_library() "
98+ "could load the expected library path. This can indicate a mismatched or stale "
99+ "libdeepmd_op_pt.so in the current process.\n "
100+ f"Expected library path: { module_file } \n "
101+ f"Environment override ({ _OP_LIB_OVERRIDE_ENV } { os .environ .get (_OP_LIB_OVERRIDE_ENV , '(unset)' )} \n "
102+ f"Loaded-library hints:\n { hint_text } \n "
103+ f"If this preloaded setup is intentional, set { _ALLOW_PRELOADED_ENV }
104+ )
105+
106+ try :
107+ torch .ops .load_library (module_file )
108+ except OSError as e :
109+ # check: CXX11_ABI_FLAG; version
110+ # from our op
111+ PT_VERSION = GLOBAL_CONFIG ["pt_version" ]
112+ PT_CXX11_ABI_FLAG = int (GLOBAL_CONFIG ["pt_cxx11_abi_flag" ])
113+ # from torch
114+ # strip the local version
115+ pt_py_version = Version (torch .__version__ ).public
116+ pt_cxx11_abi_flag = int (torch .compiled_with_cxx11_abi ())
117+
118+ if PT_CXX11_ABI_FLAG != pt_cxx11_abi_flag :
119+ raise RuntimeError (
120+ "This deepmd-kit package was compiled with "
121+ f"CXX11_ABI_FLAG={ PT_CXX11_ABI_FLAG }
122+ f"with CXX11_ABI_FLAG={ pt_cxx11_abi_flag }
123+ f"incompatible and thus an error is raised when loading { module_name }
93124 "You need to rebuild deepmd-kit against this PyTorch "
94125 "runtime."
126+ ) from e
127+
128+ # different versions may cause incompatibility, see TF
129+ if PT_VERSION != pt_py_version :
130+ raise RuntimeError (
131+ "The version of PyTorch used to compile this "
132+ f"deepmd-kit package is { PT_VERSION }
133+ f"runtime you are using is { pt_py_version }
134+ f"incompatible and thus an error is raised when loading { module_name }
135+ f"You need to install PyTorch { PT_VERSION }
136+ f"against PyTorch { pt_py_version } \n If you are using a wheel from "
137+ "PyPI, you may consider to install deepmd-kit execuating "
138+ "`DP_ENABLE_PYTORCH=1 pip install deepmd-kit --no-binary deepmd-kit` "
139+ "instead."
140+ ) from e
141+ error_message = (
142+ "This deepmd-kit package is inconsistent with PyTorch "
143+ f"Runtime, thus an error is raised when loading { module_name }
144+ "You need to rebuild deepmd-kit against this PyTorch "
145+ "runtime."
146+ )
147+ if PT_CXX11_ABI_FLAG == 1 :
148+ # #1791
149+ error_message += (
150+ "\n WARNING: devtoolset on RHEL6 and RHEL7 does not support _GLIBCXX_USE_CXX11_ABI=1. "
151+ "See https://bugzilla.redhat.com/show_bug.cgi?id=1546704"
95152 )
96- if PT_CXX11_ABI_FLAG == 1 :
97- # #1791
98- error_message += (
99- "\n WARNING: devtoolset on RHEL6 and RHEL7 does not support _GLIBCXX_USE_CXX11_ABI=1. "
100- "See https://bugzilla.redhat.com/show_bug.cgi?id=1546704"
101- )
102- raise RuntimeError (error_message ) from e
103- return True
104- return False
153+ raise RuntimeError (error_message ) from e
154+ return True
105155
106156
107157def load_mpi_library () -> None :
0 commit comments