pub_soft/dgl
3
0
Fork 0
mirror of https://github.com/dmlc/dgl.git synced 2026-06-03 19:34:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

[GraphBolt][CUDA] Warn only when cuda is available. (#7643)

Browse Source
This commit is contained in:
Muhammed Fatih BALIN
2024-08-06 09:18:06 -04:00
committed by GitHub
parent 90e7edd466
commit 40d32455be
1 changed files with 29 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

50
python/dgl/graphbolt/__init__.py
View File

@@ -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)
)