Skip to content

Commit e647684

Browse files
authored
[MetaSchedule] Replace xgboost.rabit with xgboost.collective because it's deprecated (#17166)
* use collective instead of rabit * can work with xgb==1.4.2 in CI
1 parent 162d43a commit e647684

File tree

1 file changed

+7
-2
lines changed

1 file changed

+7
-2
lines changed

python/tvm/meta_schedule/cost_model/xgb_model.py

+7-2
Original file line numberDiff line numberDiff line change
@@ -755,7 +755,12 @@ def _fmt_metric(value, show_stdv=True):
755755
raise ValueError("wrong metric value", value)
756756

757757
import xgboost as xgb
758-
from xgboost import rabit # type: ignore
758+
759+
# make it compatible with xgboost<1.7
760+
try:
761+
from xgboost import rabit as collective # type: ignore
762+
except ImportError:
763+
from xgboost import collective # type: ignore
759764

760765
try:
761766
from xgboost.training import aggcv # type: ignore
@@ -841,7 +846,7 @@ def _fmt_metric(value, show_stdv=True):
841846
elif epoch - best_iteration >= self.early_stopping_rounds:
842847
best_msg = self.state["best_msg"]
843848

844-
if self.verbose_eval and rabit.get_rank() == 0:
849+
if self.verbose_eval and collective.get_rank() == 0:
845850
logger.debug("XGB stopped. Best iteration: %s ", best_msg)
846851
# instead of raising EarlyStopException, returning True to end the training
847852
return True

0 commit comments

Comments
 (0)