diff --git a/projects/rccl-tests/README.md b/projects/rccl-tests/README.md index 2258db761fb..f81fb75b43c 100644 --- a/projects/rccl-tests/README.md +++ b/projects/rccl-tests/README.md @@ -125,6 +125,20 @@ For MPI (using MPICH), you need to use a command similar to the following: mpirun.mpich -np 8 -env NCCL_DEBUG=VERSION -env HSA_NO_SCRATCH_RECLAIM=1 ./build/all_reduce_perf -b 8M -e 128M -i 8388608 -g 1 -d bfloat16 ``` +#### Socket encryption (TLS) + +`NCCL_TESTS_TLS_PSK=` makes every process call `ncclSetEncryption()` with this pre-shared key +before the communicator is created, so RCCL's bootstrap, proxy and socket-transport TCP traffic is +TLS-encrypted. The key must be at least 32 bytes and identical on all ranks, and RCCL must be +built with `-DTLS_BACKEND=OPENSSL3` (otherwise communicator creation fails instead of falling back +to plaintext). Requires RCCL 2.32 or newer, both at build time and at run time; with older +headers or an older library the test exits with an error. The variable must reach every rank: +if only some ranks have the key, bootstrap hangs instead of failing. + +```shell +NCCL_TESTS_TLS_PSK= mpirun -np 8 -x NCCL_TESTS_TLS_PSK ./build/all_reduce_perf -b 8 -e 128M -f 2 +``` + #### Grouped-broadcast (AllGatherV) mode `RCCL_TESTS_ALLGATHERV_ENABLE=1` makes `broadcast_perf` issue `nranks` broadcasts with distinct diff --git a/projects/rccl-tests/src/CMakeLists.txt b/projects/rccl-tests/src/CMakeLists.txt index d0c5e1b3a26..92a09e31652 100644 --- a/projects/rccl-tests/src/CMakeLists.txt +++ b/projects/rccl-tests/src/CMakeLists.txt @@ -136,6 +136,11 @@ add_custom_target(hipify DEPENDS ${HIP_COMMON_SOURCES}) add_library(rccl_common OBJECT ${HIP_COMMON_SOURCES}) add_dependencies(rccl_common hipify git_version_check) target_link_libraries(rccl_common roc::rccl hip::device Threads::Threads dl nccl_test_os nccl_tests_options verifiable) +# dlsym(RTLD_DEFAULT) cannot see symbols of a statically linked RCCL. +get_target_property(RCCL_LIBRARY_TYPE roc::rccl TYPE) +if(RCCL_LIBRARY_TYPE STREQUAL "STATIC_LIBRARY") + target_compile_definitions(rccl_common PRIVATE RCCL_TESTS_STATIC_RCCL) +endif() if(USE_MPI) target_link_libraries(rccl_common MPI::MPI_CXX) endif() diff --git a/projects/rccl-tests/src/common.cu b/projects/rccl-tests/src/common.cu index 2fd8f695d8a..60cc50e6570 100644 --- a/projects/rccl-tests/src/common.cu +++ b/projects/rccl-tests/src/common.cu @@ -2539,6 +2539,37 @@ testResult_t run() { if (proc == 0) printf("#\n# Reducing maxBytes to %ld due to memory limitation\n", maxBytes); } + const char* tlsPsk = getenv(NCCL_TESTS_TLS_PSK_ENV); + if (tlsPsk) { +#if NCCL_VERSION_CODE >= NCCL_VERSION(2,32,0) + typedef ncclResult_t (*ncclSetEncryption_t)(const ncclEncryptionConfig_t*); +#ifdef RCCL_TESTS_STATIC_RCCL + ncclSetEncryption_t setEncryption = ncclSetEncryption; +#else + // Looked up at run time: a direct reference would keep the binary from loading against + // an older librccl when it is linked with -z now or run with LD_BIND_NOW. + ncclSetEncryption_t setEncryption = (ncclSetEncryption_t) dlsym(RTLD_DEFAULT, "ncclSetEncryption"); + if (setEncryption == NULL) { + fprintf(stderr, NCCL_TESTS_TLS_PSK_ENV " is set, but the RCCL library (%d) has no ncclSetEncryption; 2.32 or newer is required.\n", + test_ncclVersion); + return testInvalidUsage; + } +#endif + ncclEncryptionConfig_t encryption = NCCL_ENCRYPTION_CONFIG_INITIALIZER; + encryption.mode = NCCL_ENCRYPTION_MODE_PSK; + encryption.psk = tlsPsk; + NCCLCHECK(setEncryption(&encryption)); + // RCCL keeps its own copy of the key. Dropping the variable keeps it out of RCCL's RAS + // environment report, which collects every NCCL_* entry, and out of child processes. + unsetenv(NCCL_TESTS_TLS_PSK_ENV); + if (proc == 0) printf("# Socket encryption: TLS (PSK)\n"); +#else + fprintf(stderr, NCCL_TESTS_TLS_PSK_ENV " is set, but rccl-tests was built against RCCL %d headers; 2.32 or newer is required.\n", + NCCL_VERSION_CODE); + return testInvalidUsage; +#endif + } + ncclUniqueId ncclId; if (ncclProc == 0) { NCCLCHECK(ncclGetUniqueId(&ncclId)); diff --git a/projects/rccl-tests/src/common.h b/projects/rccl-tests/src/common.h index 01247e60bc3..3a7fdf61ab8 100644 --- a/projects/rccl-tests/src/common.h +++ b/projects/rccl-tests/src/common.h @@ -9,6 +9,8 @@ #define __COMMON_H__ #define NCCL_TESTS_VERSION "2.19.6" +// Socket encryption pre-shared key; util.cu redacts it from the JSON report by this name. +#define NCCL_TESTS_TLS_PSK_ENV "NCCL_TESTS_TLS_PSK" #include "rccl/rccl.h" // nccl_device.h provides the device-API public types referenced below diff --git a/projects/rccl-tests/src/util.cu b/projects/rccl-tests/src/util.cu index a02ab0d9722..0f7e8b9a3e0 100644 --- a/projects/rccl-tests/src/util.cu +++ b/projects/rccl-tests/src/util.cu @@ -373,7 +373,12 @@ void jsonOutputInit(const char *in_path, jsonKey("env"); jsonStartList(); for(char **e = envp; *e; e++) { - jsonStr(*e); + // Reports get shared; the TLS pre-shared key must not end up in them. + if (strncmp(*e, NCCL_TESTS_TLS_PSK_ENV "=", strlen(NCCL_TESTS_TLS_PSK_ENV "=")) == 0) { + jsonStr(NCCL_TESTS_TLS_PSK_ENV "="); + } else { + jsonStr(*e); + } } jsonFinishList(); jsonKey("nccl_version"); jsonInt(test_ncclVersion);