From 40d32455bec79b4747901871db14e96097ef9412 Mon Sep 17 00:00:00 2001 From: Muhammed Fatih BALIN Date: Tue, 6 Aug 2024 09:18:06 -0400 Subject: [PATCH] [GraphBolt][CUDA] Warn only when cuda is available. (#7643) --- python/dgl/graphbolt/__init__.py | 50 ++++++++++++++++++-------------- 1 file changed, 29 insertions(+), 21 deletions(-) diff --git a/python/dgl/graphbolt/__init__.py b/python/dgl/graphbolt/__init__.py index a954f4f82f..2c980dd532 100644 --- a/python/dgl/graphbolt/__init__.py +++ b/python/dgl/graphbolt/__init__.py @@ -9,33 +9,36 @@ An experimental feature for CUDA allocations is turned on for better allocation pattern resulting in better memory usage for minibatch GNN training workloads. See https://pytorch.org/docs/stable/notes/cuda.html#optimizing-memory-usage-with-pytorch-cuda-alloc-conf, and set the environment variable `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False` -if you want to disable it. +if you want to disable it and set it True to acknowledge and disable the warning. """ cuda_allocator_env = os.getenv("PYTORCH_CUDA_ALLOC_CONF") -if cuda_allocator_env is None: - os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "expandable_segments:True" - gb_warning(CUDA_ALLOCATOR_ENV_WARNING_STR) -else: - configs = { +WARNING_STR_TO_BE_SHOWN = None +configs = ( + {} + if cuda_allocator_env is None + else { kv_pair.split(":")[0]: kv_pair.split(":")[1] for kv_pair in cuda_allocator_env.split(",") } - if "expandable_segments" in configs: - if configs["expandable_segments"] != "True": - gb_warning( - "You should consider `expandable_segments:True` in the" - " environment variable `PYTORCH_CUDA_ALLOC_CONF` for lower" - " memory usage. See " - "https://pytorch.org/docs/stable/notes/cuda.html" - "#optimizing-memory-usage-with-pytorch-cuda-alloc-conf" - ) - else: - configs["expandable_segments"] = "True" - os.environ["PYTORCH_CUDA_ALLOC_CONF"] = ",".join( - [k + ":" + v for k, v in configs.items()] +) +if "expandable_segments" in configs: + if configs["expandable_segments"] != "True": + WARNING_STR_TO_BE_SHOWN = ( + "You should consider `expandable_segments:True` in the" + " environment variable `PYTORCH_CUDA_ALLOC_CONF` for lower" + " memory usage. See " + "https://pytorch.org/docs/stable/notes/cuda.html" + "#optimizing-memory-usage-with-pytorch-cuda-alloc-conf" ) - gb_warning(CUDA_ALLOCATOR_ENV_WARNING_STR) - +else: + configs["expandable_segments"] = "True" + os.environ["PYTORCH_CUDA_ALLOC_CONF"] = ",".join( + [k + ":" + v for k, v in configs.items()] + ) + WARNING_STR_TO_BE_SHOWN = CUDA_ALLOCATOR_ENV_WARNING_STR +del configs +del cuda_allocator_env +del CUDA_ALLOCATOR_ENV_WARNING_STR # pylint: disable=wrong-import-position, wrong-import-order import torch @@ -100,6 +103,11 @@ if torch.cuda.is_available() and not built_with_cuda(): "is installed. Consider reinstalling GraphBolt with CUDA support, see " "installation instructions at https://www.dgl.ai/pages/start.html" ) + +if torch.cuda.is_available() and WARNING_STR_TO_BE_SHOWN is not None: + gb_warning(WARNING_STR_TO_BE_SHOWN) +del WARNING_STR_TO_BE_SHOWN + torch.ops.graphbolt.set_num_io_uring_threads( min((torch.get_num_threads() + 1) // 2, 8) )