use RabitContext intead of init/finalize (#7911)

2022-05-16 21:15:41 -07:00
parent 4fcfd9c96e
commit 77d4a53c32
6 changed files with 137 additions and 153 deletions
--- a/tests/python/test_tracker.py
+++ b/tests/python/test_tracker.py
@@ -16,10 +16,9 @@ def test_rabit_tracker():
    rabit_env = []
    for k, v in worker_env.items():
        rabit_env.append(f"{k}={v}".encode())
-    xgb.rabit.init(rabit_env)
-    ret = xgb.rabit.broadcast('test1234', 0)
-    assert str(ret) == 'test1234'
-    xgb.rabit.finalize()
+    with xgb.rabit.RabitContext(rabit_env):
+        ret = xgb.rabit.broadcast('test1234', 0)
+        assert str(ret) == 'test1234'


 def run_rabit_ops(client, n_workers):