Fix pyspark with verbosity=3. (#10172)

This commit is contained in:
Jiaming Yuan 2024-04-09 23:18:56 +08:00 committed by GitHub
parent a99bb38bd2
commit f0a138f33a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 19 additions and 15 deletions

View File

@ -1052,6 +1052,7 @@ class _SparkXGBEstimator(Estimator, _SparkXGBParams, MLReadable, MLWritable):
dev_ordinal = None dev_ordinal = None
use_qdm = _can_use_qdm(booster_params.get("tree_method", None)) use_qdm = _can_use_qdm(booster_params.get("tree_method", None))
verbosity = booster_params.get("verbosity", 1)
msg = "Training on CPUs" msg = "Training on CPUs"
if run_on_gpu: if run_on_gpu:
dev_ordinal = ( dev_ordinal = (
@ -1089,6 +1090,7 @@ class _SparkXGBEstimator(Estimator, _SparkXGBParams, MLReadable, MLWritable):
evals_result: Dict[str, Any] = {} evals_result: Dict[str, Any] = {}
with CommunicatorContext(context, **_rabit_args): with CommunicatorContext(context, **_rabit_args):
with xgboost.config_context(verbosity=verbosity):
dtrain, dvalid = create_dmatrix_from_partitions( dtrain, dvalid = create_dmatrix_from_partitions(
pandas_df_iter, pandas_df_iter,
feature_prop.features_cols_names, feature_prop.features_cols_names,

View File

@ -307,14 +307,16 @@ class MemoryLogger {
void RegisterDeallocation(void *ptr, size_t n, int current_device) { void RegisterDeallocation(void *ptr, size_t n, int current_device) {
auto itr = device_allocations.find(ptr); auto itr = device_allocations.find(ptr);
if (itr == device_allocations.end()) { if (itr == device_allocations.end()) {
LOG(WARNING) << "Attempting to deallocate " << n << " bytes on device " LOG(WARNING) << "Attempting to deallocate " << n << " bytes on device " << current_device
<< current_device << " that was never allocated "; << " that was never allocated\n"
} << dmlc::StackTrace();
} else {
num_deallocations++; num_deallocations++;
CHECK_LE(num_deallocations, num_allocations); CHECK_LE(num_deallocations, num_allocations);
currently_allocated_bytes -= itr->second; currently_allocated_bytes -= itr->second;
device_allocations.erase(itr); device_allocations.erase(itr);
} }
}
}; };
DeviceStats stats_; DeviceStats stats_;
std::mutex mutex_; std::mutex mutex_;