Add use_rmm flag to global configuration (#6656)

* Ensure RMM is 0.18 or later

* Add use_rmm flag to global configuration

* Modify XGBCachingDeviceAllocatorImpl to skip CUB when use_rmm=True

* Update the demo

* [CI] Pin NumPy to 1.19.4, since NumPy 1.19.5 doesn't work with latest Shap
This commit is contained in:
Philip Hyunsu Cho
2021-03-09 14:53:05 -08:00
committed by GitHub
parent e4894111ba
commit 366f3cb9d8
12 changed files with 117 additions and 20 deletions

View File

@@ -19,4 +19,14 @@ TEST(GlobalConfiguration, Verbosity) {
EXPECT_EQ(get<String>(current_config["verbosity"]), "0");
}
TEST(GlobalConfiguration, UseRMM) {
Json config{JsonObject()};
config["use_rmm"] = String("true");
auto& global_config = *GlobalConfigThreadLocalStore::Get();
FromJson(config, &global_config);
// GetConfig() should return updated use_rmm flag
Json current_config { ToJson(*GlobalConfigThreadLocalStore::Get()) };
EXPECT_EQ(get<String>(current_config["use_rmm"]), "1");
}
} // namespace xgboost