Skip to content

Commit b7b6982

Browse files
apboseapbose
authored
setup_nccl_for_torch_tensorrt change (#4303)
1 parent 32fb12b commit b7b6982

2 files changed

Lines changed: 98 additions & 118 deletions

File tree

‎py/torch_tensorrt/distributed/_nccl_utils.py‎

Lines changed: 98 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -28,9 +28,9 @@
2828
symlink workarounds.
2929
"""
3030

31-
import ctypes
3231
import logging
3332
import os
33+
import subprocess
3434
from typing import Optional
3535

3636
import torch.distributed as dist
@@ -101,88 +101,122 @@ def ensure_nccl_symlink(nccl_lib_dir: str) -> bool:
101101
return False
102102

103103

104-
def check_nccl_library_path() -> bool:
105-
"""
106-
Check if LD_LIBRARY_PATH includes PyTorch's NCCL directory.
104+
def _sys_libdir_on_ldso_path() -> str:
105+
"""Pick a system library directory that ld.so searches by default.
107106
108-
Returns:
109-
True if configuration is correct, False if LD_LIBRARY_PATH needs updating.
107+
Returns the first existing directory from a portability-ordered list:
108+
Debian/Ubuntu x86_64 multiarch → ARM64 multiarch → RHEL/CentOS lib64 →
109+
bare /usr/lib (always on ld.so's search path as a final fallback).
110110
"""
111-
nccl_lib_dir = get_nccl_library_path()
112-
113-
if nccl_lib_dir is None:
114-
# System NCCL - no action needed
115-
return True
116-
117-
ld_library_path = os.environ.get("LD_LIBRARY_PATH", "")
118-
return nccl_lib_dir in ld_library_path
111+
for d in (
112+
"/usr/lib/x86_64-linux-gnu", # Debian / Ubuntu x86_64
113+
"/usr/lib/aarch64-linux-gnu", # Debian / Ubuntu ARM64 (Jetson)
114+
"/usr/lib64", # RHEL / CentOS / Fedora x86_64
115+
):
116+
if os.path.isdir(d):
117+
return d
118+
return "/usr/lib"
119119

120120

121121
def setup_nccl_for_torch_tensorrt() -> None:
122122
"""
123-
Setup NCCL library for TensorRT distributed inference.
124-
125-
This function:
126-
1. Detects if nvidia.nccl pip package is installed
127-
2. Creates libnccl.so symlink if needed
128-
3. Pre-loads libnccl.so via ctypes (helps Python runtime path)
129-
4. Updates LD_LIBRARY_PATH for dynamic loaders
130-
131-
Note: TRT's internal loader (libLoader.cpp) reads LD_LIBRARY_PATH at
132-
process launch time, not when updated via os.environ. For the C++ TRT
133-
runtime path, LD_LIBRARY_PATH must be set before the process starts:
134-
135-
NCCL_LIB=$(python -c "from torch_tensorrt.distributed._nccl_utils import get_nccl_library_path; print(get_nccl_library_path())")
136-
LD_LIBRARY_PATH="$NCCL_LIB:$LD_LIBRARY_PATH" python script.py
137-
138-
For NGC containers (system NCCL), this is a no-op.
123+
Point a `libnccl.so` symlink on ld.so's default search path at PyTorch's
124+
libnccl.so.2 so TRT and PyTorch share a single NCCL library in the process.
125+
126+
What this function does:
127+
1. Locate the nvidia.nccl pip package's libnccl.so.2 via
128+
get_nccl_library_path(). If pip's nccl isn't installed (NGC /
129+
system-NCCL environments) returns immediately — no action needed.
130+
2. Pick a system library directory that ld.so already searches by
131+
default via _sys_libdir_on_ldso_path() (Debian/Ubuntu multiarch,
132+
RHEL lib64, or /usr/lib fallback).
133+
3. If <sys_libdir>/libnccl.so already points at that libnccl.so.2,
134+
return.
135+
4. Otherwise, atomically install a fresh symlink:
136+
<sys_libdir>/libnccl.so → <pip>/libnccl.so.2
137+
via "symlink to a unique per-pid temp name, then os.replace onto
138+
the target." This is multi-process safe: when several ranks of a
139+
distributed test call this function concurrently, none of them
140+
crash on FileExistsError, and the final on-disk state is the same
141+
regardless of execution order.
142+
5. Run `ldconfig` to refresh /etc/ld.so.cache.
143+
6. Guarded by a module-global flag so subsequent calls in the same
144+
process are a no-op.
145+
146+
Requires write access to the chosen sys_libdir (root inside Docker is
147+
the common case). On OSError the function raises RuntimeError with
148+
documented LD_PRELOAD / LD_LIBRARY_PATH workarounds for non-root setups.
139149
"""
140150
global _nccl_setup_checked
141-
142-
# Only check once per process
143151
if _nccl_setup_checked:
144152
return
145153
_nccl_setup_checked = True
146154

147155
nccl_lib_dir = get_nccl_library_path()
148-
149156
if nccl_lib_dir is None:
150-
# NGC container or system NCCL - no action needed
151157
logger.debug(
152-
"nvidia.nccl package not found. "
153-
"Assuming system NCCL is used by both PyTorch and TensorRT."
158+
"nvidia.nccl package not found; assuming system NCCL is shared by PyTorch and TensorRT."
154159
)
155160
return
156161

157-
logger.debug(f"Found nvidia.nccl package at: {nccl_lib_dir}")
158-
159-
# Ensure symlink exists
160-
symlink_ok = ensure_nccl_symlink(nccl_lib_dir)
161-
162-
# Ensure LD_LIBRARY_PATH includes the NCCL directory so TRT's dlopen("libnccl.so")
163-
# finds the same library PyTorch already loaded. dlopen() reads LD_LIBRARY_PATH
164-
# dynamically, so updating os.environ here takes effect for subsequent loads.
165-
ld_library_path = os.environ.get("LD_LIBRARY_PATH", "")
166-
if nccl_lib_dir not in ld_library_path:
167-
os.environ["LD_LIBRARY_PATH"] = (
168-
f"{nccl_lib_dir}:{ld_library_path}" if ld_library_path else nccl_lib_dir
162+
nccl_so_2 = os.path.join(nccl_lib_dir, "libnccl.so.2")
163+
if not os.path.isfile(nccl_so_2):
164+
logger.warning(
165+
f"Expected {nccl_so_2} to exist but it doesn't; skipping NCCL setup."
169166
)
170-
logger.debug(f"Added NCCL directory to LD_LIBRARY_PATH: {nccl_lib_dir}")
171-
else:
172-
logger.debug(f"LD_LIBRARY_PATH already includes NCCL directory: {nccl_lib_dir}")
167+
return
173168

174-
if symlink_ok:
175-
# Pre-load libnccl.so into the process with RTLD_GLOBAL so that TRT's
176-
# subsequent dlopen("libnccl.so") inside setCommunicator() finds the
177-
# already-loaded library rather than searching LD_LIBRARY_PATH again.
178-
nccl_so = os.path.join(nccl_lib_dir, "libnccl.so")
179-
try:
180-
ctypes.CDLL(nccl_so, mode=ctypes.RTLD_GLOBAL)
181-
logger.debug(f"Pre-loaded NCCL library: {nccl_so}")
182-
except OSError as e:
183-
logger.warning(f"Failed to pre-load NCCL library {nccl_so}: {e}")
169+
sys_libdir = _sys_libdir_on_ldso_path()
170+
target = os.path.join(sys_libdir, "libnccl.so")
184171

185-
logger.debug("NCCL library setup complete")
172+
# Fast path: already set up by a prior process or rank.
173+
if os.path.islink(target) and os.readlink(target) == nccl_so_2:
174+
logger.debug(f"{target} already points at {nccl_so_2}; nothing to do.")
175+
return
176+
177+
# Race-safe symlink swap. Multiple ranks may enter this function
178+
# concurrently (e.g. MultiProcessTestCase forks 2 children that each call
179+
# setup_nccl_for_torch_tensorrt simultaneously). Using `os.remove` +
180+
# `os.symlink` opens a window where one rank's symlink call races with
181+
# another's, raising FileExistsError on the loser.
182+
#
183+
# Instead: create the new symlink under a unique per-pid temp name
184+
# (no contention possible — different filenames), then atomically rename
185+
# it onto `target`. `os.replace` is a single POSIX rename(2) call:
186+
# it overwrites unconditionally and is observable as a single transition,
187+
# never a missing or half-written state. All ranks converge on the same
188+
# final symlink without any of them crashing.
189+
tmp = f"{target}.torchtrt-{os.getpid()}"
190+
try:
191+
if os.path.lexists(tmp):
192+
os.remove(tmp)
193+
os.symlink(nccl_so_2, tmp)
194+
os.replace(tmp, target)
195+
subprocess.run(["ldconfig"], check=False)
196+
logger.info(
197+
f"NCCL: linked {target} -> {nccl_so_2} so TRT and PyTorch share one libnccl."
198+
)
199+
except OSError as e:
200+
# Clean up our temp link if we left one behind.
201+
if os.path.lexists(tmp):
202+
try:
203+
os.remove(tmp)
204+
except OSError:
205+
pass
206+
# If another rank already produced the correct final symlink while
207+
# we were failing, accept that as success — the end state we wanted
208+
# is in place.
209+
if os.path.islink(target) and os.readlink(target) == nccl_so_2:
210+
return
211+
raise RuntimeError(
212+
f"setup_nccl_for_torch_tensorrt(): cannot write {target} "
213+
f"(needed so TRT's dlopen('libnccl.so') resolves to PyTorch's libnccl.so.2). "
214+
f"Workarounds without root: relaunch python with "
215+
f"LD_PRELOAD={nccl_so_2} ; or pre-set "
216+
f"LD_LIBRARY_PATH={nccl_lib_dir}:$LD_LIBRARY_PATH before python starts "
217+
f"(and create a libnccl.so symlink in that dir first). "
218+
f"Original error: {e}"
219+
) from e
186220

187221

188222
def initialize_nccl_comm(device: Optional[int] = None) -> None:
@@ -253,29 +287,11 @@ def initialize_nccl_comm(device: Optional[int] = None) -> None:
253287

254288

255289
def check_nccl_engine_requirements() -> None:
256-
"""Warn if an requires_native_multidevice TRT engine's NCCL prerequisites are not satisfied.
257-
258-
Checks two conditions and logs a warning for each:
259-
1. LD_LIBRARY_PATH does not include PyTorch's NCCL lib dir (too late to fix,
260-
must be set before process launch — use torchtrtrun).
261-
2. torch.distributed is not initialized or world_size == 1.
290+
"""Warn if a requires_native_multidevice TRT engine's NCCL prerequisites are not satisfied.
262291
263-
Call this from both TorchTensorRTModule and PythonTorchTensorRTModule after
292+
Called from TorchTensorRTModule and PythonTorchTensorRTModule after
264293
confirming the engine has NCCL collective ops.
265294
"""
266-
if get_nccl_library_path() is not None and not check_nccl_library_path():
267-
logger.warning(
268-
"This TRT engine contains NCCL collective ops but "
269-
"LD_LIBRARY_PATH does not include PyTorch's NCCL library directory. "
270-
"TRT may load a different NCCL instance than PyTorch, causing "
271-
"communicator sharing to fail. Use torchtrtrun to launch distributed "
272-
"scripts, or set LD_PRELOAD and LD_LIBRARY_PATH before process start:\n"
273-
" NCCL_LIB=$(python -c 'from torch_tensorrt.distributed._nccl_utils "
274-
"import get_nccl_library_path; print(get_nccl_library_path())')\n"
275-
" LD_PRELOAD=$NCCL_LIB/libnccl.so.2 "
276-
"LD_LIBRARY_PATH=$NCCL_LIB:$LD_LIBRARY_PATH python ..."
277-
)
278-
279295
if not (
280296
dist.is_available() and dist.is_initialized() and dist.get_world_size() > 1
281297
):

‎tests/py/dynamo/distributed/test_native_nccl.py‎

Lines changed: 0 additions & 36 deletions
Original file line numberDiff line numberDiff line change
@@ -713,22 +713,6 @@ def test_get_nccl_library_path_returns_none_or_string(self) -> None:
713713
f"libnccl.so.2 not found in {result}",
714714
)
715715

716-
def test_check_nccl_library_path_system_nccl(self) -> None:
717-
"""check_nccl_library_path returns True when nvidia.nccl not installed."""
718-
from torch_tensorrt.distributed._nccl_utils import (
719-
check_nccl_library_path,
720-
get_nccl_library_path,
721-
)
722-
723-
nccl_lib_dir = get_nccl_library_path()
724-
if nccl_lib_dir is None:
725-
# System NCCL path — must return True
726-
self.assertTrue(check_nccl_library_path())
727-
else:
728-
# nvidia.nccl installed — result depends on LD_LIBRARY_PATH
729-
result = check_nccl_library_path()
730-
self.assertIsInstance(result, bool)
731-
732716
def test_setup_nccl_for_torch_tensorrt_idempotent(self) -> None:
733717
"""Calling setup_nccl_for_torch_tensorrt() multiple times is safe."""
734718
from torch_tensorrt.distributed import _nccl_utils
@@ -749,26 +733,6 @@ def test_ensure_nccl_symlink_nonexistent_dir(self) -> None:
749733
# libnccl.so.2 doesn't exist there → returns False
750734
self.assertFalse(result)
751735

752-
def test_check_nccl_library_path_detects_missing_ld_path(self) -> None:
753-
"""check_nccl_library_path returns False when LD_LIBRARY_PATH is absent."""
754-
from torch_tensorrt.distributed._nccl_utils import get_nccl_library_path
755-
756-
nccl_lib_dir = get_nccl_library_path()
757-
if nccl_lib_dir is None:
758-
self.skipTest("nvidia.nccl not installed; system NCCL path is always OK")
759-
760-
from torch_tensorrt.distributed._nccl_utils import check_nccl_library_path
761-
762-
original = os.environ.get("LD_LIBRARY_PATH", "")
763-
# Remove nccl_lib_dir from LD_LIBRARY_PATH
764-
paths = [p for p in original.split(":") if p and p != nccl_lib_dir]
765-
os.environ["LD_LIBRARY_PATH"] = ":".join(paths)
766-
try:
767-
result = check_nccl_library_path()
768-
self.assertFalse(result)
769-
finally:
770-
os.environ["LD_LIBRARY_PATH"] = original
771-
772736

773737
# ============================================================================
774738
# Section 4 — fuse_distributed_ops graph pass (no GPU, no dist) [was Section 3]

0 commit comments

Comments
 (0)