Unverified Commit dbb4caaf authored by KrishnanPrash's avatar KrishnanPrash Committed by GitHub
Browse files

docs: Remove `TRTLLM_USE_NIXL_KVCACHE` and `TRTLLM_USE_UCX_KVCACHE` environment variables (#2231)

parent e5fa8b62
......@@ -62,9 +62,3 @@ To enable NIXL for KV cache transfer in disaggregated serving:
See [client](./README.md#client) section to learn how to send the request to deployment.
**Important:** Ensure that ETCD and NATS services are running before starting the service.
\ No newline at end of file
The container will automatically configure the appropriate environment variables (`TRTLLM_USE_NIXL_KVCACHE=1`) when built with the NIXL flag. The same container image can be used to use UCX for KV cache transfer.
```bash
unset TRTLLM_USE_NIXL_KVCACHE
export TRTLLM_USE_UCX_KVCACHE=1
```
\ No newline at end of file
......@@ -22,13 +22,6 @@ if [[ -z ${ENGINE_CONFIG} ]]; then
exit 1
fi
# NOTE: When this script is run directly from srun, the environment variables
# for TRTLLM KV cache are not set. So we need to set them here.
# Related issue: https://github.com/ai-dynamo/dynamo/issues/1743
if [[ -z ${TRTLLM_USE_UCX_KVCACHE} ]] && [[ -z ${TRTLLM_USE_NIXL_KVCACHE} ]]; then
export TRTLLM_USE_UCX_KVCACHE=1
fi
EXTRA_ARGS=""
if [[ -n ${DISAGGREGATION_MODE} ]]; then
EXTRA_ARGS+="--disaggregation-mode ${DISAGGREGATION_MODE} "
......
......@@ -332,26 +332,6 @@ RUN pip install dist/ai_dynamo_runtime*cp312*.whl && \
pip install dist/ai_dynamo*any.whl
ENV DYNAMO_HOME=/workspace
# Use UCX for TRTLLM KV Cache Transfer
ARG TRTLLM_USE_NIXL_KVCACHE_EXPERIMENTAL
ENV TRTLLM_USE_UCX_KVCACHE=1
# Create a script that sets the environment variables and source it.
# If TRTLLM_USE_NIXL_KVCACHE_EXPERIMENTAL is set to 1, unset TRTLLM_USE_UCX_KVCACHE
# and set TRTLLM_USE_NIXL_KVCACHE to 1.
RUN echo '#!/bin/bash' > /usr/local/bin/set_trtllm_env.sh && \
if [ "$TRTLLM_USE_NIXL_KVCACHE_EXPERIMENTAL" = "1" ]; then \
echo 'unset TRTLLM_USE_UCX_KVCACHE' >> /usr/local/bin/set_trtllm_env.sh; \
echo 'export TRTLLM_USE_NIXL_KVCACHE=1' >> /usr/local/bin/set_trtllm_env.sh; \
else \
echo '# TRTLLM_USE_UCX_KVCACHE should already be set to 1' >> /usr/local/bin/set_trtllm_env.sh; \
fi && \
chmod +x /usr/local/bin/set_trtllm_env.sh
# Source the script in bashrc
RUN echo 'source /usr/local/bin/set_trtllm_env.sh' >> /root/.bashrc
# Copy launch banner
RUN --mount=type=bind,source=./container/launch_message.txt,target=/workspace/launch_message.txt \
sed '/^#\s/d' /workspace/launch_message.txt > ~/.launch_screen && \
......@@ -501,11 +481,6 @@ RUN uv pip install --extra-index-url "${TENSORRTLLM_INDEX_URL}" \
"${TENSORRTLLM_PIP_WHEEL}" && \
uv pip install ai-dynamo nixl --find-links wheelhouse
# Setup TRTLLM environment variables, same as in dev image
ENV TRTLLM_USE_UCX_KVCACHE=1
COPY --from=dev /usr/local/bin/set_trtllm_env.sh /usr/local/bin/set_trtllm_env.sh
RUN echo 'source /usr/local/bin/set_trtllm_env.sh' >> /root/.bashrc
# Copy benchmarks, backends and tests for CI
# TODO: Remove this once we have a functional CI image built on top of the runtime image
COPY tests /workspace/tests
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment