From e192296e0df9ba1857fe517677cfa33ed07dcc99 Mon Sep 17 00:00:00 2001 From: Masahiro Hiramori Date: Wed, 17 Jul 2024 10:16:00 +0900 Subject: [PATCH 1/2] use collective instead of rabit --- python/tvm/meta_schedule/cost_model/xgb_model.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/tvm/meta_schedule/cost_model/xgb_model.py b/python/tvm/meta_schedule/cost_model/xgb_model.py index 6b6b7a2dc1ed..38feb3252f75 100644 --- a/python/tvm/meta_schedule/cost_model/xgb_model.py +++ b/python/tvm/meta_schedule/cost_model/xgb_model.py @@ -755,7 +755,7 @@ def _fmt_metric(value, show_stdv=True): raise ValueError("wrong metric value", value) import xgboost as xgb - from xgboost import rabit # type: ignore + from xgboost import collective # type: ignore try: from xgboost.training import aggcv # type: ignore @@ -841,7 +841,7 @@ def _fmt_metric(value, show_stdv=True): elif epoch - best_iteration >= self.early_stopping_rounds: best_msg = self.state["best_msg"] - if self.verbose_eval and rabit.get_rank() == 0: + if self.verbose_eval and collective.get_rank() == 0: logger.debug("XGB stopped. Best iteration: %s ", best_msg) # instead of raising EarlyStopException, returning True to end the training return True From 5e849b6c5e1196f3c73fad990ec8f0ceb9d4f316 Mon Sep 17 00:00:00 2001 From: Masahiro Hiramori Date: Wed, 17 Jul 2024 17:23:46 +0900 Subject: [PATCH 2/2] can work with xgb==1.4.2 in CI --- python/tvm/meta_schedule/cost_model/xgb_model.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/python/tvm/meta_schedule/cost_model/xgb_model.py b/python/tvm/meta_schedule/cost_model/xgb_model.py index 38feb3252f75..aaee58fc94c8 100644 --- a/python/tvm/meta_schedule/cost_model/xgb_model.py +++ b/python/tvm/meta_schedule/cost_model/xgb_model.py @@ -755,7 +755,12 @@ def _fmt_metric(value, show_stdv=True): raise ValueError("wrong metric value", value) import xgboost as xgb - from xgboost import collective # type: ignore + + # make it compatible with xgboost<1.7 + try: + from xgboost import rabit as collective # type: ignore + except ImportError: + from xgboost import collective # type: ignore try: from xgboost.training import aggcv # type: ignore