[TRTLLMINF-336][infra] BOLT profile merge: skip runtime tensorrt import gate (libs-only) (#18395)

Signed-off-by: Matt Lefebvre <mlefebvre@nvidia.com>

Author: mlefeb01

scripts/bolt/internal/slurm_merge.sh

@@ -117,7 +117,11 @@ srun --ntasks=1 --ntasks-per-node=1 --nodes=1 \
         #    No instrument -- the collect jobs already produced the .fdata.
         _extract='"$RUNDIR"'/extract
         mkdir -p "$_extract" && tar xzf /builds/'"$TARBALL_NAME"' -C "$_extract" 2>/dev/null
-        /workspace/bolt/setup_env.sh "$_extract"
+        # Libs-only: the merge only needs the wheel ELFs on disk to convert
+        # .fdata -> .yaml (no runtime import), so skip setup_env.sh runtime
+        # tensorrt gate -- the profiling base image has no importable tensorrt
+        # before install and would otherwise fail here.
+        BOLT_SETUP_LIBS_ONLY=1 /workspace/bolt/setup_env.sh "$_extract"
         source /workspace/bolt/bolt_lib.sh
         bolt_run_stages setup_directories backup_libraries
 

scripts/bolt/setup_env.sh

@@ -25,17 +25,32 @@ set -euo pipefail
 EXTRACT="${1:?usage: setup_env.sh <dir containing TensorRT-LLM/ (wheel)>}"
 PYTHON="${PYTHON:-python3}"
 
+# BOLT_SETUP_LIBS_ONLY=1: the caller only needs the wheel's ELF .so files ON DISK
+# (e.g. the merge job reconstructs the pre-instrument originals to convert
+# .fdata -> .yaml -- see internal/slurm_merge.sh). That path never imports
+# tensorrt / tensorrt_llm at runtime (library discovery uses
+# importlib.util.find_spec, which does NOT execute the package), so the
+# runtime-tensorrt gate below is irrelevant and would spuriously fail on a
+# profiling base image whose `import tensorrt` isn't wired up before install.
+# Skip the import checks in that mode but still install the wheel (--no-deps).
+LIBS_ONLY="${BOLT_SETUP_LIBS_ONLY:-0}"
+
 WHEEL=$(ls "$EXTRACT"/TensorRT-LLM/tensorrt_llm-*.whl 2>/dev/null | head -1 || true)
 [[ -f "$WHEEL" ]] || { echo "[ERROR] No tensorrt_llm wheel under $EXTRACT/TensorRT-LLM/"; exit 1; }
 
 # 0. The container must already have a working tensorrt (runtime-capable image).
-if ! "$PYTHON" -c "import tensorrt" 2>/dev/null; then
-    echo "[ERROR] 'import tensorrt' fails in this container BEFORE install." >&2
-    echo "        Use a runtime-capable image (devel/release) with TensorRT intact;" >&2
-    echo "        a bare CUDA/pytorch base or an image with a stripped tensorrt will not work." >&2
-    exit 1
+#    Skipped in libs-only mode (no runtime import happens downstream).
+if [[ "$LIBS_ONLY" != "1" ]]; then
+    if ! "$PYTHON" -c "import tensorrt" 2>/dev/null; then
+        echo "[ERROR] 'import tensorrt' fails in this container BEFORE install." >&2
+        echo "        Use a runtime-capable image (devel/release) with TensorRT intact;" >&2
+        echo "        a bare CUDA/pytorch base or an image with a stripped tensorrt will not work." >&2
+        exit 1
+    fi
+    echo "[INFO] tensorrt OK: $("$PYTHON" -c 'import tensorrt; print(tensorrt.__version__)')"
+else
+    echo "[INFO] BOLT_SETUP_LIBS_ONLY=1: skipping runtime tensorrt import checks (libs-only)."
 fi
-echo "[INFO] tensorrt OK: $("$PYTHON" -c 'import tensorrt; print(tensorrt.__version__)')"
 
 # 1. Wheel WITHOUT deps -> never touches tensorrt. --force-reinstall so the BOLT
 #    libs overlay any preinstalled trtllm (release image).
@@ -44,5 +59,9 @@ pip install --no-deps --force-reinstall "$WHEEL"
 
 # 2. Verify from a neutral cwd so the extracted source tree doesn't shadow the
 #    installed package (a source `tensorrt_llm/` has no compiled bindings).
-( cd /tmp && "$PYTHON" -c "import tensorrt, tensorrt_llm; print('[INFO] runtime + trtllm ok:', tensorrt_llm.__file__)" )
+#    Skipped in libs-only mode: only the on-disk .so files are needed, and the
+#    runtime import may legitimately fail on a libs-only base image.
+if [[ "$LIBS_ONLY" != "1" ]]; then
+    ( cd /tmp && "$PYTHON" -c "import tensorrt, tensorrt_llm; print('[INFO] runtime + trtllm ok:', tensorrt_llm.__file__)" )
+fi
 echo "[SUCCESS] Environment ready for BOLT flow."