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
...@@ -61,10 +61,4 @@ To enable NIXL for KV cache transfer in disaggregated serving: ...@@ -61,10 +61,4 @@ To enable NIXL for KV cache transfer in disaggregated serving:
4. **Send the request:** 4. **Send the request:**
See [client](./README.md#client) section to learn how to send the request to deployment. 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. **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 ...@@ -22,13 +22,6 @@ if [[ -z ${ENGINE_CONFIG} ]]; then
exit 1 exit 1
fi 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="" EXTRA_ARGS=""
if [[ -n ${DISAGGREGATION_MODE} ]]; then if [[ -n ${DISAGGREGATION_MODE} ]]; then
EXTRA_ARGS+="--disaggregation-mode ${DISAGGREGATION_MODE} " EXTRA_ARGS+="--disaggregation-mode ${DISAGGREGATION_MODE} "
......
...@@ -332,26 +332,6 @@ RUN pip install dist/ai_dynamo_runtime*cp312*.whl && \ ...@@ -332,26 +332,6 @@ RUN pip install dist/ai_dynamo_runtime*cp312*.whl && \
pip install dist/ai_dynamo*any.whl pip install dist/ai_dynamo*any.whl
ENV DYNAMO_HOME=/workspace 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 # Copy launch banner
RUN --mount=type=bind,source=./container/launch_message.txt,target=/workspace/launch_message.txt \ RUN --mount=type=bind,source=./container/launch_message.txt,target=/workspace/launch_message.txt \
sed '/^#\s/d' /workspace/launch_message.txt > ~/.launch_screen && \ sed '/^#\s/d' /workspace/launch_message.txt > ~/.launch_screen && \
...@@ -501,11 +481,6 @@ RUN uv pip install --extra-index-url "${TENSORRTLLM_INDEX_URL}" \ ...@@ -501,11 +481,6 @@ RUN uv pip install --extra-index-url "${TENSORRTLLM_INDEX_URL}" \
"${TENSORRTLLM_PIP_WHEEL}" && \ "${TENSORRTLLM_PIP_WHEEL}" && \
uv pip install ai-dynamo nixl --find-links wheelhouse 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 # Copy benchmarks, backends and tests for CI
# TODO: Remove this once we have a functional CI image built on top of the runtime image # TODO: Remove this once we have a functional CI image built on top of the runtime image
COPY tests /workspace/tests 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