Hendrik/xgboost
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Speed up tests with federated learning enabled (#8350)

Browse Source 
* Speed up tests with federated learning enabled

* Re-enable timeouts

Co-authored-by: Hyunsu Philip Cho <chohyu01@cs.washington.edu>
This commit is contained in:
Rong Ou
2022-10-17 15:17:04 -07:00
committed by GitHub
parent 031d66ec27
commit 8f3dee58be
5 changed files with 10 additions and 14 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
src/c_api/c_api.cu
View File

@@ -42,19 +42,15 @@ void XGBBuildInfoDevice(Json *p_info) {
}
void XGBoostAPIGuard::SetGPUAttribute() {
try {
device_id_ = dh::CurrentDevice();
} catch (dmlc::Error const&) {
// do nothing, running on CPU only machine
}
// Not calling `safe_cuda` to avoid unnecessary exception handling overhead.
// If errors, do nothing, assuming running on CPU only machine.
cudaGetDevice(&device_id_);
}
void XGBoostAPIGuard::RestoreGPUAttribute() {
try {
dh::safe_cuda(cudaSetDevice(device_id_));
} catch (dmlc::Error const&) {
// do nothing, running on CPU only machine
}
// Not calling `safe_cuda` to avoid unnecessary exception handling overhead.
// If errors, do nothing, assuming running on CPU only machine.
cudaSetDevice(device_id_);
}
} // namespace xgboost