build: support custom TRTLLM build for commits not on main branch (#2021)

container/Dockerfile.tensorrt_llm

@@ -151,8 +151,7 @@ RUN [ -f /etc/pip/constraint.txt ] && : > /etc/pip/constraint.txt || true && \
         fi; \
     else \
          # Install TensorRT-LLM wheel from the provided index URL, allow dependencies from PyPI
-         pip install --index-url "${TENSORRTLLM_INDEX_URL}" \
-         --extra-index-url https://pypi.org/simple \
+         pip install --extra-index-url "${TENSORRTLLM_INDEX_URL}" \
          "${TENSORRTLLM_PIP_WHEEL}" ; \
     fi
 
@@ -466,8 +465,10 @@ ARG TENSORRTLLM_INDEX_URL="https://pypi.python.org/simple"
 COPY --from=dev /workspace/wheels/nixl/*.whl wheelhouse/
 COPY --from=wheel_builder /workspace/dist/*.whl wheelhouse/
 
-RUN uv pip install --index-url "${TENSORRTLLM_INDEX_URL}" \
-    --extra-index-url https://pypi.org/simple \
+# NOTE: If a package (tensorrt_llm) exists on both --index-url and --extra-index-url,
+# uv will prioritize the --extra-index-url, unless --index-strategy unsafe-best-match
+# is also specified. So set the configurable index as a --extra-index-url for prioritization.
+RUN uv pip install --extra-index-url "${TENSORRTLLM_INDEX_URL}" \
     "${TENSORRTLLM_PIP_WHEEL}" && \
     uv pip install ai-dynamo nixl --find-links wheelhouse
 

container/build_trtllm_wheel.sh

@@ -44,7 +44,7 @@ MAIN_DIR=$(dirname "$(readlink -f "$0")")
 (cd /tmp && \
 # Clone the TensorRT-LLM repository.
 if [ ! -d "TensorRT-LLM" ]; then
-  git clone --single-branch --branch main https://github.com/NVIDIA/TensorRT-LLM.git
+  git clone https://github.com/NVIDIA/TensorRT-LLM.git
 fi
 
 cd TensorRT-LLM
@@ -120,4 +120,4 @@ rm -rf $OUTPUT_DIR/commit.txt
 echo ${ARCH}_${TRTLLM_COMMIT} > $OUTPUT_DIR/commit.txt
 
 echo "TRT-LLM wheel built successfully."
-ls -al $OUTPUT_DIR
+ls -al $OUTPUT_DIR
\ No newline at end of file