46 lines
2.2 KiB
Bash
Executable File
46 lines
2.2 KiB
Bash
Executable File
#!/bin/bash
|
|
# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
#
|
|
# Framework-specific environment variables and paths
|
|
# Only add paths that exist to avoid cluttering environment
|
|
|
|
# TensorRT-LLM specific variables
|
|
if [ -d /usr/local/tensorrt/targets ]; then
|
|
export TENSORRT_LIB_DIR=/usr/local/tensorrt/targets/$(uname -m)-linux-gnu/lib
|
|
[ -d "$TENSORRT_LIB_DIR" ] && export LD_LIBRARY_PATH="${TENSORRT_LIB_DIR}:${LD_LIBRARY_PATH}"
|
|
fi
|
|
|
|
if [ -d /opt/hpcx/ompi ]; then
|
|
export OPAL_PREFIX=/opt/hpcx/ompi
|
|
export OMPI_MCA_coll_ucc_enable=0
|
|
export PATH="/opt/hpcx/ompi/bin:${PATH}"
|
|
export LD_LIBRARY_PATH="/opt/hpcx/ompi/lib:${LD_LIBRARY_PATH}"
|
|
fi
|
|
|
|
[ -d /opt/hpcx/ucc/lib ] && export LD_LIBRARY_PATH="/opt/hpcx/ucc/lib:${LD_LIBRARY_PATH}"
|
|
[ -f /etc/shinit_v2 ] && export ENV="${ENV:-/etc/shinit_v2}"
|
|
[ -d /usr/local/ucx/bin ] && export PATH="/usr/local/ucx/bin:${PATH}"
|
|
[ -d /usr/local/cuda/bin ] && export PATH="/usr/local/cuda/bin:${PATH}"
|
|
[ -d /usr/local/cuda/nvvm/bin ] && export PATH="/usr/local/cuda/nvvm/bin:${PATH}"
|
|
|
|
# vLLM nvshmem
|
|
[ -d /opt/vllm/tools/ep_kernels/ep_kernels_workspace/nvshmem_install/lib ] && \
|
|
export LD_LIBRARY_PATH="/opt/vllm/tools/ep_kernels/ep_kernels_workspace/nvshmem_install/lib:${LD_LIBRARY_PATH}"
|
|
|
|
# System nvshmem (TRT-LLM)
|
|
ARCH_ALT=$(uname -m | sed 's/aarch64/aarch64/;s/x86_64/x86_64/')
|
|
[ -d "/usr/lib/${ARCH_ALT}-linux-gnu/nvshmem/13" ] && \
|
|
export LD_LIBRARY_PATH="/usr/lib/${ARCH_ALT}-linux-gnu/nvshmem/13:${LD_LIBRARY_PATH}"
|
|
|
|
# PyTorch libraries (TRT-LLM)
|
|
# PYTHON_VERSION should be set via ENV in container; fail early if missing
|
|
if [ -z "${PYTHON_VERSION}" ]; then
|
|
echo "WARNING: PYTHON_VERSION not set, defaulting to 3.12" >&2
|
|
PYTHON_VERSION=3.12
|
|
fi
|
|
[ -d "/opt/dynamo/venv/lib/python${PYTHON_VERSION}/site-packages/torch/lib" ] && \
|
|
export LD_LIBRARY_PATH="/opt/dynamo/venv/lib/python${PYTHON_VERSION}/site-packages/torch/lib:${LD_LIBRARY_PATH}"
|
|
[ -d "/opt/dynamo/venv/lib/python${PYTHON_VERSION}/site-packages/torch_tensorrt/lib" ] && \
|
|
export LD_LIBRARY_PATH="/opt/dynamo/venv/lib/python${PYTHON_VERSION}/site-packages/torch_tensorrt/lib:${LD_LIBRARY_PATH}"
|