allow multi call on init (#92)

This commit is contained in:
Nan Zhu 2019-04-26 18:41:02 -07:00 committed by GitHub
parent 5c3b36f346
commit a429748e24
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23

View File

@ -45,11 +45,11 @@ typedef ThreadLocalStore<ThreadLocalEntry> EngineThreadLocal;
/*! \brief intiialize the synchronization module */ /*! \brief intiialize the synchronization module */
void Init(int argc, char *argv[]) { void Init(int argc, char *argv[]) {
ThreadLocalEntry* e = EngineThreadLocal::Get(); ThreadLocalEntry* e = EngineThreadLocal::Get();
utils::Check(e->engine.get() == nullptr, if (e->engine.get() == nullptr) {
"rabit::Init is already called in this thread"); e->initialized = true;
e->initialized = true; e->engine.reset(new Manager());
e->engine.reset(new Manager()); e->engine->Init(argc, argv);
e->engine->Init(argc, argv); }
} }
/*! \brief finalize syncrhonization module */ /*! \brief finalize syncrhonization module */
@ -59,6 +59,7 @@ void Finalize() {
"rabit::Finalize engine is not initialized or already been finalized."); "rabit::Finalize engine is not initialized or already been finalized.");
e->engine->Shutdown(); e->engine->Shutdown();
e->engine.reset(nullptr); e->engine.reset(nullptr);
e->initialized = false;
} }
/*! \brief singleton method to get engine */ /*! \brief singleton method to get engine */
@ -68,8 +69,7 @@ IEngine *GetEngine() {
ThreadLocalEntry* e = EngineThreadLocal::Get(); ThreadLocalEntry* e = EngineThreadLocal::Get();
IEngine* ptr = e->engine.get(); IEngine* ptr = e->engine.get();
if (ptr == nullptr) { if (ptr == nullptr) {
utils::Check(!e->initialized, utils::Check(!e->initialized, "the rabit has not been initialized");
"Doing rabit call after Finalize");
return &default_manager; return &default_manager;
} else { } else {
return ptr; return ptr;