From 905612b7b43fe5372b990b51f47bc9894a682565 Mon Sep 17 00:00:00 2001
From: Benson Ma <bensonma415@meta.com>
Date: Wed, 8 Jan 2025 13:11:30 -0800
Subject: [PATCH] Set LD_LIBRARY_PATH in install script (#2671)

Summary:
- Set LD_LIBRARY_PATH in the fbgemm_gpu install script to fix the issue with fbgem_gpu being unable to locate libnvrtc.so on start

Pull Request resolved: https://github.com/pytorch/torchrec/pull/2671

Reviewed By: q10, iamzainhuda

Differential Revision: D67949409

Pulled By: TroyGarden

fbshipit-source-id: dabef28755ee123d9560d99848b2764598e9b3a6
---
 .github/scripts/install_fbgemm.sh | 14 ++++++++++----
 1 file changed, 10 insertions(+), 4 deletions(-)

diff --git a/.github/scripts/install_fbgemm.sh b/.github/scripts/install_fbgemm.sh
index 08765d62f..0b0b3e347 100644
--- a/.github/scripts/install_fbgemm.sh
+++ b/.github/scripts/install_fbgemm.sh
@@ -5,11 +5,17 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-echo "CU_VERSION"
-echo "$CU_VERSION"
+echo "CU_VERSION: ${CU_VERSION}"
+echo "CHANNEL: ${CHANNEL}"
+echo "CONDA_ENV: ${CONDA_ENV}"
 
-echo "CHANNEL"
-echo "$CHANNEL"
+if [[ $CU_VERSION = cu* ]]; then
+    # Setting LD_LIBRARY_PATH fixes the runtime error with fbgemm_gpu not
+    # being able to locate libnvrtc.so
+    echo "[NOVA] Setting LD_LIBRARY_PATH ..."
+    conda env config vars set -p ${CONDA_ENV}  \
+        LD_LIBRARY_PATH="/usr/local/lib:${CUDA_HOME}/lib64:${CONDA_ENV}/lib:${LD_LIBRARY_PATH}"
+fi
 
 if [ "$CHANNEL" = "nightly" ]; then
     ${CONDA_RUN} pip install fbgemm-gpu --index-url https://download.pytorch.org/whl/nightly/"$CU_VERSION"