From 212ffca2757e57644e1439186c82ec4e52f8d1db Mon Sep 17 00:00:00 2001 From: Jiaming Yuan Date: Fri, 25 Oct 2024 14:19:05 +0800 Subject: [PATCH] [EM] Small fixes for the example - Setup rmm in the initializer. --- demo/guide-python/distributed_extmem_basic.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/demo/guide-python/distributed_extmem_basic.py b/demo/guide-python/distributed_extmem_basic.py index 007bc733c965..3207572bc8e9 100644 --- a/demo/guide-python/distributed_extmem_basic.py +++ b/demo/guide-python/distributed_extmem_basic.py @@ -14,7 +14,6 @@ - cupy - rmm -- python-cuda """ @@ -100,7 +99,6 @@ def reset(self) -> None: def setup_rmm() -> None: """Setup RMM for GPU-based external memory training.""" import rmm - from cuda import cudart from rmm.allocators.cupy import rmm_cupy_allocator if not xgboost.build_info()["USE_RMM"]: @@ -170,7 +168,8 @@ def initializer(device: str) -> None: if device == "cuda": lop, sidx = mp.current_process().name.split("-") idx = int(sidx) # 1-based indexing from loky - os.environ["CUDA_VISIBLE_DEVICES"] = str(int(sidx) - 1) + os.environ["CUDA_VISIBLE_DEVICES"] = str(int(idx) - 1) + setup_rmm() with get_reusable_executor( max_workers=n_workers, initargs=(args.device,), initializer=initializer