Test adding unload method for shared providers

This commit is contained in:
Ryan Hill 2021-04-29 22:41:11 -07:00
parent 9df9325fcb
commit c427b78799
2 changed files with 11 additions and 3 deletions

View file

@ -938,6 +938,7 @@ void addGlobalMethods(py::module& m, Environment& env) {
throw std::runtime_error("Error when creating and registering allocator: " + st.ErrorMessage());
}
});
m.def("unload_shared_providers", &UnloadSharedProviders, "Unload the shared providers (needs to be done for a safe shutdown)");
#ifdef USE_NUPHAR
// TODO remove deprecated global config
@ -2097,9 +2098,9 @@ PYBIND11_MODULE(onnxruntime_pybind11_state, m) {
LOGS(default_logger, WARNING) << "Init provider bridge failed.";
}
atexit([] {
UnloadSharedProviders();
});
// atexit([] {
// UnloadSharedProviders();
// });
#endif
#ifdef ENABLE_TRAINING

View file

@ -930,5 +930,12 @@ class TestInferenceSession(unittest.TestCase):
set())
print("Create session with customize execution provider successfully!")
def testUnloadSharedProviders(self):
print("Unloading shared providers")
from onnxruntime.capi import _pybind_state as C
C.unload_shared_providers()
print("Finished unloading")
if __name__ == '__main__':
unittest.main()