diff --git a/.github/scripts/check_lmdeploy.py b/.github/scripts/check_lmdeploy.py index 808fe7c4b3..e46a131670 100644 --- a/.github/scripts/check_lmdeploy.py +++ b/.github/scripts/check_lmdeploy.py @@ -20,6 +20,8 @@ def check_module_init(root: str): continue elif d.startswith('lmdeploy/monitoring'): continue + elif d.startswith('lmdeploy/router'): + continue elif d.startswith('lmdeploy/serve/turbomind/triton_models'): continue elif d.startswith('lmdeploy/serve/turbomind/triton_python_backend'): diff --git a/.github/workflows/pr_ete_test.yml b/.github/workflows/pr_ete_test.yml index b2943f8be7..5f94e92ca8 100644 --- a/.github/workflows/pr_ete_test.yml +++ b/.github/workflows/pr_ete_test.yml @@ -48,12 +48,15 @@ jobs: steps: - name: Clone repository uses: actions/checkout@v6 + - name: Install Rust toolchain + uses: dtolnay/rust-toolchain@stable - name: Install lmdeploy run: | python3 -m pip install --upgrade -r requirements/runtime_cuda.txt python3 -m pip install -r requirements/lite.txt python3 -m pip install -r requirements/test.txt python3 -m pip install -e . + python3 -m pip install -e lmdeploy/router - name: Check env run: | python3 -m pip list @@ -65,147 +68,290 @@ jobs: run: | pytest autotest -m 'pr_test and gpu_num_2' -x --alluredir=${{env.REPORT_DIR}} --clean-alluredir pytest autotest -m 'pr_test and gpu_num_1' -n 2 -x --alluredir=${{env.REPORT_DIR}} + - name: Start lmdeploy-router + run: | + lmdeploy-router --host 127.0.0.1 --port 23333 --policy cache_aware > ${{env.SERVER_LOG}}/router.log 2>&1 & + router_pid=$! + for i in $(seq 1 30) + do + sleep 1 + if curl -f -s http://127.0.0.1:23333/health > /dev/null 2>&1; then + echo "lmdeploy-router is ready" + exit 0 + fi + done + echo "lmdeploy-router failed to start" + cat ${{env.SERVER_LOG}}/router.log + exit 1 - name: Test restful server - pytorch Qwen3.5-35B-A3B run: | - CUDA_VISIBLE_DEVICES=6,7 lmdeploy serve api_server /nvme/qa_test_models/Qwen/Qwen3.5-35B-A3B --tp 2 --backend pytorch --logprobs-mode raw_logprobs --enable-return-routed-experts --reasoning-parser qwen-qwq --tool-call-parser qwen3coder --speculative-algorithm qwen3_5_mtp --speculative-num-draft-tokens 4 --max-batch-size 256 --allow-terminate-by-client > ${{env.SERVER_LOG}}/pytorch_Qwen3.5-35B-A3B_start_restful.log 2>&1 & + CUDA_VISIBLE_DEVICES=6,7 lmdeploy serve api_server /nvme/qa_test_models/Qwen/Qwen3.5-35B-A3B --server-name 127.0.0.1 --server-port 23401 --tp 2 --backend pytorch --logprobs-mode raw_logprobs --enable-return-routed-experts --reasoning-parser qwen-qwq --tool-call-parser qwen3coder --speculative-algorithm qwen3_5_mtp --speculative-num-draft-tokens 4 --max-batch-size 256 --allow-terminate-by-client --proxy-url http://127.0.0.1:23333 > ${{env.SERVER_LOG}}/pytorch_Qwen3.5-35B-A3B_start_restful.log 2>&1 & echo "restful_pid=$!" for i in $(seq 1 180) do sleep 5 echo "health check try $i" - if curl -f -s http://127.0.0.1:23333/health > /dev/null 2>&1; then + if curl -f -s http://127.0.0.1:23401/health > /dev/null 2>&1 && \ + curl -f -s http://127.0.0.1:23333/nodes/status | grep -q 'http://127.0.0.1:23401'; then pytest autotest/interface/restful/test_restful_chat_completions_v1.py -n 20 -k 'Qwen/Qwen3.5-35B-A3B and pytorch and not Qwen/Qwen3.5-35B-A3B-FP8' -m 'not not_pytorch and pr_test' --alluredir=${{env.REPORT_DIR}} pytest autotest/interface/restful/test_restful_generate.py -n 20 -k 'Qwen/Qwen3.5-35B-A3B and pytorch and not Qwen/Qwen3.5-35B-A3B-FP8' -m 'not not_pytorch' --alluredir=${{env.REPORT_DIR}} pytest autotest/interface/restful/tool_parser/test_tool_call_basic.py -n 20 -k 'Qwen/Qwen3.5-35B-A3B and pytorch and not Qwen/Qwen3.5-35B-A3B-FP8' --alluredir=${{env.REPORT_DIR}} - curl -f -s http://127.0.0.1:23333/terminate > /dev/null 2>&1 + curl -f -s http://127.0.0.1:23333/nodes/remove -H 'Content-Type: application/json' -d '{"url":"http://127.0.0.1:23401"}' > /dev/null 2>&1 || true + curl -f -s http://127.0.0.1:23401/terminate > /dev/null 2>&1 + wait $restful_pid exit 0 fi done echo "health check fail" - curl -f -s http://127.0.0.1:23333/terminate > /dev/null 2>&1 + curl -f -s http://127.0.0.1:23401/terminate > /dev/null 2>&1 + wait $restful_pid cat ${{env.SERVER_LOG}}/pytorch_Qwen3.5-35B-A3B_start_restful.log exit 1 + - name: Start lmdeploy-router for PD + run: | + lmdeploy-router --host 127.0.0.1 --port 23410 --prometheus-port 29001 --policy round_robin --lmdeploy-pd-disaggregation > ${{env.SERVER_LOG}}/pd_router.log 2>&1 & + pd_router_pid=$! + echo "PD_ROUTER_PID=$pd_router_pid" >> "$GITHUB_ENV" + for i in $(seq 1 30) + do + sleep 1 + if curl -f -s http://127.0.0.1:23410/health > /dev/null 2>&1; then + echo "PD lmdeploy-router is ready" + exit 0 + fi + done + echo "PD lmdeploy-router failed to start" + cat ${{env.SERVER_LOG}}/pd_router.log + exit 1 + - name: Check pytorch PD migration environment + shell: bash + run: | + echo "::group::Python migration dependencies" + python3 - <<'PY' + try: + import dlslime + print(f'dlslime available: {dlslime.__file__}') + except Exception as error: + print(f'dlslime import failed: {error!r}') + + try: + from dlslime import available_nic + print(f'dlslime available NICs: {available_nic()}') + except Exception as error: + print(f'dlslime NIC discovery failed: {error!r}') + PY + echo "::endgroup::" + + echo "::group::RDMA devices" + ibv_devices || true + ls -l /dev/infiniband || true + echo "::endgroup::" + - name: Test restful server - pytorch pd Qwen3.5-35B-A3B + shell: bash + run: | + if ! python3 -c "from dlslime import available_nic; assert available_nic()" 2>/dev/null || [ ! -d /dev/infiniband ]; then + echo "Skip pytorch PD test: this runner has no RDMA NIC or /dev/infiniband." + exit 0 + fi + + CUDA_VISIBLE_DEVICES=6,7 lmdeploy serve api_server /nvme/qa_test_models/Qwen/Qwen3.5-35B-A3B --server-name 127.0.0.1 --server-port 23411 --tp 2 --backend pytorch --role Prefill --proxy-url http://127.0.0.1:23410 > ${{env.SERVER_LOG}}/pytorch_pd_Qwen3.5-35B-A3B_prefill.log 2>&1 & + prefill_pid=$! + CUDA_VISIBLE_DEVICES=4,5 lmdeploy serve api_server /nvme/qa_test_models/Qwen/Qwen3.5-35B-A3B --server-name 127.0.0.1 --server-port 23412 --tp 2 --backend pytorch --role Decode --proxy-url http://127.0.0.1:23410 > ${{env.SERVER_LOG}}/pytorch_pd_Qwen3.5-35B-A3B_decode.log 2>&1 & + decode_pid=$! + for i in $(seq 1 240) + do + sleep 5 + echo "health check try $i" + if curl -f -s http://127.0.0.1:23411/health > /dev/null 2>&1 && \ + curl -f -s http://127.0.0.1:23412/health > /dev/null 2>&1 && \ + curl -f -s http://127.0.0.1:23410/nodes/status | grep -q 'http://127.0.0.1:23411' && \ + curl -f -s http://127.0.0.1:23410/nodes/status | grep -q 'http://127.0.0.1:23412'; then + if ! LMDEPLOY_PORT=23410 pytest autotest/interface/restful/test_restful_chat_completions_v1.py -n 8 -k 'Qwen/Qwen3.5-35B-A3B and pytorch and not Qwen/Qwen3.5-35B-A3B-FP8' -m 'not not_pytorch and pr_test' --alluredir=${{env.REPORT_DIR}}; then + pd_test_status=1 + break + fi + if ! LMDEPLOY_PORT=23410 pytest autotest/interface/restful/tool_parser/test_tool_call_basic.py -n 8 -k 'Qwen/Qwen3.5-35B-A3B and pytorch and not Qwen/Qwen3.5-35B-A3B-FP8' --alluredir=${{env.REPORT_DIR}}; then + pd_test_status=1 + break + fi + curl -f -s http://127.0.0.1:23410/nodes/remove -H 'Content-Type: application/json' -d '{"url":"http://127.0.0.1:23412"}' > /dev/null 2>&1 || true + curl -f -s http://127.0.0.1:23410/nodes/remove -H 'Content-Type: application/json' -d '{"url":"http://127.0.0.1:23411"}' > /dev/null 2>&1 || true + curl -f -s http://127.0.0.1:23412/terminate > /dev/null 2>&1 + curl -f -s http://127.0.0.1:23411/terminate > /dev/null 2>&1 + wait $prefill_pid + wait $decode_pid + if [ "$pd_test_status" -ne 0 ]; then + cat ${{env.SERVER_LOG}}/pytorch_pd_Qwen3.5-35B-A3B_prefill.log || true + cat ${{env.SERVER_LOG}}/pytorch_pd_Qwen3.5-35B-A3B_decode.log || true + cat ${{env.SERVER_LOG}}/pd_router.log || true + exit 1 + fi + exit 0 + fi + done + + echo "health check fail" + curl -f -s http://127.0.0.1:23412/terminate > /dev/null 2>&1 + curl -f -s http://127.0.0.1:23411/terminate > /dev/null 2>&1 + wait $prefill_pid + wait $decode_pid + cat ${{env.SERVER_LOG}}/pytorch_pd_Qwen3.5-35B-A3B_prefill.log + cat ${{env.SERVER_LOG}}/pytorch_pd_Qwen3.5-35B-A3B_decode.log + exit 1 + - name: Stop lmdeploy-router for PD + shell: bash + run: | + kill "${PD_ROUTER_PID:?PD router PID was not exported}" || true + wait "${PD_ROUTER_PID}" || true - name: Test restful server - turbomind Qwen3-32B run: | - CUDA_VISIBLE_DEVICES=6,7 lmdeploy serve api_server /nvme/qa_test_models/Qwen/Qwen3-32B --tp 2 --backend turbomind --logprobs-mode raw_logprobs --allow-terminate-by-client > ${{env.SERVER_LOG}}/turbomind_Qwen3-32B_start_restful.log 2>&1 & + CUDA_VISIBLE_DEVICES=6,7 lmdeploy serve api_server /nvme/qa_test_models/Qwen/Qwen3-32B --server-name 127.0.0.1 --server-port 23402 --tp 2 --backend turbomind --logprobs-mode raw_logprobs --allow-terminate-by-client --proxy-url http://127.0.0.1:23333 > ${{env.SERVER_LOG}}/turbomind_Qwen3-32B_start_restful.log 2>&1 & echo "restful_pid=$!" for i in $(seq 1 180) do sleep 5 echo "health check try $i" - if curl -f -s http://127.0.0.1:23333/health > /dev/null 2>&1; then + if curl -f -s http://127.0.0.1:23402/health > /dev/null 2>&1 && \ + curl -f -s http://127.0.0.1:23333/nodes/status | grep -q 'http://127.0.0.1:23402'; then pytest autotest/interface/restful/test_restful_chat_completions_v1.py -n 20 -k 'Qwen/Qwen3-32B and turbomind' -m 'not not_turbomind and pr_test' --alluredir=${{env.REPORT_DIR}} pytest autotest/interface/restful/test_restful_generate.py -n 20 -k 'Qwen/Qwen3-32B and turbomind' -m 'not not_turbomind and not experts' --alluredir=${{env.REPORT_DIR}} - curl -f -s http://127.0.0.1:23333/terminate > /dev/null 2>&1 + curl -f -s http://127.0.0.1:23333/nodes/remove -H 'Content-Type: application/json' -d '{"url":"http://127.0.0.1:23402"}' > /dev/null 2>&1 || true + curl -f -s http://127.0.0.1:23402/terminate > /dev/null 2>&1 + wait $restful_pid exit 0 fi done echo "health check fail" - curl -f -s http://127.0.0.1:23333/terminate > /dev/null 2>&1 + curl -f -s http://127.0.0.1:23402/terminate > /dev/null 2>&1 + wait $restful_pid cat ${{env.SERVER_LOG}}/turbomind_Qwen3-32B_start_restful.log exit 1 - name: Test restful server - turbomind InternVL3-38B run: | - CUDA_VISIBLE_DEVICES=6,7 lmdeploy serve api_server /nvme/qa_test_models/OpenGVLab/InternVL3-38B --tp 2 --backend turbomind --logprobs-mode raw_logprobs --allow-terminate-by-client --trust-remote-code > ${{env.SERVER_LOG}}/turbomind_InternVL3-38B_start_restful.log 2>&1 & + CUDA_VISIBLE_DEVICES=6,7 lmdeploy serve api_server /nvme/qa_test_models/OpenGVLab/InternVL3-38B --server-name 127.0.0.1 --server-port 23403 --tp 2 --backend turbomind --logprobs-mode raw_logprobs --allow-terminate-by-client --trust-remote-code --proxy-url http://127.0.0.1:23333 > ${{env.SERVER_LOG}}/turbomind_InternVL3-38B_start_restful.log 2>&1 & echo "restful_pid=$!" for i in $(seq 1 180) do sleep 5 echo "health check try $i" - if curl -f -s http://127.0.0.1:23333/health > /dev/null 2>&1; then + if curl -f -s http://127.0.0.1:23403/health > /dev/null 2>&1 && \ + curl -f -s http://127.0.0.1:23333/nodes/status | grep -q 'http://127.0.0.1:23403'; then pytest autotest/interface/restful/test_restful_chat_completions_v1.py -n 20 -k 'OpenGVLab/InternVL3-38B and turbomind' -m 'not not_turbomind and pr_test' --alluredir=${{env.REPORT_DIR}} pytest autotest/interface/restful/test_restful_generate.py -n 20 -k 'OpenGVLab/InternVL3-38B and turbomind' -m 'not not_turbomind and not experts' --alluredir=${{env.REPORT_DIR}} - curl -f -s http://127.0.0.1:23333/terminate > /dev/null 2>&1 + curl -f -s http://127.0.0.1:23333/nodes/remove -H 'Content-Type: application/json' -d '{"url":"http://127.0.0.1:23403"}' > /dev/null 2>&1 || true + curl -f -s http://127.0.0.1:23403/terminate > /dev/null 2>&1 + wait $restful_pid exit 0 fi done echo "health check fail" - curl -f -s http://127.0.0.1:23333/terminate > /dev/null 2>&1 + curl -f -s http://127.0.0.1:23403/terminate > /dev/null 2>&1 + wait $restful_pid cat ${{env.SERVER_LOG}}/turbomind_InternVL3-38B_start_restful.log exit 1 - name: Test restful server - turbomind Qwen3-30B-A3B run: | - CUDA_VISIBLE_DEVICES=6,7 lmdeploy serve api_server /nvme/qa_test_models/Qwen/Qwen3-30B-A3B --tp 2 --backend turbomind --logprobs-mode raw_logprobs --allow-terminate-by-client> ${{env.SERVER_LOG}}/turbomind_Qwen3-30B-A3B_start_restful.log 2>&1 & + CUDA_VISIBLE_DEVICES=6,7 lmdeploy serve api_server /nvme/qa_test_models/Qwen/Qwen3-30B-A3B --server-name 127.0.0.1 --server-port 23404 --tp 2 --backend turbomind --logprobs-mode raw_logprobs --allow-terminate-by-client --proxy-url http://127.0.0.1:23333 > ${{env.SERVER_LOG}}/turbomind_Qwen3-30B-A3B_start_restful.log 2>&1 & echo "restful_pid=$!" for i in $(seq 1 180) do sleep 5 echo "health check try $i" - if curl -f -s http://127.0.0.1:23333/health > /dev/null 2>&1; then + if curl -f -s http://127.0.0.1:23404/health > /dev/null 2>&1 && \ + curl -f -s http://127.0.0.1:23333/nodes/status | grep -q 'http://127.0.0.1:23404'; then pytest autotest/interface/restful/test_restful_chat_completions_v1.py -n 20 -k 'Qwen/Qwen3-30B-A3B and turbomind' -m 'not not_turbomind and pr_test' --alluredir=${{env.REPORT_DIR}} pytest autotest/interface/restful/test_restful_generate.py -n 20 -k 'Qwen/Qwen3-30B-A3B and turbomind' -m 'not not_turbomind and not experts' --alluredir=${{env.REPORT_DIR}} - curl -f -s http://127.0.0.1:23333/terminate > /dev/null 2>&1 + curl -f -s http://127.0.0.1:23333/nodes/remove -H 'Content-Type: application/json' -d '{"url":"http://127.0.0.1:23404"}' > /dev/null 2>&1 || true + curl -f -s http://127.0.0.1:23404/terminate > /dev/null 2>&1 + wait $restful_pid exit 0 fi done echo "health check fail" - curl -f -s http://127.0.0.1:23333/terminate > /dev/null 2>&1 + curl -f -s http://127.0.0.1:23404/terminate > /dev/null 2>&1 + wait $restful_pid cat ${{env.SERVER_LOG}}/turbomind_Qwen3-30B-A3B_start_restful.log exit 1 - name: Test restful server - pytorch Qwen3-30B-A3B run: | - CUDA_VISIBLE_DEVICES=6,7 lmdeploy serve api_server /nvme/qa_test_models/Qwen/Qwen3-30B-A3B --tp 2 --backend pytorch --logprobs-mode raw_logprobs --enable-return-routed-experts --allow-terminate-by-client > ${{env.SERVER_LOG}}/pytorch_Qwen3-30B-A3B_start_restful.log 2>&1 & + CUDA_VISIBLE_DEVICES=6,7 lmdeploy serve api_server /nvme/qa_test_models/Qwen/Qwen3-30B-A3B --server-name 127.0.0.1 --server-port 23405 --tp 2 --backend pytorch --logprobs-mode raw_logprobs --enable-return-routed-experts --allow-terminate-by-client --proxy-url http://127.0.0.1:23333 > ${{env.SERVER_LOG}}/pytorch_Qwen3-30B-A3B_start_restful.log 2>&1 & echo "restful_pid=$!" for i in $(seq 1 180) do sleep 5 echo "health check try $i" - if curl -f -s http://127.0.0.1:23333/health > /dev/null 2>&1; then + if curl -f -s http://127.0.0.1:23405/health > /dev/null 2>&1 && \ + curl -f -s http://127.0.0.1:23333/nodes/status | grep -q 'http://127.0.0.1:23405'; then pytest autotest/interface/restful/test_restful_chat_completions_v1.py -n 20 -k 'Qwen/Qwen3-30B-A3B and pytorch' -m 'not not_pytorch and pr_test' --alluredir=${{env.REPORT_DIR}} pytest autotest/interface/restful/test_restful_generate.py -n 20 -k 'Qwen/Qwen3-30B-A3B and pytorch' -m 'not not_pytorch' --alluredir=${{env.REPORT_DIR}} - curl -f -s http://127.0.0.1:23333/terminate > /dev/null 2>&1 + curl -f -s http://127.0.0.1:23333/nodes/remove -H 'Content-Type: application/json' -d '{"url":"http://127.0.0.1:23405"}' > /dev/null 2>&1 || true + curl -f -s http://127.0.0.1:23405/terminate > /dev/null 2>&1 + wait $restful_pid exit 0 fi done echo "health check fail" - curl -f -s http://127.0.0.1:23333/terminate > /dev/null 2>&1 + curl -f -s http://127.0.0.1:23405/terminate > /dev/null 2>&1 + wait $restful_pid cat ${{env.SERVER_LOG}}/pytorch_Qwen3-30B-A3B_start_restful.log exit 1 - name: Test restful server - pytorch Qwen3-VL-30B-A3B-Instruct run: | - CUDA_VISIBLE_DEVICES=6,7 lmdeploy serve api_server /nvme/qa_test_models/Qwen/Qwen3-VL-30B-A3B-Instruct --tp 2 --backend pytorch --logprobs-mode raw_logprobs --allow-terminate-by-client > ${{env.SERVER_LOG}}/pytorch_Qwen3-VL-30B-A3B-Instruct_start_restful.log 2>&1 & + CUDA_VISIBLE_DEVICES=6,7 lmdeploy serve api_server /nvme/qa_test_models/Qwen/Qwen3-VL-30B-A3B-Instruct --server-name 127.0.0.1 --server-port 23406 --tp 2 --backend pytorch --logprobs-mode raw_logprobs --allow-terminate-by-client --proxy-url http://127.0.0.1:23333 > ${{env.SERVER_LOG}}/pytorch_Qwen3-VL-30B-A3B-Instruct_start_restful.log 2>&1 & echo "restful_pid=$!" for i in $(seq 1 180) do sleep 5 echo "health check try $i" - if curl -f -s http://127.0.0.1:23333/health > /dev/null 2>&1; then + if curl -f -s http://127.0.0.1:23406/health > /dev/null 2>&1 && \ + curl -f -s http://127.0.0.1:23333/nodes/status | grep -q 'http://127.0.0.1:23406'; then pytest autotest/interface/restful/test_restful_chat_completions_v1.py -n 20 -k 'Qwen/Qwen3-VL-30B-A3B-Instruct and pytorch' -m 'not not_pytorch and pr_test' --alluredir=${{env.REPORT_DIR}} pytest autotest/interface/restful/test_restful_generate.py -n 20 -k 'Qwen/Qwen3-VL-30B-A3B-Instruct and pytorch' -m 'not not_pytorch and not experts' --alluredir=${{env.REPORT_DIR}} - curl -f -s http://127.0.0.1:23333/terminate > /dev/null 2>&1 + curl -f -s http://127.0.0.1:23333/nodes/remove -H 'Content-Type: application/json' -d '{"url":"http://127.0.0.1:23406"}' > /dev/null 2>&1 || true + curl -f -s http://127.0.0.1:23406/terminate > /dev/null 2>&1 + wait $restful_pid exit 0 fi done echo "health check fail" - curl -f -s http://127.0.0.1:23333/terminate > /dev/null 2>&1 + curl -f -s http://127.0.0.1:23406/terminate > /dev/null 2>&1 + wait $restful_pid cat ${{env.SERVER_LOG}}/pytorch_Qwen3-VL-30B-A3B-Instruct_start_restful.log exit 1 - name: Test restful server - pytorch InternVL3_5-30B-A3B run: | - CUDA_VISIBLE_DEVICES=6,7 lmdeploy serve api_server /nvme/qa_test_models/OpenGVLab/InternVL3_5-30B-A3B --tp 2 --backend pytorch --logprobs-mode raw_logprobs --allow-terminate-by-client --trust-remote-code > ${{env.SERVER_LOG}}/pytorch_InternVL3_5-30B-A3B_start_restful.log 2>&1 & + CUDA_VISIBLE_DEVICES=6,7 lmdeploy serve api_server /nvme/qa_test_models/OpenGVLab/InternVL3_5-30B-A3B --server-name 127.0.0.1 --server-port 23407 --tp 2 --backend pytorch --logprobs-mode raw_logprobs --allow-terminate-by-client --trust-remote-code --proxy-url http://127.0.0.1:23333 > ${{env.SERVER_LOG}}/pytorch_InternVL3_5-30B-A3B_start_restful.log 2>&1 & echo "restful_pid=$!" for i in $(seq 1 180) do sleep 5 echo "health check try $i" - if curl -f -s http://127.0.0.1:23333/health > /dev/null 2>&1; then + if curl -f -s http://127.0.0.1:23407/health > /dev/null 2>&1 && \ + curl -f -s http://127.0.0.1:23333/nodes/status | grep -q 'http://127.0.0.1:23407'; then pytest autotest/interface/restful/test_restful_chat_completions_v1.py -n 20 -k 'OpenGVLab/InternVL3_5-30B-A3B and pytorch' -m 'not not_pytorch and pr_test' --alluredir=${{env.REPORT_DIR}} pytest autotest/interface/restful/test_restful_generate.py -n 20 -k 'OpenGVLab/InternVL3_5-30B-A3B and pytorch' -m 'not not_pytorch and not experts' --alluredir=${{env.REPORT_DIR}} - curl -f -s http://127.0.0.1:23333/terminate > /dev/null 2>&1 + curl -f -s http://127.0.0.1:23333/nodes/remove -H 'Content-Type: application/json' -d '{"url":"http://127.0.0.1:23407"}' > /dev/null 2>&1 || true + curl -f -s http://127.0.0.1:23407/terminate > /dev/null 2>&1 + wait $restful_pid exit 0 fi done echo "health check fail" - curl -f -s http://127.0.0.1:23333/terminate > /dev/null 2>&1 + curl -f -s http://127.0.0.1:23407/terminate > /dev/null 2>&1 + wait $restful_pid cat ${{env.SERVER_LOG}}/pytorch_InternVL3_5-30B-A3B_start_restful.log exit 1 + - name: Stop lmdeploy-router + if: always() + run: | + kill "$router_pid" 2>/dev/null || true - name: Clear workfile if: always() run: | diff --git a/.github/workflows/router_release.yml b/.github/workflows/router_release.yml new file mode 100644 index 0000000000..3a317ac297 --- /dev/null +++ b/.github/workflows/router_release.yml @@ -0,0 +1,64 @@ +name: router-release + +on: + workflow_dispatch: + inputs: + upload: + description: 'Upload lmdeploy-router distributions to PyPI' + required: true + type: boolean + default: false + +permissions: + contents: read + +jobs: + build: + runs-on: ubuntu-latest + timeout-minutes: 30 + defaults: + run: + working-directory: lmdeploy/router + steps: + - name: Clone repository + uses: actions/checkout@v4 + - name: Install Rust toolchain + uses: dtolnay/rust-toolchain@stable + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.12' + - name: Build router wheel + run: | + python -m pip install --upgrade build setuptools-rust wheel + python -m build --wheel --outdir dist . + - name: Inspect distributions + run: ls -lh dist/* + - name: Upload build artifacts + uses: actions/upload-artifact@v4 + with: + name: lmdeploy-router-${{ github.sha }} + path: lmdeploy/router/dist/* + if-no-files-found: error + retention-days: 14 + + publish: + if: ${{ inputs.upload }} + needs: build + runs-on: ubuntu-latest + timeout-minutes: 10 + environment: prod + steps: + - name: Download router artifacts + uses: actions/download-artifact@v4 + with: + name: lmdeploy-router-${{ github.sha }} + path: dist + - name: Publish lmdeploy-router to PyPI + env: + TWINE_USERNAME: __token__ + TWINE_PASSWORD: ${{ secrets.ROUTER_PYPI_TOKEN }} + run: | + python -m pip install --upgrade twine + twine check dist/* + twine upload dist/* diff --git a/.github/workflows/router_test.yml b/.github/workflows/router_test.yml new file mode 100644 index 0000000000..4ee438e90d --- /dev/null +++ b/.github/workflows/router_test.yml @@ -0,0 +1,64 @@ +name: router-test + +on: + pull_request: + paths: + - ".github/workflows/router_test.yml" + - "lmdeploy/router/**" + push: + branches: + - main + paths: + - ".github/workflows/router_test.yml" + - "lmdeploy/router/**" + +concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }} + cancel-in-progress: true + +jobs: + rust: + runs-on: ubuntu-latest + timeout-minutes: 45 + defaults: + run: + working-directory: lmdeploy/router + steps: + - name: Clone repository + uses: actions/checkout@v6 + - name: Install Rust toolchain + uses: dtolnay/rust-toolchain@stable + with: + components: clippy, rustfmt + - name: Cache Rust artifacts + uses: Swatinem/rust-cache@v2 + with: + workspaces: lmdeploy/router + - name: Check formatting + run: cargo fmt --check + - name: Static analysis + run: cargo clippy --all-targets --locked -- -D warnings + - name: Rust tests + run: cargo test --all-targets --locked + + python: + runs-on: ubuntu-latest + timeout-minutes: 15 + defaults: + run: + working-directory: lmdeploy/router + steps: + - name: Clone repository + uses: actions/checkout@v6 + - name: Set up Python + uses: actions/setup-python@v5 + with: + python-version: '3.12' + cache: pip + - name: Install test dependencies + run: python -m pip install -e .[dev] + - name: Python unit tests + run: python -m pytest py_test/unit + - name: Python mock integration tests + run: >- + python -m pytest py_test/integration -m 'integration and not pd_engine and not lmdeploy' diff --git a/.gitignore b/.gitignore index 973e1666f9..029c5df68d 100644 --- a/.gitignore +++ b/.gitignore @@ -58,6 +58,9 @@ lmdeploy/turbomind/_turbomind*.so lmdeploy/turbomind/_turbomind*.pyd lmdeploy/turbomind/lib/ dist/ +lmdeploy/router/target/ +lmdeploy/router/.pytest_cache/ +lmdeploy/router/.tokenizer_cache/ examples/cpp/llama/*.csv *.npy *.weight @@ -74,6 +77,8 @@ wscripts/ *.bin *config.json *generate_config.json +!lmdeploy/router/examples/configs/*.json +!lmdeploy/router/test_prompt_input.json !lmdeploy/turbomind/hf_repo/config.json # Pytorch diff --git a/docker/Dockerfile b/docker/Dockerfile index 733cad77e9..a510fbde39 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -1,6 +1,7 @@ # Base images ARG IMAGE_TYPE=final ARG CUDA_VERSION=cu128 +ARG BUILD_ROUTER=1 FROM nvidia/cuda:13.0.3-devel-ubuntu24.04 AS cu130 ENV CUDA_VERSION_SHORT=cu130 @@ -19,8 +20,9 @@ FROM cu128 AS cu12 # Builder image FROM ${CUDA_VERSION} AS dev ARG PYTHON_VERSION=3.12 +ARG BUILD_ROUTER=1 -ENV PATH=/opt/py3/bin:/root/.local/bin:${PATH} +ENV PATH=/opt/py3/bin:/root/.local/bin:/root/.cargo/bin:${PATH} ENV DEBIAN_FRONTEND=noninteractive ENV TZ=Etc/UTC @@ -41,7 +43,10 @@ RUN --mount=type=cache,target=/root/.cache \ rm -rf /var/lib/apt/lists/* && \ wget -qO- https://astral.sh/uv/install.sh | sh && \ uv venv -p python${PYTHON_VERSION} --seed /opt/py3 && \ - pip install --upgrade pip build + pip install --upgrade pip build && \ + if [ "${BUILD_ROUTER}" != "0" ]; then \ + curl --proto "=https" --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal; \ + fi # Pip deps only (cache until requirements/common.txt or runtime_cuda.txt changes). No full repo COPY yet. FROM dev AS runtime-deps-pip @@ -79,10 +84,30 @@ WORKDIR /opt/lmdeploy RUN --mount=type=cache,target=/root/.cache \ docker/build.sh +# Build the standalone Rust router in the same image by default. Pass +# --build-arg BUILD_ROUTER=0 to skip Rust installation and compilation. The +# disabled branch posts a small marker because later bind mounts require the +# same artifact path in both variants. +FROM dev AS router-builder +ARG BUILD_ROUTER=1 +ENV PATH=/opt/py3/bin:/root/.local/bin:/root/.cargo/bin:${PATH} +WORKDIR /src +COPY lmdeploy/router/Cargo.toml lmdeploy/router/Cargo.lock ./ +COPY lmdeploy/router/src ./src +RUN if [ "${BUILD_ROUTER}" != "0" ]; then \ + cargo build --release --locked --no-default-features --bin lmdeploy-router && \ + install -D -m 0755 target/release/lmdeploy-router /src/router/lmdeploy-router; \ + else \ + mkdir -p /src/router && printf 'router disabled\n' > /src/router/lmdeploy-router; \ + fi + FROM runtime-deps-pip AS wheels COPY --from=thirdparty-wheels /wheels /wheels COPY --from=lmdeploy-wheel /wheels /wheels +# Keep a stable artifact path across enabled and disabled builds. +RUN --mount=type=bind,target=/router-build-stage,from=router-builder,source=/src/router \ + cp /router-build-stage/lmdeploy-router /wheels/lmdeploy-router # Runtime image FROM nvidia/cuda:13.0.3-base-ubuntu24.04 AS cu130-base @@ -100,6 +125,7 @@ FROM cu128-base AS cu12-base FROM ${CUDA_VERSION}-base AS final ARG PYTHON_VERSION=3.12 +ARG BUILD_ROUTER=1 ENV PYTHON_VERSION=${PYTHON_VERSION} # Install the CUDA runtime requirements explicitly so torch is selected from @@ -114,4 +140,9 @@ RUN --mount=type=cache,target=/root/.cache \ # explicitly set ptxas path for triton ENV PATH=/opt/py3/bin:$PATH ENV TRITON_PTXAS_PATH=/usr/local/cuda/bin/ptxas +RUN --mount=type=bind,target=/wheels-build-stage,from=wheels,source=/wheels \ + if [ "${BUILD_ROUTER}" != "0" ]; then \ + install -m 0755 /wheels-build-stage/lmdeploy-router /opt/lmdeploy-router; \ + fi +ENV PATH=/opt:${PATH} FROM ${IMAGE_TYPE} diff --git a/docs/en/advance/metrics.md b/docs/en/advance/metrics.md index 99a081c7f3..cec6976000 100644 --- a/docs/en/advance/metrics.md +++ b/docs/en/advance/metrics.md @@ -72,7 +72,7 @@ As an example, we use the model `Qwen/Qwen2.5-7B-Instruct` with `DP=2, TP=2`. St ```bash # Proxy server -lmdeploy serve proxy --server-port 8000 --routing-strategy 'min_expected_latency' --serving-strategy Hybrid --log-level INFO +lmdeploy serve proxy --server-port 8000 --routing-strategy 'cache_aware' --serving-strategy Hybrid --log-level INFO # API server LMDEPLOY_DP_MASTER_ADDR=127.0.0.1 \ @@ -87,7 +87,7 @@ lmdeploy serve api_server \ --node-rank 0 ``` -You should be able to see multiple API servers added to the proxy server list. Details can be found in `lmdeploy/serve/proxy/proxy_config.json`. +You should be able to see multiple API servers in the router node list at `http://0.0.0.0:8000/nodes/status`. For example, you may have the following API servers: diff --git a/docs/en/api/openapi.rst b/docs/en/api/openapi.rst index f1f0a8bed4..a781c5d0d0 100644 --- a/docs/en/api/openapi.rst +++ b/docs/en/api/openapi.rst @@ -14,11 +14,3 @@ OpenAI Compatible API Endpoints .. openapi:: ../_static/openai.yaml :request: :paths: /v1/responses - - -Proxy Server API ----------------- - -.. openapi:: ../_static/proxy.yaml - :request: - :examples: diff --git a/docs/en/conf.py b/docs/en/conf.py index a54c642bd6..900eca4f40 100644 --- a/docs/en/conf.py +++ b/docs/en/conf.py @@ -24,7 +24,6 @@ from lmdeploy.serve.openai.api_server import ServerContext # noqa: E402 from lmdeploy.serve.openai.endpoints import create_openai_router # noqa: E402 from lmdeploy.serve.openai.responses import create_responses_router # noqa: E402 -from lmdeploy.serve.proxy.proxy import app as proxy_server # noqa: E402 version_file = '../../lmdeploy/version.py' with open(version_file) as f: @@ -74,9 +73,6 @@ def metrics(): with open(spec_dir / 'openai.yaml', 'w', encoding='utf-8') as f: f.write(safe_dump(openai_server.openapi())) -with open(spec_dir / 'proxy.yaml', 'w', encoding='utf-8') as f: - f.write(safe_dump(proxy_server.openapi())) - # -- General configuration --------------------------------------------------- # If your documentation needs a minimal Sphinx version, state it here. diff --git a/docs/en/llm/api_server.md b/docs/en/llm/api_server.md index d708b7f4dc..c5da0c1f1d 100644 --- a/docs/en/llm/api_server.md +++ b/docs/en/llm/api_server.md @@ -2,6 +2,7 @@ This article primarily discusses the deployment of a single LLM model across multiple GPUs on a single node, providing a service that is compatible with the OpenAI interface, as well as the usage of the service API. For the sake of convenience, we refer to this service as `api_server`. Regarding parallel services with multiple models, please refer to the guide about [Request Distribution Server](proxy_server.md). +When using the router entry point `lmdeploy serve proxy`, install the separate `lmdeploy-router` package first. In the following sections, we will first introduce methods for starting the service, choosing the appropriate one based on your application scenario. diff --git a/docs/en/llm/proxy_server.md b/docs/en/llm/proxy_server.md index 5c7502a776..2c4b9426a2 100644 --- a/docs/en/llm/proxy_server.md +++ b/docs/en/llm/proxy_server.md @@ -1,23 +1,33 @@ -# Request Distributor Server +# Request Router Server -The request distributor service can parallelize multiple api_server services. Users only need to access the proxy URL, and they can indirectly access different api_server services. The proxy service will automatically distribute requests internally, achieving load balancing. +LMDeploy uses the Rust-based `lmdeploy-router` to distribute requests across multiple `api_server` services. Users access the router URL, and the router forwards requests to available `api_server` instances with load balancing. ## Startup -Start the proxy service: +Install the router package, then start the router with the compatibility `lmdeploy serve proxy` entry point: ```shell -lmdeploy serve proxy --server-name {server_name} --server-port {server_port} --routing-strategy "min_expected_latency" --serving-strategy Hybrid +pip install lmdeploy-router +lmdeploy serve proxy --server-name {server_name} --server-port {server_port} --routing-strategy "cache_aware" --serving-strategy Hybrid ``` -After startup is successful, the URL of the proxy service will also be printed by the script. Access this URL in your browser to open the Swagger UI. -Subsequently, users can add it directly to the proxy service when starting the `api_server` service by using the `--proxy-url` command. For example: -`lmdeploy serve api_server InternLM/internlm2-chat-1_8b --proxy-url http://0.0.0.0:8000`。 -In this way, users can access the services of the `api_server` through the proxy node, and the usage of the proxy node is exactly the same as that of the `api_server`, both of which are compatible with the OpenAI format. +The legacy `min_expected_latency` and `min_observed_latency` strategies remain accepted and are mapped to `cache_aware`. For the full router configuration, use the standalone `lmdeploy-router` CLI described in [`lmdeploy/router/README.md`](../../../lmdeploy/router/README.md). -- /v1/models -- /v1/chat/completions -- /v1/completions +After startup, register workers by starting `api_server` with `--proxy-url`, or use the node-management APIs described below: + +```shell +lmdeploy serve api_server InternLM/internlm2-chat-1_8b \ + --server-name 127.0.0.1 \ + --server-port 23333 \ + --proxy-url http://127.0.0.1:8000 +``` + +Clients send OpenAI-compatible requests to the router instead of an `api_server`: + +- `/health` +- `/v1/models` +- `/v1/chat/completions` +- `/v1/completions` ## Node Management @@ -49,9 +59,10 @@ curl -X 'POST' \ ```shell curl -X 'POST' \ - 'http://localhost:8000/nodes/remove?node_url=http://0.0.0.0:23333' \ + 'http://localhost:8000/nodes/remove' \ -H 'accept: application/json' \ - -d '' + -H 'Content-Type: application/json' \ + -d '{"url": "http://127.0.0.1:23333"}' ``` ### Node Management through python @@ -82,9 +93,9 @@ print(response.text) # delete a node import requests url = 'http://localhost:8000/nodes/remove' -headers = {'accept': 'application/json',} -params = {'node_url': 'http://0.0.0.0:23333',} -response = requests.post(url, headers=headers, data='', params=params) +headers = {'accept': 'application/json'} +data = {'url': 'http://127.0.0.1:23333'} +response = requests.post(url, headers=headers, json=data) print(response.text) ``` @@ -97,8 +108,60 @@ LMDeploy currently supports two serving strategies: ## Dispatch Strategy -The current distribution strategies of the proxy service are as follows: +The compatibility entry point supports the following router strategies: + +- `random`: selects a healthy worker at random. +- `round_robin`: selects healthy workers in round-robin order. +- `cache_aware`: routes requests with shared prefixes to the worker with the most useful cache state. +- `power_of_two`: samples two healthy workers and selects the less-loaded one. +- `consistent_hash`: routes requests with the same routing key to the same worker. +- `rendezvous_hash`: scores every healthy worker with the same routing key and selects the highest score; it preserves session affinity with near-even distribution. + +## PD Disaggregation + +For DistServe, start the router first, then register Prefill and Decode services. The router expects +Prefill workers to handle the cache-prefill phase and Decode workers to return the OpenAI-compatible +response. It establishes the required P2P connection between each pair automatically. + +```shell +lmdeploy-router \ + --host 127.0.0.1 \ + --port 8000 \ + --policy power_of_two \ + --lmdeploy-pd-disaggregation +``` + +```shell +lmdeploy serve api_server InternLM/internlm2-chat-1_8b \ + --server-name 127.0.0.1 --server-port 23333 \ + --role Prefill --proxy-url http://127.0.0.1:8000 + +lmdeploy serve api_server InternLM/internlm2-chat-1_8b \ + --server-name 127.0.0.1 --server-port 23334 \ + --role Decode --proxy-url http://127.0.0.1:8000 +``` + +RDMA, DLSlime, and Mooncake infrastructure must be provisioned outside the router. + +## Migration From the Legacy Proxy + +`lmdeploy serve proxy` remains a compatibility entry point that starts `lmdeploy-router`. + +| Legacy command | Router command | +| ------------------------------------------------------------ | ----------------------------------------- | +| `lmdeploy serve proxy --server-name HOST --server-port PORT` | `lmdeploy-router --host HOST --port PORT` | +| `--routing-strategy POLICY` | `--policy POLICY` | +| `--serving-strategy Hybrid` | Omit `--lmdeploy-pd-disaggregation` | +| `--serving-strategy DistServe` | Add `--lmdeploy-pd-disaggregation` | +| `--migration-protocol RDMA` | `--lmdeploy-migration-protocol rdma` | +| `--migration-protocol NVLINK` | `--lmdeploy-migration-protocol nvlink` | +| `--link-type RoCE` | `--lmdeploy-rdma-link-type roce` | +| `--link-type IB` | `--lmdeploy-rdma-link-type ib` | + +The legacy `min_expected_latency` and `min_observed_latency` values are accepted only by the compatibility +entry point; both map to `cache_aware`. -- random： dispatches based on the ability of each api_server node provided by the user to process requests. The greater the request throughput, the more likely it is to be allocated. Nodes that do not provide throughput are treated according to the average throughput of other nodes. -- min_expected_latency： allocates based on the number of requests currently waiting to be processed on each node, and the throughput capability of each node, calculating the expected time required to complete the response. The shortest one gets allocated. Nodes that do not provide throughput are treated similarly. -- min_observed_latency： allocates based on the average time required to handle a certain number of past requests on each node. The one with the shortest time gets allocated. +The old Python implementation in `lmdeploy.serve.proxy` was removed. Importing its internal modules, +including `lmdeploy.serve.proxy.proxy`, `lmdeploy.serve.proxy.utils`, and +`lmdeploy.serve.proxy.streaming_response`, is no longer supported. Use `lmdeploy-router` or the public +`lmdeploy_router` Python package instead. diff --git a/docs/zh_cn/advance/metrics.md b/docs/zh_cn/advance/metrics.md index 1f7bc3588d..8174a6bebc 100644 --- a/docs/zh_cn/advance/metrics.md +++ b/docs/zh_cn/advance/metrics.md @@ -69,7 +69,7 @@ docker compose up ```bash # Proxy server -lmdeploy serve proxy --server-port 8000 --routing-strategy 'min_expected_latency' --serving-strategy Hybrid --log-level INFO +lmdeploy serve proxy --server-port 8000 --routing-strategy 'cache_aware' --serving-strategy Hybrid --log-level INFO # API server LMDEPLOY_DP_MASTER_ADDR=127.0.0.1 \ @@ -84,7 +84,7 @@ lmdeploy serve api_server \ --node-rank 0 ``` -您应该能在代理服务器列表中看到多个 API 服务实例。详细信息可以在 `lmdeploy/serve/proxy/proxy_config.json` 中找到。 +您可以通过 `http://0.0.0.0:8000/nodes/status` 查看 router 节点列表中的多个 API 服务实例。 例如，您可能会看到如下 API 服务地址： diff --git a/docs/zh_cn/api/openapi.rst b/docs/zh_cn/api/openapi.rst index 774f2dde4c..b3aa8e2fae 100644 --- a/docs/zh_cn/api/openapi.rst +++ b/docs/zh_cn/api/openapi.rst @@ -14,11 +14,3 @@ OpenAI 兼容服务器接口 .. openapi:: ../_static/openai.yaml :request: :paths: /v1/responses - - -Proxy 服务器接口 ------------------ - -.. openapi:: ../_static/proxy.yaml - :request: - :examples: diff --git a/docs/zh_cn/conf.py b/docs/zh_cn/conf.py index 81d6ae28c2..1fb93c803e 100644 --- a/docs/zh_cn/conf.py +++ b/docs/zh_cn/conf.py @@ -24,7 +24,6 @@ from lmdeploy.serve.openai.api_server import ServerContext # noqa: E402 from lmdeploy.serve.openai.endpoints import create_openai_router # noqa: E402 from lmdeploy.serve.openai.responses import create_responses_router # noqa: E402 -from lmdeploy.serve.proxy.proxy import app as proxy_server # noqa: E402 version_file = '../../lmdeploy/version.py' with open(version_file) as f: @@ -74,9 +73,6 @@ def metrics(): with open(spec_dir / 'openai.yaml', 'w', encoding='utf-8') as f: f.write(safe_dump(openai_server.openapi())) -with open(spec_dir / 'proxy.yaml', 'w', encoding='utf-8') as f: - f.write(safe_dump(proxy_server.openapi())) - # -- General configuration --------------------------------------------------- # If your documentation needs a minimal Sphinx version, state it here. diff --git a/docs/zh_cn/llm/api_server.md b/docs/zh_cn/llm/api_server.md index 935a808ed3..d9f2c93ba6 100644 --- a/docs/zh_cn/llm/api_server.md +++ b/docs/zh_cn/llm/api_server.md @@ -1,6 +1,7 @@ # 部署 LLM 类 openai 服务 本文主要介绍单个模型在单机多卡环境下，部署兼容 openai 接口服务的方式，以及服务接口的用法。为行文方便，我们把该服务名称为 `api_server`。对于多模型的并行服务，请阅读[请求分发服务器](./proxy_server.md)一文。 +使用 `lmdeploy serve proxy` 路由入口时，请先单独安装 `lmdeploy-router` 包。 在这篇文章中， 我们首先介绍服务启动的两种方法，你可以根据应用场景，选择合适的。 diff --git a/docs/zh_cn/llm/proxy_server.md b/docs/zh_cn/llm/proxy_server.md index eef70aef37..d4122f5e65 100644 --- a/docs/zh_cn/llm/proxy_server.md +++ b/docs/zh_cn/llm/proxy_server.md @@ -1,22 +1,33 @@ -# 请求分发服务器 +# 请求路由服务器 -请求分发服务可以将多个 api_server 服务，进行并联。用户可以只需要访问代理 URL，就可以间接访问不同的 api_server 服务。代理服务内部会自动分发请求，做到负载均衡。 +LMDeploy 使用基于 Rust 的 `lmdeploy-router` 为多个 `api_server` 服务分发请求。用户只需访问 router URL，router 会按负载均衡策略将请求转发到可用的 `api_server` 实例。 ## 启动 -启动代理服务： +安装 router 包后，通过兼容入口 `lmdeploy serve proxy` 启动： ```shell -lmdeploy serve proxy --server-name {server_name} --server-port {server_port} --routing-strategy "min_expected_latency" --serving-strategy Hybrid +pip install lmdeploy-router +lmdeploy serve proxy --server-name {server_name} --server-port {server_port} --routing-strategy "cache_aware" --serving-strategy Hybrid ``` -启动成功后，代理服务的 URL 也会被脚本打印。浏览器访问这个 URL，可以打开 Swagger UI。 -随后，用户可以在启动 api_server 服务的时候，通过 `--proxy-url` 命令将其直接添加到代理服务中。例如：`lmdeploy serve api_server InternLM/internlm2-chat-1_8b --proxy-url http://0.0.0.0:8000`。 -这样，用户可以通过代理节点访问 api_server 的服务，代理节点的使用方式和 api_server 一模一样，都是兼容 OpenAI 的形式。 +旧的 `min_expected_latency` 和 `min_observed_latency` 策略仍可传入，并会映射为 `cache_aware`。完整 router 配置请使用 [`lmdeploy/router/README.md`](../../../lmdeploy/router/README.md) 中的独立 `lmdeploy-router` CLI。 -- /v1/models -- /v1/chat/completions -- /v1/completions +启动后，可在启动 `api_server` 时通过 `--proxy-url` 注册 worker，也可以使用下文的节点管理接口： + +```shell +lmdeploy serve api_server InternLM/internlm2-chat-1_8b \ + --server-name 127.0.0.1 \ + --server-port 23333 \ + --proxy-url http://127.0.0.1:8000 +``` + +客户端直接访问 router 的 OpenAI 兼容接口，而不是某个 `api_server`： + +- `/health` +- `/v1/models` +- `/v1/chat/completions` +- `/v1/completions` ## 节点管理 @@ -48,9 +59,10 @@ curl -X 'POST' \ ```shell curl -X 'POST' \ - 'http://localhost:8000/nodes/remove?node_url=http://0.0.0.0:23333' \ + 'http://localhost:8000/nodes/remove' \ -H 'accept: application/json' \ - -d '' + -H 'Content-Type: application/json' \ + -d '{"url": "http://127.0.0.1:23333"}' ``` ### 通过 python 脚本增删查 @@ -81,9 +93,9 @@ print(response.text) # 删除某个节点 import requests url = 'http://localhost:8000/nodes/remove' -headers = {'accept': 'application/json',} -params = {'node_url': 'http://0.0.0.0:23333',} -response = requests.post(url, headers=headers, data='', params=params) +headers = {'accept': 'application/json'} +data = {'url': 'http://127.0.0.1:23333'} +response = requests.post(url, headers=headers, json=data) print(response.text) ``` @@ -96,8 +108,58 @@ LMDeploy 当前支持混合部署服务（Hybrid），以及 PD 分离部署服 ## 分发策略 -代理服务目前的分发策略如下： +兼容入口支持以下 router 策略： + +- `random`：随机选择健康 worker。 +- `round_robin`：按轮询顺序选择健康 worker。 +- `cache_aware`：优先将共享前缀的请求路由到缓存命中最多的 worker。 +- `power_of_two`：采样两个健康 worker，并选择负载较低者。 +- `consistent_hash`：将相同路由键的请求发送到同一 worker。 +- `rendezvous_hash`：用同一路由键为每个健康 worker 计算分数并选择最高分；保持会话粘性的同时分布更均匀。 + +## PD 分离部署 + +使用 DistServe 时，先启动 router，再注册 Prefill 与 Decode 服务。router 会把缓存 prefill 阶段发给 +Prefill worker，把 OpenAI 兼容响应返回给客户端，并自动为每组 Prefill/Decode worker 建立 P2P 连接。 + +```shell +lmdeploy-router \ + --host 127.0.0.1 \ + --port 8000 \ + --policy power_of_two \ + --lmdeploy-pd-disaggregation +``` + +```shell +lmdeploy serve api_server InternLM/internlm2-chat-1_8b \ + --server-name 127.0.0.1 --server-port 23333 \ + --role Prefill --proxy-url http://127.0.0.1:8000 + +lmdeploy serve api_server InternLM/internlm2-chat-1_8b \ + --server-name 127.0.0.1 --server-port 23334 \ + --role Decode --proxy-url http://127.0.0.1:8000 +``` + +RDMA、DLSlime 与 Mooncake 等基础设施仍需在 router 外部准备。 + +## 从旧 Proxy 迁移 + +`lmdeploy serve proxy` 是兼容入口，内部启动 `lmdeploy-router`。 + +| 旧命令 | Router 命令 | +| ------------------------------------------------------------ | ----------------------------------------- | +| `lmdeploy serve proxy --server-name HOST --server-port PORT` | `lmdeploy-router --host HOST --port PORT` | +| `--routing-strategy POLICY` | `--policy POLICY` | +| `--serving-strategy Hybrid` | 省略 `--lmdeploy-pd-disaggregation` | +| `--serving-strategy DistServe` | 增加 `--lmdeploy-pd-disaggregation` | +| `--migration-protocol RDMA` | `--lmdeploy-migration-protocol rdma` | +| `--migration-protocol NVLINK` | `--lmdeploy-migration-protocol nvlink` | +| `--link-type RoCE` | `--lmdeploy-rdma-link-type roce` | +| `--link-type IB` | `--lmdeploy-rdma-link-type ib` | + +旧值 `min_expected_latency` 和 `min_observed_latency` 只在兼容入口中接受，两者都会映射为 `cache_aware`。 -- random： 根据用户提供的各个 api_server 节点的处理请求的能力，进行有权重的随机。处理请求的吞吐量越大，就越有可能被分配。部分节点没有提供吞吐量，将按照其他节点的平均吞吐量对待。 -- min_expected_latency： 根据每个节点现有的待处理完的请求，和各个节点吞吐能力，计算预期完成响应所需时间，时间最短的将被分配。未提供吞吐量的节点，同上。 -- min_observed_latency： 根据每个节点过去一定数量的请求，处理完成所需的平均用时，用时最短的将被分配。 +旧 Python 实现 `lmdeploy.serve.proxy` 已移除。不再支持导入其内部模块，包括 +`lmdeploy.serve.proxy.proxy`、`lmdeploy.serve.proxy.utils` 和 +`lmdeploy.serve.proxy.streaming_response`。请改用 `lmdeploy-router`，或使用公开的 +`lmdeploy_router` Python 包。 diff --git a/lmdeploy/cli/serve.py b/lmdeploy/cli/serve.py index 4be4e9317c..e1655700da 100644 --- a/lmdeploy/cli/serve.py +++ b/lmdeploy/cli/serve.py @@ -1,4 +1,6 @@ # Copyright (c) OpenMMLab. All rights reserved. +import argparse + from lmdeploy.pytorch.disagg.config import EngineRole, MigrationBackend from lmdeploy.utils import get_max_batch_size @@ -6,7 +8,6 @@ from .utils import ( ArgumentHelper, DefaultsAndTypesHelpFormatter, - convert_args, get_chat_template, get_lora_adapters, get_speculative_config, @@ -189,7 +190,7 @@ def add_parser_api_server(): @staticmethod def add_parser_proxy(): - """Add parser for proxy server command.""" + """Add parser for router server command.""" parser = SubCliServe.subparsers.add_parser('proxy', formatter_class=DefaultsAndTypesHelpFormatter, description=SubCliServe.proxy.__doc__, @@ -197,6 +198,16 @@ def add_parser_proxy(): parser.set_defaults(run=SubCliServe.proxy) parser.add_argument('--server-name', type=str, default='0.0.0.0', help='Host ip for proxy serving') parser.add_argument('--server-port', type=int, default=8000, help='Server port of the proxy') + parser.add_argument( + '--routing-strategy', + type=str, + choices=[ + 'random', 'round_robin', 'cache_aware', 'power_of_two', 'consistent_hash', + 'rendezvous_hash', 'min_expected_latency', 'min_observed_latency' + ], + default='cache_aware', + help='The strategy to dispatch requests to nodes. min_expected_latency and ' + 'min_observed_latency are mapped to cache_aware for compatibility.') parser.add_argument('--serving-strategy', type=str, choices=['Hybrid', 'DistServe'], @@ -204,16 +215,9 @@ def add_parser_proxy(): help='the strategy to serve, Hybrid for colocating Prefill and Decode' 'workloads into same engine, DistServe for Prefill-Decode Disaggregation') parser.add_argument('--dummy-prefill', action='store_true', help='dummy prefill for performance profiler') - parser.add_argument('--routing-strategy', - type=str, - choices=['random', 'min_expected_latency', 'min_observed_latency'], - default='min_expected_latency', - help='the strategy to dispatch requests to nodes') parser.add_argument('--disable-cache-status', action='store_true', - help='Whether to disable cache status of the ' - 'proxy. If set, the proxy will forget the status ' - 'of the previous time') + help=argparse.SUPPRESS) # For Disaggregation parser.add_argument('--migration-protocol', @@ -387,10 +391,42 @@ def api_server(args): @staticmethod def proxy(args): - """Proxy server that manages distributed api_server nodes.""" - from lmdeploy.serve.proxy.proxy import proxy - kwargs = convert_args(args) - proxy(**kwargs) + """Router server that manages distributed api_server nodes.""" + try: + from lmdeploy_router.launch_router import launch_router + except ImportError as exc: + raise ImportError( + 'lmdeploy serve proxy requires the lmdeploy-router package. ' + 'Install it with `pip install lmdeploy-router`.') from exc + + launch_router(SubCliServe.build_router_args(args)) + + @staticmethod + def build_router_args(args): + """Translate the legacy proxy CLI arguments into router arguments.""" + if args.ssl: + raise ValueError('lmdeploy-router does not currently support TLS at the router ingress. ' + 'Terminate TLS in a front proxy or use the standalone lmdeploy-router CLI.') + if args.api_keys and len(args.api_keys) > 1: + raise ValueError('lmdeploy-router currently supports one API key; ' + 'pass only one --api-keys value.') + + routing_strategy = args.routing_strategy + if routing_strategy in ('min_expected_latency', 'min_observed_latency'): + routing_strategy = 'cache_aware' + + return argparse.Namespace( + host=args.server_name, + port=args.server_port, + lmdeploy_pd_disaggregation=args.serving_strategy == 'DistServe', + policy=routing_strategy, + lmdeploy_migration_protocol=args.migration_protocol.lower(), + lmdeploy_rdma_link_type=args.link_type.lower(), + lmdeploy_dummy_prefill=args.dummy_prefill, + lmdeploy_disable_gdr=args.disable_gdr, + api_key=args.api_keys[0] if args.api_keys else None, + log_level=args.log_level.lower(), + ) @staticmethod def add_parsers(): diff --git a/lmdeploy/pytorch/disagg/conn/engine_conn.py b/lmdeploy/pytorch/disagg/conn/engine_conn.py index eb051c19a0..9055e03ed5 100644 --- a/lmdeploy/pytorch/disagg/conn/engine_conn.py +++ b/lmdeploy/pytorch/disagg/conn/engine_conn.py @@ -95,7 +95,16 @@ async def handle_zmq_recv(self, remote_engine_id: str): if not self.engine.end_session(session_id=session_id): logger.error(f'invalid free, {remote_engine_id}, {session_id}') - def zmq_disconnect(self, remote_engine_id: str): - self.p2p_receiver[remote_engine_id].close() - self.p2p_sender[remote_engine_id].close() - self.p2p_conn_ctx[remote_engine_id].term() + def zmq_disconnect(self, remote_engine_id: str) -> bool: + """Disconnect a peer and return whether it was connected.""" + + if remote_engine_id not in self.p2p_receiver: + return False + + receiver = self.p2p_receiver.pop(remote_engine_id) + sender = self.p2p_sender.pop(remote_engine_id) + context = self.p2p_conn_ctx.pop(remote_engine_id) + receiver.close() + sender.close() + context.term() + return True diff --git a/lmdeploy/router/.cargo/config.toml b/lmdeploy/router/.cargo/config.toml new file mode 100644 index 0000000000..305d57fa3b --- /dev/null +++ b/lmdeploy/router/.cargo/config.toml @@ -0,0 +1,15 @@ +[build] +rustflags = [] +incremental = true + +[target.aarch64-apple-darwin] +rustflags = [ + "-C", "link-arg=-undefined", + "-C", "link-arg=dynamic_lookup", +] + +[target.x86_64-apple-darwin] +rustflags = [ + "-C", "link-arg=-undefined", + "-C", "link-arg=dynamic_lookup", +] diff --git a/lmdeploy/router/.coveragerc b/lmdeploy/router/.coveragerc new file mode 100644 index 0000000000..29469c9ca1 --- /dev/null +++ b/lmdeploy/router/.coveragerc @@ -0,0 +1,9 @@ +[run] +source = py_src/lmdeploy_router +omit = + py_src/lmdeploy_router/mini_lb.py + +[report] +fail_under = 80 +omit = + py_src/lmdeploy_router/mini_lb.py diff --git a/lmdeploy/router/Cargo.lock b/lmdeploy/router/Cargo.lock new file mode 100644 index 0000000000..6b758a16ab --- /dev/null +++ b/lmdeploy/router/Cargo.lock @@ -0,0 +1,5460 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + +[[package]] +name = "ahash" +version = "0.8.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a15f179cd60c4584b8a8c596927aadc462e27f2ca70c04e0071964a73ba7a75" +dependencies = [ + "cfg-if", + "getrandom 0.3.4", + "once_cell", + "serde", + "version_check", + "zerocopy", +] + +[[package]] +name = "aho-corasick" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ddd31a130427c27518df266943a5308ed92d4b226cc639f5a8f1002816174301" +dependencies = [ + "memchr", +] + +[[package]] +name = "allocator-api2" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "683d7910e743518b0e34f1186f92494becacb047c7b6bf616c96772180fef923" + +[[package]] +name = "android_system_properties" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "819e7219dbd41043ac279b19830f2efc897156490d7fd6ea916720117ee66311" +dependencies = [ + "libc", +] + +[[package]] +name = "anes" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4b46cbb362ab8752921c97e041f5e366ee6297bd428a31275b9fcf1e380f7299" + +[[package]] +name = "anstream" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "824a212faf96e9acacdbd09febd34438f8f711fb84e09a8916013cd7815ca28d" +dependencies = [ + "anstyle", + "anstyle-parse", + "anstyle-query", + "anstyle-wincon", + "colorchoice", + "is_terminal_polyfill", + "utf8parse", +] + +[[package]] +name = "anstyle" +version = "1.0.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "940b3a0ca603d1eade50a4846a2afffd5ef57a9feac2c0e2ec2e14f9ead76000" + +[[package]] +name = "anstyle-parse" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "52ce7f38b242319f7cabaa6813055467063ecdc9d355bbb4ce0c68908cd8130e" +dependencies = [ + "utf8parse", +] + +[[package]] +name = "anstyle-query" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "40c48f72fd53cd289104fc64099abca73db4166ad86ea0b4341abe65af83dadc" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "anstyle-wincon" +version = "3.0.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "291e6a250ff86cd4a820112fb8898808a366d8f9f58ce16d1f538353ad55747d" +dependencies = [ + "anstyle", + "once_cell_polyfill", + "windows-sys 0.61.2", +] + +[[package]] +name = "anyhow" +version = "1.0.102" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7f202df86484c868dbad7eaa557ef785d5c66295e41b460ef922eca0723b842c" + +[[package]] +name = "async-broadcast" +version = "0.7.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "435a87a52755b8f27fcf321ac4f04b2802e337c8c4872923137471ec39c37532" +dependencies = [ + "event-listener 5.4.1", + "event-listener-strategy", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-channel" +version = "1.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81953c529336010edd6d8e358f886d9581267795c61b19475b71314bffa46d35" +dependencies = [ + "concurrent-queue", + "event-listener 2.5.3", + "futures-core", +] + +[[package]] +name = "async-channel" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "924ed96dd52d1b75e9c1a3e6275715fd320f5f9439fb5a4a11fa51f4221158d2" +dependencies = [ + "concurrent-queue", + "event-listener-strategy", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-compression" +version = "0.4.41" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0f9ee0f6e02ffd7ad5816e9464499fba7b3effd01123b515c41d1697c43dad1" +dependencies = [ + "compression-codecs", + "compression-core", + "pin-project-lite", + "tokio", +] + +[[package]] +name = "async-executor" +version = "1.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c96bf972d85afc50bf5ab8fe2d54d1586b4e0b46c97c50a0c9e71e2f7bcd812a" +dependencies = [ + "async-task", + "concurrent-queue", + "fastrand", + "futures-lite", + "pin-project-lite", + "slab", +] + +[[package]] +name = "async-global-executor" +version = "2.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05b1b633a2115cd122d73b955eadd9916c18c8f510ec9cd1686404c60ad1c29c" +dependencies = [ + "async-channel 2.5.0", + "async-executor", + "async-io", + "async-lock", + "blocking", + "futures-lite", + "once_cell", +] + +[[package]] +name = "async-io" +version = "2.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "456b8a8feb6f42d237746d4b3e9a178494627745c3c56c6ea55d92ba50d026fc" +dependencies = [ + "autocfg", + "cfg-if", + "concurrent-queue", + "futures-io", + "futures-lite", + "parking", + "polling", + "rustix", + "slab", + "windows-sys 0.61.2", +] + +[[package]] +name = "async-lock" +version = "3.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "290f7f2596bd5b78a9fec8088ccd89180d7f9f55b94b0576823bbbdc72ee8311" +dependencies = [ + "event-listener 5.4.1", + "event-listener-strategy", + "pin-project-lite", +] + +[[package]] +name = "async-process" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc50921ec0055cdd8a16de48773bfeec5c972598674347252c0399676be7da75" +dependencies = [ + "async-channel 2.5.0", + "async-io", + "async-lock", + "async-signal", + "async-task", + "blocking", + "cfg-if", + "event-listener 5.4.1", + "futures-lite", + "rustix", +] + +[[package]] +name = "async-signal" +version = "0.2.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43c070bbf59cd3570b6b2dd54cd772527c7c3620fce8be898406dd3ed6adc64c" +dependencies = [ + "async-io", + "async-lock", + "atomic-waker", + "cfg-if", + "futures-core", + "futures-io", + "rustix", + "signal-hook-registry", + "slab", + "windows-sys 0.61.2", +] + +[[package]] +name = "async-std" +version = "1.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2c8e079a4ab67ae52b7403632e4618815d6db36d2a010cfe41b02c1b1578f93b" +dependencies = [ + "async-channel 1.9.0", + "async-global-executor", + "async-io", + "async-lock", + "async-process", + "crossbeam-utils", + "futures-channel", + "futures-core", + "futures-io", + "futures-lite", + "gloo-timers", + "kv-log-macro", + "log", + "memchr", + "once_cell", + "pin-project-lite", + "pin-utils", + "slab", + "wasm-bindgen-futures", +] + +[[package]] +name = "async-stream" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b5a71a6f37880a80d1d7f19efd781e4b5de42c88f0722cc13bcb6cc2cfe8476" +dependencies = [ + "async-stream-impl", + "futures-core", + "pin-project-lite", +] + +[[package]] +name = "async-stream-impl" +version = "0.3.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7c24de15d275a1ecfd47a380fb4d5ec9bfe0933f309ed5e705b775596a3574d" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "async-task" +version = "4.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b75356056920673b02621b35afd0f7dda9306d03c79a30f5c56c44cf256e3de" + +[[package]] +name = "async-trait" +version = "0.1.89" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9035ad2d096bed7955a320ee7e2230574d28fd3c3a0f186cbea1ff3c7eed5dbb" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "atomic-waker" +version = "1.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1505bd5d3d116872e7271a6d4e16d81d0c8570876c8de68093a09ac269d8aac0" + +[[package]] +name = "autocfg" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08606f8c3cbf4ce6ec8e28fb0014a2c086708fe954eaa885384a6165172e7e8" + +[[package]] +name = "aws-lc-rs" +version = "1.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94bffc006df10ac2a68c83692d734a465f8ee6c5b384d8545a636f81d858f4bf" +dependencies = [ + "aws-lc-sys", + "zeroize", +] + +[[package]] +name = "aws-lc-sys" +version = "0.38.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4321e568ed89bb5a7d291a7f37997c2c0df89809d7b6d12062c81ddb54aa782e" +dependencies = [ + "cc", + "cmake", + "dunce", + "fs_extra", +] + +[[package]] +name = "axum" +version = "0.7.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "edca88bc138befd0323b20752846e6587272d3b03b0343c8ea28a6f819e6e71f" +dependencies = [ + "async-trait", + "axum-core 0.4.5", + "bytes", + "futures-util", + "http", + "http-body", + "http-body-util", + "itoa", + "matchit 0.7.3", + "memchr", + "mime", + "percent-encoding", + "pin-project-lite", + "rustversion", + "serde", + "sync_wrapper", + "tower 0.5.3", + "tower-layer", + "tower-service", +] + +[[package]] +name = "axum" +version = "0.8.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b52af3cb4058c895d37317bb27508dccc8e5f2d39454016b297bf4a400597b8" +dependencies = [ + "axum-core 0.5.6", + "axum-macros", + "base64 0.22.1", + "bytes", + "form_urlencoded", + "futures-util", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-util", + "itoa", + "matchit 0.8.4", + "memchr", + "mime", + "percent-encoding", + "pin-project-lite", + "serde_core", + "serde_json", + "serde_path_to_error", + "serde_urlencoded", + "sha1", + "sync_wrapper", + "tokio", + "tokio-tungstenite", + "tower 0.5.3", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "axum-core" +version = "0.4.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09f2bd6146b97ae3359fa0cc6d6b376d9539582c7b4220f041a33ec24c226199" +dependencies = [ + "async-trait", + "bytes", + "futures-util", + "http", + "http-body", + "http-body-util", + "mime", + "pin-project-lite", + "rustversion", + "sync_wrapper", + "tower-layer", + "tower-service", +] + +[[package]] +name = "axum-core" +version = "0.5.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08c78f31d7b1291f7ee735c1c6780ccde7785daae9a9206026862dab7d8792d1" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "http-body-util", + "mime", + "pin-project-lite", + "sync_wrapper", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "axum-macros" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "604fde5e028fea851ce1d8570bbdc034bec850d157f7569d10f347d06808c05c" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "backoff" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b62ddb9cb1ec0a098ad4bbf9344d0713fa193ae1a80af55febcff2627b6a00c1" +dependencies = [ + "futures-core", + "getrandom 0.2.17", + "instant", + "pin-project-lite", + "rand 0.8.5", + "tokio", +] + +[[package]] +name = "backon" +version = "1.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cffb0e931875b666fc4fcb20fee52e9bbd1ef836fd9e9e04ec21555f9f85f7ef" +dependencies = [ + "fastrand", + "gloo-timers", + "tokio", +] + +[[package]] +name = "base64" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9e1b586273c5702936fe7b7d6896644d8be71e6314cfe09d3167c95f712589e8" + +[[package]] +name = "base64" +version = "0.22.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72b3254f16251a8381aa12e40e3c4d2f0199f8c6508fbecb9d91f575e0fbb8c6" + +[[package]] +name = "bit-set" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0700ddab506f33b20a03b13996eccd309a48e5ff77d0d95926aa0210fb4e95f1" +dependencies = [ + "bit-vec", +] + +[[package]] +name = "bit-vec" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "349f9b6a179ed607305526ca489b34ad0a41aed5f7980fa90eb03160b69598fb" + +[[package]] +name = "bitflags" +version = "2.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843867be96c8daad0d758b57df9392b6d8d271134fce549de6ce169ff98a92af" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "blocking" +version = "1.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e83f8d02be6967315521be875afa792a316e28d57b5a2d401897e2a7921b7f21" +dependencies = [ + "async-channel 2.5.0", + "async-task", + "futures-io", + "futures-lite", + "piper", +] + +[[package]] +name = "bstr" +version = "1.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63044e1ae8e69f3b5a92c736ca6269b8d12fa7efe39bf34ddb06d102cf0e2cab" +dependencies = [ + "memchr", + "regex-automata", + "serde", +] + +[[package]] +name = "bumpalo" +version = "3.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d20789868f4b01b2f2caec9f5c4e0213b41e3e5702a50157d699ae31ced2fcb" + +[[package]] +name = "byteorder" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fd0f2584146f6f2ef48085050886acf353beff7305ebd1ae69500e27c67f64b" + +[[package]] +name = "bytes" +version = "1.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e748733b7cbc798e1434b6ac524f0c1ff2ab456fe201501e6497c8417a4fc33" + +[[package]] +name = "cast" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37b2a672a2cb129a2e41c10b1224bb368f9f37a2b16b612598138befd7b37eb5" + +[[package]] +name = "castaway" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dec551ab6e7578819132c713a93c022a05d60159dc86e7a7050223577484c55a" +dependencies = [ + "rustversion", +] + +[[package]] +name = "cc" +version = "1.2.57" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7a0dd1ca384932ff3641c8718a02769f1698e7563dc6974ffd03346116310423" +dependencies = [ + "find-msvc-tools", + "jobserver", + "libc", + "shlex", +] + +[[package]] +name = "cesu8" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6d43a04d8753f35258c91f8ec639f792891f748a1edbd759cf1dcea3382ad83c" + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "cfg_aliases" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" + +[[package]] +name = "chrono" +version = "0.4.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c673075a2e0e5f4a1dde27ce9dee1ea4558c7ffe648f576438a20ca1d2acc4b0" +dependencies = [ + "iana-time-zone", + "js-sys", + "num-traits", + "serde", + "wasm-bindgen", + "windows-link", +] + +[[package]] +name = "ciborium" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42e69ffd6f0917f5c029256a24d0161db17cea3997d185db0d35926308770f0e" +dependencies = [ + "ciborium-io", + "ciborium-ll", + "serde", +] + +[[package]] +name = "ciborium-io" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "05afea1e0a06c9be33d539b876f1ce3692f4afea2cb41f740e7743225ed1c757" + +[[package]] +name = "ciborium-ll" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57663b653d948a338bfb3eeba9bb2fd5fcfaecb9e199e87e1eda4d9e8b240fd9" +dependencies = [ + "ciborium-io", + "half", +] + +[[package]] +name = "clap" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b193af5b67834b676abd72466a96c1024e6a6ad978a1f484bd90b85c94041351" +dependencies = [ + "clap_builder", + "clap_derive", +] + +[[package]] +name = "clap_builder" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "714a53001bf66416adb0e2ef5ac857140e7dc3a0c48fb28b2f10762fc4b5069f" +dependencies = [ + "anstream", + "anstyle", + "clap_lex", + "strsim", +] + +[[package]] +name = "clap_derive" +version = "4.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1110bd8a634a1ab8cb04345d8d878267d57c3cf1b38d91b71af6686408bbca6a" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "clap_lex" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8d4a3bb8b1e0c1050499d1815f5ab16d04f0959b233085fb31653fbfc9d98f9" + +[[package]] +name = "cmake" +version = "0.1.57" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75443c44cd6b379beb8c5b45d85d0773baf31cce901fe7bb252f4eff3008ef7d" +dependencies = [ + "cc", +] + +[[package]] +name = "colorchoice" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570" + +[[package]] +name = "combine" +version = "4.6.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba5a308b75df32fe02788e748662718f03fde005016435c444eea572398219fd" +dependencies = [ + "bytes", + "memchr", +] + +[[package]] +name = "compact_str" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3fdb1325a1cece981e8a296ab8f0f9b63ae357bd0784a9faaf548cc7b480707a" +dependencies = [ + "castaway", + "cfg-if", + "itoa", + "rustversion", + "ryu", + "serde", + "static_assertions", +] + +[[package]] +name = "compression-codecs" +version = "0.4.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eb7b51a7d9c967fc26773061ba86150f19c50c0d65c887cb1fbe295fd16619b7" +dependencies = [ + "compression-core", + "flate2", + "memchr", +] + +[[package]] +name = "compression-core" +version = "0.4.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75984efb6ed102a0d42db99afb6c1948f0380d1d91808d5529916e6c08b49d8d" + +[[package]] +name = "concurrent-queue" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "console" +version = "0.15.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "054ccb5b10f9f2cbf51eb355ca1d05c2d279ce1804688d0db74b4733a5aeafd8" +dependencies = [ + "encode_unicode", + "libc", + "once_cell", + "unicode-width", + "windows-sys 0.59.0", +] + +[[package]] +name = "console" +version = "0.16.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d64e8af5551369d19cf50138de61f1c42074ab970f74e99be916646777f8fc87" +dependencies = [ + "encode_unicode", + "libc", + "unicode-width", + "windows-sys 0.61.2", +] + +[[package]] +name = "core-foundation" +version = "0.9.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91e195e091a93c46f7102ec7818a2aa394e1e1771c3ab4825963fa03e45afb8f" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b2a6cd9ae233e7f62ba4e9353e81a88df7fc8a5987b8d445b4d90c879bd156f6" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "core-foundation-sys" +version = "0.8.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773648b94d0e5d620f64f280777445740e61fe701025087ec8b57f45c791888b" + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "criterion" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2b12d017a929603d80db1831cd3a24082f8137ce19c69e6447f54f5fc8d692f" +dependencies = [ + "anes", + "cast", + "ciborium", + "clap", + "criterion-plot", + "is-terminal", + "itertools 0.10.5", + "num-traits", + "once_cell", + "oorandom", + "plotters", + "rayon", + "regex", + "serde", + "serde_derive", + "serde_json", + "tinytemplate", + "walkdir", +] + +[[package]] +name = "criterion-plot" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6b50826342786a51a89e2da3a28f1c32b06e387201bc2d19791f622c673706b1" +dependencies = [ + "cast", + "itertools 0.10.5", +] + +[[package]] +name = "crossbeam-channel" +version = "0.5.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "82b8f8f868b36967f9606790d1903570de9ceaf870a7bf9fbbd3016d636a2cb2" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" + +[[package]] +name = "crunchy" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "darling" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc7f46116c46ff9ab3eb1597a45688b6715c6e628b5c133e288e709a29bcb4ee" +dependencies = [ + "darling_core", + "darling_macro", +] + +[[package]] +name = "darling_core" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d00b9596d185e565c2207a0b01f8bd1a135483d02d9b7b0a54b11da8d53412e" +dependencies = [ + "fnv", + "ident_case", + "proc-macro2", + "quote", + "strsim", + "syn", +] + +[[package]] +name = "darling_macro" +version = "0.20.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc34b93ccb385b40dc71c6fceac4b2ad23662c7eeb248cf10d529b7e055b6ead" +dependencies = [ + "darling_core", + "quote", + "syn", +] + +[[package]] +name = "dary_heap" +version = "0.3.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06d2e3287df1c007e74221c49ca10a95d557349e54b3a75dc2fb14712c751f04" +dependencies = [ + "serde", +] + +[[package]] +name = "dashmap" +version = "6.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5041cc499144891f3790297212f32a74fb938e5136a14943f338ef9e0ae276cf" +dependencies = [ + "cfg-if", + "crossbeam-utils", + "hashbrown 0.14.5", + "lock_api", + "once_cell", + "parking_lot_core", +] + +[[package]] +name = "data-encoding" +version = "2.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7a1e2f27636f116493b8b860f5546edb47c8d8f8ea73e1d2a20be88e28d1fea" + +[[package]] +name = "deadpool" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0be2b1d1d6ec8d846f05e137292d0b89133caf95ef33695424c09568bdd39b1b" +dependencies = [ + "deadpool-runtime", + "lazy_static", + "num_cpus", + "tokio", +] + +[[package]] +name = "deadpool-runtime" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "092966b41edc516079bdf31ec78a2e0588d1d0c08f78b91d8307215928642b2b" +dependencies = [ + "tokio", +] + +[[package]] +name = "deranged" +version = "0.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cd812cc2bc1d69d4764bd80df88b4317eaef9e773c75226407d9bc0876b211c" +dependencies = [ + "powerfmt", +] + +[[package]] +name = "derive_builder" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "507dfb09ea8b7fa618fcf76e953f4f5e192547945816d5358edffe39f6f94947" +dependencies = [ + "derive_builder_macro", +] + +[[package]] +name = "derive_builder_core" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2d5bcf7b024d6835cfb3d473887cd966994907effbe9227e8c8219824d06c4e8" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "derive_builder_macro" +version = "0.20.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab63b0e2bf4d5928aff72e83a7dace85d7bba5fe12dcc3c5a572d78caffd3f3c" +dependencies = [ + "derive_builder_core", + "syn", +] + +[[package]] +name = "derive_more" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d751e9e49156b02b44f9c1815bcb94b984cdcc4396ecc32521c739452808b134" +dependencies = [ + "derive_more-impl", +] + +[[package]] +name = "derive_more-impl" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "799a97264921d8623a957f6c3b9011f3b5492f557bbb7a5a19b7fa6d06ba8dcb" +dependencies = [ + "proc-macro2", + "quote", + "rustc_version", + "syn", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "dirs" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3e8aa94d75141228480295a7d0e7feb620b1a5ad9f12bc40be62411e38cce4e" +dependencies = [ + "dirs-sys", +] + +[[package]] +name = "dirs-sys" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e01a3366d27ee9890022452ee61b2b63a67e6f13f58900b651ff5665f0bb1fab" +dependencies = [ + "libc", + "option-ext", + "redox_users", + "windows-sys 0.61.2", +] + +[[package]] +name = "displaydoc" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97369cbbc041bc366949bc74d34658d6cda5621039731c6310521892a3a20ae0" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "dotenvy" +version = "0.15.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1aaf95b3e5c8f23aa320147307562d361db0ae0d51242340f558153b4eb2439b" + +[[package]] +name = "dunce" +version = "1.0.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92773504d58c093f6de2459af4af33faa518c13451eb8f2b5698ed3d36e7c813" + +[[package]] +name = "dyn-clone" +version = "1.0.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" + +[[package]] +name = "educe" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d7bc049e1bd8cdeb31b68bbd586a9464ecf9f3944af3958a7a9d0f8b9799417" +dependencies = [ + "enum-ordinalize", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "either" +version = "1.15.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" + +[[package]] +name = "encode_unicode" +version = "1.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34aa73646ffb006b8f5147f3dc182bd4bcb190227ce861fc4a4844bf8e3cb2c0" + +[[package]] +name = "encoding_rs" +version = "0.8.35" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75030f3c4f45dafd7586dd6780965a8c7e8e285a5ecb86713e63a79c5b2766f3" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "enum-ordinalize" +version = "4.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4a1091a7bb1f8f2c4b28f1fe2cef4980ca2d410a3d727d67ecc3178c9b0800f0" +dependencies = [ + "enum-ordinalize-derive", +] + +[[package]] +name = "enum-ordinalize-derive" +version = "4.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ca9601fb2d62598ee17836250842873a413586e5d7ed88b356e38ddbb0ec631" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "equivalent" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877a4ace8713b0bcf2a4e7eec82529c029f1d0619886d18145fea96c3ffe5c0f" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "esaxx-rs" +version = "0.1.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d817e038c30374a4bcb22f94d0a8a0e216958d4c3dcde369b1439fec4bdda6e6" +dependencies = [ + "cc", +] + +[[package]] +name = "event-listener" +version = "2.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0206175f82b8d6bf6652ff7d71a1e27fd2e4efde587fd368662814d6ec1d9ce0" + +[[package]] +name = "event-listener" +version = "5.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13b66accf52311f30a0db42147dadea9850cb48cd070028831ae5f5d4b856ab" +dependencies = [ + "concurrent-queue", + "parking", + "pin-project-lite", +] + +[[package]] +name = "event-listener-strategy" +version = "0.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8be9f3dfaaffdae2972880079a491a1a8bb7cbed0b8dd7a347f668b4150a3b93" +dependencies = [ + "event-listener 5.4.1", + "pin-project-lite", +] + +[[package]] +name = "fancy-regex" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "531e46835a22af56d1e3b66f04844bed63158bc094a628bec1d321d9b4c44bf2" +dependencies = [ + "bit-set", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "fastrand" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "37909eebbb50d72f9059c3b6d82c0463f2ff062c9e95845c43a6c9c0355411be" + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "flate2" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +dependencies = [ + "crc32fast", + "miniz_oxide", +] + +[[package]] +name = "fnv" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f9eec918d3f24069decb9af1554cad7c880e2da24a9afd88aca000531ab82c1" + +[[package]] +name = "foldhash" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9c4f5dac5e15c24eb999c26181a6ca40b39fe946cbe4c263c7209467bc83af2" + +[[package]] +name = "foldhash" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77ce24cb58228fbb8aa041425bb1050850ac19177686ea6e0f41a70416f56fdb" + +[[package]] +name = "foreign-types" +version = "0.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6f339eb8adc052cd2ca78910fda869aefa38d22d5cb648e6485e4d3fc06f3b1" +dependencies = [ + "foreign-types-shared", +] + +[[package]] +name = "foreign-types-shared" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "00b0228411908ca8685dba7fc2cdd70ec9990a6e753e89b6ac91a84c40fbaf4b" + +[[package]] +name = "form_urlencoded" +version = "1.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb4cb245038516f5f85277875cdaa4f7d2c9a0fa0468de06ed190163b1581fcf" +dependencies = [ + "percent-encoding", +] + +[[package]] +name = "fs_extra" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42703706b716c37f96a77aea830392ad231f44c9e9a67872fa5548707e11b11c" + +[[package]] +name = "futures" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b147ee9d1f6d097cef9ce628cd2ee62288d963e16fb287bd9286455b241382d" +dependencies = [ + "futures-channel", + "futures-core", + "futures-executor", + "futures-io", + "futures-sink", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-channel" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "07bbe89c50d7a535e539b8c17bc0b49bdb77747034daa8087407d655f3f7cc1d" +dependencies = [ + "futures-core", + "futures-sink", +] + +[[package]] +name = "futures-core" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e3450815272ef58cec6d564423f6e755e25379b217b0bc688e295ba24df6b1d" + +[[package]] +name = "futures-executor" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "baf29c38818342a3b26b5b923639e7b1f4a61fc5e76102d4b1981c6dc7a7579d" +dependencies = [ + "futures-core", + "futures-task", + "futures-util", +] + +[[package]] +name = "futures-io" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cecba35d7ad927e23624b22ad55235f2239cfa44fd10428eecbeba6d6a717718" + +[[package]] +name = "futures-lite" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f78e10609fe0e0b3f4157ffab1876319b5b0db102a2c60dc4626306dc46b44ad" +dependencies = [ + "fastrand", + "futures-core", + "futures-io", + "parking", + "pin-project-lite", +] + +[[package]] +name = "futures-macro" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e835b70203e41293343137df5c0664546da5745f82ec9b84d40be8336958447b" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "futures-sink" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c39754e157331b013978ec91992bde1ac089843443c49cbc7f46150b0fad0893" + +[[package]] +name = "futures-task" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "037711b3d59c33004d3856fbdc83b99d4ff37a24768fa1be9ce3538a1cde4393" + +[[package]] +name = "futures-util" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "389ca41296e6190b48053de0321d02a77f32f8a5d2461dd38762c0593805c6d6" +dependencies = [ + "futures-channel", + "futures-core", + "futures-io", + "futures-macro", + "futures-sink", + "futures-task", + "memchr", + "pin-project-lite", + "slab", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "wasi", + "wasm-bindgen", +] + +[[package]] +name = "getrandom" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "899def5c37c4fd7b2664648c28120ecec138e4d395b459e5ca34f9cce2dd77fd" +dependencies = [ + "cfg-if", + "js-sys", + "libc", + "r-efi 5.3.0", + "wasip2", + "wasm-bindgen", +] + +[[package]] +name = "getrandom" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0de51e6874e94e7bf76d726fc5d13ba782deca734ff60d5bb2fb2607c7406555" +dependencies = [ + "cfg-if", + "libc", + "r-efi 6.0.0", + "wasip2", + "wasip3", +] + +[[package]] +name = "glob" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" + +[[package]] +name = "gloo-timers" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbb143cf96099802033e0d4f4963b19fd2e0b728bcf076cd9cf7f6634f092994" +dependencies = [ + "futures-channel", + "futures-core", + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "h2" +version = "0.4.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f44da3a8150a6703ed5d34e164b875fd14c2cdab9af1252a9a1020bde2bdc54" +dependencies = [ + "atomic-waker", + "bytes", + "fnv", + "futures-core", + "futures-sink", + "http", + "indexmap 2.13.0", + "slab", + "tokio", + "tokio-util", + "tracing", +] + +[[package]] +name = "half" +version = "2.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ea2d84b969582b4b1864a92dc5d27cd2b77b622a8d79306834f1be5ba20d84b" +dependencies = [ + "cfg-if", + "crunchy", + "zerocopy", +] + +[[package]] +name = "hashbrown" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a9ee70c43aaf417c914396645a0fa852624801b24ebb7ae78fe8272889ac888" + +[[package]] +name = "hashbrown" +version = "0.14.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5274423e17b7c9fc20b6e7e208532f9b19825d82dfd615708b70edd83df41f1" + +[[package]] +name = "hashbrown" +version = "0.15.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9229cfe53dfd69f0609a49f65461bd93001ea1ef889cd5529dd176593f5338a1" +dependencies = [ + "allocator-api2", + "equivalent", + "foldhash 0.1.5", +] + +[[package]] +name = "hashbrown" +version = "0.16.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" +dependencies = [ + "foldhash 0.2.0", +] + +[[package]] +name = "hdrhistogram" +version = "7.5.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "765c9198f173dd59ce26ff9f95ef0aafd0a0fe01fb9d72841bc5066a4c06511d" +dependencies = [ + "byteorder", + "num-traits", +] + +[[package]] +name = "headers" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b3314d5adb5d94bcdf56771f2e50dbbc80bb4bdf88967526706205ac9eff24eb" +dependencies = [ + "base64 0.22.1", + "bytes", + "headers-core", + "http", + "httpdate", + "mime", + "sha1", +] + +[[package]] +name = "headers-core" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "54b4a22553d4242c49fddb9ba998a99962b5cc6f22cb5a3482bec22522403ce4" +dependencies = [ + "http", +] + +[[package]] +name = "heck" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2304e00983f87ffb38b55b444b5e3b60a884b5d30c0fca7d82fe33449bbe55ea" + +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + +[[package]] +name = "hf-hub" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "629d8f3bbeda9d148036d6b0de0a3ab947abd08ce90626327fc3547a49d59d97" +dependencies = [ + "dirs", + "futures", + "http", + "indicatif 0.17.11", + "libc", + "log", + "native-tls", + "num_cpus", + "rand 0.9.2", + "reqwest 0.12.28", + "serde", + "serde_json", + "thiserror 2.0.18", + "tokio", + "ureq", + "windows-sys 0.60.2", +] + +[[package]] +name = "home" +version = "0.5.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cc627f471c528ff0c4a49e1d5e60450c8f6461dd6d10ba9dcd3a61d3dff7728d" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "hostname" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "617aaa3557aef3810a6369d0a99fac8a080891b68bd9f9812a1eeda0c0730cbd" +dependencies = [ + "cfg-if", + "libc", + "windows-link", +] + +[[package]] +name = "http" +version = "1.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3ba2a386d7f85a81f119ad7498ebe444d2e22c2af0b86b069416ace48b3311a" +dependencies = [ + "bytes", + "itoa", +] + +[[package]] +name = "http-body" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1efedce1fb8e6913f23e0c92de8e62cd5b772a67e7b3946df930a62566c93184" +dependencies = [ + "bytes", + "http", +] + +[[package]] +name = "http-body-util" +version = "0.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b021d93e26becf5dc7e1b75b1bed1fd93124b374ceb73f43d4d4eafec896a64a" +dependencies = [ + "bytes", + "futures-core", + "http", + "http-body", + "pin-project-lite", +] + +[[package]] +name = "httparse" +version = "1.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6dbf3de79e51f3d586ab4cb9d5c3e2c14aa28ed23d180cf89b4df0454a69cc87" + +[[package]] +name = "httpdate" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" + +[[package]] +name = "hyper" +version = "1.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2ab2d4f250c3d7b1c9fcdff1cece94ea4e2dfbec68614f7b87cb205f24ca9d11" +dependencies = [ + "atomic-waker", + "bytes", + "futures-channel", + "futures-core", + "h2", + "http", + "http-body", + "httparse", + "httpdate", + "itoa", + "pin-project-lite", + "pin-utils", + "smallvec", + "tokio", + "want", +] + +[[package]] +name = "hyper-http-proxy" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ad4b0a1e37510028bc4ba81d0e38d239c39671b0f0ce9e02dfa93a8133f7c08" +dependencies = [ + "bytes", + "futures-util", + "headers", + "http", + "hyper", + "hyper-rustls", + "hyper-util", + "pin-project-lite", + "rustls-native-certs 0.7.3", + "tokio", + "tokio-rustls", + "tower-service", +] + +[[package]] +name = "hyper-rustls" +version = "0.27.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3c93eb611681b207e1fe55d5a71ecf91572ec8a6705cdb6857f7d8d5242cf58" +dependencies = [ + "http", + "hyper", + "hyper-util", + "log", + "rustls", + "rustls-native-certs 0.8.3", + "rustls-pki-types", + "tokio", + "tokio-rustls", + "tower-service", +] + +[[package]] +name = "hyper-timeout" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b90d566bffbce6a75bd8b09a05aa8c2cb1fabb6cb348f8840c9e4c90a0d83b0" +dependencies = [ + "hyper", + "hyper-util", + "pin-project-lite", + "tokio", + "tower-service", +] + +[[package]] +name = "hyper-tls" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70206fc6890eaca9fde8a0bf71caa2ddfc9fe045ac9e5c70df101a7dbde866e0" +dependencies = [ + "bytes", + "http-body-util", + "hyper", + "hyper-util", + "native-tls", + "tokio", + "tokio-native-tls", + "tower-service", +] + +[[package]] +name = "hyper-util" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "96547c2556ec9d12fb1578c4eaf448b04993e7fb79cbaad930a656880a6bdfa0" +dependencies = [ + "base64 0.22.1", + "bytes", + "futures-channel", + "futures-util", + "http", + "http-body", + "hyper", + "ipnet", + "libc", + "percent-encoding", + "pin-project-lite", + "socket2 0.6.3", + "system-configuration", + "tokio", + "tower-service", + "tracing", + "windows-registry", +] + +[[package]] +name = "iana-time-zone" +version = "0.1.65" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e31bc9ad994ba00e440a8aa5c9ef0ec67d5cb5e5cb0cc7f8b744a35b389cc470" +dependencies = [ + "android_system_properties", + "core-foundation-sys", + "iana-time-zone-haiku", + "js-sys", + "log", + "wasm-bindgen", + "windows-core", +] + +[[package]] +name = "iana-time-zone-haiku" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f31827a206f56af32e590ba56d5d2d085f558508192593743f16b2306495269f" +dependencies = [ + "cc", +] + +[[package]] +name = "icu_collections" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c6b649701667bbe825c3b7e6388cb521c23d88644678e83c0c4d0a621a34b43" +dependencies = [ + "displaydoc", + "potential_utf", + "yoke", + "zerofrom", + "zerovec", +] + +[[package]] +name = "icu_locale_core" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "edba7861004dd3714265b4db54a3c390e880ab658fec5f7db895fae2046b5bb6" +dependencies = [ + "displaydoc", + "litemap", + "tinystr", + "writeable", + "zerovec", +] + +[[package]] +name = "icu_normalizer" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5f6c8828b67bf8908d82127b2054ea1b4427ff0230ee9141c54251934ab1b599" +dependencies = [ + "icu_collections", + "icu_normalizer_data", + "icu_properties", + "icu_provider", + "smallvec", + "zerovec", +] + +[[package]] +name = "icu_normalizer_data" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7aedcccd01fc5fe81e6b489c15b247b8b0690feb23304303a9e560f37efc560a" + +[[package]] +name = "icu_properties" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "020bfc02fe870ec3a66d93e677ccca0562506e5872c650f893269e08615d74ec" +dependencies = [ + "icu_collections", + "icu_locale_core", + "icu_properties_data", + "icu_provider", + "zerotrie", + "zerovec", +] + +[[package]] +name = "icu_properties_data" +version = "2.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "616c294cf8d725c6afcd8f55abc17c56464ef6211f9ed59cccffe534129c77af" + +[[package]] +name = "icu_provider" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85962cf0ce02e1e0a629cc34e7ca3e373ce20dda4c4d7294bbd0bf1fdb59e614" +dependencies = [ + "displaydoc", + "icu_locale_core", + "writeable", + "yoke", + "zerofrom", + "zerotrie", + "zerovec", +] + +[[package]] +name = "id-arena" +version = "2.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3d3067d79b975e8844ca9eb072e16b31c3c1c36928edf9c6789548c524d0d954" + +[[package]] +name = "ident_case" +version = "1.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e0384b61958566e926dc50660321d12159025e767c18e043daf26b70104c39" + +[[package]] +name = "idna" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3b0875f23caa03898994f6ddc501886a45c7d3d62d04d2d90788d47be1b1e4de" +dependencies = [ + "idna_adapter", + "smallvec", + "utf8_iter", +] + +[[package]] +name = "idna_adapter" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3acae9609540aa318d1bc588455225fb2085b9ed0c4f6bd0d9d5bcd86f1a0344" +dependencies = [ + "icu_normalizer", + "icu_properties", +] + +[[package]] +name = "indexmap" +version = "1.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bd070e393353796e801d209ad339e89596eb4c8d430d18ede6a1cced8fafbd99" +dependencies = [ + "autocfg", + "hashbrown 0.12.3", +] + +[[package]] +name = "indexmap" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7714e70437a7dc3ac8eb7e6f8df75fd8eb422675fc7678aff7364301092b1017" +dependencies = [ + "equivalent", + "hashbrown 0.16.1", + "serde", + "serde_core", +] + +[[package]] +name = "indicatif" +version = "0.17.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "183b3088984b400f4cfac3620d5e076c84da5364016b4f49473de574b2586235" +dependencies = [ + "console 0.15.11", + "number_prefix", + "portable-atomic", + "unicode-width", + "web-time", +] + +[[package]] +name = "indicatif" +version = "0.18.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25470f23803092da7d239834776d653104d551bc4d7eacaf31e6837854b8e9eb" +dependencies = [ + "console 0.16.3", + "portable-atomic", + "unicode-width", + "unit-prefix", + "web-time", +] + +[[package]] +name = "indoc" +version = "2.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "79cf5c93f93228cf8efb3ba362535fb11199ac548a09ce117c9b1adc3030d706" +dependencies = [ + "rustversion", +] + +[[package]] +name = "instant" +version = "0.1.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e0242819d153cba4b4b05a5a8f2a7e9bbf97b6055b2a002b395c96b5ff3c0222" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "ipnet" +version = "2.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" + +[[package]] +name = "iri-string" +version = "0.7.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c91338f0783edbd6195decb37bae672fd3b165faffb89bf7b9e6942f8b1a731a" +dependencies = [ + "memchr", + "serde", +] + +[[package]] +name = "is-terminal" +version = "0.4.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3640c1c38b8e4e43584d8df18be5fc6b0aa314ce6ebf51b53313d4306cca8e46" +dependencies = [ + "hermit-abi", + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "is_terminal_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6cb138bb79a146c1bd460005623e142ef0181e3d0219cb493e02f7d08a35695" + +[[package]] +name = "itertools" +version = "0.10.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0fd2260e829bddf4cb6ea802289de2f86d6a7a690192fbe91b3f46e0f2c8473" +dependencies = [ + "either", +] + +[[package]] +name = "itertools" +version = "0.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b192c782037fadd9cfa75548310488aabdbf3d2da73885b31bd0abd03351285" +dependencies = [ + "either", +] + +[[package]] +name = "itoa" +version = "1.0.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "92ecc6618181def0457392ccd0ee51198e065e016d1d527a7ac1b6dc7c1f09d2" + +[[package]] +name = "jni" +version = "0.21.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1a87aa2bb7d2af34197c04845522473242e1aa17c12f4935d5856491a7fb8c97" +dependencies = [ + "cesu8", + "cfg-if", + "combine", + "jni-sys", + "log", + "thiserror 1.0.69", + "walkdir", + "windows-sys 0.45.0", +] + +[[package]] +name = "jni-sys" +version = "0.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8eaf4bc02d17cbdd7ff4c7438cafcdf7fb9a4613313ad11b4f8fefe7d3fa0130" + +[[package]] +name = "jobserver" +version = "0.1.34" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9afb3de4395d6b3e67a780b6de64b51c978ecf11cb9a462c66be7d4ca9039d33" +dependencies = [ + "getrandom 0.3.4", + "libc", +] + +[[package]] +name = "js-sys" +version = "0.3.91" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b49715b7073f385ba4bc528e5747d02e66cb39c6146efb66b781f131f0fb399c" +dependencies = [ + "once_cell", + "wasm-bindgen", +] + +[[package]] +name = "json-patch" +version = "4.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f300e415e2134745ef75f04562dd0145405c2f7fd92065db029ac4b16b57fe90" +dependencies = [ + "jsonptr", + "serde", + "serde_json", + "thiserror 1.0.69", +] + +[[package]] +name = "jsonpath-rust" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c00ae348f9f8fd2d09f82a98ca381c60df9e0820d8d79fce43e649b4dc3128b" +dependencies = [ + "pest", + "pest_derive", + "regex", + "serde_json", + "thiserror 2.0.18", +] + +[[package]] +name = "jsonptr" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a5a3cc660ba5d72bce0b3bb295bf20847ccbb40fd423f3f05b61273672e561fe" +dependencies = [ + "serde", + "serde_json", +] + +[[package]] +name = "k8s-openapi" +version = "0.25.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "aa60a41b57ae1a0a071af77dbcf89fc9819cfe66edaf2beeb204c34459dcf0b2" +dependencies = [ + "base64 0.22.1", + "chrono", + "serde", + "serde_json", +] + +[[package]] +name = "kube" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "778f98664beaf4c3c11372721e14310d1ae00f5e2d9aabcf8906c881aa4e9f51" +dependencies = [ + "k8s-openapi", + "kube-client", + "kube-core", + "kube-derive", + "kube-runtime", +] + +[[package]] +name = "kube-client" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7cb276b85b6e94ded00ac8ea2c68fcf4697ea0553cb25fddc35d4a0ab718db8d" +dependencies = [ + "base64 0.22.1", + "bytes", + "chrono", + "either", + "futures", + "home", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-http-proxy", + "hyper-rustls", + "hyper-timeout", + "hyper-util", + "jsonpath-rust", + "k8s-openapi", + "kube-core", + "pem", + "rustls", + "secrecy", + "serde", + "serde_json", + "serde_yaml", + "thiserror 2.0.18", + "tokio", + "tokio-util", + "tower 0.5.3", + "tower-http", + "tracing", +] + +[[package]] +name = "kube-core" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3c56ff45deb0031f2a476017eed60c06872251f271b8387ad8020b8fef60960" +dependencies = [ + "chrono", + "derive_more", + "form_urlencoded", + "http", + "json-patch", + "k8s-openapi", + "schemars", + "serde", + "serde-value", + "serde_json", + "thiserror 2.0.18", +] + +[[package]] +name = "kube-derive" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "079fc8c1c397538628309cfdee20696ebdcc26745f9fb17f89b78782205bd995" +dependencies = [ + "darling", + "proc-macro2", + "quote", + "serde", + "serde_json", + "syn", +] + +[[package]] +name = "kube-runtime" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f1326e946fadf6248febdf8a1c001809c3899ccf48cb9768cbc536b741040dc" +dependencies = [ + "ahash", + "async-broadcast", + "async-stream", + "backon", + "educe", + "futures", + "hashbrown 0.15.5", + "hostname", + "json-patch", + "k8s-openapi", + "kube-client", + "parking_lot", + "pin-project", + "serde", + "serde_json", + "thiserror 2.0.18", + "tokio", + "tokio-util", + "tracing", +] + +[[package]] +name = "kv-log-macro" +version = "1.0.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0de8b303297635ad57c9f5059fd9cee7a47f8e8daa09df0fcd07dd39fb22977f" +dependencies = [ + "log", +] + +[[package]] +name = "lazy_static" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbd2bcb4c963f2ddae06a2efc7e9f3591312473c50c6685e1f298068316e66fe" + +[[package]] +name = "leb128fmt" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09edd9e8b54e49e587e4f6295a7d29c3ea94d469cb40ab8ca70b288248a81db2" + +[[package]] +name = "libc" +version = "0.2.183" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5b646652bf6661599e1da8901b3b9522896f01e736bad5f723fe7a3a27f899d" + +[[package]] +name = "libredox" +version = "0.1.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1744e39d1d6a9948f4f388969627434e31128196de472883b39f148769bfe30a" +dependencies = [ + "libc", +] + +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + +[[package]] +name = "litemap" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6373607a59f0be73a39b6fe456b8192fcc3585f602af20751600e974dd455e77" + +[[package]] +name = "lmdeploy_router_rs" +version = "0.0.3" +dependencies = [ + "anyhow", + "async-trait", + "axum 0.8.8", + "backoff", + "bytes", + "chrono", + "clap", + "criterion", + "dashmap", + "deadpool", + "dotenvy", + "futures", + "futures-util", + "hf-hub", + "http", + "http-body-util", + "k8s-openapi", + "kube", + "lazy_static", + "metrics", + "metrics-exporter-prometheus", + "minijinja", + "once_cell", + "opentelemetry", + "opentelemetry-http", + "opentelemetry-otlp", + "opentelemetry_sdk", + "parking_lot", + "portpicker", + "pyo3", + "rand 0.9.2", + "regex", + "reqwest 0.13.2", + "rustls", + "serde", + "serde_json", + "strum", + "tempfile", + "thiserror 2.0.18", + "tiktoken-rs", + "tokenizers", + "tokio", + "tokio-stream", + "tower 0.5.3", + "tower-http", + "tracing", + "tracing-appender", + "tracing-log", + "tracing-opentelemetry", + "tracing-subscriber", + "ulid", + "url", + "uuid", +] + +[[package]] +name = "lock_api" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "224399e74b87b5f3557511d98dff8b14089b3dadafcab6bb93eab67d3aace965" +dependencies = [ + "scopeguard", +] + +[[package]] +name = "log" +version = "0.4.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5e5032e24019045c762d3c0f28f5b6b8bbf38563a65908389bf7978758920897" +dependencies = [ + "value-bag", +] + +[[package]] +name = "lru-slab" +version = "0.1.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "112b39cec0b298b6c1999fee3e31427f74f676e4cb9879ed1a121b43661a4154" + +[[package]] +name = "macro_rules_attribute" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "65049d7923698040cd0b1ddcced9b0eb14dd22c5f86ae59c3740eab64a676520" +dependencies = [ + "macro_rules_attribute-proc_macro", + "paste", +] + +[[package]] +name = "macro_rules_attribute-proc_macro" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "670fdfda89751bc4a84ac13eaa63e205cf0fd22b4c9a5fbfa085b63c1f1d3a30" + +[[package]] +name = "matchers" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1525a2a28c7f4fa0fc98bb91ae755d1e2d1505079e05539e35bc876b5d65ae9" +dependencies = [ + "regex-automata", +] + +[[package]] +name = "matchit" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0e7465ac9959cc2b1404e8e2367b43684a6d13790fe23056cc8c6c5a6b7bcb94" + +[[package]] +name = "matchit" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47e1ffaa40ddd1f3ed91f717a33c8c0ee23fff369e3aa8772b9605cc1d22f4c3" + +[[package]] +name = "memchr" +version = "2.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8ca58f447f06ed17d5fc4043ce1b10dd205e060fb3ce5b979b8ed8e59ff3f79" + +[[package]] +name = "memo-map" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "38d1115007560874e373613744c6fba374c17688327a71c1476d1a5954cc857b" + +[[package]] +name = "memoffset" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "488016bfae457b036d996092f6cb448677611ce4449e970ceaf42695203f218a" +dependencies = [ + "autocfg", +] + +[[package]] +name = "metrics" +version = "0.24.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d5312e9ba3771cfa961b585728215e3d972c950a3eed9252aa093d6301277e8" +dependencies = [ + "ahash", + "portable-atomic", +] + +[[package]] +name = "metrics-exporter-prometheus" +version = "0.17.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2b166dea96003ee2531cf14833efedced545751d800f03535801d833313f8c15" +dependencies = [ + "base64 0.22.1", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "indexmap 2.13.0", + "ipnet", + "metrics", + "metrics-util", + "quanta", + "thiserror 2.0.18", + "tokio", + "tracing", +] + +[[package]] +name = "metrics-util" +version = "0.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdfb1365fea27e6dd9dc1dbc19f570198bc86914533ad639dae939635f096be4" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", + "hashbrown 0.16.1", + "metrics", + "quanta", + "rand 0.9.2", + "rand_xoshiro", + "sketches-ddsketch", +] + +[[package]] +name = "mime" +version = "0.3.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6877bb514081ee2a7ff5ef9de3281f14a4dd4bceac4c09388074a6b5df8a139a" + +[[package]] +name = "minijinja" +version = "2.18.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "328251e58ad8e415be6198888fc207502727dc77945806421ab34f35bf012e7d" +dependencies = [ + "memo-map", + "serde", +] + +[[package]] +name = "minimal-lexical" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68354c5c6bd36d73ff3feceb05efa59b6acb7626617f4962be322a825e61f79a" + +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + +[[package]] +name = "mio" +version = "1.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a69bcab0ad47271a0234d9422b131806bf3968021e5dc9328caf2d4cd58557fc" +dependencies = [ + "libc", + "wasi", + "windows-sys 0.61.2", +] + +[[package]] +name = "monostate" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3341a273f6c9d5bef1908f17b7267bbab0e95c9bf69a0d4dcf8e9e1b2c76ef67" +dependencies = [ + "monostate-impl", + "serde", + "serde_core", +] + +[[package]] +name = "monostate-impl" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e4db6d5580af57bf992f59068d4ea26fd518574ff48d7639b255a36f9de6e7e9" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "native-tls" +version = "0.2.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "465500e14ea162429d264d44189adc38b199b62b1c21eea9f69e4b73cb03bbf2" +dependencies = [ + "libc", + "log", + "openssl", + "openssl-probe 0.2.1", + "openssl-sys", + "schannel", + "security-framework 3.7.0", + "security-framework-sys", + "tempfile", +] + +[[package]] +name = "nom" +version = "7.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d273983c5a657a70a3e8f2a01329822f3b8c8172b73826411a55751e404a0a4a" +dependencies = [ + "memchr", + "minimal-lexical", +] + +[[package]] +name = "nu-ansi-term" +version = "0.50.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7957b9740744892f114936ab4a57b3f487491bbeafaf8083688b16841a4240e5" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "num-conv" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cf97ec579c3c42f953ef76dbf8d55ac91fb219dde70e49aa4a6b7d74e9919050" + +[[package]] +name = "num-traits" +version = "0.2.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "071dfc062690e90b734c0b2273ce72ad0ffa95f0c74596bc250dcfd960262841" +dependencies = [ + "autocfg", +] + +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + +[[package]] +name = "number_prefix" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "830b246a0e5f20af87141b25c173cd1b609bd7779a4617d6ec582abaf90870f3" + +[[package]] +name = "once_cell" +version = "1.21.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" + +[[package]] +name = "once_cell_polyfill" +version = "1.70.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "384b8ab6d37215f3c5301a95a4accb5d64aa607f1fcb26a11b5303878451b4fe" + +[[package]] +name = "onig" +version = "6.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "336b9c63443aceef14bea841b899035ae3abe89b7c486aaf4c5bd8aafedac3f0" +dependencies = [ + "bitflags", + "libc", + "once_cell", + "onig_sys", +] + +[[package]] +name = "onig_sys" +version = "69.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7f86c6eef3d6df15f23bcfb6af487cbd2fed4e5581d58d5bf1f5f8b7f6727dc" +dependencies = [ + "cc", + "pkg-config", +] + +[[package]] +name = "oorandom" +version = "11.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6790f58c7ff633d8771f42965289203411a5e5c68388703c06e14f24770b41e" + +[[package]] +name = "openssl" +version = "0.10.76" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "951c002c75e16ea2c65b8c7e4d3d51d5530d8dfa7d060b4776828c88cfb18ecf" +dependencies = [ + "bitflags", + "cfg-if", + "foreign-types", + "libc", + "once_cell", + "openssl-macros", + "openssl-sys", +] + +[[package]] +name = "openssl-macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a948666b637a0f465e8564c73e89d4dde00d72d4d473cc972f390fc3dcee7d9c" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "openssl-probe" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d05e27ee213611ffe7d6348b942e8f942b37114c00cc03cec254295a4a17852e" + +[[package]] +name = "openssl-probe" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7c87def4c32ab89d880effc9e097653c8da5d6ef28e6b539d313baaacfbafcbe" + +[[package]] +name = "openssl-sys" +version = "0.9.112" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57d55af3b3e226502be1526dfdba67ab0e9c96fc293004e79576b2b9edb0dbdb" +dependencies = [ + "cc", + "libc", + "pkg-config", + "vcpkg", +] + +[[package]] +name = "opentelemetry" +version = "0.27.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab70038c28ed37b97d8ed414b6429d343a8bbf44c9f79ec854f3a643029ba6d7" +dependencies = [ + "futures-core", + "futures-sink", + "js-sys", + "pin-project-lite", + "thiserror 1.0.69", + "tracing", +] + +[[package]] +name = "opentelemetry-http" +version = "0.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10a8a7f5f6ba7c1b286c2fbca0454eaba116f63bbe69ed250b642d36fbb04d80" +dependencies = [ + "async-trait", + "bytes", + "http", + "opentelemetry", +] + +[[package]] +name = "opentelemetry-otlp" +version = "0.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91cf61a1868dacc576bf2b2a1c3e9ab150af7272909e80085c3173384fe11f76" +dependencies = [ + "async-trait", + "futures-core", + "http", + "opentelemetry", + "opentelemetry-proto", + "opentelemetry_sdk", + "prost", + "thiserror 1.0.69", + "tokio", + "tonic", + "tracing", +] + +[[package]] +name = "opentelemetry-proto" +version = "0.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a6e05acbfada5ec79023c85368af14abd0b307c015e9064d249b2a950ef459a6" +dependencies = [ + "opentelemetry", + "opentelemetry_sdk", + "prost", + "tonic", +] + +[[package]] +name = "opentelemetry_sdk" +version = "0.27.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "231e9d6ceef9b0b2546ddf52335785ce41252bc7474ee8ba05bfad277be13ab8" +dependencies = [ + "async-std", + "async-trait", + "futures-channel", + "futures-executor", + "futures-util", + "glob", + "opentelemetry", + "percent-encoding", + "rand 0.8.5", + "serde_json", + "thiserror 1.0.69", + "tokio", + "tokio-stream", + "tracing", +] + +[[package]] +name = "option-ext" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "04744f49eae99ab78e0d5c0b603ab218f515ea8cfe5a456d7629ad883a3b6e7d" + +[[package]] +name = "ordered-float" +version = "2.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68f19d67e5a2795c94e73e0bb1cc1a7edeb2e28efd39e2e1c9b7a40c1108b11c" +dependencies = [ + "num-traits", +] + +[[package]] +name = "parking" +version = "2.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba" + +[[package]] +name = "parking_lot" +version = "0.12.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93857453250e3077bd71ff98b6a65ea6621a19bb0f559a85248955ac12c45a1a" +dependencies = [ + "lock_api", + "parking_lot_core", +] + +[[package]] +name = "parking_lot_core" +version = "0.9.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" +dependencies = [ + "cfg-if", + "libc", + "redox_syscall", + "smallvec", + "windows-link", +] + +[[package]] +name = "paste" +version = "1.0.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" + +[[package]] +name = "pem" +version = "3.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d30c53c26bc5b31a98cd02d20f25a7c8567146caf63ed593a9d87b2775291be" +dependencies = [ + "base64 0.22.1", + "serde_core", +] + +[[package]] +name = "percent-encoding" +version = "2.3.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b4f627cb1b25917193a259e49bdad08f671f8d9708acfd5fe0a8c1455d87220" + +[[package]] +name = "pest" +version = "2.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e0848c601009d37dfa3430c4666e147e49cdcf1b92ecd3e63657d8a5f19da662" +dependencies = [ + "memchr", + "ucd-trie", +] + +[[package]] +name = "pest_derive" +version = "2.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "11f486f1ea21e6c10ed15d5a7c77165d0ee443402f0780849d1768e7d9d6fe77" +dependencies = [ + "pest", + "pest_generator", +] + +[[package]] +name = "pest_generator" +version = "2.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8040c4647b13b210a963c1ed407c1ff4fdfa01c31d6d2a098218702e6664f94f" +dependencies = [ + "pest", + "pest_meta", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "pest_meta" +version = "2.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "89815c69d36021a140146f26659a81d6c2afa33d216d736dd4be5381a7362220" +dependencies = [ + "pest", + "sha2", +] + +[[package]] +name = "pin-project" +version = "1.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f1749c7ed4bcaf4c3d0a3efc28538844fb29bcdd7d2b67b2be7e20ba861ff517" +dependencies = [ + "pin-project-internal", +] + +[[package]] +name = "pin-project-internal" +version = "1.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d9b20ed30f105399776b9c883e68e536ef602a16ae6f596d2c473591d6ad64c6" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "pin-project-lite" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd" + +[[package]] +name = "pin-utils" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8b870d8c151b6f2fb93e84a13146138f05d02ed11c7e7c54f8826aaaf7c9f184" + +[[package]] +name = "piper" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c835479a4443ded371d6c535cbfd8d31ad92c5d23ae9770a61bc155e4992a3c1" +dependencies = [ + "atomic-waker", + "fastrand", + "futures-io", +] + +[[package]] +name = "pkg-config" +version = "0.3.32" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7edddbd0b52d732b21ad9a5fab5c704c14cd949e5e9a1ec5929a24fded1b904c" + +[[package]] +name = "plotters" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5aeb6f403d7a4911efb1e33402027fc44f29b5bf6def3effcc22d7bb75f2b747" +dependencies = [ + "num-traits", + "plotters-backend", + "plotters-svg", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "plotters-backend" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "df42e13c12958a16b3f7f4386b9ab1f3e7933914ecea48da7139435263a4172a" + +[[package]] +name = "plotters-svg" +version = "0.3.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "51bae2ac328883f7acdfea3d66a7c35751187f870bc81f94563733a154d7a670" +dependencies = [ + "plotters-backend", +] + +[[package]] +name = "polling" +version = "3.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5d0e4f59085d47d8241c88ead0f274e8a0cb551f3625263c05eb8dd897c34218" +dependencies = [ + "cfg-if", + "concurrent-queue", + "hermit-abi", + "pin-project-lite", + "rustix", + "windows-sys 0.61.2", +] + +[[package]] +name = "portable-atomic" +version = "1.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c33a9471896f1c69cecef8d20cbe2f7accd12527ce60845ff44c153bb2a21b49" + +[[package]] +name = "portpicker" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be97d76faf1bfab666e1375477b23fde79eccf0276e9b63b92a39d676a889ba9" +dependencies = [ + "rand 0.8.5", +] + +[[package]] +name = "potential_utf" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b73949432f5e2a09657003c25bca5e19a0e9c84f8058ca374f49e0ebe605af77" +dependencies = [ + "zerovec", +] + +[[package]] +name = "powerfmt" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "439ee305def115ba05938db6eb1644ff94165c5ab5e9420d1c1bcedbba909391" + +[[package]] +name = "ppv-lite86" +version = "0.2.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85eae3c4ed2f50dcfe72643da4befc30deadb458a9b590d720cde2f2b1e97da9" +dependencies = [ + "zerocopy", +] + +[[package]] +name = "prettyplease" +version = "0.2.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "479ca8adacdd7ce8f1fb39ce9ecccbfe93a3f1344b3d0d97f20bc0196208f62b" +dependencies = [ + "proc-macro2", + "syn", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "prost" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2796faa41db3ec313a31f7624d9286acf277b52de526150b7e69f3debf891ee5" +dependencies = [ + "bytes", + "prost-derive", +] + +[[package]] +name = "prost-derive" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" +dependencies = [ + "anyhow", + "itertools 0.14.0", + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "pyo3" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ba0117f4212101ee6544044dae45abe1083d30ce7b29c4b5cbdfa2354e07383" +dependencies = [ + "indoc", + "libc", + "memoffset", + "once_cell", + "portable-atomic", + "pyo3-build-config", + "pyo3-ffi", + "pyo3-macros", + "unindent", +] + +[[package]] +name = "pyo3-build-config" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fc6ddaf24947d12a9aa31ac65431fb1b851b8f4365426e182901eabfb87df5f" +dependencies = [ + "target-lexicon", +] + +[[package]] +name = "pyo3-ffi" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "025474d3928738efb38ac36d4744a74a400c901c7596199e20e45d98eb194105" +dependencies = [ + "libc", + "pyo3-build-config", +] + +[[package]] +name = "pyo3-macros" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e64eb489f22fe1c95911b77c44cc41e7c19f3082fc81cce90f657cdc42ffded" +dependencies = [ + "proc-macro2", + "pyo3-macros-backend", + "quote", + "syn", +] + +[[package]] +name = "pyo3-macros-backend" +version = "0.26.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "100246c0ecf400b475341b8455a9213344569af29a3c841d29270e53102e0fcf" +dependencies = [ + "heck", + "proc-macro2", + "pyo3-build-config", + "quote", + "syn", +] + +[[package]] +name = "quanta" +version = "0.12.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3ab5a9d756f0d97bdc89019bd2e4ea098cf9cde50ee7564dde6b81ccc8f06c7" +dependencies = [ + "crossbeam-utils", + "libc", + "once_cell", + "raw-cpuid", + "wasi", + "web-sys", + "winapi", +] + +[[package]] +name = "quinn" +version = "0.11.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9e20a958963c291dc322d98411f541009df2ced7b5a4f2bd52337638cfccf20" +dependencies = [ + "bytes", + "cfg_aliases", + "pin-project-lite", + "quinn-proto", + "quinn-udp", + "rustc-hash 2.1.1", + "rustls", + "socket2 0.6.3", + "thiserror 2.0.18", + "tokio", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-proto" +version = "0.11.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "434b42fec591c96ef50e21e886936e66d3cc3f737104fdb9b737c40ffb94c098" +dependencies = [ + "aws-lc-rs", + "bytes", + "getrandom 0.3.4", + "lru-slab", + "rand 0.9.2", + "ring", + "rustc-hash 2.1.1", + "rustls", + "rustls-pki-types", + "slab", + "thiserror 2.0.18", + "tinyvec", + "tracing", + "web-time", +] + +[[package]] +name = "quinn-udp" +version = "0.5.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "addec6a0dcad8a8d96a771f815f0eaf55f9d1805756410b39f5fa81332574cbd" +dependencies = [ + "cfg_aliases", + "libc", + "once_cell", + "socket2 0.6.3", + "tracing", + "windows-sys 0.60.2", +] + +[[package]] +name = "quote" +version = "1.0.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41f2619966050689382d2b44f664f4bc593e129785a36d6ee376ddf37259b924" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "r-efi" +version = "5.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "69cdb34c158ceb288df11e18b4bd39de994f6657d83847bdffdbd7f346754b0f" + +[[package]] +name = "r-efi" +version = "6.0.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" + +[[package]] +name = "rand" +version = "0.8.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "34af8d1a0e25924bc5b7c43c079c942339d8f0a8b57c39049bef581b46327404" +dependencies = [ + "libc", + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + +[[package]] +name = "rand" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6db2770f06117d490610c7488547d543617b21bfa07796d7a12f6f1bd53850d1" +dependencies = [ + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core 0.6.4", +] + +[[package]] +name = "rand_chacha" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" +dependencies = [ + "ppv-lite86", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.17", +] + +[[package]] +name = "rand_core" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "76afc826de14238e6e8c374ddcc1fa19e374fd8dd986b0d2af0d02377261d83c" +dependencies = [ + "getrandom 0.3.4", +] + +[[package]] +name = "rand_xoshiro" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f703f4665700daf5512dcca5f43afa6af89f09db47fb56be587f80636bda2d41" +dependencies = [ + "rand_core 0.9.5", +] + +[[package]] +name = "raw-cpuid" +version = "11.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "498cd0dc59d73224351ee52a95fee0f1a617a2eae0e7d9d720cc622c73a54186" +dependencies = [ + "bitflags", +] + +[[package]] +name = "rayon" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "368f01d005bf8fd9b1206fb6fa653e6c4a81ceb1466406b81792d87c5677a58f" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-cond" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2964d0cf57a3e7a06e8183d14a8b527195c706b7983549cd5462d5aa3747438f" +dependencies = [ + "either", + "itertools 0.14.0", + "rayon", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "redox_syscall" +version = "0.5.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ed2bf2547551a7053d6fdfafda3f938979645c44812fbfcda098faae3f1a362d" +dependencies = [ + "bitflags", +] + +[[package]] +name = "redox_users" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4e608c6638b9c18977b00b475ac1f28d14e84b27d8d42f70e0bf1e3dec127ac" +dependencies = [ + "getrandom 0.2.17", + "libredox", + "thiserror 2.0.18", +] + +[[package]] +name = "regex" +version = "1.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e10754a14b9137dd7b1e3e5b0493cc9171fdd105e0ab477f51b72e7f3ac0e276" +dependencies = [ + "aho-corasick", + "memchr", + "regex-automata", + "regex-syntax", +] + +[[package]] +name = "regex-automata" +version = "0.4.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6e1dd4122fc1595e8162618945476892eefca7b88c52820e74af6262213cae8f" +dependencies = [ + "aho-corasick", + "memchr", + "regex-syntax", +] + +[[package]] +name = "regex-syntax" +version = "0.8.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dc897dd8d9e8bd1ed8cdad82b5966c3e0ecae09fb1907d58efaa013543185d0a" + +[[package]] +name = "reqwest" +version = "0.12.28" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eddd3ca559203180a307f12d114c268abf583f59b03cb906fd0b3ff8646c1147" +dependencies = [ + "base64 0.22.1", + "bytes", + "encoding_rs", + "futures-core", + "futures-util", + "h2", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-tls", + "hyper-util", + "js-sys", + "log", + "mime", + "native-tls", + "percent-encoding", + "pin-project-lite", + "rustls-pki-types", + "serde", + "serde_json", + "serde_urlencoded", + "sync_wrapper", + "tokio", + "tokio-native-tls", + "tokio-util", + "tower 0.5.3", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams 0.4.2", + "web-sys", +] + +[[package]] +name = "reqwest" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ab3f43e3283ab1488b624b44b0e988d0acea0b3214e694730a055cb6b2efa801" +dependencies = [ + "base64 0.22.1", + "bytes", + "encoding_rs", + "futures-channel", + "futures-core", + "futures-util", + "h2", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-rustls", + "hyper-util", + "js-sys", + "log", + "mime", + "percent-encoding", + "pin-project-lite", + "quinn", + "rustls", + "rustls-pki-types", + "rustls-platform-verifier", + "serde", + "serde_json", + "sync_wrapper", + "tokio", + "tokio-rustls", + "tokio-util", + "tower 0.5.3", + "tower-http", + "tower-service", + "url", + "wasm-bindgen", + "wasm-bindgen-futures", + "wasm-streams 0.5.0", + "web-sys", +] + +[[package]] +name = "ring" +version = "0.17.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" +dependencies = [ + "cc", + "cfg-if", + "getrandom 0.2.17", + "libc", + "untrusted", + "windows-sys 0.52.0", +] + +[[package]] +name = "rustc-hash" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "08d43f7aa6b08d49f382cde6a7982047c3426db949b1424bc4b7ec9ae12c6ce2" + +[[package]] +name = "rustc-hash" +version = "2.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "357703d41365b4b27c590e3ed91eabb1b663f07c4c084095e60cbed4362dff0d" + +[[package]] +name = "rustc_version" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cfcb3a22ef46e85b45de6ee7e79d063319ebb6594faafcf1c225ea92ab6e9b92" +dependencies = [ + "semver", +] + +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + +[[package]] +name = "rustls" +version = "0.23.37" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "758025cb5fccfd3bc2fd74708fd4682be41d99e5dff73c377c0646c6012c73a4" +dependencies = [ + "aws-lc-rs", + "log", + "once_cell", + "ring", + "rustls-pki-types", + "rustls-webpki", + "subtle", + "zeroize", +] + +[[package]] +name = "rustls-native-certs" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e5bfb394eeed242e909609f56089eecfe5fda225042e8b171791b9c95f5931e5" +dependencies = [ + "openssl-probe 0.1.6", + "rustls-pemfile", + "rustls-pki-types", + "schannel", + "security-framework 2.11.1", +] + +[[package]] +name = "rustls-native-certs" +version = "0.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "612460d5f7bea540c490b2b6395d8e34a953e52b491accd6c86c8164c5932a63" +dependencies = [ + "openssl-probe 0.2.1", + "rustls-pki-types", + "schannel", + "security-framework 3.7.0", +] + +[[package]] +name = "rustls-pemfile" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dce314e5fee3f39953d46bb63bb8a46d40c2f8fb7cc5a3b6cab2bde9721d6e50" +dependencies = [ + "rustls-pki-types", +] + +[[package]] +name = "rustls-pki-types" +version = "1.14.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be040f8b0a225e40375822a563fa9524378b9d63112f53e19ffff34df5d33fdd" +dependencies = [ + "web-time", + "zeroize", +] + +[[package]] +name = "rustls-platform-verifier" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1d99feebc72bae7ab76ba994bb5e121b8d83d910ca40b36e0921f53becc41784" +dependencies = [ + "core-foundation 0.10.1", + "core-foundation-sys", + "jni", + "log", + "once_cell", + "rustls", + "rustls-native-certs 0.8.3", + "rustls-platform-verifier-android", + "rustls-webpki", + "security-framework 3.7.0", + "security-framework-sys", + "webpki-root-certs", + "windows-sys 0.61.2", +] + +[[package]] +name = "rustls-platform-verifier-android" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f87165f0995f63a9fbeea62b64d10b4d9d8e78ec6d7d51fb2125fda7bb36788f" + +[[package]] +name = "rustls-webpki" +version = "0.103.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7df23109aa6c1567d1c575b9952556388da57401e4ace1d15f79eedad0d8f53" +dependencies = [ + "aws-lc-rs", + "ring", + "rustls-pki-types", + "untrusted", +] + +[[package]] +name = "rustversion" +version = "1.0.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b39cdef0fa800fc44525c84ccb54a029961a8215f9619753635a9c0d2538d46d" + +[[package]] +name = "ryu" +version = "1.0.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" + +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "schannel" +version = "0.1.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91c1b7e4904c873ef0710c1f407dde2e6287de2bebc1bbbf7d430bb7cbffd939" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "schemars" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3fbf2ae1b8bc8e02df939598064d22402220cd5bbcca1c76f7d6a310974d5615" +dependencies = [ + "dyn-clone", + "schemars_derive", + "serde", + "serde_json", +] + +[[package]] +name = "schemars_derive" +version = "0.8.22" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32e265784ad618884abaea0600a9adf15393368d840e0222d101a072f3f7534d" +dependencies = [ + "proc-macro2", + "quote", + "serde_derive_internals", + "syn", +] + +[[package]] +name = "scopeguard" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" + +[[package]] +name = "secrecy" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e891af845473308773346dc847b2c23ee78fe442e0472ac50e22a18a93d3ae5a" +dependencies = [ + "zeroize", +] + +[[package]] +name = "security-framework" +version = "2.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "897b2245f0b511c87893af39b033e5ca9cce68824c4d7e7630b5a1d339658d02" +dependencies = [ + "bitflags", + "core-foundation 0.9.4", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework" +version = "3.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7f4bc775c73d9a02cde8bf7b2ec4c9d12743edf609006c7facc23998404cd1d" +dependencies = [ + "bitflags", + "core-foundation 0.10.1", + "core-foundation-sys", + "libc", + "security-framework-sys", +] + +[[package]] +name = "security-framework-sys" +version = "2.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2691df843ecc5d231c0b14ece2acc3efb62c0a398c7e1d875f3983ce020e3" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "semver" +version = "1.0.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d767eb0aabc880b29956c35734170f26ed551a859dbd361d140cdbeca61ab1e2" + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde-value" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3a1a3341211875ef120e117ea7fd5228530ae7e7036a779fdc9117be6b3282c" +dependencies = [ + "ordered-float", + "serde", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_derive_internals" +version = "0.29.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18d26a20a969b9e3fdf2fc2d9f21eda6c40e2de84c9408bb5d3b05d499aae711" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.149" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "83fc039473c5595ace860d8c4fafa220ff474b3fc6bfdb4293327f1a37e94d86" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "serde_path_to_error" +version = "0.1.20" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10a9ff822e371bb5403e391ecd83e182e0e77ba7f6fe0160b795797109d1b457" +dependencies = [ + "itoa", + "serde", + "serde_core", +] + +[[package]] +name = "serde_urlencoded" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3491c14715ca2294c4d6a88f15e84739788c1d030eed8c110436aafdaa2f3fd" +dependencies = [ + "form_urlencoded", + "itoa", + "ryu", + "serde", +] + +[[package]] +name = "serde_yaml" +version = "0.9.34+deprecated" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6a8b1a1a2ebf674015cc02edccce75287f1a0130d394307b36743c2f5d504b47" +dependencies = [ + "indexmap 2.13.0", + "itoa", + "ryu", + "serde", + "unsafe-libyaml", +] + +[[package]] +name = "sha1" +version = "0.10.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e3bf829a2d51ab4a5ddf1352d8470c140cadc8301b2ae1789db023f01cedd6ba" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "sharded-slab" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f40ca3c46823713e0d4209592e8d6e826aa57e928f09752619fc696c499637f6" +dependencies = [ + "lazy_static", +] + +[[package]] +name = "shlex" +version = "1.3.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0fda2ff0d084019ba4d7c6f371c95d8fd75ce3524c3cb8fb653a3023f6323e64" + +[[package]] +name = "signal-hook-registry" +version = "1.4.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b" +dependencies = [ + "errno", + "libc", +] + +[[package]] +name = "simd-adler32" +version = "0.3.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e320a6c5ad31d271ad523dcf3ad13e2767ad8b1cb8f047f75a8aeaf8da139da2" + +[[package]] +name = "sketches-ddsketch" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c6f73aeb92d671e0cc4dca167e59b2deb6387c375391bc99ee743f326994a2b" + +[[package]] +name = "slab" +version = "0.4.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c790de23124f9ab44544d7ac05d60440adc586479ce501c1d6d7da3cd8c9cf5" + +[[package]] +name = "smallvec" +version = "1.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "67b1b7a3b5fe4f1376887184045fcf45c69e92af734b7aaddc05fb777b6fbd03" + +[[package]] +name = "socket2" +version = "0.5.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e22376abed350d73dd1cd119b57ffccad95b4e585a7cda43e286245ce23c0678" +dependencies = [ + "libc", + "windows-sys 0.52.0", +] + +[[package]] +name = "socket2" +version = "0.6.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a766e1110788c36f4fa1c2b71b387a7815aa65f88ce0229841826633d93723e" +dependencies = [ + "libc", + "windows-sys 0.61.2", +] + +[[package]] +name = "socks" +version = "0.3.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0c3dbbd9ae980613c6dd8e28a9407b50509d3803b57624d5dfe8315218cd58b" +dependencies = [ + "byteorder", + "libc", + "winapi", +] + +[[package]] +name = "spm_precompiled" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5851699c4033c63636f7ea4cf7b7c1f1bf06d0cc03cfb42e711de5a5c46cf326" +dependencies = [ + "base64 0.13.1", + "nom", + "serde", + "unicode-segmentation", +] + +[[package]] +name = "stable_deref_trait" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6ce2be8dc25455e1f91df71bfa12ad37d7af1092ae736f3a6cd0e37bc7810596" + +[[package]] +name = "static_assertions" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" + +[[package]] +name = "strsim" +version = "0.11.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7da8b5736845d9f2fcb837ea5d9e2628564b3b043a70948a3f0b778838c5fb4f" + +[[package]] +name = "strum" +version = "0.26.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fec0f0aef304996cf250b31b5a10dee7980c85da9d759361292b8bca5a18f06" +dependencies = [ + "strum_macros", +] + +[[package]] +name = "strum_macros" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c6bee85a5a24955dc440386795aa378cd9cf82acd5f764469152d2270e581be" +dependencies = [ + "heck", + "proc-macro2", + "quote", + "rustversion", + "syn", +] + +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + +[[package]] +name = "syn" +version = "2.0.117" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e665b8803e7b1d2a727f4023456bbbbe74da67099c585258af0ad9c5013b9b99" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "sync_wrapper" +version = "1.0.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0bf256ce5efdfa370213c1dabab5935a12e49f2c58d15e9eac2870d3b4f27263" +dependencies = [ + "futures-core", +] + +[[package]] +name = "synstructure" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "728a70f3dbaf5bab7f0c4b1ac8d7ae5ea60a4b5549c8a5914361c99147a709d2" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "system-configuration" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a13f3d0daba03132c0aa9767f98351b3488edc2c100cda2d2ec2b04f3d8d3c8b" +dependencies = [ + "bitflags", + "core-foundation 0.9.4", + "system-configuration-sys", +] + +[[package]] +name = "system-configuration-sys" +version = "0.6.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e1d1b10ced5ca923a1fcb8d03e96b8d3268065d724548c0211415ff6ac6bac4" +dependencies = [ + "core-foundation-sys", + "libc", +] + +[[package]] +name = "target-lexicon" +version = "0.13.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "adb6935a6f5c20170eeceb1a3835a49e12e19d792f6dd344ccc76a985ca5a6ca" + +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.4.2", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + +[[package]] +name = "thiserror" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6aaf5339b578ea85b50e080feb250a3e8ae8cfcdff9a461c9ec2904bc923f52" +dependencies = [ + "thiserror-impl 1.0.69", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl 2.0.18", +] + +[[package]] +name = "thiserror-impl" +version = "1.0.69" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4fee6c4efc90059e10f81e6d42c60a18f76588c3d74cb83a0b242a2b6c7504c1" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "thread_local" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f60246a4944f24f6e018aa17cdeffb7818b76356965d03b07d6a9886e8962185" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "tiktoken-rs" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "25563eeba904d770acf527e8b370fe9a5547bacd20ff84a0b6c3bc41288e5625" +dependencies = [ + "anyhow", + "base64 0.22.1", + "bstr", + "fancy-regex", + "lazy_static", + "regex", + "rustc-hash 1.1.0", +] + +[[package]] +name = "time" +version = "0.3.47" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "743bd48c283afc0388f9b8827b976905fb217ad9e647fae3a379a9283c4def2c" +dependencies = [ + "deranged", + "itoa", + "num-conv", + "powerfmt", + "serde_core", + "time-core", + "time-macros", +] + +[[package]] +name = "time-core" +version = "0.1.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7694e1cfe791f8d31026952abf09c69ca6f6fa4e1a1229e18988f06a04a12dca" + +[[package]] +name = "time-macros" +version = "0.2.27" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2e70e4c5a0e0a8a4823ad65dfe1a6930e4f4d756dcd9dd7939022b5e8c501215" +dependencies = [ + "num-conv", + "time-core", +] + +[[package]] +name = "tinystr" +version = "0.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "42d3e9c45c09de15d06dd8acf5f4e0e399e85927b7f00711024eb7ae10fa4869" +dependencies = [ + "displaydoc", + "zerovec", +] + +[[package]] +name = "tinytemplate" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "be4d6b5f19ff7664e8c98d03e2139cb510db9b0a60b55f8e8709b689d939b6bc" +dependencies = [ + "serde", + "serde_json", +] + +[[package]] +name = "tinyvec" +version = "1.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3e61e67053d25a4e82c844e8424039d9745781b3fc4f32b8d55ed50f5f667ef3" +dependencies = [ + "tinyvec_macros", +] + +[[package]] +name = "tinyvec_macros" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f3ccbac311fea05f86f61904b462b55fb3df8837a366dfc601a0161d0532f20" + +[[package]] +name = "tokenizers" +version = "0.22.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b238e22d44a15349529690fb07bd645cf58149a1b1e44d6cb5bd1641ff1a6223" +dependencies = [ + "ahash", + "aho-corasick", + "compact_str", + "dary_heap", + "derive_builder", + "esaxx-rs", + "getrandom 0.3.4", + "indicatif 0.18.4", + "itertools 0.14.0", + "log", + "macro_rules_attribute", + "monostate", + "onig", + "paste", + "rand 0.9.2", + "rayon", + "rayon-cond", + "regex", + "regex-syntax", + "serde", + "serde_json", + "spm_precompiled", + "thiserror 2.0.18", + "unicode-normalization-alignments", + "unicode-segmentation", + "unicode_categories", +] + +[[package]] +name = "tokio" +version = "1.50.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "27ad5e34374e03cfffefc301becb44e9dc3c17584f414349ebe29ed26661822d" +dependencies = [ + "bytes", + "libc", + "mio", + "parking_lot", + "pin-project-lite", + "signal-hook-registry", + "socket2 0.6.3", + "tokio-macros", + "windows-sys 0.61.2", +] + +[[package]] +name = "tokio-macros" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c55a2eff8b69ce66c84f85e1da1c233edc36ceb85a2058d11b0d6a3c7e7569c" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tokio-native-tls" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bbae76ab933c85776efabc971569dd6119c580d8f5d448769dec1764bf796ef2" +dependencies = [ + "native-tls", + "tokio", +] + +[[package]] +name = "tokio-rustls" +version = "0.26.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1729aa945f29d91ba541258c8df89027d5792d85a8841fb65e8bf0f4ede4ef61" +dependencies = [ + "rustls", + "tokio", +] + +[[package]] +name = "tokio-stream" +version = "0.1.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32da49809aab5c3bc678af03902d4ccddea2a87d028d86392a4b1560c6906c70" +dependencies = [ + "futures-core", + "pin-project-lite", + "tokio", + "tokio-util", +] + +[[package]] +name = "tokio-tungstenite" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d25a406cddcc431a75d3d9afc6a7c0f7428d4891dd973e4d54c56b46127bf857" +dependencies = [ + "futures-util", + "log", + "tokio", + "tungstenite", +] + +[[package]] +name = "tokio-util" +version = "0.7.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ae9cec805b01e8fc3fd2fe289f89149a9b66dd16786abd8b19cfa7b48cb0098" +dependencies = [ + "bytes", + "futures-core", + "futures-sink", + "pin-project-lite", + "slab", + "tokio", +] + +[[package]] +name = "tonic" +version = "0.12.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "877c5b330756d856ffcc4553ab34a5684481ade925ecc54bcd1bf02b1d0d4d52" +dependencies = [ + "async-stream", + "async-trait", + "axum 0.7.9", + "base64 0.22.1", + "bytes", + "h2", + "http", + "http-body", + "http-body-util", + "hyper", + "hyper-timeout", + "hyper-util", + "percent-encoding", + "pin-project", + "prost", + "socket2 0.5.10", + "tokio", + "tokio-stream", + "tower 0.4.13", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tower" +version = "0.4.13" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8fa9be0de6cf49e536ce1851f987bd21a43b771b09473c3549a6c853db37c1c" +dependencies = [ + "futures-core", + "futures-util", + "indexmap 1.9.3", + "pin-project", + "pin-project-lite", + "rand 0.8.5", + "slab", + "tokio", + "tokio-util", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tower" +version = "0.5.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebe5ef63511595f1344e2d5cfa636d973292adc0eec1f0ad45fae9f0851ab1d4" +dependencies = [ + "futures-core", + "futures-util", + "hdrhistogram", + "indexmap 2.13.0", + "pin-project-lite", + "slab", + "sync_wrapper", + "tokio", + "tokio-util", + "tower-layer", + "tower-service", + "tracing", +] + +[[package]] +name = "tower-http" +version = "0.6.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4e6559d53cc268e5031cd8429d05415bc4cb4aefc4aa5d6cc35fbf5b924a1f8" +dependencies = [ + "async-compression", + "base64 0.22.1", + "bitflags", + "bytes", + "futures-core", + "futures-util", + "http", + "http-body", + "http-body-util", + "iri-string", + "mime", + "pin-project-lite", + "tokio", + "tokio-util", + "tower 0.5.3", + "tower-layer", + "tower-service", + "tracing", + "uuid", +] + +[[package]] +name = "tower-layer" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "121c2a6cda46980bb0fcd1647ffaf6cd3fc79a013de288782836f6df9c48780e" + +[[package]] +name = "tower-service" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8df9b6e13f2d32c91b9bd719c00d1958837bc7dec474d94952798cc8e69eeec3" + +[[package]] +name = "tracing" +version = "0.1.44" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63e71662fa4b2a2c3a26f570f037eb95bb1f85397f3cd8076caed2f026a6d100" +dependencies = [ + "log", + "pin-project-lite", + "tracing-attributes", + "tracing-core", +] + +[[package]] +name = "tracing-appender" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "786d480bce6247ab75f005b14ae1624ad978d3029d9113f0a22fa1ac773faeaf" +dependencies = [ + "crossbeam-channel", + "thiserror 2.0.18", + "time", + "tracing-subscriber", +] + +[[package]] +name = "tracing-attributes" +version = "0.1.31" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7490cfa5ec963746568740651ac6781f701c9c5ea257c58e057f3ba8cf69e8da" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "tracing-core" +version = "0.1.36" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "db97caf9d906fbde555dd62fa95ddba9eecfd14cb388e4f491a66d74cd5fb79a" +dependencies = [ + "once_cell", + "valuable", +] + +[[package]] +name = "tracing-log" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee855f1f400bd0e5c02d150ae5de3840039a3f54b025156404e34c23c03f47c3" +dependencies = [ + "log", + "once_cell", + "tracing-core", +] + +[[package]] +name = "tracing-opentelemetry" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97a971f6058498b5c0f1affa23e7ea202057a7301dbff68e968b2d578bcbd053" +dependencies = [ + "js-sys", + "once_cell", + "opentelemetry", + "opentelemetry_sdk", + "smallvec", + "tracing", + "tracing-core", + "tracing-log", + "tracing-subscriber", + "web-time", +] + +[[package]] +name = "tracing-serde" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "704b1aeb7be0d0a84fc9828cae51dab5970fee5088f83d1dd7ee6f6246fc6ff1" +dependencies = [ + "serde", + "tracing-core", +] + +[[package]] +name = "tracing-subscriber" +version = "0.3.23" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cb7f578e5945fb242538965c2d0b04418d38ec25c79d160cd279bf0731c8d319" +dependencies = [ + "chrono", + "matchers", + "nu-ansi-term", + "once_cell", + "regex-automata", + "serde", + "serde_json", + "sharded-slab", + "smallvec", + "thread_local", + "tracing", + "tracing-core", + "tracing-log", + "tracing-serde", +] + +[[package]] +name = "try-lock" +version = "0.2.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" + +[[package]] +name = "tungstenite" +version = "0.28.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8628dcc84e5a09eb3d8423d6cb682965dea9133204e8fb3efee74c2a0c259442" +dependencies = [ + "bytes", + "data-encoding", + "http", + "httparse", + "log", + "rand 0.9.2", + "sha1", + "thiserror 2.0.18", + "utf-8", +] + +[[package]] +name = "typenum" +version = "1.19.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "562d481066bde0658276a35467c4af00bdc6ee726305698a55b86e61d7ad82bb" + +[[package]] +name = "ucd-trie" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2896d95c02a80c6d6a5d6e953d479f5ddf2dfdb6a244441010e373ac0fb88971" + +[[package]] +name = "ulid" +version = "1.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "470dbf6591da1b39d43c14523b2b469c86879a53e8b758c8e090a470fe7b1fbe" +dependencies = [ + "rand 0.9.2", + "web-time", +] + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "unicode-normalization-alignments" +version = "0.1.12" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "43f613e4fa046e69818dd287fdc4bc78175ff20331479dab6e1b0f98d57062de" +dependencies = [ + "smallvec", +] + +[[package]] +name = "unicode-segmentation" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f6ccf251212114b54433ec949fd6a7841275f9ada20dddd2f29e9ceea4501493" + +[[package]] +name = "unicode-width" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" + +[[package]] +name = "unicode-xid" +version = "0.2.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc1c04c71510c7f702b52b7c350734c9ff1295c464a03335b00bb84fc54f853" + +[[package]] +name = "unicode_categories" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39ec24b3121d976906ece63c9daad25b85969647682eee313cb5779fdd69e14e" + +[[package]] +name = "unindent" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7264e107f553ccae879d21fbea1d6724ac785e8c3bfc762137959b5802826ef3" + +[[package]] +name = "unit-prefix" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "81e544489bf3d8ef66c953931f56617f423cd4b5494be343d9b9d3dda037b9a3" + +[[package]] +name = "unsafe-libyaml" +version = "0.2.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "673aac59facbab8a9007c7f6108d11f63b603f7cabff99fabf650fea5c32b861" + +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + +[[package]] +name = "ureq" +version = "2.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02d1a66277ed75f640d608235660df48c8e3c19f3b4edb6a263315626cc3c01d" +dependencies = [ + "base64 0.22.1", + "flate2", + "log", + "native-tls", + "once_cell", + "rustls", + "rustls-pki-types", + "serde", + "serde_json", + "socks", + "url", + "webpki-roots 0.26.11", +] + +[[package]] +name = "url" +version = "2.5.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff67a8a4397373c3ef660812acab3268222035010ab8680ec4215f38ba3d0eed" +dependencies = [ + "form_urlencoded", + "idna", + "percent-encoding", + "serde", +] + +[[package]] +name = "utf-8" +version = "0.7.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09cc8ee72d2a9becf2f2febe0205bbed8fc6615b7cb429ad062dc7b7ddd036a9" + +[[package]] +name = "utf8_iter" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" + +[[package]] +name = "utf8parse" +version = "0.2.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" + +[[package]] +name = "uuid" +version = "1.22.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a68d3c8f01c0cfa54a75291d83601161799e4a89a39e0929f4b0354d88757a37" +dependencies = [ + "getrandom 0.4.2", + "js-sys", + "serde_core", + "wasm-bindgen", +] + +[[package]] +name = "valuable" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ba73ea9cf16a25df0c8caa16c51acb937d5712a8429db78a3ee29d5dcacd3a65" + +[[package]] +name = "value-bag" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7ba6f5989077681266825251a52748b8c1d8a4ad098cc37e440103d0ea717fc0" + +[[package]] +name = "vcpkg" +version = "0.2.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "accd4ea62f7bb7a82fe23066fb0957d48ef677f6eeb8215f372f52e48bb32426" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + +[[package]] +name = "want" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bfa7760aed19e106de2c7c0b581b509f2f25d3dacaf737cb82ac61bc6d760b0e" +dependencies = [ + "try-lock", +] + +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + +[[package]] +name = "wasip2" +version = "1.0.2+wasi-0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9517f9239f02c069db75e65f174b3da828fe5f5b945c4dd26bd25d89c03ebcf5" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wasip3" +version = "0.4.0+wasi-0.3.0-rc-2026-01-06" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5428f8bf88ea5ddc08faddef2ac4a67e390b88186c703ce6dbd955e1c145aca5" +dependencies = [ + "wit-bindgen", +] + +[[package]] +name = "wasm-bindgen" +version = "0.2.114" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6532f9a5c1ece3798cb1c2cfdba640b9b3ba884f5db45973a6f442510a87d38e" +dependencies = [ + "cfg-if", + "once_cell", + "rustversion", + "wasm-bindgen-macro", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-futures" +version = "0.4.64" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e9c5522b3a28661442748e09d40924dfb9ca614b21c00d3fd135720e48b67db8" +dependencies = [ + "cfg-if", + "futures-util", + "js-sys", + "once_cell", + "wasm-bindgen", + "web-sys", +] + +[[package]] +name = "wasm-bindgen-macro" +version = "0.2.114" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "18a2d50fcf105fb33bb15f00e7a77b772945a2ee45dcf454961fd843e74c18e6" +dependencies = [ + "quote", + "wasm-bindgen-macro-support", +] + +[[package]] +name = "wasm-bindgen-macro-support" +version = "0.2.114" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "03ce4caeaac547cdf713d280eda22a730824dd11e6b8c3ca9e42247b25c631e3" +dependencies = [ + "bumpalo", + "proc-macro2", + "quote", + "syn", + "wasm-bindgen-shared", +] + +[[package]] +name = "wasm-bindgen-shared" +version = "0.2.114" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75a326b8c223ee17883a4251907455a2431acc2791c98c26279376490c378c16" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "wasm-encoder" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "990065f2fe63003fe337b932cfb5e3b80e0b4d0f5ff650e6985b1048f62c8319" +dependencies = [ + "leb128fmt", + "wasmparser", +] + +[[package]] +name = "wasm-metadata" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "bb0e353e6a2fbdc176932bbaab493762eb1255a7900fe0fea1a2f96c296cc909" +dependencies = [ + "anyhow", + "indexmap 2.13.0", + "wasm-encoder", + "wasmparser", +] + +[[package]] +name = "wasm-streams" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + +[[package]] +name = "wasm-streams" +version = "0.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d1ec4f6517c9e11ae630e200b2b65d193279042e28edd4a2cda233e46670bbb" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + +[[package]] +name = "wasmparser" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "47b807c72e1bac69382b3a6fb3dbe8ea4c0ed87ff5629b8685ae6b9a611028fe" +dependencies = [ + "bitflags", + "hashbrown 0.15.5", + "indexmap 2.13.0", + "semver", +] + +[[package]] +name = "web-sys" +version = "0.3.91" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "854ba17bb104abfb26ba36da9729addc7ce7f06f5c0f90f3c391f8461cca21f9" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "web-time" +version = "1.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5a6580f308b1fad9207618087a65c04e7a10bc77e02c8e84e9b00dd4b12fa0bb" +dependencies = [ + "js-sys", + "wasm-bindgen", +] + +[[package]] +name = "webpki-root-certs" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "804f18a4ac2676ffb4e8b5b5fa9ae38af06df08162314f96a68d2a363e21a8ca" +dependencies = [ + "rustls-pki-types", +] + +[[package]] +name = "webpki-roots" +version = "0.26.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "521bc38abb08001b01866da9f51eb7c5d647a19260e00054a8c7fd5f9e57f7a9" +dependencies = [ + "webpki-roots 1.0.6", +] + +[[package]] +name = "webpki-roots" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22cfaf3c063993ff62e73cb4311efde4db1efb31ab78a3e5c457939ad5cc0bed" +dependencies = [ + "rustls-pki-types", +] + +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys 0.61.2", +] + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + +[[package]] +name = "windows-core" +version = "0.62.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8e83a14d34d0623b51dce9581199302a221863196a1dde71a7663a4c2be9deb" +dependencies = [ + "windows-implement", + "windows-interface", + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-implement" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "053e2e040ab57b9dc951b72c264860db7eb3b0200ba345b4e4c3b14f67855ddf" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "windows-interface" +version = "0.59.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3f316c4a2570ba26bbec722032c4099d8c8bc095efccdc15688708623367e358" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-registry" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "02752bf7fbdcce7f2a27a742f798510f3e5ad88dbe84871e5168e2120c3d5720" +dependencies = [ + "windows-link", + "windows-result", + "windows-strings", +] + +[[package]] +name = "windows-result" +version = "0.4.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7781fa89eaf60850ac3d2da7af8e5242a5ea78d1a11c49bf2910bb5a73853eb5" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-strings" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7837d08f69c77cf6b07689544538e017c1bfcf57e34b4c0ff58e6c2cd3b37091" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-sys" +version = "0.45.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "75283be5efb2831d37ea142365f009c02ec203cd29a3ebecbc093d52315b66d0" +dependencies = [ + "windows-targets 0.42.2", +] + +[[package]] +name = "windows-sys" +version = "0.52.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "282be5f36a8ce781fad8c8ae18fa3f9beff57ec1b52cb3de0789201425d9a33d" +dependencies = [ + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-sys" +version = "0.59.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e38bc4d79ed67fd075bcc251a1c39b32a1776bbe92e5bef1f0bf1f8c531853b" +dependencies = [ + "windows-targets 0.52.6", +] + +[[package]] +name = "windows-sys" +version = "0.60.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2f500e4d28234f72040990ec9d39e3a6b950f9f22d3dba18416c35882612bcb" +dependencies = [ + "windows-targets 0.53.5", +] + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "windows-targets" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e5180c00cd44c9b1c88adb3693291f1cd93605ded80c250a75d472756b4d071" +dependencies = [ + "windows_aarch64_gnullvm 0.42.2", + "windows_aarch64_msvc 0.42.2", + "windows_i686_gnu 0.42.2", + "windows_i686_msvc 0.42.2", + "windows_x86_64_gnu 0.42.2", + "windows_x86_64_gnullvm 0.42.2", + "windows_x86_64_msvc 0.42.2", +] + +[[package]] +name = "windows-targets" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9b724f72796e036ab90c1021d4780d4d3d648aca59e491e6b98e725b84e99973" +dependencies = [ + "windows_aarch64_gnullvm 0.52.6", + "windows_aarch64_msvc 0.52.6", + "windows_i686_gnu 0.52.6", + "windows_i686_gnullvm 0.52.6", + "windows_i686_msvc 0.52.6", + "windows_x86_64_gnu 0.52.6", + "windows_x86_64_gnullvm 0.52.6", + "windows_x86_64_msvc 0.52.6", +] + +[[package]] +name = "windows-targets" +version = "0.53.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4945f9f551b88e0d65f3db0bc25c33b8acea4d9e41163edf90dcd0b19f9069f3" +dependencies = [ + "windows-link", + "windows_aarch64_gnullvm 0.53.1", + "windows_aarch64_msvc 0.53.1", + "windows_i686_gnu 0.53.1", + "windows_i686_gnullvm 0.53.1", + "windows_i686_msvc 0.53.1", + "windows_x86_64_gnu 0.53.1", + "windows_x86_64_gnullvm 0.53.1", + "windows_x86_64_msvc 0.53.1", +] + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "597a5118570b68bc08d8d59125332c54f1ba9d9adeedeef5b99b02ba2b0698f8" + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a4622180e7a0ec044bb555404c800bc9fd9ec262ec147edd5989ccd0c02cd3" + +[[package]] +name = "windows_aarch64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a9d8416fa8b42f5c947f8482c43e7d89e73a173cead56d044f6a56104a6d1b53" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e08e8864a60f06ef0d0ff4ba04124db8b0fb3be5776a5cd47641e942e58c4d43" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "09ec2a7bb152e2252b53fa7803150007879548bc709c039df7627cabbd05d469" + +[[package]] +name = "windows_aarch64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b9d782e804c2f632e395708e99a94275910eb9100b2114651e04744e9b125006" + +[[package]] +name = "windows_i686_gnu" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c61d927d8da41da96a81f029489353e68739737d3beca43145c8afec9a31a84f" + +[[package]] +name = "windows_i686_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8e9b5ad5ab802e97eb8e295ac6720e509ee4c243f69d781394014ebfe8bbfa0b" + +[[package]] +name = "windows_i686_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "960e6da069d81e09becb0ca57a65220ddff016ff2d6af6a223cf372a506593a3" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0eee52d38c090b3caa76c563b86c3a4bd71ef1a819287c19d586d7334ae8ed66" + +[[package]] +name = "windows_i686_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fa7359d10048f68ab8b09fa71c3daccfb0e9b559aed648a8f95469c27057180c" + +[[package]] +name = "windows_i686_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "44d840b6ec649f480a41c8d80f9c65108b92d89345dd94027bfe06ac444d1060" + +[[package]] +name = "windows_i686_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "240948bc05c5e7c6dabba28bf89d89ffce3e303022809e73deaefe4f6ec56c66" + +[[package]] +name = "windows_i686_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1e7ac75179f18232fe9c285163565a57ef8d3c89254a30685b57d83a38d326c2" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8de912b8b8feb55c064867cf047dda097f92d51efad5b491dfb98f6bbb70cb36" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "147a5c80aabfbf0c7d901cb5895d1de30ef2907eb21fbbab29ca94c5b08b1a78" + +[[package]] +name = "windows_x86_64_gnu" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9c3842cdd74a865a8066ab39c8a7a473c0778a3f29370b5fd6b4b9aa7df4a499" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "26d41b46a36d453748aedef1486d5c7a85db22e56aff34643984ea85514e94a3" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d5b23dc417412679681396f2b49f3de8c1473deb516bd34410872eff51ed0d" + +[[package]] +name = "windows_x86_64_gnullvm" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0ffa179e2d07eee8ad8f57493436566c7cc30ac536a3379fdf008f47f6bb7ae1" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.42.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9aec5da331524158c6d1a4ac0ab1541149c0b9505fde06423b02f5ef0106b9f0" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.52.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "589f6da84c646204747d1270a2a5661ea66ed1cced2631d546fdfb155959f9ec" + +[[package]] +name = "windows_x86_64_msvc" +version = "0.53.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" + +[[package]] +name = "wit-bindgen" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d7249219f66ced02969388cf2bb044a09756a083d0fab1e566056b04d9fbcaa5" +dependencies = [ + "wit-bindgen-rust-macro", +] + +[[package]] +name = "wit-bindgen-core" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ea61de684c3ea68cb082b7a88508a8b27fcc8b797d738bfc99a82facf1d752dc" +dependencies = [ + "anyhow", + "heck", + "wit-parser", +] + +[[package]] +name = "wit-bindgen-rust" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b7c566e0f4b284dd6561c786d9cb0142da491f46a9fbed79ea69cdad5db17f21" +dependencies = [ + "anyhow", + "heck", + "indexmap 2.13.0", + "prettyplease", + "syn", + "wasm-metadata", + "wit-bindgen-core", + "wit-component", +] + +[[package]] +name = "wit-bindgen-rust-macro" +version = "0.51.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0c0f9bfd77e6a48eccf51359e3ae77140a7f50b1e2ebfe62422d8afdaffab17a" +dependencies = [ + "anyhow", + "prettyplease", + "proc-macro2", + "quote", + "syn", + "wit-bindgen-core", + "wit-bindgen-rust", +] + +[[package]] +name = "wit-component" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d66ea20e9553b30172b5e831994e35fbde2d165325bec84fc43dbf6f4eb9cb2" +dependencies = [ + "anyhow", + "bitflags", + "indexmap 2.13.0", + "log", + "serde", + "serde_derive", + "serde_json", + "wasm-encoder", + "wasm-metadata", + "wasmparser", + "wit-parser", +] + +[[package]] +name = "wit-parser" +version = "0.244.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ecc8ac4bc1dc3381b7f59c34f00b67e18f910c2c0f50015669dde7def656a736" +dependencies = [ + "anyhow", + "id-arena", + "indexmap 2.13.0", + "log", + "semver", + "serde", + "serde_derive", + "serde_json", + "unicode-xid", + "wasmparser", +] + +[[package]] +name = "writeable" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9edde0db4769d2dc68579893f2306b26c6ecfbe0ef499b013d731b7b9247e0b9" + +[[package]] +name = "yoke" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "72d6e5c6afb84d73944e5cedb052c4680d5657337201555f9f2a16b7406d4954" +dependencies = [ + "stable_deref_trait", + "yoke-derive", + "zerofrom", +] + +[[package]] +name = "yoke-derive" +version = "0.8.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b659052874eb698efe5b9e8cf382204678a0086ebf46982b79d6ca3182927e5d" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zerocopy" +version = "0.8.42" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f2578b716f8a7a858b7f02d5bd870c14bf4ddbbcf3a4c05414ba6503640505e3" +dependencies = [ + "zerocopy-derive", +] + +[[package]] +name = "zerocopy-derive" +version = "0.8.42" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7e6cc098ea4d3bd6246687de65af3f920c430e236bee1e3bf2e441463f08a02f" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zerofrom" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "50cc42e0333e05660c3587f3bf9d0478688e15d870fab3346451ce7f8c9fbea5" +dependencies = [ + "zerofrom-derive", +] + +[[package]] +name = "zerofrom-derive" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d71e5d6e06ab090c67b5e44993ec16b72dcbaabc526db883a360057678b48502" +dependencies = [ + "proc-macro2", + "quote", + "syn", + "synstructure", +] + +[[package]] +name = "zeroize" +version = "1.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b97154e67e32c85465826e8bcc1c59429aaaf107c1e4a9e53c8d8ccd5eff88d0" + +[[package]] +name = "zerotrie" +version = "0.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a59c17a5562d507e4b54960e8569ebee33bee890c70aa3fe7b97e85a9fd7851" +dependencies = [ + "displaydoc", + "yoke", + "zerofrom", +] + +[[package]] +name = "zerovec" +version = "0.11.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c28719294829477f525be0186d13efa9a3c602f7ec202ca9e353d310fb9a002" +dependencies = [ + "yoke", + "zerofrom", + "zerovec-derive", +] + +[[package]] +name = "zerovec-derive" +version = "0.11.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "eadce39539ca5cb3985590102671f2567e659fca9666581ad3411d59207951f3" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" diff --git a/lmdeploy/router/Cargo.toml b/lmdeploy/router/Cargo.toml new file mode 100644 index 0000000000..282a36b12a --- /dev/null +++ b/lmdeploy/router/Cargo.toml @@ -0,0 +1,129 @@ +[package] +name = "lmdeploy_router_rs" +version = "0.0.3" +edition = "2021" + +[lib] +name = "lmdeploy_router_rs" +# Pure Rust library: Just omit crate-type (defaults to rlib) +# Python/C binding + Rust library: Use ["cdylib", "rlib"] +crate-type = ["cdylib", "rlib"] + +[features] +default = ["python"] +python = ["dep:pyo3"] + +[[bin]] +name = "lmdeploy-router" +path = "src/main.rs" + +[dependencies] +clap = { version = "4", features = ["derive"] } +axum = { version = "0.8.4", features = ["macros", "ws", "tracing"] } +tower = { version = "0.5", features = ["full"] } +tower-http = { version = "0.6", features = [ + "trace", + "compression-gzip", + "cors", + "timeout", + "limit", + "request-id", + "util", +] } +serde = { version = "1.0", features = ["derive"] } +serde_json = "1.0" +bytes = "1.8.0" +rand = "0.9.2" +reqwest = { version = "0.13", features = ["stream", "blocking", "json"] } +futures-util = "0.3" +futures = "0.3" +pyo3 = { version = "0.26", features = ["extension-module"], optional = true } +dashmap = "6.1.0" +http = "1.1.0" +tokio = { version = "1.42.0", features = ["full"] } +async-trait = "0.1" +tracing = "0.1" +tracing-subscriber = { version = "0.3", features = [ + "env-filter", + "json", + "chrono", +] } +tracing-log = "0.2" +tracing-appender = "0.2.3" +opentelemetry = "0.27" +opentelemetry_sdk = { version = "0.27", features = ["trace", "rt-tokio"] } +opentelemetry-otlp = { version = "0.27", features = ["trace", "grpc-tonic"] } +tracing-opentelemetry = "0.28" +opentelemetry-http = "0.27" +chrono = "0.4" +kube = { version = "1.1.0", features = ["runtime", "derive"] } +k8s-openapi = { version = "0.25.0", features = ["v1_33"] } +metrics = "0.24.2" +metrics-exporter-prometheus = "0.17.0" +uuid = { version = "1.10", features = ["v4", "serde"] } +ulid = "1.2.1" +parking_lot = "0.12.4" +thiserror = "2.0.12" +regex = "1.10" +url = "2.5.4" +tokio-stream = { version = "0.1", features = ["sync"] } +anyhow = "1.0" +dotenvy = "0.15" +tokenizers = { version = "0.22.2" } +tiktoken-rs = { version = "0.7.0" } +minijinja = { version = "2.0" } +rustls = { version = "0.23", default-features = false, features = [ + "ring", + "std", +] } +hf-hub = { version = "0.4.3", features = ["tokio"] } + +deadpool = { version = "0.12", features = ["managed", "rt_tokio_1"] } +backoff = { version = "0.4", features = ["tokio"] } +strum = { version = "0.26", features = ["derive"] } +once_cell = "1.21.3" + +[dev-dependencies] +criterion = { version = "0.5", features = ["html_reports"] } +tower = { version = "0.5", features = ["util"] } +http-body-util = "0.1" +portpicker = "0.1" +tempfile = "3.8" +lazy_static = "1.4" +opentelemetry_sdk = { version = "0.27", features = ["testing"] } + +[[bench]] +name = "request_processing" +harness = false +path = "benches/request_processing.rs" + +[[bench]] +name = "tokenizer_benchmark" +harness = false +path = "benches/tokenizer_benchmark.rs" + +[[bench]] +name = "otel_disabled_path" +harness = false +path = "benches/otel_disabled_path.rs" + +[profile.release] +lto = "thin" +codegen-units = 1 +strip = true + +[profile.dev] +opt-level = 0 +debug = true +split-debuginfo = "unpacked" +incremental = true + + +[profile.dev.build-override] +opt-level = 3 +codegen-units = 1 +strip = true + +[profile.dev-opt] +inherits = "dev" +opt-level = 1 diff --git a/lmdeploy/router/MANIFEST.in b/lmdeploy/router/MANIFEST.in new file mode 100644 index 0000000000..e1d6e7a901 --- /dev/null +++ b/lmdeploy/router/MANIFEST.in @@ -0,0 +1,3 @@ +# Must include: +include Cargo.toml # Rust project configuration +recursive-include src *.rs # Rust source files diff --git a/lmdeploy/router/Makefile b/lmdeploy/router/Makefile new file mode 100644 index 0000000000..18c69f6dc4 --- /dev/null +++ b/lmdeploy/router/Makefile @@ -0,0 +1,126 @@ +# LMDeploy Router Makefile +# Provides convenient shortcuts for common development tasks + +# Check if sccache is available and set RUSTC_WRAPPER accordingly +SCCACHE := $(shell which sccache 2>/dev/null) +ifdef SCCACHE + export RUSTC_WRAPPER := $(SCCACHE) + $(info Using sccache for compilation caching) +else + $(info sccache not found. Install it for faster builds: cargo install sccache) +endif + +.PHONY: help test build clean docs check fmt dev-setup bench-ci bench-report bench-clean perf-monitor setup-sccache sccache-stats sccache-clean sccache-stop + +help: ## Show this help message + @echo "LMDeploy Router Development Commands" + @echo "==================================" + @echo "" + @grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | sort | awk 'BEGIN {FS = ":.*?## "}; {printf " \033[36m%-20s\033[0m %s\n", $$1, $$2}' + @echo "" + +build: ## Build the project in release mode + @echo "Building LMDeploy Router..." + @cargo build --release + +test: ## Run all tests + @echo "Running tests..." + @cargo test + +bench: ## Run full benchmark suite + @cargo bench + +bench-quick: ## Run quick benchmark benchmarks + @cargo bench --bench request_processing -- --quick + +bench-baseline: ## Save current performance as baseline + @cargo bench -- --save-baseline main + +bench-compare: ## Compare with saved baseline + @cargo bench -- --baseline main + +bench-ci: ## Run benchmarks suitable for CI + @cargo bench --bench request_processing -- --quick + +clean: ## Clean build artifacts + @echo "Cleaning build artifacts..." + @cargo clean + +docs: ## Generate and open documentation + @echo "Generating documentation..." + @cargo doc --open + +check: ## Run cargo check and clippy with CI settings + @echo "Running cargo check..." + @cargo check --locked --all-targets + @echo "Running clippy..." + @cargo clippy --all-targets --locked -- -D warnings + +fmt: ## Format code with rustfmt + @echo "Formatting code..." + @cargo fmt + +# Development workflow shortcuts +dev-setup: build test ## Set up development environment + @echo "Development environment ready!" + +pre-commit: fmt check test bench-quick ## Run router pre-commit checks + @echo "Pre-commit checks passed!" + +# Benchmark analysis shortcuts +bench-report: ## Open benchmark HTML report + @if [ -f "target/criterion/request_processing/report/index.html" ]; then \ + echo "Opening benchmark report..."; \ + if command -v xdg-open >/dev/null 2>&1; then \ + xdg-open target/criterion/request_processing/report/index.html; \ + elif command -v open >/dev/null 2>&1; then \ + open target/criterion/request_processing/report/index.html; \ + else \ + echo "Please open target/criterion/request_processing/report/index.html in your browser"; \ + fi \ + else \ + echo "No benchmark report found. Run 'make bench' first."; \ + fi + +bench-clean: ## Clean benchmark results + @echo "Cleaning benchmark results..." + @rm -rf target/criterion + +# Performance monitoring +perf-monitor: ## Run continuous performance monitoring + @echo "Starting performance monitoring..." + @if command -v watch >/dev/null 2>&1; then \ + watch -n 300 'make bench-quick'; \ + else \ + echo "Warning: 'watch' command not found. Install it or run 'make bench-quick' manually."; \ + fi + +# sccache management targets +setup-sccache: ## Install and configure sccache + @echo "Setting up sccache..." + @./scripts/setup-sccache.sh + +sccache-stats: ## Show sccache statistics + @if [ -n "$(SCCACHE)" ]; then \ + echo "sccache statistics:"; \ + sccache -s; \ + else \ + echo "sccache not installed. Run 'make setup-sccache' to install it."; \ + fi + +sccache-clean: ## Clear sccache cache + @if [ -n "$(SCCACHE)" ]; then \ + echo "Clearing sccache cache..."; \ + sccache -C; \ + echo "sccache cache cleared"; \ + else \ + echo "sccache not installed"; \ + fi + +sccache-stop: ## Stop the sccache server + @if [ -n "$(SCCACHE)" ]; then \ + echo "Stopping sccache server..."; \ + sccache --stop-server || true; \ + else \ + echo "sccache not installed"; \ + fi diff --git a/lmdeploy/router/README.md b/lmdeploy/router/README.md new file mode 100644 index 0000000000..032cc28fd0 --- /dev/null +++ b/lmdeploy/router/README.md @@ -0,0 +1,288 @@ +# LMDeploy Router + +A high-performance and lightweight request forwarding system for LMDeploy deployments, providing advanced load balancing methods and prefill/decode disaggregation support. + +### Key Features + +- **Core Architecture**: Request routing framework and async processing patterns +- **Load Balancing**: Multiple algorithms (cache-aware, power of two, consistent hashing, random, round robin) +- **Prefill-Decode Disaggregation**: Specialized routing for separated processing phases +- **Service Discovery**: Kubernetes-native worker management and health monitoring +- **Enterprise Features**: Circuit breakers, retry logic, metrics collection + +## Quick Start + +### Prerequisites + +**Rust and Cargo:** + +```bash +# Install rustup (Rust installer and version manager) +curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh + +# Follow the installation prompts, then reload your shell +source $HOME/.cargo/env + +# Verify installation +rustc --version +cargo --version + +``` + +**Python with pip installed** + +### Installation & Basic Usage + +#### Unified Image + +The main LMDeploy image builds `lmdeploy-router` by default. If you need the +standard image without the Rust toolchain and router binary, disable that +optional build: + +```bash +docker build -f docker/Dockerfile --build-arg BUILD_ROUTER=0 . +``` + +### Rust Binary + +```bash +# Build Rust components +cargo build --release +``` + +#### Python Package + +Install from PyPI: + +```bash +pip install lmdeploy-router +``` + +Build a wheel from source. The package uses one release configuration at the +repository root. It contains the Python launcher, the `lmdeploy_router_rs` +PyO3 extension, and the standalone Rust binary as `lmdeploy-router-bin`. + +```bash +python -m pip install -U pip +python -m pip install 'build>=1.2' 'setuptools>=64' 'setuptools-rust>=1.5.2' wheel + +# Optional: if you already built the release binary, explicitly reuse it. +export LMDEPLOY_ROUTER_BIN="${PWD}/target/release/lmdeploy-router" + +python -m build --wheel --outdir dist . +``` + +`lmdeploy_router.ROUTER_AVAILABLE` reports whether the Rust API is available. +A Python-only wheel can be built for packaging checks with +`LMDEPLOY_ROUTER_BUILD_NO_RUST=1`; it contains neither the PyO3 extension nor +the Rust binary. + +Install and verify it with: + +```bash +python -m pip install --force-reinstall dist/lmdeploy_router-*.whl + +python -c 'import lmdeploy_router_rs, lmdeploy_router; print(lmdeploy_router.__version__)' +lmdeploy-router-bin --version +lmdeploy-router --help +``` + +### Usage Examples + +#### Standard Data Parallelism Routing + +```bash +# Launch router with one worker per URL +./target/release/lmdeploy-router \ + --worker-urls http://worker1:8000 http://worker2:8000 \ + --policy consistent_hash + +# Alternative: using cargo run +cargo run --release -- \ + --worker-urls http://worker1:8000 http://worker2:8000 \ + --policy consistent_hash + +# Alternative: using python launcher +lmdeploy-router \ + --worker-urls http://worker1:8000 http://worker2:8000 \ + --policy consistent_hash +``` + +#### Prefill-Decode Disaggregation + +LMDeploy Prefill/Decode services can be registered statically or through the native `--proxy-url` registration flow. RDMA, DLSlime, and Mooncake infrastructure must be provisioned outside the router. + +```bash +cargo run --release --bin lmdeploy-router -- \ + --policy round_robin \ + --lmdeploy-pd-disaggregation \ + --prefill http://127.0.0.1:8081 \ + --decode http://127.0.0.1:8083 \ + --host 127.0.0.1 \ + --port 8090 \ + --prefill-policy round_robin \ + --decode-policy round_robin +``` + +For the native registration flow, start the router without static PD URLs and start LMDeploy services with `--role Prefill` or `--role Decode` plus `--proxy-url`. + +## Configuration + +### Authentication + +Enable bearer-token validation by listing validation URLs (comma-separated) in `.env` via `API_KEY_VALIDATION_URLS` or passing `--api-key-validation-urls`. +When set, all HTTP endpoints require `Authorization: Bearer ` and tokens are validated with HTTP 200 responses. + +```bash +# .env +API_KEY_VALIDATION_URLS=https://codebase.helmholtz.cloud/api/v4/user + +# CLI override +lmdeploy-router --api-key-validation-urls https://codebase.helmholtz.cloud/api/v4/user +``` + +### Metrics + +Prometheus metrics endpoint available at `127.0.0.1:29000` by default. + +```bash +# Custom metrics configuration +lmdeploy-router \ + --worker-urls http://localhost:8080 http://localhost:8081 \ + --prometheus-host 0.0.0.0 \ + --prometheus-port 9000 +``` + +### Retries and Circuit Breakers + +#### Retry Configuration + +Retries are enabled by default with exponential backoff and jitter: + +```bash +lmdeploy-router \ + --worker-urls http://localhost:8080 http://localhost:8081 \ + --retry-max-retries 3 \ + --retry-initial-backoff-ms 100 \ + --retry-max-backoff-ms 10000 \ + --retry-backoff-multiplier 2.0 \ + --retry-jitter-factor 0.1 +``` + +#### Circuit Breaker Configuration + +Circuit breakers protect workers and provide automatic recovery: + +```bash +lmdeploy-router \ + --worker-urls http://localhost:8080 http://localhost:8081 \ + --cb-failure-threshold 5 \ + --cb-success-threshold 2 \ + --cb-timeout-duration-secs 30 \ + --cb-window-duration-secs 60 +``` + +**Circuit Breaker State Machine:** + +- `Closed` → `Open` after N consecutive failures (failure-threshold) +- `Open` → `HalfOpen` after timeout (timeout-duration-secs) +- `HalfOpen` → `Closed` after M consecutive successes (success-threshold) + +**Retry Policy:** Retries on HTTP status codes 408/429/500/502/503/504, with backoff/jitter between attempts. + +### Request ID Tracking + +Track requests across distributed systems with configurable headers: + +```bash +# Use custom request ID headers +lmdeploy-router \ + --worker-urls http://localhost:8080 \ + --request-id-headers x-trace-id x-request-id +``` + +**Default headers:** `x-request-id`, `x-correlation-id`, `x-trace-id`, `request-id` + +### Load Balancing Policies + +The router supports multiple load balancing policies: + +| Policy | Description | Session Affinity | Use Case | +| ----------------- | ---------------------------------------- | ---------------- | ---------------------------------- | +| `round_robin` | Sequential distribution across workers | No | General purpose, even distribution | +| `random` | Uniform random selection | No | Simple deployments | +| `consistent_hash` | Routes same session/user to same worker | Yes | Multi-turn chat, KV cache reuse | +| `power_of_two` | Picks least loaded of two random workers | No | Load-sensitive workloads | +| `cache_aware` | Optimizes for prefix cache hits | Yes | Repeated prompts, few-shot | + +```bash +# Example: Using consistent_hash with HTTP header for session affinity +curl -X POST http://router:8000/v1/chat/completions \ + -H "X-Session-ID: my-session-123" \ + -H "Content-Type: application/json" \ + -d '{"model": "llama-3", "messages": [{"role": "user", "content": "Hello!"}]}' +``` + +For detailed configuration options, hash key priorities, and usage examples, see [Load Balancing Documentation](docs/load_balancing/README.md). + +## Advanced Features + +### Kubernetes Service Discovery + +Automatic worker discovery and management in Kubernetes environments. + +#### Basic Service Discovery + +```bash +lmdeploy-router \ + --service-discovery \ + --selector app=inference-worker role=inference \ + --service-discovery-namespace default +``` + +### Command Line Arguments Reference + +#### Service Discovery + +- `--service-discovery`: Enable Kubernetes service discovery +- `--service-discovery-port`: Port for worker URLs (default: 8000) +- `--service-discovery-namespace`: Kubernetes namespace to watch +- `--selector`: Label selectors for regular mode (format: `key1=value1 key2=value2`) + +## Development + +### Troubleshooting + +**VSCode Rust Analyzer Issues:** +Set `rust-analyzer.linkedProjects` to the absolute path of `Cargo.toml`: + +```json +{ + "rust-analyzer.linkedProjects": ["/workspaces/lmdeploy/lmdeploy-router/Cargo.toml"] +} +``` + +### Release Model + +The router wheel has its own `pyproject.toml` version and is published independently from the main `lmdeploy` wheel. The manual `router-release` workflow builds artifacts by default; select the upload option only after those artifacts are approved. Publication uses the separate `ROUTER_PYPI_TOKEN` secret and the `prod` GitHub environment. + +### CI/CD Pipeline + +The continuous integration pipeline includes comprehensive testing, benchmarking, and publishing: + +#### Build & Test + +1. **Build Wheels**: Uses `cibuildwheel` for manylinux x86_64 packages +2. **Build Source Distribution**: Creates source distribution for pip fallback +3. **Rust HTTP Server Benchmarking**: Performance testing of router overhead +4. **Basic Inference Testing**: End-to-end validation through the router +5. **PD Disaggregation Testing**: Benchmark and sanity checks for prefill-decode load balancing + +#### Publishing + +- **PyPI Publishing**: Wheels and source distributions published when version changes in `pyproject.toml` +- **Container Images**: Docker images published using `/docker/Dockerfile.router` + +## Acknowledgement + +This project is a fork of [vLLM Router](https://github.com/vllm-project/router), and we would like to explicitly acknowledge and thank the original authors for their work. At this stage, our fork includes only minimal changes to preserve the existing interface and ensure compatibility with LMDeploy. We anticipate further divergence as we pursue the roadmap we have in mind, which is the reason for creating the fork. diff --git a/lmdeploy/router/benches/otel_disabled_path.rs b/lmdeploy/router/benches/otel_disabled_path.rs new file mode 100644 index 0000000000..ff68498372 --- /dev/null +++ b/lmdeploy/router/benches/otel_disabled_path.rs @@ -0,0 +1,70 @@ +use axum::http::{HeaderMap, HeaderValue, Request}; +use criterion::{black_box, criterion_group, criterion_main, BatchSize, Criterion}; +use lmdeploy_router_rs::{middleware::RequestSpan, otel_trace, routers::header_utils}; +use tower_http::trace::MakeSpan; + +fn bench_request_span_disabled(c: &mut Criterion) { + otel_trace::shutdown_otel(); + + let request = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header( + "traceparent", + "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", + ) + .body(()) + .unwrap(); + let mut request_span = RequestSpan; + + c.bench_function("otel_request_span_disabled", |b| { + b.iter(|| { + let span = request_span.make_span(black_box(&request)); + black_box(span.is_none()); + }); + }); +} + +fn bench_trace_header_forwarding_disabled(c: &mut Criterion) { + otel_trace::shutdown_otel(); + + let client = reqwest::Client::new(); + let mut incoming = HeaderMap::new(); + incoming.insert( + "traceparent", + HeaderValue::from_static("00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"), + ); + incoming.insert("tracestate", HeaderValue::from_static("vendor=value")); + incoming.insert("baggage", HeaderValue::from_static("userId=alice")); + + for index in 0..32 { + let name = format!("x-noise-{index}"); + incoming.insert( + axum::http::HeaderName::from_bytes(name.as_bytes()).unwrap(), + HeaderValue::from_static("1"), + ); + } + + c.bench_function("otel_trace_header_forwarding_disabled", |b| { + b.iter_batched( + || client.post("http://example.com"), + |request_builder| { + let request = header_utils::propagate_trace_headers( + request_builder, + Some(black_box(&incoming)), + ) + .build() + .unwrap(); + black_box(request.headers().len()); + }, + BatchSize::SmallInput, + ); + }); +} + +criterion_group!( + otel_disabled_path_benches, + bench_request_span_disabled, + bench_trace_header_forwarding_disabled +); +criterion_main!(otel_disabled_path_benches); diff --git a/lmdeploy/router/benches/request_processing.rs b/lmdeploy/router/benches/request_processing.rs new file mode 100644 index 0000000000..df881e02fb --- /dev/null +++ b/lmdeploy/router/benches/request_processing.rs @@ -0,0 +1,525 @@ +use criterion::{black_box, criterion_group, criterion_main, BenchmarkId, Criterion, Throughput}; +use serde_json::{from_str, to_string, to_value, to_vec}; +use std::time::Instant; + +use lmdeploy_router_rs::protocols::spec::{ + ChatCompletionRequest, ChatMessage, CompletionRequest, GenerateParameters, GenerateRequest, + PromptInput, SamplingParams, UserMessageContent, +}; + +/// Create a default GenerateRequest for benchmarks with minimal fields set +fn default_generate_request() -> GenerateRequest { + GenerateRequest { + text: None, + prompt: None, + input_ids: None, + stream: false, + parameters: None, + sampling_params: None, + return_logprob: false, + rid: None, + other: serde_json::Map::new(), + } +} + +/// Create a default ChatCompletionRequest for benchmarks with minimal fields set +fn default_chat_completion_request() -> ChatCompletionRequest { + ChatCompletionRequest { + model: None, + messages: vec![], + input_ids: None, + max_tokens: None, + max_completion_tokens: None, + temperature: None, + top_p: None, + n: None, + stream: false, + stream_options: None, + stop: None, + presence_penalty: None, + frequency_penalty: None, + logit_bias: None, + logprobs: false, + top_logprobs: None, + user: None, + response_format: None, + seed: None, + tools: None, + tool_choice: None, + parallel_tool_calls: None, + function_call: None, + functions: None, + // Backend extensions + top_k: None, + min_p: None, + min_tokens: None, + repetition_penalty: None, + regex: None, + ebnf: None, + stop_token_ids: None, + no_stop_trim: false, + ignore_eos: false, + add_generation_prompt: true, + continue_final_message: false, + skip_special_tokens: true, + separate_reasoning: true, + stream_reasoning: true, + chat_template_kwargs: None, + echo: None, + reasoning_effort: None, + include_reasoning: true, + structured_outputs: None, + other: serde_json::Map::new(), + } +} + +/// Create a default CompletionRequest for benchmarks with minimal fields set +fn default_completion_request() -> CompletionRequest { + CompletionRequest { + model: None, + prompt: PromptInput::String(String::new()), + suffix: None, + max_tokens: None, + temperature: None, + top_p: None, + n: None, + stream: false, + stream_options: None, + logprobs: None, + echo: false, + stop: None, + presence_penalty: None, + frequency_penalty: None, + best_of: None, + logit_bias: None, + user: None, + seed: None, + // Backend extensions + top_k: None, + min_p: None, + min_tokens: None, + repetition_penalty: None, + regex: None, + ebnf: None, + json_schema: None, + stop_token_ids: None, + no_stop_trim: false, + ignore_eos: false, + skip_special_tokens: true, + other: serde_json::Map::new(), + } +} + +// Sample request data for benchmarks +fn create_sample_generate_request() -> GenerateRequest { + GenerateRequest { + text: Some("Write a story about artificial intelligence".to_string()), + parameters: Some(GenerateParameters { + max_new_tokens: Some(100), + temperature: Some(0.8), + top_p: Some(0.9), + top_k: Some(50), + repetition_penalty: Some(1.0), + ..Default::default() + }), + sampling_params: Some(SamplingParams { + temperature: Some(0.8), + top_p: Some(0.9), + top_k: Some(50), + frequency_penalty: Some(0.0), + presence_penalty: Some(0.0), + repetition_penalty: Some(1.0), + ..Default::default() + }), + ..default_generate_request() + } +} + +fn create_sample_chat_completion_request() -> ChatCompletionRequest { + ChatCompletionRequest { + model: Some("gpt-3.5-turbo".to_string()), + messages: vec![ + ChatMessage::System { + role: "system".to_string(), + content: "You are a helpful assistant".to_string(), + name: None, + }, + ChatMessage::User { + role: "user".to_string(), + content: UserMessageContent::Text( + "Explain quantum computing in simple terms".to_string(), + ), + name: None, + }, + ], + max_tokens: Some(150), + max_completion_tokens: Some(150), + temperature: Some(0.7), + top_p: Some(1.0), + n: Some(1), + presence_penalty: Some(0.0), + frequency_penalty: Some(0.0), + parallel_tool_calls: Some(true), + ..default_chat_completion_request() + } +} + +fn create_sample_completion_request() -> CompletionRequest { + CompletionRequest { + model: Some("text-davinci-003".to_string()), + prompt: PromptInput::String("Complete this sentence: The future of AI is".to_string()), + max_tokens: Some(50), + temperature: Some(0.8), + top_p: Some(1.0), + n: Some(1), + presence_penalty: Some(0.0), + frequency_penalty: Some(0.0), + best_of: Some(1), + ..default_completion_request() + } +} + +fn create_large_chat_completion_request() -> ChatCompletionRequest { + let mut messages = vec![ChatMessage::System { + role: "system".to_string(), + content: "You are a helpful assistant with extensive knowledge.".to_string(), + name: None, + }]; + + // Add many user/assistant pairs to simulate a long conversation + for i in 0..50 { + messages.push(ChatMessage::User { + role: "user".to_string(), + content: UserMessageContent::Text(format!("Question {}: What do you think about topic number {} which involves complex reasoning about multiple interconnected systems and their relationships?", i, i)), + name: None, + }); + messages.push(ChatMessage::Assistant { + role: "assistant".to_string(), + content: Some(format!("Answer {}: This is a detailed response about topic {} that covers multiple aspects and provides comprehensive analysis of the interconnected systems you mentioned.", i, i)), + name: None, + tool_calls: None, + function_call: None, + reasoning: None, + }); + } + + ChatCompletionRequest { + model: Some("gpt-4".to_string()), + messages, + max_tokens: Some(1000), + max_completion_tokens: Some(1000), + temperature: Some(0.7), + top_p: Some(0.95), + n: Some(1), + presence_penalty: Some(0.1), + frequency_penalty: Some(0.1), + top_logprobs: Some(5), + user: Some("benchmark_user".to_string()), + seed: Some(42), + parallel_tool_calls: Some(true), + ..default_chat_completion_request() + } +} + +// Benchmark JSON serialization +fn bench_json_serialization(c: &mut Criterion) { + let mut group = c.benchmark_group("json_serialization"); + + let generate_req = create_sample_generate_request(); + let chat_req = create_sample_chat_completion_request(); + let completion_req = create_sample_completion_request(); + let large_chat_req = create_large_chat_completion_request(); + + group.bench_function("generate_request", |b| { + b.iter(|| { + let json = to_string(black_box(&generate_req)).unwrap(); + black_box(json); + }); + }); + + group.bench_function("chat_completion_request", |b| { + b.iter(|| { + let json = to_string(black_box(&chat_req)).unwrap(); + black_box(json); + }); + }); + + group.bench_function("completion_request", |b| { + b.iter(|| { + let json = to_string(black_box(&completion_req)).unwrap(); + black_box(json); + }); + }); + + group.bench_function("large_chat_completion_request", |b| { + b.iter(|| { + let json = to_string(black_box(&large_chat_req)).unwrap(); + black_box(json); + }); + }); + + group.bench_function("generate_request_to_bytes", |b| { + b.iter(|| { + let bytes = to_vec(black_box(&generate_req)).unwrap(); + black_box(bytes); + }); + }); + + group.finish(); +} + +// Benchmark JSON deserialization +fn bench_json_deserialization(c: &mut Criterion) { + let mut group = c.benchmark_group("json_deserialization"); + + let generate_json = to_string(&create_sample_generate_request()).unwrap(); + let chat_json = to_string(&create_sample_chat_completion_request()).unwrap(); + let completion_json = to_string(&create_sample_completion_request()).unwrap(); + let large_chat_json = to_string(&create_large_chat_completion_request()).unwrap(); + + group.bench_function("generate_request", |b| { + b.iter(|| { + let req: GenerateRequest = from_str(black_box(&generate_json)).unwrap(); + black_box(req); + }); + }); + + group.bench_function("chat_completion_request", |b| { + b.iter(|| { + let req: ChatCompletionRequest = from_str(black_box(&chat_json)).unwrap(); + black_box(req); + }); + }); + + group.bench_function("completion_request", |b| { + b.iter(|| { + let req: CompletionRequest = from_str(black_box(&completion_json)).unwrap(); + black_box(req); + }); + }); + + group.bench_function("large_chat_completion_request", |b| { + b.iter(|| { + let req: ChatCompletionRequest = from_str(black_box(&large_chat_json)).unwrap(); + black_box(req); + }); + }); + + group.finish(); +} + +// Benchmark direct JSON routing (replaces regular routing) +fn bench_direct_json_routing(c: &mut Criterion) { + let mut group = c.benchmark_group("direct_json_routing"); + + let generate_req = create_sample_generate_request(); + let chat_req = create_sample_chat_completion_request(); + let completion_req = create_sample_completion_request(); + + group.bench_function("generate_to_json", |b| { + b.iter(|| { + let json = to_value(black_box(&generate_req)).unwrap(); + black_box(json); + }); + }); + + group.bench_function("generate_to_json_string", |b| { + b.iter(|| { + let json = to_string(black_box(&generate_req)).unwrap(); + black_box(json); + }); + }); + + group.bench_function("generate_to_bytes", |b| { + b.iter(|| { + let bytes = to_vec(black_box(&generate_req)).unwrap(); + black_box(bytes); + }); + }); + + group.bench_function("chat_completion_to_json", |b| { + b.iter(|| { + let json = to_value(black_box(&chat_req)).unwrap(); + black_box(json); + }); + }); + + group.bench_function("chat_completion_to_json_string", |b| { + b.iter(|| { + let json = to_string(black_box(&chat_req)).unwrap(); + black_box(json); + }); + }); + + group.bench_function("completion_to_json", |b| { + b.iter(|| { + let json = to_value(black_box(&completion_req)).unwrap(); + black_box(json); + }); + }); + + group.finish(); +} + +// Benchmark throughput with different request sizes +fn bench_throughput_by_size(c: &mut Criterion) { + let mut group = c.benchmark_group("throughput_by_size"); + + // Create requests of different sizes + let small_generate = GenerateRequest { + text: Some("Hi".to_string()), + ..default_generate_request() + }; + + let medium_generate = GenerateRequest { + text: Some("Write a medium length story about AI".repeat(10)), + ..default_generate_request() + }; + + let large_generate = GenerateRequest { + text: Some("Write a very long and detailed story about artificial intelligence and its impact on society".repeat(100)), + ..default_generate_request() + }; + + for (name, req) in [ + ("small", &small_generate), + ("medium", &medium_generate), + ("large", &large_generate), + ] { + let json = to_string(req).unwrap(); + let size_bytes = json.len(); + + group.throughput(Throughput::Bytes(size_bytes as u64)); + group.bench_with_input(BenchmarkId::new("serialize", name), &req, |b, req| { + b.iter(|| { + let json = to_string(black_box(req)).unwrap(); + black_box(json); + }); + }); + + group.bench_with_input( + BenchmarkId::new("deserialize", name), + &json, + |b, json_str| { + b.iter(|| { + let req: GenerateRequest = black_box(from_str(json_str)).unwrap(); + black_box(req); + }); + }, + ); + } + + group.finish(); +} + +// Benchmark full round-trip: deserialize -> route JSON -> serialize +fn bench_full_round_trip(c: &mut Criterion) { + let mut group = c.benchmark_group("full_round_trip"); + + let generate_json = to_string(&create_sample_generate_request()).unwrap(); + let chat_json = to_string(&create_sample_chat_completion_request()).unwrap(); + let completion_json = to_string(&create_sample_completion_request()).unwrap(); + + group.bench_function("generate_direct_json_pipeline", |b| { + b.iter(|| { + // Deserialize OpenAI request + let req: GenerateRequest = from_str(black_box(&generate_json)).unwrap(); + // Convert to JSON for direct routing + let routing_json = to_value(&req).unwrap(); + let json_string = to_string(&routing_json).unwrap(); + black_box(json_string); + }); + }); + + group.bench_function("chat_completion_direct_json_pipeline", |b| { + b.iter(|| { + let req: ChatCompletionRequest = from_str(black_box(&chat_json)).unwrap(); + let routing_json = to_value(&req).unwrap(); + let json_string = to_string(&routing_json).unwrap(); + black_box(json_string); + }); + }); + + group.bench_function("completion_direct_json_pipeline", |b| { + b.iter(|| { + let req: CompletionRequest = from_str(black_box(&completion_json)).unwrap(); + let routing_json = to_value(&req).unwrap(); + let json_string = to_string(&routing_json).unwrap(); + black_box(json_string); + }); + }); + + group.finish(); +} + +fn benchmark_summary(c: &mut Criterion) { + let group = c.benchmark_group("benchmark_summary"); + + println!("\nLMDeploy Router Performance Benchmark Suite"); + println!("============================================="); + + // Quick performance overview + let generate_req = create_sample_generate_request(); + + println!("\nQuick Performance Overview:"); + + // Measure serialization + let start = Instant::now(); + for _ in 0..1000 { + let _ = black_box(to_string(&generate_req).unwrap()); + } + let serialize_time = start.elapsed().as_nanos() / 1000; + println!(" * Serialization (avg): {:>8} ns/req", serialize_time); + + // Measure deserialization + let json = to_string(&generate_req).unwrap(); + let start = Instant::now(); + for _ in 0..1000 { + let _: GenerateRequest = black_box(from_str(&json).unwrap()); + } + let deserialize_time = start.elapsed().as_nanos() / 1000; + println!( + " * Deserialization (avg): {:>8} ns/req", + deserialize_time + ); + + // Calculate ratios + let total_pipeline = serialize_time + deserialize_time; + println!(" * Total Pipeline (avg): {:>8} ns/req", total_pipeline); + + println!("\nPerformance Insights:"); + if deserialize_time > serialize_time * 2 { + println!(" • Deserialization is significantly faster than serialization"); + } + if total_pipeline < 100_000 { + println!(" • Total pipeline latency is excellent (< 100μs)"); + } + + println!("\nSimplification Benefits:"); + println!(" • Eliminated complex type conversion layer"); + println!(" • Reduced memory allocations"); + println!(" • Automatic field preservation (no manual mapping)"); + println!(" • Direct JSON manipulation improves performance"); + + println!("\nRecommendations:"); + if serialize_time > deserialize_time { + println!(" • Focus optimization efforts on serialization rather than deserialization"); + } + println!(" • PD mode overhead is minimal - safe to use for latency-sensitive workloads"); + println!(" • Consider batching small requests to improve overall throughput"); + + println!("\n{}", "=".repeat(50)); + + group.finish(); +} + +criterion_group!( + benches, + benchmark_summary, + bench_json_serialization, + bench_json_deserialization, + bench_direct_json_routing, + bench_throughput_by_size, + bench_full_round_trip +); +criterion_main!(benches); diff --git a/lmdeploy/router/benches/tokenizer_benchmark.rs b/lmdeploy/router/benches/tokenizer_benchmark.rs new file mode 100644 index 0000000000..dc4ed5e68b --- /dev/null +++ b/lmdeploy/router/benches/tokenizer_benchmark.rs @@ -0,0 +1,1400 @@ +//! Comprehensive tokenizer benchmark with clean summary output +//! Each test adds a row to the final summary table + +use criterion::{black_box, criterion_group, BenchmarkId, Criterion, Throughput}; +use lmdeploy_router_rs::tokenizer::{ + huggingface::HuggingFaceTokenizer, sequence::Sequence, stop::*, stream::DecodeStream, traits::*, +}; +use std::collections::BTreeMap; +use std::path::PathBuf; +use std::sync::atomic::{AtomicBool, AtomicU64, Ordering}; +use std::sync::{Arc, Mutex, OnceLock}; +use std::thread; +use std::time::{Duration, Instant}; + +// Include the common test utilities +#[path = "../tests/common/mod.rs"] +mod common; +use common::ensure_tokenizer_cached; + +// Cache the tokenizer path for the entire benchmark run +static TOKENIZER_PATH: OnceLock = OnceLock::new(); + +fn get_tokenizer_path() -> &'static PathBuf { + TOKENIZER_PATH.get_or_init(ensure_tokenizer_cached) +} + +// Production target: 100k tokens per second +const TARGET_TOKENS_PER_SECOND: u64 = 100_000; + +// Typical prompt sizes +const SHORT_PROMPT: &str = "What is the capital of France?"; +const MEDIUM_PROMPT: &str = "Write a detailed explanation of quantum computing, including its principles, current applications, and future potential. Be sure to cover both the theoretical foundations and practical implementations."; +const LONG_PROMPT: &str = "You are an expert software engineer. Review the following code and provide detailed feedback on performance optimizations, potential bugs, and architectural improvements. Consider scalability, maintainability, and best practices. The code implements a distributed caching system with the following requirements: 1) High availability across multiple regions, 2) Sub-millisecond latency for cache hits, 3) Automatic failover and recovery, 4) Support for both LRU and LFU eviction policies, 5) Real-time monitoring and alerting. Please analyze each component thoroughly and suggest concrete improvements with code examples where appropriate."; + +// System prompts can be quite large +fn generate_system_prompt(size: usize) -> String { + let base = "You are a helpful AI assistant with expertise in "; + let domains = vec![ + "mathematics", + "physics", + "chemistry", + "biology", + "computer science", + "engineering", + "medicine", + "law", + "economics", + "philosophy", + ]; + + let mut prompt = base.to_string(); + while prompt.len() < size { + for domain in &domains { + prompt.push_str(domain); + prompt.push_str(", "); + if prompt.len() >= size { + break; + } + } + } + prompt +} + +// Global results storage +lazy_static::lazy_static! { + static ref BENCHMARK_RESULTS: Mutex> = Mutex::new(BTreeMap::new()); +} + +fn add_result(category: &str, result: String) { + let mut results = BENCHMARK_RESULTS.lock().unwrap(); + let index = results.len(); + results.insert(format!("{:03}_{}", index, category), result); +} + +fn bench_encode_throughput(c: &mut Criterion) { + let tokenizer_path = get_tokenizer_path(); + let tokenizer = Arc::new( + HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap()) + .expect("Failed to load tokenizer"), + ); + + // Pre-generate system prompts + let system_1k = generate_system_prompt(1000); + let system_4k = generate_system_prompt(4000); + let system_16k = generate_system_prompt(16000); + + let test_cases = vec![ + ("short_30B", SHORT_PROMPT), + ("medium_230B", MEDIUM_PROMPT), + ("long_670B", LONG_PROMPT), + ("system_1KB", system_1k.as_str()), + ("system_4KB", system_4k.as_str()), + ("system_16KB", system_16k.as_str()), + ]; + + let mut group = c.benchmark_group("encode_throughput"); + + for (name, prompt) in test_cases { + let prompt_len = prompt.len(); + let tokenizer_clone = tokenizer.clone(); + + // Get token count once + let encoding = tokenizer.encode(prompt).unwrap(); + let token_count = encoding.token_ids().len(); + + // Track if metrics have been printed for this test case + let printed = Arc::new(AtomicBool::new(false)); + + group.throughput(Throughput::Bytes(prompt_len as u64)); + group.bench_function(name, |b| { + let printed_clone = printed.clone(); + let tokenizer = tokenizer_clone.clone(); + + b.iter_custom(|iters| { + let start = Instant::now(); + for _ in 0..iters { + black_box(tokenizer.encode(prompt).unwrap()); + } + let duration = start.elapsed(); + + // Store result only once per test case + if !printed_clone.load(Ordering::Relaxed) { + let ops_per_sec = iters as f64 / duration.as_secs_f64(); + let chars_per_sec = (iters as f64 * prompt_len as f64) / duration.as_secs_f64(); + let tokens_per_sec = + (iters as f64 * token_count as f64) / duration.as_secs_f64(); + + let result = format!( + "{:<15} | {:>8} | {:>8} | {:>12.0} | {:>12.0} | {:>10.0} | {:>10}", + name, + prompt_len, + token_count, + chars_per_sec, + tokens_per_sec, + ops_per_sec, + 1 + ); + add_result("encode", result); + + printed_clone.store(true, Ordering::Relaxed); + } + + duration + }); + }); + } + + group.finish(); +} + +fn bench_batch_encode(c: &mut Criterion) { + let tokenizer_path = get_tokenizer_path(); + let tokenizer = Arc::new( + HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap()) + .expect("Failed to load tokenizer"), + ); + + let batch_sizes = vec![1, 8, 16, 32, 64, 128]; + let prompt = MEDIUM_PROMPT; + let prompt_len = prompt.len(); + let encoding = tokenizer.encode(prompt).unwrap(); + let token_count = encoding.token_ids().len(); + + let mut group = c.benchmark_group("batch_encode"); + + for batch_size in batch_sizes { + let prompts: Vec<&str> = vec![prompt; batch_size]; + let printed = Arc::new(AtomicBool::new(false)); + let tokenizer_clone = tokenizer.clone(); + + group.throughput(Throughput::Elements(batch_size as u64)); + group.bench_with_input( + BenchmarkId::from_parameter(batch_size), + &batch_size, + |b, &size| { + let printed_clone = printed.clone(); + let tokenizer = tokenizer_clone.clone(); + + b.iter_custom(|iters| { + let start = Instant::now(); + for _ in 0..iters { + black_box(tokenizer.encode_batch(&prompts).unwrap()); + } + let duration = start.elapsed(); + + if !printed_clone.load(Ordering::Relaxed) { + let prompts_per_sec = (iters as f64 * size as f64) / duration.as_secs_f64(); + let tokens_per_sec = prompts_per_sec * token_count as f64; + let chars_per_sec = prompts_per_sec * prompt_len as f64; + + let result = format!( + "{:<15} | {:>8} | {:>8} | {:>12.0} | {:>12.0} | {:>10.0} | {:>10}", + format!("batch_{}", size), + prompt_len * size, + token_count * size, + prompts_per_sec, + tokens_per_sec, + chars_per_sec, + 1 + ); + add_result("batch", result); + + printed_clone.store(true, Ordering::Relaxed); + } + + duration + }); + }, + ); + } + + group.finish(); +} + +fn bench_concurrent_encode(c: &mut Criterion) { + let tokenizer = Arc::new( + HuggingFaceTokenizer::from_file(get_tokenizer_path().to_str().unwrap()) + .expect("Failed to load tokenizer"), + ); + + let client_counts = vec![1, 4, 8, 16, 32]; + + let mut group = c.benchmark_group("concurrent_encode"); + group.measurement_time(Duration::from_secs(2)); + + for num_clients in client_counts { + let printed = Arc::new(AtomicBool::new(false)); + let tokenizer_clone = tokenizer.clone(); + + group.bench_with_input( + BenchmarkId::from_parameter(num_clients), + &num_clients, + |b, &clients| { + let printed_clone = printed.clone(); + + b.iter_custom(|_iters| { + let tokenizer = tokenizer_clone.clone(); + let total_operations = Arc::new(AtomicU64::new(0)); + let total_chars = Arc::new(AtomicU64::new(0)); + let start = Instant::now(); + + let handles: Vec<_> = (0..clients) + .map(|client_id| { + let tokenizer = tokenizer.clone(); + let total_ops = total_operations.clone(); + let total_ch = total_chars.clone(); + + thread::spawn(move || { + let prompts = [SHORT_PROMPT, MEDIUM_PROMPT, LONG_PROMPT]; + let prompt = prompts[client_id % prompts.len()]; + let mut local_ops = 0u64; + let mut local_chars = 0u64; + + while start.elapsed() < Duration::from_millis(500) { + let _ = tokenizer.encode(prompt).unwrap(); + local_ops += 1; + local_chars += prompt.len() as u64; + } + + total_ops.fetch_add(local_ops, Ordering::Relaxed); + total_ch.fetch_add(local_chars, Ordering::Relaxed); + }) + }) + .collect(); + + for handle in handles { + handle.join().unwrap(); + } + + let duration = start.elapsed(); + + if !printed_clone.load(Ordering::Relaxed) { + let total_ops = total_operations.load(Ordering::Relaxed); + let total_ch = total_chars.load(Ordering::Relaxed); + let ops_per_sec = total_ops as f64 / duration.as_secs_f64(); + let chars_per_sec = total_ch as f64 / duration.as_secs_f64(); + let per_client = ops_per_sec / clients as f64; + + let result = format!( + "{:<15} | {:>10} | {:>12.0} | {:>12.0} | {:>15.0}", + format!("{}_clients", clients), + total_ops, + ops_per_sec, + chars_per_sec, + per_client + ); + add_result("concurrent", result); + + printed_clone.store(true, Ordering::Relaxed); + } + + duration + }); + }, + ); + } + + group.finish(); +} + +fn bench_decode_performance(c: &mut Criterion) { + let tokenizer = Arc::new( + HuggingFaceTokenizer::from_file(get_tokenizer_path().to_str().unwrap()) + .expect("Failed to load tokenizer"), + ); + + let test_text = "The quick brown fox jumps over the lazy dog. ".repeat(10); + let encoding = tokenizer.encode(&test_text).unwrap(); + let tokens = encoding.token_ids(); + let num_tokens = tokens.len(); + + let mut group = c.benchmark_group("decode_performance"); + + // Test direct decode + let printed_direct = Arc::new(AtomicBool::new(false)); + group.bench_function("direct_decode", |b| { + let printed = printed_direct.clone(); + let tokenizer = tokenizer.clone(); + + b.iter_custom(|iters| { + let start = Instant::now(); + for _ in 0..iters { + black_box(tokenizer.decode(tokens, false).unwrap()); + } + let duration = start.elapsed(); + + if !printed.load(Ordering::Relaxed) { + let ops_per_sec = iters as f64 / duration.as_secs_f64(); + let tokens_per_sec = ops_per_sec * num_tokens as f64; + + let result = format!( + "{:<20} | {:>10} | {:>12.0} | {:>12.0} | {:>10}", + "Direct", num_tokens, tokens_per_sec, ops_per_sec, 1 + ); + add_result("decode", result); + + printed.store(true, Ordering::Relaxed); + } + + duration + }); + }); + + // Test DecodeStream + let printed_stream = Arc::new(AtomicBool::new(false)); + group.bench_function("decode_stream", |b| { + let printed = printed_stream.clone(); + let tokenizer = tokenizer.clone(); + + b.iter_custom(|iters| { + let start = Instant::now(); + for _ in 0..iters { + let mut decoder = DecodeStream::new(tokenizer.clone(), &[], false); + let mut output = String::new(); + for token in tokens { + if let Some(text) = decoder.step(*token).unwrap() { + output.push_str(&text); + } + } + black_box(output); + } + let duration = start.elapsed(); + + if !printed.load(Ordering::Relaxed) { + let ops_per_sec = iters as f64 / duration.as_secs_f64(); + let tokens_per_sec = ops_per_sec * num_tokens as f64; + + let result = format!( + "{:<20} | {:>10} | {:>12.0} | {:>12.0} | {:>10}", + "DecodeStream", num_tokens, tokens_per_sec, ops_per_sec, 1 + ); + add_result("decode", result); + + printed.store(true, Ordering::Relaxed); + } + + duration + }); + }); + + // Test Sequence + let printed_seq = Arc::new(AtomicBool::new(false)); + group.bench_function("sequence_decode", |b| { + let printed = printed_seq.clone(); + let tokenizer = tokenizer.clone(); + + b.iter_custom(|iters| { + let start = Instant::now(); + for _ in 0..iters { + let mut sequence = Sequence::new(tokenizer.clone()); + let mut output = String::new(); + for token in tokens { + let text = sequence.append_token(*token).unwrap(); + output.push_str(&text); + } + black_box(output); + } + let duration = start.elapsed(); + + if !printed.load(Ordering::Relaxed) { + let ops_per_sec = iters as f64 / duration.as_secs_f64(); + let tokens_per_sec = ops_per_sec * num_tokens as f64; + + let result = format!( + "{:<20} | {:>10} | {:>12.0} | {:>12.0} | {:>10}", + "Sequence", num_tokens, tokens_per_sec, ops_per_sec, 1 + ); + add_result("decode", result); + + printed.store(true, Ordering::Relaxed); + } + + duration + }); + }); + + group.finish(); +} + +fn bench_streaming_decode_100k(c: &mut Criterion) { + let tokenizer = Arc::new( + HuggingFaceTokenizer::from_file(get_tokenizer_path().to_str().unwrap()) + .expect("Failed to load tokenizer"), + ); + + let sample_text = "The quick brown fox jumps over the lazy dog. ".repeat(1000); + let encoding = tokenizer.encode(&sample_text).unwrap(); + let all_tokens = encoding.token_ids(); + + let mut group = c.benchmark_group("streaming_100k"); + group.measurement_time(Duration::from_secs(1)); + + // Test DecodeStream + let printed_stream = Arc::new(AtomicBool::new(false)); + group.bench_function("decode_stream_100k", |b| { + let printed = printed_stream.clone(); + let tokenizer = tokenizer.clone(); + + b.iter_custom(|_iters| { + let start = Instant::now(); + let mut decoder = DecodeStream::new(tokenizer.clone(), &[], false); + let mut output = String::new(); + let mut tokens_processed = 0u64; + + for token in all_tokens.iter().cycle() { + if start.elapsed() >= Duration::from_millis(500) { + break; + } + + if let Some(text) = decoder.step(*token).unwrap() { + output.push_str(&text); + } + tokens_processed += 1; + } + + let duration = start.elapsed(); + + if !printed.load(Ordering::Relaxed) { + let tokens_per_sec = tokens_processed as f64 / duration.as_secs_f64(); + let status = if tokens_per_sec >= TARGET_TOKENS_PER_SECOND as f64 { + "PASS" + } else { + "BELOW" + }; + + let result = format!( + "{:<20} | {:>12} | {:>12.0} | {:>12} | {:>10} | {:>12}", + "DecodeStream", + tokens_processed, + tokens_per_sec, + TARGET_TOKENS_PER_SECOND, + 1, + status + ); + add_result("streaming_100k", result); + + printed.store(true, Ordering::Relaxed); + } + + duration + }); + }); + + // Test Sequence + let printed_seq = Arc::new(AtomicBool::new(false)); + group.bench_function("sequence_100k", |b| { + let printed = printed_seq.clone(); + let tokenizer = tokenizer.clone(); + + b.iter_custom(|_iters| { + let start = Instant::now(); + let mut sequence = Sequence::new(tokenizer.clone()); + let mut output = String::new(); + let mut tokens_processed = 0u64; + + for token in all_tokens.iter().cycle() { + if start.elapsed() >= Duration::from_millis(500) { + break; + } + + let text = sequence.append_token(*token).unwrap(); + output.push_str(&text); + tokens_processed += 1; + } + + let duration = start.elapsed(); + + if !printed.load(Ordering::Relaxed) { + let tokens_per_sec = tokens_processed as f64 / duration.as_secs_f64(); + let status = if tokens_per_sec >= TARGET_TOKENS_PER_SECOND as f64 { + "PASS" + } else { + "BELOW" + }; + + let result = format!( + "{:<20} | {:>12} | {:>12.0} | {:>12} | {:>10} | {:>12}", + "Sequence", + tokens_processed, + tokens_per_sec, + TARGET_TOKENS_PER_SECOND, + 1, + status + ); + add_result("streaming_100k", result); + + printed.store(true, Ordering::Relaxed); + } + + duration + }); + }); + + group.finish(); +} + +fn bench_latency_distribution(c: &mut Criterion) { + let tokenizer = Arc::new( + HuggingFaceTokenizer::from_file(get_tokenizer_path().to_str().unwrap()) + .expect("Failed to load tokenizer"), + ); + + // Test latency for individual token processing + let sample_tokens = vec![1, 450, 6635, 3290, 491, 278, 3474, 29892]; + + let mut group = c.benchmark_group("latency"); + + // Encode latency + let system_4k = generate_system_prompt(4000); + let test_cases = vec![ + ("encode_short", SHORT_PROMPT), + ("encode_medium", MEDIUM_PROMPT), + ("encode_long", LONG_PROMPT), + ("encode_4KB", system_4k.as_str()), + ]; + + for (name, prompt) in test_cases { + let printed = Arc::new(AtomicBool::new(false)); + group.bench_function(name, |b| { + let printed_clone = printed.clone(); + let tokenizer = tokenizer.clone(); + + b.iter_custom(|iters| { + // Only collect detailed latency on first iteration + let total_duration = if !printed_clone.load(Ordering::Relaxed) { + let mut latencies = Vec::new(); + + // Warm up + for _ in 0..100 { + let _ = tokenizer.encode(prompt).unwrap(); + } + + // Measure for statistics + for _ in 0..1000 { + let start = Instant::now(); + let _ = tokenizer.encode(prompt).unwrap(); + let latency = start.elapsed(); + latencies.push(latency); + } + + latencies.sort(); + let p50 = latencies[latencies.len() / 2]; + let p95 = latencies[latencies.len() * 95 / 100]; + let p99 = latencies[latencies.len() * 99 / 100]; + let max = latencies.last().unwrap(); + + let result = format!( + "{:<20} | {:>10.1} | {:>10.1} | {:>10.1} | {:>10.1} | {:>10}", + name, + p50.as_micros() as f64, + p95.as_micros() as f64, + p99.as_micros() as f64, + max.as_micros() as f64, + 1000 + ); + add_result("latency", result); + + printed_clone.store(true, Ordering::Relaxed); + + // Return median for consistency + p50 * iters as u32 + } else { + // Regular benchmark iterations + let start = Instant::now(); + for _ in 0..iters { + black_box(tokenizer.encode(prompt).unwrap()); + } + start.elapsed() + }; + + total_duration + }); + }); + } + + // Decode token latency + let printed_decode = Arc::new(AtomicBool::new(false)); + group.bench_function("decode_token", |b| { + let printed_clone = printed_decode.clone(); + let tokenizer = tokenizer.clone(); + let tokens = sample_tokens.clone(); + + b.iter_custom(|iters| { + let total_duration = if !printed_clone.load(Ordering::Relaxed) { + let mut latencies = Vec::new(); + let mut decoder = DecodeStream::new(tokenizer.clone(), &[], false); + + for token in tokens.iter().cycle().take(1000) { + let start = Instant::now(); + let _ = decoder.step(*token).unwrap(); + let latency = start.elapsed(); + latencies.push(latency); + } + + latencies.sort(); + let p50 = latencies[latencies.len() / 2]; + let p95 = latencies[latencies.len() * 95 / 100]; + let p99 = latencies[latencies.len() * 99 / 100]; + let max = latencies.last().unwrap(); + + let result = format!( + "{:<20} | {:>10.1} | {:>10.1} | {:>10.1} | {:>10.1} | {:>10}", + "decode_token", + p50.as_micros() as f64, + p95.as_micros() as f64, + p99.as_micros() as f64, + max.as_micros() as f64, + 1000 + ); + add_result("latency", result); + + // Check target latency + let target_latency = Duration::from_micros(10); + if p50 > target_latency { + let warning = format!( + "WARNING: P50 latency exceeds target of {:?} for 100k tokens/sec", + target_latency + ); + add_result("latency_warning", warning); + } + + printed_clone.store(true, Ordering::Relaxed); + + // Return approximate time for consistency + p50 * iters as u32 + } else { + // Regular benchmark iterations + let start = Instant::now(); + for _ in 0..iters { + let mut decoder = DecodeStream::new(tokenizer.clone(), &[], false); + for token in tokens.iter().take(10) { + black_box(decoder.step(*token).unwrap()); + } + } + start.elapsed() + }; + + total_duration + }); + }); + + group.finish(); +} + +fn bench_concurrent_streaming(c: &mut Criterion) { + let tokenizer = Arc::new( + HuggingFaceTokenizer::from_file(get_tokenizer_path().to_str().unwrap()) + .expect("Failed to load tokenizer"), + ); + + let num_sequences = 16; + let tokens_per_sequence = 10_000; + + let sample_text = "The quick brown fox jumps over the lazy dog. ".repeat(100); + let encoding = tokenizer.encode(&sample_text).unwrap(); + let token_batch: Vec = encoding.token_ids().to_vec(); + + let mut group = c.benchmark_group("concurrent_streaming"); + group.measurement_time(Duration::from_secs(2)); + + let printed = Arc::new(AtomicBool::new(false)); + group.bench_function("concurrent_16_sequences", |b| { + let printed_clone = printed.clone(); + let tokenizer = tokenizer.clone(); + let tokens = token_batch.clone(); + + b.iter_custom(|_iters| { + let total_tokens = Arc::new(AtomicU64::new(0)); + let start = Instant::now(); + + let handles: Vec<_> = (0..num_sequences) + .map(|_seq_id| { + let tokenizer = tokenizer.clone(); + let tokens = tokens.clone(); + let total_tokens = total_tokens.clone(); + + thread::spawn(move || { + let mut decoder = DecodeStream::new(tokenizer, &[], false); + let mut output = String::new(); + let mut local_count = 0u64; + + for token in tokens.iter().cycle().take(tokens_per_sequence) { + if let Some(text) = decoder.step(*token).unwrap() { + output.push_str(&text); + } + local_count += 1; + } + + total_tokens.fetch_add(local_count, Ordering::Relaxed); + }) + }) + .collect(); + + for handle in handles { + handle.join().unwrap(); + } + + let duration = start.elapsed(); + + if !printed_clone.load(Ordering::Relaxed) { + let total = total_tokens.load(Ordering::Relaxed); + let throughput = total as f64 / duration.as_secs_f64(); + let per_seq = throughput / num_sequences as f64; + + let result = format!( + "{:<20} | {:>10} | {:>12.0} | {:>15.0} | {:>15}", + format!("{}_sequences", num_sequences), + total, + throughput, + per_seq, + num_sequences + ); + add_result("concurrent_streaming", result); + + printed_clone.store(true, Ordering::Relaxed); + } + + duration + }); + }); + + group.finish(); +} + +fn bench_stop_sequences(c: &mut Criterion) { + let tokenizer = Arc::new( + HuggingFaceTokenizer::from_file(get_tokenizer_path().to_str().unwrap()) + .expect("Failed to load tokenizer"), + ); + + let config = StopSequenceConfig::default() + .with_stop_sequence("") + .with_stop_sequence("\n\n") + .with_stop_sequence("###") + .with_stop_token(2); + + let sample_text = "Hello world! This is a test. ### Stop here. Continue after.".repeat(100); + let encoding = tokenizer.encode(&sample_text).unwrap(); + let tokens = encoding.token_ids(); + + let mut group = c.benchmark_group("stop_sequences"); + + // No stops + let printed_no_stop = Arc::new(AtomicBool::new(false)); + group.bench_function("no_stops", |b| { + let printed_clone = printed_no_stop.clone(); + let tokenizer = tokenizer.clone(); + + b.iter_custom(|iters| { + let start = Instant::now(); + let mut total_tokens = 0u64; + + for _ in 0..iters { + let mut decoder = StopSequenceDecoder::new( + tokenizer.clone(), + StopSequenceConfig::default(), + false, + ); + for token in tokens { + let _ = decoder.process_token(*token).unwrap(); + total_tokens += 1; + } + } + + let duration = start.elapsed(); + + if !printed_clone.load(Ordering::Relaxed) { + let tokens_per_sec = total_tokens as f64 / duration.as_secs_f64(); + let seq_per_sec = iters as f64 / duration.as_secs_f64(); + + let result = format!( + "{:<20} | {:>10} | {:>12} | {:>12.0} | {:>10.0}", + "No stops", iters, total_tokens, tokens_per_sec, seq_per_sec + ); + add_result("stop_sequences", result); + + printed_clone.store(true, Ordering::Relaxed); + } + + duration + }); + }); + + // With stops + let printed_with_stops = Arc::new(AtomicBool::new(false)); + group.bench_function("with_stops", |b| { + let printed_clone = printed_with_stops.clone(); + let tokenizer = tokenizer.clone(); + let config = config.clone(); + + b.iter_custom(|iters| { + let start = Instant::now(); + let mut total_tokens = 0u64; + let mut total_sequences = 0u64; + + for _ in 0..iters { + let mut decoder = + StopSequenceDecoder::new(tokenizer.clone(), config.clone(), false); + let mut sequence_tokens = 0u64; + + for token in tokens { + let result = decoder.process_token(*token).unwrap(); + sequence_tokens += 1; + + if matches!( + result, + SequenceDecoderOutput::Stopped | SequenceDecoderOutput::StoppedWithText(_) + ) { + break; + } + } + + total_tokens += sequence_tokens; + total_sequences += 1; + } + + let duration = start.elapsed(); + + if !printed_clone.load(Ordering::Relaxed) { + let tokens_per_sec = total_tokens as f64 / duration.as_secs_f64(); + let seq_per_sec = total_sequences as f64 / duration.as_secs_f64(); + + let result = format!( + "{:<20} | {:>10} | {:>12} | {:>12.0} | {:>10.0}", + "With stops", total_sequences, total_tokens, tokens_per_sec, seq_per_sec + ); + add_result("stop_sequences", result); + + printed_clone.store(true, Ordering::Relaxed); + } + + duration + }); + }); + + group.finish(); +} + +fn bench_multithreaded_encode(c: &mut Criterion) { + let tokenizer = Arc::new( + HuggingFaceTokenizer::from_file(get_tokenizer_path().to_str().unwrap()) + .expect("Failed to load tokenizer"), + ); + + let thread_counts = vec![1, 2, 4, 8, 16]; + let operations_per_thread = 1000; + + // Test with medium-sized prompt for balanced workload + let test_prompt = MEDIUM_PROMPT; + + let mut group = c.benchmark_group("multithreaded_encode"); + group.measurement_time(Duration::from_secs(2)); + + let mut baseline_throughput = 0.0; + + for num_threads in thread_counts { + let printed = Arc::new(AtomicBool::new(false)); + let tokenizer_clone = tokenizer.clone(); + + group.bench_with_input( + BenchmarkId::from_parameter(num_threads), + &num_threads, + |b, &threads| { + let printed_clone = printed.clone(); + let tokenizer = tokenizer_clone.clone(); + + b.iter_custom(|_iters| { + let total_operations = Arc::new(AtomicU64::new(0)); + let total_tokens = Arc::new(AtomicU64::new(0)); + let start = Instant::now(); + + let handles: Vec<_> = (0..threads) + .map(|_| { + let tokenizer = tokenizer.clone(); + let total_ops = total_operations.clone(); + let total_tok = total_tokens.clone(); + + thread::spawn(move || { + for _ in 0..operations_per_thread { + let encoding = tokenizer.encode(test_prompt).unwrap(); + total_tok.fetch_add( + encoding.token_ids().len() as u64, + Ordering::Relaxed, + ); + } + total_ops + .fetch_add(operations_per_thread as u64, Ordering::Relaxed); + }) + }) + .collect(); + + for handle in handles { + handle.join().unwrap(); + } + + let duration = start.elapsed(); + + if !printed_clone.load(Ordering::Relaxed) { + let total_ops = total_operations.load(Ordering::Relaxed); + let total_tok = total_tokens.load(Ordering::Relaxed); + let ops_per_sec = total_ops as f64 / duration.as_secs_f64(); + let tokens_per_sec = total_tok as f64 / duration.as_secs_f64(); + + if threads == 1 { + baseline_throughput = tokens_per_sec; + } + + let efficiency = if threads == 1 { + 100.0 + } else { + (tokens_per_sec / (baseline_throughput * threads as f64)) * 100.0 + }; + + let result = format!( + "{:<20} | {:>10} | {:>12.0} | {:>12.0} | {:>10} | {:>11.1}%", + format!("encode_{}_threads", threads), + total_ops, + ops_per_sec, + tokens_per_sec, + threads, + efficiency + ); + add_result("mt_encode", result); + + printed_clone.store(true, Ordering::Relaxed); + } + + duration + }); + }, + ); + } + + group.finish(); +} + +fn bench_multithreaded_decode(c: &mut Criterion) { + let tokenizer = Arc::new( + HuggingFaceTokenizer::from_file(get_tokenizer_path().to_str().unwrap()) + .expect("Failed to load tokenizer"), + ); + + let thread_counts = vec![1, 2, 4, 8, 16]; + let tokens_per_thread = 5000; + + // Generate tokens for decoding + let test_text = "The quick brown fox jumps over the lazy dog. ".repeat(100); + let encoding = tokenizer.encode(&test_text).unwrap(); + let test_tokens: Vec = encoding.token_ids().to_vec(); + + let mut group = c.benchmark_group("multithreaded_decode"); + group.measurement_time(Duration::from_secs(2)); + + let mut baseline_throughput = 0.0; + + for num_threads in thread_counts { + let printed = Arc::new(AtomicBool::new(false)); + let tokenizer_clone = tokenizer.clone(); + let tokens = test_tokens.clone(); + + group.bench_with_input( + BenchmarkId::from_parameter(num_threads), + &num_threads, + |b, &threads| { + let printed_clone = printed.clone(); + let tokenizer = tokenizer_clone.clone(); + let tokens = tokens.clone(); + + b.iter_custom(|_iters| { + let total_tokens = Arc::new(AtomicU64::new(0)); + let start = Instant::now(); + + let handles: Vec<_> = (0..threads) + .map(|_| { + let tokenizer = tokenizer.clone(); + let tokens = tokens.clone(); + let total_tok = total_tokens.clone(); + + thread::spawn(move || { + let mut decoder = DecodeStream::new(tokenizer, &[], false); + let mut output = String::new(); + let mut local_tokens = 0u64; + + for token in tokens.iter().cycle().take(tokens_per_thread) { + if let Some(text) = decoder.step(*token).unwrap() { + output.push_str(&text); + } + local_tokens += 1; + } + + total_tok.fetch_add(local_tokens, Ordering::Relaxed); + }) + }) + .collect(); + + for handle in handles { + handle.join().unwrap(); + } + + let duration = start.elapsed(); + + if !printed_clone.load(Ordering::Relaxed) { + let total = total_tokens.load(Ordering::Relaxed); + let tokens_per_sec = total as f64 / duration.as_secs_f64(); + + if threads == 1 { + baseline_throughput = tokens_per_sec; + } + + let efficiency = if threads == 1 { + 100.0 + } else { + (tokens_per_sec / (baseline_throughput * threads as f64)) * 100.0 + }; + + let result = format!( + "{:<20} | {:>12} | {:>12.0} | {:>10} | {:>11.1}%", + format!("decode_{}_threads", threads), + total, + tokens_per_sec, + threads, + efficiency + ); + add_result("mt_decode", result); + + printed_clone.store(true, Ordering::Relaxed); + } + + duration + }); + }, + ); + } + + group.finish(); +} + +fn bench_memory_efficiency(c: &mut Criterion) { + let tokenizer = Arc::new( + HuggingFaceTokenizer::from_file(get_tokenizer_path().to_str().unwrap()) + .expect("Failed to load tokenizer"), + ); + + let large_text = "The quick brown fox jumps over the lazy dog. ".repeat(1000); + let encoding = tokenizer.encode(&large_text).unwrap(); + + let mut group = c.benchmark_group("memory"); + + // Track owned baseline time + let mut owned_time_ns = 0.0; + + // Owned + let printed_owned = Arc::new(AtomicBool::new(false)); + group.bench_function("token_ids_owned", |b| { + let printed_clone = printed_owned.clone(); + let encoding = encoding.clone(); + + b.iter_custom(|iters| { + let start = Instant::now(); + for _ in 0..iters { + let _ = black_box(encoding.token_ids()); + } + let duration = start.elapsed(); + + if !printed_clone.load(Ordering::Relaxed) { + let ops_per_sec = iters as f64 / duration.as_secs_f64(); + let time_per_call = duration.as_nanos() as f64 / iters as f64; + owned_time_ns = time_per_call; + + let result = format!( + "{:<20} | {:>12.0} | {:>11.0}ns | {:>12}", + "token_ids(owned)", ops_per_sec, time_per_call, "baseline" + ); + add_result("memory", result); + + printed_clone.store(true, Ordering::Relaxed); + } + + duration + }); + }); + + // Reference + let printed_ref = Arc::new(AtomicBool::new(false)); + + group.bench_function("token_ids_ref", |b| { + let printed_clone = printed_ref.clone(); + let encoding = encoding.clone(); + + b.iter_custom(|iters| { + let start = Instant::now(); + for _ in 0..iters { + let _ = black_box(encoding.token_ids()); + } + let duration = start.elapsed(); + + if !printed_clone.load(Ordering::Relaxed) { + let ops_per_sec = iters as f64 / duration.as_secs_f64(); + let time_per_call = duration.as_nanos() as f64 / iters as f64; + + // Calculate improvement + let improvement = if owned_time_ns > 0.0 { + format!("{:.1}x faster", owned_time_ns / time_per_call) + } else { + "N/A".to_string() + }; + + let result = format!( + "{:<20} | {:>12.0} | {:>11.0}ns | {:>12}", + "token_ids_ref(ref)", ops_per_sec, time_per_call, improvement + ); + add_result("memory", result); + + printed_clone.store(true, Ordering::Relaxed); + } + + duration + }); + }); + + group.finish(); +} + +fn bench_scaling_characteristics(c: &mut Criterion) { + let tokenizer = Arc::new( + HuggingFaceTokenizer::from_file(get_tokenizer_path().to_str().unwrap()) + .expect("Failed to load tokenizer"), + ); + + let thread_counts = vec![1, 2, 4, 8, 16]; + let tokens_per_thread = 10000; + + let mut group = c.benchmark_group("scaling"); + group.measurement_time(Duration::from_secs(2)); + + let mut baseline_throughput = 0.0; + + for num_threads in thread_counts { + let printed = Arc::new(AtomicBool::new(false)); + + group.bench_with_input( + BenchmarkId::from_parameter(num_threads), + &num_threads, + |b, &threads| { + let printed_clone = printed.clone(); + let tokenizer = tokenizer.clone(); + + b.iter_custom(|_iters| { + let total_tokens = Arc::new(AtomicU64::new(0)); + let start = Instant::now(); + + let handles: Vec<_> = (0..threads) + .map(|_| { + let tokenizer = tokenizer.clone(); + let total_tokens = total_tokens.clone(); + + thread::spawn(move || { + let mut decoder = DecodeStream::new(tokenizer, &[], false); + let sample_tokens = [1, 450, 6635, 3290, 491]; + + for token in sample_tokens.iter().cycle().take(tokens_per_thread) { + let _ = decoder.step(*token).unwrap(); + } + + total_tokens.fetch_add(tokens_per_thread as u64, Ordering::Relaxed); + }) + }) + .collect(); + + for handle in handles { + handle.join().unwrap(); + } + + let duration = start.elapsed(); + + if !printed_clone.load(Ordering::Relaxed) { + let total = total_tokens.load(Ordering::Relaxed); + let throughput = total as f64 / duration.as_secs_f64(); + + if threads == 1 { + baseline_throughput = throughput; + } + + let efficiency = if threads == 1 { + 100.0 + } else { + (throughput / (baseline_throughput * threads as f64)) * 100.0 + }; + + let result = format!( + "{:<15} | {:>12} | {:>12.0} | {:>11.1}%", + format!("{}_threads", threads), + total, + throughput, + efficiency + ); + add_result("scaling", result); + + printed_clone.store(true, Ordering::Relaxed); + } + + duration + }); + }, + ); + } + + group.finish(); +} + +// Print final summary table +fn print_summary() { + println!("\n{}", "=".repeat(120)); + println!("TOKENIZER BENCHMARK SUMMARY"); + println!("{}", "=".repeat(120)); + + let results = BENCHMARK_RESULTS.lock().unwrap(); + + let mut current_category = String::new(); + for (key, value) in results.iter() { + let category = key.split('_').skip(1).collect::>().join("_"); + + if category != current_category { + current_category = category.clone(); + + // Print section header based on category + println!("\n{}", "-".repeat(120)); + match category.as_str() { + "encode" => { + println!("ENCODING THROUGHPUT"); + println!( + "{:<15} | {:>8} | {:>8} | {:>12} | {:>12} | {:>10} | {:>10}", + "Test Case", + "Size(B)", + "Tokens", + "Chars/sec", + "Tokens/sec", + "Ops/sec", + "Thread" + ); + } + "batch" => { + println!("BATCH ENCODING"); + println!( + "{:<15} | {:>8} | {:>8} | {:>12} | {:>12} | {:>10} | {:>10}", + "Batch Size", + "Size(B)", + "Tokens", + "Prompts/sec", + "Tokens/sec", + "Chars/sec", + "Thread" + ); + } + "concurrent" => { + println!("CONCURRENT ENCODING"); + println!( + "{:<15} | {:>10} | {:>12} | {:>12} | {:>15}", + "Clients", "Total Ops", "Ops/sec", "Chars/sec", "Per-Client/sec" + ); + } + "mt_encode" => { + println!("MULTI-THREADED ENCODING"); + println!( + "{:<20} | {:>10} | {:>12} | {:>12} | {:>10} | {:>12}", + "Configuration", + "Total Ops", + "Ops/sec", + "Tokens/sec", + "Threads", + "Efficiency" + ); + } + "decode" => { + println!("DECODE PERFORMANCE"); + println!( + "{:<20} | {:>10} | {:>12} | {:>12} | {:>10}", + "Method", "Tokens", "Tokens/sec", "Ops/sec", "Thread" + ); + } + "mt_decode" => { + println!("MULTI-THREADED DECODING"); + println!( + "{:<20} | {:>12} | {:>12} | {:>10} | {:>12}", + "Configuration", "Total Tokens", "Tokens/sec", "Threads", "Efficiency" + ); + } + "streaming_100k" => { + println!("STREAMING DECODE (100K Target)"); + println!( + "{:<20} | {:>12} | {:>12} | {:>12} | {:>10} | {:>12}", + "Method", "Tokens", "Tokens/sec", "Target", "Thread", "Status" + ); + } + "concurrent_streaming" => { + println!("CONCURRENT STREAMING"); + println!( + "{:<20} | {:>10} | {:>12} | {:>15} | {:>15}", + "Sequences", "Total", "Aggregate/sec", "Per-Seq/sec", "Threads" + ); + } + "stop_sequences" => { + println!("STOP SEQUENCE PERFORMANCE"); + println!( + "{:<20} | {:>10} | {:>12} | {:>12} | {:>10}", + "Config", "Sequences", "Tokens", "Tokens/sec", "Seq/sec" + ); + } + "latency" => { + println!("LATENCY DISTRIBUTION"); + println!( + "{:<20} | {:>10} | {:>10} | {:>10} | {:>10} | {:>10}", + "Operation", "P50(µs)", "P95(µs)", "P99(µs)", "Max(µs)", "Samples" + ); + } + "scaling" => { + println!("SCALING CHARACTERISTICS"); + println!( + "{:<15} | {:>12} | {:>12} | {:>12}", + "Threads", "Total Tokens", "Tokens/sec", "Efficiency" + ); + } + "memory" => { + println!("MEMORY EFFICIENCY"); + println!( + "{:<20} | {:>12} | {:>12} | {:>12}", + "Operation", "Calls/sec", "Time/call", "Improvement" + ); + } + _ => {} + } + println!("{}", "-".repeat(120)); + } + + println!("{}", value); + } + + println!("\n{}", "=".repeat(120)); +} + +fn run_benchmarks(c: &mut Criterion) { + bench_encode_throughput(c); + bench_batch_encode(c); + bench_concurrent_encode(c); + bench_multithreaded_encode(c); + bench_decode_performance(c); + bench_multithreaded_decode(c); + bench_streaming_decode_100k(c); + bench_concurrent_streaming(c); + bench_stop_sequences(c); + bench_latency_distribution(c); + bench_scaling_characteristics(c); + bench_memory_efficiency(c); + + // Print summary at the end + print_summary(); +} + +criterion_group!(benches, run_benchmarks); +criterion::criterion_main!(benches); diff --git a/lmdeploy/router/docs/lmdeploy_support_guide.md b/lmdeploy/router/docs/lmdeploy_support_guide.md new file mode 100644 index 0000000000..8129def3e8 --- /dev/null +++ b/lmdeploy/router/docs/lmdeploy_support_guide.md @@ -0,0 +1,869 @@ +# Router 支持 LMDeploy 推理与 PD 分离 — 编译与测试指南 + +## 目录 + +01. [概述](#1-概述) +02. [代码改动清单](#2-代码改动清单) +03. [环境准备](#3-环境准备) +04. [编译与单元测试](#4-编译与单元测试) +05. [集成测试环境](#5-集成测试环境) +06. [负载均衡策略测试](#6-负载均衡策略测试) +07. [适配 lmdeploy 测试套件](#7-适配-lmdeploy-测试套件) +08. [测试结果汇总](#8-测试结果汇总) +09. [完整复现步骤](#9-完整复现步骤) +10. [Python wheel 打包与安装](#10-python-wheel-打包与安装) + +______________________________________________________________________ + +## 1. 概述 + +本指南描述如何为 lmdeploy-router 添加 LMDeploy 推理引擎支持，包括： + +- **Phase 1（普通模式）**：Router 天然支持 LMDeploy 的 OpenAI 兼容端点（`/v1/chat/completions`、`/v1/completions`、`/v1/models`、`/generate`），无需代码改动。 +- **Phase 2（PD 分离）**：新增 `LMDeployPDRouter`，支持 LMDeploy 的 DistServe PD 分离协议（`migration_request` + `/distserve/p2p_initialize` + `/distserve/p2p_connect` 两阶段流程）。 +- **动态注册**：Router 可以不指定任何后端 URL 先启动；LMDeploy API Server 再通过原生 `--proxy-url` 参数向 Router 的 `/nodes/add` 注册。普通 Hybrid 和 PD Prefill/Decode 两种部署均支持。 +- **Token-in-Token-out**：支持 LMDeploy 的 `input_ids`/`output_ids` 透传，并为所有负载均衡策略提供基于 token-id 的路由键。 + +### PD 协议差异 + +| 维度 | 传统 PD 语义 | LMDeploy | +| ---------------- | ------------------------------------------------------- | ----------------------------------------------------------------------------------------- | +| PD 协调字段 | `kv_transfer_params`（请求体内嵌） | `migration_request`（decode 请求体） | +| KV 传输 | nixl/mooncake/moriio engine-to-engine | RDMA (DLSlime/Mooncake) out-of-band | +| P2P 建立 | bootstrap 信息交换（无显式端点） | `/distserve/p2p_initialize` + `/distserve/p2p_connect`（显式 HTTP） | +| Prefill 特殊参数 | `max_tokens=1, stream=false` | `max_tokens=1, with_cache=true, preserve_cache=true` | +| 独立 token API | `/inference/v1/generate` + `GenerateRequest{token_ids}` | `/generate` + `GenerateReqInput{input_ids}`；PD 模式的 token-in/out 使用 Chat Completions | + +______________________________________________________________________ + +## 2. 代码改动清单 + +核心实现涉及以下文件（包含初始 LMDeploy 适配和动态注册功能）： + +| 文件 | 类型 | 改动说明 | +| ---------------------------------------- | -------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | +| `src/config/types.rs` | 修改 | 新增 `LMDeployMigrationProtocol` 枚举、`LMDeployRdmaConfig` 结构体、`RoutingMode::LMDeployPrefillDecode` 变体；更新 `is_pd_mode`/`is_lmdeploy_pd_mode`/`worker_count`/`get_prefill_policy`/`get_decode_policy`/`mode_type`；新增序列化测试 | +| `src/config/validation.rs` | 修改 | 新增 `LMDeployPrefillDecode` 的 `validate_mode`/`validate_discovery`/`validate_compatibility` 三个分支 | +| `src/routers/http/mod.rs` | 修改 | 注册 `pub mod lmdeploy_pd_router;` | +| `src/routers/http/lmdeploy_pd_router.rs` | **新建** | `LMDeployPDRouter` 实现 | +| `src/routers/factory.rs` | 修改 | 新增 `LMDeployPrefillDecode` match 分支 + `create_lmdeploy_pd_router()` | +| `src/main.rs` | 修改 | LMDeploy PD CLI flags + mode 选择分支 + `pd_mode` 联动；允许普通/LMDeploy PD 模式零 URL 启动 | +| `src/lib.rs` | 修改 | PyO3 Router 配置字段 + 默认值 + `new()` 入参 + `to_router_config` mode 分支 | +| `src/protocols/spec.rs` | 修改 | 为 Chat 和 Generate 请求显式建模 `input_ids`，直接生成 cache-aware/consistent-hash 路由键；为 Chat 普通/流式响应和 Generate 响应建模 `output_ids`；保留 `GenerateRequest` 的 `#[serde(flatten)] other` 以透明转发 LMDeploy 顶层扩展参数 | +| `src/server.rs` | 修改 | 新增与 LMDeploy 原生 proxy 兼容的 `/nodes/add`、`/nodes/remove`、`/nodes/status` 接口，并按 Hybrid/Prefill/Decode role 分发注册 | +| `src/routers/http/router.rs` | 修改 | 普通 Router 支持 LMDeploy startup callback 阶段的无探活注册 | +| `src/routers/http/pd_router.rs` | 修改 | PD Router 支持 Prefill/Decode startup callback 阶段的无探活注册 | +| `tests/lmdeploy_registration_test.rs` | 新建 | 覆盖普通和 PD 动态注册、重复注册、状态查询、注销及非法 role | + +### 关键改动详解 + +#### 2.1 `RoutingMode::LMDeployPrefillDecode` (types.rs) + +```rust +LMDeployPrefillDecode { + prefill_urls: Vec, + decode_urls: Vec, + prefill_policy: Option, + decode_policy: Option, + migration_protocol: LMDeployMigrationProtocol, + rdma_config: Option, + dummy_prefill: bool, +} +``` + +serde tag: `"lmdeploy_prefill_decode"` + +#### 2.2 `GenerateRequest` flatten 修复 (spec.rs) + +该字段用于普通 Router 的透明转发。LMDeploy 当前的 `/generate` 和 `/v1/responses` +端点不暴露 PD 所需的 cache migration 元数据，因此 LMDeploy PD Router 会对这两个 +端点返回 `501 Not Implemented`；PD 模式只支持 `/v1/chat/completions` 和 +`/v1/completions`。 + +**Fix**: 在 `GenerateRequest` 结构体末尾添加： + +```rust +#[serde(flatten)] +pub other: serde_json::Map, +``` + +#### 2.3 Token-in 路由键 (spec.rs) + +`ChatCompletionRequest.input_ids` 使用显式的 `Option>`，与 LMDeploy 的 +`list[int] | None` 一致；`GenerateRequest.input_ids` 使用 `Option`，同时兼容 +LMDeploy 单请求和 Router 的 batch 格式。`extract_text_for_routing` 直接读取 typed +字段并拼接为稳定路由键，优先级为：`session_id` -> `input_ids` -> 空字符串，不再从 +`#[serde(flatten)] other` 中二次解析 JSON。 + +#### 2.4 CLI flags (main.rs) + +| Flag | 类型 | 默认值 | 说明 | +| ------------------------------- | ---- | ------ | ------------------------- | +| `--lmdeploy-pd-disaggregation` | bool | false | 启用 LMDeploy PD 分离模式 | +| `--prefill` | Vec | - | prefill 节点 URL | +| `--decode` | Vec | - | decode 节点 URL | +| `--lmdeploy-migration-protocol` | enum | rdma | 迁移协议：rdma/nvlink | +| `--lmdeploy-rdma-link-type` | enum | roce | RDMA 链路类型：roce/ib | +| `--lmdeploy-disable-gdr` | bool | false | 禁用 GPU Direct RDMA | +| `--lmdeploy-dummy-prefill` | bool | false | 是否使用 dummy prefill | + +RDMA 设备选择不通过 Router 请求传递。DLSlime 请使用 +`SLIME_VISIBLE_DEVICES`，Mooncake 使用 LMDeploy 后端发现的 RDMA 网卡。 + +______________________________________________________________________ + +## 3. 环境准备 + +### 3.1 安装 Rust 工具链 + +```bash +# 安装 rustup +curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y +source $HOME/.cargo/env + +# 验证 +rustc --version # 1.97.1 +cargo --version # 1.97.1 +``` + +### 3.2 配置 Cargo（磁盘空间 + 代理） + +以下命令均从 Router 仓库根目录执行。默认使用仓库内的相对构建和测试产物目录；如磁盘 +空间不足，可在执行命令前自行将 `CARGO_TARGET_DIR` 指向其他磁盘： + +```bash +export ROUTER_ROOT="$(pwd)" +export CARGO_TARGET_DIR="${CARGO_TARGET_DIR:-target}" +export CARGO_HTTP_MULTIPLEXING=false # 代理环境下 HTTP/2 多路复用会挂起 +export LMDEPLOY_ROUTER_BIN="${CARGO_TARGET_DIR}/release/lmdeploy-router" +export LMDEPLOY_TEST_ARTIFACT_DIR="${LMDEPLOY_TEST_ARTIFACT_DIR:-.test-artifacts/lmdeploy}" +mkdir -p "${LMDEPLOY_TEST_ARTIFACT_DIR}" +``` + +### 3.3 配置 crates.io 镜像（中国大陆） + +```bash +mkdir -p "${CARGO_HOME:-$HOME/.cargo}" +cat > "${CARGO_HOME:-$HOME/.cargo}/config.toml" << 'EOF' +[source.crates-io] +replace-with = "rsproxy-sparse" + +[source.rsproxy-sparse] +registry = "sparse+https://rsproxy.cn/index/" + +[net] +git-fetch-with-cli = true +EOF +``` + +### 3.4 安装 Python 测试依赖 + +```bash +python -m pip install -e '.[dev]' +``` + +测试使用仓库的 `pytest.ini` 和 `py_test/integration/conftest.py`，无需从 `/tmp` 复制或 +单独维护脚本。 + +______________________________________________________________________ + +## 4. 编译与单元测试 + +### 4.1 编译检查 + +```bash +cargo check --lib +``` + +预期输出： + +``` + Finished dev profile [unoptimized + debuginfo] target(s) in 11.71s +``` + +命令应无编译错误；warning 应逐项确认，文档不将 warning 视为默认可忽略项。 + +### 4.2 运行单元测试 + +```bash +cargo test --lib +``` + +本次验证结果：**524 passed, 0 failed**。 + +关键测试： + +- `config::types::test_lmdeploy_prefill_decode_serialization` +- `config::types::test_lmdeploy_migration_protocol_serde_variants` +- `protocols::spec::test_chat_completion_extract_routing_with_input_ids` +- `protocols::spec::test_chat_completion_preserves_typed_input_ids_on_serialize` +- `protocols::spec::test_lmdeploy_chat_response_preserves_output_ids` +- `protocols::spec::test_lmdeploy_generate_response_preserves_output_ids` +- `routers::http::lmdeploy_pd_router::test_build_migration_request_extracts_fields` +- `routers::http::lmdeploy_pd_router::test_prepare_prefill_request_sets_with_cache_and_preserve_cache` + +### 4.3 编译二进制 + +```bash +cargo build --release --bin lmdeploy-router +``` + +预期输出： + +``` + Finished release profile [optimized] target(s) in ... +``` + +二进制位置：`${CARGO_TARGET_DIR}/release/lmdeploy-router`。后续命令统一使用 +`${LMDEPLOY_ROUTER_BIN}`，不依赖某台机器的仓库绝对路径。 + +______________________________________________________________________ + +## 5. 集成测试环境 + +### 5.1 后端部署 + +两个 LMDeploy api_server 实例（非 PD 分离，普通模式）： + +```bash +export MODEL_PATH="${MODEL_PATH:-qwen3-06b}" + +# Backend 1 (port 9000) +lmdeploy serve api_server "${MODEL_PATH}" \ + --server-name 0.0.0.0 --server-port 9000 \ + --logprobs-mode raw_logprobs + +# Backend 2 (port 9003；与 PD Prefill 的 9001 错开) +lmdeploy serve api_server "${MODEL_PATH}" \ + --server-name 0.0.0.0 --server-port 9003 \ + --logprobs-mode raw_logprobs +``` + +验证后端： + +```bash +curl http://10.102.98.154:9000/v1/models +# {"id":"qwen3-06b","owned_by":"lmdeploy"} + +curl http://10.102.98.154:9003/v1/models +# {"id":"qwen3-06b","owned_by":"lmdeploy"} +``` + +### 5.2 启动 Router + +```bash +# 普通模式（非 PD） +setsid "${LMDEPLOY_ROUTER_BIN}" \ + --host 0.0.0.0 \ + --port 30000 \ + --worker-urls http://10.102.98.154:9000 http://10.102.98.154:9003 \ + --policy random \ + --prometheus-port 19001 \ + > "${LMDEPLOY_TEST_ARTIFACT_DIR}/router.log" 2>&1 < /dev/null & +disown +``` + +验证 Router： + +```bash +curl http://127.0.0.1:30000/health +# All servers healthy + +curl http://127.0.0.1:30000/v1/models +# {"data":[{"id":"qwen3-06b",...}]} + +curl http://127.0.0.1:30000/list_workers +# {"urls":["http://10.102.98.154:9000","http://10.102.98.154:9003"]} +``` + +### 5.3 PD 分离模式静态启动（可选） + +```bash +setsid "${LMDEPLOY_ROUTER_BIN}" \ + --host 0.0.0.0 \ + --port 30000 \ + --lmdeploy-pd-disaggregation \ + --prefill http://prefill-host:23333 \ + --decode http://decode-host:23333 \ + --lmdeploy-migration-protocol rdma \ + --lmdeploy-rdma-link-type roce \ + --prometheus-port 19001 \ + > "${LMDEPLOY_TEST_ARTIFACT_DIR}/router_pd.log" 2>&1 < /dev/null & +disown +``` + +### 5.4 LMDeploy 原生动态注册概览 + +除静态传入 `--worker-urls`、`--prefill` 和 `--decode` 外，Router 还兼容 +LMDeploy 原生 proxy 的启动方式： + +1. 先启动一个没有后端 URL 的 Router。 +2. 启动 LMDeploy API Server，并传入 `--proxy-url http://:`。 +3. LMDeploy 在 FastAPI startup callback 中向 Router 发送 `POST /nodes/add`。 +4. Router 根据注册请求中的 role，将节点加入普通、Prefill 或 Decode worker pool。 +5. API Server 开始监听后，Router 的健康检查器负责维护节点可用状态。 + +LMDeploy 启动时发送的请求体如下： + +```json +{ + "url": "http://10.102.98.154:9001", + "status": { + "models": ["qwen3_06B"], + "role": 2 + } +} +``` + +role 与 Router 模式的对应关系： + +| role 数值 | LMDeploy role | Router 启动模式 | 注册目标 | +| --------- | ------------- | ------------------------------ | ------------------- | +| `1` | `Hybrid` | 普通模式 | Regular worker pool | +| `2` | `Prefill` | `--lmdeploy-pd-disaggregation` | Prefill worker pool | +| `3` | `Decode` | `--lmdeploy-pd-disaggregation` | Decode worker pool | + +普通模式只接受 Hybrid，LMDeploy PD 模式只接受 Prefill/Decode。role 与 Router +模式不匹配时，`/nodes/add` 返回 `400 Bad Request`。相同 URL 和 role 的重复注册是 +幂等操作，不会产生重复 worker。 + +> **为什么动态注册不在 `/nodes/add` 中同步探活？** +> +> LMDeploy 从 FastAPI startup callback 发起注册，此时 Uvicorn 尚未开始接受 +> `/health` 请求。如果 Router 在处理 `/nodes/add` 时等待该节点健康，两端会互相 +> 等待直到超时。动态注册路径因此先写入 worker registry，随后由 Router 的后台 +> 健康检查器接管；原有 `/add_worker` 和静态 URL 启动流程仍保留探活行为。 + +### 5.5 普通模式动态注册（Hybrid） + +先启动不带 `--worker-urls` 的 Router： + +```bash +setsid "${LMDEPLOY_ROUTER_BIN}" \ + --host 0.0.0.0 \ + --port 30000 \ + --policy round_robin \ + --prometheus-port 19001 \ + > "${LMDEPLOY_TEST_ARTIFACT_DIR}/router_dynamic.log" 2>&1 < /dev/null & +disown +``` + +然后启动 LMDeploy Hybrid 服务。`--role Hybrid` 是默认值，建议显式写出以便排查： + +```bash +CUDA_VISIBLE_DEVICES=0 lmdeploy serve api_server "${MODEL_PATH}" \ + --tp 1 \ + --backend pytorch \ + --server-name 10.102.98.154 \ + --server-port 9000 \ + --model-name qwen3_06B \ + --role Hybrid \ + --proxy-url http://ROUTER_IP:30000 +``` + +验证注册状态和数据面： + +```bash +curl http://ROUTER_IP:30000/nodes/status +# { +# "http://10.102.98.154:9000": { +# "role": 1, +# "models": ["qwen3_06B"], +# "unfinished": 0, +# "latency": [], +# "speed": null +# } +# } + +curl http://ROUTER_IP:30000/list_workers +# {"urls":["http://10.102.98.154:9000"]} + +curl -X POST http://ROUTER_IP:30000/v1/chat/completions \ + -H 'Content-Type: application/json' \ + -d '{ + "model": "qwen3_06B", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 32, + "stream": false + }' +``` + +可以启动多个 Hybrid API Server，并将它们的 `--proxy-url` 指向同一个 Router；每个 +实例都会独立注册，Router 按 `--policy` 选择节点。 + +### 5.6 PD 分离模式动态注册（Prefill/Decode） + +先启动不带 `--prefill`、`--decode` 的 LMDeploy PD Router： + +```bash +setsid "${LMDEPLOY_ROUTER_BIN}" \ + --host 0.0.0.0 \ + --port 30000 \ + --lmdeploy-pd-disaggregation \ + --lmdeploy-migration-protocol rdma \ + --lmdeploy-rdma-link-type roce \ + --policy round_robin \ + --prefill-policy round_robin \ + --decode-policy round_robin \ + --prometheus-port 19001 \ + > "${LMDEPLOY_TEST_ARTIFACT_DIR}/router_pd_dynamic.log" 2>&1 < /dev/null & +disown +``` + +再分别启动 Prefill 和 Decode。两端的 `--migration-backend` 必须一致，并与实际环境 +配置匹配；以下示例使用 DLSlime： + +```bash +export SLIME_VISIBLE_DEVICES="mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_6,mlx5_7" +export SLIME_GID_INDEX=3 + +# Prefill: GPU 1 / port 9001 +CUDA_VISIBLE_DEVICES=1 lmdeploy serve api_server "${MODEL_PATH}" \ + --tp 1 \ + --backend pytorch \ + --server-name 10.102.98.154 \ + --server-port 9001 \ + --model-name qwen3_06B \ + --role Prefill \ + --migration-backend DLSlime \ + --proxy-url http://ROUTER_IP:30000 + +# Decode: GPU 2 / port 9002 +CUDA_VISIBLE_DEVICES=2 lmdeploy serve api_server "${MODEL_PATH}" \ + --tp 1 \ + --backend pytorch \ + --server-name 10.102.98.154 \ + --server-port 9002 \ + --model-name qwen3_06B \ + --role Decode \ + --migration-backend DLSlime \ + --proxy-url http://ROUTER_IP:30000 +``` + +Mooncake 部署只需将两端改为 `--migration-backend Mooncake`，并完成对应的 Mooncake +网卡和存储配置。Router 的 `--lmdeploy-migration-protocol` 与 +`--lmdeploy-rdma-link-type` 描述 Router 发给 LMDeploy 的 migration request；它们不 +替代 DLSlime/Mooncake 自身的网卡环境变量。 + +验证两个 role 均已注册： + +```bash +curl http://ROUTER_IP:30000/nodes/status +# { +# "http://10.102.98.154:9001": { +# "role": 2, +# "models": ["qwen3_06B"], +# "unfinished": 0, +# "latency": [], +# "speed": null +# }, +# "http://10.102.98.154:9002": { +# "role": 3, +# "models": ["qwen3_06B"], +# "unfinished": 0, +# "latency": [], +# "speed": null +# } +# } +``` + +此后向 Router 的 `/v1/chat/completions` 或 `/v1/completions` 发请求即可触发 +Prefill、KV cache migration、Decode 两阶段流程。LMDeploy PD 模式下 `/generate` +和 `/v1/responses` 不具备 migration 元数据，Router 会返回 `501 Not Implemented`。 + +### 5.7 注册管理接口 + +#### 查询节点 + +```bash +curl http://ROUTER_IP:30000/nodes/status +``` + +响应是以节点 URL 为 key、LMDeploy `Status` 为 value 的 JSON 对象，字段包括 +`role`、`models`、`unfinished`、`latency` 和 `speed`。 + +#### 手动注册 + +正常情况下无需手动调用；以下命令可用于调试或接入自定义启动器： + +```bash +curl -X POST http://ROUTER_IP:30000/nodes/add \ + -H 'Content-Type: application/json' \ + -d '{ + "url": "http://10.102.98.154:9000", + "status": {"models": ["qwen3_06B"], "role": 1} + }' +# "Added successfully" +``` + +#### 注销节点 + +```bash +curl -X POST http://ROUTER_IP:30000/nodes/remove \ + -H 'Content-Type: application/json' \ + -d '{"url": "http://10.102.98.154:9000"}' +# "Deleted successfully" +``` + +注销是幂等操作。当前 LMDeploy API Server 的 shutdown callback 不会主动请求 +`/nodes/remove`；进程退出后，Router 的健康检查会将 worker 标记为不可用。如需立即 +从 `/nodes/status` 和 worker registry 删除节点，应由进程管理器调用 +`/nodes/remove`。同一 URL 的服务重启后可以直接重新注册。 + +### 5.8 网络、安全与排障 + +动态注册要求网络双向可达： + +- LMDeploy API Server 必须能访问 `--proxy-url` 指向的 Router。 +- Router 必须能访问注册请求体 `url` 指向的 LMDeploy API Server。 +- 跨机器部署时，`--server-name` 必须使用 Router 可路由的节点 IP 或主机名，不能使用 + `0.0.0.0`、`127.0.0.1` 或仅在容器内部可解析的地址。 +- `--server-name` 还决定 Uvicorn 的监听地址；该 IP 必须存在于 LMDeploy 所在主机或 + 容器的网络命名空间中。 + +LMDeploy 原生 `--proxy-url` 注册请求不携带 Authorization header，因此 +`/nodes/add`、`/nodes/remove` 和 `/nodes/status` 是未鉴权的控制面接口。生产环境应通过 +安全组、防火墙或反向代理，只允许可信 LMDeploy 节点访问这些路径。 + +常见问题： + +| 现象 | 检查项 | +| ------------------------------------------ | ---------------------------------------------------------------------------------------------------------------- | +| API Server 已启动，但 `/nodes/status` 为空 | 检查 LMDeploy 日志是否出现 `Service registration failed`；从 API Server 所在机器 `curl /nodes/status` | +| 注册成功，但请求报连接失败 | 检查 `--server-name` 是否为 Router 可回连地址；确认注册 JSON 中没有 `0.0.0.0` 或 `127.0.0.1` | +| `/nodes/add` 长时间卡住 | 确认使用的是包含动态注册修复的 Router；注册接口不应同步等待 API Server `/health` | +| PD 注册返回 `400` | Router 必须带 `--lmdeploy-pd-disaggregation`，Prefill/Decode 的 role 分别为 `2`/`3` | +| 普通注册返回 `400` | 普通 Router 只接受 Hybrid role `1` | +| PD 已注册但迁移失败 | 检查两端 migration backend、RDMA 网卡/GID、Router migration protocol/link type 是否一致 | + +静态 URL 与动态注册可以共存。例如可用 `--prefill` 配置固定 Prefill，同时让新增的 +Decode 节点通过 `--proxy-url` 注册；不过生产环境建议采用一致的节点生命周期管理方式， +以便排障和容量统计。 + +______________________________________________________________________ + +## 6. 负载均衡策略测试 + +### 6.1 pytest 测试位置 + +测试位于 `py_test/integration/lmdeploy/test_load_balancing.py`，使用 Router 现有的 pytest +规范和 fixture： + +- 每个策略自动分配 Router 与 Prometheus 端口，并在用例结束后回收进程。 +- 每个策略发送 10 个 messages 请求和 10 个 `input_ids` 请求。 +- 通过 LMDeploy 每个 worker 独立递增的数值响应 `id` 判断实际 worker。 +- `temperature=0`，严格比较 Router 响应与对应直连后端的文本；token-in 请求同时比较 + `output_ids`。 +- 日志默认写入相对目录 `.test-artifacts/lmdeploy/`。 + +### 6.2 配置真实后端 + +```bash +export LMDEPLOY_BACKEND_URLS="http://10.102.98.154:9000,http://10.102.98.154:9003" +export LMDEPLOY_ROUTER_BIN="${CARGO_TARGET_DIR}/release/lmdeploy-router" +``` + +`LMDEPLOY_MODEL` 默认从所有后端的 `/v1/models` 自动发现；只有后端暴露多个模型且需要 +指定其中一个时才设置，例如 `export LMDEPLOY_MODEL="qwen3-06b"`。 + +如模型 tokenizer 不同，可用 JSON 数组覆盖 token-in 测试数据： + +```bash +export LMDEPLOY_TOKEN_IDS='[151644,8948,198,2610]' +``` + +### 6.3 运行测试 + +```bash +pytest -v py_test/integration/lmdeploy/test_load_balancing.py +``` + +未设置 `LMDEPLOY_BACKEND_URLS` 时测试会 skip，不会把开发机地址硬编码进默认 CI。 + +### 6.4 策略断言 + +| 策略 | 验证逻辑 | +| --------------- | --------------------------------------------------- | +| random | 两个 worker 均处理请求 | +| round_robin | messages 请求严格交替 | +| power_of_two | 两个 worker 均处理请求 | +| consistent_hash | 相同 messages/input_ids 路由键分别固定到同一 worker | +| cache_aware | 相同 input_ids 和递增 token 前缀固定到同一 worker | + +______________________________________________________________________ + +## 7. 适配 lmdeploy 测试套件 + +### 7.1 测试来源 + +适配自 LMDeploy 的 autotest 测试套件： + +- `/autotest/interface/restful/test_restful_chat_completions_v1.py` +- `/autotest/interface/restful/test_restful_completions_v1.py` +- `/autotest/interface/restful/test_restful_generate.py` +- `/autotest/utils/restful_return_check.py`（断言辅助函数） + +### 7.2 Router 测试 + +适配测试位于 `py_test/integration/lmdeploy/test_rest_api.py`，不再使用 `/tmp` 中的独立 +main 脚本。每一个验证点都是独立 pytest case，失败时可直接使用 `-k` 选择并由 pytest +输出 traceback。 + +主要改动： + +- Router 地址和端口由 fixture 动态分配。 +- 后端、模型、超时和 token IDs 通过环境变量配置。 +- 跳过 router 不代理的端点测试（`/encode`、tokenizer 相关） +- 新增 strict content match 测试（temp=0 下比较 router vs 直连后端） +- 新增 token-in-token-out 测试（`input_ids` + `return_token_ids` -> `output_ids`） +- 流式响应按 SSE event 独立校验。 + +### 7.3 运行测试 + +```bash +pytest -v py_test/integration/lmdeploy/test_rest_api.py +``` + +### 7.4 测试矩阵（28 项） + +| 类别 | 测试项 | 数量 | 结果 | +| ---------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ---- | ---------- | +| 基础端点 | /v1/models, /health, /list_workers | 3 | 3/3 PASS | +| Chat Completions | basic, streaming, single_stopword, array_stopwords, max_tokens, max_completion_tokens, input_ids(token-in), input_ids_vs_direct, temp=0_determinism, logprobs, invalid_model, strict_content_match | 12 | 12/12 PASS | +| Completions | basic, streaming, stopword, max_tokens, strict_content_match | 5 | 5/5 PASS | +| Generate | basic, input_ids(token-in), streaming, stop_token_ids, session_id, conflict_prompt+input_ids(400), empty_prompt(400), strict_content_match | 8 | 8/8 PASS | + +### 7.5 注册与 PD 测试 + +`py_test/integration/lmdeploy/test_registration.py` 覆盖： + +- Hybrid role=1 动态注册、重复注册幂等、完整状态字段、真实推理和幂等注销。 +- 普通 Router 拒绝 Prefill role，PD Router 拒绝 Hybrid role。 +- Prefill role=2 与 Decode role=3 的 PD 动态注册、重复注册和真实推理。 +- 通过 `--prefill`/`--decode` 静态配置的 PD 真实推理。 +- 静态和动态 PD 模式下 `/generate`、`/v1/responses` 均返回 `501`，避免绕过 KV + migration 流程。 + +PD 真实测试需要额外配置： + +```bash +export LMDEPLOY_PD_PREFILL_URL="http://10.102.98.154:9001" +export LMDEPLOY_PD_DECODE_URL="http://10.102.98.154:9002" +pytest -v py_test/integration/lmdeploy/test_registration.py +``` + +只配置 Hybrid 后端而未配置两个 PD 变量时，两个 PD 真实用例 skip。完整套件共 37 个 +pytest case：28 个 REST、5 个负载均衡策略和 4 个注册/PD 场景。 + +______________________________________________________________________ + +## 8. 测试结果汇总 + +### 8.1 单元测试 + +| 项目 | 结果 | +| --------------------------------------------- | -------------------- | +| `cargo check --lib` | PASS | +| `cargo test --lib` | 524 passed, 0 failed | +| `cargo build --release --bin lmdeploy-router` | PASS | + +### 8.2 集成测试 + +2026-08-04 在 H200 节点上使用两个 Hybrid 服务（9000/9003）和一组 DLSlime PD 服务 +（Prefill 9001、Decode 9002）执行完整命令，结果为 **37 passed in 14.61s**。PD 使用 +RDMA/RoCE、`SLIME_VISIBLE_DEVICES=mlx5_0,...,mlx5_7` 和 `SLIME_GID_INDEX=3`。四个 +LMDeploy 日志均未出现 `ERROR`、`Traceback` 或 `Invalid Free session`。 +容器仍会提示当前 Transformers 5.12.1 超出 LMDeploy 声明的 `[4.33.0, 5.3.0)` 范围； +本轮功能测试未受影响，但正式环境应将依赖版本调整到 LMDeploy 支持范围。 + +| 项目 | 结果 | +| ------------------ | ----------------------------------------------------------------- | +| 5 LB 策略测试 | 5/5 PASS | +| 28 适配测试 | 28/28 PASS | +| Hybrid 动态注册 | PASS（零 worker 启动、role=1 注册、真实推理） | +| PD 动态注册 | PASS（role=2/3 注册、Prefill → RDMA/RoCE → Decode 真实推理） | +| 注册与 PD 场景 | 4/4 PASS（静态/动态、幂等、查询、注销、非法 role、PD 不支持端点） | +| Token-in-token-out | PASS（input_ids 透传 + output_ids 返回） | +| 严格内容匹配 | PASS（temp=0 下 router == 直连后端） | +| logprobs | PASS（`--logprobs-mode raw_logprobs`） | + +### 8.3 发现并修复的 Bug + +| Bug | 文件 | 修复 | +| --------------------------------------------------------------------------------------------- | ---------------------------- | ------------------------------------------------------------------------------ | +| `GenerateRequest` 缺少 `#[serde(flatten)] other` 字段，导致 lmdeploy /generate 参数被静默丢弃 | `src/protocols/spec.rs:1993` | 添加 `#[serde(flatten)] pub other: serde_json::Map` | + +______________________________________________________________________ + +## 9. 完整复现步骤 + +```bash +# ========== 1. 安装 Rust ========== +curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y +source $HOME/.cargo/env + +# ========== 2. 进入仓库并配置相对目录 ========== +cd router +export ROUTER_ROOT="$(pwd)" +export CARGO_TARGET_DIR="${CARGO_TARGET_DIR:-target}" +export CARGO_HTTP_MULTIPLEXING=false +export LMDEPLOY_ROUTER_BIN="${CARGO_TARGET_DIR}/release/lmdeploy-router" +export LMDEPLOY_TEST_ARTIFACT_DIR="${LMDEPLOY_TEST_ARTIFACT_DIR:-.test-artifacts/lmdeploy}" +mkdir -p "${LMDEPLOY_TEST_ARTIFACT_DIR}" + +# ========== 3. 编译 ========== +python -m pip install -e '.[dev]' +cargo check --lib +cargo test --lib +cargo build --release --bin lmdeploy-router + +# ========== 4. 启动后端（两个 lmdeploy 实例） ========== +# Backend 1: 10.102.98.154:9000 +# Backend 2: 10.102.98.154:9003 +# 模型: qwen3-06b +# 需开启 --logprobs-mode raw_logprobs + +# ========== 5. 配置测试后端 ========== +export LMDEPLOY_BACKEND_URLS="http://10.102.98.154:9000,http://10.102.98.154:9003" +# 可选：后端暴露多个模型时指定；否则从 /v1/models 自动发现 +# export LMDEPLOY_MODEL="qwen3-06b" + +# 可选：启用真实 PD 动态注册与推理测试 +export LMDEPLOY_PD_PREFILL_URL="http://10.102.98.154:9001" +export LMDEPLOY_PD_DECODE_URL="http://10.102.98.154:9002" + +# ========== 6. 运行全部 LMDeploy pytest ========== +# fixture 自动启动/停止 Router，并为每个实例分配空闲端口 +pytest -v py_test/integration/lmdeploy +``` + +______________________________________________________________________ + +## 10. Python wheel 打包与安装 + +### 10.1 打包方式 + +`lmdeploy-router` 由仓库根目录的 `pyproject.toml` 与 `setup.py` 通过 setuptools-rust 打包 +成一个 wheel（与上游 vllm-router 的打包方式一致）。该 wheel 同时包含三部分： + +- `lmdeploy_router_rs`：PyO3 扩展模块（`src/lib.rs`），提供 `PolicyType` 与 `Router`。 +- `lmdeploy_router`：纯 Python 包（`py_src/lmdeploy_router/`），提供 `RouterArgs`、 + `launch_router` 和 MiniLB。 +- `lmdeploy-router`：入口命令，等价于 `python -m lmdeploy_router.launch_router`。 + +因此安装同一个 wheel 之后，既可以 `import lmdeploy_router_rs`，也可以直接运行 +`lmdeploy-router` 命令；二者参数完全一致。PyO3 扩展按 abi3 构建 +（`py_limited_api = "cp38"`），同一个 `cp38-abi3` wheel 可用于 CPython 3.8 及以上版本。 +不同操作系统或 CPU 架构仍需分别构建。 + +### 10.2 构建 wheel + +```bash +python -m pip install build 'setuptools-rust>=1.5.2' +python -m build --wheel --outdir dist . +``` + +如需快速验证 Python 打包结构，可设置 `LMDEPLOY_ROUTER_BUILD_NO_RUST=1` 构建 +不含 Rust 扩展和二进制的 wheel；可通过 `lmdeploy_router.ROUTER_AVAILABLE` +判断 Rust API 是否可用。完整发布仍必须构建默认 wheel。 + +产物统一放在仓库根目录的 `dist/` 中： + +```text +dist/lmdeploy_router-0.0.3-cp38-abi3-linux_x86_64.whl +``` + +Rust 依赖较多，首次构建需要数分钟；通过 `CARGO_HOME` 复用缓存可以显著加速。`dist/` 已加入 +`.gitignore`，适合作为本地或 CI 构建产物目录。正式发布时应将 wheel 上传到内部 PyPI、 +制品库或 Release，而不是提交进 Git。wheel 文件名遵循 Python Wheel 规范，不要将其重命名为 +`lmdeploy-router.whl`；pip 会校验文件名中的包名、版本和兼容性标签。 + +`pyproject.toml`、`Cargo.toml` 与 `py_src/lmdeploy_router/version.py` 的版本号必须保持一致。 + +### 10.3 安装与使用 + +在目标 Python 环境中安装构建出的 wheel： + +```bash +python -m pip install ./dist/lmdeploy_router-0.0.3-cp38-abi3-linux_x86_64.whl +lmdeploy-router --help +python -c "from lmdeploy_router_rs import PolicyType; print(PolicyType)" +``` + +入口命令走的是 Python `RouterArgs`，它与 Rust CLI 接受同一套 PD 参数： +`--lmdeploy-pd-disaggregation`、`--lmdeploy-migration-protocol`、 +`--lmdeploy-rdma-link-type`、`--lmdeploy-disable-gdr` 和 `--lmdeploy-dummy-prefill`。 +启动参数与直接运行 Cargo release 二进制完全相同，例如先启动 Router、等待 LMDeploy 服务 +通过 `--proxy-url` 动态注册： + +```bash +lmdeploy-router \ + --host 0.0.0.0 \ + --port 30000 \ + --lmdeploy-pd-disaggregation \ + --lmdeploy-migration-protocol rdma \ + --lmdeploy-rdma-link-type roce +``` + +### 10.4 验证多 Python 版本 + +扩展为 abi3，同一个 wheel 可直接安装在 CPython 3.8 及以上的环境中： + +```bash +WHEEL=dist/lmdeploy_router-0.0.3-cp38-abi3-linux_x86_64.whl + +for PYTHON in python3.10 python3.11 python3.12; do + ENV_DIR=".test-artifacts/wheel/${PYTHON}" + "${PYTHON}" -m venv "${ENV_DIR}" + "${ENV_DIR}/bin/python" -m pip install --no-deps "${WHEEL}" + "${ENV_DIR}/bin/lmdeploy-router" --help +done +``` + +除安装检查外，发布流水线还应在目标 Linux 发行版上运行一次真实的 Router 启动和请求 +转发测试。wheel 的平台标签保证的是 Linux 基础系统兼容范围，不代替运行时功能验证。 + +### 10.5 在 caikun-lmdeploy 容器中验收 + +目标容器当前为 Linux x86_64、Python 3.12，安装及启动验证可执行： + +```bash +WHEEL=dist/lmdeploy_router-0.0.1-py3-none-manylinux_2_17_x86_64.whl +WHEEL_NAME="$(basename "${WHEEL}")" +TARGET=gpu-l-lg-cmc-h-h200-0382.host.h.pjlab.org.cn + +scp "${WHEEL}" "ailab@${TARGET}:/tmp/" +ssh "ailab@${TARGET}" \ + "docker cp /tmp/${WHEEL_NAME} caikun-lmdeploy:/tmp/${WHEEL_NAME} && \ + docker exec caikun-lmdeploy python3 -m pip install --no-deps --force-reinstall /tmp/${WHEEL_NAME} && \ + docker exec caikun-lmdeploy lmdeploy-router --version" +``` + +验收时还应使用待部署的 LMDeploy 参数短暂启动 Router，并请求 `/list_workers` 或 `/health` +确认服务监听成功。不要只检查 pip 的安装返回码。 + +______________________________________________________________________ + +## 附录：文件索引 + +| 文件 | 说明 | +| ----------------------------------------------------- | --------------------------------------------------------------------------------- | +| `src/config/types.rs` | LMDeployMigrationProtocol, LMDeployRdmaConfig, RoutingMode::LMDeployPrefillDecode | +| `src/config/validation.rs` | LMDeployPrefillDecode 验证分支 | +| `src/routers/http/lmdeploy_pd_router.rs` | LMDeployPDRouter 实现 | +| `src/routers/http/router.rs` | 普通模式动态 worker 注册 | +| `src/routers/http/pd_router.rs` | Prefill/Decode 动态 worker 注册 | +| `src/routers/factory.rs` | create_lmdeploy_pd_router 工厂方法 | +| `src/routers/http/mod.rs` | 模块注册 | +| `src/main.rs` | CLI flags + mode 分支 | +| `src/server.rs` | `/nodes/add`、`/nodes/remove`、`/nodes/status` 注册接口 | +| `src/lib.rs` | PyO3 Router 结构 | +| `src/protocols/spec.rs` | typed input_ids/output_ids + GenerateRequest flatten | +| `tests/lmdeploy_registration_test.rs` | 普通/PD 动态注册 Rust 回归测试 | +| `py_test/integration/lmdeploy/conftest.py` | 真实后端配置、Router 生命周期和测试产物 fixture | +| `py_test/integration/lmdeploy/test_load_balancing.py` | 五种 LB 策略测试 | +| `py_test/integration/lmdeploy/test_rest_api.py` | 28 项 LMDeploy REST 适配测试 | +| `py_test/integration/lmdeploy/test_registration.py` | Hybrid/PD 动态注册测试 | +| `pyproject.toml` + `setup.py` | Python wheel 打包配置（setuptools-rust） | +| `py_src/lmdeploy_router/` | Python 包：`RouterArgs`、`launch_router`、MiniLB | +| `${CARGO_TARGET_DIR}/release/lmdeploy-router` | 编译产物（release 二进制） | diff --git a/lmdeploy/router/docs/load_balancing/README.md b/lmdeploy/router/docs/load_balancing/README.md new file mode 100644 index 0000000000..2e5ac3b4c9 --- /dev/null +++ b/lmdeploy/router/docs/load_balancing/README.md @@ -0,0 +1,299 @@ +# Load Balancing Policies + +The LMDeploy Router supports multiple load balancing policies for distributing requests across backend workers. Each policy is designed for different use cases and can be configured based on your deployment requirements. + +## Available Policies + +| Policy | Best For | Session Affinity | Load Aware | +| ----------------- | ---------------------------------------- | ----------------- | ---------- | +| `round_robin` | General purpose, even distribution | No | No | +| `random` | Simple deployments | No | No | +| `consistent_hash` | Multi-turn conversations, KV cache reuse | Yes | No | +| `power_of_two` | Load-sensitive workloads | No | Yes | +| `cache_aware` | Prefix caching optimization | Yes (cache-based) | Yes | + +______________________________________________________________________ + +## Consistent Hash + +The `consistent_hash` policy routes requests with the same session/user identifier to the same backend worker. This is essential for: + +- **Multi-turn conversations**: Ensures conversation history stays on the same worker +- **KV cache reuse**: Maximizes cache hits by routing related requests together +- **Session affinity**: Maintains user-specific state on a single worker + +### Configuration + +```bash +# Using CLI +lmdeploy-router --policy consistent_hash --worker-urls http://worker1:8000,http://worker2:8000 + +# Using Python +from lmdeploy_router import Router +router = Router( + policy="consistent_hash", + worker_urls=["http://worker1:8000", "http://worker2:8000"] +) +``` + +### Hash Key Priority + +The consistent hash policy extracts a routing key in the following priority order: + +| Priority | Source | Header/Field | Example | +| -------- | ------------ | --------------------------- | ------------------------------------------- | +| 1 | HTTP Header | `X-Session-ID` | `X-Session-ID: session-abc-123` | +| 2 | HTTP Header | `X-User-ID` | `X-User-ID: user-456` | +| 3 | HTTP Header | `X-Tenant-ID` | `X-Tenant-ID: tenant-xyz` | +| 4 | HTTP Header | `X-Request-ID` | `X-Request-ID: req-789` | +| 5 | HTTP Header | `X-Correlation-ID` | `X-Correlation-ID: corr-001` | +| 6 | HTTP Header | `X-Trace-ID` | `X-Trace-ID: trace-002` | +| 7 | Request Body | `session_params.session_id` | `{"session_params": {"session_id": "..."}}` | +| 8 | Request Body | `user` | `{"user": "..."}` (OpenAI format) | +| 9 | Request Body | `session_id` | `{"session_id": "..."}` (legacy) | +| 10 | Request Body | `user_id` | `{"user_id": "..."}` (legacy) | +| 11 | Fallback | Request body hash | Hash of entire request body | + +### Usage Examples + +#### Recommended: Using HTTP Headers + +HTTP headers are the **recommended approach** for session affinity because: + +- No JSON body parsing required (faster routing) +- Works with any request format +- Compatible with standard infrastructure tools (Nginx, Envoy, K8s Ingress) + +```bash +# Using X-Session-ID header (recommended) +curl -X POST http://router:8000/v1/chat/completions \ + -H "X-Session-ID: conversation-12345" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "llama-3", + "messages": [{"role": "user", "content": "Hello!"}] + }' + +# Using X-User-ID header for user-based routing +curl -X POST http://router:8000/v1/chat/completions \ + -H "X-User-ID: user-67890" \ + -H "Content-Type: application/json" \ + -d '{ + "model": "llama-3", + "messages": [{"role": "user", "content": "Hello!"}] + }' +``` + +#### Alternative: Using Request Body + +For backward compatibility, you can also include session information in the request body: + +```bash +# Using session_params.session_id in body +curl -X POST http://router:8000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "llama-3", + "session_params": {"session_id": "conversation-12345"}, + "messages": [{"role": "user", "content": "Hello!"}] + }' + +# Using OpenAI user field +curl -X POST http://router:8000/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "llama-3", + "user": "user-67890", + "messages": [{"role": "user", "content": "Hello!"}] + }' +``` + +### Behavior + +- **Consistency**: Same session ID always routes to the same worker +- **Unhealthy fallback**: If the target worker is unhealthy, falls back to the first healthy worker +- **Virtual nodes**: Uses 160 virtual nodes per worker for even distribution +- **DP-aware routing**: Supports data-parallel worker URLs (e.g., `http://worker:8000@0`) + +______________________________________________________________________ + +## Round Robin + +The `round_robin` policy distributes requests evenly across all healthy workers in sequential order. + +### Configuration + +```bash +lmdeploy-router --policy round_robin --worker-urls http://worker1:8000,http://worker2:8000 +``` + +### Behavior + +- Cycles through workers: worker1 → worker2 → worker3 → worker1 → ... +- Skips unhealthy workers automatically +- Simple and predictable distribution +- No session affinity (each request may go to a different worker) + +### Best For + +- Stateless workloads +- Single-turn requests +- When even distribution is more important than cache locality + +______________________________________________________________________ + +## Random + +The `random` policy selects a random healthy worker for each request. + +### Configuration + +```bash +lmdeploy-router --policy random --worker-urls http://worker1:8000,http://worker2:8000 +``` + +### Behavior + +- Uniform random selection among healthy workers +- No session affinity +- Statistically even distribution over many requests + +### Best For + +- Simple deployments +- When you want to avoid any sequential patterns +- Testing and development + +______________________________________________________________________ + +## Power of Two Choices + +The `power_of_two` policy randomly selects two workers and routes to the one with lower load. This provides good load distribution with minimal coordination overhead. + +### Configuration + +```bash +lmdeploy-router --policy power_of_two --worker-urls http://worker1:8000,http://worker2:8000,http://worker3:8000 +``` + +### Behavior + +1. Randomly pick two healthy workers +2. Query their current load (pending requests) +3. Route to the worker with lower load + +### Best For + +- Load-sensitive workloads +- When request processing times vary significantly +- Avoiding hot spots without full load tracking overhead + +### Note + +Requires at least 2 workers. With only 1 worker, behaves like direct routing. + +______________________________________________________________________ + +## Cache Aware + +The `cache_aware` policy optimizes for prefix caching by maintaining an approximate radix tree of request prefixes per worker. + +### Configuration + +```bash +lmdeploy-router --policy cache_aware \ + --cache-threshold 0.5 \ + --balance-abs-threshold 32 \ + --balance-rel-threshold 1.1 \ + --worker-urls http://worker1:8000,http://worker2:8000 +``` + +### Parameters + +| Parameter | Default | Description | +| ------------------------ | ------- | ----------------------------------------------------- | +| `cache_threshold` | 0.5 | Minimum prefix match ratio to use cache-based routing | +| `balance_abs_threshold` | 32 | Absolute load difference threshold for load balancing | +| `balance_rel_threshold` | 1.1 | Relative load ratio threshold for load balancing | +| `eviction_interval_secs` | 30 | Interval for cache eviction | +| `max_tree_size` | 10000 | Maximum nodes per radix tree | + +### Behavior + +1. **Balanced mode** (when load is even): + + - Find worker with highest prefix match for the request + - If match rate > `cache_threshold`: route to that worker (cache hit) + - Otherwise: route to worker with smallest tree (most cache capacity) + +2. **Imbalanced mode** (when load is skewed): + + - Route to worker with lowest load (shortest queue) + - Still updates the tree to maintain cache state + +### Best For + +- Workloads with repeated prompt prefixes (system prompts, few-shot examples) +- When prefix caching is enabled on backend workers +- Multi-tenant deployments with distinct prompt patterns + +______________________________________________________________________ + +## Choosing a Policy + +``` + ┌─────────────────────┐ + │ Need session │ + │ affinity? │ + └─────────┬───────────┘ + │ + ┌───────────────┴───────────────┐ + │ │ + Yes No + │ │ + ▼ ▼ + ┌─────────────────┐ ┌─────────────────┐ + │ Multi-turn or │ │ Load-sensitive │ + │ KV cache reuse?│ │ workload? │ + └────────┬────────┘ └────────┬────────┘ + │ │ + ┌───────────┴───────────┐ ┌───────────┴───────────┐ + │ │ │ │ + Yes No Yes No + │ │ │ │ + ▼ ▼ ▼ ▼ + ┌────────────────┐ ┌────────────────┐ ┌────────────────┐ ┌────────────────┐ + │ consistent_hash│ │ cache_aware │ │ power_of_two │ │ round_robin │ + │ │ │ │ │ │ │ or random │ + └────────────────┘ └────────────────┘ └────────────────┘ └────────────────┘ +``` + +### Quick Reference + +| Scenario | Recommended Policy | +| ------------------------------------------- | ------------------ | +| Chat applications with conversation history | `consistent_hash` | +| Batch inference with no state | `round_robin` | +| Variable request complexity | `power_of_two` | +| Repeated system prompts / few-shot | `cache_aware` | +| Simple testing / development | `random` | + +______________________________________________________________________ + +## PD (Prefill-Decode) Mode + +In prefill-decode disaggregated mode, you can configure separate policies for prefill and decode workers: + +```bash +lmdeploy-router \ + --mode pd \ + --prefill-policy consistent_hash \ + --decode-policy round_robin \ + --prefill-workers http://prefill1:8000,http://prefill2:8000 \ + --decode-workers http://decode1:8000,http://decode2:8000 +``` + +This allows optimizing each stage independently: + +- **Prefill**: Use `consistent_hash` for cache locality +- **Decode**: Use `round_robin` or `power_of_two` for load distribution diff --git a/lmdeploy/router/examples/configs/round_robin_config.json b/lmdeploy/router/examples/configs/round_robin_config.json new file mode 100644 index 0000000000..2002a90d1b --- /dev/null +++ b/lmdeploy/router/examples/configs/round_robin_config.json @@ -0,0 +1,15 @@ +{ + "route": { + "type": "RoundRobinRoute", + "servers": [ + { + "host": "localhost", + "port": 8000 + }, + { + "host": "localhost", + "port": 8001 + } + ] + } +} diff --git a/lmdeploy/router/examples/configs/single_server_config.json b/lmdeploy/router/examples/configs/single_server_config.json new file mode 100644 index 0000000000..f8c5e84bb8 --- /dev/null +++ b/lmdeploy/router/examples/configs/single_server_config.json @@ -0,0 +1,7 @@ +{ + "route": { + "type": "SingleServerRoute", + "host": "localhost", + "port": 8000 + } +} diff --git a/lmdeploy/router/examples/simulate_consistent_hash.rs b/lmdeploy/router/examples/simulate_consistent_hash.rs new file mode 100644 index 0000000000..931ea04a22 --- /dev/null +++ b/lmdeploy/router/examples/simulate_consistent_hash.rs @@ -0,0 +1,376 @@ +//! Simulation: Evaluate consistent hash balance for PD disagg decode workers +//! +//! Replicates the router's exact consistent hash algorithm to measure +//! request distribution balance across decode workers with data parallelism. +//! +//! Usage: +//! cargo run --example simulate_consistent_hash +//! cargo run --example simulate_consistent_hash -- --num-sessions 500000 --dp-size 4 +//! cargo run --example simulate_consistent_hash -- --trials 100 +//! cargo run --example simulate_consistent_hash -- --mode rendezvous --trials 100 + +use std::collections::{BTreeMap, HashMap}; + +use clap::{Parser, ValueEnum}; +use uuid::Uuid; + +use lmdeploy_router_rs::policies::ConsistentHashPolicy; +use lmdeploy_router_rs::policies::VIRTUAL_NODES_PER_WORKER; + +#[derive(Clone, ValueEnum)] +enum Mode { + /// Ring-based consistent hashing (current router implementation) + Ring, + /// Rendezvous hashing (Highest Random Weight) + Rendezvous, +} + +#[derive(Parser)] +#[command(name = "simulate_consistent_hash")] +#[command(about = "Simulate consistent hash distribution for PD disagg decode workers")] +struct Args { + /// Number of random session UUIDs to simulate + #[arg(long, default_value_t = 256)] + num_sessions: usize, + + /// Number of physical decode workers + #[arg(long, default_value_t = 2)] + num_workers: usize, + + /// Intra-node data parallel size (GPUs per worker) + #[arg(long, default_value_t = 4)] + dp_size: usize, + + /// Virtual nodes per logical worker on the hash ring (only used in ring mode) + #[arg(long, default_value_t = VIRTUAL_NODES_PER_WORKER)] + virtual_nodes: u32, + + /// Header name used for hash key (e.g. x-session-id, x-correlation-id) + #[arg(long, default_value = "x-session-id")] + header_name: String, + + /// Number of independent trials to run + #[arg(long, default_value_t = 1)] + trials: usize, + + /// Hashing mode: ring (consistent hash ring) or rendezvous (HRW) + #[arg(long, value_enum, default_value_t = Mode::Ring)] + mode: Mode, +} + +/// Stats collected from a single trial +struct TrialStats { + /// Coefficient of variation for physical workers + worker_cv: f64, + /// Imbalance ratio (max/min) for physical workers + worker_imbalance: f64, + /// Coefficient of variation for DP ranks + rank_cv: f64, + /// Imbalance ratio (max/min) for DP ranks + rank_imbalance: f64, +} + +fn run_single_trial(args: &Args, verbose: bool) -> TrialStats { + let total_logical = args.num_workers * args.dp_size; + + // Build worker URLs + let mut worker_urls: Vec = Vec::new(); + for w in 0..args.num_workers { + for rank in 0..args.dp_size { + worker_urls.push(format!("http://decode-worker{}:8000@{}", w, rank)); + } + } + + // Build hash ring (only used in ring mode) + let mut ring: BTreeMap = BTreeMap::new(); + if matches!(args.mode, Mode::Ring) { + for url in &worker_urls { + for i in 0..args.virtual_nodes { + let virtual_key = format!("{}:{}", url, i); + let hash_value = ConsistentHashPolicy::fbi_hash(&virtual_key); + ring.insert(hash_value, url.clone()); + } + } + } + + if verbose && matches!(args.mode, Mode::Ring) { + analyze_ring_distribution(&ring, &worker_urls, args.num_workers, args.dp_size); + } + + // Generate session UUIDs using proper UUID v4 + let mut session_ids: Vec = Vec::with_capacity(args.num_sessions); + for _ in 0..args.num_sessions { + session_ids.push(Uuid::new_v4().to_string()); + } + + // Simulate routing + let mut per_url_count: HashMap = HashMap::new(); + for url in &worker_urls { + per_url_count.insert(url.clone(), 0); + } + + for session_id in &session_ids { + let hash_key = format!("header:{}:{}", args.header_name, session_id); + + let selected = match args.mode { + Mode::Ring => { + let hash_value = ConsistentHashPolicy::fbi_hash(&hash_key); + ring.range(hash_value..) + .next() + .or_else(|| ring.iter().next()) + .map(|(_, url)| url.clone()) + .unwrap() + } + Mode::Rendezvous => { + // Rendezvous hashing: score each worker, pick highest + worker_urls + .iter() + .max_by_key(|url| { + let score_key = format!("{}:{}", hash_key, url); + ConsistentHashPolicy::fbi_hash(&score_key) + }) + .unwrap() + .clone() + } + }; + + *per_url_count.get_mut(&selected).unwrap() += 1; + } + + // Aggregate per physical worker + let mut per_worker_count: Vec = vec![0; args.num_workers]; + for (w, count) in per_worker_count.iter_mut().enumerate() { + for rank in 0..args.dp_size { + let url = format!("http://decode-worker{}:8000@{}", w, rank); + *count += per_url_count[&url]; + } + } + + // Physical worker stats + let mean = args.num_sessions as f64 / args.num_workers as f64; + let variance: f64 = per_worker_count + .iter() + .map(|&c| (c as f64 - mean).powi(2)) + .sum::() + / args.num_workers as f64; + let std_dev = variance.sqrt(); + let worker_cv = std_dev / mean; + let min_count = *per_worker_count.iter().min().unwrap(); + let max_count = *per_worker_count.iter().max().unwrap(); + let worker_imbalance = max_count as f64 / mean; + + // Per-DP-rank stats + let rank_counts: Vec = worker_urls.iter().map(|u| per_url_count[u]).collect(); + let rank_min = *rank_counts.iter().min().unwrap(); + let rank_max = *rank_counts.iter().max().unwrap(); + let rank_mean = args.num_sessions as f64 / total_logical as f64; + let rank_variance: f64 = rank_counts + .iter() + .map(|&c| (c as f64 - rank_mean).powi(2)) + .sum::() + / total_logical as f64; + let rank_std_dev = rank_variance.sqrt(); + let rank_cv = rank_std_dev / rank_mean; + let rank_imbalance = rank_max as f64 / rank_mean; + + if verbose { + // Print detailed per-URL distribution + println!("\n=== Per-DP-Rank Distribution ==="); + println!("{:<45} {:>8} {:>8}", "Worker URL", "Count", "%"); + println!("{}", "-".repeat(65)); + for url in &worker_urls { + let count = per_url_count[url]; + let pct = 100.0 * count as f64 / args.num_sessions as f64; + println!("{:<45} {:>8} {:>7.2}%", url, count, pct); + } + + println!("\n=== Per-Physical-Worker Distribution ==="); + println!("{:<30} {:>8} {:>8}", "Worker", "Count", "%"); + println!("{}", "-".repeat(50)); + for (w, count) in per_worker_count.iter().enumerate() { + let pct = 100.0 * *count as f64 / args.num_sessions as f64; + println!("decode-worker{:<17} {:>8} {:>7.2}%", w, count, pct); + } + + println!("\n=== Balance Statistics (Physical Workers) ==="); + println!("Expected per worker: {:.1}", mean); + println!("Min: {}", min_count); + println!("Max: {}", max_count); + println!( + "Imbalance ratio: {:.4} (max/expected)", + worker_imbalance + ); + println!("Std deviation: {:.2}", std_dev); + println!("Coeff of variation: {:.4} (lower is better)", worker_cv); + + println!("\n=== Balance Statistics (Per DP Rank) ==="); + println!("Expected per rank: {:.1}", rank_mean); + println!("Min: {}", rank_min); + println!("Max: {}", rank_max); + println!("Imbalance ratio: {:.4} (max/expected)", rank_imbalance); + println!("Std deviation: {:.2}", rank_std_dev); + println!("Coeff of variation: {:.4} (lower is better)", rank_cv); + } + + TrialStats { + worker_cv, + worker_imbalance, + rank_cv, + rank_imbalance, + } +} + +fn main() { + let args = Args::parse(); + + let total_logical = args.num_workers * args.dp_size; + let mode_name = match args.mode { + Mode::Ring => "ring", + Mode::Rendezvous => "rendezvous", + }; + println!("=== Consistent Hash Balance Simulation ==="); + println!("Mode: {}", mode_name); + println!("Physical decode workers: {}", args.num_workers); + println!("DP size (GPUs per worker): {}", args.dp_size); + println!("Logical worker URLs: {}", total_logical); + if matches!(args.mode, Mode::Ring) { + println!("Virtual nodes per logical worker: {}", args.virtual_nodes); + println!( + "Total virtual nodes on ring: {}", + total_logical as u32 * args.virtual_nodes + ); + } + println!("Sessions to simulate: {}", args.num_sessions); + println!("Trials: {}", args.trials); + println!(); + + if args.trials == 1 { + // Single trial: verbose output + run_single_trial(&args, true); + } else { + // Multiple trials: collect stats and aggregate + let mut all_stats: Vec = Vec::with_capacity(args.trials); + + for _ in 0..args.trials { + let stats = run_single_trial(&args, false); + all_stats.push(stats); + } + + // Aggregate + let worker_cvs: Vec = all_stats.iter().map(|s| s.worker_cv).collect(); + let worker_imbs: Vec = all_stats.iter().map(|s| s.worker_imbalance).collect(); + let rank_cvs: Vec = all_stats.iter().map(|s| s.rank_cv).collect(); + let rank_imbs: Vec = all_stats.iter().map(|s| s.rank_imbalance).collect(); + + println!("=== Aggregated Stats over {} Trials ===", args.trials); + println!(); + print_aggregated("Physical Worker CV", &worker_cvs); + print_aggregated("Physical Worker Imbalance (max/expected)", &worker_imbs); + print_aggregated("DP Rank CV", &rank_cvs); + print_aggregated("DP Rank Imbalance (max/expected)", &rank_imbs); + } +} + +fn print_aggregated(label: &str, values: &[f64]) { + let n = values.len() as f64; + let mean = values.iter().sum::() / n; + let mut sorted = values.to_vec(); + sorted.sort_by(|a, b| a.partial_cmp(b).unwrap()); + let min = sorted[0]; + let max = sorted[sorted.len() - 1]; + let p50 = percentile(&sorted, 50.0); + let p95 = percentile(&sorted, 95.0); + let p99 = percentile(&sorted, 99.0); + let std_dev = (values.iter().map(|v| (v - mean).powi(2)).sum::() / n).sqrt(); + + println!(" {}", label); + println!( + " mean={:.4} std={:.4} min={:.4} p50={:.4} p95={:.4} p99={:.4} max={:.4}", + mean, std_dev, min, p50, p95, p99, max + ); +} + +fn percentile(sorted: &[f64], pct: f64) -> f64 { + let idx = (pct / 100.0 * (sorted.len() - 1) as f64).round() as usize; + sorted[idx.min(sorted.len() - 1)] +} + +/// Compute arc ownership per URL on the ring, returning both totals and per-vnode arcs +/// Analyze how virtual nodes are distributed on the hash ring +fn analyze_ring_distribution( + ring: &BTreeMap, + worker_urls: &[String], + num_workers: usize, + dp_size: usize, +) { + println!("=== Hash Ring Virtual Node Analysis ==="); + + let mut vnodes_per_url: HashMap = HashMap::new(); + for url in worker_urls { + vnodes_per_url.insert(url.clone(), 0); + } + for url in ring.values() { + *vnodes_per_url.get_mut(url).unwrap() += 1; + } + + let mut vnodes_per_worker: Vec = vec![0; num_workers]; + for (w, count) in vnodes_per_worker.iter_mut().enumerate() { + for rank in 0..dp_size { + let url = format!("http://decode-worker{}:8000@{}", w, rank); + *count += vnodes_per_url[&url]; + } + } + + let total_vnodes: usize = ring.len(); + println!( + "Total virtual nodes on ring: {} (expected: {})", + total_vnodes, + worker_urls.len() as u32 * VIRTUAL_NODES_PER_WORKER + ); + + for (w, &vnode_count) in vnodes_per_worker.iter().enumerate() { + let pct = 100.0 * vnode_count as f64 / total_vnodes as f64; + println!(" decode-worker{}: {} vnodes ({:.1}%)", w, vnode_count, pct); + } + + let hashes: Vec = ring.keys().copied().collect(); + let owners: Vec = ring.values().cloned().collect(); + + let mut arc_per_worker: Vec = vec![0.0; num_workers]; + let n = hashes.len(); + for i in 0..n { + let next = (i + 1) % n; + let arc = if next > i { + hashes[next] - hashes[i] + } else { + (u64::MAX - hashes[i]) + hashes[next] + 1 + }; + let arc_frac = arc as f64 / u64::MAX as f64; + let url = &owners[i]; + for (w, arc) in arc_per_worker.iter_mut().enumerate() { + let prefix = format!("http://decode-worker{}:8000@", w); + if url.starts_with(&prefix) { + *arc += arc_frac; + break; + } + } + } + + println!("\nHash ring arc ownership (fraction of keyspace):"); + for (w, &arc) in arc_per_worker.iter().enumerate() { + println!(" decode-worker{}: {:.4} ({:.2}%)", w, arc, arc * 100.0); + } + let arc_min = arc_per_worker.iter().cloned().fold(f64::INFINITY, f64::min); + let arc_max = arc_per_worker + .iter() + .cloned() + .fold(f64::NEG_INFINITY, f64::max); + println!( + " Arc imbalance ratio: {:.4} (max/min)", + if arc_min > 0.0 { + arc_max / arc_min + } else { + f64::INFINITY + } + ); +} diff --git a/lmdeploy/router/py_src/lmdeploy_router/__init__.py b/lmdeploy/router/py_src/lmdeploy_router/__init__.py new file mode 100644 index 0000000000..327b4cd53b --- /dev/null +++ b/lmdeploy/router/py_src/lmdeploy_router/__init__.py @@ -0,0 +1,10 @@ +# Copyright (c) OpenMMLab. All rights reserved. +from lmdeploy_router.router import ROUTER_AVAILABLE +from lmdeploy_router.version import __version__ + +__all__ = ['__version__'] + +if ROUTER_AVAILABLE: + from lmdeploy_router.router import Router as Router + + __all__.append('Router') diff --git a/lmdeploy/router/py_src/lmdeploy_router/launch_router.py b/lmdeploy/router/py_src/lmdeploy_router/launch_router.py new file mode 100644 index 0000000000..7c184721d3 --- /dev/null +++ b/lmdeploy/router/py_src/lmdeploy_router/launch_router.py @@ -0,0 +1,102 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import argparse +import logging +import os +import sys + +import setproctitle + +from lmdeploy_router.mini_lb import MiniLoadBalancer +from lmdeploy_router.router_args import RouterArgs + +logger = logging.getLogger('router') + +try: + from lmdeploy_router.router import Router +except ImportError: + Router = None + logger.warning( + 'Rust Router is not installed, only python MiniLB (debugging only) is available' + ) + + +def launch_router(args: argparse.Namespace) -> Router | None: + """Launch the LMDeploy router with the configuration from parsed arguments. + + Args: + args: Namespace object containing router configuration + Can be either raw argparse.Namespace or converted RouterArgs + + Returns: + Router instance if successful, None if failed + """ + setproctitle.setproctitle('lmdeploy::router') + + packaged_router_bin = os.path.join(sys.prefix, 'bin', 'lmdeploy-router-bin') + if os.path.isfile(packaged_router_bin): + os.environ.setdefault('LMDPLOY_PYTHON_ROUTER_BIN', packaged_router_bin) + try: + # Convert to RouterArgs if needed + if not isinstance(args, RouterArgs): + router_args = RouterArgs.from_cli_args(args) + else: + router_args = args + + if router_args.mini_lb: + mini_lb = MiniLoadBalancer(router_args) + mini_lb.start() + else: + if Router is None: + raise RuntimeError('Rust Router is not installed') + router_args._validate_router_args() + router = Router.from_args(router_args) + router.start() + + except Exception as e: + logger.error(f'Error starting router: {e}') + raise e + + +class CustomHelpFormatter( + argparse.RawDescriptionHelpFormatter, argparse.ArgumentDefaultsHelpFormatter +): + """Custom formatter that preserves both description formatting and shows + defaults.""" + + pass + + +def parse_router_args(args: list[str]) -> RouterArgs: + """Parse command line arguments and return RouterArgs instance.""" + parser = argparse.ArgumentParser( + description="""LMDeploy Router - High-performance request distribution across worker nodes + +Usage: +This launcher enables starting a router with individual worker instances. It is useful for +multi-node setups or when you want to start workers and router separately. + +Examples: + # Regular mode + lmdeploy-router --worker-urls http://worker1:8000 http://worker2:8000 + + # PD disaggregated mode with different policies for prefill and decode + lmdeploy-router --lmdeploy-pd-disaggregation \\ + --prefill http://prefill1:8000 --prefill http://prefill2:8000 \\ + --decode http://decode1:8001 --decode http://decode2:8001 \\ + --prefill-policy cache_aware --decode-policy power_of_two + + """, + formatter_class=CustomHelpFormatter, + ) + + RouterArgs.add_cli_args(parser, use_router_prefix=False) + return RouterArgs.from_cli_args(parser.parse_args(args), use_router_prefix=False) + + +def main() -> None: + router_args = parse_router_args(sys.argv[1:]) + launch_router(router_args) + + +if __name__ == '__main__': + main() diff --git a/lmdeploy/router/py_src/lmdeploy_router/mini_lb.py b/lmdeploy/router/py_src/lmdeploy_router/mini_lb.py new file mode 100644 index 0000000000..4f6ee46ae7 --- /dev/null +++ b/lmdeploy/router/py_src/lmdeploy_router/mini_lb.py @@ -0,0 +1,328 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""Minimal HTTP load balancer for prefill and decode servers for testing.""" + +import asyncio +import logging +import random +from http import HTTPStatus +from itertools import chain + +import aiohttp +import orjson +import uvicorn +from fastapi import FastAPI, HTTPException +from fastapi.responses import ORJSONResponse, Response, StreamingResponse + +from lmdeploy_router.router_args import RouterArgs + +logger = logging.getLogger(__name__) + +AIOHTTP_STREAM_READ_CHUNK_SIZE = ( + 1024 * 64 +) # 64KB, to prevent aiohttp's "Chunk too big" error + + +class MiniLoadBalancer: + def __init__( + self, + router_args: RouterArgs, + ): + self._validate_router_args(router_args) + + self.host = router_args.host + self.port = router_args.port + self.timeout = router_args.request_timeout_secs + self.prefill_urls = router_args.prefill_urls + self.decode_urls = router_args.decode_urls + + def _validate_router_args(self, router_args: RouterArgs): + logger.warning( + '\x1b[33mMiniLB is only for debugging purposes, it only supports random policy!\033[0m' + ) + + # NOTE: too many arguments unsupported, just validate some important ones + if router_args.policy != 'random': + logger.warning('[MiniLB] Overriding policy to random') + router_args.policy = 'random' + + if not router_args.lmdeploy_pd_disaggregation: + raise ValueError('MiniLB only supports PD disaggregation mode') + + if len(router_args.prefill_urls) == 0 or len(router_args.decode_urls) == 0: + raise ValueError( + 'MiniLB requires at least one prefill and one decode server' + ) + + def start(self): + global lb + lb = self + uvicorn.run(app, host=self.host, port=self.port) + + def select_pair(self): + assert len(self.prefill_urls) > 0, 'No prefill servers available' + assert len(self.decode_urls) > 0, 'No decode servers available' + pidx = random.randint(0, len(self.prefill_urls) - 1) + didx = random.randint(0, len(self.decode_urls) - 1) + return self.prefill_urls[pidx], self.decode_urls[didx] + + async def generate( + self, modified_request, prefill_server, decode_server, endpoint + ) -> ORJSONResponse: + assert endpoint[0] != '/', f"Endpoint should not start with '/': {endpoint}" + + async with aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout( + total=self.timeout + ) # Add timeout for request reliability + ) as session: + tasks = [ + session.post(f'{prefill_server}/{endpoint}', json=modified_request), + session.post(f'{decode_server}/{endpoint}', json=modified_request), + ] + + # Wait for both responses to complete. Prefill should end first. + prefill_response, decode_response = await asyncio.gather(*tasks) + + if 'return_logprob' in modified_request: + + prefill_json = await prefill_response.json() + ret_json = await decode_response.json() + + # merge `meta_info.input_token_logprobs` from prefill to decode + if 'meta_info' in ret_json: + if 'input_token_logprobs' in ret_json['meta_info']: + ret_json['meta_info']['input_token_logprobs'] = ( + prefill_json['meta_info']['input_token_logprobs'] + + ret_json['meta_info']['input_token_logprobs'] + ) + else: + ret_json = await decode_response.json() + + return ORJSONResponse( + content=ret_json, + status_code=decode_response.status, + ) + + async def generate_stream( + self, modified_request, prefill_server, decode_server, endpoint='generate' + ): + assert endpoint[0] != '/', f"Endpoint should not start with '/': {endpoint}" + + async def stream_results(): + async with aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout( + total=self.timeout + ) # Add timeout for request reliability + ) as session: + # Create the tasks for both prefill and decode requests + tasks = [ + session.post(f'{prefill_server}/{endpoint}', json=modified_request), + session.post(f'{decode_server}/{endpoint}', json=modified_request), + ] + # Wait for both responses to complete. Since this is streaming, they return immediately. + prefill_response, decode_response = await asyncio.gather(*tasks) + + if modified_request.get('return_logprob', False): + prefill_chunks = [] + async for chunk in prefill_response.content: + prefill_chunks.append(chunk) + + first_prefill_chunk = ( + prefill_chunks[0].decode('utf-8')[5:].strip('\n') + ) + first_prefill_chunk_json = orjson.loads(first_prefill_chunk) + + async for chunk in decode_response.content: + # Note: This is inefficient + # merge prefill input_token_logprobs, output_token_logprobs to decode + decoded_chunk = chunk.decode('utf-8') + if ( + decoded_chunk + and decoded_chunk.startswith('data:') + and '[DONE]' not in decoded_chunk + ): + ret_json = orjson.loads(decoded_chunk[5:].strip('\n')) + ret_json['meta_info']['input_token_logprobs'] = ( + first_prefill_chunk_json['meta_info'][ + 'input_token_logprobs' + ] + + ret_json['meta_info']['input_token_logprobs'] + ) + + yield b'data: ' + orjson.dumps(ret_json) + b'\n\n' + else: + yield chunk + else: + async for chunk in decode_response.content.iter_chunked( + AIOHTTP_STREAM_READ_CHUNK_SIZE + ): + yield chunk + + return StreamingResponse( + stream_results(), + media_type='text/event-stream', + ) + + +app = FastAPI() +lb: MiniLoadBalancer | None = None + + +@app.get('/health') +async def health_check(): + return Response(status_code=200) + + +@app.get('/health_generate') +async def health_generate(): + async with aiohttp.ClientSession() as session: + # Create the tasks + tasks = [] + for server in chain(lb.prefill_urls, lb.decode_urls): + tasks.append(session.get(f'{server}/health_generate')) + for i, response in enumerate(asyncio.as_completed(tasks)): + await response + return Response(status_code=200) + + +@app.post('/flush_cache') +async def flush_cache(): + async with aiohttp.ClientSession() as session: + # Create the tasks + tasks = [] + for server in chain(lb.prefill_urls, lb.decode_urls): + tasks.append(session.post(f'{server}/flush_cache')) + for i, response in enumerate(asyncio.as_completed(tasks)): + await response + return Response(status_code=200) + + +@app.get('/get_server_info') +async def get_server_info(): + prefill_infos = [] + decode_infos = [] + all_internal_states = [] + + async with aiohttp.ClientSession() as session: + for server in lb.prefill_urls: + server_info = await session.get(f'{server}/get_server_info') + prefill_infos.append(await server_info.json()) + for server in lb.decode_urls: + server_info = await session.get(f'{server}/get_server_info') + info_json = await server_info.json() + decode_infos.append(info_json) + # Extract internal_states from decode servers + if 'internal_states' in info_json: + all_internal_states.extend(info_json['internal_states']) + + # Return format expected by bench_one_batch_server.py + if all_internal_states: + return { + 'internal_states': all_internal_states, + 'prefill': prefill_infos, + 'decode': decode_infos, + } + else: + # Fallback with dummy data if no internal states found + return { + 'internal_states': [ + { + 'last_gen_throughput': 0.0, + 'avg_spec_accept_length': None, + } + ], + 'prefill': prefill_infos, + 'decode': decode_infos, + } + + +@app.get('/get_model_info') +async def get_model_info(): + if not lb or not lb.prefill_urls: + raise HTTPException( + status_code=HTTPStatus.SERVICE_UNAVAILABLE, + detail='There is no server registered', + ) + + target_server_url = lb.prefill_urls[0] + endpoint_url = f'{target_server_url}/get_model_info' + + async with aiohttp.ClientSession() as session: + try: + async with session.get(endpoint_url) as response: + if response.status != 200: + error_text = await response.text() + raise HTTPException( + status_code=HTTPStatus.BAD_GATEWAY, + detail=( + f'Failed to get model info from {target_server_url}' + f'Status: {response.status}, Response: {error_text}' + ), + ) + + model_info_json = await response.json() + return ORJSONResponse(content=model_info_json) + + except aiohttp.ClientError: + raise HTTPException( + status_code=HTTPStatus.SERVICE_UNAVAILABLE, + detail='Failed to get model info from backend', + ) + + +@app.post('/generate') +async def handle_generate_request(request_data: dict): + prefill_server, decode_server = lb.select_pair() + + if request_data.get('stream', False): + return await lb.generate_stream( + request_data, prefill_server, decode_server, 'generate' + ) + else: + return await lb.generate( + request_data, prefill_server, decode_server, 'generate' + ) + + +async def _forward_to_backend(request_data: dict, endpoint_name: str): + prefill_server, decode_server = lb.select_pair() + + if request_data.get('stream', False): + return await lb.generate_stream( + request_data, + prefill_server, + decode_server, + endpoint=endpoint_name, + ) + else: + return await lb.generate( + request_data, + prefill_server, + decode_server, + endpoint=endpoint_name, + ) + + +@app.post('/v1/chat/completions') +async def handle_chat_completion_request(request_data: dict): + return await _forward_to_backend(request_data, 'v1/chat/completions') + + +@app.post('/v1/completions') +async def handle_completion_request(request_data: dict): + return await _forward_to_backend(request_data, 'v1/completions') + +@app.get('/v1/models') +async def get_models(): + prefill_server = lb.prefill_urls[0] # Get the first prefill server + async with aiohttp.ClientSession() as session: + try: + response = await session.get(f'{prefill_server}/v1/models') + if response.status != 200: + raise HTTPException( + status_code=response.status, + detail=f'Prefill server error: Status {response.status}', + ) + return ORJSONResponse(content=await response.json()) + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) diff --git a/lmdeploy/router/py_src/lmdeploy_router/router.py b/lmdeploy/router/py_src/lmdeploy_router/router.py new file mode 100644 index 0000000000..6dccaa8401 --- /dev/null +++ b/lmdeploy/router/py_src/lmdeploy_router/router.py @@ -0,0 +1,155 @@ +# Copyright (c) OpenMMLab. All rights reserved. + +from lmdeploy_router.router_args import RouterArgs + +try: + from lmdeploy_router_rs import PolicyType + from lmdeploy_router_rs import Router as _Router +except ImportError as exc: # No-Rust development mode + PolicyType = None + _Router = None + _ROUTER_IMPORT_ERROR = exc + +ROUTER_AVAILABLE = _Router is not None + + +def policy_from_str(policy_str: str | None) -> PolicyType: + """Convert policy string to PolicyType enum.""" + if PolicyType is None: + raise _ROUTER_IMPORT_ERROR + if policy_str is None: + return None + policy_map = { + 'random': PolicyType.Random, + 'round_robin': PolicyType.RoundRobin, + 'cache_aware': PolicyType.CacheAware, + 'power_of_two': PolicyType.PowerOfTwo, + 'consistent_hash': PolicyType.ConsistentHash, + 'rendezvous_hash': PolicyType.RendezvousHash, + } + return policy_map[policy_str] + + +class Router: + """A high-performance router for distributing requests across worker nodes. + + Args: + worker_urls: List of URLs for worker nodes that will handle requests. Each URL should include + the protocol, host, and port (e.g., ['http://worker1:8000', 'http://worker2:8000']) + policy: Load balancing policy to use. Options: + - PolicyType.Random: Randomly select workers + - PolicyType.RoundRobin: Distribute requests in round-robin fashion + - PolicyType.CacheAware: Distribute requests based on cache state and load balance + - PolicyType.PowerOfTwo: Select best of two random workers based on load (PD mode only) + host: Host address to bind the router server. Default: '127.0.0.1' + port: Port number to bind the router server. Default: 3001 + worker_startup_timeout_secs: Timeout in seconds for worker startup. Default: 300 + worker_startup_check_interval: Interval in seconds between checks for worker initialization. Default: 10 + cache_threshold: Cache threshold (0.0-1.0) for cache-aware routing. Routes to cached worker + if the match rate exceeds threshold, otherwise routes to the worker with the smallest + tree. Default: 0.5 + balance_abs_threshold: Load balancing is triggered when (max_load - min_load) > abs_threshold + AND max_load > min_load * rel_threshold. Otherwise, use cache aware. Default: 32 + balance_rel_threshold: Load balancing is triggered when (max_load - min_load) > abs_threshold + AND max_load > min_load * rel_threshold. Otherwise, use cache aware. Default: 1.0001 + eviction_interval_secs: Interval in seconds between cache eviction operations in cache-aware + routing. Default: 60 + max_payload_size: Maximum payload size in bytes. Default: 256MB + max_tree_size: Maximum size of the approximation tree for cache-aware routing. Default: 2^24 + enable_igw: Enable IGW (Inference-Gateway) mode for multi-model support. When enabled, + the router can manage multiple models simultaneously with per-model load balancing + policies. Default: False + api_key: API key for authorization with workers. Default: None + log_dir: Directory to store log files. If None, logs are only output to console. Default: None + log_level: Logging level. Options: 'debug', 'info', 'warning', 'error', 'critical'. + service_discovery: Enable Kubernetes service discovery. When enabled, the router will + automatically discover worker pods based on the selector. Default: False + selector: Dictionary mapping of label keys to values for Kubernetes pod selection. + Example: {"app": "inference-worker"}. Default: {} + service_discovery_port: Port to use for service discovery. The router will generate + worker URLs using this port. Default: 80 + service_discovery_namespace: Kubernetes namespace to watch for pods. If not provided, + watches pods across all namespaces (requires cluster-wide permissions). Default: None + prefill_selector: Dictionary mapping of label keys to values for Kubernetes pod selection + for prefill servers (PD mode only). Default: {} + decode_selector: Dictionary mapping of label keys to values for Kubernetes pod selection + for decode servers (PD mode only). Default: {} + prometheus_port: Port to expose Prometheus metrics. Default: None + prometheus_host: Host address to bind the Prometheus metrics server. Default: None + prefill_urls: List of prefill server URLs (PD mode only) + decode_urls: List of URLs for decode servers (PD mode only) + prefill_policy: Specific load balancing policy for prefill nodes (PD mode only). + If not specified, uses the main policy. Default: None + decode_policy: Specific load balancing policy for decode nodes (PD mode only). + If not specified, uses the main policy. Default: None + request_id_headers: List of HTTP headers to check for request IDs. If not specified, + uses common defaults: ['x-request-id', 'x-correlation-id', 'x-trace-id', 'request-id']. + Example: ['x-my-request-id', 'x-custom-trace-id']. Default: None + request_timeout_secs: Request timeout in seconds. Default: 600 + max_concurrent_requests: Maximum number of concurrent requests allowed for rate limiting. Default: 256 + queue_size: Queue size for pending requests when max concurrent limit reached + (0 = no queue, return 429 immediately). Default: 100 + queue_timeout_secs: Maximum time (in seconds) a request can wait in queue before timing out. Default: 60 + rate_limit_tokens_per_second: Token bucket refill rate (tokens per second). If not set, + defaults to max_concurrent_requests. Default: None + cors_allowed_origins: List of allowed origins for CORS. Empty list allows all origins. Default: [] + health_failure_threshold: Number of consecutive health check failures before marking + worker unhealthy. Default: 3 + health_success_threshold: Number of consecutive health check successes before marking worker healthy. Default: 2 + health_check_timeout_secs: Timeout in seconds for health check requests. Default: 5 + health_check_interval_secs: Interval in seconds between runtime health checks. Default: 60 + health_check_endpoint: Health check endpoint path. Default: '/health' + """ + + def __init__(self, router: _Router | None = None, **kwargs): + """Initialize Router either from a _Router instance or keyword + arguments. + + Args: + router: Optional _Router instance. If provided, kwargs are ignored. + **kwargs: Keyword arguments to pass to _Router constructor if router is None. + """ + if router is not None: + self._router = router + else: + # Create _Router from kwargs + self._router = _Router(**kwargs) + + @staticmethod + def from_args(args: RouterArgs) -> 'Router': + """Create a router from a RouterArgs instance.""" + + args_dict = vars(args) + # Convert RouterArgs to _Router parameters + args_dict['worker_urls'] = ( + [] + if args_dict['service_discovery'] or args_dict['lmdeploy_pd_disaggregation'] + else args_dict['worker_urls'] + ) + args_dict['policy'] = policy_from_str(args_dict['policy']) + args_dict['prefill_urls'] = ( + args_dict['prefill_urls'] + if args_dict['lmdeploy_pd_disaggregation'] + else None + ) + args_dict['decode_urls'] = ( + args_dict['decode_urls'] + if args_dict['lmdeploy_pd_disaggregation'] + else None + ) + args_dict['prefill_policy'] = policy_from_str(args_dict['prefill_policy']) + args_dict['decode_policy'] = policy_from_str(args_dict['decode_policy']) + + # remove mini_lb parameter + args_dict.pop('mini_lb') + # Rust exposes lmdeploy_with_gdr; the CLI flag disables it + args_dict['lmdeploy_with_gdr'] = not args_dict.pop('lmdeploy_disable_gdr') + + return Router(router=_Router(**args_dict)) + + def start(self) -> None: + """Start the router server. + + This method blocks until the server is shut down. + """ + self._router.start() diff --git a/lmdeploy/router/py_src/lmdeploy_router/router_args.py b/lmdeploy/router/py_src/lmdeploy_router/router_args.py new file mode 100644 index 0000000000..31f51be9ab --- /dev/null +++ b/lmdeploy/router/py_src/lmdeploy_router/router_args.py @@ -0,0 +1,590 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import argparse +import dataclasses +import logging + +logger = logging.getLogger(__name__) + + +@dataclasses.dataclass +class RouterArgs: + # Worker configuration + worker_urls: list[str] = dataclasses.field(default_factory=list) + host: str = '127.0.0.1' + port: int = 30000 + + # PD-specific configuration + mini_lb: bool = False + lmdeploy_pd_disaggregation: bool = False + lmdeploy_migration_protocol: str = 'rdma' + lmdeploy_rdma_link_type: str = 'roce' + lmdeploy_disable_gdr: bool = False + lmdeploy_dummy_prefill: bool = False + prefill_urls: list[str] = dataclasses.field(default_factory=list) + decode_urls: list[str] = dataclasses.field(default_factory=list) + + # Routing policy + policy: str = 'cache_aware' + prefill_policy: str | None = None # Specific policy for prefill nodes in PD mode + decode_policy: str | None = None # Specific policy for decode nodes in PD mode + worker_startup_timeout_secs: int = 600 + worker_startup_check_interval: int = 30 + cache_threshold: float = 0.3 + balance_abs_threshold: int = 64 + balance_rel_threshold: float = 1.5 + eviction_interval_secs: int = 120 + max_tree_size: int = 2**26 + max_payload_size: int = 512 * 1024 * 1024 # 512MB default for large batches + enable_igw: bool = False # Enable IGW (Inter-Gateway) mode for multi-model support + api_key: str | None = None + log_dir: str | None = None + log_level: str | None = None + # Service discovery configuration + service_discovery: bool = False + selector: dict[str, str] = dataclasses.field(default_factory=dict) + service_discovery_port: int = 80 + service_discovery_namespace: str | None = None + # PD service discovery configuration + prefill_selector: dict[str, str] = dataclasses.field(default_factory=dict) + decode_selector: dict[str, str] = dataclasses.field(default_factory=dict) + # KV connector for PD disaggregation (nixl pull-based or mooncake push-based) + kv_connector: str = 'nixl' + # Prometheus configuration + prometheus_port: int | None = None + prometheus_host: str | None = None + # Request ID headers configuration + request_id_headers: list[str] | None = None + # Request timeout in seconds + request_timeout_secs: int = 1800 + # Max concurrent requests for rate limiting + max_concurrent_requests: int = 32768 + # Queue size for pending requests when max concurrent limit reached + queue_size: int = 100 + # Maximum time (in seconds) a request can wait in queue before timing out + queue_timeout_secs: int = 60 + # Token bucket refill rate (tokens per second). If not set, defaults to max_concurrent_requests + rate_limit_tokens_per_second: int | None = None + # CORS allowed origins + cors_allowed_origins: list[str] = dataclasses.field(default_factory=list) + # Retry configuration + retry_max_retries: int = 5 + retry_initial_backoff_ms: int = 50 + retry_max_backoff_ms: int = 30_000 + retry_backoff_multiplier: float = 1.5 + retry_jitter_factor: float = 0.2 + disable_retries: bool = False + # Health check configuration + health_failure_threshold: int = 3 + health_success_threshold: int = 2 + health_check_timeout_secs: int = 5 + health_check_interval_secs: int = 60 + health_check_endpoint: str = '/health' + # Circuit breaker configuration + cb_failure_threshold: int = 10 + cb_success_threshold: int = 3 + cb_timeout_duration_secs: int = 60 + cb_window_duration_secs: int = 120 + disable_circuit_breaker: bool = False + + @staticmethod + def add_cli_args( + parser: argparse.ArgumentParser, + use_router_prefix: bool = False, + exclude_host_port: bool = False, + ): + """Add router-specific arguments to an argument parser. + + Args: + parser: The argument parser to add arguments to + use_router_prefix: If True, prefix all arguments with 'router-' to avoid conflicts + exclude_host_port: If True, don't add host and port arguments (used when inheriting from server) + """ + prefix = 'router-' if use_router_prefix else '' + + # Worker configuration + if not exclude_host_port: + parser.add_argument( + '--host', + type=str, + default=RouterArgs.host, + help='Host address to bind the router server', + ) + parser.add_argument( + '--port', + type=int, + default=RouterArgs.port, + help='Port number to bind the router server', + ) + + parser.add_argument( + '--worker-urls', + type=str, + nargs='*', + default=[], + help='List of worker URLs (e.g., http://worker1:8000 http://worker2:8000)', + ) + + # Routing policy configuration + parser.add_argument( + f'--{prefix}policy', + type=str, + default=RouterArgs.policy, + choices=[ + 'random', + 'round_robin', + 'cache_aware', + 'power_of_two', + 'consistent_hash', + 'rendezvous_hash', + ], + help='Load balancing policy to use. In PD mode, this is used for both prefill and decode unless overridden', + ) + parser.add_argument( + f'--{prefix}prefill-policy', + type=str, + default=None, + choices=[ + 'random', + 'round_robin', + 'cache_aware', + 'power_of_two', + 'consistent_hash', + 'rendezvous_hash', + ], + help='Specific policy for prefill nodes in PD mode. If not specified, uses the main policy', + ) + parser.add_argument( + f'--{prefix}decode-policy', + type=str, + default=None, + choices=[ + 'random', + 'round_robin', + 'cache_aware', + 'power_of_two', + 'consistent_hash', + 'rendezvous_hash', + ], + help='Specific policy for decode nodes in PD mode. If not specified, uses the main policy', + ) + + # PD-specific arguments + parser.add_argument( + f'--{prefix}mini-lb', + action='store_true', + help='Enable MiniLB', + ) + parser.add_argument( + f'--{prefix}lmdeploy-pd-disaggregation', + action='store_true', + help='Enable LMDeploy PD (Prefill-Decode) disaggregated mode', + ) + parser.add_argument( + f'--{prefix}lmdeploy-migration-protocol', + type=str, + default=RouterArgs.lmdeploy_migration_protocol, + choices=['rdma', 'nvlink'], + help='Migration protocol for LMDeploy PD disaggregation', + ) + parser.add_argument( + f'--{prefix}lmdeploy-rdma-link-type', + type=str, + default=RouterArgs.lmdeploy_rdma_link_type, + choices=['ib', 'roce'], + help='RDMA link type for LMDeploy PD disaggregation (used when migration protocol is rdma)', + ) + parser.add_argument( + f'--{prefix}lmdeploy-disable-gdr', + action='store_true', + help='Disable GPU Direct RDMA for LMDeploy PD disaggregation', + ) + parser.add_argument( + f'--{prefix}lmdeploy-dummy-prefill', + action='store_true', + help='Use dummy prefill for LMDeploy PD disaggregation', + ) + parser.add_argument( + f'--{prefix}prefill', + nargs=1, + action='append', + metavar=('URL',), + help='Prefill server URL. Can be specified multiple times. ' + 'Format: --prefill URL.', + ) + parser.add_argument( + f'--{prefix}decode', + nargs=1, + action='append', + metavar=('URL',), + help='Decode server URL. Can be specified multiple times.', + ) + parser.add_argument( + f'--{prefix}worker-startup-timeout-secs', + type=int, + default=RouterArgs.worker_startup_timeout_secs, + help='Timeout in seconds for worker startup', + ) + parser.add_argument( + f'--{prefix}worker-startup-check-interval', + type=int, + default=RouterArgs.worker_startup_check_interval, + help='Interval in seconds between checks for worker startup', + ) + parser.add_argument( + f'--{prefix}cache-threshold', + type=float, + default=RouterArgs.cache_threshold, + help='Cache threshold (0.0-1.0) for cache-aware routing', + ) + parser.add_argument( + f'--{prefix}balance-abs-threshold', + type=int, + default=RouterArgs.balance_abs_threshold, + help='Load balancing is triggered when (max_load - min_load) > abs_threshold ' + 'AND max_load > min_load * rel_threshold. Otherwise, use cache aware', + ) + parser.add_argument( + f'--{prefix}balance-rel-threshold', + type=float, + default=RouterArgs.balance_rel_threshold, + help='Load balancing is triggered when (max_load - min_load) > abs_threshold ' + 'AND max_load > min_load * rel_threshold. Otherwise, use cache aware', + ) + parser.add_argument( + f'--{prefix}eviction-interval-secs', + type=int, + default=RouterArgs.eviction_interval_secs, + help='Interval in seconds between cache eviction operations', + ) + parser.add_argument( + f'--{prefix}max-tree-size', + type=int, + default=RouterArgs.max_tree_size, + help='Maximum size of the approximation tree for cache-aware routing', + ) + parser.add_argument( + f'--{prefix}max-payload-size', + type=int, + default=RouterArgs.max_payload_size, + help='Maximum payload size in bytes', + ) + parser.add_argument( + f'--{prefix}enable-igw', + action='store_true', + help='Enable IGW (Inference-Gateway) mode for multi-model support', + ) + parser.add_argument( + f'--{prefix}api-key', + type=str, + default=None, + help='The API key used for authorization with workers.', + ) + parser.add_argument( + f'--{prefix}log-dir', + type=str, + default=None, + help='Directory to store log files. If not specified, logs are only output to console.', + ) + parser.add_argument( + f'--{prefix}log-level', + type=str, + default='info', + choices=['debug', 'info', 'warning', 'error', 'critical'], + help='Set the logging level. If not specified, defaults to INFO.', + ) + parser.add_argument( + f'--{prefix}service-discovery', + action='store_true', + help='Enable Kubernetes service discovery', + ) + parser.add_argument( + f'--{prefix}selector', + type=str, + nargs='+', + default={}, + help='Label selector for Kubernetes service discovery (format: key1=value1 key2=value2)', + ) + parser.add_argument( + f'--{prefix}service-discovery-port', + type=int, + default=RouterArgs.service_discovery_port, + help='Port to use for discovered worker pods', + ) + parser.add_argument( + f'--{prefix}service-discovery-namespace', + type=str, + help='Kubernetes namespace to watch for pods. If not provided, watches all namespaces ' + '(requires cluster-wide permissions)', + ) + parser.add_argument( + f'--{prefix}prefill-selector', + type=str, + nargs='+', + default={}, + help='Label selector for prefill server pods in PD mode (format: key1=value1 key2=value2)', + ) + parser.add_argument( + f'--{prefix}decode-selector', + type=str, + nargs='+', + default={}, + help='Label selector for decode server pods in PD mode (format: key1=value1 key2=value2)', + ) + parser.add_argument( + f'--{prefix}kv-connector', + type=str, + default=RouterArgs.kv_connector, + choices=['nixl', 'mooncake', 'moriio'], + help="KV connector type for PD disaggregation. 'nixl' (default) uses NIXL's " + "pull-based KV transfer; 'mooncake' uses Mooncake's push-based protocol " + "'moriio' uses the MoRI-IO connector (either READ or WRITE modes).", + ) + # Prometheus configuration + parser.add_argument( + f'--{prefix}prometheus-port', + type=int, + default=29000, + help='Port to expose Prometheus metrics. If not specified, Prometheus metrics are disabled', + ) + parser.add_argument( + f'--{prefix}prometheus-host', + type=str, + default='127.0.0.1', + help='Host address to bind the Prometheus metrics server', + ) + parser.add_argument( + f'--{prefix}request-id-headers', + type=str, + nargs='*', + help='Custom HTTP headers to check for request IDs (e.g., x-request-id x-trace-id). ' + 'If not specified, uses common defaults.', + ) + parser.add_argument( + f'--{prefix}request-timeout-secs', + type=int, + default=RouterArgs.request_timeout_secs, + help='Request timeout in seconds', + ) + # Retry configuration + parser.add_argument( + f'--{prefix}retry-max-retries', + type=int, + default=RouterArgs.retry_max_retries, + ) + parser.add_argument( + f'--{prefix}retry-initial-backoff-ms', + type=int, + default=RouterArgs.retry_initial_backoff_ms, + ) + parser.add_argument( + f'--{prefix}retry-max-backoff-ms', + type=int, + default=RouterArgs.retry_max_backoff_ms, + ) + parser.add_argument( + f'--{prefix}retry-backoff-multiplier', + type=float, + default=RouterArgs.retry_backoff_multiplier, + ) + parser.add_argument( + f'--{prefix}retry-jitter-factor', + type=float, + default=RouterArgs.retry_jitter_factor, + ) + parser.add_argument( + f'--{prefix}disable-retries', + action='store_true', + help='Disable retries (equivalent to setting retry_max_retries=1)', + ) + # Circuit breaker configuration + parser.add_argument( + f'--{prefix}cb-failure-threshold', + type=int, + default=RouterArgs.cb_failure_threshold, + ) + parser.add_argument( + f'--{prefix}cb-success-threshold', + type=int, + default=RouterArgs.cb_success_threshold, + ) + parser.add_argument( + f'--{prefix}cb-timeout-duration-secs', + type=int, + default=RouterArgs.cb_timeout_duration_secs, + ) + parser.add_argument( + f'--{prefix}cb-window-duration-secs', + type=int, + default=RouterArgs.cb_window_duration_secs, + ) + parser.add_argument( + f'--{prefix}disable-circuit-breaker', + action='store_true', + help='Disable circuit breaker (equivalent to setting cb_failure_threshold to u32::MAX)', + ) + # Health check configuration + parser.add_argument( + f'--{prefix}health-failure-threshold', + type=int, + default=RouterArgs.health_failure_threshold, + help='Number of consecutive health check failures before marking worker unhealthy', + ) + parser.add_argument( + f'--{prefix}health-success-threshold', + type=int, + default=RouterArgs.health_success_threshold, + help='Number of consecutive health check successes before marking worker healthy', + ) + parser.add_argument( + f'--{prefix}health-check-timeout-secs', + type=int, + default=RouterArgs.health_check_timeout_secs, + help='Timeout in seconds for health check requests', + ) + parser.add_argument( + f'--{prefix}health-check-interval-secs', + type=int, + default=RouterArgs.health_check_interval_secs, + help='Interval in seconds between runtime health checks', + ) + parser.add_argument( + f'--{prefix}health-check-endpoint', + type=str, + default=RouterArgs.health_check_endpoint, + help='Health check endpoint path', + ) + parser.add_argument( + f'--{prefix}max-concurrent-requests', + type=int, + default=RouterArgs.max_concurrent_requests, + help='Maximum number of concurrent requests allowed (for rate limiting)', + ) + parser.add_argument( + f'--{prefix}queue-size', + type=int, + default=RouterArgs.queue_size, + help='Queue size for pending requests when max concurrent limit reached ' + '(0 = no queue, return 429 immediately)', + ) + parser.add_argument( + f'--{prefix}queue-timeout-secs', + type=int, + default=RouterArgs.queue_timeout_secs, + help='Maximum time (in seconds) a request can wait in queue before timing out', + ) + parser.add_argument( + f'--{prefix}rate-limit-tokens-per-second', + type=int, + default=RouterArgs.rate_limit_tokens_per_second, + help='Token bucket refill rate (tokens per second). If not set, defaults to max_concurrent_requests', + ) + parser.add_argument( + f'--{prefix}cors-allowed-origins', + type=str, + nargs='*', + default=[], + help='CORS allowed origins (e.g., http://localhost:3000 https://example.com)', + ) + + @classmethod + def from_cli_args( + cls, args: argparse.Namespace, use_router_prefix: bool = False + ) -> 'RouterArgs': + """Create RouterArgs instance from parsed command line arguments. + + Args: + args: Parsed command line arguments + use_router_prefix: If True, look for arguments with 'router-' prefix + """ + prefix = 'router_' if use_router_prefix else '' + cli_args_dict = vars(args) + args_dict = {} + + for attr in dataclasses.fields(cls): + # Auto strip prefix from args + if f'{prefix}{attr.name}' in cli_args_dict: + args_dict[attr.name] = cli_args_dict[f'{prefix}{attr.name}'] + elif attr.name in cli_args_dict: + args_dict[attr.name] = cli_args_dict[attr.name] + + # parse special arguments and remove "--prefill" and "--decode" from cli_args_dict + args_dict['prefill_urls'] = cls._parse_prefill_urls( + cli_args_dict.get(f'{prefix}prefill', None) + ) + args_dict['decode_urls'] = cls._parse_decode_urls( + cli_args_dict.get(f'{prefix}decode', None) + ) + args_dict['selector'] = cls._parse_selector( + cli_args_dict.get(f'{prefix}selector', None) + ) + args_dict['prefill_selector'] = cls._parse_selector( + cli_args_dict.get(f'{prefix}prefill_selector', None) + ) + args_dict['decode_selector'] = cls._parse_selector( + cli_args_dict.get(f'{prefix}decode_selector', None) + ) + + return cls(**args_dict) + + def _validate_router_args(self): + # Validate configuration based on mode + if self.lmdeploy_pd_disaggregation: + if self.prefill_policy and self.decode_policy and self.policy: + logger.warning( + 'Both --prefill-policy and --decode-policy are specified. ' + 'The main --policy flag will be ignored for PD mode.' + ) + elif self.prefill_policy and not self.decode_policy and self.policy: + logger.info( + f"Using --prefill-policy '{self.prefill_policy}' for prefill nodes " + f"and --policy '{self.policy}' for decode nodes." + ) + elif self.decode_policy and not self.prefill_policy and self.policy: + logger.info( + f"Using --policy '{self.policy}' for prefill nodes " + f"and --decode-policy '{self.decode_policy}' for decode nodes." + ) + + @staticmethod + def _parse_selector(selector_list): + if not selector_list: + return {} + + selector = {} + for item in selector_list: + if '=' in item: + key, value = item.split('=', 1) + selector[key] = value + return selector + + @staticmethod + def _parse_prefill_urls(prefill_list): + """Parse prefill URLs from --prefill arguments. + + Format: --prefill URL + Example: + --prefill http://prefill1:8080 + --prefill http://prefill2:8080 + """ + if not prefill_list: + return [] + + prefill_urls = [] + for prefill_args in prefill_list: + url = prefill_args[0] + prefill_urls.append(url) + + return prefill_urls + + @staticmethod + def _parse_decode_urls(decode_list): + """Parse decode URLs from --decode arguments. + + Format: --decode URL + Example: --decode http://decode1:8081 --decode http://decode2:8081 + """ + if not decode_list: + return [] + + # decode_list is a list of single-element lists due to nargs=1 + return [url[0] for url in decode_list] diff --git a/lmdeploy/serve/proxy/__init__.py b/lmdeploy/router/py_src/lmdeploy_router/version.py similarity index 68% rename from lmdeploy/serve/proxy/__init__.py rename to lmdeploy/router/py_src/lmdeploy_router/version.py index ef101fec61..9b25ac57d9 100644 --- a/lmdeploy/serve/proxy/__init__.py +++ b/lmdeploy/router/py_src/lmdeploy_router/version.py @@ -1 +1,2 @@ # Copyright (c) OpenMMLab. All rights reserved. +__version__ = '0.0.3' diff --git a/lmdeploy/router/py_test/__init__.py b/lmdeploy/router/py_test/__init__.py new file mode 100644 index 0000000000..e68c257b60 --- /dev/null +++ b/lmdeploy/router/py_test/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""Test package root for router Python tests.""" diff --git a/lmdeploy/router/py_test/conftest.py b/lmdeploy/router/py_test/conftest.py new file mode 100644 index 0000000000..8d04654918 --- /dev/null +++ b/lmdeploy/router/py_test/conftest.py @@ -0,0 +1,9 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import sys +from pathlib import Path + +# Ensure local sources in py_src are importable ahead of any installed package +_ROOT = Path(__file__).resolve().parents[1] +_SRC = _ROOT / 'py_src' +if str(_SRC) not in sys.path: + sys.path.insert(0, str(_SRC)) diff --git a/lmdeploy/router/py_test/e2e/conftest.py b/lmdeploy/router/py_test/e2e/conftest.py new file mode 100644 index 0000000000..501795dc84 --- /dev/null +++ b/lmdeploy/router/py_test/e2e/conftest.py @@ -0,0 +1,217 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import os +import socket +import subprocess +import time +import urllib.error +import urllib.request +from pathlib import Path +from types import SimpleNamespace + +import pytest +import requests + +DEFAULT_MODEL = 'Qwen/Qwen2.5-0.5B-Instruct' +DEFAULT_STARTUP_TIMEOUT = 300.0 + + +def _find_available_port() -> int: + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock: + sock.bind(('127.0.0.1', 0)) + return sock.getsockname()[1] + + +def _wait_for_health(url: str, process: subprocess.Popen, timeout: float) -> None: + deadline = time.monotonic() + timeout + last_error = 'no response' + while time.monotonic() < deadline: + if process.poll() is not None: + raise RuntimeError(f'process exited with status {process.returncode}') + try: + with urllib.request.urlopen(f'{url}/health', timeout=2) as response: + if response.status == 200: + return + last_error = f'HTTP {response.status}' + except (OSError, urllib.error.URLError) as exc: + last_error = str(exc) + time.sleep(1) + raise RuntimeError(f'timed out waiting for {url}: {last_error}') + + +def _start(command: list[str], log_path: Path) -> subprocess.Popen: + log_path.parent.mkdir(parents=True, exist_ok=True) + with log_path.open('w', encoding='utf-8') as log_file: + log_file.write('COMMAND: ' + ' '.join(command) + '\n\n') + return subprocess.Popen( + command, + stdin=subprocess.DEVNULL, + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + + +def _stop(process: subprocess.Popen | None) -> None: + if process is None or process.poll() is not None: + return + process.terminate() + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + process.kill() + process.wait(timeout=10) + + +def _api_server_command( + model: str, + port: int, + *, + model_name: str | None = None, + role: str = 'Hybrid', + proxy_url: str | None = None, + tp: int = 1, +) -> list[str]: + command = [ + 'lmdeploy', + 'serve', + 'api_server', + model, + '--server-name', + '127.0.0.1', + '--server-port', + str(port), + '--role', + role, + '--tp', + str(tp), + ] + if model_name: + command.extend(['--model-name', model_name]) + if proxy_url: + command.extend(['--proxy-url', proxy_url]) + return command + + +def _discover_model(url: str) -> str: + response = requests.get(f'{url}/v1/models', timeout=10) + response.raise_for_status() + payload = response.json() + for item in payload.get('data', []): + if isinstance(item, dict) and isinstance(item.get('id'), str): + return item['id'] + if isinstance(payload.get('id'), str): + return payload['id'] + raise RuntimeError(f'{url}/v1/models returned no model ID') + + +def pytest_configure(config): + config.addinivalue_line('markers', 'e2e: requires a real LMDeploy server') + + +@pytest.fixture(scope='session') +def e2e_model() -> str: + return os.getenv('LMDEPLOY_E2E_MODEL', DEFAULT_MODEL) + + +@pytest.fixture(scope='session') +def e2e_model_name() -> str: + return os.getenv('LMDEPLOY_E2E_MODEL_NAME', Path(e2e_model.__wrapped__()).name) + + +@pytest.fixture(scope='session') +def e2e_embedding_model() -> str | None: + return os.getenv('LMDEPLOY_E2E_EMBEDDING_MODEL') or None + + +@pytest.fixture(scope='session') +def e2e_startup_timeout() -> float: + return float(os.getenv('LMDEPLOY_E2E_STARTUP_TIMEOUT', str(DEFAULT_STARTUP_TIMEOUT))) + + +@pytest.fixture(scope='session') +def e2e_artifact_dir() -> Path: + directory = Path(os.getenv('LMDEPLOY_E2E_ARTIFACT_DIR', '/tmp/lmdeploy-router-e2e')) + directory.mkdir(parents=True, exist_ok=True) + return directory + + +@pytest.fixture +def e2e_router_only_rr(e2e_artifact_dir): + port = _find_available_port() + url = f'http://127.0.0.1:{port}' + process = _start( + ['lmdeploy-router', '--host', '127.0.0.1', '--port', str(port), '--policy', 'round_robin'], + e2e_artifact_dir / f'router-{port}.log', + ) + try: + _wait_for_health(url, process, 60) + yield SimpleNamespace(process=process, url=url) + finally: + _stop(process) + + +@pytest.fixture(scope='session') +def e2e_primary_worker(e2e_model, e2e_model_name, e2e_startup_timeout, e2e_artifact_dir): + if not os.getenv('LMDEPLOY_E2E_RUN_REGULAR'): + pytest.skip('set LMDEPLOY_E2E_RUN_REGULAR=1 to run the managed LMDeploy worker e2e') + port = _find_available_port() + url = f'http://127.0.0.1:{port}' + command = _api_server_command(e2e_model, port, model_name=e2e_model_name) + process = _start(command, e2e_artifact_dir / 'regular-worker.log') + try: + _wait_for_health(url, process, e2e_startup_timeout) + yield SimpleNamespace(process=process, url=url) + finally: + _stop(process) + + +@pytest.fixture(scope='session') +def e2e_embedding_worker(e2e_embedding_model, e2e_startup_timeout, e2e_artifact_dir): + if not e2e_embedding_model: + pytest.skip('set LMDEPLOY_E2E_EMBEDDING_MODEL to run embeddings e2e') + port = _find_available_port() + url = f'http://127.0.0.1:{port}' + command = _api_server_command(e2e_embedding_model, port) + process = _start(command, e2e_artifact_dir / 'embedding-worker.log') + try: + _wait_for_health(url, process, e2e_startup_timeout) + yield SimpleNamespace(process=process, url=url) + finally: + _stop(process) + + +def _split_urls(value: str) -> tuple[str, ...]: + return tuple(url.strip().rstrip('/') for url in value.split(',') if url.strip()) + + +@pytest.fixture(scope='module') +def e2e_pd_cluster(e2e_model, e2e_artifact_dir): + prefill_urls = _split_urls(os.getenv('LMDEPLOY_E2E_PREFILL_URLS', '')) + decode_urls = _split_urls(os.getenv('LMDEPLOY_E2E_DECODE_URLS', '')) + if not prefill_urls or not decode_urls: + pytest.skip( + 'set LMDEPLOY_E2E_PREFILL_URLS and LMDEPLOY_E2E_DECODE_URLS to run LMDeploy PD e2e' + ) + + port = _find_available_port() + url = f'http://127.0.0.1:{port}' + command = [ + 'lmdeploy-router', '--host', '127.0.0.1', '--port', str(port), '--policy', 'round_robin', + '--lmdeploy-pd-disaggregation', + ] + for prefill_url in prefill_urls: + command.extend(['--prefill', prefill_url]) + for decode_url in decode_urls: + command.extend(['--decode', decode_url]) + process = _start(command, e2e_artifact_dir / f'pd-router-{port}.log') + try: + _wait_for_health(url, process, 60) + try: + model_id = _discover_model(prefill_urls[0]) + except RuntimeError: + model_id = e2e_model + yield SimpleNamespace( + process=process, url=url, model=model_id, prefill_urls=prefill_urls, decode_urls=decode_urls + ) + finally: + _stop(process) diff --git a/lmdeploy/router/py_test/e2e/pd_disagg_lmdeploy/README.md b/lmdeploy/router/py_test/e2e/pd_disagg_lmdeploy/README.md new file mode 100644 index 0000000000..c732cc5e6e --- /dev/null +++ b/lmdeploy/router/py_test/e2e/pd_disagg_lmdeploy/README.md @@ -0,0 +1,58 @@ +# LMDeploy P/D Disaggregation Tests + +These tests validate that LMDeploy Router routes OpenAI-compatible traffic +between externally managed LMDeploy Prefill and Decode API servers. + +## Architecture + +The tests do not manage RDMA, DLSlime, or Mooncake resources. Those must be +provisioned before the run: + +- LMDeploy Prefill API servers started with `--role Prefill`. +- LMDeploy Decode API servers started with `--role Decode`. +- Both roles using the same `--migration-backend`. +- The model name reported by all servers matching `MODEL_NAME`. + +The router is started with `--lmdeploy-pd-disaggregation`, `--prefill`, and +`--decode`. + +## Files + +- `run_accuracy_test.sh`: starts the router and runs completion/streaming validation. +- `test_pd_accuracy.py`: checks health, non-streaming completions, and streaming completions. +- `test_lm_eval_accuracy.py`: optional LM-Eval evaluation through the router. +- `tp_config_sweep_test.sh`: runs the accuracy script for externally provisioned TP configurations. + +## Run + +```bash +export LMDEPLOY_ROUTER_BIN="$PWD/target/release/lmdeploy-router" +export PREFILL_URLS=http://prefill-host:9001 +export DECODE_URLS=http://decode-host:9002 +export MODEL_PATH=Qwen/Qwen2.5-0.5B-Instruct +export MODEL_NAME=Qwen2.5-0.5B-Instruct + +py_test/e2e/pd_disagg_lmdeploy/run_accuracy_test.sh +``` + +For a tensor-parallel sweep, provision TP=1 and TP=2 servers and provide their +URLs explicitly: + +```bash +TP_SIZES="1 2" \ +PREFILL_TP1_URLS=http://prefill-tp1:9001 \ +DECODE_TP1_URLS=http://decode-tp1:9002 \ +PREFILL_TP2_URLS=http://prefill-tp2:9003 \ +DECODE_TP2_URLS=http://decode-tp2:9004 \ +py_test/e2e/pd_disagg_lmdeploy/tp_config_sweep_test.sh +``` + +Optional LM-Eval: + +```bash +./run_accuracy_test.sh + +python3 test_lm_eval_accuracy.py \ + --router-url http://127.0.0.1:8300 \ + --model "$MODEL_NAME" +``` diff --git a/lmdeploy/router/py_test/e2e/pd_disagg_lmdeploy/run_accuracy_test.sh b/lmdeploy/router/py_test/e2e/pd_disagg_lmdeploy/run_accuracy_test.sh new file mode 100755 index 0000000000..2db3145d62 --- /dev/null +++ b/lmdeploy/router/py_test/e2e/pd_disagg_lmdeploy/run_accuracy_test.sh @@ -0,0 +1,79 @@ +#!/usr/bin/env bash +# Run LMDeploy P/D disaggregation router accuracy tests against externally +# managed LMDeploy Prefill and Decode API servers. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)" + +MODEL_PATH=${MODEL_PATH:-"${LMDEPLOY_E2E_MODEL:-Qwen/Qwen2.5-0.5B-Instruct}"} +MODEL_NAME=${MODEL_NAME:-LMDEPLOY_MODEL_NAME:-${MODEL_PATH##*/}} +NUM_REQUESTS=${NUM_REQUESTS:-20} +NUM_STREAMING_REQUESTS=${NUM_STREAMING_REQUESTS:-5} +MAX_TOKENS=${MAX_TOKENS:-30} +ROUTER_PORT=${ROUTER_PORT:-8300} +PREFILL_URLS=${PREFILL_URLS:?set PREFILL_URLS to comma-separated LMDeploy Prefill URLs} +DECODE_URLS=${DECODE_URLS:?set DECODE_URLS to comma-separated LMDeploy Decode URLs} +MIGRATION_PROTOCOL=${MIGRATION_PROTOCOL:-rdma} +RDMA_LINK_TYPE=${RDMA_LINK_TYPE:-roce} + +ROUTER_BIN=${LMDEPLOY_ROUTER_BIN:-${REPO_ROOT}/target/debug/lmdeploy-router} +if [[ ! -x "${ROUTER_BIN}" ]]; then + echo "Router binary not found: ${ROUTER_BIN}" >&2 + exit 1 +fi + +PREFILL_ARGS=() +for url in ${PREFILL_URLS//,/ }; do + PREFILL_ARGS+=("--prefill" "${url}") +done +DECODE_ARGS=() +for url in ${DECODE_URLS//,/ }; do + DECODE_ARGS+=("--decode" "${url}") +done + +ROUTER_LOG=${ROUTER_LOG:-/tmp/lmdeploy-pd-router.log} +"${ROUTER_BIN}" \ + --host 127.0.0.1 \ + --port "${ROUTER_PORT}" \ + --policy round_robin \ + --lmdeploy-pd-disaggregation \ + --lmdeploy-migration-protocol "${MIGRATION_PROTOCOL}" \ + --lmdeploy-rdma-link-type "${RDMA_LINK_TYPE}" \ + --worker-startup-check-interval 1 \ + "${PREFILL_ARGS[@]}" \ + "${DECODE_ARGS[@]}" \ + >"${ROUTER_LOG}" 2>&1 & +ROUTER_PID=$! + +cleanup() { + kill "${ROUTER_PID}" 2>/dev/null || true + wait "${ROUTER_PID}" 2>/dev/null || true +} +trap cleanup EXIT INT TERM + +for _ in $(seq 1 60); do + if curl --fail --silent "http://127.0.0.1:${ROUTER_PORT}/health" >/dev/null; then + break + fi + if ! kill -0 "${ROUTER_PID}" 2>/dev/null; then + echo "Router exited during startup. Log tail:" >&2 + tail -100 "${ROUTER_LOG}" >&2 + exit 1 + fi + sleep 1 +done + +if ! curl --fail --silent "http://127.0.0.1:${ROUTER_PORT}/health" >/dev/null; then + echo "Router failed to become healthy. Log tail:" >&2 + tail -100 "${ROUTER_LOG}" >&2 + exit 1 +fi + +python3 "${SCRIPT_DIR}/test_pd_accuracy.py" \ + --router-url "http://127.0.0.1:${ROUTER_PORT}" \ + --model "${MODEL_NAME}" \ + --num-requests "${NUM_REQUESTS}" \ + --num-streaming-requests "${NUM_STREAMING_REQUESTS}" \ + --max-tokens "${MAX_TOKENS}" diff --git a/lmdeploy/router/py_test/e2e/pd_disagg_lmdeploy/test_lm_eval_accuracy.py b/lmdeploy/router/py_test/e2e/pd_disagg_lmdeploy/test_lm_eval_accuracy.py new file mode 100644 index 0000000000..8bcc46d54f --- /dev/null +++ b/lmdeploy/router/py_test/e2e/pd_disagg_lmdeploy/test_lm_eval_accuracy.py @@ -0,0 +1,295 @@ +# Copyright (c) OpenMMLab. All rights reserved. +#!/usr/bin/env python3 +"""P/D Disaggregation LM-Eval Accuracy Test for LMDeploy Router. + +This script measures LM-Eval accuracy through the LMDeploy Prefill/Decode router. It is optional and requires the router +to be already running. + +Uses the LM Evaluation Harness (lm-eval) to measure accuracy on the gsm8k task. +""" + +import argparse +import os +import sys + +import lm_eval +import openai + +# Test configuration +TASK = 'gsm8k' +FILTER = 'exact_match,strict-match' +RTOL = 0.03 # Relative tolerance for accuracy comparison + +# Optional model-specific expected values. Set MODEL_NAME to the model ID +# exposed by the LMDeploy servers; add a local baseline when one is known. +EXPECTED_VALUES = { + 'meta-llama/Llama-3.2-1B-Instruct': 0.33, # Lowered to accept >30% accuracy + 'Qwen/Qwen3-0.6B': 0.41, + 'deepseek-ai/deepseek-vl2-small': 0.59, + 'deepseek-ai/deepseek-vl2-tiny': 0.19, + 'deepseek-ai/DeepSeek-V2-Lite-Chat': 0.65, +} + +# Simple prompt for connectivity test +SIMPLE_PROMPT = ( + 'LMDeploy serves models with an active open-source community, which means' +) + + +class Colors: + """Terminal colors for output.""" + + GREEN = '\033[92m' + RED = '\033[91m' + YELLOW = '\033[93m' + BLUE = '\033[94m' + RESET = '\033[0m' + + +def print_success(msg: str): + print(f'{Colors.GREEN}✓ {msg}{Colors.RESET}') + + +def print_error(msg: str): + print(f'{Colors.RED}✗ {msg}{Colors.RESET}') + + +def print_info(msg: str): + print(f'{Colors.BLUE}ℹ {msg}{Colors.RESET}') + + +def print_warning(msg: str): + print(f'{Colors.YELLOW}⚠ {msg}{Colors.RESET}') + + +def run_simple_prompt(base_url: str, model_name: str) -> bool: + """Run a simple prompt to verify connectivity before running full + evaluation. + + Args: + base_url: Base URL for the router API + model_name: Model name to test + + Returns: + True if successful, False otherwise + """ + print_info('Running connectivity test with simple prompt...') + + try: + # Use a very long timeout (10 minutes) for connectivity test + client = openai.OpenAI(api_key='EMPTY', base_url=base_url, timeout=600.0) + # Use chat completions for Instruct models + completion = client.chat.completions.create( + model=model_name, + messages=[{'role': 'user', 'content': SIMPLE_PROMPT}], + max_tokens=50, + ) + + output = completion.choices[0].message.content if completion.choices else '' + + print('-' * 60) + print_info(f'Connectivity Test Results for {model_name}:') + print(f'Prompt: {SIMPLE_PROMPT}') + print(f'Output: {output}') + print('-' * 60) + + if not output or len(output.strip()) == 0: + print_error('Connectivity test failed: Empty output') + return False + + print_success('Connectivity test passed') + return True + + except Exception as e: + print_error(f'Connectivity test failed: {e}') + return False + + +def run_accuracy_evaluation( + base_url: str, + model_name: str, + num_concurrent: int = 20, +) -> dict: + """Run LM Evaluation Harness on gsm8k task. + + Args: + base_url: Base URL for the router API (should be http://host:port/v1) + model_name: Model name to evaluate + num_concurrent: Number of concurrent requests + + Returns: + Dictionary containing evaluation results + """ + print_info(f'Running LM-Eval accuracy test on {TASK} task...') + print_info('This may take several minutes...') + + # Use Chat Completions API for Instruct models (native format) + # This fixes 422 errors from endpoint mismatch + try: + results = lm_eval.simple_evaluate( + model='local-chat-completions', + model_args={ + 'model': model_name, + 'base_url': f'{base_url}/chat/completions', + 'num_concurrent': num_concurrent, + 'max_retries': 3, + 'tokenized_requests': False, + }, + tasks=TASK, + num_fewshot=5, + limit=500, + apply_chat_template=True, # Enable chat template for Instruct models + fewshot_as_multiturn=True, # Format few-shot examples as conversation turns + log_samples=False, + ) + return results + + except Exception as e: + print_error(f'LM-Eval failed: {e}') + raise + + +def validate_accuracy( + results: dict, + model_name: str, +) -> bool: + """Validate that accuracy meets expected thresholds. + + Args: + results: LM-Eval results dictionary + model_name: Model name being evaluated + + Returns: + True if accuracy is within acceptable range, False otherwise + """ + measured_value = results['results'][TASK][FILTER] + expected_value = EXPECTED_VALUES.get(model_name) + + print() + print('=' * 60) + print_info('Accuracy Results:') + print_info(f' Model: {model_name}') + print_info(f' Task: {TASK}') + print_info(f' Metric: {FILTER}') + print_info(f' Measured Accuracy: {measured_value:.4f}') + + if expected_value is None: + print_warning( + f'No expected baseline found for {model_name}. ' 'Cannot validate accuracy.' + ) + print_info( + 'If this is the first time testing this model, ' + 'you may want to add the measured value to EXPECTED_VALUES.' + ) + print('=' * 60) + return True # Pass if no baseline (assume correct) + + print_info(f' Expected Accuracy: {expected_value:.4f}') + print_info(f' Tolerance: ±{RTOL:.4f}') + + lower_bound = expected_value - RTOL + # Higher accuracy is always acceptable, so we only enforce lower bound + minimum_threshold = 0.3 + + print_info(f' Minimum Threshold: {minimum_threshold:.4f}') + print_info(f' Lower Bound: {lower_bound:.4f}') + print('=' * 60) + print() + + if measured_value >= lower_bound: + print_success( + f'Accuracy meets requirements! ' + f'({measured_value:.4f} vs expected {expected_value:.4f})' + ) + return True + else: + print_error( + f'Accuracy below acceptable threshold! ' + f'({measured_value:.4f} vs expected {expected_value:.4f})' + ) + print_error( + f'Shortfall: {lower_bound - measured_value:.4f} ' + f'(minimum: {lower_bound:.4f})' + ) + return False + + +def main(): + parser = argparse.ArgumentParser( + description='Test P/D disaggregation accuracy using LM-Eval' + ) + parser.add_argument( + '--router-url', + type=str, + required=True, + help='URL of the router (e.g., http://localhost:8300)', + ) + parser.add_argument( + '--model', + type=str, + help='Model name to use for testing (can also use TEST_MODEL env var)', + ) + parser.add_argument( + '--num-concurrent', + type=int, + default=20, + help='Number of concurrent requests (default: 20)', + ) + parser.add_argument( + '--skip-connectivity', + action='store_true', + help='Skip initial connectivity test', + ) + + args = parser.parse_args() + + # Get model name from args or environment + model_name = args.model or os.environ.get('TEST_MODEL') + if not model_name: + print_error('Model name must be provided via --model or TEST_MODEL env var') + return 1 + + # Construct base URL for OpenAI API + base_url = f'{args.router_url}/v1' + + print() + print('=' * 60) + print_info('P/D Disaggregation LM-Eval Accuracy Test') + print('=' * 60) + print_info(f'Router URL: {args.router_url}') + print_info(f'API Base URL: {base_url}') + print_info(f'Model: {model_name}') + print_info(f'Task: {TASK}') + print_info(f'Concurrent Reqs: {args.num_concurrent}') + print('=' * 60) + print() + + # Step 1: Connectivity test (optional) + if not args.skip_connectivity: + if not run_simple_prompt(base_url, model_name): + print_error('Connectivity test failed. Aborting evaluation.') + return 1 + print() + + # Step 2: Run LM-Eval + try: + results = run_accuracy_evaluation( + base_url=base_url, + model_name=model_name, + num_concurrent=args.num_concurrent, + ) + except Exception as e: + print_error(f'Evaluation failed: {e}') + return 1 + + # Step 3: Validate accuracy + if not validate_accuracy(results, model_name): + print_error('Accuracy validation failed!') + return 1 + + print_success('All accuracy tests PASSED!') + return 0 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/lmdeploy/router/py_test/e2e/pd_disagg_lmdeploy/test_pd_accuracy.py b/lmdeploy/router/py_test/e2e/pd_disagg_lmdeploy/test_pd_accuracy.py new file mode 100755 index 0000000000..78a758ccf2 --- /dev/null +++ b/lmdeploy/router/py_test/e2e/pd_disagg_lmdeploy/test_pd_accuracy.py @@ -0,0 +1,415 @@ +# Copyright (c) OpenMMLab. All rights reserved. +#!/usr/bin/env python3 +"""P/D Disaggregation Accuracy Test for LMDeploy Router. + +This script validates that the router correctly routes requests through prefill and decode instances with proper output +accuracy. +""" + +import argparse +import json +import sys + +import requests + + +class Colors: + """Terminal colors for output.""" + + GREEN = '\033[92m' + RED = '\033[91m' + YELLOW = '\033[93m' + BLUE = '\033[94m' + RESET = '\033[0m' + + +def print_success(msg: str): + print(f'{Colors.GREEN}✓ {msg}{Colors.RESET}') + + +def print_error(msg: str): + print(f'{Colors.RED}✗ {msg}{Colors.RESET}') + + +def print_info(msg: str): + print(f'{Colors.BLUE}ℹ {msg}{Colors.RESET}') + + +def print_warning(msg: str): + print(f'{Colors.YELLOW}⚠ {msg}{Colors.RESET}') + + +SAMPLE_PROMPTS = [ + 'Red Hat is the best company in the world to work for because', + 'We hold these truths to be self-evident, that all men are created equal,', + 'The quick brown fox jumps over the lazy dog and', + 'In the beginning was the Word, and the Word was with', + 'To be or not to be, that is the question:', + 'Four score and seven years ago our fathers brought forth on this continent,', + 'It was the best of times, it was the worst of times,', + 'Call me Ishmael. Some years ago—never mind how long precisely—', + 'All happy families are alike; each unhappy family is unhappy in', + 'It is a truth universally acknowledged, that a single man in possession', +] + + +def test_completion_accuracy( + router_url: str, + model: str, + num_requests: int = 20, + max_tokens: int = 30, + temperature: float = 0.0, +) -> bool: + """Test that completions return valid outputs through P/D disaggregation. + + Args: + router_url: URL of the router + model: Model name to use + num_requests: Number of test requests to send + max_tokens: Maximum tokens to generate + temperature: Sampling temperature (0 for deterministic) + + Returns: + True if all tests pass, False otherwise + """ + print_info(f'Testing completion accuracy with {num_requests} requests') + + session = requests.Session() + failures = [] + + for i in range(num_requests): + prompt = SAMPLE_PROMPTS[i % len(SAMPLE_PROMPTS)] + + try: + response = session.post( + f'{router_url}/v1/completions', + json={ + 'model': model, + 'prompt': prompt, + 'max_tokens': max_tokens, + 'temperature': temperature, + 'stream': False, + }, + timeout=60, + ) + + if response.status_code != 200: + failures.append(f'Request {i}: HTTP {response.status_code}') + print_error(f'Request {i} failed with status {response.status_code}') + continue + + data = response.json() + + # Validate response structure + if 'choices' not in data or not data['choices']: + failures.append(f'Request {i}: No choices in response') + print_error(f'Request {i}: No choices in response') + continue + + choice = data['choices'][0] + if 'text' not in choice: + failures.append(f'Request {i}: No text in choice') + print_error(f'Request {i}: No text in choice') + continue + + output_text = choice['text'] + + # Validate output is non-empty + if not output_text or len(output_text.strip()) == 0: + failures.append(f'Request {i}: Empty output') + print_error(f'Request {i}: Empty output') + continue + + # Validate usage information + if 'usage' not in data: + print_warning(f'Request {i}: No usage information') + + print_success(f'Request {i}: Generated {len(output_text)} chars') + + except requests.RequestException as e: + failures.append(f'Request {i}: {str(e)}') + print_error(f'Request {i} failed with exception: {e}') + except json.JSONDecodeError as e: + failures.append(f'Request {i}: Invalid JSON response') + print_error(f'Request {i}: JSON decode error: {e}') + except Exception as e: + failures.append(f'Request {i}: Unexpected error: {str(e)}') + print_error(f'Request {i}: Unexpected error: {e}') + + # Summary + success_count = num_requests - len(failures) + print() + print_info('=' * 60) + print_info('Completion Accuracy Test Results:') + print_info(f' Total Requests: {num_requests}') + print_info(f' Successful: {success_count}') + print_info(f' Failed: {len(failures)}') + print_info('=' * 60) + + if failures: + print() + print_error('Failures:') + for failure in failures: + print_error(f' - {failure}') + return False + + return True + + +def test_streaming_accuracy( + router_url: str, + model: str, + num_requests: int = 5, + max_tokens: int = 20, + temperature: float = 0.0, +) -> bool: + """Test that streaming completions work correctly through P/D + disaggregation. + + Args: + router_url: URL of the router + model: Model name to use + num_requests: Number of test requests to send + max_tokens: Maximum tokens to generate + temperature: Sampling temperature + + Returns: + True if all tests pass, False otherwise + """ + print_info(f'Testing streaming accuracy with {num_requests} requests') + + session = requests.Session() + failures = [] + + for i in range(num_requests): + prompt = SAMPLE_PROMPTS[i % len(SAMPLE_PROMPTS)] + + try: + response = session.post( + f'{router_url}/v1/completions', + json={ + 'model': model, + 'prompt': prompt, + 'max_tokens': max_tokens, + 'temperature': temperature, + 'stream': True, + }, + timeout=60, + stream=True, + ) + + if response.status_code != 200: + failures.append(f'Stream {i}: HTTP {response.status_code}') + print_error(f'Stream {i} failed with status {response.status_code}') + continue + + chunks = [] + for line in response.iter_lines(): + if not line: + continue + + line_str = line.decode('utf-8') + if line_str.startswith('data: '): + data_str = line_str[6:] + if data_str.strip() == '[DONE]': + break + + try: + chunk = json.loads(data_str) + chunks.append(chunk) + except json.JSONDecodeError: + continue + + if not chunks: + failures.append(f'Stream {i}: No chunks received') + print_error(f'Stream {i}: No chunks received') + continue + + # Reconstruct full text from chunks + full_text = '' + for chunk in chunks: + if 'choices' in chunk and chunk['choices']: + choice = chunk['choices'][0] + if 'text' in choice: + full_text += choice['text'] + + if not full_text: + failures.append(f'Stream {i}: No text in chunks') + print_error(f'Stream {i}: No text in chunks') + continue + + print_success( + f'Stream {i}: Received {len(chunks)} chunks, {len(full_text)} chars' + ) + + except requests.RequestException as e: + failures.append(f'Stream {i}: {str(e)}') + print_error(f'Stream {i} failed with exception: {e}') + except Exception as e: + failures.append(f'Stream {i}: Unexpected error: {str(e)}') + print_error(f'Stream {i}: Unexpected error: {e}') + + # Summary + success_count = num_requests - len(failures) + print() + print_info('=' * 60) + print_info('Streaming Accuracy Test Results:') + print_info(f' Total Requests: {num_requests}') + print_info(f' Successful: {success_count}') + print_info(f' Failed: {len(failures)}') + print_info('=' * 60) + + if failures: + print() + print_error('Failures:') + for failure in failures: + print_error(f' - {failure}') + return False + + return True + + +def test_router_health(router_url: str) -> bool: + """Test that the router is healthy and responding. + + Args: + router_url: URL of the router + + Returns: + True if healthy, False otherwise + """ + print_info('Checking router health...') + + try: + # Try health endpoint + response = requests.get(f'{router_url}/health', timeout=10) + if response.status_code == 200: + print_success('Router health check passed') + return True + except requests.RequestException: + pass + + try: + # Try models endpoint as fallback + response = requests.get(f'{router_url}/v1/models', timeout=10) + if response.status_code == 200: + print_success('Router models endpoint accessible') + return True + except requests.RequestException as e: + print_error(f'Router health check failed: {e}') + return False + + print_error('Router is not healthy') + return False + + +def main(): + parser = argparse.ArgumentParser( + description='Test P/D disaggregation accuracy through LMDeploy router' + ) + parser.add_argument( + '--router-url', + type=str, + required=True, + help='URL of the router (e.g., http://localhost:8300)', + ) + parser.add_argument( + '--model', + type=str, + required=True, + help='Model name to use for testing', + ) + parser.add_argument( + '--num-requests', + type=int, + default=20, + help='Number of completion requests to test (default: 20)', + ) + parser.add_argument( + '--num-streaming-requests', + type=int, + default=5, + help='Number of streaming requests to test (default: 5)', + ) + parser.add_argument( + '--max-tokens', + type=int, + default=30, + help='Maximum tokens to generate (default: 30)', + ) + parser.add_argument( + '--skip-streaming', + action='store_true', + help='Skip streaming tests', + ) + + args = parser.parse_args() + + print() + print_info('=' * 60) + print_info('P/D Disaggregation Accuracy Test') + print_info('=' * 60) + print_info(f'Router URL: {args.router_url}') + print_info(f'Model: {args.model}') + print_info(f'Completion Tests: {args.num_requests}') + print_info(f'Streaming Tests: {args.num_streaming_requests}') + print_info('=' * 60) + print() + + # Test 1: Router Health + if not test_router_health(args.router_url): + print_error('Router health check failed. Aborting tests.') + return 1 + + print() + + # Test 2: Completion Accuracy + completion_success = test_completion_accuracy( + router_url=args.router_url, + model=args.model, + num_requests=args.num_requests, + max_tokens=args.max_tokens, + ) + + print() + + # Test 3: Streaming Accuracy (optional) + streaming_success = True + if not args.skip_streaming: + streaming_success = test_streaming_accuracy( + router_url=args.router_url, + model=args.model, + num_requests=args.num_streaming_requests, + max_tokens=args.max_tokens, + ) + print() + + # Final Summary + print_info('=' * 60) + print_info('Final Results:') + print_info('=' * 60) + + if completion_success: + print_success('Completion Tests: PASSED') + else: + print_error('Completion Tests: FAILED') + + if not args.skip_streaming: + if streaming_success: + print_success('Streaming Tests: PASSED') + else: + print_error('Streaming Tests: FAILED') + + print_info('=' * 60) + print() + + if completion_success and streaming_success: + print_success('All P/D disaggregation accuracy tests PASSED!') + return 0 + else: + print_error('Some P/D disaggregation accuracy tests FAILED!') + return 1 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/lmdeploy/router/py_test/e2e/pd_disagg_lmdeploy/tp_config_sweep_test.sh b/lmdeploy/router/py_test/e2e/pd_disagg_lmdeploy/tp_config_sweep_test.sh new file mode 100755 index 0000000000..eed90cc602 --- /dev/null +++ b/lmdeploy/router/py_test/e2e/pd_disagg_lmdeploy/tp_config_sweep_test.sh @@ -0,0 +1,24 @@ +#!/usr/bin/env bash +# Sweep LMDeploy P/D tensor-parallel combinations. The caller must provision +# one prefill and one decode server for each requested TP size. + +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +TEST_SCRIPT="${SCRIPT_DIR}/run_accuracy_test.sh" + +for tp_size in ${TP_SIZES:-1 2}; do + prefill_urls_var="PREFILL_TP${tp_size}_URLS" + decode_urls_var="DECODE_TP${tp_size}_URLS" + prefill_urls=${!prefill_urls_var:-} + decode_urls=${!decode_urls_var:-} + if [[ -z "${prefill_urls}" || -z "${decode_urls}" ]]; then + echo "Skipping TP=${tp_size}: set ${prefill_urls_var} and ${decode_urls_var}." >&2 + continue + fi + + echo "Running LMDeploy P/D accuracy test with TP=${tp_size}" + PREFILL_URLS="${prefill_urls}" \ + DECODE_URLS="${decode_urls}" \ + bash "${TEST_SCRIPT}" +done diff --git a/lmdeploy/router/py_test/e2e/test_e2e_embeddings.py b/lmdeploy/router/py_test/e2e/test_e2e_embeddings.py new file mode 100644 index 0000000000..389ab9f097 --- /dev/null +++ b/lmdeploy/router/py_test/e2e/test_e2e_embeddings.py @@ -0,0 +1,27 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import pytest +import requests + + +@pytest.mark.e2e +def test_embeddings_basic(e2e_router_only_rr, e2e_embedding_worker, e2e_embedding_model): + router_url = e2e_router_only_rr.url + response = requests.post( + f'{router_url}/add_worker', + params={'url': e2e_embedding_worker.url}, + timeout=30, + ) + response.raise_for_status() + + response = requests.post( + f'{router_url}/v1/embeddings', + json={ + 'model': e2e_embedding_model, + 'input': ['the quick brown fox', 'jumps over the lazy dog'], + }, + timeout=120, + ) + assert response.status_code == 200, response.text + data = response.json()['data'] + assert len(data) == 2 + assert all(item['embedding'] for item in data) diff --git a/lmdeploy/router/py_test/e2e/test_pd_router.py b/lmdeploy/router/py_test/e2e/test_pd_router.py new file mode 100644 index 0000000000..621d110ea3 --- /dev/null +++ b/lmdeploy/router/py_test/e2e/test_pd_router.py @@ -0,0 +1,41 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import pytest +import requests + + +@pytest.mark.e2e +def test_lmdeploy_pd_chat(e2e_pd_cluster): + response = requests.post( + f'{e2e_pd_cluster.url}/v1/chat/completions', + json={ + 'model': e2e_pd_cluster.model, + 'messages': [{'role': 'user', 'content': 'Say PD works.'}], + 'max_tokens': 8, + 'temperature': 0.0, + 'stream': False, + }, + timeout=180, + ) + assert response.status_code == 200, response.text + choice = response.json()['choices'][0] + assert isinstance(choice['message']['content'], str) + + +@pytest.mark.e2e +def test_lmdeploy_pd_streaming(e2e_pd_cluster): + response = requests.post( + f'{e2e_pd_cluster.url}/v1/chat/completions', + json={ + 'model': e2e_pd_cluster.model, + 'messages': [{'role': 'user', 'content': 'Say PD streaming works.'}], + 'max_tokens': 8, + 'temperature': 0.0, + 'stream': True, + }, + stream=True, + timeout=180, + ) + assert response.status_code == 200, response.text + chunks = [line for line in response.iter_lines() if line] + assert chunks + assert any(line == b'data: [DONE]' for line in chunks) diff --git a/lmdeploy/router/py_test/e2e/test_regular_router.py b/lmdeploy/router/py_test/e2e/test_regular_router.py new file mode 100644 index 0000000000..8dffeee61b --- /dev/null +++ b/lmdeploy/router/py_test/e2e/test_regular_router.py @@ -0,0 +1,64 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import pytest +import requests +from py_test.e2e.conftest import _discover_model + + +@pytest.mark.e2e +def test_regular_router_serving( + e2e_router_only_rr, e2e_primary_worker +): + router_url = e2e_router_only_rr.url + response = requests.post( + f'{router_url}/add_worker', + params={'url': e2e_primary_worker.url}, + timeout=30, + ) + response.raise_for_status() + + model = _discover_model(e2e_primary_worker.url) + response = requests.post( + f'{router_url}/v1/chat/completions', + json={ + 'model': model, + 'messages': [{'role': 'user', 'content': 'Say router works.'}], + 'max_tokens': 8, + 'temperature': 0.0, + 'stream': False, + }, + timeout=120, + ) + assert response.status_code == 200, response.text + choice = response.json()['choices'][0] + assert isinstance(choice['message']['content'], str) + + +@pytest.mark.e2e +def test_regular_router_streaming( + e2e_router_only_rr, e2e_primary_worker +): + router_url = e2e_router_only_rr.url + response = requests.post( + f'{router_url}/add_worker', + params={'url': e2e_primary_worker.url}, + timeout=30, + ) + response.raise_for_status() + model = _discover_model(e2e_primary_worker.url) + + response = requests.post( + f'{router_url}/v1/chat/completions', + json={ + 'model': model, + 'messages': [{'role': 'user', 'content': 'Say streaming works.'}], + 'max_tokens': 8, + 'temperature': 0.0, + 'stream': True, + }, + stream=True, + timeout=120, + ) + assert response.status_code == 200, response.text + chunks = [line for line in response.iter_lines() if line] + assert chunks + assert any(line == b'data: [DONE]' for line in chunks) diff --git a/lmdeploy/router/py_test/fixtures/__init__.py b/lmdeploy/router/py_test/fixtures/__init__.py new file mode 100644 index 0000000000..3c61155d6a --- /dev/null +++ b/lmdeploy/router/py_test/fixtures/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""Shared fixtures for router integration tests.""" diff --git a/lmdeploy/router/py_test/fixtures/mock_worker.py b/lmdeploy/router/py_test/fixtures/mock_worker.py new file mode 100644 index 0000000000..94194e280f --- /dev/null +++ b/lmdeploy/router/py_test/fixtures/mock_worker.py @@ -0,0 +1,286 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""Lightweight mock worker HTTP server for router integration tests. + +Implements minimal endpoints used by the router: +- GET /health, /health_generate +- POST /generate, /v1/completions, /v1/chat/completions +- POST /flush_cache +- GET /get_server_info, /get_model_info, /v1/models +- GET /distserve/engine_info +- POST /distserve/p2p_initialize, /distserve/p2p_connect, /distserve/p2p_drop_connect + +Behavior knobs are controlled via CLI flags to simulate failures, latency, and load. +""" + +import argparse +import asyncio +import json +import os +import random +import signal +import sys +import time +from contextlib import asynccontextmanager + +import uvicorn +from fastapi import FastAPI, HTTPException, Request +from fastapi.responses import JSONResponse, PlainTextResponse, StreamingResponse + +# Global state (per-process) +_inflight = 0 +_failures_seen = 0 + + +def _parse_args() -> argparse.Namespace: + p = argparse.ArgumentParser() + p.add_argument('--host', default='127.0.0.1') + p.add_argument('--port', type=int, required=True) + p.add_argument('--worker-id', default=None) + p.add_argument('--latency-ms', type=int, default=0) + p.add_argument('--timeout', action='store_true') + p.add_argument('--status-code', type=int, default=200) + p.add_argument('--fail-first-n', type=int, default=0) + p.add_argument('--random-fail-rate', type=float, default=0.0) + p.add_argument('--require-api-key', action='store_true') + p.add_argument('--api-key', default=None) + p.add_argument('--max-payload-bytes', type=int, default=10 * 1024 * 1024) + p.add_argument('--stream', action='store_true') + p.add_argument('--dp-size', type=int, default=1) + p.add_argument('--crash-on-request', action='store_true') + p.add_argument('--health-fail-after-ms', type=int, default=0) + return p.parse_args() + + +def _extract_worker_id(args: argparse.Namespace) -> str: + if args.worker_id: + return str(args.worker_id) + # default to port (unique enough for tests) + return f'worker-{args.port}' + + +def create_app(args: argparse.Namespace) -> FastAPI: + app = FastAPI() + worker_id = _extract_worker_id(args) + start_ts = time.time() + crashed = {'done': False} + + async def maybe_delay(): + if args.latency_ms > 0: + await asyncio.sleep(args.latency_ms / 1000.0) + + def should_fail() -> int | None: + global _failures_seen + # Fail first N requests (500) + if args.fail_first_n > 0 and _failures_seen < args.fail_first_n: + _failures_seen += 1 + return 500 + # Random failure probability (500) + if args.random_fail_rate > 0.0 and random.random() < args.random_fail_rate: + return 500 + # Forced status code override (non-200) for all responses + if args.status_code != 200: + return int(args.status_code) + return None + + def check_api_key(request: Request): + if not args.require_api_key: + return + auth = request.headers.get('Authorization') + if not auth or not auth.startswith('Bearer '): + raise HTTPException(status_code=401, detail='Unauthorized') + key = auth.split(' ', 1)[1] + if args.api_key and key != args.api_key: + raise HTTPException(status_code=401, detail='Unauthorized') + + @asynccontextmanager + async def track_inflight(): + global _inflight + _inflight += 1 + try: + yield + finally: + _inflight -= 1 + + @app.get('/health') + async def health(): + if ( + args.health_fail_after_ms + and (time.time() - start_ts) * 1000.0 >= args.health_fail_after_ms + ): + return PlainTextResponse('bad', status_code=500) + return PlainTextResponse('ok', status_code=200) + + @app.get('/health_generate') + async def health_generate(): + return PlainTextResponse('ok', status_code=200) + + @app.post('/flush_cache') + async def flush_cache(): + return PlainTextResponse('ok', status_code=200) + + @app.get('/get_model_info') + async def get_model_info(): + return JSONResponse({'model': 'mock', 'vocab_size': 32000}) + + @app.get('/v1/models') + async def list_models(): + return JSONResponse({'data': [{'id': 'mock', 'object': 'model'}]}) + + @app.get('/get_server_info') + async def get_server_info(request: Request): + # Enforce API key on server info when required (used by router for worker discovery) + check_api_key(request) + return JSONResponse( + { + 'worker_id': worker_id, + 'load_in_flight': _inflight, + 'cache': {'size': 0, 'hit_rate': 0.0}, + 'dp_size': int(args.dp_size), + } + ) + + @app.get('/get_load') + async def get_load(): + return JSONResponse({'load': _inflight}) + + @app.get('/distserve/engine_info') + async def distserve_engine_info(): + return JSONResponse({'tp_size': 1, 'dp_size': int(args.dp_size)}) + + @app.post('/distserve/p2p_initialize') + async def distserve_p2p_initialize(): + return JSONResponse( + { + 'status': 1, + 'engine_endpoint_info': {'url': f'{args.host}:{args.port}'}, + 'kvtransfer_endpoint_info': [], + } + ) + + @app.post('/distserve/p2p_connect') + async def distserve_p2p_connect(): + return JSONResponse({'status': 1}) + + @app.post('/distserve/p2p_drop_connect') + async def distserve_p2p_drop_connect(): + return JSONResponse({'status': 1}) + + def make_json_response(obj: dict, status_code: int = 200) -> JSONResponse: + resp = JSONResponse(obj, status_code=status_code) + resp.headers['X-Worker-Id'] = worker_id + return resp + + async def handle_text_request(request: Request): + # Authorization + check_api_key(request) + + # Payload limit + body = await request.body() + if len(body) > args.max_payload_bytes: + return make_json_response({'error': 'payload too large'}, status_code=413) + + # Simulate crash on first request + if args.crash_on_request and not crashed['done']: + crashed['done'] = True + os._exit(1) + + # Optional timeout (simulate hang) + if args.timeout: + await asyncio.sleep(3600) + + # Optional latency + await maybe_delay() + + # Optional failures + fail_code = should_fail() + if fail_code is not None and fail_code != 200: + return make_json_response( + {'error': f'mock failure {fail_code}'}, status_code=fail_code + ) + + # Build response echoing minimal shape + try: + data = await request.json() + except (json.JSONDecodeError, ValueError): + data = {} + + now = time.time() + is_prefill = ( + data.get('with_cache') is True and data.get('preserve_cache') is True + ) + if is_prefill: + ret = { + 'id': int(now * 1000), + 'cache_block_ids': [1, 2, 3], + 'remote_token_ids': [0, 1, 3], + } + return make_json_response(ret, status_code=200) + + ret = { + 'id': f'cmpl-{int(now*1000)}', + 'object': 'text_completion', + 'created': int(now), + 'model': 'mock', + 'choices': [ + { + 'text': 'ok', + 'index': 0, + 'finish_reason': 'stop', + } + ], + 'worker_id': worker_id, + 'echo': data, + } + return make_json_response(ret, status_code=200) + + async def handle_stream_request(request: Request): + check_api_key(request) + + async def gen(): + # minimal 2-chunk stream then [DONE] + for i in range(2): + await asyncio.sleep(0.01) + chunk = { + 'choices': [{'delta': {'content': 'x'}}], + 'worker_id': worker_id, + } + yield f'data: {json.dumps(chunk)}\n\n' + yield 'data: [DONE]\n\n' + + headers = {'X-Worker-Id': worker_id} + return StreamingResponse(gen(), media_type='text/event-stream', headers=headers) + + @app.post('/generate') + async def generate(request: Request): + async with track_inflight(): + if args.stream: + return await handle_stream_request(request) + return await handle_text_request(request) + + @app.post('/v1/completions') + async def completions(request: Request): + async with track_inflight(): + if args.stream: + return await handle_stream_request(request) + return await handle_text_request(request) + + @app.post('/v1/chat/completions') + async def chat_completions(request: Request): + async with track_inflight(): + if args.stream: + return await handle_stream_request(request) + return await handle_text_request(request) + + return app + + +def main() -> None: + args = _parse_args() + app = create_app(args) + # Handle SIGTERM gracefully for fast test teardown + signal.signal(signal.SIGTERM, lambda *_: sys.exit(0)) + uvicorn.run(app, host=args.host, port=args.port, log_level='warning') + + +if __name__ == '__main__': + main() diff --git a/lmdeploy/router/py_test/fixtures/ports.py b/lmdeploy/router/py_test/fixtures/ports.py new file mode 100644 index 0000000000..dbbb14dd66 --- /dev/null +++ b/lmdeploy/router/py_test/fixtures/ports.py @@ -0,0 +1,27 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import socket + + +def find_free_port() -> int: + """Return an available TCP port on localhost.""" + return find_free_ports(1)[0] + + +def find_free_ports(count: int) -> list[int]: + """Return unique available TCP ports, binding all sockets before release. + + Keeping the sockets open while allocating avoids the common race where independent port requests return two sockets + with the same ephemeral port. + """ + sockets = [] + ports: set[int] = set() + try: + while len(ports) < count: + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sockets.append(sock) + sock.bind(('127.0.0.1', 0)) + ports.add(sock.getsockname()[1]) + finally: + for sock in sockets: + sock.close() + return list(ports) diff --git a/lmdeploy/router/py_test/fixtures/router_manager.py b/lmdeploy/router/py_test/fixtures/router_manager.py new file mode 100644 index 0000000000..c20b47ed6b --- /dev/null +++ b/lmdeploy/router/py_test/fixtures/router_manager.py @@ -0,0 +1,221 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import subprocess +import time +from dataclasses import dataclass + +import requests + +from .ports import find_free_ports + + +@dataclass +class ProcHandle: + process: subprocess.Popen + url: str + + +class RouterStartupError(RuntimeError): + """Raised when the router fails before its HTTP server starts.""" + + +class RouterManager: + """Helper to spawn a router process and interact with admin endpoints.""" + + def __init__(self): + self._children: list[subprocess.Popen] = [] + + def start_router( + self, + worker_urls: list[str] | None = None, + policy: str = 'round_robin', + port: int | None = None, + extra: dict | None = None, + # PD options + lmdeploy_pd_disaggregation: bool = False, + prefill_urls: list[str] | None = None, + decode_urls: list[str] | None = None, + prefill_policy: str | None = None, + decode_policy: str | None = None, + ) -> ProcHandle: + worker_urls = worker_urls or [] + retry_count = 0 if port is not None else 2 + while True: + try: + return self._start_router_process( + worker_urls=worker_urls, + policy=policy, + port=port, + extra=extra, + lmdeploy_pd_disaggregation=lmdeploy_pd_disaggregation, + prefill_urls=prefill_urls, + decode_urls=decode_urls, + prefill_policy=prefill_policy, + decode_policy=decode_policy, + ) + except RouterStartupError: + if retry_count == 0: + raise + retry_count -= 1 + port = None + + def _start_router_process( + self, + worker_urls: list[str], + policy: str, + port: int | None, + extra: dict | None, + lmdeploy_pd_disaggregation: bool, + prefill_urls: list[str] | None, + decode_urls: list[str] | None, + prefill_policy: str | None, + decode_policy: str | None, + ) -> ProcHandle: + router_port, prom_port = ( + [port, *find_free_ports(1)] if port is not None else find_free_ports(2) + ) + cmd = [ + 'python3', + '-m', + 'lmdeploy_router.launch_router', + '--host', + '127.0.0.1', + '--port', + str(router_port), + '--policy', + policy, + ] + # Avoid Prometheus port collisions by assigning unique free ports. + cmd.extend( + ['--prometheus-port', str(prom_port), '--prometheus-host', '127.0.0.1'] + ) + if worker_urls: + cmd.extend(['--worker-urls', *worker_urls]) + + # PD routing configuration + if lmdeploy_pd_disaggregation: + cmd.append('--lmdeploy-pd-disaggregation') + if prefill_urls: + for url in prefill_urls: + cmd.extend(['--prefill', url]) + if decode_urls: + for url in decode_urls: + cmd.extend(['--decode', url]) + if prefill_policy: + cmd.extend(['--prefill-policy', prefill_policy]) + if decode_policy: + cmd.extend(['--decode-policy', decode_policy]) + + # Map supported extras to CLI flags (subset for integration) + if extra: + flag_map = { + 'max_payload_size': '--max-payload-size', + 'api_key': '--api-key', + # Health/monitoring + 'worker_startup_check_interval': '--worker-startup-check-interval', + # Cache-aware tuning + 'cache_threshold': '--cache-threshold', + 'balance_abs_threshold': '--balance-abs-threshold', + 'balance_rel_threshold': '--balance-rel-threshold', + # Retry + 'retry_max_retries': '--retry-max-retries', + 'retry_initial_backoff_ms': '--retry-initial-backoff-ms', + 'retry_max_backoff_ms': '--retry-max-backoff-ms', + 'retry_backoff_multiplier': '--retry-backoff-multiplier', + 'retry_jitter_factor': '--retry-jitter-factor', + 'disable_retries': '--disable-retries', + # Circuit breaker + 'cb_failure_threshold': '--cb-failure-threshold', + 'cb_success_threshold': '--cb-success-threshold', + 'cb_timeout_duration_secs': '--cb-timeout-duration-secs', + 'cb_window_duration_secs': '--cb-window-duration-secs', + 'disable_circuit_breaker': '--disable-circuit-breaker', + # Rate limiting + 'max_concurrent_requests': '--max-concurrent-requests', + 'queue_size': '--queue-size', + 'queue_timeout_secs': '--queue-timeout-secs', + 'rate_limit_tokens_per_second': '--rate-limit-tokens-per-second', + } + for k, v in extra.items(): + if v is None: + continue + flag = flag_map.get(k) + if not flag: + continue + if isinstance(v, bool): + if v: + cmd.append(flag) + else: + cmd.extend([flag, str(v)]) + + proc = subprocess.Popen(cmd) + self._children.append(proc) + url = f'http://127.0.0.1:{router_port}' + try: + self._wait_health(url) + except TimeoutError: + if self._router_failed_before_serving(proc): + raise RouterStartupError(f'Router at {url} failed during startup') from None + raise + return ProcHandle(process=proc, url=url) + + @staticmethod + def _router_failed_before_serving(process: subprocess.Popen) -> bool: + if process.poll() is None: + return False + stdout, _ = process.communicate(timeout=2) + return 'Router ready | workers:' not in stdout + + def _wait_health(self, base_url: str, timeout: float = 30.0): + start = time.time() + with requests.Session() as s: + while time.time() - start < timeout: + try: + r = s.get(f'{base_url}/health', timeout=2) + if r.status_code == 200: + return + except requests.RequestException: + pass + time.sleep(0.2) + raise TimeoutError(f'Router at {base_url} did not become healthy') + + def add_worker(self, base_url: str, worker_url: str) -> None: + r = requests.post(f'{base_url}/add_worker', params={'url': worker_url}) + assert r.status_code == 200, f'add_worker failed: {r.status_code} {r.text}' + + def remove_worker(self, base_url: str, worker_url: str) -> None: + r = requests.post(f'{base_url}/remove_worker', params={'url': worker_url}) + assert r.status_code == 200, f'remove_worker failed: {r.status_code} {r.text}' + + def add_lmdeploy_node(self, base_url: str, worker_url: str, role: int) -> None: + response = requests.post( + f'{base_url}/nodes/add', json={'url': worker_url, 'status': {'role': role}} + ) + assert response.status_code == 200, response.text + + def remove_lmdeploy_node(self, base_url: str, worker_url: str, role: int) -> None: + response = requests.post( + f'{base_url}/nodes/remove', + json={'url': worker_url, 'status': {'role': role}}, + ) + assert response.status_code == 200, response.text + + def lmdeploy_nodes(self, base_url: str) -> dict: + response = requests.get(f'{base_url}/nodes/status') + assert response.status_code == 200, response.text + return response.json() + + def list_workers(self, base_url: str) -> list[str]: + r = requests.get(f'{base_url}/list_workers') + assert r.status_code == 200, f'list_workers failed: {r.status_code} {r.text}' + data = r.json() + return data.get('urls', []) + + def stop_all(self): + for p in self._children: + if p.poll() is None: + p.terminate() + try: + p.wait(timeout=5) + except subprocess.TimeoutExpired: + p.kill() + self._children.clear() diff --git a/lmdeploy/router/py_test/integration/__init__.py b/lmdeploy/router/py_test/integration/__init__.py new file mode 100644 index 0000000000..240162833b --- /dev/null +++ b/lmdeploy/router/py_test/integration/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""Integration test package for the router.""" diff --git a/lmdeploy/router/py_test/integration/conftest.py b/lmdeploy/router/py_test/integration/conftest.py new file mode 100644 index 0000000000..c1af6ba338 --- /dev/null +++ b/lmdeploy/router/py_test/integration/conftest.py @@ -0,0 +1,111 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import subprocess +import time +from collections.abc import Iterable +from pathlib import Path + +import pytest +import requests + +from ..fixtures.ports import find_free_port +from ..fixtures.router_manager import RouterManager + + +def pytest_configure(config): + config.addinivalue_line('markers', 'integration: mark as router integration test') + config.addinivalue_line('markers', 'pd_engine: requires LMDeploy PD engine metadata') + config.addinivalue_line('markers', 'lmdeploy: requires externally managed LMDeploy backends') + + +@pytest.fixture +def router_manager() -> Iterable[RouterManager]: + mgr = RouterManager() + try: + yield mgr + finally: + mgr.stop_all() + + +def _spawn_mock_worker(args: list[str]) -> tuple[subprocess.Popen, str, str]: + router_root = Path(__file__).resolve().parents[2] + script = router_root / 'py_test' / 'fixtures' / 'mock_worker.py' + port = find_free_port() + worker_id = f'worker-{port}' + base_cmd = [ + 'python3', + str(script), + '--port', + str(port), + '--worker-id', + worker_id, + ] + cmd = base_cmd + args + proc = subprocess.Popen(cmd) + url = f'http://127.0.0.1:{port}' + _wait_health(url) + return proc, url, worker_id + + +def _wait_health(url: str, timeout: float = 10.0): + start = time.time() + with requests.Session() as s: + while time.time() - start < timeout: + try: + r = s.get(f'{url}/health', timeout=1) + if r.status_code == 200: + return + except requests.RequestException: + pass + time.sleep(0.1) + raise TimeoutError(f'Mock worker at {url} did not become healthy') + + +@pytest.fixture +def mock_worker(): + """Start a single healthy mock worker; yields (process, url, worker_id).""" + proc, url, worker_id = _spawn_mock_worker([]) + try: + yield proc, url, worker_id + finally: + if proc.poll() is None: + proc.terminate() + try: + proc.wait(timeout=3) + except subprocess.TimeoutExpired: + proc.kill() + + +@pytest.fixture +def mock_workers(): + """Factory to start N workers with custom args. + + Usage: + procs, urls, ids = mock_workers(n=3, args=["--latency-ms", "5"]) # same args for all + ... + """ + + procs: list[subprocess.Popen] = [] + + def _start(n: int, args: list[str] | None = None): + args = args or [] + new_procs: list[subprocess.Popen] = [] + urls: list[str] = [] + ids: list[str] = [] + for _ in range(n): + p, url, wid = _spawn_mock_worker(args) + procs.append(p) + new_procs.append(p) + urls.append(url) + ids.append(wid) + return new_procs, urls, ids + + try: + yield _start + finally: + for p in procs: + if p.poll() is None: + p.terminate() + try: + p.wait(timeout=3) + except subprocess.TimeoutExpired: + p.kill() diff --git a/lmdeploy/router/py_test/integration/lmdeploy/__init__.py b/lmdeploy/router/py_test/integration/lmdeploy/__init__.py new file mode 100644 index 0000000000..e777d7ec8a --- /dev/null +++ b/lmdeploy/router/py_test/integration/lmdeploy/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""LMDeploy integration tests for the router.""" diff --git a/lmdeploy/router/py_test/integration/lmdeploy/conftest.py b/lmdeploy/router/py_test/integration/lmdeploy/conftest.py new file mode 100644 index 0000000000..597e2632a0 --- /dev/null +++ b/lmdeploy/router/py_test/integration/lmdeploy/conftest.py @@ -0,0 +1,302 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import json +import os +import subprocess +import time +from collections.abc import Iterable, Sequence +from dataclasses import dataclass +from pathlib import Path + +import pytest +import requests + +from py_test.fixtures.ports import find_free_port + +DEFAULT_TOKEN_IDS = [ + 151644, + 8948, + 198, + 2610, + 525, + 264, + 109050, + 151645, + 198, + 151644, + 77091, + 198, +] + + +def pytest_configure(config): + config.addinivalue_line( + 'markers', 'lmdeploy: requires externally managed LMDeploy API servers' + ) + + +def _split_urls(value: str) -> tuple[str, ...]: + return tuple(url.strip().rstrip('/') for url in value.split(',') if url.strip()) + + +def _discover_models(payload: dict) -> list[str]: + if isinstance(payload.get('data'), list): + return [ + item['id'] + for item in payload['data'] + if isinstance(item, dict) and isinstance(item.get('id'), str) + ] + if isinstance(payload.get('id'), str): + return [payload['id']] + return [] + + +def _wait_for_http(url: str, timeout: float, process=None) -> None: + deadline = time.monotonic() + timeout + last_error = 'no response' + with requests.Session() as session: + while time.monotonic() < deadline: + if process is not None and process.poll() is not None: + raise RuntimeError(f'router exited with status {process.returncode}') + try: + response = session.get(url, timeout=2) + if response.status_code == 200: + return + last_error = f'HTTP {response.status_code}: {response.text[:200]}' + except requests.RequestException as exc: + last_error = str(exc) + time.sleep(0.2) + raise TimeoutError(f'timed out waiting for {url}: {last_error}') + + +@dataclass(frozen=True) +class LMDeployTestConfig: + backend_urls: tuple[str, ...] + model: str + token_ids: tuple[int, ...] + request_timeout: float + pd_prefill_url: str | None + pd_decode_url: str | None + + +@dataclass +class RouterHandle: + process: subprocess.Popen + url: str + log_path: Path + _log_file: object + + +class LMDeployRouterFactory: + def __init__( + self, + binary: Path, + config: LMDeployTestConfig, + artifact_dir: Path, + ) -> None: + self.binary = binary + self.config = config + self.artifact_dir = artifact_dir + self._handles: list[RouterHandle] = [] + + def start( + self, + *, + policy: str = 'random', + worker_urls: Sequence[str] | None = None, + lmdeploy_pd: bool = False, + prefill_urls: Sequence[str] = (), + decode_urls: Sequence[str] = (), + ) -> RouterHandle: + port = find_free_port() + prometheus_port = find_free_port() + if lmdeploy_pd: + urls = () + elif worker_urls is None: + urls = self.config.backend_urls + else: + urls = tuple(worker_urls) + mode = 'pd' if lmdeploy_pd else 'regular' + log_path = self.artifact_dir / f'router-{mode}-{policy}-{port}.log' + log_file = log_path.open('w', encoding='utf-8') + + command = [ + str(self.binary), + '--host', + '127.0.0.1', + '--port', + str(port), + '--policy', + policy, + '--prometheus-port', + str(prometheus_port), + '--prometheus-host', + '127.0.0.1', + ] + if urls: + command.extend(['--worker-urls', *urls]) + if lmdeploy_pd: + command.extend( + [ + '--lmdeploy-pd-disaggregation', + '--lmdeploy-migration-protocol', + 'rdma', + '--lmdeploy-rdma-link-type', + 'roce', + '--prefill-policy', + policy, + '--decode-policy', + policy, + ] + ) + for url in prefill_urls: + command.extend(['--prefill', url]) + for url in decode_urls: + command.extend(['--decode', url]) + + log_file.write('COMMAND: ' + ' '.join(command) + '\n\n') + log_file.flush() + process = subprocess.Popen( + command, + stdin=subprocess.DEVNULL, + stdout=log_file, + stderr=subprocess.STDOUT, + start_new_session=True, + ) + handle = RouterHandle( + process=process, + url=f'http://127.0.0.1:{port}', + log_path=log_path, + _log_file=log_file, + ) + self._handles.append(handle) + try: + _wait_for_http(f'{handle.url}/health', timeout=30, process=process) + except Exception as exc: + self.stop(handle) + tail = log_path.read_text(encoding='utf-8', errors='replace')[-4000:] + raise RuntimeError(f'failed to start router; log tail:\n{tail}') from exc + return handle + + def stop(self, handle: RouterHandle) -> None: + if handle.process.poll() is None: + handle.process.terminate() + try: + handle.process.wait(timeout=5) + except subprocess.TimeoutExpired: + handle.process.kill() + handle.process.wait(timeout=5) + if not handle._log_file.closed: + handle._log_file.close() + if handle in self._handles: + self._handles.remove(handle) + + def stop_all(self) -> None: + for handle in list(self._handles): + self.stop(handle) + + +@pytest.fixture(scope='session') +def lmdeploy_config() -> LMDeployTestConfig: + raw_urls = os.getenv('LMDEPLOY_BACKEND_URLS', '') + if not raw_urls: + pytest.skip( + 'set LMDEPLOY_BACKEND_URLS=url1,url2 to run real LMDeploy integration tests' + ) + backend_urls = _split_urls(raw_urls) + if len(backend_urls) < 2: + pytest.fail('LMDEPLOY_BACKEND_URLS must contain at least two backend URLs') + + timeout = float(os.getenv('LMDEPLOY_REQUEST_TIMEOUT', '120')) + discovered_models: list[str] = [] + backend_model_sets: list[set[str]] = [] + with requests.Session() as session: + for backend_url in backend_urls: + _wait_for_http(f'{backend_url}/health', timeout=10) + response = session.get(f'{backend_url}/v1/models', timeout=10) + response.raise_for_status() + models = _discover_models(response.json()) + if not models: + pytest.fail(f'{backend_url}/v1/models returned no model IDs') + discovered_models.extend(models) + backend_model_sets.append(set(models)) + + model = os.getenv('LMDEPLOY_MODEL') or discovered_models[0] + if any(model not in models for models in backend_model_sets): + pytest.fail( + f'model {model!r} was not reported by the configured LMDeploy backends' + ) + + raw_token_ids = os.getenv('LMDEPLOY_TOKEN_IDS') + token_ids = DEFAULT_TOKEN_IDS + if raw_token_ids: + parsed = json.loads(raw_token_ids) + if ( + not isinstance(parsed, list) + or not parsed + or not all(isinstance(item, int) for item in parsed) + ): + pytest.fail('LMDEPLOY_TOKEN_IDS must be a non-empty JSON list of integers') + token_ids = parsed + + return LMDeployTestConfig( + backend_urls=backend_urls, + model=model, + token_ids=tuple(token_ids), + request_timeout=timeout, + pd_prefill_url=os.getenv('LMDEPLOY_PD_PREFILL_URL'), + pd_decode_url=os.getenv('LMDEPLOY_PD_DECODE_URL'), + ) + + +@pytest.fixture(scope='session') +def lmdeploy_router_binary(lmdeploy_config) -> Path: + repo_root = Path(__file__).resolve().parents[3] + configured = os.getenv('LMDEPLOY_ROUTER_BIN') + if configured: + candidates = [Path(configured)] + else: + target_dir = Path(os.getenv('CARGO_TARGET_DIR', repo_root / 'target')) + candidates = [ + target_dir / 'release' / 'lmdeploy-router', + target_dir / 'debug' / 'lmdeploy-router', + ] + for candidate in candidates: + if candidate.is_file() and os.access(candidate, os.X_OK): + return candidate.resolve() + pytest.fail( + 'lmdeploy-router binary not found; run ' + '`cargo build --release --bin lmdeploy-router` or set LMDEPLOY_ROUTER_BIN' + ) + + +@pytest.fixture(scope='session') +def lmdeploy_router_factory( + lmdeploy_router_binary: Path, + lmdeploy_config: LMDeployTestConfig, +) -> Iterable[LMDeployRouterFactory]: + repo_root = Path(__file__).resolve().parents[3] + artifact_dir = Path( + os.getenv( + 'LMDEPLOY_TEST_ARTIFACT_DIR', + repo_root / '.test-artifacts' / 'lmdeploy', + ) + ) + artifact_dir.mkdir(parents=True, exist_ok=True) + factory = LMDeployRouterFactory( + binary=lmdeploy_router_binary, + config=lmdeploy_config, + artifact_dir=artifact_dir, + ) + try: + yield factory + finally: + factory.stop_all() + + +@pytest.fixture(scope='module') +def lmdeploy_router(lmdeploy_router_factory) -> Iterable[RouterHandle]: + handle = lmdeploy_router_factory.start(policy='random') + try: + yield handle + finally: + lmdeploy_router_factory.stop(handle) diff --git a/lmdeploy/router/py_test/integration/lmdeploy/test_load_balancing.py b/lmdeploy/router/py_test/integration/lmdeploy/test_load_balancing.py new file mode 100644 index 0000000000..03e54d34d9 --- /dev/null +++ b/lmdeploy/router/py_test/integration/lmdeploy/test_load_balancing.py @@ -0,0 +1,174 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""Load-balancing policy tests against two real LMDeploy workers.""" + +from collections import Counter + +import pytest +import requests + +pytestmark = [pytest.mark.integration, pytest.mark.lmdeploy] + +POLICIES = [ + 'random', + 'round_robin', + 'power_of_two', + 'consistent_hash', + 'cache_aware', +] + + +def _chat(config, base_url, payload): + response = requests.post( + f'{base_url}/v1/chat/completions', + json=payload, + timeout=config.request_timeout, + ) + response.raise_for_status() + return response.json() + + +def _message_text(response): + message = response['choices'][0]['message'] + return (message.get('reasoning_content') or '') + (message.get('content') or '') + + +def _numeric_id(response): + try: + return int(response['id']) + except (KeyError, TypeError, ValueError) as exc: + pytest.fail( + "load-balancing attribution requires LMDeploy's numeric per-worker response id" + ) + raise exc + + +def _baseline_ids(config, payload): + baselines = { + url: _numeric_id(_chat(config, url, payload)) for url in config.backend_urls + } + for _ in range(100): + values = list(baselines.values()) + if max(values) - min(values) >= 10: + return baselines + first = config.backend_urls[0] + baselines[first] = _numeric_id(_chat(config, first, payload)) + pytest.fail( + f'could not separate LMDeploy worker response-id baselines: {baselines}' + ) + + +def _attribute_worker(response, baselines): + response_id = _numeric_id(response) + exact = [url for url, baseline in baselines.items() if response_id == baseline + 1] + if len(exact) == 1: + baselines[exact[0]] = response_id + return exact[0] + + candidates = [ + (response_id - baseline, url) + for url, baseline in baselines.items() + if response_id > baseline + ] + if not candidates: + pytest.fail( + f'could not attribute response id {response_id}; baselines={baselines}' + ) + _, worker = min(candidates) + baselines[worker] = response_id + return worker + + +def _direct_outputs(config, payload): + return {url: _chat(config, url, payload) for url in config.backend_urls} + + +@pytest.mark.parametrize('policy', POLICIES) +def test_lmdeploy_policy_routes_messages_and_input_ids( + policy, + lmdeploy_config, + lmdeploy_router_factory, +): + router = lmdeploy_router_factory.start(policy=policy) + try: + warmup = { + 'model': lmdeploy_config.model, + 'messages': [{'role': 'user', 'content': 'hi'}], + 'max_tokens': 1, + 'temperature': 0, + 'stream': False, + } + baselines = _baseline_ids(lmdeploy_config, warmup) + served = Counter() + + message_payload = { + 'model': lmdeploy_config.model, + 'messages': [ + {'role': 'user', 'content': 'What is 2+2? Answer in one word.'} + ], + 'max_tokens': 30, + 'temperature': 0, + 'stream': False, + } + direct_messages = _direct_outputs(lmdeploy_config, message_payload) + baselines.update( + {url: _numeric_id(output) for url, output in direct_messages.items()} + ) + message_workers = [] + for _ in range(10): + output = _chat(lmdeploy_config, router.url, message_payload) + worker = _attribute_worker(output, baselines) + served[worker] += 1 + message_workers.append(worker) + assert _message_text(output) == _message_text(direct_messages[worker]) + + token_payload = { + 'model': lmdeploy_config.model, + 'messages': [], + 'input_ids': list(lmdeploy_config.token_ids), + 'max_tokens': 30, + 'temperature': 0, + 'return_token_ids': True, + 'stream': False, + } + direct_tokens = _direct_outputs(lmdeploy_config, token_payload) + baselines.update( + {url: _numeric_id(output) for url, output in direct_tokens.items()} + ) + token_workers = [] + for _ in range(10): + output = _chat(lmdeploy_config, router.url, token_payload) + worker = _attribute_worker(output, baselines) + served[worker] += 1 + token_workers.append(worker) + assert _message_text(output) == _message_text(direct_tokens[worker]) + assert output['choices'][0].get('output_ids') + assert ( + output['choices'][0]['output_ids'] + == direct_tokens[worker]['choices'][0]['output_ids'] + ) + + if policy in {'random', 'power_of_two'}: + assert all(served[url] > 0 for url in lmdeploy_config.backend_urls), served + elif policy == 'round_robin': + assert all( + current != previous + for previous, current in zip(message_workers, message_workers[1:]) + ), message_workers + elif policy == 'consistent_hash': + assert len(set(message_workers)) == 1, message_workers + assert len(set(token_workers)) == 1, token_workers + elif policy == 'cache_aware': + assert len(set(token_workers)) == 1, token_workers + + growing_workers = [] + for length in range(3, 8): + growing_payload = { + **token_payload, + 'input_ids': list(range(1, length + 1)), + 'max_tokens': 2, + } + output = _chat(lmdeploy_config, router.url, growing_payload) + growing_workers.append(_attribute_worker(output, baselines)) + assert len(set(growing_workers)) == 1, growing_workers + finally: + lmdeploy_router_factory.stop(router) diff --git a/lmdeploy/router/py_test/integration/lmdeploy/test_registration.py b/lmdeploy/router/py_test/integration/lmdeploy/test_registration.py new file mode 100644 index 0000000000..5d51de605c --- /dev/null +++ b/lmdeploy/router/py_test/integration/lmdeploy/test_registration.py @@ -0,0 +1,216 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""LMDeploy-compatible dynamic registration tests with real workers.""" + +import time + +import pytest +import requests + +pytestmark = [pytest.mark.integration, pytest.mark.lmdeploy] + + +def _add_node(router_url, worker_url, model, role): + return requests.post( + f'{router_url}/nodes/add', + json={ + 'url': worker_url, + 'status': {'models': [model], 'role': role}, + }, + timeout=30, + ) + + +def _wait_for_inference(config, router_url, timeout=30): + deadline = time.monotonic() + timeout + last_response = None + while time.monotonic() < deadline: + last_response = requests.post( + f'{router_url}/v1/chat/completions', + json={ + 'model': config.model, + 'messages': [{'role': 'user', 'content': 'Hello'}], + 'max_tokens': 2, + 'temperature': 0, + }, + timeout=config.request_timeout, + ) + if last_response.status_code == 200: + return last_response + time.sleep(0.5) + pytest.fail( + 'registered worker never became routable: ' + + (last_response.text[:500] if last_response is not None else 'no response') + ) + + +def _assert_status_entry(status, worker_url, model, role): + entry = status[worker_url] + assert entry['role'] == role + assert model in entry['models'] + assert entry['unfinished'] == 0 + assert entry['latency'] == [] + assert entry['speed'] is None + + +def _assert_pd_unsupported_endpoints(config, router_url): + generate = requests.post( + f'{router_url}/generate', + json={'prompt': 'Hello', 'max_tokens': 2}, + timeout=config.request_timeout, + ) + assert generate.status_code == 501 + + responses = requests.post( + f'{router_url}/v1/responses', + json={'model': config.model, 'input': 'Hello', 'max_output_tokens': 2}, + timeout=config.request_timeout, + ) + assert responses.status_code == 501 + + +def test_hybrid_dynamic_registration_is_idempotent_and_routable( + lmdeploy_config, + lmdeploy_router_factory, +): + router = lmdeploy_router_factory.start(worker_urls=()) + worker_url = lmdeploy_config.backend_urls[0] + try: + first = _add_node(router.url, worker_url, lmdeploy_config.model, role=1) + first.raise_for_status() + second = _add_node(router.url, worker_url, lmdeploy_config.model, role=1) + second.raise_for_status() + + status = requests.get(f'{router.url}/nodes/status', timeout=10) + status.raise_for_status() + _assert_status_entry(status.json(), worker_url, lmdeploy_config.model, role=1) + + workers = requests.get(f'{router.url}/list_workers', timeout=10) + workers.raise_for_status() + assert workers.json().get('urls') == [worker_url] + _wait_for_inference(lmdeploy_config, router.url) + + removed = requests.post( + f'{router.url}/nodes/remove', + json={'url': worker_url}, + timeout=10, + ) + removed.raise_for_status() + assert requests.get(f'{router.url}/nodes/status', timeout=10).json() == {} + + removed_again = requests.post( + f'{router.url}/nodes/remove', + json={'url': worker_url}, + timeout=10, + ) + removed_again.raise_for_status() + finally: + lmdeploy_router_factory.stop(router) + + +def test_regular_router_rejects_prefill_role( + lmdeploy_config, + lmdeploy_router_factory, +): + router = lmdeploy_router_factory.start(worker_urls=()) + try: + response = _add_node( + router.url, + lmdeploy_config.backend_urls[0], + lmdeploy_config.model, + role=2, + ) + assert response.status_code == 400 + finally: + lmdeploy_router_factory.stop(router) + + +def test_pd_dynamic_registration_and_inference( + lmdeploy_config, + lmdeploy_router_factory, +): + if not lmdeploy_config.pd_prefill_url or not lmdeploy_config.pd_decode_url: + pytest.skip( + 'set LMDEPLOY_PD_PREFILL_URL and LMDEPLOY_PD_DECODE_URL for real PD testing' + ) + + router = lmdeploy_router_factory.start(lmdeploy_pd=True) + try: + prefill = _add_node( + router.url, + lmdeploy_config.pd_prefill_url, + lmdeploy_config.model, + role=2, + ) + decode = _add_node( + router.url, + lmdeploy_config.pd_decode_url, + lmdeploy_config.model, + role=3, + ) + prefill.raise_for_status() + decode.raise_for_status() + + status = requests.get(f'{router.url}/nodes/status', timeout=10) + status.raise_for_status() + status_payload = status.json() + _assert_status_entry( + status_payload, + lmdeploy_config.pd_prefill_url, + lmdeploy_config.model, + role=2, + ) + _assert_status_entry( + status_payload, + lmdeploy_config.pd_decode_url, + lmdeploy_config.model, + role=3, + ) + + duplicate_prefill = _add_node( + router.url, + lmdeploy_config.pd_prefill_url, + lmdeploy_config.model, + role=2, + ) + duplicate_prefill.raise_for_status() + assert len(requests.get(f'{router.url}/nodes/status', timeout=10).json()) == 2 + + hybrid = _add_node( + router.url, + lmdeploy_config.backend_urls[0], + lmdeploy_config.model, + role=1, + ) + assert hybrid.status_code == 400 + + _wait_for_inference(lmdeploy_config, router.url, timeout=120) + _assert_pd_unsupported_endpoints(lmdeploy_config, router.url) + finally: + lmdeploy_router_factory.stop(router) + + +def test_pd_static_registration_and_inference( + lmdeploy_config, + lmdeploy_router_factory, +): + if not lmdeploy_config.pd_prefill_url or not lmdeploy_config.pd_decode_url: + pytest.skip( + 'set LMDEPLOY_PD_PREFILL_URL and LMDEPLOY_PD_DECODE_URL for real PD testing' + ) + + router = lmdeploy_router_factory.start( + lmdeploy_pd=True, + prefill_urls=(lmdeploy_config.pd_prefill_url,), + decode_urls=(lmdeploy_config.pd_decode_url,), + ) + try: + assert set( + requests.get(f'{router.url}/list_workers', timeout=10).json()['urls'] + ) == { + lmdeploy_config.pd_prefill_url, + lmdeploy_config.pd_decode_url, + } + _wait_for_inference(lmdeploy_config, router.url, timeout=120) + _assert_pd_unsupported_endpoints(lmdeploy_config, router.url) + finally: + lmdeploy_router_factory.stop(router) diff --git a/lmdeploy/router/py_test/integration/lmdeploy/test_rest_api.py b/lmdeploy/router/py_test/integration/lmdeploy/test_rest_api.py new file mode 100644 index 0000000000..8a35756602 --- /dev/null +++ b/lmdeploy/router/py_test/integration/lmdeploy/test_rest_api.py @@ -0,0 +1,456 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""LMDeploy REST compatibility tests routed through lmdeploy-router. + +The cases are adapted from LMDeploy's ``autotest/interface/restful`` suite. +They require two externally managed Hybrid LMDeploy API servers; see +``docs/lmdeploy_support_guide.md`` for the environment variables. +""" + +import json +import time + +import pytest +import requests + +pytestmark = [pytest.mark.integration, pytest.mark.lmdeploy] + + +def _post(config, base_url, path, payload, *, stream=False): + return requests.post( + f'{base_url}{path}', + json=payload, + stream=stream, + timeout=config.request_timeout, + ) + + +def _chat(config, base_url, payload, *, stream=False): + return _post(config, base_url, '/v1/chat/completions', payload, stream=stream) + + +def _completion(config, base_url, payload, *, stream=False): + return _post(config, base_url, '/v1/completions', payload, stream=stream) + + +def _generate(config, base_url, payload, *, stream=False): + return _post(config, base_url, '/generate', payload, stream=stream) + + +def _sse_payloads(response): + payloads = [] + for line in response.iter_lines(): + if not line: + continue + text = line.decode('utf-8').strip() + if text.startswith('data: ') and text != 'data: [DONE]': + payloads.append(json.loads(text[6:])) + return payloads + + +def _message_text(choice): + message = choice.get('message') or {} + return (message.get('reasoning_content') or '') + (message.get('content') or '') + + +def _assert_usage(usage): + assert usage.get('prompt_tokens', 0) > 0 + assert usage.get('completion_tokens', 0) > 0 + assert usage.get('total_tokens', 0) > 0 + + +def _assert_chat_response(output, model): + assert output.get('id') is not None + assert output.get('object') == 'chat.completion' + assert output.get('model') == model + _assert_usage(output.get('usage', {})) + choices = output.get('choices', []) + assert len(choices) == 1 + choice = choices[0] + assert choice.get('index') == 0 + assert choice.get('finish_reason') in {'stop', 'length'} + assert (choice.get('message') or {}).get('role') == 'assistant' + assert _message_text(choice) + + +def _assert_completion_response(output, model): + assert output.get('id') is not None + assert output.get('object') == 'text_completion' + assert output.get('model') == model + _assert_usage(output.get('usage', {})) + choices = output.get('choices', []) + assert len(choices) == 1 + assert choices[0].get('index') == 0 + assert choices[0].get('finish_reason') in {'stop', 'length'} + assert choices[0].get('text') + + +def _assert_generate_response(output): + assert isinstance(output.get('text'), str) and output['text'] + assert isinstance(output.get('output_ids'), list) and output['output_ids'] + assert isinstance(output.get('meta_info'), dict) + + +def _router_url(lmdeploy_router): + return lmdeploy_router.url + + +def test_models(lmdeploy_config, lmdeploy_router): + response = requests.get(f'{_router_url(lmdeploy_router)}/v1/models', timeout=10) + response.raise_for_status() + payload = response.json() + assert payload.get('object') == 'list' + assert any( + item.get('id') == lmdeploy_config.model for item in payload.get('data', []) + ) + + +def test_health(lmdeploy_router): + response = requests.get(f'{_router_url(lmdeploy_router)}/health', timeout=10) + response.raise_for_status() + assert 'healthy' in response.text.lower() + + +def test_list_workers(lmdeploy_config, lmdeploy_router): + response = requests.get(f'{_router_url(lmdeploy_router)}/list_workers', timeout=10) + response.raise_for_status() + assert set(response.json().get('urls', [])) == set(lmdeploy_config.backend_urls) + + +def test_chat_basic(lmdeploy_config, lmdeploy_router): + payload = { + 'model': lmdeploy_config.model, + 'messages': [{'role': 'user', 'content': '用一句话介绍你自己'}], + 'max_tokens': 30, + 'temperature': 0.01, + } + response = _chat(lmdeploy_config, _router_url(lmdeploy_router), payload) + response.raise_for_status() + _assert_chat_response(response.json(), lmdeploy_config.model) + + +def test_chat_streaming(lmdeploy_config, lmdeploy_router): + payload = { + 'model': lmdeploy_config.model, + 'messages': [{'role': 'user', 'content': '介绍你自己'}], + 'max_tokens': 30, + 'temperature': 0.01, + 'stream': True, + } + response = _chat( + lmdeploy_config, _router_url(lmdeploy_router), payload, stream=True + ) + response.raise_for_status() + chunks = _sse_payloads(response) + assert len(chunks) >= 2 + choices = [chunk['choices'][0] for chunk in chunks if chunk.get('choices')] + assert choices + assert choices[-1].get('finish_reason') in {'stop', 'length'} + assert all(choice.get('index') == 0 for choice in choices) + + +@pytest.mark.parametrize('stop', ['6', ['goodbye', 'farewell', 'bye']]) +def test_chat_stopwords(lmdeploy_config, lmdeploy_router, stop): + payload = { + 'model': lmdeploy_config.model, + 'messages': [{'role': 'user', 'content': 'Count to 10: 1, 2, 3, 4, 5, 6'}], + 'max_tokens': 100, + 'stop': stop, + 'temperature': 0.01, + } + response = _chat(lmdeploy_config, _router_url(lmdeploy_router), payload) + response.raise_for_status() + output = response.json() + _assert_chat_response(output, lmdeploy_config.model) + text = _message_text(output['choices'][0]) + stops = [stop] if isinstance(stop, str) else stop + assert all(item not in text for item in stops) + + +@pytest.mark.parametrize('field', ['max_tokens', 'max_completion_tokens']) +def test_chat_token_limit(lmdeploy_config, lmdeploy_router, field): + payload = { + 'model': lmdeploy_config.model, + 'messages': [{'role': 'user', 'content': '介绍人工智能'}], + field: 5, + 'temperature': 0.01, + } + response = _chat(lmdeploy_config, _router_url(lmdeploy_router), payload) + response.raise_for_status() + output = response.json() + _assert_chat_response(output, lmdeploy_config.model) + assert output['choices'][0]['finish_reason'] == 'length' + + +def test_chat_input_ids_token_in_out(lmdeploy_config, lmdeploy_router): + payload = { + 'model': lmdeploy_config.model, + 'messages': [], + 'input_ids': list(lmdeploy_config.token_ids), + 'max_tokens': 20, + 'temperature': 0.01, + 'return_token_ids': True, + } + response = _chat(lmdeploy_config, _router_url(lmdeploy_router), payload) + response.raise_for_status() + output = response.json() + _assert_chat_response(output, lmdeploy_config.model) + assert output['choices'][0].get('output_ids') + + +def test_chat_input_ids_matches_direct_backend(lmdeploy_config, lmdeploy_router): + payload = { + 'model': lmdeploy_config.model, + 'messages': [], + 'input_ids': list(lmdeploy_config.token_ids), + 'max_tokens': 15, + 'temperature': 0, + 'return_token_ids': True, + } + direct = [] + for backend_url in lmdeploy_config.backend_urls: + response = _chat(lmdeploy_config, backend_url, payload) + if response.status_code == 200: + direct.append(response.json()) + assert direct + + response = _chat(lmdeploy_config, _router_url(lmdeploy_router), payload) + response.raise_for_status() + routed = response.json()['choices'][0] + matches = [ + item['choices'][0] + for item in direct + if _message_text(item['choices'][0]) == _message_text(routed) + ] + assert matches + assert any(item.get('output_ids') == routed.get('output_ids') for item in matches) + + +def test_chat_temperature_zero_is_deterministic(lmdeploy_config, lmdeploy_router): + payload = { + 'model': lmdeploy_config.model, + 'messages': [{'role': 'user', 'content': 'What is 2+2? Answer in one word.'}], + 'max_tokens': 10, + 'temperature': 0, + } + first = _chat(lmdeploy_config, _router_url(lmdeploy_router), payload) + second = _chat(lmdeploy_config, _router_url(lmdeploy_router), payload) + first.raise_for_status() + second.raise_for_status() + assert _message_text(first.json()['choices'][0]) == _message_text( + second.json()['choices'][0] + ) + + +def test_chat_logprobs(lmdeploy_config, lmdeploy_router): + payload = { + 'model': lmdeploy_config.model, + 'messages': [{'role': 'user', 'content': 'Hello'}], + 'max_tokens': 5, + 'temperature': 0.01, + 'logprobs': True, + 'top_logprobs': 5, + } + response = _chat(lmdeploy_config, _router_url(lmdeploy_router), payload) + if response.status_code == 400: + pytest.skip('LMDeploy backend was not started with raw logprobs enabled') + response.raise_for_status() + content = response.json()['choices'][0].get('logprobs', {}).get('content') + assert content + + +def test_chat_invalid_model(lmdeploy_config, lmdeploy_router): + payload = { + 'model': 'model-that-does-not-exist', + 'messages': [{'role': 'user', 'content': 'Hello'}], + 'max_tokens': 5, + } + response = _chat(lmdeploy_config, _router_url(lmdeploy_router), payload) + assert response.status_code in {400, 404, 500} + + +def test_chat_matches_direct_backend(lmdeploy_config, lmdeploy_router): + payload = { + 'model': lmdeploy_config.model, + 'messages': [{'role': 'user', 'content': '什么是人工智能?用一句话回答'}], + 'max_tokens': 20, + 'temperature': 0, + } + direct_texts = { + _message_text(_chat(lmdeploy_config, url, payload).json()['choices'][0]) + for url in lmdeploy_config.backend_urls + } + response = _chat(lmdeploy_config, _router_url(lmdeploy_router), payload) + response.raise_for_status() + assert _message_text(response.json()['choices'][0]) in direct_texts + + +def test_completion_basic(lmdeploy_config, lmdeploy_router): + payload = { + 'model': lmdeploy_config.model, + 'prompt': 'Shanghai is', + 'max_tokens': 16, + 'temperature': 0.01, + } + response = _completion(lmdeploy_config, _router_url(lmdeploy_router), payload) + response.raise_for_status() + _assert_completion_response(response.json(), lmdeploy_config.model) + + +def test_completion_streaming(lmdeploy_config, lmdeploy_router): + payload = { + 'model': lmdeploy_config.model, + 'prompt': 'Shanghai is', + 'max_tokens': 16, + 'temperature': 0.01, + 'stream': True, + } + response = _completion( + lmdeploy_config, _router_url(lmdeploy_router), payload, stream=True + ) + response.raise_for_status() + chunks = _sse_payloads(response) + assert len(chunks) >= 2 + choices = [chunk['choices'][0] for chunk in chunks if chunk.get('choices')] + assert choices[-1].get('finish_reason') in {'stop', 'length'} + + +def test_completion_stopword(lmdeploy_config, lmdeploy_router): + payload = { + 'model': lmdeploy_config.model, + 'prompt': 'Shanghai is', + 'max_tokens': 200, + 'stop': ' city', + 'temperature': 0.01, + } + response = _completion(lmdeploy_config, _router_url(lmdeploy_router), payload) + response.raise_for_status() + assert ' city' not in response.json()['choices'][0]['text'] + + +def test_completion_max_tokens(lmdeploy_config, lmdeploy_router): + payload = { + 'model': lmdeploy_config.model, + 'prompt': '介绍人工智能', + 'max_tokens': 5, + 'temperature': 0.01, + } + response = _completion(lmdeploy_config, _router_url(lmdeploy_router), payload) + response.raise_for_status() + assert response.json()['choices'][0]['finish_reason'] == 'length' + + +def test_completion_matches_direct_backend(lmdeploy_config, lmdeploy_router): + payload = { + 'model': lmdeploy_config.model, + 'prompt': 'The capital of France is', + 'max_tokens': 5, + 'temperature': 0, + } + direct_texts = { + _completion(lmdeploy_config, url, payload).json()['choices'][0]['text'] + for url in lmdeploy_config.backend_urls + } + response = _completion(lmdeploy_config, _router_url(lmdeploy_router), payload) + response.raise_for_status() + assert response.json()['choices'][0]['text'] in direct_texts + + +def test_generate_basic(lmdeploy_config, lmdeploy_router): + payload = {'prompt': 'The sky is', 'max_tokens': 5, 'temperature': 0.01} + response = _generate(lmdeploy_config, _router_url(lmdeploy_router), payload) + response.raise_for_status() + _assert_generate_response(response.json()) + + +def test_generate_input_ids(lmdeploy_config, lmdeploy_router): + payload = { + 'input_ids': list(lmdeploy_config.token_ids), + 'max_tokens': 10, + 'temperature': 0.01, + } + response = _generate(lmdeploy_config, _router_url(lmdeploy_router), payload) + response.raise_for_status() + _assert_generate_response(response.json()) + + +def test_generate_streaming(lmdeploy_config, lmdeploy_router): + payload = { + 'prompt': 'Count to 10: 1, 2, 3,', + 'max_tokens': 8, + 'temperature': 0.01, + 'stream': True, + } + response = _generate( + lmdeploy_config, _router_url(lmdeploy_router), payload, stream=True + ) + response.raise_for_status() + chunks = _sse_payloads(response) + assert len(chunks) >= 2 + assert ''.join(chunk.get('text', '') for chunk in chunks).strip() + assert any(chunk.get('output_ids') for chunk in chunks) + + +def test_generate_stop_token_ids(lmdeploy_config, lmdeploy_router): + payload = { + 'prompt': 'Hello world', + 'max_tokens': 50, + 'stop_token_ids': [13], + 'temperature': 0.01, + } + response = _generate(lmdeploy_config, _router_url(lmdeploy_router), payload) + response.raise_for_status() + _assert_generate_response(response.json()) + + +def test_generate_session_id(lmdeploy_config, lmdeploy_router): + session_id = int(time.time()) % 100000 + first = _generate( + lmdeploy_config, + _router_url(lmdeploy_router), + {'prompt': 'First:', 'session_id': session_id, 'max_tokens': 5}, + ) + second = _generate( + lmdeploy_config, + _router_url(lmdeploy_router), + {'prompt': 'Second:', 'session_id': session_id, 'max_tokens': 5}, + ) + first.raise_for_status() + second.raise_for_status() + + +def test_generate_rejects_prompt_with_input_ids(lmdeploy_config, lmdeploy_router): + response = _generate( + lmdeploy_config, + _router_url(lmdeploy_router), + {'prompt': 'Hello', 'input_ids': [1, 2, 3], 'max_tokens': 5}, + ) + assert response.status_code in {400, 422} + + +def test_generate_rejects_empty_prompt(lmdeploy_config, lmdeploy_router): + response = _generate( + lmdeploy_config, + _router_url(lmdeploy_router), + {'prompt': '', 'max_tokens': 5}, + ) + assert response.status_code in {400, 422} + + +def test_generate_matches_direct_backend(lmdeploy_config, lmdeploy_router): + payload = { + 'prompt': 'The capital of France is', + 'max_tokens': 5, + 'temperature': 0, + } + direct = [] + for backend_url in lmdeploy_config.backend_urls: + response = _generate(lmdeploy_config, backend_url, payload) + response.raise_for_status() + direct.append(response.json()) + + response = _generate(lmdeploy_config, _router_url(lmdeploy_router), payload) + response.raise_for_status() + routed = response.json() + matches = [item for item in direct if item['text'] == routed['text']] + assert matches + assert any(item.get('output_ids') == routed.get('output_ids') for item in matches) diff --git a/lmdeploy/router/py_test/integration/load_balancing/__init__.py b/lmdeploy/router/py_test/integration/load_balancing/__init__.py new file mode 100644 index 0000000000..5becd69080 --- /dev/null +++ b/lmdeploy/router/py_test/integration/load_balancing/__init__.py @@ -0,0 +1,2 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""Load balancing integration tests.""" diff --git a/lmdeploy/router/py_test/integration/load_balancing/test_cache_aware.py b/lmdeploy/router/py_test/integration/load_balancing/test_cache_aware.py new file mode 100644 index 0000000000..88f067e7e1 --- /dev/null +++ b/lmdeploy/router/py_test/integration/load_balancing/test_cache_aware.py @@ -0,0 +1,74 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import collections +import concurrent.futures +import uuid + +import pytest +import requests + + +@pytest.mark.integration +def test_cache_aware_affinity(mock_workers, router_manager): + # Two workers; same prompt should stick to one due to cache tree + _, urls, ids = mock_workers(n=2) + rh = router_manager.start_router(worker_urls=urls, policy='cache_aware') + + counts = collections.Counter() + with requests.Session() as s: + for i in range(12): + r = s.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': 'repeated prompt for cache', + 'max_tokens': 1, + 'stream': False, + }, + ) + assert r.status_code == 200 + wid = r.headers.get('X-Worker-Id') or r.json().get('worker_id') + counts[wid] += 1 + + # Expect strong skew toward one worker (tree match); majority > 80% + top = max(counts.values()) + assert top >= 10, counts + + +@pytest.mark.integration +def test_cache_aware_diverse_prompts_balances(mock_workers, router_manager): + # Add latency so concurrent requests overlap and influence load-based selection + _, urls, ids = mock_workers(n=3, args=['--latency-ms', '30']) + rh = router_manager.start_router( + worker_urls=urls, + policy='cache_aware', + extra={ + 'cache_threshold': 0.99, + 'balance_abs_threshold': 0, + 'balance_rel_threshold': 1.0, + }, + ) + + counts = collections.Counter() + + def call(i): + # Use diverse, unrelated prompts to avoid prefix matches entirely + prompt = str(uuid.uuid4()) + r = requests.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': prompt, + 'max_tokens': 1, + 'stream': False, + }, + timeout=5, + ) + assert r.status_code == 200 + return r.headers.get('X-Worker-Id') or r.json().get('worker_id') + + with concurrent.futures.ThreadPoolExecutor(max_workers=16) as ex: + for wid in ex.map(call, range(40)): + counts[wid] += 1 + + # Expect participation of at least two workers + assert sum(1 for v in counts.values() if v > 0) >= 2, counts diff --git a/lmdeploy/router/py_test/integration/load_balancing/test_hash_policies.py b/lmdeploy/router/py_test/integration/load_balancing/test_hash_policies.py new file mode 100644 index 0000000000..c896e33ed8 --- /dev/null +++ b/lmdeploy/router/py_test/integration/load_balancing/test_hash_policies.py @@ -0,0 +1,44 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import collections + +import pytest +import requests + + +@pytest.mark.integration +@pytest.mark.parametrize('policy', ['consistent_hash', 'rendezvous_hash']) +def test_hash_policy_routes_sessions_consistently_and_distributes_sessions( + policy, mock_workers, router_manager +): + _, worker_urls, worker_ids = mock_workers(n=3) + router = router_manager.start_router( + worker_urls=worker_urls, + policy=policy, + ) + + selected = [] + for session_seq in range(30): + session_id = f'session-{session_seq}' + for _ in range(2): + response = requests.post( + f'{router.url}/v1/completions', + json={ + 'model': 'mock', + 'prompt': 'hello', + 'user': session_id, + }, + timeout=5, + ) + assert response.status_code == 200 + worker_id = response.headers.get('X-Worker-Id') + assert worker_id in worker_ids + selected.append((session_id, worker_id)) + + by_session = collections.defaultdict(set) + for session_id, worker_id in selected: + by_session[session_id].add(worker_id) + assert all( + len(worker_ids_for_session) == 1 + for worker_ids_for_session in by_session.values() + ) + assert len(set(frozenset(worker_ids) for worker_ids in by_session.values())) >= 2 diff --git a/lmdeploy/router/py_test/integration/load_balancing/test_power_of_two.py b/lmdeploy/router/py_test/integration/load_balancing/test_power_of_two.py new file mode 100644 index 0000000000..a55492ba41 --- /dev/null +++ b/lmdeploy/router/py_test/integration/load_balancing/test_power_of_two.py @@ -0,0 +1,91 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import collections +import concurrent.futures +import time + +import pytest +import requests + + +@pytest.mark.integration +def test_power_of_two_prefers_less_loaded(mock_workers, router_manager): + # Start two workers: one slow (higher inflight), one fast + # Router monitors /get_load and Power-of-Two uses cached loads to choose + # Start one slow and one fast worker using the fixture factory + procs_slow, urls_slow, ids_slow = mock_workers(n=1, args=['--latency-ms', '200']) + procs_fast, urls_fast, ids_fast = mock_workers(n=1, args=['--latency-ms', '0']) + _procs = procs_slow + procs_fast # Keep processes alive + urls = urls_slow + urls_fast + ids = ids_slow + ids_fast + slow_id = ids_slow[0] + slow_url = urls_slow[0] + + rh = router_manager.start_router( + worker_urls=urls, + policy='power_of_two', + extra={'worker_startup_check_interval': 1}, + ) + + # Prime: fire a burst to create measurable load on slow worker, then wait for monitor tick + + def _prime_call(i): + try: + requests.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': f'warm-{i}', + 'max_tokens': 1, + 'stream': False, + }, + timeout=5, + ) + except Exception: + pass + + with concurrent.futures.ThreadPoolExecutor(max_workers=32) as ex: + list(ex.map(_prime_call, range(128))) + time.sleep(2) + + # Apply direct background load on the slow worker to amplify load diff + def _direct_load(i): + try: + requests.post( + f'{slow_url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': f'bg-{i}', + 'max_tokens': 1, + 'stream': False, + }, + timeout=5, + ) + except Exception: + pass + + with concurrent.futures.ThreadPoolExecutor(max_workers=32) as ex: + list(ex.map(_direct_load, range(128))) + time.sleep(1) + + def call(i): + r = requests.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': f'p{i}', + 'max_tokens': 1, + 'stream': False, + }, + timeout=5, + ) + assert r.status_code == 200 + return r.headers.get('X-Worker-Id') or r.json().get('worker_id') + + counts = collections.Counter() + with concurrent.futures.ThreadPoolExecutor(max_workers=32) as ex: + for wid in ex.map(call, range(200)): + counts[wid] += 1 + + # Expect the slow worker (higher latency/inflight) to receive fewer requests + fast_worker_id = [i for i in ids if i != slow_id][0] + assert counts[slow_id] < counts[fast_worker_id], counts diff --git a/lmdeploy/router/py_test/integration/load_balancing/test_random.py b/lmdeploy/router/py_test/integration/load_balancing/test_random.py new file mode 100644 index 0000000000..0bd27c289e --- /dev/null +++ b/lmdeploy/router/py_test/integration/load_balancing/test_random.py @@ -0,0 +1,33 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import collections + +import pytest +import requests + + +@pytest.mark.integration +def test_random_distribution(mock_workers, router_manager): + procs, urls, ids = mock_workers(n=4) + rh = router_manager.start_router(worker_urls=urls, policy='random') + + counts = collections.Counter() + N = 200 + with requests.Session() as s: + for i in range(N): + r = s.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': f'p{i}', + 'max_tokens': 1, + 'stream': False, + }, + ) + assert r.status_code == 200 + wid = r.headers.get('X-Worker-Id') or r.json().get('worker_id') + counts[wid] += 1 + + # simple statistical tolerance: each worker should be within ±50% of mean + mean = N / len(ids) + for wid in ids: + assert 0.5 * mean <= counts[wid] <= 1.5 * mean, counts diff --git a/lmdeploy/router/py_test/integration/load_balancing/test_round_robin.py b/lmdeploy/router/py_test/integration/load_balancing/test_round_robin.py new file mode 100644 index 0000000000..5e159fba6d --- /dev/null +++ b/lmdeploy/router/py_test/integration/load_balancing/test_round_robin.py @@ -0,0 +1,34 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import collections + +import pytest +import requests + + +@pytest.mark.integration +def test_round_robin_distribution(mock_workers, router_manager): + procs, urls, ids = mock_workers(n=3) + + rh = router_manager.start_router(worker_urls=urls, policy='round_robin') + + counts = collections.Counter() + with requests.Session() as s: + for i in range(30): + r = s.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': f'hello {i}', + 'max_tokens': 1, + 'stream': False, + }, + ) + assert r.status_code == 200 + wid = r.headers.get('X-Worker-Id') or r.json().get('worker_id') + assert wid in ids + counts[wid] += 1 + + # Expect near-even distribution across 3 workers + # 30 requests -> ideally 10 each; allow small tolerance ±3 + for wid in ids: + assert 7 <= counts[wid] <= 13, counts diff --git a/lmdeploy/router/py_test/integration/test_api_auth.py b/lmdeploy/router/py_test/integration/test_api_auth.py new file mode 100644 index 0000000000..aa18e276be --- /dev/null +++ b/lmdeploy/router/py_test/integration/test_api_auth.py @@ -0,0 +1,39 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import pytest +import requests + + +@pytest.mark.integration +def test_router_api_key_enforcement(router_manager, mock_workers): + # Start backend requiring API key; router should forward Authorization header transparently + _, urls, _ = mock_workers( + n=1, args=['--require-api-key', '--api-key', 'correct_api_key'] + ) + rh = router_manager.start_router( + worker_urls=urls, + policy='round_robin', + extra={}, + ) + + # No auth -> 401 + r = requests.post( + f'{rh.url}/v1/completions', + json={'model': 'test-model', 'prompt': 'x', 'max_tokens': 1, 'stream': False}, + ) + assert r.status_code == 401 + + # Invalid auth -> 401 + r = requests.post( + f'{rh.url}/v1/completions', + json={'model': 'test-model', 'prompt': 'x', 'max_tokens': 1, 'stream': False}, + headers={'Authorization': 'Bearer wrong'}, + ) + assert r.status_code == 401 + + # Correct auth -> 200 + r = requests.post( + f'{rh.url}/v1/completions', + json={'model': 'test-model', 'prompt': 'x', 'max_tokens': 1, 'stream': False}, + headers={'Authorization': 'Bearer correct_api_key'}, + ) + assert r.status_code == 200 diff --git a/lmdeploy/router/py_test/integration/test_circuit_breaker.py b/lmdeploy/router/py_test/integration/test_circuit_breaker.py new file mode 100644 index 0000000000..23d61a2c27 --- /dev/null +++ b/lmdeploy/router/py_test/integration/test_circuit_breaker.py @@ -0,0 +1,236 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import time + +import pytest +import requests + + +@pytest.mark.integration +def test_circuit_breaker_opens_and_recovers(router_manager, mock_workers): + # A single worker that fails first 3 requests, then succeeds + _, [wurl], _ = mock_workers(n=1, args=['--fail-first-n', '3']) # fails first 3 + rh = router_manager.start_router( + worker_urls=[wurl], + policy='round_robin', + extra={ + 'cb_failure_threshold': 3, + 'cb_success_threshold': 2, + 'cb_timeout_duration_secs': 3, + 'cb_window_duration_secs': 10, + 'disable_retries': True, + }, + ) + + def post_once(): + return requests.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': 'trigger', + 'max_tokens': 1, + 'stream': False, + }, + timeout=3, + ) + + saw_503 = False + for _ in range(8): + r = post_once() + if r.status_code == 503: + saw_503 = True + break + assert saw_503, 'circuit breaker did not open to return 503' + + time.sleep(4) + r1 = post_once() + r2 = post_once() + assert r1.status_code == 200 and r2.status_code == 200 + + +@pytest.mark.integration +def test_circuit_breaker_half_open_failure_reopens(router_manager, mock_workers): + _, [wurl], _ = mock_workers(n=1, args=['--status-code', '500']) # always fail + rh = router_manager.start_router( + worker_urls=[wurl], + policy='round_robin', + extra={ + 'cb_failure_threshold': 2, + 'cb_success_threshold': 2, + 'cb_timeout_duration_secs': 2, + 'cb_window_duration_secs': 5, + 'disable_retries': True, + }, + ) + + def post_once(): + return requests.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': 'x', + 'max_tokens': 1, + 'stream': False, + }, + timeout=3, + ) + + opened = False + for _ in range(8): + r = post_once() + if r.status_code == 503: + opened = True + break + assert opened, 'circuit breaker did not open' + + time.sleep(3) + r = post_once() + assert r.status_code == 500 + r2 = post_once() + assert r2.status_code == 503 + + +@pytest.mark.integration +def test_circuit_breaker_disable_flag(router_manager, mock_workers): + _, [wurl], _ = mock_workers(n=1, args=['--status-code', '500']) # always fail + rh = router_manager.start_router( + worker_urls=[wurl], + policy='round_robin', + extra={ + 'disable_circuit_breaker': True, + 'disable_retries': True, + }, + ) + r = requests.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': 'x', + 'max_tokens': 1, + 'stream': False, + }, + timeout=3, + ) + assert r.status_code == 500 + + +@pytest.mark.integration +def test_circuit_breaker_per_worker_isolation(router_manager, mock_workers): + _, [fail_url], _ = mock_workers(n=1, args=['--status-code', '500']) # always fail + _, [ok_url], _ = mock_workers(n=1) + rh = router_manager.start_router( + worker_urls=[fail_url, ok_url], + policy='round_robin', + extra={ + 'cb_failure_threshold': 2, + 'cb_success_threshold': 1, + 'cb_timeout_duration_secs': 2, + 'cb_window_duration_secs': 10, + 'disable_retries': True, + }, + ) + + def post_once(): + return requests.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': 'y', + 'max_tokens': 1, + 'stream': False, + }, + timeout=3, + ) + + failures = 0 + successes_after_open = 0 + opened = False + for _ in range(30): + r = post_once() + if not opened: + if r.status_code == 500: + failures += 1 + if failures >= 2: + _ = post_once() + _ = post_once() + opened = True + else: + if r.status_code == 200: + successes_after_open += 1 + else: + assert False, f'Unexpected non-200 after CB open: {r.status_code}' + assert opened and successes_after_open >= 5 + + +@pytest.mark.integration +def test_circuit_breaker_with_retries(router_manager, mock_workers): + _, [fail_url], _ = mock_workers(n=1, args=['--status-code', '500']) # always fail + _, [ok_url], _ = mock_workers(n=1) + rh = router_manager.start_router( + worker_urls=[fail_url, ok_url], + policy='round_robin', + extra={ + 'retry_max_retries': 3, + 'retry_initial_backoff_ms': 10, + 'retry_max_backoff_ms': 50, + 'cb_failure_threshold': 2, + 'cb_success_threshold': 1, + 'cb_timeout_duration_secs': 2, + 'cb_window_duration_secs': 10, + }, + ) + + r = requests.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': 'z', + 'max_tokens': 1, + 'stream': False, + }, + timeout=5, + ) + assert r.status_code == 200 + + +@pytest.mark.integration +def test_circuit_breaker_ignores_4xx_client_errors(router_manager, mock_workers): + """4xx client errors (e.g., 400 Bad Request) should NOT count as circuit + breaker failures. Only 5xx server errors indicate worker health issues. + + This test ensures that many 400 responses do not trip the circuit breaker, while 500 responses still do. + """ + # Worker that returns 400 Bad Request for all requests + _, [client_error_url], _ = mock_workers(n=1, args=['--status-code', '400']) + rh = router_manager.start_router( + worker_urls=[client_error_url], + policy='round_robin', + extra={ + 'cb_failure_threshold': 2, # Low threshold - would trip quickly if 400s counted + 'cb_success_threshold': 1, + 'cb_timeout_duration_secs': 2, + 'cb_window_duration_secs': 10, + 'disable_retries': True, + }, + ) + + def post_once(): + return requests.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': 'test', + 'max_tokens': 1, + 'stream': False, + }, + timeout=3, + ) + + # Send many requests that return 400 + # If 4xx counted as failures, circuit would open after 2 requests + for i in range(10): + r = post_once() + # Should get 400 from worker, NOT 503 from circuit breaker + assert r.status_code == 400, ( + f'Request {i+1}: Expected 400 (client error passthrough), ' + f'got {r.status_code}. Circuit breaker incorrectly opened on 4xx errors.' + ) diff --git a/lmdeploy/router/py_test/integration/test_fault_tolerance.py b/lmdeploy/router/py_test/integration/test_fault_tolerance.py new file mode 100644 index 0000000000..8469dc3587 --- /dev/null +++ b/lmdeploy/router/py_test/integration/test_fault_tolerance.py @@ -0,0 +1,33 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import pytest +import requests + + +@pytest.mark.integration +def test_worker_crash_reroute_with_retries(router_manager, mock_workers): + # Start one healthy and one that will crash on first request + _, [ok_url], _ = mock_workers(n=1) + _, [crash_url], _ = mock_workers(n=1, args=['--crash-on-request']) + rh = router_manager.start_router( + worker_urls=[crash_url, ok_url], + policy='round_robin', + extra={ + 'retry_max_retries': 3, + 'retry_initial_backoff_ms': 10, + 'retry_max_backoff_ms': 50, + }, + ) + + # A single request should succeed via retry to the healthy worker + r = requests.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': 'crash', + 'max_tokens': 1, + 'stream': False, + }, + timeout=5, + ) + assert r.status_code == 200 + # mock_workers fixture handles cleanup diff --git a/lmdeploy/router/py_test/integration/test_payload_size.py b/lmdeploy/router/py_test/integration/test_payload_size.py new file mode 100644 index 0000000000..36c84a6903 --- /dev/null +++ b/lmdeploy/router/py_test/integration/test_payload_size.py @@ -0,0 +1,34 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import pytest +import requests + + +@pytest.mark.integration +def test_payload_size_limit(router_manager, mock_workers): + # Start one backend and a router with a 1MB payload limit + _, urls, _ = mock_workers(n=1) + rh = router_manager.start_router( + worker_urls=urls, + policy='round_robin', + extra={'max_payload_size': 1 * 1024 * 1024}, # 1MB + ) + + # Payload just under 1MB should succeed + payload_small = { + 'model': 'test-model', + 'prompt': 'x' * int(0.5 * 1024 * 1024), # ~0.5MB + 'max_tokens': 1, + 'stream': False, + } + r = requests.post(f'{rh.url}/v1/completions', json=payload_small) + assert r.status_code == 200 + + # Payload over 1MB should fail with 413 + payload_large = { + 'model': 'test-model', + 'prompt': 'x' * int(1.2 * 1024 * 1024), # ~1.2MB + 'max_tokens': 1, + 'stream': False, + } + r = requests.post(f'{rh.url}/v1/completions', json=payload_large) + assert r.status_code == 413 diff --git a/lmdeploy/router/py_test/integration/test_pd_registration.py b/lmdeploy/router/py_test/integration/test_pd_registration.py new file mode 100644 index 0000000000..c897a23666 --- /dev/null +++ b/lmdeploy/router/py_test/integration/test_pd_registration.py @@ -0,0 +1,50 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import pytest +import requests + + +@pytest.mark.integration +def test_pd_dynamic_registration_and_request_routing(mock_workers, router_manager): + _, prefill_urls, _ = mock_workers(n=1) + _, decode_urls, decode_ids = mock_workers(n=2) + router = router_manager.start_router( + policy='round_robin', + lmdeploy_pd_disaggregation=True, + prefill_urls=[], + decode_urls=[], + ) + + router_manager.add_lmdeploy_node(router.url, prefill_urls[0], role=2) + for decode_url in decode_urls: + router_manager.add_lmdeploy_node(router.url, decode_url, role=3) + + nodes = router_manager.lmdeploy_nodes(router.url) + assert nodes[prefill_urls[0]]['role'] == 2 + assert all(nodes[url]['role'] == 3 for url in decode_urls) + + selected_decode_ids = set() + for request_seq in range(4): + response = requests.post( + f'{router.url}/v1/completions', + json={ + 'model': 'mock', + 'prompt': f'hello-{request_seq}', + 'max_tokens': 2, + }, + timeout=10, + ) + assert response.status_code == 200 + worker_id = response.headers.get('X-Worker-Id') + assert worker_id in decode_ids + selected_decode_ids.add(worker_id) + + assert len(selected_decode_ids) == 2 + + router_manager.remove_lmdeploy_node(router.url, decode_urls[0], role=3) + response = requests.post( + f'{router.url}/v1/completions', + json={'model': 'mock', 'prompt': 'after-removal', 'max_tokens': 1}, + timeout=10, + ) + assert response.status_code == 200 + assert response.headers.get('X-Worker-Id') == decode_ids[1] diff --git a/lmdeploy/router/py_test/integration/test_pd_routing.py b/lmdeploy/router/py_test/integration/test_pd_routing.py new file mode 100644 index 0000000000..6e5ff7f663 --- /dev/null +++ b/lmdeploy/router/py_test/integration/test_pd_routing.py @@ -0,0 +1,157 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import collections +import concurrent.futures +import time + +import pytest +import requests + + +@pytest.mark.integration +@pytest.mark.pd_engine +def test_pd_power_of_two_decode_attribution(router_manager, mock_workers): + # Start two prefill and three decode mock workers via fixture + _, prefill_urls_raw, prefill_ids = mock_workers(n=2) + _, decode_urls_raw, decode_ids_list = mock_workers(n=3) + prefill_urls = list(prefill_urls_raw) + decode_urls = list(decode_urls_raw) + decode_ids = set(decode_ids_list) + + rh = router_manager.start_router( + policy='power_of_two', + lmdeploy_pd_disaggregation=True, + prefill_urls=prefill_urls, + decode_urls=decode_urls, + extra={'worker_startup_check_interval': 1}, + ) + + counts = collections.Counter() + with requests.Session() as s: + for i in range(30): + r = s.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': f'p{i}', + 'max_tokens': 1, + 'stream': False, + }, + ) + assert r.status_code == 200 + wid = r.headers.get('X-Worker-Id') or r.json().get('worker_id') + assert wid in decode_ids + counts[wid] += 1 + + assert sum(1 for v in counts.values() if v > 0) >= 2 + + +@pytest.mark.integration +@pytest.mark.pd_engine +def test_pd_power_of_two_skews_to_faster_decode(router_manager, mock_workers): + # Start two prefill workers (fast) + _, prefill_urls_raw, _ = mock_workers(n=2) + + # Start two decode workers: one very slow, one fast + # Use 2000ms latency to create a clear load difference + _, [decode_slow_url], [slow_id] = mock_workers( + n=1, args=['--latency-ms', '2000'] + ) # 2 second latency - very slow + _, [decode_fast_url], [fast_id] = mock_workers(n=1) + decode_urls_raw = [decode_slow_url, decode_fast_url] + + prefill_urls = list(prefill_urls_raw) + decode_urls = list(decode_urls_raw) + + rh = router_manager.start_router( + policy='power_of_two', + lmdeploy_pd_disaggregation=True, + prefill_urls=prefill_urls, + decode_urls=decode_urls, + extra={'worker_startup_check_interval': 1}, + ) + + # Prime the router with some initial requests + def _prime_call(i): + try: + requests.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': f'warm-{i}', + 'max_tokens': 1, + 'stream': False, + }, + timeout=10, + ) + except Exception: + pass + + with concurrent.futures.ThreadPoolExecutor(max_workers=16) as ex: + list(ex.map(_prime_call, range(50))) + time.sleep(1) + + # Create sustained load on slow worker by sending requests in background + # These will take 2 seconds each, keeping the slow worker loaded + stop_background_load = False + + def _background_load(): + i = 0 + while not stop_background_load: + try: + requests.post( + f'{decode_slow_url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': f'bg-{i}', + 'max_tokens': 1, + 'stream': False, + }, + timeout=10, + ) + i += 1 + except Exception: + pass + + # Start background load in separate threads + import threading + + bg_threads = [ + threading.Thread(target=_background_load, daemon=True) for _ in range(4) + ] + for t in bg_threads: + t.start() + + # Wait for slow worker to accumulate load + time.sleep(3) + + # Now send test requests - power-of-two should strongly prefer the fast worker + def call(i): + r = requests.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': f'p{i}', + 'max_tokens': 1, + 'stream': False, + }, + timeout=10, + ) + assert r.status_code == 200 + return r.headers.get('X-Worker-Id') or r.json().get('worker_id') + + counts = collections.Counter() + with concurrent.futures.ThreadPoolExecutor(max_workers=32) as ex: + for wid in ex.map(call, range(200)): + counts[wid] += 1 + + # Stop background load + stop_background_load = True + + # With sustained load on slow worker, fast worker should get significantly more requests + # Allow for some variance, but expect at least 60/40 split favoring fast worker + assert ( + counts[fast_id] > counts[slow_id] + ), f'Expected fast worker to handle more requests, got {counts}' + assert ( + counts[fast_id] >= counts[slow_id] * 1.2 + ), f'Expected fast worker to handle at least 20% more requests than slow worker, got {counts}' diff --git a/lmdeploy/router/py_test/integration/test_rate_limiting.py b/lmdeploy/router/py_test/integration/test_rate_limiting.py new file mode 100644 index 0000000000..51bdf056bd --- /dev/null +++ b/lmdeploy/router/py_test/integration/test_rate_limiting.py @@ -0,0 +1,92 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import concurrent.futures + +import pytest +import requests + + +@pytest.mark.integration +def test_rate_limit_and_queue(router_manager, mock_workers): + # Add latency to ensure requests overlap and trigger rate limiting + # Without latency, requests may complete too fast to trigger concurrency limits + _, urls, _ = mock_workers(n=1, args=['--latency-ms', '100']) + rh = router_manager.start_router( + worker_urls=urls, + policy='round_robin', + extra={ + 'max_concurrent_requests': 2, + 'queue_size': 0, # no queue -> immediate 429 when limit exceeded + }, + ) + + def call_once(i): + try: + r = requests.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': f'p{i}', + 'max_tokens': 1, + 'stream': False, + }, + timeout=3, + ) + return r.status_code + except Exception: + return 599 + + # Fire a burst of concurrent requests + with concurrent.futures.ThreadPoolExecutor(max_workers=16) as ex: + results = list(ex.map(call_once, range(16))) + + # Expect some to succeed and some to be rate limited (429) + assert any(code == 200 for code in results) + assert any(code == 429 for code in results) + + +@pytest.mark.integration +def test_rate_limit_queue_and_timeout(router_manager, mock_workers): + # Slow backend: ~2s per request ensures queue wait > timeout + _, urls, _ = mock_workers(n=1, args=['--latency-ms', '2000']) # 2.0s per request + + # Allow 1 concurrent, queue up to 1, with 1s queue timeout + rh = router_manager.start_router( + worker_urls=urls, + policy='round_robin', + extra={ + 'max_concurrent_requests': 1, + 'queue_size': 1, + 'queue_timeout_secs': 1, + }, + ) + + def call_once(i): + try: + r = requests.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': f'q{i}', + 'max_tokens': 1, + 'stream': False, + }, + timeout=5, + ) + return r.status_code + except Exception: + return 599 + + # Fire 4 concurrent requests: 1 runs (~2s), 1 queued (times out at 1s -> 408), 2 overflow -> 429 + import concurrent.futures + + with concurrent.futures.ThreadPoolExecutor(max_workers=8) as ex: + results = list(ex.map(call_once, range(4))) + + # We expect: + # - Some 200s (processed) + # - At least one 408 (queued too long and timed out) + # - Remaining non-200s are either 429 (queue overflow) or additional 408s depending on scheduling + assert any(code == 200 for code in results) + assert any(code == 408 for code in results), results + non200 = [c for c in results if c != 200] + assert len(non200) >= 2 and all(c in (408, 429) for c in non200), results diff --git a/lmdeploy/router/py_test/integration/test_rest_api.py b/lmdeploy/router/py_test/integration/test_rest_api.py new file mode 100644 index 0000000000..84230d02eb --- /dev/null +++ b/lmdeploy/router/py_test/integration/test_rest_api.py @@ -0,0 +1,46 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import pytest +import requests + + +@pytest.mark.integration +def test_health_models_and_inference_endpoints(mock_workers, router_manager): + _, worker_urls, _ = mock_workers(n=2) + router = router_manager.start_router( + worker_urls=worker_urls, + policy='round_robin', + ) + + health = requests.get(f'{router.url}/health', timeout=5) + assert health.status_code == 200 + assert health.text == 'All servers healthy' + + models = requests.get(f'{router.url}/v1/models', timeout=5) + assert models.status_code == 200 + assert models.json()['data'][0]['id'] == 'mock' + + completion = requests.post( + f'{router.url}/v1/completions', + json={ + 'model': 'mock', + 'prompt': 'hello', + 'max_tokens': 1, + 'stream': False, + }, + timeout=5, + ) + assert completion.status_code == 200 + assert completion.json()['choices'][0]['text'] == 'ok' + + chat_completion = requests.post( + f'{router.url}/v1/chat/completions', + json={ + 'model': 'mock', + 'messages': [{'role': 'user', 'content': 'hello'}], + 'max_tokens': 1, + 'stream': False, + }, + timeout=5, + ) + assert chat_completion.status_code == 200 + assert chat_completion.json()['choices'][0]['text'] == 'ok' diff --git a/lmdeploy/router/py_test/integration/test_retries.py b/lmdeploy/router/py_test/integration/test_retries.py new file mode 100644 index 0000000000..4d4bc3c5fb --- /dev/null +++ b/lmdeploy/router/py_test/integration/test_retries.py @@ -0,0 +1,62 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import pytest +import requests + + +@pytest.mark.integration +def test_retry_reroutes_to_healthy_worker(router_manager, mock_workers): + # Worker A always 500; Worker B healthy + # Worker A always 500; Worker B/C healthy + _, [url_a], [id_a] = mock_workers(n=1, args=['--status-code', '500']) # fail + _, [url_b], [id_b] = mock_workers(n=1) + _, [url_c], [id_c] = mock_workers(n=1) + rh = router_manager.start_router( + worker_urls=[url_a, url_b, url_c], + policy='round_robin', + extra={ + 'retry_max_retries': 3, + 'retry_initial_backoff_ms': 10, + 'retry_max_backoff_ms': 50, + }, + ) + + r = requests.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': 'x', + 'max_tokens': 1, + 'stream': False, + }, + timeout=5, + ) + assert r.status_code == 200 + wid = r.headers.get('X-Worker-Id') or r.json().get('worker_id') + assert wid in [id_b, id_c] # should have retried onto a healthy worker (B or C) + # mock_workers fixture handles cleanup + + +@pytest.mark.integration +def test_disable_retries_surfaces_failure(router_manager, mock_workers): + # Single failing worker, retries disabled -> should return 500 + _, [url], [wid] = mock_workers(n=1, args=['--status-code', '500']) # always fail + rh = router_manager.start_router( + worker_urls=[url], + policy='round_robin', + extra={ + 'disable_retries': True, + }, + ) + + r = requests.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': 'x', + 'max_tokens': 1, + 'stream': False, + }, + timeout=5, + ) + assert r.status_code == 500 + # mock_workers fixture handles cleanup diff --git a/lmdeploy/router/py_test/integration/test_service_discovery_shim.py b/lmdeploy/router/py_test/integration/test_service_discovery_shim.py new file mode 100644 index 0000000000..2a77544ed1 --- /dev/null +++ b/lmdeploy/router/py_test/integration/test_service_discovery_shim.py @@ -0,0 +1,37 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import pytest +import requests + + +@pytest.mark.integration +def test_discovery_shim_add_remove(router_manager, mock_workers): + # Start router without workers + rh = router_manager.start_router(worker_urls=[], policy='round_robin') + + # Initially empty + urls = router_manager.list_workers(rh.url) + assert urls == [] + + # Add a worker (simulate discovery event) + _, [wurl], [wid] = mock_workers(n=1) + router_manager.add_worker(rh.url, wurl) + urls = router_manager.list_workers(rh.url) + assert wurl in urls + + # Can serve a request + r = requests.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': 'hi', + 'max_tokens': 1, + 'stream': False, + }, + ) + assert r.status_code == 200 + + # Remove worker (simulate pod deletion) + router_manager.remove_worker(rh.url, wurl) + urls = router_manager.list_workers(rh.url) + assert wurl not in urls + # mock_workers fixture handles cleanup diff --git a/lmdeploy/router/py_test/integration/test_worker_management.py b/lmdeploy/router/py_test/integration/test_worker_management.py new file mode 100644 index 0000000000..181ae30bbb --- /dev/null +++ b/lmdeploy/router/py_test/integration/test_worker_management.py @@ -0,0 +1,58 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import pytest +import requests + + +@pytest.mark.integration +def test_add_and_remove_worker(mock_worker, router_manager, mock_workers): + # Start with a single worker + proc1, url1, id1 = mock_worker + rh = router_manager.start_router(worker_urls=[url1], policy='round_robin') + + # Add a second worker + + procs2, urls2, ids2 = mock_workers(n=1) + url2 = urls2[0] + id2 = ids2[0] + router_manager.add_worker(rh.url, url2) + + # Send some requests and ensure both workers are seen + seen = set() + with requests.Session() as s: + for i in range(20): + r = s.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': f'x{i}', + 'max_tokens': 1, + 'stream': False, + }, + ) + assert r.status_code == 200 + wid = r.headers.get('X-Worker-Id') or r.json().get('worker_id') + seen.add(wid) + if len(seen) == 2: + break + + assert id1 in seen and id2 in seen + + # Now remove the second worker + router_manager.remove_worker(rh.url, url2) + + # After removal, subsequent requests should only come from first worker + with requests.Session() as s: + for i in range(10): + r = s.post( + f'{rh.url}/v1/completions', + json={ + 'model': 'test-model', + 'prompt': f'y{i}', + 'max_tokens': 1, + 'stream': False, + }, + ) + assert r.status_code == 200 + wid = r.headers.get('X-Worker-Id') or r.json().get('worker_id') + assert wid == id1 + # mock_workers fixture handles cleanup diff --git a/lmdeploy/router/py_test/test_consistent_hash_policy.py b/lmdeploy/router/py_test/test_consistent_hash_policy.py new file mode 100755 index 0000000000..9e44776167 --- /dev/null +++ b/lmdeploy/router/py_test/test_consistent_hash_policy.py @@ -0,0 +1,500 @@ +# Copyright (c) OpenMMLab. All rights reserved. +#!/usr/bin/env python3 +"""Test script for the new Consistent Hash load balancing policy. + +This demonstrates how the consistent hash policy routes requests based on session_id or user_id, ensuring that requests +from the same user/session are consistently routed to the same worker for better cache locality and stateful processing. +""" + +import argparse +import random +import string +import time + +import requests + + +class ConsistentHashTester: + def __init__(self, router_url: str = 'http://localhost:30000'): + self.router_url = router_url.rstrip('/') + self.session = requests.Session() + + def log(self, message: str, level: str = 'INFO'): + """Log a message with timestamp.""" + timestamp = time.strftime('%H:%M:%S') + print(f'[{timestamp}] [{level}] {message}') + + def generate_session_id(self) -> str: + """Generate a random session ID.""" + return ''.join(random.choices(string.ascii_letters + string.digits, k=12)) + + def generate_user_id(self) -> str: + """Generate a random user ID.""" + return f'user_{random.randint(1000, 9999)}' + + def make_request( + self, prompt: str, session_id: str = None, user_id: str = None + ) -> tuple[bool, str]: + """Make a request to the router and return (success, response_text). + + Returns the full response for analysis. + """ + # Build request with session_id or user_id in session_params + request_data = { + 'text': prompt, + 'sampling_params': {'temperature': 0.7, 'max_new_tokens': 50}, + } + + # Add session_id to session_params OR user at top level (proper locations per OpenAI spec) + if session_id: + request_data['session_params'] = {'session_id': session_id} + elif user_id: + request_data['user'] = user_id # Top-level user field (OpenAI standard) + + try: + response = self.session.post( + f'{self.router_url}/generate', json=request_data, timeout=30 + ) + + if response.status_code == 200: + return True, response.text + else: + return False, f'HTTP {response.status_code}: {response.text}' + + except Exception as e: + return False, str(e) + + def test_session_consistency(self, num_requests: int = 10) -> bool: + """Test that requests with the same session_id always go to the same + worker.""" + self.log('Testing session consistency...') + + session_id = self.generate_session_id() + worker_responses = [] + + for i in range(num_requests): + prompt = f'Request {i+1} for session {session_id}' + success, response = self.make_request(prompt, session_id=session_id) + + if success: + worker_responses.append(response) + self.log(f' Request {i+1}: ✅') + else: + self.log(f' Request {i+1}: ❌ {response}', 'ERROR') + return False + + # Extract worker information from responses (this would depend on your router's response format) + # For now, we'll assume consistency if all requests succeed + self.log( + f'✅ Session consistency test passed: {num_requests} requests with session_id ' + f"'{session_id}' completed successfully" + ) + return True + + def test_user_consistency(self, num_requests: int = 10) -> bool: + """Test that requests with the same user_id always go to the same + worker.""" + self.log('Testing user consistency...') + + user_id = self.generate_user_id() + worker_responses = [] + + for i in range(num_requests): + prompt = f'Request {i+1} for user {user_id}' + success, response = self.make_request(prompt, user_id=user_id) + + if success: + worker_responses.append(response) + self.log(f' Request {i+1}: ✅') + else: + self.log(f' Request {i+1}: ❌ {response}', 'ERROR') + return False + + self.log( + f"✅ User consistency test passed: {num_requests} requests with user_id '{user_id}' completed successfully" + ) + return True + + def test_session_priority_over_user(self) -> bool: + """Test that session_id takes priority over user_id for routing.""" + self.log('Testing session_id priority over user_id...') + + session_id = self.generate_session_id() + user_id = self.generate_user_id() + + # Make requests with both session_id and user_id in session_params + responses = [] + for i in range(5): + prompt = f'Priority test request {i+1}' + request_data = { + 'text': prompt, + 'sampling_params': {'temperature': 0.7, 'max_new_tokens': 30}, + 'session_params': { + 'session_id': session_id, + 'user_id': user_id, # This should be ignored in favor of session_id + }, + } + + try: + response = self.session.post( + f'{self.router_url}/generate', json=request_data + ) + if response.status_code == 200: + responses.append(response.text) + self.log(f' Priority test request {i+1}: ✅') + else: + self.log( + f' Priority test request {i+1}: ❌ HTTP {response.status_code}', + 'ERROR', + ) + return False + except Exception as e: + self.log(f' Priority test request {i+1}: ❌ {e}', 'ERROR') + return False + + self.log( + '✅ Session priority test passed: All requests with both session_id and user_id completed successfully' + ) + return True + + def test_distribution_across_workers(self, num_sessions: int = 20) -> bool: + """Test that different sessions are distributed across multiple + workers.""" + self.log( + f'Testing distribution across workers with {num_sessions} different sessions...' + ) + + successful_requests = 0 + + for i in range(num_sessions): + session_id = self.generate_session_id() + prompt = f'Distribution test for session {session_id}' + + success, response = self.make_request(prompt, session_id=session_id) + + if success: + successful_requests += 1 + self.log(f' Session {i+1}: ✅') + else: + self.log(f' Session {i+1}: ❌ {response}', 'ERROR') + + success_rate = successful_requests / num_sessions + if success_rate >= 0.8: # Allow for some failures + self.log( + f'✅ Distribution test passed: {successful_requests}/{num_sessions} sessions succeeded' + f'({success_rate:.1%})' + ) + return True + else: + self.log( + f'❌ Distribution test failed: Only {successful_requests}/{num_sessions} sessions succeeded' + f'({success_rate:.1%})', + 'ERROR', + ) + return False + + def test_fallback_without_session_or_user(self) -> bool: + """Test that requests without session_id or user_id still work + (fallback behavior).""" + self.log('Testing fallback behavior without session_id or user_id...') + + for i in range(5): + prompt = f'Fallback test request {i+1}' + success, response = self.make_request(prompt) # No session_id or user_id + + if success: + self.log(f' Fallback request {i+1}: ✅') + else: + self.log(f' Fallback request {i+1}: ❌ {response}', 'ERROR') + return False + + self.log( + '✅ Fallback test passed: All requests without session_id/user_id completed successfully' + ) + return True + + def test_openai_user_field_routing(self) -> bool: + """Test routing with OpenAI-style user field instead of + session_params.""" + self.log('Testing OpenAI user field routing...') + + user = self.generate_user_id() + + # Make requests using user field (OpenAI ChatCompletion/Completion style) + for i in range(5): + prompt = f'OpenAI user field test request {i+1}' + request_data = { + 'text': prompt, + 'user': user, # OpenAI-style user field (not in session_params) + 'sampling_params': {'temperature': 0.7, 'max_new_tokens': 30}, + } + + try: + response = self.session.post( + f'{self.router_url}/generate', json=request_data + ) + if response.status_code == 200: + self.log(f' OpenAI user request {i+1}: ✅') + else: + self.log( + f' OpenAI user request {i+1}: ❌ HTTP {response.status_code}', + 'ERROR', + ) + return False + except Exception as e: + self.log(f' OpenAI user request {i+1}: ❌ {e}', 'ERROR') + return False + + self.log('✅ OpenAI user field routing test passed') + return True + + def test_priority_routing(self) -> bool: + """ + Test routing priority: session_params.session_id > session_params.user_id > user field. + """ + self.log('Testing routing priority...') + + session_id = self.generate_session_id() + user_id = self.generate_user_id() + user = f'openai_{self.generate_user_id()}' + + # Test with all three: session_params should win + request_data = { + 'text': 'Priority test', + 'user': user, # Should be ignored + 'session_params': { + 'session_id': session_id, # Should take priority + 'user_id': user_id, # Should be ignored in favor of session_id + }, + 'sampling_params': {'temperature': 0.7, 'max_new_tokens': 30}, + } + + try: + response = self.session.post( + f'{self.router_url}/generate', json=request_data + ) + if response.status_code == 200: + self.log('✅ Priority routing test passed') + return True + else: + self.log( + f'❌ Priority routing test failed: HTTP {response.status_code}', + 'ERROR', + ) + return False + except Exception as e: + self.log(f'❌ Priority routing test failed: {e}', 'ERROR') + return False + + def test_concurrent_requests(self, num_concurrent: int = 10) -> bool: + """Test concurrent requests with the same session_id to verify thread + safety.""" + import threading + import time + + self.log( + f'Testing {num_concurrent} concurrent requests with same session_id...' + ) + + session_id = self.generate_session_id() + results = [] + + def make_concurrent_request(request_id: int): + prompt = f'Concurrent request {request_id} for session {session_id}' + success, response = self.make_request(prompt, session_id=session_id) + results.append((request_id, success, response)) + + # Start all threads + threads = [] + start_time = time.time() + + for i in range(num_concurrent): + thread = threading.Thread(target=make_concurrent_request, args=(i + 1,)) + thread.start() + threads.append(thread) + + # Wait for all threads to complete + for thread in threads: + thread.join() + + end_time = time.time() + + # Check results + successful = sum(1 for _, success, _ in results if success) + success_rate = successful / num_concurrent + + if success_rate >= 0.8: # Allow for some failures in concurrent scenario + self.log( + f'✅ Concurrent test passed: {successful}/{num_concurrent} requests succeeded' + f'({success_rate:.1%}) in {end_time - start_time:.2f}s' + ) + return True + else: + self.log( + f'❌ Concurrent test failed: Only {successful}/{num_concurrent} requests succeeded' + f'({success_rate:.1%})', + 'ERROR', + ) + return False + + def test_different_request_formats(self) -> bool: + """Test consistent hashing with different request formats (JSON + variations).""" + self.log('Testing different request formats...') + + session_id = self.generate_session_id() + + # Test different JSON formats + formats = [ + # Standard format + { + 'text': 'Format test 1', + 'session_id': session_id, + 'sampling_params': {'max_new_tokens': 20}, + }, + # Different field order + { + 'session_id': session_id, + 'text': 'Format test 2', + 'sampling_params': {'temperature': 0.5, 'max_new_tokens': 20}, + }, + # With extra fields + { + 'text': 'Format test 3', + 'session_id': session_id, + 'extra_field': 'ignored', + 'sampling_params': {'max_new_tokens': 20}, + }, + ] + + for i, request_data in enumerate(formats): + try: + response = self.session.post( + f'{self.router_url}/generate', json=request_data + ) + if response.status_code == 200: + self.log(f' Format {i+1}: ✅') + else: + self.log(f' Format {i+1}: ❌ HTTP {response.status_code}', 'ERROR') + return False + except Exception as e: + self.log(f' Format {i+1}: ❌ {e}', 'ERROR') + return False + + self.log('✅ Different request formats test passed') + return True + + def run_comprehensive_test(self) -> bool: + """Run all test cases and return overall success.""" + self.log('=' * 80) + self.log('Starting Consistent Hash Policy Comprehensive Test') + self.log(f'Router URL: {self.router_url}') + self.log('=' * 80) + + tests = [ + ('Session Consistency', lambda: self.test_session_consistency(10)), + ('User Consistency', lambda: self.test_user_consistency(10)), + ('Session Priority', self.test_session_priority_over_user), + ('OpenAI User Field', self.test_openai_user_field_routing), + ('Priority Routing', self.test_priority_routing), + ('Distribution', lambda: self.test_distribution_across_workers(20)), + ('Fallback Behavior', self.test_fallback_without_session_or_user), + ('Concurrent Requests', lambda: self.test_concurrent_requests(10)), + ('Request Formats', self.test_different_request_formats), + ] + + passed = 0 + failed = 0 + + for test_name, test_func in tests: + self.log('') + self.log(f'🧪 Running: {test_name}') + try: + if test_func(): + passed += 1 + else: + failed += 1 + except Exception as e: + self.log(f"❌ Test '{test_name}' failed with exception: {e}", 'ERROR') + failed += 1 + + self.log('') + self.log('=' * 80) + self.log('Test Results Summary') + self.log('=' * 80) + + total = passed + failed + success_rate = passed / total if total > 0 else 0 + + self.log(f'Total Tests: {total}') + self.log(f'Passed: {passed}') + self.log(f'Failed: {failed}') + self.log(f'Success Rate: {success_rate:.1%}') + + if failed == 0: + self.log('🎉 All tests passed! Consistent hashing is working correctly.') + return True + else: + self.log( + f'⚠️ {failed} test(s) failed. Please check the router configuration and try again.' + ) + return False + + +def main(): + parser = argparse.ArgumentParser( + description='Test consistent hash policy for the router' + ) + parser.add_argument( + '--router-url', + default='http://localhost:30000', + help='Router URL (default: http://localhost:30000)', + ) + parser.add_argument( + '--test', + choices=[ + 'session', + 'user', + 'priority', + 'distribution', + 'fallback', + 'dp_routing', + 'concurrent', + 'formats', + 'all', + ], + default='all', + help='Run specific test (default: all)', + ) + + args = parser.parse_args() + + tester = ConsistentHashTester(router_url=args.router_url) + + if args.test == 'all': + success = tester.run_comprehensive_test() + return 0 if success else 1 + else: + # Run specific test + test_methods = { + 'session': lambda: tester.test_session_consistency(10), + 'user': lambda: tester.test_user_consistency(10), + 'priority': tester.test_session_priority_over_user, + 'distribution': lambda: tester.test_distribution_across_workers(20), + 'fallback': tester.test_fallback_without_session_or_user, + 'concurrent': lambda: tester.test_concurrent_requests(10), + 'formats': tester.test_different_request_formats, + } + + if args.test in test_methods: + success = test_methods[args.test]() + return 0 if success else 1 + else: + print(f'Unknown test: {args.test}') + return 1 + + +if __name__ == '__main__': + import sys + + sys.exit(main()) diff --git a/lmdeploy/router/py_test/test_launch_router.py b/lmdeploy/router/py_test/test_launch_router.py new file mode 100644 index 0000000000..a6ad35012d --- /dev/null +++ b/lmdeploy/router/py_test/test_launch_router.py @@ -0,0 +1,323 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import multiprocessing +import time +import unittest +from types import SimpleNamespace + + +def terminate_process(process: multiprocessing.Process, timeout: float = 1.0) -> None: + """Terminate a process gracefully, with forced kill as fallback. + + Args: + process: The process to terminate + timeout: Seconds to wait for graceful termination before forcing kill + """ + if not process.is_alive(): + return + + process.terminate() + process.join(timeout=timeout) + if process.is_alive(): + process.kill() # Force kill if terminate didn't work + process.join() + + +def run_router(args): + try: + from lmdeploy_router.launch_router import launch_router + + router = launch_router(args) + if router is None: + return 1 + return 0 + except Exception as e: + print(e) + return 1 + + +class TestLaunchRouter(unittest.TestCase): + def setUp(self): + """Set up default arguments for router tests.""" + self.default_args = SimpleNamespace( + host='127.0.0.1', + port=30000, + policy='cache_aware', + worker_startup_timeout_secs=600, + worker_startup_check_interval=10, + cache_threshold=0.5, + balance_abs_threshold=32, + balance_rel_threshold=1.0001, + eviction_interval=60, + max_tree_size=2**24, + max_payload_size=256 * 1024 * 1024, # 256MB + verbose=False, + log_dir=None, + log_level=None, + service_discovery=False, + selector=None, + service_discovery_port=80, + service_discovery_namespace=None, + prometheus_port=None, + prometheus_host=None, + request_timeout_secs=60, + max_concurrent_requests=64, + cors_allowed_origins=[], + lmdeploy_pd_disaggregation=False, + prefill=None, + decode=None, + worker_urls=[], + retry_max_retries=3, + retry_initial_backoff_ms=100, + retry_max_backoff_ms=10_000, + retry_backoff_multiplier=2.0, + retry_jitter_factor=0.1, + cb_failure_threshold=5, + cb_success_threshold=2, + cb_timeout_duration_secs=30, + cb_window_duration_secs=60, + disable_retries=False, + disable_circuit_breaker=False, + ) + + def create_router_args(self, **kwargs): + """Create router arguments by updating default args with provided + kwargs.""" + args_dict = vars(self.default_args).copy() + args_dict.update(kwargs) + return SimpleNamespace(**args_dict) + + def run_router_process(self, args): + """Run router in a separate process and verify it starts + successfully.""" + process = multiprocessing.Process(target=run_router, args=(args,)) + try: + process.start() + # Wait 3 seconds + time.sleep(3) + # Process is still running means router started successfully + self.assertTrue(process.is_alive()) + finally: + terminate_process(process) + + def test_launch_router_common(self): + args = self.create_router_args(worker_urls=['http://localhost:8000']) + self.run_router_process(args) + + def test_launch_router_with_empty_worker_urls(self): + args = self.create_router_args(worker_urls=[]) + self.run_router_process( + args + ) # Should start successfully with empty worker list + + def test_launch_router_with_service_discovery(self): + # Test router startup with service discovery enabled but no selectors + args = self.create_router_args( + worker_urls=[], service_discovery=True, selector=['app=test-worker'] + ) + self.run_router_process(args) + + def test_launch_router_with_service_discovery_namespace(self): + # Test router startup with service discovery enabled and namespace specified + args = self.create_router_args( + worker_urls=[], + service_discovery=True, + selector=['app=test-worker'], + service_discovery_namespace='test-namespace', + ) + self.run_router_process(args) + + def test_launch_router_pd_mode_basic(self): + """Test basic PD router functionality without actually starting + servers.""" + # This test just verifies the PD router can be created and configured + # without actually starting it (which would require real prefill/decode servers) + from lmdeploy_router import Router + from lmdeploy_router.launch_router import RouterArgs + from lmdeploy_router_rs import PolicyType + + # Test RouterArgs parsing for PD mode + # Simulate the parsed args structure from argparse with action="append" + args = self.create_router_args( + lmdeploy_pd_disaggregation=True, + policy='power_of_two', # PowerOfTwo is only valid in PD mode + prefill=[ + ['http://prefill1:8080'], + ['http://prefill2:8080'], + ], + decode=[ + ['http://decode1:8081'], + ['http://decode2:8081'], + ], + worker_urls=[], # Empty for PD mode + ) + + router_args = RouterArgs.from_cli_args(args) + self.assertTrue(router_args.lmdeploy_pd_disaggregation) + self.assertEqual(router_args.policy, 'power_of_two') + self.assertEqual(len(router_args.prefill_urls), 2) + self.assertEqual(len(router_args.decode_urls), 2) + + # Verify the parsed URLs + self.assertEqual(router_args.prefill_urls[0], 'http://prefill1:8080') + self.assertEqual(router_args.prefill_urls[1], 'http://prefill2:8080') + self.assertEqual(router_args.decode_urls[0], 'http://decode1:8081') + self.assertEqual(router_args.decode_urls[1], 'http://decode2:8081') + + # Test Router creation in PD mode + router = Router( + worker_urls=[], # Empty for PD mode + lmdeploy_pd_disaggregation=True, + prefill_urls=['http://prefill1:8080', 'http://prefill2:8080'], + decode_urls=['http://decode1:8081', 'http://decode2:8081'], + policy=PolicyType.CacheAware, + host='127.0.0.1', + port=3001, + ) + self.assertIsNotNone(router) + + def test_policy_validation(self): + """Test that policy validation works correctly for PD and regular + modes.""" + from lmdeploy_router.launch_router import launch_router + + # Test 1: PowerOfTwo requires at least 2 workers + args = self.create_router_args( + lmdeploy_pd_disaggregation=False, + policy='power_of_two', + worker_urls=['http://localhost:8000'], # Only 1 worker + ) + + # Should raise error + with self.assertRaises(ValueError) as cm: + launch_router(args) + self.assertIn( + 'Power-of-two policy requires at least 2 workers', + str(cm.exception), + ) + + # Test 2: PowerOfTwo with sufficient workers should succeed + args = self.create_router_args( + lmdeploy_pd_disaggregation=False, + policy='power_of_two', + worker_urls=['http://localhost:8000', 'http://localhost:8001'], # 2 workers + ) + # This should not raise an error (validation passes) + + # Test 3: All policies now work in both modes + # Regular mode with RoundRobin + args = self.create_router_args( + lmdeploy_pd_disaggregation=False, + policy='round_robin', + worker_urls=['http://localhost:8000'], + ) + # This should not raise validation error + + # PD mode with RoundRobin (now supported!) + args = self.create_router_args( + lmdeploy_pd_disaggregation=True, + policy='round_robin', + prefill=[['http://prefill1:8080', '9000']], + decode=[['http://decode1:8081']], + worker_urls=[], + ) + # This should not raise validation error + + def test_pd_service_discovery_args_parsing(self): + """Test PD service discovery CLI argument parsing.""" + import argparse + + from lmdeploy_router.launch_router import RouterArgs + + parser = argparse.ArgumentParser() + RouterArgs.add_cli_args(parser) + + args = parser.parse_args( + [ + '--lmdeploy-pd-disaggregation', + '--service-discovery', + '--prefill-selector', + 'app=lmdeploy', + 'component=prefill', + '--decode-selector', + 'app=lmdeploy', + 'component=decode', + '--service-discovery-port', + '8000', + '--service-discovery-namespace', + 'production', + '--policy', + 'cache_aware', + ] + ) + + router_args = RouterArgs.from_cli_args(args) + + self.assertTrue(router_args.lmdeploy_pd_disaggregation) + self.assertTrue(router_args.service_discovery) + self.assertEqual( + router_args.prefill_selector, + {'app': 'lmdeploy', 'component': 'prefill'}, + ) + self.assertEqual( + router_args.decode_selector, + {'app': 'lmdeploy', 'component': 'decode'}, + ) + self.assertEqual(router_args.service_discovery_port, 8000) + self.assertEqual(router_args.service_discovery_namespace, 'production') + + def test_regular_service_discovery_args_parsing(self): + """Test regular mode service discovery CLI argument parsing.""" + import argparse + + from lmdeploy_router.launch_router import RouterArgs + + parser = argparse.ArgumentParser() + RouterArgs.add_cli_args(parser) + + args = parser.parse_args( + [ + '--service-discovery', + '--selector', + 'app=lmdeploy-worker', + 'environment=staging', + '--service-discovery-port', + '8000', + '--policy', + 'round_robin', + ] + ) + + router_args = RouterArgs.from_cli_args(args) + + self.assertFalse(router_args.lmdeploy_pd_disaggregation) + self.assertTrue(router_args.service_discovery) + self.assertEqual( + router_args.selector, + {'app': 'lmdeploy-worker', 'environment': 'staging'}, + ) + self.assertEqual(router_args.prefill_selector, {}) + self.assertEqual(router_args.decode_selector, {}) + + def test_empty_worker_urls_args_parsing(self): + """Test that router accepts no worker URLs and defaults to empty + list.""" + import argparse + + from lmdeploy_router.launch_router import RouterArgs + + parser = argparse.ArgumentParser() + RouterArgs.add_cli_args(parser) + + # Test with no --worker-urls argument at all + args = parser.parse_args(['--policy', 'random', '--port', '30000']) + router_args = RouterArgs.from_cli_args(args) + self.assertEqual(router_args.worker_urls, []) + + # Test with explicit empty --worker-urls + args = parser.parse_args(['--worker-urls', '--policy', 'random']) + router_args = RouterArgs.from_cli_args(args) + self.assertEqual(router_args.worker_urls, []) + + +if __name__ == '__main__': + unittest.main() diff --git a/lmdeploy/router/py_test/unit/__init__.py b/lmdeploy/router/py_test/unit/__init__.py new file mode 100644 index 0000000000..f5b33bf606 --- /dev/null +++ b/lmdeploy/router/py_test/unit/__init__.py @@ -0,0 +1,6 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""Unit tests for lmdeploy_router. + +This package contains fast, isolated unit tests for Python components of the LMDeploy router. These tests focus on +testing individual functions and classes in isolation without starting actual router instances. +""" diff --git a/lmdeploy/router/py_test/unit/test_arg_parser.py b/lmdeploy/router/py_test/unit/test_arg_parser.py new file mode 100644 index 0000000000..68fd5ba882 --- /dev/null +++ b/lmdeploy/router/py_test/unit/test_arg_parser.py @@ -0,0 +1,598 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""Unit tests for argument parsing functionality in lmdeploy_router. + +These tests focus on testing the argument parsing logic in isolation, without starting actual router instances. +""" + +from types import SimpleNamespace + +import pytest +from lmdeploy_router.launch_router import RouterArgs, parse_router_args +from lmdeploy_router.router import policy_from_str + + +class TestRouterArgs: + """Test RouterArgs dataclass and its methods.""" + + def test_default_values(self): + """Test that RouterArgs has correct default values.""" + args = RouterArgs() + + # Test basic defaults + assert args.host == '127.0.0.1' + assert args.port == 30000 + assert args.policy == 'cache_aware' + assert args.worker_urls == [] + assert args.prefill_urls == [] + assert args.decode_urls == [] + + # Test PD-specific defaults + assert args.prefill_policy is None + assert args.decode_policy is None + + # Test service discovery defaults + assert args.service_discovery is False + assert args.selector == {} + assert args.service_discovery_port == 80 + assert args.service_discovery_namespace is None + + # Test retry and circuit breaker defaults + assert args.retry_max_retries == 5 + assert args.cb_failure_threshold == 10 + assert args.disable_retries is False + assert args.disable_circuit_breaker is False + + def test_parse_selector_valid(self): + """Test parsing valid selector arguments.""" + # Test single key-value pair + result = RouterArgs._parse_selector(['app=worker']) + assert result == {'app': 'worker'} + + # Test multiple key-value pairs + result = RouterArgs._parse_selector(['app=worker', 'env=prod', 'version=v1']) + assert result == {'app': 'worker', 'env': 'prod', 'version': 'v1'} + + # Test empty list + result = RouterArgs._parse_selector([]) + assert result == {} + + # Test None + result = RouterArgs._parse_selector(None) + assert result == {} + + def test_parse_selector_invalid(self): + """Test parsing invalid selector arguments.""" + # Test malformed selector (no equals sign) + result = RouterArgs._parse_selector(['app']) + assert result == {} + + # Test multiple equals signs (should use first one) + result = RouterArgs._parse_selector(['app=worker=extra']) + assert result == {'app': 'worker=extra'} + + def test_parse_prefill_urls_valid(self): + """Test parsing valid prefill URL arguments.""" + # Test a single prefill URL + result = RouterArgs._parse_prefill_urls([['http://prefill1:8000']]) + assert result == ['http://prefill1:8000'] + + # Test multiple prefill URLs + result = RouterArgs._parse_prefill_urls( + [ + ['http://prefill1:8000'], + ['http://prefill2:8000'], + ['http://prefill3:8000'], + ] + ) + expected = [ + 'http://prefill1:8000', + 'http://prefill2:8000', + 'http://prefill3:8000', + ] + assert result == expected + + # Test empty list + result = RouterArgs._parse_prefill_urls([]) + assert result == [] + + # Test None + result = RouterArgs._parse_prefill_urls(None) + assert result == [] + + def test_parse_decode_urls_valid(self): + """Test parsing valid decode URL arguments.""" + # Test single decode URL + result = RouterArgs._parse_decode_urls([['http://decode1:8001']]) + assert result == ['http://decode1:8001'] + + # Test multiple decode URLs + result = RouterArgs._parse_decode_urls( + [['http://decode1:8001'], ['http://decode2:8001']] + ) + assert result == ['http://decode1:8001', 'http://decode2:8001'] + + # Test empty list + result = RouterArgs._parse_decode_urls([]) + assert result == [] + + # Test None + result = RouterArgs._parse_decode_urls(None) + assert result == [] + + def test_from_cli_args_basic(self): + """Test creating RouterArgs from basic CLI arguments.""" + args = SimpleNamespace( + host='0.0.0.0', + port=30001, + worker_urls=['http://worker1:8000', 'http://worker2:8000'], + policy='round_robin', + prefill=None, + decode=None, + router_policy='round_robin', + router_prefill_policy=None, + router_decode_policy=None, + router_worker_startup_timeout_secs=300, + router_worker_startup_check_interval=15, + router_cache_threshold=0.7, + router_balance_abs_threshold=128, + router_balance_rel_threshold=2.0, + router_eviction_interval=180, + router_max_tree_size=2**28, + router_max_payload_size=1024 * 1024 * 1024, # 1GB + router_api_key='test-key', + router_log_dir='/tmp/logs', + router_log_level='debug', + router_service_discovery=True, + router_selector=['app=worker', 'env=test'], + router_service_discovery_port=8080, + router_service_discovery_namespace='default', + router_prefill_selector=['app=prefill'], + router_decode_selector=['app=decode'], + router_prometheus_port=29000, + router_prometheus_host='0.0.0.0', + router_request_id_headers=['x-request-id', 'x-trace-id'], + router_request_timeout_secs=1200, + router_max_concurrent_requests=512, + router_queue_size=200, + router_queue_timeout_secs=120, + router_rate_limit_tokens_per_second=100, + router_cors_allowed_origins=['http://localhost:3000'], + router_retry_max_retries=3, + router_retry_initial_backoff_ms=100, + router_retry_max_backoff_ms=10000, + router_retry_backoff_multiplier=2.0, + router_retry_jitter_factor=0.1, + router_cb_failure_threshold=5, + router_cb_success_threshold=2, + router_cb_timeout_duration_secs=30, + router_cb_window_duration_secs=60, + router_disable_retries=False, + router_disable_circuit_breaker=False, + router_health_failure_threshold=2, + router_health_success_threshold=1, + router_health_check_timeout_secs=3, + router_health_check_interval_secs=30, + router_health_check_endpoint='/healthz', + ) + + router_args = RouterArgs.from_cli_args(args, use_router_prefix=True) + + # Test basic configuration + assert router_args.host == '0.0.0.0' + assert router_args.port == 30001 + assert router_args.worker_urls == ['http://worker1:8000', 'http://worker2:8000'] + assert router_args.policy == 'round_robin' + + # Test PD configuration + assert router_args.prefill_urls == [] + assert router_args.decode_urls == [] + + # Test service discovery + assert router_args.service_discovery is True + assert router_args.selector == {'app': 'worker', 'env': 'test'} + assert router_args.service_discovery_port == 8080 + assert router_args.service_discovery_namespace == 'default' + assert router_args.prefill_selector == {'app': 'prefill'} + assert router_args.decode_selector == {'app': 'decode'} + + # Test other configurations + assert router_args.api_key == 'test-key' + assert router_args.log_dir == '/tmp/logs' + assert router_args.log_level == 'debug' + assert router_args.prometheus_port == 29000 + assert router_args.prometheus_host == '0.0.0.0' + assert router_args.request_id_headers == ['x-request-id', 'x-trace-id'] + assert router_args.request_timeout_secs == 1200 + assert router_args.max_concurrent_requests == 512 + assert router_args.queue_size == 200 + assert router_args.queue_timeout_secs == 120 + assert router_args.rate_limit_tokens_per_second == 100 + assert router_args.cors_allowed_origins == ['http://localhost:3000'] + + # Test retry configuration + assert router_args.retry_max_retries == 3 + assert router_args.retry_initial_backoff_ms == 100 + assert router_args.retry_max_backoff_ms == 10000 + assert router_args.retry_backoff_multiplier == 2.0 + assert router_args.retry_jitter_factor == 0.1 + + # Test circuit breaker configuration + assert router_args.cb_failure_threshold == 5 + assert router_args.cb_success_threshold == 2 + assert router_args.cb_timeout_duration_secs == 30 + assert router_args.cb_window_duration_secs == 60 + assert router_args.disable_retries is False + assert router_args.disable_circuit_breaker is False + + # Test health check configuration + assert router_args.health_failure_threshold == 2 + assert router_args.health_success_threshold == 1 + assert router_args.health_check_timeout_secs == 3 + assert router_args.health_check_interval_secs == 30 + assert router_args.health_check_endpoint == '/healthz' + + def test_from_cli_args_pd_mode(self): + """Test creating RouterArgs from CLI arguments in PD mode.""" + args = SimpleNamespace( + host='127.0.0.1', + port=30000, + worker_urls=[], + policy='cache_aware', + prefill=[ + ['http://prefill1:8000'], + ['http://prefill2:8000'], + ], + decode=[['http://decode1:8001'], ['http://decode2:8001']], + router_prefill=[ + ['http://prefill1:8000'], + ['http://prefill2:8000'], + ], + router_decode=[['http://decode1:8001'], ['http://decode2:8001']], + router_policy='cache_aware', + router_lmdeploy_pd_disaggregation=True, + router_prefill_policy='power_of_two', + router_decode_policy='round_robin', + # Include all required fields with defaults + router_worker_startup_timeout_secs=600, + router_worker_startup_check_interval=30, + router_cache_threshold=0.3, + router_balance_abs_threshold=64, + router_balance_rel_threshold=1.5, + router_eviction_interval=120, + router_max_tree_size=2**26, + router_max_payload_size=512 * 1024 * 1024, + router_api_key=None, + router_log_dir=None, + router_log_level=None, + router_service_discovery=False, + router_selector=None, + router_service_discovery_port=80, + router_service_discovery_namespace=None, + router_prefill_selector=None, + router_decode_selector=None, + router_prometheus_port=None, + router_prometheus_host=None, + router_request_id_headers=None, + router_request_timeout_secs=1800, + router_max_concurrent_requests=256, + router_queue_size=100, + router_queue_timeout_secs=60, + router_rate_limit_tokens_per_second=None, + router_cors_allowed_origins=[], + router_retry_max_retries=5, + router_retry_initial_backoff_ms=50, + router_retry_max_backoff_ms=30000, + router_retry_backoff_multiplier=1.5, + router_retry_jitter_factor=0.2, + router_cb_failure_threshold=10, + router_cb_success_threshold=3, + router_cb_timeout_duration_secs=60, + router_cb_window_duration_secs=120, + router_disable_retries=False, + router_disable_circuit_breaker=False, + router_health_failure_threshold=3, + router_health_success_threshold=2, + router_health_check_timeout_secs=5, + router_health_check_interval_secs=60, + router_health_check_endpoint='/health', + ) + + router_args = RouterArgs.from_cli_args(args, use_router_prefix=True) + + # Test PD configuration + assert router_args.lmdeploy_pd_disaggregation is True + assert router_args.prefill_urls == [ + 'http://prefill1:8000', + 'http://prefill2:8000', + ] + assert router_args.decode_urls == ['http://decode1:8001', 'http://decode2:8001'] + assert router_args.prefill_policy == 'power_of_two' + assert router_args.decode_policy == 'round_robin' + assert router_args.policy == 'cache_aware' # Main policy still set + + def test_from_cli_args_without_prefix(self): + """Test creating RouterArgs from CLI arguments without router + prefix.""" + args = SimpleNamespace( + host='127.0.0.1', + port=30000, + worker_urls=['http://worker1:8000'], + policy='random', + prefill=None, + decode=None, + prefill_policy=None, + decode_policy=None, + worker_startup_timeout_secs=600, + worker_startup_check_interval=30, + cache_threshold=0.3, + balance_abs_threshold=64, + balance_rel_threshold=1.5, + eviction_interval=120, + max_tree_size=2**26, + max_payload_size=512 * 1024 * 1024, + api_key=None, + log_dir=None, + log_level=None, + service_discovery=False, + selector=None, + service_discovery_port=80, + service_discovery_namespace=None, + prefill_selector=None, + decode_selector=None, + prometheus_port=None, + prometheus_host=None, + request_id_headers=None, + request_timeout_secs=1800, + max_concurrent_requests=256, + queue_size=100, + queue_timeout_secs=60, + rate_limit_tokens_per_second=None, + cors_allowed_origins=[], + retry_max_retries=5, + retry_initial_backoff_ms=50, + retry_max_backoff_ms=30000, + retry_backoff_multiplier=1.5, + retry_jitter_factor=0.2, + cb_failure_threshold=10, + cb_success_threshold=3, + cb_timeout_duration_secs=60, + cb_window_duration_secs=120, + disable_retries=False, + disable_circuit_breaker=False, + health_failure_threshold=3, + health_success_threshold=2, + health_check_timeout_secs=5, + health_check_interval_secs=60, + health_check_endpoint='/health', + ) + + router_args = RouterArgs.from_cli_args(args, use_router_prefix=False) + + assert router_args.host == '127.0.0.1' + assert router_args.port == 30000 + assert router_args.worker_urls == ['http://worker1:8000'] + assert router_args.policy == 'random' + + +class TestPolicyFromStr: + """Test policy string to enum conversion.""" + + def test_valid_policies(self): + """Test conversion of valid policy strings.""" + from lmdeploy_router_rs import PolicyType + + assert policy_from_str('random') == PolicyType.Random + assert policy_from_str('round_robin') == PolicyType.RoundRobin + assert policy_from_str('cache_aware') == PolicyType.CacheAware + assert policy_from_str('power_of_two') == PolicyType.PowerOfTwo + assert policy_from_str('consistent_hash') == PolicyType.ConsistentHash + + def test_invalid_policy(self): + """Test conversion of invalid policy string.""" + with pytest.raises(KeyError): + policy_from_str('invalid_policy') + + +class TestParseRouterArgs: + """Test the parse_router_args function.""" + + def test_parse_basic_args(self): + """Test parsing basic router arguments.""" + args = [ + '--host', + '0.0.0.0', + '--port', + '30001', + '--worker-urls', + 'http://worker1:8000', + 'http://worker2:8000', + '--policy', + 'round_robin', + ] + + router_args = parse_router_args(args) + + assert router_args.host == '0.0.0.0' + assert router_args.port == 30001 + assert router_args.worker_urls == ['http://worker1:8000', 'http://worker2:8000'] + assert router_args.policy == 'round_robin' + + def test_parse_pd_args(self): + """Test parsing PD disaggregated mode arguments.""" + args = [ + '--lmdeploy-pd-disaggregation', + '--prefill', + 'http://prefill1:8000', + '--prefill', + 'http://prefill2:8000', + '--decode', + 'http://decode1:8001', + '--decode', + 'http://decode2:8001', + '--prefill-policy', + 'power_of_two', + '--decode-policy', + 'round_robin', + ] + + router_args = parse_router_args(args) + + assert router_args.lmdeploy_pd_disaggregation is True + assert router_args.prefill_urls == [ + 'http://prefill1:8000', + 'http://prefill2:8000', + ] + assert router_args.decode_urls == ['http://decode1:8001', 'http://decode2:8001'] + assert router_args.prefill_policy == 'power_of_two' + assert router_args.decode_policy == 'round_robin' + + def test_parse_service_discovery_args(self): + """Test parsing service discovery arguments.""" + args = [ + '--service-discovery', + '--selector', + 'app=worker', + 'env=prod', + '--service-discovery-port', + '8080', + '--service-discovery-namespace', + 'default', + ] + + router_args = parse_router_args(args) + + assert router_args.service_discovery is True + assert router_args.selector == {'app': 'worker', 'env': 'prod'} + assert router_args.service_discovery_port == 8080 + assert router_args.service_discovery_namespace == 'default' + + def test_parse_retry_and_circuit_breaker_args(self): + """Test parsing retry and circuit breaker arguments.""" + args = [ + '--retry-max-retries', + '3', + '--retry-initial-backoff-ms', + '100', + '--retry-max-backoff-ms', + '10000', + '--retry-backoff-multiplier', + '2.0', + '--retry-jitter-factor', + '0.1', + '--disable-retries', + '--cb-failure-threshold', + '5', + '--cb-success-threshold', + '2', + '--cb-timeout-duration-secs', + '30', + '--cb-window-duration-secs', + '60', + '--disable-circuit-breaker', + ] + + router_args = parse_router_args(args) + + # Test retry configuration + assert router_args.retry_max_retries == 3 + assert router_args.retry_initial_backoff_ms == 100 + assert router_args.retry_max_backoff_ms == 10000 + assert router_args.retry_backoff_multiplier == 2.0 + assert router_args.retry_jitter_factor == 0.1 + assert router_args.disable_retries is True + + # Test circuit breaker configuration + assert router_args.cb_failure_threshold == 5 + assert router_args.cb_success_threshold == 2 + assert router_args.cb_timeout_duration_secs == 30 + assert router_args.cb_window_duration_secs == 60 + assert router_args.disable_circuit_breaker is True + + def test_parse_rate_limiting_args(self): + """Test parsing rate limiting arguments.""" + args = [ + '--max-concurrent-requests', + '512', + '--queue-size', + '200', + '--queue-timeout-secs', + '120', + '--rate-limit-tokens-per-second', + '100', + ] + + router_args = parse_router_args(args) + + assert router_args.max_concurrent_requests == 512 + assert router_args.queue_size == 200 + assert router_args.queue_timeout_secs == 120 + assert router_args.rate_limit_tokens_per_second == 100 + + def test_parse_health_check_args(self): + """Test parsing health check arguments.""" + args = [ + '--health-failure-threshold', + '2', + '--health-success-threshold', + '1', + '--health-check-timeout-secs', + '3', + '--health-check-interval-secs', + '30', + '--health-check-endpoint', + '/healthz', + ] + + router_args = parse_router_args(args) + + assert router_args.health_failure_threshold == 2 + assert router_args.health_success_threshold == 1 + assert router_args.health_check_timeout_secs == 3 + assert router_args.health_check_interval_secs == 30 + assert router_args.health_check_endpoint == '/healthz' + + def test_parse_cors_args(self): + """Test parsing CORS arguments.""" + args = [ + '--cors-allowed-origins', + 'http://localhost:3000', + 'https://example.com', + ] + + router_args = parse_router_args(args) + + assert router_args.cors_allowed_origins == [ + 'http://localhost:3000', + 'https://example.com', + ] + + def test_parse_tokenizer_args(self): + """Test parsing tokenizer arguments.""" + # Note: model-path and tokenizer-path arguments are not available in current implementation + # This test is skipped until those arguments are added + pytest.skip('Tokenizer arguments not available in current implementation') + + def test_parse_invalid_args(self): + """Test parsing invalid arguments.""" + # Test invalid policy + with pytest.raises(SystemExit): + parse_router_args(['--policy', 'invalid_policy']) + + # --prefill only accepts a single URL + with pytest.raises(SystemExit): + parse_router_args( + [ + '--lmdeploy-pd-disaggregation', + '--prefill', + 'http://prefill1:8000', + 'invalid_port', + ] + ) + + def test_help_output(self): + """Test that help output is generated correctly.""" + with pytest.raises(SystemExit) as exc_info: + parse_router_args(['--help']) + + # SystemExit with code 0 indicates help was displayed + assert exc_info.value.code == 0 diff --git a/lmdeploy/router/py_test/unit/test_router_config.py b/lmdeploy/router/py_test/unit/test_router_config.py new file mode 100644 index 0000000000..d69784c99d --- /dev/null +++ b/lmdeploy/router/py_test/unit/test_router_config.py @@ -0,0 +1,417 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""Unit tests for router configuration validation and setup. + +These tests focus on testing the router configuration logic in isolation, including validation of configuration +parameters and their interactions. +""" + +from unittest.mock import MagicMock, patch + +import pytest +from lmdeploy_router.launch_router import RouterArgs, launch_router +from lmdeploy_router.router import policy_from_str +from lmdeploy_router_rs import PolicyType + + +class TestRouterConfigValidation: + """Test router configuration validation logic.""" + + def test_valid_basic_config(self): + """Test that a valid basic configuration passes validation.""" + args = RouterArgs( + host='127.0.0.1', + port=30000, + worker_urls=['http://worker1:8000', 'http://worker2:8000'], + policy='cache_aware', + ) + + # Should not raise any exceptions + assert args.host == '127.0.0.1' + assert args.port == 30000 + assert args.worker_urls == ['http://worker1:8000', 'http://worker2:8000'] + assert args.policy == 'cache_aware' + + def test_valid_pd_config(self): + """Test that a valid PD configuration passes validation.""" + args = RouterArgs( + host='127.0.0.1', + port=30000, + lmdeploy_pd_disaggregation=True, + prefill_urls=[ + 'http://prefill1:8000', + 'http://prefill2:8000', + ], + decode_urls=['http://decode1:8001', 'http://decode2:8001'], + policy='cache_aware', + ) + + assert args.lmdeploy_pd_disaggregation is True + assert args.prefill_urls == [ + 'http://prefill1:8000', + 'http://prefill2:8000', + ] + assert args.decode_urls == ['http://decode1:8001', 'http://decode2:8001'] + assert args.policy == 'cache_aware' + + def test_pd_config_without_urls_allows_dynamic_registration(self): + """Test that PD mode without URLs allows worker registration.""" + args = RouterArgs( + lmdeploy_pd_disaggregation=True, + prefill_urls=[], + decode_urls=[], + service_discovery=False, + ) + + with patch('lmdeploy_router.launch_router.Router') as router_mod: + mock_router_instance = MagicMock() + router_mod.from_args = MagicMock(return_value=mock_router_instance) + + launch_router(args) + router_mod.from_args.assert_called_once() + + def test_pd_config_with_service_discovery_allows_empty_urls(self): + """Test that PD mode with service discovery allows empty URLs.""" + args = RouterArgs( + lmdeploy_pd_disaggregation=True, + prefill_urls=[], + decode_urls=[], + service_discovery=True, + ) + + # Should not raise validation error when service discovery is enabled + with patch('lmdeploy_router.launch_router.Router') as router_mod: + mock_router_instance = MagicMock() + router_mod.from_args = MagicMock(return_value=mock_router_instance) + + launch_router(args) + + # Should create router instance via from_args + router_mod.from_args.assert_called_once() + + def test_regular_mode_without_workers_allows_empty_urls(self): + """Test that regular mode allows empty worker URLs.""" + args = RouterArgs(worker_urls=[], service_discovery=False) + + # Should not raise validation error + with patch('lmdeploy_router.launch_router.Router') as router_mod: + mock_router_instance = MagicMock() + router_mod.from_args = MagicMock(return_value=mock_router_instance) + + launch_router(args) + + # Should create router instance via from_args + router_mod.from_args.assert_called_once() + + def test_cache_threshold_validation(self): + """Test cache threshold validation.""" + # Valid cache threshold + args = RouterArgs(cache_threshold=0.5) + assert args.cache_threshold == 0.5 + + # Edge cases + args = RouterArgs(cache_threshold=0.0) + assert args.cache_threshold == 0.0 + + args = RouterArgs(cache_threshold=1.0) + assert args.cache_threshold == 1.0 + + def test_balance_threshold_validation(self): + """Test load balancing threshold validation.""" + # Valid thresholds + args = RouterArgs(balance_abs_threshold=64, balance_rel_threshold=1.5) + assert args.balance_abs_threshold == 64 + assert args.balance_rel_threshold == 1.5 + + # Edge cases + args = RouterArgs(balance_abs_threshold=0, balance_rel_threshold=1.0) + assert args.balance_abs_threshold == 0 + assert args.balance_rel_threshold == 1.0 + + def test_timeout_validation(self): + """Test timeout parameter validation.""" + # Valid timeouts + args = RouterArgs( + worker_startup_timeout_secs=600, + worker_startup_check_interval=30, + request_timeout_secs=1800, + queue_timeout_secs=60, + ) + assert args.worker_startup_timeout_secs == 600 + assert args.worker_startup_check_interval == 30 + assert args.request_timeout_secs == 1800 + assert args.queue_timeout_secs == 60 + + def test_retry_config_validation(self): + """Test retry configuration validation.""" + # Valid retry config + args = RouterArgs( + retry_max_retries=5, + retry_initial_backoff_ms=50, + retry_max_backoff_ms=30000, + retry_backoff_multiplier=1.5, + retry_jitter_factor=0.2, + disable_retries=False, + ) + assert args.retry_max_retries == 5 + assert args.retry_initial_backoff_ms == 50 + assert args.retry_max_backoff_ms == 30000 + assert args.retry_backoff_multiplier == 1.5 + assert args.retry_jitter_factor == 0.2 + assert args.disable_retries is False + + def test_circuit_breaker_config_validation(self): + """Test circuit breaker configuration validation.""" + # Valid circuit breaker config + args = RouterArgs( + cb_failure_threshold=10, + cb_success_threshold=3, + cb_timeout_duration_secs=60, + cb_window_duration_secs=120, + disable_circuit_breaker=False, + ) + assert args.cb_failure_threshold == 10 + assert args.cb_success_threshold == 3 + assert args.cb_timeout_duration_secs == 60 + assert args.cb_window_duration_secs == 120 + assert args.disable_circuit_breaker is False + + def test_health_check_config_validation(self): + """Test health check configuration validation.""" + # Valid health check config + args = RouterArgs( + health_failure_threshold=3, + health_success_threshold=2, + health_check_timeout_secs=5, + health_check_interval_secs=60, + health_check_endpoint='/health', + ) + assert args.health_failure_threshold == 3 + assert args.health_success_threshold == 2 + assert args.health_check_timeout_secs == 5 + assert args.health_check_interval_secs == 60 + assert args.health_check_endpoint == '/health' + + def test_rate_limiting_config_validation(self): + """Test rate limiting configuration validation.""" + # Valid rate limiting config + args = RouterArgs( + max_concurrent_requests=256, + queue_size=100, + queue_timeout_secs=60, + rate_limit_tokens_per_second=100, + ) + assert args.max_concurrent_requests == 256 + assert args.queue_size == 100 + assert args.queue_timeout_secs == 60 + assert args.rate_limit_tokens_per_second == 100 + + def test_service_discovery_config_validation(self): + """Test service discovery configuration validation.""" + # Valid service discovery config + args = RouterArgs( + service_discovery=True, + selector={'app': 'worker', 'env': 'prod'}, + service_discovery_port=8080, + service_discovery_namespace='default', + ) + assert args.service_discovery is True + assert args.selector == {'app': 'worker', 'env': 'prod'} + assert args.service_discovery_port == 8080 + assert args.service_discovery_namespace == 'default' + + def test_pd_service_discovery_config_validation(self): + """Test PD service discovery configuration validation.""" + # Valid PD service discovery config + args = RouterArgs( + lmdeploy_pd_disaggregation=True, + service_discovery=True, + prefill_selector={'app': 'prefill'}, + decode_selector={'app': 'decode'}, + ) + assert args.lmdeploy_pd_disaggregation is True + assert args.service_discovery is True + assert args.prefill_selector == {'app': 'prefill'} + assert args.decode_selector == {'app': 'decode'} + + def test_prometheus_config_validation(self): + """Test Prometheus configuration validation.""" + # Valid Prometheus config + args = RouterArgs(prometheus_port=29000, prometheus_host='127.0.0.1') + assert args.prometheus_port == 29000 + assert args.prometheus_host == '127.0.0.1' + + def test_cors_config_validation(self): + """Test CORS configuration validation.""" + # Valid CORS config + args = RouterArgs( + cors_allowed_origins=['http://localhost:3000', 'https://example.com'] + ) + assert args.cors_allowed_origins == [ + 'http://localhost:3000', + 'https://example.com', + ] + + def test_tokenizer_config_validation(self): + """Test tokenizer configuration validation.""" + pytest.skip('Tokenizer configuration not available in current implementation') + + def test_api_key_config_validation(self): + """Test API key configuration remains independent of routing mode.""" + args = RouterArgs(api_key='test-api-key') + assert args.api_key == 'test-api-key' + + def test_request_id_headers_validation(self): + """Test request ID headers configuration validation.""" + # Valid request ID headers config + args = RouterArgs( + request_id_headers=['x-request-id', 'x-trace-id', 'x-correlation-id'] + ) + assert args.request_id_headers == [ + 'x-request-id', + 'x-trace-id', + 'x-correlation-id', + ] + + def test_policy_consistency_validation(self): + """Test policy consistency validation in PD mode.""" + # Test with both prefill and decode policies specified + args = RouterArgs( + lmdeploy_pd_disaggregation=True, + prefill_urls=['http://prefill1:8000'], + decode_urls=['http://decode1:8001'], + policy='cache_aware', + prefill_policy='power_of_two', + decode_policy='round_robin', + ) + + # Should not raise validation error + with patch('lmdeploy_router.launch_router.Router') as router_mod: + mock_router_instance = MagicMock() + router_mod.from_args = MagicMock(return_value=mock_router_instance) + + launch_router(args) + + # Should create router instance via from_args + router_mod.from_args.assert_called_once() + + def test_policy_fallback_validation(self): + """Test policy fallback validation in PD mode.""" + # Test with only prefill policy specified + args = RouterArgs( + lmdeploy_pd_disaggregation=True, + prefill_urls=['http://prefill1:8000'], + decode_urls=['http://decode1:8001'], + policy='cache_aware', + prefill_policy='power_of_two', + decode_policy=None, + ) + + # Should not raise validation error + with patch('lmdeploy_router.launch_router.Router') as router_mod: + mock_router_instance = MagicMock() + router_mod.from_args = MagicMock(return_value=mock_router_instance) + + launch_router(args) + + # Should create router instance via from_args + router_mod.from_args.assert_called_once() + + def test_policy_enum_conversion(self): + """Test policy string to enum conversion.""" + # Test all valid policy conversions + assert policy_from_str('random') == PolicyType.Random + assert policy_from_str('round_robin') == PolicyType.RoundRobin + assert policy_from_str('cache_aware') == PolicyType.CacheAware + assert policy_from_str('power_of_two') == PolicyType.PowerOfTwo + assert policy_from_str('consistent_hash') == PolicyType.ConsistentHash + + def test_invalid_policy_enum_conversion(self): + """Test invalid policy string to enum conversion.""" + with pytest.raises(KeyError): + policy_from_str('invalid_policy') + + def test_config_immutability(self): + """Test that configuration objects are properly immutable.""" + args = RouterArgs( + host='127.0.0.1', port=30000, worker_urls=['http://worker1:8000'] + ) + + # Test that we can't modify the configuration after creation + # (This is more of a design test - dataclasses are mutable by default) + original_host = args.host + args.host = '0.0.0.0' + assert args.host == '0.0.0.0' # Dataclasses are mutable + assert args.host != original_host + + def test_config_defaults_consistency(self): + """Test that configuration defaults are consistent.""" + args1 = RouterArgs() + args2 = RouterArgs() + + # Both instances should have the same defaults + assert args1.host == args2.host + assert args1.port == args2.port + assert args1.policy == args2.policy + assert args1.worker_urls == args2.worker_urls + assert args1.lmdeploy_pd_disaggregation == args2.lmdeploy_pd_disaggregation + + def test_config_serialization(self): + """Test that configuration can be serialized/deserialized.""" + args = RouterArgs( + host='127.0.0.1', + port=30000, + worker_urls=['http://worker1:8000'], + policy='cache_aware', + cache_threshold=0.5, + ) + + # Test that we can access all attributes + assert hasattr(args, 'host') + assert hasattr(args, 'port') + assert hasattr(args, 'worker_urls') + assert hasattr(args, 'policy') + assert hasattr(args, 'cache_threshold') + + def test_config_with_none_values(self): + """Test configuration with None values.""" + args = RouterArgs( + api_key=None, + log_dir=None, + log_level=None, + prometheus_port=None, + prometheus_host=None, + request_id_headers=None, + rate_limit_tokens_per_second=None, + service_discovery_namespace=None, + ) + + # All None values should be preserved + assert args.api_key is None + assert args.log_dir is None + assert args.log_level is None + assert args.prometheus_port is None + assert args.prometheus_host is None + assert args.request_id_headers is None + assert args.rate_limit_tokens_per_second is None + assert args.service_discovery_namespace is None + + def test_config_with_empty_lists(self): + """Test configuration with empty lists.""" + args = RouterArgs( + worker_urls=[], prefill_urls=[], decode_urls=[], cors_allowed_origins=[] + ) + + # All empty lists should be preserved + assert args.worker_urls == [] + assert args.prefill_urls == [] + assert args.decode_urls == [] + assert args.cors_allowed_origins == [] + + def test_config_with_empty_dicts(self): + """Test configuration with empty dictionaries.""" + args = RouterArgs(selector={}, prefill_selector={}, decode_selector={}) + + # All empty dictionaries should be preserved + assert args.selector == {} + assert args.prefill_selector == {} + assert args.decode_selector == {} diff --git a/lmdeploy/router/py_test/unit/test_startup_sequence.py b/lmdeploy/router/py_test/unit/test_startup_sequence.py new file mode 100644 index 0000000000..12fd2b84ee --- /dev/null +++ b/lmdeploy/router/py_test/unit/test_startup_sequence.py @@ -0,0 +1,820 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""Unit tests for startup sequence logic in lmdeploy_router. + +These tests focus on testing the startup sequence logic in isolation, including router initialization, configuration +validation, and startup flow. +""" + +import logging +from unittest.mock import MagicMock, patch + +import pytest +from lmdeploy_router.launch_router import RouterArgs, launch_router +from lmdeploy_router.router import policy_from_str +from lmdeploy_router_rs import PolicyType + + +# Local helper mirroring the router logger setup used in production +def setup_logger(): + logger = logging.getLogger('router') + logger.setLevel(logging.INFO) + if not logger.handlers: + formatter = logging.Formatter( + '[Router (Python)] %(asctime)s - %(levelname)s - %(message)s', + datefmt='%Y-%m-%d %H:%M:%S', + ) + handler = logging.StreamHandler() + handler.setFormatter(formatter) + logger.addHandler(handler) + return logger + + +class TestSetupLogger: + """Test logger setup functionality.""" + + def test_setup_logger_returns_logger(self): + """Test that setup_logger returns a logger instance.""" + logger = setup_logger() + + assert isinstance(logger, logging.Logger) + assert logger.name == 'router' + assert logger.level == logging.INFO + + def test_setup_logger_has_handler(self): + """Test that setup_logger configures a handler.""" + logger = setup_logger() + + assert len(logger.handlers) > 0 + handler = logger.handlers[0] + assert isinstance(handler, logging.StreamHandler) + + def test_setup_logger_has_formatter(self): + """Test that setup_logger configures a formatter.""" + logger = setup_logger() + + handler = logger.handlers[0] + formatter = handler.formatter + + assert formatter is not None + assert '[Router (Python)]' in formatter._fmt + + def test_setup_logger_multiple_calls(self): + """Test that multiple calls to setup_logger work correctly.""" + logger1 = setup_logger() + logger2 = setup_logger() + + # Should return the same logger instance + assert logger1 is logger2 + + +class TestPolicyFromStr: + """Test policy string to enum conversion in startup context.""" + + def test_policy_conversion_in_startup(self): + """Test policy conversion during startup sequence.""" + # Test all valid policies + policies = [ + 'random', + 'round_robin', + 'cache_aware', + 'power_of_two', + 'consistent_hash', + ] + expected_enums = [ + PolicyType.Random, + PolicyType.RoundRobin, + PolicyType.CacheAware, + PolicyType.PowerOfTwo, + PolicyType.ConsistentHash, + ] + + for policy_str, expected_enum in zip(policies, expected_enums): + result = policy_from_str(policy_str) + assert result == expected_enum + + def test_invalid_policy_in_startup(self): + """Test handling of invalid policy during startup.""" + with pytest.raises(KeyError): + policy_from_str('invalid_policy') + + +class TestRouterInitialization: + """Test router initialization logic.""" + + def test_router_initialization_basic(self): + """Test basic router initialization.""" + args = RouterArgs( + host='127.0.0.1', + port=30000, + worker_urls=['http://worker1:8000'], + policy='cache_aware', + ) + + with patch('lmdeploy_router.launch_router.Router') as router_mod: + captured_args = {} + + mock_router_instance = MagicMock() + + def fake_from_args(router_args): + # capture needed fields from RouterArgs + captured_args.update( + dict( + host=router_args.host, + port=router_args.port, + worker_urls=router_args.worker_urls, + policy=policy_from_str(router_args.policy), + ) + ) + return mock_router_instance + + router_mod.from_args = MagicMock(side_effect=fake_from_args) + + launch_router(args) + + # Verify Router.from_args was called and captured fields match + router_mod.from_args.assert_called_once() + assert captured_args['host'] == '127.0.0.1' + assert captured_args['port'] == 30000 + assert captured_args['worker_urls'] == ['http://worker1:8000'] + assert captured_args['policy'] == PolicyType.CacheAware + + # Verify router.start() was called + mock_router_instance.start.assert_called_once() + + # Function returns None; ensure start was invoked + + def test_router_initialization_pd_mode(self): + """Test router initialization in PD mode.""" + args = RouterArgs( + lmdeploy_pd_disaggregation=True, + prefill_urls=['http://prefill1:8000'], + decode_urls=['http://decode1:8001'], + policy='power_of_two', + ) + + with patch('lmdeploy_router.launch_router.Router') as router_mod: + captured_args = {} + mock_router_instance = MagicMock() + + def fake_from_args(router_args): + captured_args.update( + dict( + lmdeploy_pd_disaggregation=router_args.lmdeploy_pd_disaggregation, + prefill_urls=router_args.prefill_urls, + decode_urls=router_args.decode_urls, + policy=policy_from_str(router_args.policy), + ) + ) + return mock_router_instance + + router_mod.from_args = MagicMock(side_effect=fake_from_args) + + launch_router(args) + + # Verify Router.from_args was called with PD parameters + router_mod.from_args.assert_called_once() + assert captured_args['lmdeploy_pd_disaggregation'] is True + assert captured_args['prefill_urls'] == ['http://prefill1:8000'] + assert captured_args['decode_urls'] == ['http://decode1:8001'] + assert captured_args['policy'] == PolicyType.PowerOfTwo + + # Verify router.start() was called + mock_router_instance.start.assert_called_once() + + # Function returns None; ensure start was invoked + + def test_router_initialization_with_service_discovery(self): + """Test router initialization with service discovery.""" + args = RouterArgs( + service_discovery=True, + selector={'app': 'worker', 'env': 'prod'}, + service_discovery_port=8080, + service_discovery_namespace='default', + ) + + with patch('lmdeploy_router.launch_router.Router') as router_mod: + captured_args = {} + mock_router_instance = MagicMock() + + def fake_from_args(router_args): + captured_args.update( + dict( + service_discovery=router_args.service_discovery, + selector=router_args.selector, + service_discovery_port=router_args.service_discovery_port, + service_discovery_namespace=router_args.service_discovery_namespace, + ) + ) + return mock_router_instance + + router_mod.from_args = MagicMock(side_effect=fake_from_args) + + launch_router(args) + + # Verify Router.from_args was called with service discovery parameters + router_mod.from_args.assert_called_once() + assert captured_args['service_discovery'] is True + assert captured_args['selector'] == {'app': 'worker', 'env': 'prod'} + assert captured_args['service_discovery_port'] == 8080 + assert captured_args['service_discovery_namespace'] == 'default' + + # Verify router.start() was called + mock_router_instance.start.assert_called_once() + + # Function returns None; ensure start was invoked + + def test_router_initialization_with_retry_config(self): + """Test router initialization with retry configuration.""" + args = RouterArgs( + retry_max_retries=3, + retry_initial_backoff_ms=100, + retry_max_backoff_ms=10000, + retry_backoff_multiplier=2.0, + retry_jitter_factor=0.1, + disable_retries=False, + ) + + with patch('lmdeploy_router.launch_router.Router') as router_mod: + captured_args = {} + mock_router_instance = MagicMock() + + def fake_from_args(router_args): + captured_args.update( + dict( + retry_max_retries=router_args.retry_max_retries, + retry_initial_backoff_ms=router_args.retry_initial_backoff_ms, + retry_max_backoff_ms=router_args.retry_max_backoff_ms, + retry_backoff_multiplier=router_args.retry_backoff_multiplier, + retry_jitter_factor=router_args.retry_jitter_factor, + disable_retries=router_args.disable_retries, + ) + ) + return mock_router_instance + + router_mod.from_args = MagicMock(side_effect=fake_from_args) + + launch_router(args) + + # Verify router was created with retry parameters + router_mod.from_args.assert_called_once() + assert captured_args['retry_max_retries'] == 3 + assert captured_args['retry_initial_backoff_ms'] == 100 + assert captured_args['retry_max_backoff_ms'] == 10000 + assert captured_args['retry_backoff_multiplier'] == 2.0 + assert captured_args['retry_jitter_factor'] == 0.1 + assert captured_args['disable_retries'] is False + + # Verify router.start() was called + mock_router_instance.start.assert_called_once() + + # Function returns None; ensure start was invoked + + def test_router_initialization_with_circuit_breaker_config(self): + """Test router initialization with circuit breaker configuration.""" + args = RouterArgs( + cb_failure_threshold=5, + cb_success_threshold=2, + cb_timeout_duration_secs=30, + cb_window_duration_secs=60, + disable_circuit_breaker=False, + ) + + with patch('lmdeploy_router.launch_router.Router') as router_mod: + captured_args = {} + mock_router_instance = MagicMock() + + def fake_from_args(router_args): + captured_args.update( + dict( + cb_failure_threshold=router_args.cb_failure_threshold, + cb_success_threshold=router_args.cb_success_threshold, + cb_timeout_duration_secs=router_args.cb_timeout_duration_secs, + cb_window_duration_secs=router_args.cb_window_duration_secs, + disable_circuit_breaker=router_args.disable_circuit_breaker, + ) + ) + return mock_router_instance + + router_mod.from_args = MagicMock(side_effect=fake_from_args) + + launch_router(args) + + # Verify router was created with circuit breaker parameters + router_mod.from_args.assert_called_once() + assert captured_args['cb_failure_threshold'] == 5 + assert captured_args['cb_success_threshold'] == 2 + assert captured_args['cb_timeout_duration_secs'] == 30 + assert captured_args['cb_window_duration_secs'] == 60 + assert captured_args['disable_circuit_breaker'] is False + + # Verify router.start() was called + mock_router_instance.start.assert_called_once() + + # Function returns None; ensure start was invoked + + def test_router_initialization_with_rate_limiting_config(self): + """Test router initialization with rate limiting configuration.""" + args = RouterArgs( + max_concurrent_requests=512, + queue_size=200, + queue_timeout_secs=120, + rate_limit_tokens_per_second=100, + ) + + with patch('lmdeploy_router.launch_router.Router') as router_mod: + captured_args = {} + mock_router_instance = MagicMock() + + def fake_from_args(router_args): + captured_args.update( + dict( + max_concurrent_requests=router_args.max_concurrent_requests, + queue_size=router_args.queue_size, + queue_timeout_secs=router_args.queue_timeout_secs, + rate_limit_tokens_per_second=router_args.rate_limit_tokens_per_second, + ) + ) + return mock_router_instance + + router_mod.from_args = MagicMock(side_effect=fake_from_args) + + launch_router(args) + + # Verify router was created with rate limiting parameters + router_mod.from_args.assert_called_once() + assert captured_args['max_concurrent_requests'] == 512 + assert captured_args['queue_size'] == 200 + assert captured_args['queue_timeout_secs'] == 120 + assert captured_args['rate_limit_tokens_per_second'] == 100 + + # Verify router.start() was called + mock_router_instance.start.assert_called_once() + + # Function returns None; ensure start was invoked + + def test_router_initialization_with_health_check_config(self): + """Test router initialization with health check configuration.""" + args = RouterArgs( + health_failure_threshold=2, + health_success_threshold=1, + health_check_timeout_secs=3, + health_check_interval_secs=30, + health_check_endpoint='/healthz', + ) + + with patch('lmdeploy_router.launch_router.Router') as router_mod: + captured_args = {} + mock_router_instance = MagicMock() + + def fake_from_args(router_args): + captured_args.update( + dict( + health_failure_threshold=router_args.health_failure_threshold, + health_success_threshold=router_args.health_success_threshold, + health_check_timeout_secs=router_args.health_check_timeout_secs, + health_check_interval_secs=router_args.health_check_interval_secs, + health_check_endpoint=router_args.health_check_endpoint, + ) + ) + return mock_router_instance + + router_mod.from_args = MagicMock(side_effect=fake_from_args) + + launch_router(args) + + # Verify router was created with health check parameters + router_mod.from_args.assert_called_once() + assert captured_args['health_failure_threshold'] == 2 + assert captured_args['health_success_threshold'] == 1 + assert captured_args['health_check_timeout_secs'] == 3 + assert captured_args['health_check_interval_secs'] == 30 + assert captured_args['health_check_endpoint'] == '/healthz' + + # Verify router.start() was called + mock_router_instance.start.assert_called_once() + + # Function returns None; ensure start was invoked + + def test_router_initialization_with_prometheus_config(self): + """Test router initialization with Prometheus configuration.""" + args = RouterArgs(prometheus_port=29000, prometheus_host='127.0.0.1') + + with patch('lmdeploy_router.launch_router.Router') as router_mod: + captured_args = {} + mock_router_instance = MagicMock() + + def fake_from_args(router_args): + captured_args.update( + dict( + prometheus_port=router_args.prometheus_port, + prometheus_host=router_args.prometheus_host, + ) + ) + return mock_router_instance + + router_mod.from_args = MagicMock(side_effect=fake_from_args) + + launch_router(args) + + # Verify router was created with Prometheus parameters + router_mod.from_args.assert_called_once() + assert captured_args['prometheus_port'] == 29000 + assert captured_args['prometheus_host'] == '127.0.0.1' + + # Verify router.start() was called + mock_router_instance.start.assert_called_once() + + # Function returns None; ensure start was invoked + + def test_router_initialization_with_cors_config(self): + """Test router initialization with CORS configuration.""" + args = RouterArgs( + cors_allowed_origins=['http://localhost:3000', 'https://example.com'] + ) + + with patch('lmdeploy_router.launch_router.Router') as router_mod: + captured_args = {} + mock_router_instance = MagicMock() + + def fake_from_args(router_args): + captured_args.update( + dict(cors_allowed_origins=router_args.cors_allowed_origins) + ) + return mock_router_instance + + router_mod.from_args = MagicMock(side_effect=fake_from_args) + + launch_router(args) + + # Verify router was created with CORS parameters + router_mod.from_args.assert_called_once() + assert captured_args['cors_allowed_origins'] == [ + 'http://localhost:3000', + 'https://example.com', + ] + + # Verify router.start() was called + mock_router_instance.start.assert_called_once() + + # Function returns None; ensure start was invoked + + def test_router_initialization_with_tokenizer_config(self): + """Test router initialization with tokenizer configuration.""" + pytest.skip('Tokenizer configuration not available in current implementation') + + +class TestStartupValidation: + """Test startup validation logic.""" + + def test_pd_mode_without_urls_allows_dynamic_registration(self): + """Test that PD mode allows workers to register after startup.""" + args = RouterArgs( + lmdeploy_pd_disaggregation=True, + prefill_urls=[], + decode_urls=[], + service_discovery=False, + ) + + with patch('lmdeploy_router.launch_router.Router') as router_mod: + mock_router_instance = MagicMock() + router_mod.from_args = MagicMock(return_value=mock_router_instance) + + launch_router(args) + router_mod.from_args.assert_called_once() + + def test_pd_mode_with_service_discovery_validation(self): + """Test PD mode with service discovery validation during startup.""" + args = RouterArgs( + lmdeploy_pd_disaggregation=True, + prefill_urls=[], + decode_urls=[], + service_discovery=True, + ) + + # Should not raise validation error + with patch('lmdeploy_router.launch_router.Router') as router_mod: + mock_router_instance = MagicMock() + + router_mod.from_args = MagicMock(return_value=mock_router_instance) + + launch_router(args) + + # Should create router instance + router_mod.from_args.assert_called_once() + + def test_policy_warning_during_startup(self): + """Test policy warning during startup in PD mode.""" + args = RouterArgs( + lmdeploy_pd_disaggregation=True, + prefill_urls=['http://prefill1:8000'], + decode_urls=['http://decode1:8001'], + policy='cache_aware', + prefill_policy='power_of_two', + decode_policy='round_robin', + ) + + with patch('lmdeploy_router.launch_router.Router') as router_mod: + mock_router_instance = MagicMock() + router_mod.from_args = MagicMock(return_value=mock_router_instance) + + # The policy messages are emitted by router_args logger + with patch('lmdeploy_router.router_args.logger') as mock_logger: + launch_router(args) + + # Should log warning about policy usage + mock_logger.warning.assert_called_once() + warning_call = mock_logger.warning.call_args[0][0] + assert ( + 'Both --prefill-policy and --decode-policy are specified' + in warning_call + ) + + # Should create router instance + router_mod.from_args.assert_called_once() + + def test_policy_info_during_startup(self): + """Test policy info logging during startup in PD mode.""" + # Test with only prefill policy specified + args = RouterArgs( + lmdeploy_pd_disaggregation=True, + prefill_urls=['http://prefill1:8000'], + decode_urls=['http://decode1:8001'], + policy='cache_aware', + prefill_policy='power_of_two', + decode_policy=None, + ) + + with patch('lmdeploy_router.launch_router.Router') as router_mod: + mock_router_instance = MagicMock() + router_mod.from_args = MagicMock(return_value=mock_router_instance) + + # The policy messages are emitted by router_args logger + with patch('lmdeploy_router.router_args.logger') as mock_logger: + launch_router(args) + + # Should log info about policy usage + mock_logger.info.assert_called_once() + info_call = mock_logger.info.call_args[0][0] + assert "Using --prefill-policy 'power_of_two'" in info_call + assert "and --policy 'cache_aware'" in info_call + + # Should create router instance + router_mod.from_args.assert_called_once() + + def test_policy_info_decode_only_during_startup(self): + """Test policy info logging during startup with only decode policy + specified.""" + args = RouterArgs( + lmdeploy_pd_disaggregation=True, + prefill_urls=['http://prefill1:8000'], + decode_urls=['http://decode1:8001'], + policy='cache_aware', + prefill_policy=None, + decode_policy='round_robin', + ) + + with patch('lmdeploy_router.launch_router.Router') as router_mod: + mock_router_instance = MagicMock() + router_mod.from_args = MagicMock(return_value=mock_router_instance) + + # The policy messages are emitted by router_args logger + with patch('lmdeploy_router.router_args.logger') as mock_logger: + launch_router(args) + + # Should log info about policy usage + mock_logger.info.assert_called_once() + info_call = mock_logger.info.call_args[0][0] + assert "Using --policy 'cache_aware'" in info_call + assert "and --decode-policy 'round_robin'" in info_call + + # Should create router instance + router_mod.from_args.assert_called_once() + + +class TestStartupErrorHandling: + """Test startup error handling logic.""" + + def test_router_creation_error_handling(self): + """Test error handling when router creation fails.""" + args = RouterArgs( + host='127.0.0.1', port=30000, worker_urls=['http://worker1:8000'] + ) + + with patch('lmdeploy_router.launch_router.Router') as router_mod: + # Simulate router creation failure in from_args + router_mod.from_args = MagicMock( + side_effect=Exception('Router creation failed') + ) + + with patch('lmdeploy_router.launch_router.logger') as mock_logger: + with pytest.raises(Exception, match='Router creation failed'): + launch_router(args) + + # Should log error + mock_logger.error.assert_called_once() + error_call = mock_logger.error.call_args[0][0] + assert 'Error starting router: Router creation failed' in error_call + + def test_router_start_error_handling(self): + """Test error handling when router start fails.""" + args = RouterArgs( + host='127.0.0.1', port=30000, worker_urls=['http://worker1:8000'] + ) + + with patch('lmdeploy_router.launch_router.Router') as router_mod: + mock_router_instance = MagicMock() + router_mod.from_args = MagicMock(return_value=mock_router_instance) + + # Simulate router start failure + mock_router_instance.start.side_effect = Exception('Router start failed') + + with patch('lmdeploy_router.launch_router.logger') as mock_logger: + with pytest.raises(Exception, match='Router start failed'): + launch_router(args) + + # Should log error + mock_logger.error.assert_called_once() + error_call = mock_logger.error.call_args[0][0] + assert 'Error starting router: Router start failed' in error_call + + +class TestStartupFlow: + """Test complete startup flow.""" + + def test_complete_startup_flow_basic(self): + """Test complete startup flow for basic configuration.""" + args = RouterArgs( + host='127.0.0.1', + port=30000, + worker_urls=['http://worker1:8000', 'http://worker2:8000'], + policy='cache_aware', + cache_threshold=0.5, + balance_abs_threshold=32, + balance_rel_threshold=1.5, + ) + + with patch('lmdeploy_router.launch_router.Router') as router_mod: + mock_router_instance = MagicMock() + router_mod.from_args = MagicMock(return_value=mock_router_instance) + + launch_router(args) + + # Verify complete flow + router_mod.from_args.assert_called_once() + mock_router_instance.start.assert_called_once() + + def test_complete_startup_flow_pd_mode(self): + """Test complete startup flow for PD mode configuration.""" + args = RouterArgs( + lmdeploy_pd_disaggregation=True, + prefill_urls=[ + 'http://prefill1:8000', + 'http://prefill2:8000', + ], + decode_urls=['http://decode1:8001', 'http://decode2:8001'], + policy='power_of_two', + prefill_policy='cache_aware', + decode_policy='round_robin', + ) + + with patch('lmdeploy_router.launch_router.Router') as router_mod: + mock_router_instance = MagicMock() + router_mod.from_args = MagicMock(return_value=mock_router_instance) + + with patch('lmdeploy_router.router_args.logger') as mock_logger: + launch_router(args) + + # Verify complete flow + router_mod.from_args.assert_called_once() + mock_router_instance.start.assert_called_once() + + # Verify policy warning was logged + mock_logger.warning.assert_called_once() + + def test_complete_startup_flow_with_all_features(self): + """Test complete startup flow with all features enabled.""" + args = RouterArgs( + host='0.0.0.0', + port=30001, + worker_urls=['http://worker1:8000'], + policy='round_robin', + service_discovery=True, + selector={'app': 'worker'}, + service_discovery_port=8080, + service_discovery_namespace='default', + api_key='test-key', + log_dir='/tmp/logs', + log_level='debug', + prometheus_port=29000, + prometheus_host='0.0.0.0', + request_id_headers=['x-request-id', 'x-trace-id'], + request_timeout_secs=1200, + max_concurrent_requests=512, + queue_size=200, + queue_timeout_secs=120, + rate_limit_tokens_per_second=100, + cors_allowed_origins=['http://localhost:3000'], + retry_max_retries=3, + retry_initial_backoff_ms=100, + retry_max_backoff_ms=10000, + retry_backoff_multiplier=2.0, + retry_jitter_factor=0.1, + cb_failure_threshold=5, + cb_success_threshold=2, + cb_timeout_duration_secs=30, + cb_window_duration_secs=60, + health_failure_threshold=2, + health_success_threshold=1, + health_check_timeout_secs=3, + health_check_interval_secs=30, + health_check_endpoint='/healthz', + ) + + with patch('lmdeploy_router.launch_router.Router') as router_mod: + captured_args = {} + mock_router_instance = MagicMock() + + def fake_from_args(router_args): + captured_args.update( + dict( + host=router_args.host, + port=router_args.port, + worker_urls=router_args.worker_urls, + policy=policy_from_str(router_args.policy), + service_discovery=router_args.service_discovery, + selector=router_args.selector, + service_discovery_port=router_args.service_discovery_port, + service_discovery_namespace=router_args.service_discovery_namespace, + api_key=router_args.api_key, + log_dir=router_args.log_dir, + log_level=router_args.log_level, + prometheus_port=router_args.prometheus_port, + prometheus_host=router_args.prometheus_host, + request_id_headers=router_args.request_id_headers, + request_timeout_secs=router_args.request_timeout_secs, + max_concurrent_requests=router_args.max_concurrent_requests, + queue_size=router_args.queue_size, + queue_timeout_secs=router_args.queue_timeout_secs, + rate_limit_tokens_per_second=router_args.rate_limit_tokens_per_second, + cors_allowed_origins=router_args.cors_allowed_origins, + retry_max_retries=router_args.retry_max_retries, + retry_initial_backoff_ms=router_args.retry_initial_backoff_ms, + retry_max_backoff_ms=router_args.retry_max_backoff_ms, + retry_backoff_multiplier=router_args.retry_backoff_multiplier, + retry_jitter_factor=router_args.retry_jitter_factor, + cb_failure_threshold=router_args.cb_failure_threshold, + cb_success_threshold=router_args.cb_success_threshold, + cb_timeout_duration_secs=router_args.cb_timeout_duration_secs, + cb_window_duration_secs=router_args.cb_window_duration_secs, + health_failure_threshold=router_args.health_failure_threshold, + health_success_threshold=router_args.health_success_threshold, + health_check_timeout_secs=router_args.health_check_timeout_secs, + health_check_interval_secs=router_args.health_check_interval_secs, + health_check_endpoint=router_args.health_check_endpoint, + ) + ) + return mock_router_instance + + router_mod.from_args = MagicMock(side_effect=fake_from_args) + + launch_router(args) + + # Verify complete flow + router_mod.from_args.assert_called_once() + mock_router_instance.start.assert_called_once() + + # Verify key parameters were propagated into RouterArgs + assert captured_args['host'] == '0.0.0.0' + assert captured_args['port'] == 30001 + assert captured_args['worker_urls'] == ['http://worker1:8000'] + assert captured_args['policy'] == PolicyType.RoundRobin + assert captured_args['service_discovery'] is True + assert captured_args['selector'] == {'app': 'worker'} + assert captured_args['service_discovery_port'] == 8080 + assert captured_args['service_discovery_namespace'] == 'default' + assert captured_args['api_key'] == 'test-key' + assert captured_args['log_dir'] == '/tmp/logs' + assert captured_args['log_level'] == 'debug' + assert captured_args['prometheus_port'] == 29000 + assert captured_args['prometheus_host'] == '0.0.0.0' + assert captured_args['request_id_headers'] == ['x-request-id', 'x-trace-id'] + assert captured_args['request_timeout_secs'] == 1200 + assert captured_args['max_concurrent_requests'] == 512 + assert captured_args['queue_size'] == 200 + assert captured_args['queue_timeout_secs'] == 120 + assert captured_args['rate_limit_tokens_per_second'] == 100 + assert captured_args['cors_allowed_origins'] == ['http://localhost:3000'] + assert captured_args['retry_max_retries'] == 3 + assert captured_args['retry_initial_backoff_ms'] == 100 + assert captured_args['retry_max_backoff_ms'] == 10000 + assert captured_args['retry_backoff_multiplier'] == 2.0 + assert captured_args['retry_jitter_factor'] == 0.1 + assert captured_args['cb_failure_threshold'] == 5 + assert captured_args['cb_success_threshold'] == 2 + assert captured_args['cb_timeout_duration_secs'] == 30 + assert captured_args['cb_window_duration_secs'] == 60 + assert captured_args['health_failure_threshold'] == 2 + assert captured_args['health_success_threshold'] == 1 + assert captured_args['health_check_timeout_secs'] == 3 + assert captured_args['health_check_interval_secs'] == 30 + assert captured_args['health_check_endpoint'] == '/healthz' diff --git a/lmdeploy/router/py_test/unit/test_validation.py b/lmdeploy/router/py_test/unit/test_validation.py new file mode 100644 index 0000000000..05d0442445 --- /dev/null +++ b/lmdeploy/router/py_test/unit/test_validation.py @@ -0,0 +1,496 @@ +# Copyright (c) OpenMMLab. All rights reserved. +"""Unit tests for validation logic in lmdeploy_router. + +These tests focus on testing the validation logic in isolation, including parameter validation, URL validation, and +configuration validation. +""" + +from unittest.mock import MagicMock, patch + +import pytest +from lmdeploy_router.launch_router import RouterArgs, launch_router + + +class TestURLValidation: + """Test URL validation logic.""" + + def test_valid_worker_urls(self): + """Test validation of valid worker URLs.""" + valid_urls = [ + 'http://worker1:8000', + 'https://worker2:8000', + 'http://localhost:8000', + 'http://127.0.0.1:8000', + 'http://192.168.1.100:8000', + 'http://worker.example.com:8000', + ] + + for url in valid_urls: + args = RouterArgs(worker_urls=[url]) + # Should not raise any validation errors + assert url in args.worker_urls + + def test_valid_prefill_urls(self): + """Test validation of valid prefill URLs.""" + valid_prefill_urls = [ + 'http://prefill1:8000', + 'https://prefill2:8000', + 'http://localhost:8000', + 'http://127.0.0.1:8000', + ] + + for url in valid_prefill_urls: + args = RouterArgs(prefill_urls=[url]) + # Should not raise any validation errors + assert url in args.prefill_urls + + def test_valid_decode_urls(self): + """Test validation of valid decode URLs.""" + valid_decode_urls = [ + 'http://decode1:8001', + 'https://decode2:8001', + 'http://localhost:8001', + 'http://127.0.0.1:8001', + ] + + for url in valid_decode_urls: + args = RouterArgs(decode_urls=[url]) + # Should not raise any validation errors + assert url in args.decode_urls + + def test_malformed_urls(self): + """Test handling of malformed URLs.""" + # Note: The current implementation doesn't validate URL format + # This test documents the current behavior + malformed_urls = [ + 'not-a-url', + 'ftp://worker1:8000', # Wrong protocol + 'http://', # Missing host + ':8000', # Missing protocol and host + 'http://worker1', # Missing port + ] + + for url in malformed_urls: + args = RouterArgs(worker_urls=[url]) + # Currently, malformed URLs are accepted + # This might be something to improve in the future + assert url in args.worker_urls + + +class TestPortValidation: + """Test port validation logic.""" + + def test_valid_ports(self): + """Test validation of valid port numbers.""" + valid_ports = [1, 80, 8000, 30000, 65535] + + for port in valid_ports: + args = RouterArgs(port=port) + assert args.port == port + + def test_invalid_ports(self): + """Test handling of invalid port numbers.""" + # Note: The current implementation doesn't validate port ranges + # This test documents the current behavior + invalid_ports = [0, -1, 65536, 70000] + + for port in invalid_ports: + args = RouterArgs(port=port) + # Currently, invalid ports are accepted + # This might be something to improve in the future + assert args.port == port + +class TestParameterValidation: + """Test parameter validation logic.""" + + def test_cache_threshold_validation(self): + """Test cache threshold parameter validation.""" + # Valid cache thresholds + valid_thresholds = [0.0, 0.1, 0.5, 0.9, 1.0] + + for threshold in valid_thresholds: + args = RouterArgs(cache_threshold=threshold) + assert args.cache_threshold == threshold + + def test_balance_threshold_validation(self): + """Test load balancing threshold parameter validation.""" + # Valid absolute thresholds + valid_abs_thresholds = [0, 1, 32, 64, 128, 1000] + for threshold in valid_abs_thresholds: + args = RouterArgs(balance_abs_threshold=threshold) + assert args.balance_abs_threshold == threshold + + # Valid relative thresholds + valid_rel_thresholds = [1.0, 1.1, 1.5, 2.0, 10.0] + for threshold in valid_rel_thresholds: + args = RouterArgs(balance_rel_threshold=threshold) + assert args.balance_rel_threshold == threshold + + def test_timeout_validation(self): + """Test timeout parameter validation.""" + # Valid timeouts + valid_timeouts = [1, 30, 60, 300, 600, 1800, 3600] + + for timeout in valid_timeouts: + args = RouterArgs( + worker_startup_timeout_secs=timeout, + worker_startup_check_interval=timeout, + request_timeout_secs=timeout, + queue_timeout_secs=timeout, + ) + assert args.worker_startup_timeout_secs == timeout + assert args.worker_startup_check_interval == timeout + assert args.request_timeout_secs == timeout + assert args.queue_timeout_secs == timeout + + def test_retry_parameter_validation(self): + """Test retry parameter validation.""" + # Valid retry parameters + valid_retry_counts = [0, 1, 3, 5, 10] + for count in valid_retry_counts: + args = RouterArgs(retry_max_retries=count) + assert args.retry_max_retries == count + + # Valid backoff parameters + valid_backoff_ms = [1, 50, 100, 1000, 30000] + for backoff in valid_backoff_ms: + args = RouterArgs( + retry_initial_backoff_ms=backoff, retry_max_backoff_ms=backoff + ) + assert args.retry_initial_backoff_ms == backoff + assert args.retry_max_backoff_ms == backoff + + # Valid multiplier parameters + valid_multipliers = [1.0, 1.5, 2.0, 3.0] + for multiplier in valid_multipliers: + args = RouterArgs(retry_backoff_multiplier=multiplier) + assert args.retry_backoff_multiplier == multiplier + + # Valid jitter parameters + valid_jitter = [0.0, 0.1, 0.2, 0.5] + for jitter in valid_jitter: + args = RouterArgs(retry_jitter_factor=jitter) + assert args.retry_jitter_factor == jitter + + def test_circuit_breaker_parameter_validation(self): + """Test circuit breaker parameter validation.""" + # Valid failure thresholds + valid_failure_thresholds = [1, 3, 5, 10, 20] + for threshold in valid_failure_thresholds: + args = RouterArgs(cb_failure_threshold=threshold) + assert args.cb_failure_threshold == threshold + + # Valid success thresholds + valid_success_thresholds = [1, 2, 3, 5] + for threshold in valid_success_thresholds: + args = RouterArgs(cb_success_threshold=threshold) + assert args.cb_success_threshold == threshold + + # Valid timeout durations + valid_timeouts = [10, 30, 60, 120, 300] + for timeout in valid_timeouts: + args = RouterArgs( + cb_timeout_duration_secs=timeout, cb_window_duration_secs=timeout + ) + assert args.cb_timeout_duration_secs == timeout + assert args.cb_window_duration_secs == timeout + + def test_health_check_parameter_validation(self): + """Test health check parameter validation.""" + # Valid failure thresholds + valid_failure_thresholds = [1, 2, 3, 5, 10] + for threshold in valid_failure_thresholds: + args = RouterArgs(health_failure_threshold=threshold) + assert args.health_failure_threshold == threshold + + # Valid success thresholds + valid_success_thresholds = [1, 2, 3, 5] + for threshold in valid_success_thresholds: + args = RouterArgs(health_success_threshold=threshold) + assert args.health_success_threshold == threshold + + # Valid timeouts and intervals + valid_times = [1, 5, 10, 30, 60, 120] + for time_val in valid_times: + args = RouterArgs( + health_check_timeout_secs=time_val, health_check_interval_secs=time_val + ) + assert args.health_check_timeout_secs == time_val + assert args.health_check_interval_secs == time_val + + def test_rate_limiting_parameter_validation(self): + """Test rate limiting parameter validation.""" + # Valid concurrent request limits + valid_limits = [1, 10, 64, 256, 512, 1000] + for limit in valid_limits: + args = RouterArgs(max_concurrent_requests=limit) + assert args.max_concurrent_requests == limit + + # Valid queue sizes + valid_queue_sizes = [0, 10, 50, 100, 500, 1000] + for size in valid_queue_sizes: + args = RouterArgs(queue_size=size) + assert args.queue_size == size + + # Valid token rates + valid_rates = [1, 10, 50, 100, 500, 1000] + for rate in valid_rates: + args = RouterArgs(rate_limit_tokens_per_second=rate) + assert args.rate_limit_tokens_per_second == rate + + def test_tree_size_validation(self): + """Test tree size parameter validation.""" + # Valid tree sizes (powers of 2) + valid_sizes = [2**10, 2**20, 2**24, 2**26, 2**28, 2**30] + + for size in valid_sizes: + args = RouterArgs(max_tree_size=size) + assert args.max_tree_size == size + + def test_payload_size_validation(self): + """Test payload size parameter validation.""" + # Valid payload sizes + valid_sizes = [ + 1024, # 1KB + 1024 * 1024, # 1MB + 10 * 1024 * 1024, # 10MB + 100 * 1024 * 1024, # 100MB + 512 * 1024 * 1024, # 512MB + 1024 * 1024 * 1024, # 1GB + ] + + for size in valid_sizes: + args = RouterArgs(max_payload_size=size) + assert args.max_payload_size == size + + +class TestConfigurationValidation: + """Test configuration validation logic.""" + + def test_pd_mode_validation(self): + """Test PD mode configuration validation.""" + # Valid PD configuration + args = RouterArgs( + lmdeploy_pd_disaggregation=True, + prefill_urls=['http://prefill1:8000'], + decode_urls=['http://decode1:8001'], + ) + + assert args.lmdeploy_pd_disaggregation is True + assert len(args.prefill_urls) > 0 + assert len(args.decode_urls) > 0 + + def test_service_discovery_validation(self): + """Test service discovery configuration validation.""" + # Valid service discovery configuration + args = RouterArgs( + service_discovery=True, + selector={'app': 'worker', 'env': 'prod'}, + service_discovery_port=8080, + service_discovery_namespace='default', + ) + + assert args.service_discovery is True + assert args.selector == {'app': 'worker', 'env': 'prod'} + assert args.service_discovery_port == 8080 + assert args.service_discovery_namespace == 'default' + + def test_pd_service_discovery_validation(self): + """Test PD service discovery configuration validation.""" + # Valid PD service discovery configuration + args = RouterArgs( + lmdeploy_pd_disaggregation=True, + service_discovery=True, + prefill_selector={'app': 'prefill'}, + decode_selector={'app': 'decode'}, + ) + + assert args.lmdeploy_pd_disaggregation is True + assert args.service_discovery is True + assert args.prefill_selector == {'app': 'prefill'} + assert args.decode_selector == {'app': 'decode'} + + def test_policy_validation(self): + """Test policy configuration validation.""" + # Valid policies + valid_policies = ['random', 'round_robin', 'cache_aware', 'power_of_two'] + + for policy in valid_policies: + args = RouterArgs(policy=policy) + assert args.policy == policy + + def test_pd_policy_validation(self): + """Test PD policy configuration validation.""" + # Valid PD policies + valid_policies = ['random', 'round_robin', 'cache_aware', 'power_of_two'] + + for prefill_policy in valid_policies: + for decode_policy in valid_policies: + args = RouterArgs( + lmdeploy_pd_disaggregation=True, + prefill_urls=['http://prefill1:8000'], + decode_urls=['http://decode1:8001'], + prefill_policy=prefill_policy, + decode_policy=decode_policy, + ) + assert args.prefill_policy == prefill_policy + assert args.decode_policy == decode_policy + + def test_cors_validation(self): + """Test CORS configuration validation.""" + # Valid CORS origins + valid_origins = [ + [], + ['http://localhost:3000'], + ['https://example.com'], + ['http://localhost:3000', 'https://example.com'], + ['*'], # Wildcard (if supported) + ] + + for origins in valid_origins: + args = RouterArgs(cors_allowed_origins=origins) + assert args.cors_allowed_origins == origins + + def test_logging_validation(self): + """Test logging configuration validation.""" + # Valid log levels + valid_log_levels = ['debug', 'info', 'warning', 'error', 'critical'] + + for level in valid_log_levels: + args = RouterArgs(log_level=level) + assert args.log_level == level + + def test_prometheus_validation(self): + """Test Prometheus configuration validation.""" + # Valid Prometheus configuration + args = RouterArgs(prometheus_port=29000, prometheus_host='127.0.0.1') + + assert args.prometheus_port == 29000 + assert args.prometheus_host == '127.0.0.1' + + def test_tokenizer_validation(self): + """Test tokenizer configuration validation.""" + pytest.skip('Tokenizer configuration not available in current implementation') + + def test_request_id_headers_validation(self): + """Test request ID headers configuration validation.""" + # Valid request ID headers + valid_headers = [ + ['x-request-id'], + ['x-request-id', 'x-trace-id'], + ['x-request-id', 'x-trace-id', 'x-correlation-id'], + ['custom-header'], + ] + + for headers in valid_headers: + args = RouterArgs(request_id_headers=headers) + assert args.request_id_headers == headers + + +class TestLaunchValidation: + """Test launch-time validation logic.""" + + def test_pd_mode_without_urls_allows_dynamic_registration(self): + """Test that PD mode allows workers to register after startup.""" + args = RouterArgs( + lmdeploy_pd_disaggregation=True, + prefill_urls=[], + decode_urls=[], + service_discovery=False, + ) + + with patch('lmdeploy_router.launch_router.Router') as router_mod: + mock_router_instance = MagicMock() + router_mod.from_args = MagicMock(return_value=mock_router_instance) + + launch_router(args) + router_mod.from_args.assert_called_once() + + def test_pd_mode_with_service_discovery_allows_empty_urls(self): + """Test that PD mode with service discovery allows empty URLs.""" + args = RouterArgs( + lmdeploy_pd_disaggregation=True, + prefill_urls=[], + decode_urls=[], + service_discovery=True, + ) + + # Should not raise validation error + with patch('lmdeploy_router.launch_router.Router') as router_mod: + mock_router_instance = MagicMock() + router_mod.from_args = MagicMock(return_value=mock_router_instance) + + launch_router(args) + + # Should create router instance via from_args + router_mod.from_args.assert_called_once() + + def test_regular_mode_allows_empty_worker_urls(self): + """Test that regular mode allows empty worker URLs.""" + args = RouterArgs(worker_urls=[], service_discovery=False) + + # Should not raise validation error + with patch('lmdeploy_router.launch_router.Router') as router_mod: + mock_router_instance = MagicMock() + router_mod.from_args = MagicMock(return_value=mock_router_instance) + + launch_router(args) + + # Should create router instance via from_args + router_mod.from_args.assert_called_once() + + def test_launch_with_valid_config(self): + """Test launching with valid configuration.""" + args = RouterArgs( + host='127.0.0.1', + port=30000, + worker_urls=['http://worker1:8000'], + policy='cache_aware', + ) + + # Should not raise validation error + with patch('lmdeploy_router.launch_router.Router') as router_mod: + mock_router_instance = MagicMock() + router_mod.from_args = MagicMock(return_value=mock_router_instance) + + launch_router(args) + + # Should create router instance via from_args + router_mod.from_args.assert_called_once() + + def test_launch_with_pd_config(self): + """Test launching with valid PD configuration.""" + args = RouterArgs( + lmdeploy_pd_disaggregation=True, + prefill_urls=['http://prefill1:8000'], + decode_urls=['http://decode1:8001'], + policy='cache_aware', + ) + + # Should not raise validation error + with patch('lmdeploy_router.launch_router.Router') as router_mod: + mock_router_instance = MagicMock() + router_mod.from_args = MagicMock(return_value=mock_router_instance) + + launch_router(args) + + # Should create router instance via from_args + router_mod.from_args.assert_called_once() + + def test_launch_with_service_discovery_config(self): + """Test launching with valid service discovery configuration.""" + args = RouterArgs( + service_discovery=True, + selector={'app': 'worker'}, + service_discovery_port=8080, + ) + + # Should not raise validation error + with patch('lmdeploy_router.launch_router.Router') as router_mod: + mock_router_instance = MagicMock() + router_mod.from_args = MagicMock(return_value=mock_router_instance) + + launch_router(args) + + # Should create router instance via from_args + router_mod.from_args.assert_called_once() diff --git a/lmdeploy/router/pyproject.toml b/lmdeploy/router/pyproject.toml new file mode 100644 index 0000000000..c147e7c0c4 --- /dev/null +++ b/lmdeploy/router/pyproject.toml @@ -0,0 +1,45 @@ +[build-system] +requires = ["setuptools>=45", "wheel", "setuptools-rust>=1.5.2"] +build-backend = "setuptools.build_meta" + +[project] +name = "lmdeploy-router" +version = "0.0.3" +description = "High-performance Rust-based load balancer for LMDeploy with multiple routing algorithms and prefill-decode disaggregation support" +authors = [{name = "Byron Hsu", email = "byronhsu1230@gmail.com"}] +requires-python = ">=3.8" +readme = "README.md" +license = "Apache-2.0" +license-files = [] +classifiers = [ + "Programming Language :: Python :: Implementation :: CPython", + "Programming Language :: Rust", + "Programming Language :: Python :: 3", +] + +dependencies = [ + "setproctitle", + "aiohttp", + "orjson", + "uvicorn", + "fastapi", + "requests>=2.25.0", +] + +[project.optional-dependencies] +dev = [ + "pytest>=7.0.0", + "pytest-asyncio>=0.21.0", + "pytest-cov>=4.0.0", +] + +[project.scripts] +lmdeploy-router = "lmdeploy_router.launch_router:main" + + +# https://github.com/PyO3/setuptools-rust?tab=readme-ov-file +[tool.setuptools.packages] +find = { where = ["py_src"] } + +[tool.distutils.bdist_wheel] +py_limited_api = "cp38" diff --git a/lmdeploy/router/pytest.ini b/lmdeploy/router/pytest.ini new file mode 100644 index 0000000000..1fcf1d49ed --- /dev/null +++ b/lmdeploy/router/pytest.ini @@ -0,0 +1,9 @@ +[pytest] +testpaths = py_test +python_files = test_*.py +python_classes = Test* +python_functions = test_* +markers = + integration: starts the router with local mock workers + pd_engine: requires LMDeploy PD engine metadata + e2e: requires a real LMDeploy server diff --git a/lmdeploy/router/scripts/install_aws_ofi_nccl.sh b/lmdeploy/router/scripts/install_aws_ofi_nccl.sh new file mode 100755 index 0000000000..5cb2f7596a --- /dev/null +++ b/lmdeploy/router/scripts/install_aws_ofi_nccl.sh @@ -0,0 +1,115 @@ +#!/bin/bash +# AWS OFI NCCL Installation Script +# This script installs AWS OFI NCCL with all required dependencies on Ubuntu systems +# Tested on Ubuntu 24.04 with CUDA 12.8 + +set -e # Exit on error + +echo "=========================================" +echo "AWS OFI NCCL Installation Script" +echo "=========================================" + +# Step 1: Install libfabric development package (older version, will be upgraded) +echo "Step 1: Installing initial libfabric packages..." +sudo apt-get update +sudo apt-get install -y libfabric-dev + +# Step 2: Install hwloc development package (required for aws-ofi-nccl) +echo "Step 2: Installing hwloc development package..." +sudo apt-get install -y libhwloc-dev + +# Step 3: Download and install AWS EFA installer for newer libfabric (1.22.0+) +echo "Step 3: Downloading AWS EFA installer..." +cd /tmp +wget https://efa-installer.amazonaws.com/aws-efa-installer-latest.tar.gz +tar -xf aws-efa-installer-latest.tar.gz + +echo "Installing AWS EFA with newer libfabric..." +cd aws-efa-installer +# Skip kernel module and limit configuration as they may not be needed on all instances +sudo ./efa_installer.sh -y --skip-kmod --skip-limit-conf + +# Step 4: Clone aws-ofi-nccl repository +echo "Step 4: Cloning aws-ofi-nccl repository..." +cd /tmp +if [ -d "aws-ofi-nccl" ]; then + echo "Removing existing aws-ofi-nccl directory..." + rm -rf aws-ofi-nccl +fi +git clone https://github.com/aws/aws-ofi-nccl.git +cd aws-ofi-nccl + +# Step 5: Generate configure script +echo "Step 5: Generating configure script..." +autoreconf -ivf + +# Step 6: Configure aws-ofi-nccl +echo "Step 6: Configuring aws-ofi-nccl build..." +# Detect CUDA installation path +CUDA_PATH="" +if [ -d "/usr/local/cuda" ]; then + CUDA_PATH="/usr/local/cuda" +elif [ -d "/usr/local/cuda-12.8" ]; then + CUDA_PATH="/usr/local/cuda-12.8" +elif [ -d "/usr/local/cuda-12.7" ]; then + CUDA_PATH="/usr/local/cuda-12.7" +else + echo "Error: CUDA installation not found in standard locations" + exit 1 +fi + +echo "Using CUDA installation at: $CUDA_PATH" + +# Configure with AWS EFA paths +./configure \ + --with-libfabric=/opt/amazon/efa \ + --with-cuda=$CUDA_PATH \ + --with-mpi=/opt/amazon/openmpi \ + --enable-platform-aws + +# Step 7: Build aws-ofi-nccl +echo "Step 7: Building aws-ofi-nccl..." +make -j$(nproc) + +# Step 8: Install aws-ofi-nccl +echo "Step 8: Installing aws-ofi-nccl..." +sudo make install + +# Step 9: Update library cache +echo "Step 9: Updating library cache..." +sudo ldconfig + +# Step 10: Verify installation +echo "Step 10: Verifying installation..." +echo "Installed libraries:" +ls -la /usr/local/lib/libnccl* 2>/dev/null || echo "No NCCL libraries found in /usr/local/lib" + +echo "" +echo "=========================================" +echo "AWS OFI NCCL Installation Complete!" +echo "=========================================" +echo "" +echo "The following libraries have been installed:" +echo " - libnccl-net-ofi.so" +echo " - libnccl-net.so" +echo " - libnccl-ofi-tuner.so" +echo " - libnccl-tuner-ofi.so" +echo "" +echo "Location: /usr/local/lib/" +echo "" +echo "To use AWS OFI NCCL with your applications, you may need to set:" +echo " export LD_LIBRARY_PATH=/usr/local/lib:\$LD_LIBRARY_PATH" +echo " export NCCL_NET_PLUGIN=libnccl-net-ofi.so" +echo "" +echo "For EFA-enabled instances, also set:" +echo " export FI_PROVIDER=efa" +echo " export FI_EFA_USE_DEVICE_RDMA=1" +echo "" + +# Cleanup +echo "Cleaning up temporary files..." +cd / +rm -rf /tmp/aws-efa-installer* +rm -rf /tmp/aws-ofi-nccl + +echo "Installation script completed successfully!" diff --git a/lmdeploy/router/scripts/setup-sccache.sh b/lmdeploy/router/scripts/setup-sccache.sh new file mode 100755 index 0000000000..2cdfb1e58f --- /dev/null +++ b/lmdeploy/router/scripts/setup-sccache.sh @@ -0,0 +1,84 @@ +#!/usr/bin/env bash + +set -Eeuo pipefail +IFS=$'\n\t' + +echo "Setting up sccache for faster Rust compilation..." + +has_cmd() { command -v "$1" >/dev/null 2>&1; } + +install_sccache() { + echo "sccache not found." + if [[ "${AUTO_INSTALL:-0}" != "1" ]]; then + read -r -p "Install sccache now? [y/N] " response + response=${response:-N} + if [[ ! "$response" =~ ^[Yy]$ ]]; then + echo "Skipping installation. Please install sccache manually:" + echo " cargo install sccache" + echo " or" + echo " brew install sccache (macOS)" + echo " or" + echo " sudo apt-get install -y sccache (Debian/Ubuntu)" + echo " or" + echo " sudo dnf install -y sccache (RHEL/Fedora)" + echo " or" + echo " sudo pacman -S sccache (Arch)" + exit 0 + fi + fi + + if has_cmd cargo; then + echo "Installing via cargo..." + cargo install sccache --locked + elif has_cmd brew; then + echo "Installing via Homebrew..." + brew install sccache + elif has_cmd apt-get; then + echo "Installing via apt-get..." + sudo apt-get update -y && sudo apt-get install -y sccache + elif has_cmd dnf; then + echo "Installing via dnf..." + sudo dnf install -y sccache + elif has_cmd pacman; then + echo "Installing via pacman..." + sudo pacman -S --noconfirm sccache + else + echo "No supported package manager detected. Install manually:" + echo " cargo install sccache" + exit 1 + fi +} + +if ! has_cmd sccache; then + install_sccache +fi + +echo "Configuring sccache..." + +export SCCACHE_CACHE_SIZE="${SCCACHE_CACHE_SIZE:-10G}" +export SCCACHE_STATS="${SCCACHE_STATS:-1}" + +# Set RUSTC_WRAPPER to sccache for this shell session. +SCCACHE_BIN="$(command -v sccache)" +if [[ -z "${SCCACHE_BIN}" ]]; then + echo "Unexpected: sccache still not on PATH after install. Check your environment." + exit 1 +fi +export RUSTC_WRAPPER="${SCCACHE_BIN}" + +echo "sccache version: $(sccache --version || echo 'unknown')" +echo "Current cache stats:" +sccache -s || true + +# If script not sourced, remind user about persistence. +if [[ "${BASH_SOURCE[0]}" == "$0" ]]; then + echo + echo "Environment variables exported for this process only." + echo "To persist, add to your shell profile (e.g., ~/.bashrc or ~/.zshrc):" + echo ' export RUSTC_WRAPPER="$(command -v sccache 2>/dev/null || echo "")"' + echo ' export SCCACHE_CACHE_SIZE="10G"' + # echo ' export SCCACHE_DIR="$HOME/.cache/sccache"' + echo ' export SCCACHE_STATS="1"' +fi + +echo "sccache is configured." diff --git a/lmdeploy/router/setup.py b/lmdeploy/router/setup.py new file mode 100644 index 0000000000..a0b4aaaeaa --- /dev/null +++ b/lmdeploy/router/setup.py @@ -0,0 +1,61 @@ +# Copyright (c) OpenMMLab. All rights reserved. +import os +import shutil +import subprocess + +from setuptools import setup +from setuptools.command.build_py import build_py + + +class BuildPyWithRustBinary(build_py): + binary_name = 'lmdeploy-router-bin' + + def run(self): + super().run() + if os.environ.get('LMDEPLOY_ROUTER_BUILD_NO_RUST') == '1': + return + package_name = self.distribution.get_name().replace('-', '_') + data_dir = f'{package_name}-{self.distribution.get_version()}.data' + bin_dir = os.path.join(self.build_lib, data_dir, 'scripts') + os.makedirs(bin_dir, exist_ok=True) + binary = os.environ.get( + 'LMDEPLOY_ROUTER_BIN', 'target/release/lmdeploy-router' + ) + if not os.path.isfile(binary): + subprocess.run( + [ + 'cargo', + 'build', + '--release', + '--locked', + '--bin', + 'lmdeploy-router', + ], + check=True, + ) + shutil.copy(binary, os.path.join(bin_dir, self.binary_name)) + os.chmod(os.path.join(bin_dir, self.binary_name), 0o755) + + +no_rust = os.environ.get('LMDEPLOY_ROUTER_BUILD_NO_RUST') == '1' + +rust_extensions = [] +if not no_rust: + from setuptools_rust import Binding, RustExtension + + rust_extensions.append( + RustExtension( + target='lmdeploy_router_rs', + path='Cargo.toml', + binding=Binding.PyO3, + ) + ) + + +setup( + cmdclass={ + 'build_py': BuildPyWithRustBinary, + }, + rust_extensions=rust_extensions, + zip_safe=False, +) diff --git a/lmdeploy/router/src/config/mod.rs b/lmdeploy/router/src/config/mod.rs new file mode 100644 index 0000000000..4622ff7816 --- /dev/null +++ b/lmdeploy/router/src/config/mod.rs @@ -0,0 +1,28 @@ +pub mod types; +pub mod validation; + +pub use types::*; +pub use validation::*; + +/// Configuration errors +#[derive(Debug, thiserror::Error)] +pub enum ConfigError { + #[error("Validation failed: {reason}")] + ValidationFailed { reason: String }, + + #[error("Invalid value for field '{field}': {value} - {reason}")] + InvalidValue { + field: String, + value: String, + reason: String, + }, + + #[error("Incompatible configuration: {reason}")] + IncompatibleConfig { reason: String }, + + #[error("Missing required field: {field}")] + MissingRequired { field: String }, +} + +/// Result type for configuration operations +pub type ConfigResult = Result; diff --git a/lmdeploy/router/src/config/types.rs b/lmdeploy/router/src/config/types.rs new file mode 100644 index 0000000000..27175871c6 --- /dev/null +++ b/lmdeploy/router/src/config/types.rs @@ -0,0 +1,1603 @@ +use super::ConfigResult; +use crate::config::validation::ConfigValidator; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; + +/// Main router configuration +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RouterConfig { + /// Routing mode configuration + pub mode: RoutingMode, + /// Worker connection mode + #[serde(default)] + pub connection_mode: ConnectionMode, + /// Policy configuration + pub policy: PolicyConfig, + /// Server host address + pub host: String, + /// Server port + pub port: u16, + /// Maximum payload size in bytes + pub max_payload_size: usize, + /// Request timeout in seconds + pub request_timeout_secs: u64, + /// Worker startup timeout in seconds + pub worker_startup_timeout_secs: u64, + /// Worker health check interval in seconds + pub worker_startup_check_interval_secs: u64, + /// The api key used for the authorization with the worker + pub api_key: Option, + /// API key validation URLs (if set, incoming requests must validate against them) + #[serde(default)] + pub api_key_validation_urls: Vec, + /// Service discovery configuration (optional) + pub discovery: Option, + /// Metrics configuration (optional) + pub metrics: Option, + /// Log directory (None = stdout only) + pub log_dir: Option, + /// Log level (None = info) + pub log_level: Option, + /// Custom request ID headers to check (defaults to common headers) + pub request_id_headers: Option>, + /// Maximum concurrent requests allowed (for rate limiting) + pub max_concurrent_requests: usize, + /// Queue size for pending requests when max concurrent limit reached (0 = no queue, return 429 immediately) + pub queue_size: usize, + /// Maximum time (in seconds) a request can wait in queue before timing out + pub queue_timeout_secs: u64, + /// Token bucket refill rate (tokens per second). If not set, defaults to max_concurrent_requests + pub rate_limit_tokens_per_second: Option, + /// CORS allowed origins + pub cors_allowed_origins: Vec, + /// Retry configuration + pub retry: RetryConfig, + /// Circuit breaker configuration + pub circuit_breaker: CircuitBreakerConfig, + /// Disable retries (overrides retry.max_retries to 1 when true) + #[serde(default)] + pub disable_retries: bool, + /// Disable circuit breaker (overrides circuit_breaker.failure_threshold to u32::MAX when true) + #[serde(default)] + pub disable_circuit_breaker: bool, + /// Health check configuration + pub health_check: HealthCheckConfig, + /// Enable Inference Gateway mode (false = proxy mode, true = IGW mode) + #[serde(default)] + pub enable_igw: bool, + /// History backend configuration (memory or none, default: memory) + #[serde(default = "default_history_backend")] + pub history_backend: HistoryBackend, + /// Enable profiling calls to backend workers + #[serde(default)] + pub enable_profiling: bool, + /// Profiling timeout in seconds + #[serde(default = "default_profile_timeout_secs")] + pub profile_timeout_secs: u64, + /// KV connector type for PD disaggregation + #[serde(default)] + pub kv_connector: KvConnector, +} + +fn default_profile_timeout_secs() -> u64 { + 10 +} + +fn default_history_backend() -> HistoryBackend { + HistoryBackend::Memory +} + +/// History backend configuration +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +#[serde(rename_all = "lowercase")] +pub enum HistoryBackend { + /// In-memory storage (default) + Memory, + /// No history storage + None, +} + +/// KV connector type for PD disaggregation +#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default, PartialEq, Eq, clap::ValueEnum)] +#[serde(rename_all = "lowercase")] +pub enum KvConnector { + /// NIXL pull-based KV transfer (default) + #[default] + #[serde(rename = "nixl")] + #[value(name = "nixl")] + Nixl, + /// Mooncake push-based KV transfer + #[serde(rename = "mooncake")] + #[value(name = "mooncake")] + Mooncake, + /// MoRI-IO KV transfer + #[serde(rename = "moriio")] + #[value(name = "moriio")] + MoriIO, +} + +/// Migration protocol for lmdeploy PD disaggregation +#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default, PartialEq, Eq, clap::ValueEnum)] +#[serde(rename_all = "lowercase")] +pub enum LMDeployMigrationProtocol { + /// RDMA-based KV migration (default and recommended) + #[default] + #[serde(rename = "rdma")] + #[value(name = "rdma")] + Rdma, + /// NVLink-based KV migration + #[serde(rename = "nvlink")] + #[value(name = "nvlink")] + Nvlink, +} + +/// RDMA link type used by lmdeploy's migration backend. +#[derive(Debug, Clone, Copy, Serialize, Deserialize, Default, PartialEq, Eq, clap::ValueEnum)] +pub enum LMDeployRdmaLinkType { + /// InfiniBand + #[serde(rename = "IB")] + #[value(name = "ib")] + Ib, + /// RDMA over Converged Ethernet + #[default] + #[serde(rename = "RoCE")] + #[value(name = "roce")] + Roce, +} + +/// RDMA configuration for lmdeploy PD KV migration +#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)] +pub struct LMDeployRdmaConfig { + /// Whether GPU Direct RDMA is enabled. + #[serde(default = "default_lmdeploy_with_gdr")] + pub with_gdr: bool, + /// RDMA link type (RoCE by default). + #[serde(default)] + pub link_type: LMDeployRdmaLinkType, +} + +fn default_lmdeploy_with_gdr() -> bool { + true +} + +impl Default for LMDeployRdmaConfig { + fn default() -> Self { + Self { + with_gdr: true, + link_type: LMDeployRdmaLinkType::Roce, + } + } +} + +#[derive(Debug, Clone, Serialize, Deserialize, Default, PartialEq)] +#[serde(tag = "type")] +pub enum ConnectionMode { + #[default] + #[serde(rename = "http")] + Http, +} + +/// Routing mode configuration +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type")] +pub enum RoutingMode { + #[serde(rename = "regular")] + Regular { + /// List of worker URLs + worker_urls: Vec, + }, + #[serde(rename = "openai")] + OpenAI { + /// OpenAI-compatible API base(s), provided via worker URLs + worker_urls: Vec, + }, + #[serde(rename = "lmdeploy_prefill_decode")] + LMDeployPrefillDecode { + /// lmdeploy prefill worker URLs + prefill_urls: Vec, + /// lmdeploy decode worker URLs + decode_urls: Vec, + /// Optional separate policy for prefill workers + #[serde(skip_serializing_if = "Option::is_none")] + prefill_policy: Option, + /// Optional separate policy for decode workers + #[serde(skip_serializing_if = "Option::is_none")] + decode_policy: Option, + /// KV migration protocol (rdma/nvlink) + #[serde(default)] + migration_protocol: LMDeployMigrationProtocol, + /// Optional RDMA configuration + #[serde(skip_serializing_if = "Option::is_none")] + rdma_config: Option, + /// Whether to use dummy prefill (skip actual prefill compute) + #[serde(default)] + dummy_prefill: bool, + }, +} + +impl RoutingMode { + pub fn is_pd_mode(&self) -> bool { + matches!(self, RoutingMode::LMDeployPrefillDecode { .. }) + } + + /// Returns true if this is lmdeploy PD disaggregation mode + pub fn is_lmdeploy_pd_mode(&self) -> bool { + matches!(self, RoutingMode::LMDeployPrefillDecode { .. }) + } + + pub fn worker_count(&self) -> usize { + match self { + RoutingMode::Regular { worker_urls } => worker_urls.len(), + RoutingMode::LMDeployPrefillDecode { + prefill_urls, + decode_urls, + .. + } => prefill_urls.len() + decode_urls.len(), + // OpenAI mode represents a single upstream + RoutingMode::OpenAI { .. } => 1, + } + } + + /// Get the effective prefill policy for PD mode + /// Falls back to the main policy if no specific prefill policy is set + pub fn get_prefill_policy<'a>(&'a self, main_policy: &'a PolicyConfig) -> &'a PolicyConfig { + match self { + RoutingMode::LMDeployPrefillDecode { prefill_policy, .. } => { + prefill_policy.as_ref().unwrap_or(main_policy) + } + _ => main_policy, + } + } + + /// Get the effective decode policy for PD mode + /// Falls back to the main policy if no specific decode policy is set + pub fn get_decode_policy<'a>(&'a self, main_policy: &'a PolicyConfig) -> &'a PolicyConfig { + match self { + RoutingMode::LMDeployPrefillDecode { decode_policy, .. } => { + decode_policy.as_ref().unwrap_or(main_policy) + } + _ => main_policy, + } + } +} + +/// Policy configuration for routing +#[derive(Debug, Clone, Serialize, Deserialize)] +#[serde(tag = "type")] +pub enum PolicyConfig { + #[serde(rename = "random")] + Random, + + #[serde(rename = "round_robin")] + RoundRobin, + + #[serde(rename = "cache_aware")] + CacheAware { + /// Minimum prefix match ratio to use cache-based routing + cache_threshold: f32, + /// Absolute load difference threshold for load balancing + balance_abs_threshold: usize, + /// Relative load ratio threshold for load balancing + balance_rel_threshold: f32, + /// Interval between cache eviction cycles (seconds) + eviction_interval_secs: u64, + /// Maximum cache tree size per tenant + max_tree_size: usize, + }, + + #[serde(rename = "power_of_two")] + PowerOfTwo { + /// Interval for load monitoring (seconds) + load_check_interval_secs: u64, + }, + + #[serde(rename = "consistent_hash")] + ConsistentHash { + /// Number of virtual nodes per worker for better distribution + virtual_nodes: u32, + }, + + #[serde(rename = "rendezvous_hash")] + RendezvousHash, +} + +impl PolicyConfig { + pub fn name(&self) -> &'static str { + match self { + PolicyConfig::Random => "random", + PolicyConfig::RoundRobin => "round_robin", + PolicyConfig::CacheAware { .. } => "cache_aware", + PolicyConfig::PowerOfTwo { .. } => "power_of_two", + PolicyConfig::ConsistentHash { .. } => "consistent_hash", + PolicyConfig::RendezvousHash => "rendezvous_hash", + } + } +} + +/// Service discovery configuration +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct DiscoveryConfig { + /// Enable service discovery + pub enabled: bool, + /// Kubernetes namespace (None = all namespaces) + pub namespace: Option, + /// Service discovery port + pub port: u16, + /// Check interval for service discovery + pub check_interval_secs: u64, + /// Regular mode selector + pub selector: HashMap, + /// PD mode prefill selector + pub prefill_selector: HashMap, + /// PD mode decode selector + pub decode_selector: HashMap, +} + +impl Default for DiscoveryConfig { + fn default() -> Self { + Self { + enabled: false, + namespace: None, + port: 8000, + check_interval_secs: 120, + selector: HashMap::new(), + prefill_selector: HashMap::new(), + decode_selector: HashMap::new(), + } + } +} + +/// Retry configuration for request handling +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct RetryConfig { + /// Maximum number of retry attempts + pub max_retries: u32, + /// Initial backoff delay in milliseconds + pub initial_backoff_ms: u64, + /// Maximum backoff delay in milliseconds + pub max_backoff_ms: u64, + /// Backoff multiplier for exponential backoff + pub backoff_multiplier: f32, + /// Jitter factor applied to backoff (0.0 - 1.0) + /// Effective delay D' = D * (1 + U[-j, +j]) + #[serde(default = "default_retry_jitter_factor")] + pub jitter_factor: f32, +} + +impl Default for RetryConfig { + fn default() -> Self { + Self { + max_retries: 5, + initial_backoff_ms: 50, + max_backoff_ms: 30000, + backoff_multiplier: 1.5, + jitter_factor: 0.2, + } + } +} + +fn default_retry_jitter_factor() -> f32 { + 0.2 +} + +/// Health check configuration for worker monitoring +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct HealthCheckConfig { + /// Number of consecutive failures before marking unhealthy + pub failure_threshold: u32, + /// Number of consecutive successes before marking healthy + pub success_threshold: u32, + /// Timeout for health check requests in seconds + pub timeout_secs: u64, + /// Interval between health checks in seconds + pub check_interval_secs: u64, + /// Health check endpoint path + pub endpoint: String, +} + +impl Default for HealthCheckConfig { + fn default() -> Self { + Self { + failure_threshold: 3, + success_threshold: 2, + timeout_secs: 5, + check_interval_secs: 60, + endpoint: "/health".to_string(), + } + } +} + +/// Circuit breaker configuration for worker reliability +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct CircuitBreakerConfig { + /// Number of consecutive failures before opening circuit + pub failure_threshold: u32, + /// Number of consecutive successes before closing circuit + pub success_threshold: u32, + /// Time before attempting to recover from open state (in seconds) + pub timeout_duration_secs: u64, + /// Window duration for failure tracking (in seconds) + pub window_duration_secs: u64, +} + +impl Default for CircuitBreakerConfig { + fn default() -> Self { + Self { + failure_threshold: 10, + success_threshold: 3, + timeout_duration_secs: 60, + window_duration_secs: 120, + } + } +} + +/// Metrics configuration +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct MetricsConfig { + /// Prometheus metrics port + pub port: u16, + /// Prometheus metrics host + pub host: String, +} + +impl Default for MetricsConfig { + fn default() -> Self { + Self { + port: 29000, + host: "127.0.0.1".to_string(), + } + } +} + +/// OpenTelemetry tracing configuration. +/// +/// Presence of `Some(TraceConfig)` means tracing is enabled; +/// `None` means tracing is disabled. +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct TraceConfig { + /// OTLP collector endpoint (format: host:port). + /// When None, the SDK respects OTEL_EXPORTER_OTLP_ENDPOINT. + #[serde(default)] + pub otlp_traces_endpoint: Option, + /// Parent-based trace sampling ratio applied when there is no sampled parent. + #[serde(default = "TraceConfig::default_sampling_ratio")] + pub sampling_ratio: f64, + /// Exact HTTP paths whose server spans should be skipped even when tracing is enabled. + #[serde(default = "TraceConfig::default_excluded_paths")] + pub excluded_paths: Vec, +} + +impl TraceConfig { + pub fn default_sampling_ratio() -> f64 { + 1.0 + } + + pub fn default_excluded_paths() -> Vec { + vec![ + "/health".to_string(), + "/health_generate".to_string(), + "/liveness".to_string(), + "/readiness".to_string(), + ] + } +} + +impl Default for TraceConfig { + fn default() -> Self { + Self { + otlp_traces_endpoint: None, + sampling_ratio: Self::default_sampling_ratio(), + excluded_paths: Self::default_excluded_paths(), + } + } +} + +impl Default for RouterConfig { + fn default() -> Self { + Self { + mode: RoutingMode::Regular { + worker_urls: vec![], + }, + policy: PolicyConfig::Random, + host: "127.0.0.1".to_string(), + port: 3001, + max_payload_size: 536_870_912, // 512MB + request_timeout_secs: 1800, // 30 minutes + worker_startup_timeout_secs: 600, + worker_startup_check_interval_secs: 30, + api_key: None, + api_key_validation_urls: vec![], + discovery: None, + metrics: None, + log_dir: None, + log_level: None, + request_id_headers: None, + max_concurrent_requests: 32768, + queue_size: 100, + queue_timeout_secs: 60, + rate_limit_tokens_per_second: None, + cors_allowed_origins: vec![], + retry: RetryConfig::default(), + circuit_breaker: CircuitBreakerConfig::default(), + disable_retries: false, + disable_circuit_breaker: false, + health_check: HealthCheckConfig::default(), + enable_igw: false, + connection_mode: ConnectionMode::Http, + history_backend: default_history_backend(), + enable_profiling: false, + profile_timeout_secs: default_profile_timeout_secs(), + kv_connector: KvConnector::default(), + } + } +} + +impl RouterConfig { + /// Create a new configuration with mode and policy + pub fn new(mode: RoutingMode, policy: PolicyConfig) -> Self { + Self { + mode, + policy, + ..Default::default() + } + } + + /// Validate the configuration + pub fn validate(&self) -> ConfigResult<()> { + ConfigValidator::validate(self) + } + + /// Get the routing mode type as a string + pub fn mode_type(&self) -> &'static str { + match self.mode { + RoutingMode::Regular { .. } => "regular", + RoutingMode::LMDeployPrefillDecode { .. } => "lmdeploy_prefill_decode", + RoutingMode::OpenAI { .. } => "openai", + } + } + + /// Check if service discovery is enabled + pub fn has_service_discovery(&self) -> bool { + self.discovery.as_ref().is_some_and(|d| d.enabled) + } + + /// Check if metrics are enabled + pub fn has_metrics(&self) -> bool { + self.metrics.is_some() + } + + /// Compute the effective retry config considering disable flag + pub fn effective_retry_config(&self) -> RetryConfig { + let mut cfg = self.retry.clone(); + if self.disable_retries { + cfg.max_retries = 1; + } + cfg + } + + /// Compute the effective circuit breaker config considering disable flag + pub fn effective_circuit_breaker_config(&self) -> CircuitBreakerConfig { + let mut cfg = self.circuit_breaker.clone(); + if self.disable_circuit_breaker { + cfg.failure_threshold = u32::MAX; + } + cfg + } + + /// Check if running in IGW (Inference Gateway) mode + pub fn is_igw_mode(&self) -> bool { + self.enable_igw + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // ============= RouterConfig Tests ============= + + #[test] + fn test_router_config_default() { + let config = RouterConfig::default(); + + assert!( + matches!(config.mode, RoutingMode::Regular { worker_urls } if worker_urls.is_empty()) + ); + assert!(matches!(config.policy, PolicyConfig::Random)); + assert_eq!(config.host, "127.0.0.1"); + assert_eq!(config.port, 3001); + assert_eq!(config.max_payload_size, 536_870_912); + assert_eq!(config.request_timeout_secs, 1800); + assert_eq!(config.worker_startup_timeout_secs, 600); + assert_eq!(config.worker_startup_check_interval_secs, 30); + assert!(config.discovery.is_none()); + assert!(config.metrics.is_none()); + assert!(config.log_dir.is_none()); + assert!(config.log_level.is_none()); + } + + #[test] + fn test_router_config_new() { + let mode = RoutingMode::Regular { + worker_urls: vec!["http://worker1".to_string(), "http://worker2".to_string()], + }; + let policy = PolicyConfig::RoundRobin; + + let config = RouterConfig::new(mode, policy); + + match config.mode { + RoutingMode::Regular { worker_urls } => { + assert_eq!(worker_urls.len(), 2); + assert_eq!(worker_urls[0], "http://worker1"); + assert_eq!(worker_urls[1], "http://worker2"); + } + _ => panic!("Expected Regular mode"), + } + + assert!(matches!(config.policy, PolicyConfig::RoundRobin)); + // Other fields should be default + assert_eq!(config.host, "127.0.0.1"); + assert_eq!(config.port, 3001); + } + + #[test] + fn test_router_config_serialization() { + let config = RouterConfig { + mode: RoutingMode::Regular { + worker_urls: vec!["http://worker1".to_string()], + }, + policy: PolicyConfig::Random, + host: "0.0.0.0".to_string(), + port: 8080, + log_dir: Some("/var/log".to_string()), + log_level: Some("debug".to_string()), + ..Default::default() + }; + + let json = serde_json::to_string(&config).unwrap(); + let deserialized: RouterConfig = serde_json::from_str(&json).unwrap(); + + assert_eq!(config.host, deserialized.host); + assert_eq!(config.port, deserialized.port); + assert_eq!(config.max_payload_size, deserialized.max_payload_size); + assert_eq!(config.log_dir, deserialized.log_dir); + assert_eq!(config.log_level, deserialized.log_level); + // discovery and metrics are None in Default implementation + assert!(deserialized.discovery.is_none()); + assert!(deserialized.metrics.is_none()); + } + + // ============= RoutingMode Tests ============= + + #[test] + fn test_routing_mode_is_pd_mode() { + let regular = RoutingMode::Regular { + worker_urls: vec!["http://worker1".to_string()], + }; + assert!(!regular.is_pd_mode()); + + let pd = RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec!["http://prefill1".to_string()], + decode_urls: vec!["http://decode1".to_string()], + prefill_policy: None, + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::default(), + rdma_config: None, + dummy_prefill: false, + }; + assert!(pd.is_pd_mode()); + } + + #[test] + fn test_routing_mode_worker_count() { + let regular = RoutingMode::Regular { + worker_urls: vec![ + "http://worker1".to_string(), + "http://worker2".to_string(), + "http://worker3".to_string(), + ], + }; + assert_eq!(regular.worker_count(), 3); + + let pd = RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec!["http://prefill1".to_string(), "http://prefill2".to_string()], + decode_urls: vec![ + "http://decode1".to_string(), + "http://decode2".to_string(), + "http://decode3".to_string(), + ], + prefill_policy: None, + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::default(), + rdma_config: None, + dummy_prefill: false, + }; + assert_eq!(pd.worker_count(), 5); + + let empty_regular = RoutingMode::Regular { + worker_urls: vec![], + }; + assert_eq!(empty_regular.worker_count(), 0); + } + + #[test] + fn test_routing_mode_serialization() { + // Test Regular mode + let regular = RoutingMode::Regular { + worker_urls: vec!["http://worker1".to_string()], + }; + let json = serde_json::to_string(®ular).unwrap(); + assert!(json.contains("\"type\":\"regular\"")); + assert!(json.contains("\"worker_urls\"")); + + // Test LMDeployPrefillDecode mode + let pd = RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec!["http://prefill1".to_string()], + decode_urls: vec!["http://decode1".to_string()], + prefill_policy: None, + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::default(), + rdma_config: None, + dummy_prefill: false, + }; + let json = serde_json::to_string(&pd).unwrap(); + assert!(json.contains("\"type\":\"lmdeploy_prefill_decode\"")); + assert!(json.contains("\"prefill_urls\"")); + assert!(json.contains("\"decode_urls\"")); + } + + #[test] + fn test_lmdeploy_prefill_decode_serialization() { + let lmdeploy_pd = RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec!["http://prefill1".to_string()], + decode_urls: vec!["http://decode1".to_string()], + prefill_policy: None, + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::Rdma, + rdma_config: None, + dummy_prefill: false, + }; + let json = serde_json::to_string(&lmdeploy_pd).unwrap(); + assert!(json.contains("\"type\":\"lmdeploy_prefill_decode\"")); + assert!(json.contains("\"prefill_urls\"")); + assert!(json.contains("\"decode_urls\"")); + assert!(json.contains("\"migration_protocol\":\"rdma\"")); + assert!(!json.contains("\"rdma_config\"")); + assert!(json.contains("\"dummy_prefill\":false")); + + // Round-trip + let deserialized: RoutingMode = serde_json::from_str(&json).unwrap(); + assert!(deserialized.is_pd_mode()); + assert!(deserialized.is_lmdeploy_pd_mode()); + assert_eq!(deserialized.worker_count(), 2); + + // Test with rdma config and dummy_prefill=true + let lmdeploy_pd_rdma = RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec!["http://p:8000".to_string()], + decode_urls: vec!["http://d:8000".to_string()], + prefill_policy: None, + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::Rdma, + rdma_config: Some(LMDeployRdmaConfig { + with_gdr: true, + link_type: LMDeployRdmaLinkType::Roce, + }), + dummy_prefill: true, + }; + let json_rdma = serde_json::to_string(&lmdeploy_pd_rdma).unwrap(); + assert!(json_rdma.contains("\"migration_protocol\":\"rdma\"")); + assert!(json_rdma.contains("\"rdma_config\"")); + assert!(json_rdma.contains("\"with_gdr\":true")); + assert!(json_rdma.contains("\"link_type\":\"RoCE\"")); + assert!(json_rdma.contains("\"dummy_prefill\":true")); + } + + #[test] + fn test_lmdeploy_migration_protocol_default() { + assert_eq!( + LMDeployMigrationProtocol::default(), + LMDeployMigrationProtocol::Rdma + ); + } + + #[test] + fn test_lmdeploy_migration_protocol_serde_variants() { + // Router configuration uses stable, human-readable lowercase strings. + assert_eq!( + serde_json::to_string(&LMDeployMigrationProtocol::Rdma).unwrap(), + r#""rdma""# + ); + assert_eq!( + serde_json::to_string(&LMDeployMigrationProtocol::Nvlink).unwrap(), + r#""nvlink""# + ); + // Round-trip deserialization + let rdma: LMDeployMigrationProtocol = serde_json::from_str(r#""rdma""#).unwrap(); + assert_eq!(rdma, LMDeployMigrationProtocol::Rdma); + let nvlink: LMDeployMigrationProtocol = serde_json::from_str(r#""nvlink""#).unwrap(); + assert_eq!(nvlink, LMDeployMigrationProtocol::Nvlink); + } + + #[test] + fn test_lmdeploy_rdma_config_defaults() { + let config = LMDeployRdmaConfig::default(); + assert!(config.with_gdr); + assert_eq!(config.link_type, LMDeployRdmaLinkType::Roce); + + let cfg: LMDeployRdmaConfig = serde_json::from_str(r#"{}"#).unwrap(); + assert!(cfg.with_gdr); + assert_eq!(cfg.link_type, LMDeployRdmaLinkType::Roce); + } + + #[test] + fn test_lmdeploy_pd_mode_type_via_json() { + // RoutingMode doesn't have a mode_type() method; verify via JSON serialization tag. + let mode = RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec!["http://p:8000".to_string()], + decode_urls: vec!["http://d:8000".to_string()], + prefill_policy: None, + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::Rdma, + rdma_config: None, + dummy_prefill: false, + }; + let json = serde_json::to_string(&mode).unwrap(); + assert!(json.contains(r#""type":"lmdeploy_prefill_decode""#)); + } + + #[test] + fn test_lmdeploy_pd_worker_count_multiple_urls() { + let mode = RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec![ + "http://p1".to_string(), + "http://p2".to_string(), + "http://p3".to_string(), + ], + decode_urls: vec!["http://d1".to_string(), "http://d2".to_string()], + prefill_policy: None, + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::Rdma, + rdma_config: None, + dummy_prefill: false, + }; + assert_eq!(mode.worker_count(), 5); + } + + #[test] + fn test_lmdeploy_pd_policy_fallback_to_main() { + let main_policy = PolicyConfig::RoundRobin; + let mode = RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec!["http://p".to_string()], + decode_urls: vec!["http://d".to_string()], + prefill_policy: None, + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::Rdma, + rdma_config: None, + dummy_prefill: false, + }; + // When prefill_policy/decode_policy is None, should fall back to main policy. + // Compare by name() since PolicyConfig doesn't derive PartialEq. + let prefill = mode.get_prefill_policy(&main_policy); + let decode = mode.get_decode_policy(&main_policy); + assert_eq!(prefill.name(), main_policy.name()); + assert_eq!(decode.name(), main_policy.name()); + } + + #[test] + fn test_lmdeploy_pd_policy_override() { + let main_policy = PolicyConfig::RoundRobin; + let custom_policy = PolicyConfig::Random; + let mode = RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec!["http://p".to_string()], + decode_urls: vec!["http://d".to_string()], + prefill_policy: Some(custom_policy.clone()), + decode_policy: Some(custom_policy.clone()), + migration_protocol: LMDeployMigrationProtocol::Rdma, + rdma_config: None, + dummy_prefill: false, + }; + // Should use custom policy, not main. Compare by name() since PolicyConfig doesn't derive PartialEq. + let prefill = mode.get_prefill_policy(&main_policy); + let decode = mode.get_decode_policy(&main_policy); + assert_eq!(prefill.name(), "random"); + assert_eq!(decode.name(), "random"); + } + + #[test] + fn test_lmdeploy_pd_deserialize_from_json() { + let json_str = r#"{ + "type": "lmdeploy_prefill_decode", + "prefill_urls": ["http://p:23333"], + "decode_urls": ["http://d:23333"], + "migration_protocol": "rdma", + "rdma_config": {"with_gdr": true, "link_type": "RoCE"}, + "dummy_prefill": true + }"#; + let mode: RoutingMode = serde_json::from_str(json_str).unwrap(); + assert!(mode.is_lmdeploy_pd_mode()); + assert!(mode.is_pd_mode()); + assert_eq!(mode.worker_count(), 2); + } + + #[test] + fn test_lmdeploy_pd_no_rdma_config_when_protocol_is_nvlink() { + let mode = RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec!["http://p".to_string()], + decode_urls: vec!["http://d".to_string()], + prefill_policy: None, + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::Nvlink, + rdma_config: None, + dummy_prefill: false, + }; + let json = serde_json::to_string(&mode).unwrap(); + // Option::None fields are skipped by serde by default; rdma_config should be absent. + assert!(!json.contains("rdma_config")); + } + + // ============= PolicyConfig Tests ============= + + #[test] + fn test_policy_config_name() { + assert_eq!(PolicyConfig::Random.name(), "random"); + assert_eq!(PolicyConfig::RoundRobin.name(), "round_robin"); + + let cache_aware = PolicyConfig::CacheAware { + cache_threshold: 0.8, + balance_abs_threshold: 10, + balance_rel_threshold: 1.5, + eviction_interval_secs: 300, + max_tree_size: 1000, + }; + assert_eq!(cache_aware.name(), "cache_aware"); + + let power_of_two = PolicyConfig::PowerOfTwo { + load_check_interval_secs: 60, + }; + assert_eq!(power_of_two.name(), "power_of_two"); + } + + #[test] + fn test_policy_config_serialization() { + // Test Random + let random = PolicyConfig::Random; + let json = serde_json::to_string(&random).unwrap(); + assert_eq!(json, r#"{"type":"random"}"#); + + // Test CacheAware with all parameters + let cache_aware = PolicyConfig::CacheAware { + cache_threshold: 0.8, + balance_abs_threshold: 10, + balance_rel_threshold: 1.5, + eviction_interval_secs: 300, + max_tree_size: 1000, + }; + let json = serde_json::to_string(&cache_aware).unwrap(); + assert!(json.contains("\"type\":\"cache_aware\"")); + assert!(json.contains("\"cache_threshold\":0.8")); + assert!(json.contains("\"balance_abs_threshold\":10")); + + // Test PowerOfTwo + let power_of_two = PolicyConfig::PowerOfTwo { + load_check_interval_secs: 60, + }; + let json = serde_json::to_string(&power_of_two).unwrap(); + assert!(json.contains("\"type\":\"power_of_two\"")); + assert!(json.contains("\"load_check_interval_secs\":60")); + } + + #[test] + fn test_cache_aware_parameters() { + let cache_aware = PolicyConfig::CacheAware { + cache_threshold: 0.75, + balance_abs_threshold: 20, + balance_rel_threshold: 2.0, + eviction_interval_secs: 600, + max_tree_size: 5000, + }; + + match cache_aware { + PolicyConfig::CacheAware { + cache_threshold, + balance_abs_threshold, + balance_rel_threshold, + eviction_interval_secs, + max_tree_size, + } => { + assert!((cache_threshold - 0.75).abs() < 0.0001); + assert_eq!(balance_abs_threshold, 20); + assert!((balance_rel_threshold - 2.0).abs() < 0.0001); + assert_eq!(eviction_interval_secs, 600); + assert_eq!(max_tree_size, 5000); + } + _ => panic!("Expected CacheAware"), + } + } + + #[test] + fn test_power_of_two_parameters() { + let power_of_two = PolicyConfig::PowerOfTwo { + load_check_interval_secs: 120, + }; + + match power_of_two { + PolicyConfig::PowerOfTwo { + load_check_interval_secs, + } => { + assert_eq!(load_check_interval_secs, 120); + } + _ => panic!("Expected PowerOfTwo"), + } + } + + // ============= DiscoveryConfig Tests ============= + + #[test] + fn test_discovery_config_default() { + let config = DiscoveryConfig::default(); + + assert!(!config.enabled); + assert!(config.namespace.is_none()); + assert_eq!(config.port, 8000); + assert_eq!(config.check_interval_secs, 120); + assert!(config.selector.is_empty()); + assert!(config.prefill_selector.is_empty()); + assert!(config.decode_selector.is_empty()); + } + + #[test] + fn test_discovery_config_with_selectors() { + let mut selector = HashMap::new(); + selector.insert("app".to_string(), "inference".to_string()); + selector.insert("role".to_string(), "worker".to_string()); + + let config = DiscoveryConfig { + enabled: true, + namespace: Some("default".to_string()), + port: 9000, + check_interval_secs: 30, + selector: selector.clone(), + prefill_selector: selector.clone(), + decode_selector: selector.clone(), + }; + + assert!(config.enabled); + assert_eq!(config.namespace, Some("default".to_string())); + assert_eq!(config.port, 9000); + assert_eq!(config.selector.len(), 2); + assert_eq!(config.selector.get("app"), Some(&"inference".to_string())); + } + + #[test] + fn test_discovery_config_namespace() { + // Test None namespace (all namespaces) + let config = DiscoveryConfig { + namespace: None, + ..Default::default() + }; + assert!(config.namespace.is_none()); + + // Test specific namespace + let config = DiscoveryConfig { + namespace: Some("production".to_string()), + ..Default::default() + }; + assert_eq!(config.namespace, Some("production".to_string())); + } + + // ============= MetricsConfig Tests ============= + + #[test] + fn test_metrics_config_default() { + let config = MetricsConfig::default(); + + assert_eq!(config.port, 29000); + assert_eq!(config.host, "127.0.0.1"); + } + + #[test] + fn test_metrics_config_custom() { + let config = MetricsConfig { + port: 9090, + host: "0.0.0.0".to_string(), + }; + + assert_eq!(config.port, 9090); + assert_eq!(config.host, "0.0.0.0"); + } + + // ============= RouterConfig Utility Methods Tests ============= + + #[test] + fn test_mode_type() { + let config = RouterConfig { + mode: RoutingMode::Regular { + worker_urls: vec![], + }, + ..Default::default() + }; + assert_eq!(config.mode_type(), "regular"); + + let config = RouterConfig { + mode: RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec![], + decode_urls: vec![], + prefill_policy: None, + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::default(), + rdma_config: None, + dummy_prefill: false, + }, + ..Default::default() + }; + assert_eq!(config.mode_type(), "lmdeploy_prefill_decode"); + } + + #[test] + fn test_has_service_discovery() { + let config = RouterConfig::default(); + assert!(!config.has_service_discovery()); + + let config = RouterConfig { + discovery: Some(DiscoveryConfig { + enabled: false, + ..Default::default() + }), + ..Default::default() + }; + assert!(!config.has_service_discovery()); + + let config = RouterConfig { + discovery: Some(DiscoveryConfig { + enabled: true, + ..Default::default() + }), + ..Default::default() + }; + assert!(config.has_service_discovery()); + } + + #[test] + fn test_has_metrics() { + let config = RouterConfig::default(); + assert!(!config.has_metrics()); + + let config = RouterConfig { + metrics: Some(MetricsConfig::default()), + ..Default::default() + }; + assert!(config.has_metrics()); + } + + // ============= Edge Cases ============= + + #[test] + fn test_large_worker_lists() { + let large_urls: Vec = (0..1000).map(|i| format!("http://worker{}", i)).collect(); + + let mode = RoutingMode::Regular { + worker_urls: large_urls.clone(), + }; + + assert_eq!(mode.worker_count(), 1000); + + // Test serialization with large list + let config = RouterConfig { + mode, + ..Default::default() + }; + + let json = serde_json::to_string(&config).unwrap(); + let deserialized: RouterConfig = serde_json::from_str(&json).unwrap(); + + match deserialized.mode { + RoutingMode::Regular { worker_urls } => { + assert_eq!(worker_urls.len(), 1000); + } + _ => panic!("Expected Regular mode"), + } + } + + #[test] + fn test_unicode_in_config() { + let config = RouterConfig { + mode: RoutingMode::Regular { + worker_urls: vec!["http://работник1".to_string(), "http://工作者2".to_string()], + }, + log_dir: Some("/日志/目录".to_string()), + ..Default::default() + }; + + let json = serde_json::to_string(&config).unwrap(); + let deserialized: RouterConfig = serde_json::from_str(&json).unwrap(); + + match deserialized.mode { + RoutingMode::Regular { worker_urls } => { + assert_eq!(worker_urls[0], "http://работник1"); + assert_eq!(worker_urls[1], "http://工作者2"); + } + _ => panic!("Expected Regular mode"), + } + + assert_eq!(deserialized.log_dir, Some("/日志/目录".to_string())); + } + + #[test] + fn test_empty_string_fields() { + let config = RouterConfig { + host: "".to_string(), + log_dir: Some("".to_string()), + log_level: Some("".to_string()), + ..Default::default() + }; + + assert_eq!(config.host, ""); + assert_eq!(config.log_dir, Some("".to_string())); + assert_eq!(config.log_level, Some("".to_string())); + } + + // ============= Complex Configuration Tests ============= + + #[test] + fn test_full_pd_mode_config() { + let config = RouterConfig { + mode: RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec![ + "http://prefill1:8000".to_string(), + "http://prefill2:8000".to_string(), + ], + decode_urls: vec![ + "http://decode1:8000".to_string(), + "http://decode2:8000".to_string(), + ], + prefill_policy: None, + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::default(), + rdma_config: None, + dummy_prefill: false, + }, + policy: PolicyConfig::PowerOfTwo { + load_check_interval_secs: 30, + }, + host: "0.0.0.0".to_string(), + port: 3000, + max_payload_size: 1048576, + request_timeout_secs: 120, + worker_startup_timeout_secs: 60, + worker_startup_check_interval_secs: 5, + api_key: None, + api_key_validation_urls: vec![], + discovery: Some(DiscoveryConfig { + enabled: true, + namespace: Some("router".to_string()), + ..Default::default() + }), + metrics: Some(MetricsConfig { + port: 9090, + host: "0.0.0.0".to_string(), + }), + log_dir: Some("/var/log/router".to_string()), + log_level: Some("info".to_string()), + request_id_headers: None, + max_concurrent_requests: 64, + cors_allowed_origins: vec![], + retry: RetryConfig::default(), + circuit_breaker: CircuitBreakerConfig::default(), + disable_retries: false, + disable_circuit_breaker: false, + health_check: HealthCheckConfig::default(), + enable_igw: false, + queue_size: 100, + queue_timeout_secs: 60, + rate_limit_tokens_per_second: None, + connection_mode: ConnectionMode::Http, + history_backend: default_history_backend(), + enable_profiling: false, + profile_timeout_secs: default_profile_timeout_secs(), + kv_connector: KvConnector::default(), + }; + + assert!(config.mode.is_pd_mode()); + assert_eq!(config.mode.worker_count(), 4); + assert_eq!(config.policy.name(), "power_of_two"); + assert!(config.has_service_discovery()); + assert!(config.has_metrics()); + } + + #[test] + fn test_full_regular_mode_config() { + let mut selector = HashMap::new(); + selector.insert("app".to_string(), "inference".to_string()); + + let config = RouterConfig { + mode: RoutingMode::Regular { + worker_urls: vec![ + "http://worker1:8000".to_string(), + "http://worker2:8000".to_string(), + "http://worker3:8000".to_string(), + ], + }, + policy: PolicyConfig::CacheAware { + cache_threshold: 0.9, + balance_abs_threshold: 5, + balance_rel_threshold: 1.2, + eviction_interval_secs: 600, + max_tree_size: 10000, + }, + host: "0.0.0.0".to_string(), + port: 3001, + max_payload_size: 536870912, + request_timeout_secs: 300, + worker_startup_timeout_secs: 180, + worker_startup_check_interval_secs: 15, + api_key: None, + api_key_validation_urls: vec![], + discovery: Some(DiscoveryConfig { + enabled: true, + namespace: None, + port: 8080, + check_interval_secs: 45, + selector, + ..Default::default() + }), + metrics: Some(MetricsConfig::default()), + log_dir: None, + log_level: Some("debug".to_string()), + request_id_headers: None, + max_concurrent_requests: 64, + cors_allowed_origins: vec![], + retry: RetryConfig::default(), + circuit_breaker: CircuitBreakerConfig::default(), + disable_retries: false, + disable_circuit_breaker: false, + health_check: HealthCheckConfig::default(), + enable_igw: false, + queue_size: 100, + queue_timeout_secs: 60, + rate_limit_tokens_per_second: None, + connection_mode: ConnectionMode::Http, + history_backend: default_history_backend(), + enable_profiling: false, + profile_timeout_secs: default_profile_timeout_secs(), + kv_connector: KvConnector::default(), + }; + + assert!(!config.mode.is_pd_mode()); + assert_eq!(config.mode.worker_count(), 3); + assert_eq!(config.policy.name(), "cache_aware"); + assert!(config.has_service_discovery()); + assert!(config.has_metrics()); + } + + #[test] + fn test_config_with_all_options() { + let mut selectors = HashMap::new(); + selectors.insert("env".to_string(), "prod".to_string()); + selectors.insert("version".to_string(), "v1".to_string()); + + let config = RouterConfig { + mode: RoutingMode::Regular { + worker_urls: vec!["http://worker1".to_string()], + }, + policy: PolicyConfig::RoundRobin, + host: "::1".to_string(), // IPv6 + port: 8888, + max_payload_size: 1024 * 1024 * 512, // 512MB + request_timeout_secs: 900, + worker_startup_timeout_secs: 600, + worker_startup_check_interval_secs: 20, + api_key: None, + api_key_validation_urls: vec![], + discovery: Some(DiscoveryConfig { + enabled: true, + namespace: Some("production".to_string()), + port: 8443, + check_interval_secs: 120, + selector: selectors.clone(), + prefill_selector: selectors.clone(), + decode_selector: selectors, + }), + metrics: Some(MetricsConfig { + port: 9999, + host: "::".to_string(), // IPv6 any + }), + log_dir: Some("/opt/logs/router".to_string()), + log_level: Some("trace".to_string()), + request_id_headers: None, + max_concurrent_requests: 64, + cors_allowed_origins: vec![], + retry: RetryConfig::default(), + circuit_breaker: CircuitBreakerConfig::default(), + disable_retries: false, + disable_circuit_breaker: false, + health_check: HealthCheckConfig::default(), + enable_igw: false, + queue_size: 100, + queue_timeout_secs: 60, + rate_limit_tokens_per_second: None, + connection_mode: ConnectionMode::Http, + history_backend: default_history_backend(), + enable_profiling: false, + profile_timeout_secs: default_profile_timeout_secs(), + kv_connector: KvConnector::default(), + }; + + assert!(config.has_service_discovery()); + assert!(config.has_metrics()); + assert_eq!(config.mode_type(), "regular"); + + // Test round-trip serialization + let json = serde_json::to_string_pretty(&config).unwrap(); + let deserialized: RouterConfig = serde_json::from_str(&json).unwrap(); + + assert_eq!(deserialized.host, "::1"); + assert_eq!(deserialized.port, 8888); + assert_eq!( + deserialized.discovery.unwrap().namespace, + Some("production".to_string()) + ); + } + + // ============= Policy Fallback Tests ============= + + #[test] + fn test_pd_policy_fallback_both_specified() { + // When both prefill and decode policies are specified, they should be used + let pd = RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec!["http://prefill1".to_string()], + decode_urls: vec!["http://decode1".to_string()], + prefill_policy: Some(PolicyConfig::CacheAware { + cache_threshold: 0.5, + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + eviction_interval_secs: 60, + max_tree_size: 1000, + }), + decode_policy: Some(PolicyConfig::PowerOfTwo { + load_check_interval_secs: 60, + }), + migration_protocol: LMDeployMigrationProtocol::default(), + rdma_config: None, + dummy_prefill: false, + }; + + let main_policy = PolicyConfig::Random; + + // Both specific policies should be used + match pd.get_prefill_policy(&main_policy) { + PolicyConfig::CacheAware { .. } => {} // Success + _ => panic!("Expected CacheAware for prefill"), + } + + match pd.get_decode_policy(&main_policy) { + PolicyConfig::PowerOfTwo { .. } => {} // Success + _ => panic!("Expected PowerOfTwo for decode"), + } + } + + #[test] + fn test_pd_policy_fallback_only_prefill() { + // When only prefill policy is specified, decode should use main policy + let pd = RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec!["http://prefill1".to_string()], + decode_urls: vec!["http://decode1".to_string()], + prefill_policy: Some(PolicyConfig::CacheAware { + cache_threshold: 0.5, + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + eviction_interval_secs: 60, + max_tree_size: 1000, + }), + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::default(), + rdma_config: None, + dummy_prefill: false, + }; + + let main_policy = PolicyConfig::RoundRobin; + + // Prefill should use specific policy + match pd.get_prefill_policy(&main_policy) { + PolicyConfig::CacheAware { .. } => {} // Success + _ => panic!("Expected CacheAware for prefill"), + } + + // Decode should fall back to main policy + match pd.get_decode_policy(&main_policy) { + PolicyConfig::RoundRobin => {} // Success + _ => panic!("Expected RoundRobin for decode"), + } + } + + #[test] + fn test_pd_policy_fallback_only_decode() { + // When only decode policy is specified, prefill should use main policy + let pd = RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec!["http://prefill1".to_string()], + decode_urls: vec!["http://decode1".to_string()], + prefill_policy: None, + decode_policy: Some(PolicyConfig::PowerOfTwo { + load_check_interval_secs: 60, + }), + migration_protocol: LMDeployMigrationProtocol::default(), + rdma_config: None, + dummy_prefill: false, + }; + + let main_policy = PolicyConfig::Random; + + // Prefill should fall back to main policy + match pd.get_prefill_policy(&main_policy) { + PolicyConfig::Random => {} // Success + _ => panic!("Expected Random for prefill"), + } + + // Decode should use specific policy + match pd.get_decode_policy(&main_policy) { + PolicyConfig::PowerOfTwo { .. } => {} // Success + _ => panic!("Expected PowerOfTwo for decode"), + } + } + + #[test] + fn test_pd_policy_fallback_none_specified() { + // When no specific policies are specified, both should use main policy + let pd = RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec!["http://prefill1".to_string()], + decode_urls: vec!["http://decode1".to_string()], + prefill_policy: None, + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::default(), + rdma_config: None, + dummy_prefill: false, + }; + + let main_policy = PolicyConfig::CacheAware { + cache_threshold: 0.7, + balance_abs_threshold: 20, + balance_rel_threshold: 1.5, + eviction_interval_secs: 300, + max_tree_size: 2000, + }; + + // Both should fall back to main policy + match pd.get_prefill_policy(&main_policy) { + PolicyConfig::CacheAware { + cache_threshold, .. + } => { + assert!((cache_threshold - 0.7).abs() < 0.0001); + } + _ => panic!("Expected CacheAware for prefill"), + } + + match pd.get_decode_policy(&main_policy) { + PolicyConfig::CacheAware { + cache_threshold, .. + } => { + assert!((cache_threshold - 0.7).abs() < 0.0001); + } + _ => panic!("Expected CacheAware for decode"), + } + } + + #[test] + fn test_regular_mode_policy_fallback() { + // For regular mode, the helper methods should just return the main policy + let regular = RoutingMode::Regular { + worker_urls: vec!["http://worker1".to_string()], + }; + + let main_policy = PolicyConfig::RoundRobin; + + // Both methods should return main policy for regular mode + match regular.get_prefill_policy(&main_policy) { + PolicyConfig::RoundRobin => {} // Success + _ => panic!("Expected RoundRobin for regular mode"), + } + + match regular.get_decode_policy(&main_policy) { + PolicyConfig::RoundRobin => {} // Success + _ => panic!("Expected RoundRobin for regular mode"), + } + } +} diff --git a/lmdeploy/router/src/config/validation.rs b/lmdeploy/router/src/config/validation.rs new file mode 100644 index 0000000000..c7b193742b --- /dev/null +++ b/lmdeploy/router/src/config/validation.rs @@ -0,0 +1,718 @@ +use super::*; + +/// Configuration validator +pub struct ConfigValidator; + +impl ConfigValidator { + /// Validate a complete router configuration + pub fn validate(config: &RouterConfig) -> ConfigResult<()> { + let has_service_discovery = config.discovery.as_ref().is_some_and(|d| d.enabled); + + Self::validate_mode(&config.mode, has_service_discovery)?; + Self::validate_policy(&config.policy)?; + Self::validate_server_settings(config)?; + + if let Some(discovery) = &config.discovery { + Self::validate_discovery(discovery, &config.mode)?; + } + + if let Some(metrics) = &config.metrics { + Self::validate_metrics(metrics)?; + } + + Self::validate_compatibility(config)?; + + // Validate effective retry/CB configs (respect disable flags) + let retry_cfg = config.effective_retry_config(); + let cb_cfg = config.effective_circuit_breaker_config(); + Self::validate_retry(&retry_cfg)?; + Self::validate_circuit_breaker(&cb_cfg)?; + + Ok(()) + } + + /// Validate routing mode configuration + fn validate_mode(mode: &RoutingMode, _has_service_discovery: bool) -> ConfigResult<()> { + match mode { + RoutingMode::Regular { worker_urls } => { + // Validate URLs if any are provided + if !worker_urls.is_empty() { + Self::validate_urls(worker_urls)?; + } + // Note: We allow empty worker URLs even without service discovery + // to let the router start and fail at runtime when routing requests. + // This matches legacy behavior and test expectations. + } + RoutingMode::LMDeployPrefillDecode { + prefill_urls, + decode_urls, + prefill_policy, + decode_policy, + migration_protocol: _, + rdma_config: _, + dummy_prefill: _, + } => { + // Empty URL lists are valid: LMDeploy API servers can register + // later through the native-compatible /nodes/add endpoint. + if !prefill_urls.is_empty() { + Self::validate_urls(prefill_urls)?; + } + if !decode_urls.is_empty() { + Self::validate_urls(decode_urls)?; + } + if let Some(p_policy) = prefill_policy { + Self::validate_policy(p_policy)?; + } + if let Some(d_policy) = decode_policy { + Self::validate_policy(d_policy)?; + } + } + RoutingMode::OpenAI { worker_urls } => { + // Require exactly one worker URL for OpenAI router + if worker_urls.len() != 1 { + return Err(ConfigError::ValidationFailed { + reason: "OpenAI mode requires exactly one --worker-urls entry".to_string(), + }); + } + // Validate URL format + if let Err(e) = url::Url::parse(&worker_urls[0]) { + return Err(ConfigError::ValidationFailed { + reason: format!("Invalid OpenAI worker URL '{}': {}", worker_urls[0], e), + }); + } + } + } + Ok(()) + } + + /// Validate policy configuration + fn validate_policy(policy: &PolicyConfig) -> ConfigResult<()> { + match policy { + PolicyConfig::Random | PolicyConfig::RoundRobin => { + // No specific validation needed + } + PolicyConfig::CacheAware { + cache_threshold, + balance_abs_threshold: _, + balance_rel_threshold, + eviction_interval_secs, + max_tree_size, + } => { + if !(0.0..=1.0).contains(cache_threshold) { + return Err(ConfigError::InvalidValue { + field: "cache_threshold".to_string(), + value: cache_threshold.to_string(), + reason: "Must be between 0.0 and 1.0".to_string(), + }); + } + + if *balance_rel_threshold < 1.0 { + return Err(ConfigError::InvalidValue { + field: "balance_rel_threshold".to_string(), + value: balance_rel_threshold.to_string(), + reason: "Must be >= 1.0".to_string(), + }); + } + + if *eviction_interval_secs == 0 { + return Err(ConfigError::InvalidValue { + field: "eviction_interval_secs".to_string(), + value: eviction_interval_secs.to_string(), + reason: "Must be > 0".to_string(), + }); + } + + if *max_tree_size == 0 { + return Err(ConfigError::InvalidValue { + field: "max_tree_size".to_string(), + value: max_tree_size.to_string(), + reason: "Must be > 0".to_string(), + }); + } + } + PolicyConfig::PowerOfTwo { + load_check_interval_secs, + } => { + if *load_check_interval_secs == 0 { + return Err(ConfigError::InvalidValue { + field: "load_check_interval_secs".to_string(), + value: load_check_interval_secs.to_string(), + reason: "Must be > 0".to_string(), + }); + } + } + PolicyConfig::ConsistentHash { virtual_nodes } => { + if *virtual_nodes == 0 { + return Err(ConfigError::InvalidValue { + field: "virtual_nodes".to_string(), + value: virtual_nodes.to_string(), + reason: "Must be > 0".to_string(), + }); + } + } + PolicyConfig::RendezvousHash => { + // No specific validation needed + } + } + Ok(()) + } + + /// Validate server configuration + fn validate_server_settings(config: &RouterConfig) -> ConfigResult<()> { + if config.port == 0 { + return Err(ConfigError::InvalidValue { + field: "port".to_string(), + value: config.port.to_string(), + reason: "Port must be > 0".to_string(), + }); + } + + if config.max_payload_size == 0 { + return Err(ConfigError::InvalidValue { + field: "max_payload_size".to_string(), + value: config.max_payload_size.to_string(), + reason: "Must be > 0".to_string(), + }); + } + + if config.request_timeout_secs == 0 { + return Err(ConfigError::InvalidValue { + field: "request_timeout_secs".to_string(), + value: config.request_timeout_secs.to_string(), + reason: "Must be > 0".to_string(), + }); + } + + if config.worker_startup_timeout_secs == 0 { + return Err(ConfigError::InvalidValue { + field: "worker_startup_timeout_secs".to_string(), + value: config.worker_startup_timeout_secs.to_string(), + reason: "Must be > 0".to_string(), + }); + } + + if config.worker_startup_check_interval_secs == 0 { + return Err(ConfigError::InvalidValue { + field: "worker_startup_check_interval_secs".to_string(), + value: config.worker_startup_check_interval_secs.to_string(), + reason: "Must be > 0".to_string(), + }); + } + + Ok(()) + } + + /// Validate service discovery configuration + fn validate_discovery(discovery: &DiscoveryConfig, mode: &RoutingMode) -> ConfigResult<()> { + if !discovery.enabled { + return Ok(()); // No validation needed if disabled + } + + if discovery.port == 0 { + return Err(ConfigError::InvalidValue { + field: "discovery.port".to_string(), + value: discovery.port.to_string(), + reason: "Port must be > 0".to_string(), + }); + } + + if discovery.check_interval_secs == 0 { + return Err(ConfigError::InvalidValue { + field: "discovery.check_interval_secs".to_string(), + value: discovery.check_interval_secs.to_string(), + reason: "Must be > 0".to_string(), + }); + } + + // Validate selectors based on mode + match mode { + RoutingMode::Regular { .. } => { + if discovery.selector.is_empty() { + return Err(ConfigError::ValidationFailed { + reason: "Regular mode with service discovery requires a non-empty selector" + .to_string(), + }); + } + } + RoutingMode::LMDeployPrefillDecode { .. } => { + if discovery.prefill_selector.is_empty() && discovery.decode_selector.is_empty() { + return Err(ConfigError::ValidationFailed { + reason: "LMDeploy PD mode with service discovery requires at least one non-empty selector (prefill or decode)".to_string(), + }); + } + } + RoutingMode::OpenAI { .. } => { + // OpenAI mode doesn't use service discovery + return Err(ConfigError::ValidationFailed { + reason: "OpenAI mode does not support service discovery".to_string(), + }); + } + } + + Ok(()) + } + + /// Validate metrics configuration + fn validate_metrics(metrics: &MetricsConfig) -> ConfigResult<()> { + if metrics.port == 0 { + return Err(ConfigError::InvalidValue { + field: "metrics.port".to_string(), + value: metrics.port.to_string(), + reason: "Port must be > 0".to_string(), + }); + } + + if metrics.host.is_empty() { + return Err(ConfigError::InvalidValue { + field: "metrics.host".to_string(), + value: metrics.host.clone(), + reason: "Host cannot be empty".to_string(), + }); + } + + Ok(()) + } + + /// Validate retry configuration + fn validate_retry(retry: &RetryConfig) -> ConfigResult<()> { + if retry.max_retries < 1 { + return Err(ConfigError::InvalidValue { + field: "retry.max_retries".to_string(), + value: retry.max_retries.to_string(), + reason: "Must be >= 1 (set to 1 to effectively disable retries)".to_string(), + }); + } + if retry.initial_backoff_ms == 0 { + return Err(ConfigError::InvalidValue { + field: "retry.initial_backoff_ms".to_string(), + value: retry.initial_backoff_ms.to_string(), + reason: "Must be > 0".to_string(), + }); + } + if retry.max_backoff_ms < retry.initial_backoff_ms { + return Err(ConfigError::InvalidValue { + field: "retry.max_backoff_ms".to_string(), + value: retry.max_backoff_ms.to_string(), + reason: "Must be >= initial_backoff_ms".to_string(), + }); + } + if retry.backoff_multiplier < 1.0 { + return Err(ConfigError::InvalidValue { + field: "retry.backoff_multiplier".to_string(), + value: retry.backoff_multiplier.to_string(), + reason: "Must be >= 1.0".to_string(), + }); + } + if !(0.0..=1.0).contains(&retry.jitter_factor) { + return Err(ConfigError::InvalidValue { + field: "retry.jitter_factor".to_string(), + value: retry.jitter_factor.to_string(), + reason: "Must be between 0.0 and 1.0".to_string(), + }); + } + Ok(()) + } + + /// Validate circuit breaker configuration + fn validate_circuit_breaker(cb: &CircuitBreakerConfig) -> ConfigResult<()> { + if cb.failure_threshold < 1 { + return Err(ConfigError::InvalidValue { + field: "circuit_breaker.failure_threshold".to_string(), + value: cb.failure_threshold.to_string(), + reason: "Must be >= 1 (set to u32::MAX to effectively disable CB)".to_string(), + }); + } + if cb.success_threshold < 1 { + return Err(ConfigError::InvalidValue { + field: "circuit_breaker.success_threshold".to_string(), + value: cb.success_threshold.to_string(), + reason: "Must be >= 1".to_string(), + }); + } + if cb.timeout_duration_secs == 0 { + return Err(ConfigError::InvalidValue { + field: "circuit_breaker.timeout_duration_secs".to_string(), + value: cb.timeout_duration_secs.to_string(), + reason: "Must be > 0".to_string(), + }); + } + if cb.window_duration_secs == 0 { + return Err(ConfigError::InvalidValue { + field: "circuit_breaker.window_duration_secs".to_string(), + value: cb.window_duration_secs.to_string(), + reason: "Must be > 0".to_string(), + }); + } + Ok(()) + } + + /// Validate compatibility between different configuration sections + fn validate_compatibility(config: &RouterConfig) -> ConfigResult<()> { + // IGW mode is independent - skip other compatibility checks when enabled + if config.enable_igw { + return Ok(()); + } + + // All policies are now supported for both router types thanks to the unified trait design + // No mode/policy restrictions needed anymore + + let has_service_discovery = config.discovery.as_ref().is_some_and(|d| d.enabled); + + // Only validate worker counts if service discovery is disabled + if !has_service_discovery { + // Check if power-of-two policy makes sense with insufficient workers + if let PolicyConfig::PowerOfTwo { .. } = &config.policy { + let worker_count = config.mode.worker_count(); + // Zero workers is dynamic-registration mode. Validate once + // static workers have actually been supplied. + if worker_count != 0 && worker_count < 2 { + return Err(ConfigError::IncompatibleConfig { + reason: "Power-of-two policy requires at least 2 workers".to_string(), + }); + } + } + + // For LMDeploy PD mode, validate that policies have sufficient workers + if let RoutingMode::LMDeployPrefillDecode { + prefill_urls, + decode_urls, + prefill_policy, + decode_policy, + dummy_prefill, + .. + } = &config.mode + { + if !dummy_prefill { + if let Some(PolicyConfig::PowerOfTwo { .. }) = prefill_policy { + if !prefill_urls.is_empty() && prefill_urls.len() < 2 { + return Err(ConfigError::IncompatibleConfig { + reason: "Power-of-two policy for prefill requires at least 2 prefill workers".to_string(), + }); + } + } + } + if let Some(PolicyConfig::PowerOfTwo { .. }) = decode_policy { + if !decode_urls.is_empty() && decode_urls.len() < 2 { + return Err(ConfigError::IncompatibleConfig { + reason: + "Power-of-two policy for decode requires at least 2 decode workers" + .to_string(), + }); + } + } + } + } + + Ok(()) + } + + /// Validate URL format + fn validate_urls(urls: &[String]) -> ConfigResult<()> { + for url in urls { + if url.is_empty() { + return Err(ConfigError::InvalidValue { + field: "worker_url".to_string(), + value: url.clone(), + reason: "URL cannot be empty".to_string(), + }); + } + + if !url.starts_with("http://") && !url.starts_with("https://") { + return Err(ConfigError::InvalidValue { + field: "worker_url".to_string(), + value: url.clone(), + reason: "URL must start with http:// or https://".to_string(), + }); + } + + // Basic URL validation + match ::url::Url::parse(url) { + Ok(parsed) => { + // Additional validation + if parsed.host_str().is_none() { + return Err(ConfigError::InvalidValue { + field: "worker_url".to_string(), + value: url.clone(), + reason: "URL must have a valid host".to_string(), + }); + } + } + Err(e) => { + return Err(ConfigError::InvalidValue { + field: "worker_url".to_string(), + value: url.clone(), + reason: format!("Invalid URL format: {}", e), + }); + } + } + } + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_validate_regular_mode() { + let config = RouterConfig::new( + RoutingMode::Regular { + worker_urls: vec!["http://worker:8000".to_string()], + }, + PolicyConfig::Random, + ); + + assert!(ConfigValidator::validate(&config).is_ok()); + } + + #[test] + fn test_validate_empty_worker_urls() { + let config = RouterConfig::new( + RoutingMode::Regular { + worker_urls: vec![], + }, + PolicyConfig::Random, + ); + + // Empty worker URLs are now allowed to match legacy behavior + assert!(ConfigValidator::validate(&config).is_ok()); + } + + #[test] + fn test_validate_empty_worker_urls_with_service_discovery() { + let mut config = RouterConfig::new( + RoutingMode::Regular { + worker_urls: vec![], + }, + PolicyConfig::Random, + ); + + // Enable service discovery + config.discovery = Some(DiscoveryConfig { + enabled: true, + selector: vec![("app".to_string(), "test".to_string())] + .into_iter() + .collect(), + ..Default::default() + }); + + // Should pass validation since service discovery is enabled + assert!(ConfigValidator::validate(&config).is_ok()); + } + + #[test] + fn test_validate_invalid_urls() { + let config = RouterConfig::new( + RoutingMode::Regular { + worker_urls: vec!["invalid-url".to_string()], + }, + PolicyConfig::Random, + ); + + assert!(ConfigValidator::validate(&config).is_err()); + } + + #[test] + fn test_validate_cache_aware_thresholds() { + let config = RouterConfig::new( + RoutingMode::Regular { + worker_urls: vec![ + "http://worker1:8000".to_string(), + "http://worker2:8000".to_string(), + ], + }, + PolicyConfig::CacheAware { + cache_threshold: 1.5, // Invalid: > 1.0 + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + eviction_interval_secs: 60, + max_tree_size: 1000, + }, + ); + + assert!(ConfigValidator::validate(&config).is_err()); + } + + #[test] + fn test_validate_cache_aware_single_worker() { + // Cache-aware with single worker should be allowed (even if not optimal) + let config = RouterConfig::new( + RoutingMode::Regular { + worker_urls: vec!["http://worker1:8000".to_string()], + }, + PolicyConfig::CacheAware { + cache_threshold: 0.5, + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + eviction_interval_secs: 60, + max_tree_size: 1000, + }, + ); + + assert!(ConfigValidator::validate(&config).is_ok()); + } + + #[test] + fn test_validate_pd_mode() { + let config = RouterConfig::new( + RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec!["http://prefill:8000".to_string()], + decode_urls: vec!["http://decode:8000".to_string()], + prefill_policy: None, + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::Rdma, + rdma_config: None, + dummy_prefill: false, + }, + PolicyConfig::Random, + ); + + assert!(ConfigValidator::validate(&config).is_ok()); + } + + #[test] + fn test_validate_empty_lmdeploy_pd_for_dynamic_registration() { + let config = RouterConfig::new( + RoutingMode::LMDeployPrefillDecode { + prefill_urls: Vec::new(), + decode_urls: Vec::new(), + prefill_policy: None, + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::Rdma, + rdma_config: None, + dummy_prefill: false, + }, + PolicyConfig::RoundRobin, + ); + + assert!(ConfigValidator::validate(&config).is_ok()); + } + + #[test] + fn test_validate_roundrobin_with_pd_mode() { + // RoundRobin with PD mode is now supported + let config = RouterConfig::new( + RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec!["http://prefill:8000".to_string()], + decode_urls: vec!["http://decode:8000".to_string()], + prefill_policy: None, + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::Rdma, + rdma_config: None, + dummy_prefill: false, + }, + PolicyConfig::RoundRobin, + ); + + let result = ConfigValidator::validate(&config); + assert!(result.is_ok()); + } + + #[test] + fn test_validate_cache_aware_with_pd_mode() { + // CacheAware with PD mode is now supported + let config = RouterConfig::new( + RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec!["http://prefill:8000".to_string()], + decode_urls: vec!["http://decode:8000".to_string()], + prefill_policy: None, + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::Rdma, + rdma_config: None, + dummy_prefill: false, + }, + PolicyConfig::CacheAware { + cache_threshold: 0.5, + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + eviction_interval_secs: 60, + max_tree_size: 1000, + }, + ); + + let result = ConfigValidator::validate(&config); + assert!(result.is_ok()); + } + + #[test] + fn test_validate_power_of_two_with_regular_mode() { + // PowerOfTwo with Regular mode is now supported + let config = RouterConfig::new( + RoutingMode::Regular { + worker_urls: vec![ + "http://worker1:8000".to_string(), + "http://worker2:8000".to_string(), + ], + }, + PolicyConfig::PowerOfTwo { + load_check_interval_secs: 60, + }, + ); + + let result = ConfigValidator::validate(&config); + assert!(result.is_ok()); + } + + #[test] + fn test_validate_pd_mode_with_separate_policies() { + // Test PD mode with different policies for prefill and decode + let config = RouterConfig::new( + RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec![ + "http://prefill1:8000".to_string(), + "http://prefill2:8000".to_string(), + ], + decode_urls: vec![ + "http://decode1:8000".to_string(), + "http://decode2:8000".to_string(), + ], + prefill_policy: Some(PolicyConfig::CacheAware { + cache_threshold: 0.5, + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + eviction_interval_secs: 60, + max_tree_size: 1000, + }), + decode_policy: Some(PolicyConfig::PowerOfTwo { + load_check_interval_secs: 60, + }), + migration_protocol: LMDeployMigrationProtocol::Rdma, + rdma_config: None, + dummy_prefill: false, + }, + PolicyConfig::Random, // Main policy as fallback + ); + + let result = ConfigValidator::validate(&config); + assert!(result.is_ok()); + } + + #[test] + fn test_validate_pd_mode_power_of_two_insufficient_workers() { + // Test that power-of-two policy requires at least 2 workers + let config = RouterConfig::new( + RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec!["http://prefill1:8000".to_string()], // Only 1 prefill + decode_urls: vec![ + "http://decode1:8000".to_string(), + "http://decode2:8000".to_string(), + ], + prefill_policy: Some(PolicyConfig::PowerOfTwo { + load_check_interval_secs: 60, + }), // Requires 2+ workers + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::Rdma, + rdma_config: None, + dummy_prefill: false, + }, + PolicyConfig::Random, + ); + + let result = ConfigValidator::validate(&config); + assert!(result.is_err()); + if let Err(e) = result { + assert!(e.to_string().contains("prefill requires at least 2")); + } + } +} diff --git a/lmdeploy/router/src/core/circuit_breaker.rs b/lmdeploy/router/src/core/circuit_breaker.rs new file mode 100644 index 0000000000..5c374233e0 --- /dev/null +++ b/lmdeploy/router/src/core/circuit_breaker.rs @@ -0,0 +1,555 @@ +use std::sync::atomic::{AtomicU32, AtomicU64, Ordering}; +use std::sync::{Arc, RwLock}; +use std::time::{Duration, Instant}; +use tracing::info; + +/// Circuit breaker configuration +#[derive(Debug, Clone)] +pub struct CircuitBreakerConfig { + /// Number of consecutive failures to open the circuit + pub failure_threshold: u32, + /// Success threshold to close circuit from half-open + pub success_threshold: u32, + /// Duration to wait before attempting half-open + pub timeout_duration: Duration, + /// Time window for failure counting + pub window_duration: Duration, +} + +impl Default for CircuitBreakerConfig { + fn default() -> Self { + Self { + failure_threshold: 5, + success_threshold: 2, + timeout_duration: Duration::from_secs(30), + window_duration: Duration::from_secs(60), + } + } +} + +/// Circuit breaker state +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum CircuitState { + /// Normal operation - requests are allowed + Closed, + /// Circuit is open - requests are rejected + Open, + /// Testing if service has recovered - limited requests allowed + HalfOpen, +} + +impl std::fmt::Display for CircuitState { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + CircuitState::Closed => write!(f, "Closed"), + CircuitState::Open => write!(f, "Open"), + CircuitState::HalfOpen => write!(f, "HalfOpen"), + } + } +} + +/// Circuit breaker implementation +#[derive(Debug)] +pub struct CircuitBreaker { + state: Arc>, + consecutive_failures: Arc, + consecutive_successes: Arc, + total_failures: Arc, + total_successes: Arc, + last_failure_time: Arc>>, + last_state_change: Arc>, + config: CircuitBreakerConfig, +} + +impl CircuitBreaker { + /// Create a new circuit breaker with default configuration + pub fn new() -> Self { + Self::with_config(CircuitBreakerConfig::default()) + } + + /// Create a new circuit breaker with custom configuration + pub fn with_config(config: CircuitBreakerConfig) -> Self { + Self { + state: Arc::new(RwLock::new(CircuitState::Closed)), + consecutive_failures: Arc::new(AtomicU32::new(0)), + consecutive_successes: Arc::new(AtomicU32::new(0)), + total_failures: Arc::new(AtomicU64::new(0)), + total_successes: Arc::new(AtomicU64::new(0)), + last_failure_time: Arc::new(RwLock::new(None)), + last_state_change: Arc::new(RwLock::new(Instant::now())), + config, + } + } + + /// Check if a request can be executed + pub fn can_execute(&self) -> bool { + // First check if we need to transition from Open to HalfOpen + self.check_and_update_state(); + + let state = *self.state.read().unwrap(); + match state { + CircuitState::Closed => true, + CircuitState::Open => false, + CircuitState::HalfOpen => true, // Allow limited requests in half-open state + } + } + + /// Get the current state + pub fn state(&self) -> CircuitState { + self.check_and_update_state(); + *self.state.read().unwrap() + } + + /// Record the outcome of a request + pub fn record_outcome(&self, success: bool) { + if success { + self.record_success(); + } else { + self.record_failure(); + } + } + + /// Record a successful request + pub fn record_success(&self) { + self.total_successes.fetch_add(1, Ordering::Relaxed); + self.consecutive_failures.store(0, Ordering::Release); + let successes = self.consecutive_successes.fetch_add(1, Ordering::AcqRel) + 1; + // Outcome-level metrics are recorded at the worker level where the worker label is known + + let current_state = *self.state.read().unwrap(); + + match current_state { + CircuitState::HalfOpen => { + // Check if we've reached the success threshold to close the circuit + if successes >= self.config.success_threshold { + self.transition_to(CircuitState::Closed); + } + } + CircuitState::Closed => { + // Already closed, nothing to do + } + CircuitState::Open => { + // Shouldn't happen, but if it does, stay open + tracing::warn!("Success recorded while circuit is open"); + } + } + } + + /// Record a failed request + pub fn record_failure(&self) { + self.total_failures.fetch_add(1, Ordering::Relaxed); + self.consecutive_successes.store(0, Ordering::Release); + let failures = self.consecutive_failures.fetch_add(1, Ordering::AcqRel) + 1; + // Outcome-level metrics are recorded at the worker level where the worker label is known + + // Update last failure time + { + let mut last_failure = self.last_failure_time.write().unwrap(); + *last_failure = Some(Instant::now()); + } + + let current_state = *self.state.read().unwrap(); + + match current_state { + CircuitState::Closed => { + // Check if we've reached the failure threshold to open the circuit + if failures >= self.config.failure_threshold { + self.transition_to(CircuitState::Open); + } + } + CircuitState::HalfOpen => { + // Single failure in half-open state reopens the circuit + self.transition_to(CircuitState::Open); + } + CircuitState::Open => { + // Already open, nothing to do + } + } + } + + /// Check and update state based on timeout + fn check_and_update_state(&self) { + let current_state = *self.state.read().unwrap(); + + if current_state == CircuitState::Open { + // Check if timeout has expired + let last_change = *self.last_state_change.read().unwrap(); + if last_change.elapsed() >= self.config.timeout_duration { + self.transition_to(CircuitState::HalfOpen); + } + } + } + + /// Transition to a new state + fn transition_to(&self, new_state: CircuitState) { + let mut state = self.state.write().unwrap(); + let old_state = *state; + + if old_state != new_state { + *state = new_state; + + // Update last state change time + let mut last_change = self.last_state_change.write().unwrap(); + *last_change = Instant::now(); + + // Reset counters based on transition + match new_state { + CircuitState::Closed => { + self.consecutive_failures.store(0, Ordering::Release); + self.consecutive_successes.store(0, Ordering::Release); + } + CircuitState::Open => { + self.consecutive_successes.store(0, Ordering::Release); + } + CircuitState::HalfOpen => { + self.consecutive_failures.store(0, Ordering::Release); + self.consecutive_successes.store(0, Ordering::Release); + } + } + + let from = match old_state { + CircuitState::Closed => "closed", + CircuitState::Open => "open", + CircuitState::HalfOpen => "half_open", + }; + let to = match new_state { + CircuitState::Closed => "closed", + CircuitState::Open => "open", + CircuitState::HalfOpen => "half_open", + }; + info!("Circuit breaker state transition: {} -> {}", from, to); + // Transition metrics are recorded at the worker level where the worker label is known + } + } + + /// Get the number of consecutive failures + pub fn failure_count(&self) -> u32 { + self.consecutive_failures.load(Ordering::Acquire) + } + + /// Get the number of consecutive successes + pub fn success_count(&self) -> u32 { + self.consecutive_successes.load(Ordering::Acquire) + } + + /// Get total failures + pub fn total_failures(&self) -> u64 { + self.total_failures.load(Ordering::Relaxed) + } + + /// Get total successes + pub fn total_successes(&self) -> u64 { + self.total_successes.load(Ordering::Relaxed) + } + + /// Get time since last failure + pub fn time_since_last_failure(&self) -> Option { + self.last_failure_time.read().unwrap().map(|t| t.elapsed()) + } + + /// Get time since last state change + pub fn time_since_last_state_change(&self) -> Duration { + self.last_state_change.read().unwrap().elapsed() + } + + /// Check if the circuit is in a half-open state + pub fn is_half_open(&self) -> bool { + self.state() == CircuitState::HalfOpen + } + + /// Record a test success (for health check probing) + pub fn record_test_success(&self) { + if self.is_half_open() { + self.record_success(); + } + } + + /// Record a test failure (for health check probing) + pub fn record_test_failure(&self) { + if self.is_half_open() { + self.record_failure(); + } + } + + /// Reset the circuit breaker to closed state + pub fn reset(&self) { + self.transition_to(CircuitState::Closed); + self.consecutive_failures.store(0, Ordering::Release); + self.consecutive_successes.store(0, Ordering::Release); + } + + /// Force the circuit to open (for manual intervention) + pub fn force_open(&self) { + self.transition_to(CircuitState::Open); + } + + /// Get circuit breaker statistics + pub fn stats(&self) -> CircuitBreakerStats { + CircuitBreakerStats { + state: self.state(), + consecutive_failures: self.failure_count(), + consecutive_successes: self.success_count(), + total_failures: self.total_failures(), + total_successes: self.total_successes(), + time_since_last_failure: self.time_since_last_failure(), + time_since_last_state_change: self.time_since_last_state_change(), + } + } +} + +impl Clone for CircuitBreaker { + fn clone(&self) -> Self { + Self { + state: Arc::clone(&self.state), + consecutive_failures: Arc::clone(&self.consecutive_failures), + consecutive_successes: Arc::clone(&self.consecutive_successes), + total_failures: Arc::clone(&self.total_failures), + total_successes: Arc::clone(&self.total_successes), + last_failure_time: Arc::clone(&self.last_failure_time), + last_state_change: Arc::clone(&self.last_state_change), + config: self.config.clone(), + } + } +} + +impl Default for CircuitBreaker { + fn default() -> Self { + Self::new() + } +} + +/// Circuit breaker statistics +#[derive(Debug, Clone)] +pub struct CircuitBreakerStats { + pub state: CircuitState, + pub consecutive_failures: u32, + pub consecutive_successes: u32, + pub total_failures: u64, + pub total_successes: u64, + pub time_since_last_failure: Option, + pub time_since_last_state_change: Duration, +} + +#[cfg(test)] +mod tests { + use super::*; + use std::thread; + + #[test] + fn test_circuit_breaker_initial_state() { + let cb = CircuitBreaker::new(); + assert_eq!(cb.state(), CircuitState::Closed); + assert!(cb.can_execute()); + assert_eq!(cb.failure_count(), 0); + assert_eq!(cb.success_count(), 0); + } + + #[test] + fn test_circuit_opens_on_threshold() { + let config = CircuitBreakerConfig { + failure_threshold: 3, + ..Default::default() + }; + let cb = CircuitBreaker::with_config(config); + + // Record failures up to threshold + assert_eq!(cb.state(), CircuitState::Closed); + cb.record_failure(); + assert_eq!(cb.state(), CircuitState::Closed); + cb.record_failure(); + assert_eq!(cb.state(), CircuitState::Closed); + cb.record_failure(); + + // Circuit should now be open + assert_eq!(cb.state(), CircuitState::Open); + assert!(!cb.can_execute()); + assert_eq!(cb.failure_count(), 3); + } + + #[test] + fn test_circuit_half_open_after_timeout() { + let config = CircuitBreakerConfig { + failure_threshold: 1, + timeout_duration: Duration::from_millis(100), + ..Default::default() + }; + let cb = CircuitBreaker::with_config(config); + + // Open the circuit + cb.record_failure(); + assert_eq!(cb.state(), CircuitState::Open); + + // Wait for timeout + thread::sleep(Duration::from_millis(150)); + + // Circuit should be half-open + assert_eq!(cb.state(), CircuitState::HalfOpen); + assert!(cb.can_execute()); + } + + #[test] + fn test_circuit_closes_on_success_threshold() { + let config = CircuitBreakerConfig { + failure_threshold: 1, + success_threshold: 2, + timeout_duration: Duration::from_millis(50), + ..Default::default() + }; + let cb = CircuitBreaker::with_config(config); + + // Open the circuit + cb.record_failure(); + assert_eq!(cb.state(), CircuitState::Open); + + // Wait for timeout + thread::sleep(Duration::from_millis(100)); + assert_eq!(cb.state(), CircuitState::HalfOpen); + + // Record successes + cb.record_success(); + assert_eq!(cb.state(), CircuitState::HalfOpen); + cb.record_success(); + + // Circuit should now be closed + assert_eq!(cb.state(), CircuitState::Closed); + assert!(cb.can_execute()); + } + + #[test] + fn test_circuit_reopens_on_half_open_failure() { + let config = CircuitBreakerConfig { + failure_threshold: 1, + timeout_duration: Duration::from_millis(50), + ..Default::default() + }; + let cb = CircuitBreaker::with_config(config); + + // Open the circuit + cb.record_failure(); + assert_eq!(cb.state(), CircuitState::Open); + + // Wait for timeout + thread::sleep(Duration::from_millis(100)); + assert_eq!(cb.state(), CircuitState::HalfOpen); + + // Record a failure in half-open state + cb.record_failure(); + + // Circuit should reopen immediately + assert_eq!(cb.state(), CircuitState::Open); + assert!(!cb.can_execute()); + } + + #[test] + fn test_success_resets_failure_count() { + let config = CircuitBreakerConfig { + failure_threshold: 3, + ..Default::default() + }; + let cb = CircuitBreaker::with_config(config); + + // Record some failures + cb.record_failure(); + cb.record_failure(); + assert_eq!(cb.failure_count(), 2); + + // Success should reset failure count + cb.record_success(); + assert_eq!(cb.failure_count(), 0); + assert_eq!(cb.success_count(), 1); + + // Can now record more failures without opening + cb.record_failure(); + cb.record_failure(); + assert_eq!(cb.state(), CircuitState::Closed); + } + + #[test] + fn test_manual_reset() { + let config = CircuitBreakerConfig { + failure_threshold: 1, + ..Default::default() + }; + let cb = CircuitBreaker::with_config(config); + + // Open the circuit + cb.record_failure(); + assert_eq!(cb.state(), CircuitState::Open); + + // Manual reset + cb.reset(); + assert_eq!(cb.state(), CircuitState::Closed); + assert_eq!(cb.failure_count(), 0); + assert_eq!(cb.success_count(), 0); + } + + #[test] + fn test_force_open() { + let cb = CircuitBreaker::new(); + assert_eq!(cb.state(), CircuitState::Closed); + + cb.force_open(); + assert_eq!(cb.state(), CircuitState::Open); + assert!(!cb.can_execute()); + } + + #[test] + fn test_stats() { + let config = CircuitBreakerConfig { + failure_threshold: 2, + ..Default::default() + }; + let cb = CircuitBreaker::with_config(config); + + cb.record_success(); + cb.record_failure(); + cb.record_failure(); + + let stats = cb.stats(); + assert_eq!(stats.state, CircuitState::Open); + assert_eq!(stats.consecutive_failures, 2); + assert_eq!(stats.consecutive_successes, 0); + assert_eq!(stats.total_failures, 2); + assert_eq!(stats.total_successes, 1); + } + + #[test] + fn test_clone() { + let cb1 = CircuitBreaker::new(); + cb1.record_failure(); + + let cb2 = cb1.clone(); + assert_eq!(cb2.failure_count(), 1); + + // Changes to cb1 affect cb2 (shared state) + cb1.record_failure(); + assert_eq!(cb2.failure_count(), 2); + } + + #[test] + fn test_thread_safety() { + use std::sync::Arc; + + let cb = Arc::new(CircuitBreaker::new()); + let mut handles = vec![]; + + // Spawn threads that record failures + for _ in 0..10 { + let cb_clone = Arc::clone(&cb); + let handle = thread::spawn(move || { + for _ in 0..100 { + cb_clone.record_failure(); + } + }); + handles.push(handle); + } + + // Wait for all threads + for handle in handles { + handle.join().unwrap(); + } + + // Should have recorded 1000 failures + assert_eq!(cb.total_failures(), 1000); + } +} diff --git a/lmdeploy/router/src/core/error.rs b/lmdeploy/router/src/core/error.rs new file mode 100644 index 0000000000..a5eabc157d --- /dev/null +++ b/lmdeploy/router/src/core/error.rs @@ -0,0 +1,240 @@ +//! Error types for the LMDeploy router core +//! +//! This module defines error types used throughout the router for worker operations. + +use std::fmt; + +/// Worker-related errors +#[derive(Debug)] +pub enum WorkerError { + /// Health check failed + HealthCheckFailed { url: String, reason: String }, + /// Worker not found + WorkerNotFound { url: String }, + /// Invalid worker configuration + InvalidConfiguration { message: String }, + /// Network error + NetworkError { url: String, error: String }, + /// Worker is at capacity + WorkerAtCapacity { url: String }, + /// Invalid URL format + InvalidUrl { url: String }, +} + +impl fmt::Display for WorkerError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + WorkerError::HealthCheckFailed { url, reason } => { + write!(f, "Health check failed for worker {}: {}", url, reason) + } + WorkerError::WorkerNotFound { url } => { + write!(f, "Worker not found: {}", url) + } + WorkerError::InvalidConfiguration { message } => { + write!(f, "Invalid worker configuration: {}", message) + } + WorkerError::NetworkError { url, error } => { + write!(f, "Network error for worker {}: {}", url, error) + } + WorkerError::WorkerAtCapacity { url } => { + write!(f, "Worker at capacity: {}", url) + } + WorkerError::InvalidUrl { url } => { + write!(f, "Invalid URL format: {}", url) + } + } + } +} + +impl std::error::Error for WorkerError {} + +/// Result type for worker operations +pub type WorkerResult = Result; + +/// Convert from reqwest errors to worker errors +impl From for WorkerError { + fn from(err: reqwest::Error) -> Self { + WorkerError::NetworkError { + url: err.url().map(|u| u.to_string()).unwrap_or_default(), + error: err.to_string(), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::error::Error; + + #[test] + fn test_health_check_failed_display() { + let error = WorkerError::HealthCheckFailed { + url: "http://worker1:8080".to_string(), + reason: "Connection refused".to_string(), + }; + assert_eq!( + error.to_string(), + "Health check failed for worker http://worker1:8080: Connection refused" + ); + } + + #[test] + fn test_worker_not_found_display() { + let error = WorkerError::WorkerNotFound { + url: "http://worker2:8080".to_string(), + }; + assert_eq!(error.to_string(), "Worker not found: http://worker2:8080"); + } + + #[test] + fn test_invalid_configuration_display() { + let error = WorkerError::InvalidConfiguration { + message: "Missing port number".to_string(), + }; + assert_eq!( + error.to_string(), + "Invalid worker configuration: Missing port number" + ); + } + + #[test] + fn test_network_error_display() { + let error = WorkerError::NetworkError { + url: "http://worker3:8080".to_string(), + error: "Timeout after 30s".to_string(), + }; + assert_eq!( + error.to_string(), + "Network error for worker http://worker3:8080: Timeout after 30s" + ); + } + + #[test] + fn test_worker_at_capacity_display() { + let error = WorkerError::WorkerAtCapacity { + url: "http://worker4:8080".to_string(), + }; + assert_eq!(error.to_string(), "Worker at capacity: http://worker4:8080"); + } + + #[test] + fn test_worker_error_implements_std_error() { + let error = WorkerError::WorkerNotFound { + url: "http://test".to_string(), + }; + // Verify it implements Error trait + let _: &dyn Error = &error; + assert!(error.source().is_none()); + } + + #[test] + fn test_error_send_sync() { + fn assert_send_sync() {} + assert_send_sync::(); + } + + #[test] + fn test_worker_result_type_alias() { + // Test Ok variant + let result: WorkerResult = Ok(42); + assert!(matches!(result, Ok(42))); + + // Test Err variant + let error = WorkerError::WorkerNotFound { + url: "test".to_string(), + }; + let result: WorkerResult = Err(error); + assert!(result.is_err()); + } + + #[test] + fn test_empty_url_handling() { + // Test empty URLs in error variants + let error1 = WorkerError::HealthCheckFailed { + url: "".to_string(), + reason: "No connection".to_string(), + }; + assert_eq!( + error1.to_string(), + "Health check failed for worker : No connection" + ); + + let error2 = WorkerError::NetworkError { + url: "".to_string(), + error: "DNS failure".to_string(), + }; + assert_eq!(error2.to_string(), "Network error for worker : DNS failure"); + + let error3 = WorkerError::WorkerNotFound { + url: "".to_string(), + }; + assert_eq!(error3.to_string(), "Worker not found: "); + } + + #[test] + fn test_special_characters_in_messages() { + // Test with special characters + let error = WorkerError::InvalidConfiguration { + message: "Invalid JSON: {\"error\": \"test\"}".to_string(), + }; + assert_eq!( + error.to_string(), + "Invalid worker configuration: Invalid JSON: {\"error\": \"test\"}" + ); + + // Test with unicode + let error2 = WorkerError::HealthCheckFailed { + url: "http://测试:8080".to_string(), + reason: "连接被拒绝".to_string(), + }; + assert_eq!( + error2.to_string(), + "Health check failed for worker http://测试:8080: 连接被拒绝" + ); + } + + #[test] + fn test_very_long_error_messages() { + let long_message = "A".repeat(10000); + let error = WorkerError::InvalidConfiguration { + message: long_message.clone(), + }; + let display = error.to_string(); + assert!(display.contains(&long_message)); + assert_eq!( + display.len(), + "Invalid worker configuration: ".len() + long_message.len() + ); + } + + // Mock reqwest error for testing conversion + #[test] + fn test_reqwest_error_conversion() { + // Test that NetworkError is the correct variant + let network_error = WorkerError::NetworkError { + url: "http://example.com".to_string(), + error: "connection timeout".to_string(), + }; + + match network_error { + WorkerError::NetworkError { url, error } => { + assert_eq!(url, "http://example.com"); + assert_eq!(error, "connection timeout"); + } + _ => panic!("Expected NetworkError variant"), + } + } + + #[test] + fn test_error_equality() { + // WorkerError doesn't implement PartialEq, but we can test that + // the same error construction produces the same display output + let error1 = WorkerError::WorkerNotFound { + url: "http://test".to_string(), + }; + let error2 = WorkerError::WorkerNotFound { + url: "http://test".to_string(), + }; + assert_eq!(error1.to_string(), error2.to_string()); + } +} diff --git a/lmdeploy/router/src/core/mod.rs b/lmdeploy/router/src/core/mod.rs new file mode 100644 index 0000000000..a1750471e8 --- /dev/null +++ b/lmdeploy/router/src/core/mod.rs @@ -0,0 +1,26 @@ +//! Core abstractions for the LMDeploy router +//! +//! This module contains the fundamental types and traits used throughout the router: +//! - Worker trait and implementations +//! - Error types +//! - Circuit breaker for reliability +//! - Common utilities + +pub mod circuit_breaker; +pub mod error; +pub mod retry; +pub mod token_bucket; +pub mod worker; +pub mod worker_registry; + +// Re-export commonly used types at the module level +pub use circuit_breaker::{ + CircuitBreaker, CircuitBreakerConfig, CircuitBreakerStats, CircuitState, +}; +pub use error::{WorkerError, WorkerResult}; +pub use retry::{is_retryable_status, BackoffCalculator, RetryError, RetryExecutor}; +pub use worker::{ + start_health_checker, BasicWorker, ConnectionMode, HealthChecker, HealthConfig, Worker, + WorkerCollection, WorkerFactory, WorkerLoadGuard, WorkerType, +}; +pub use worker_registry::{WorkerId, WorkerRegistry, WorkerRegistryStats}; diff --git a/lmdeploy/router/src/core/retry.rs b/lmdeploy/router/src/core/retry.rs new file mode 100644 index 0000000000..8a00424a55 --- /dev/null +++ b/lmdeploy/router/src/core/retry.rs @@ -0,0 +1,409 @@ +use crate::config::types::RetryConfig; +use axum::http::StatusCode; +use axum::response::Response; +use rand::Rng; +use std::time::Duration; +use tracing::debug; + +/// Check if an HTTP status code indicates a retryable error +pub fn is_retryable_status(status: StatusCode) -> bool { + matches!( + status, + StatusCode::REQUEST_TIMEOUT + | StatusCode::TOO_MANY_REQUESTS + | StatusCode::INTERNAL_SERVER_ERROR + | StatusCode::BAD_GATEWAY + | StatusCode::SERVICE_UNAVAILABLE + | StatusCode::GATEWAY_TIMEOUT + ) +} + +/// Computes exponential backoff with optional jitter. +#[derive(Debug, Clone)] +pub struct BackoffCalculator; + +impl BackoffCalculator { + /// Calculate backoff delay for a given attempt index (0-based). + pub fn calculate_delay(config: &RetryConfig, attempt: u32) -> Duration { + // Base exponential backoff + let pow = config.backoff_multiplier.powi(attempt as i32); + let mut delay_ms = (config.initial_backoff_ms as f32 * pow) as u64; + if delay_ms > config.max_backoff_ms { + delay_ms = config.max_backoff_ms; + } + + // Apply jitter in range [-j, +j] + let jitter = config.jitter_factor.clamp(0.0, 1.0); + if jitter > 0.0 { + let mut rng = rand::rng(); + let jitter_scale: f32 = rng.random_range(-jitter..=jitter); + let jitter_ms = (delay_ms as f32 * jitter_scale) + .round() + .max(-(delay_ms as f32)); + let adjusted = (delay_ms as i64 + jitter_ms as i64).max(0) as u64; + return Duration::from_millis(adjusted); + } + + Duration::from_millis(delay_ms) + } +} + +#[derive(Debug, thiserror::Error)] +pub enum RetryError { + #[error("no available workers")] + NoAvailableWorkers, + #[error("maximum retry attempts exceeded")] + MaxRetriesExceeded, +} + +/// A thin async retry executor for generic operations. +#[derive(Debug, Clone, Default)] +pub struct RetryExecutor; + +impl RetryExecutor { + /// Execute an async operation with retries and backoff. + /// The `operation` closure is invoked each attempt with the attempt index. + pub async fn execute_with_retry( + config: &RetryConfig, + mut operation: F, + ) -> Result + where + F: FnMut(u32) -> Fut, + Fut: std::future::Future>, + { + let max = config.max_retries.max(1); + let mut attempt: u32 = 0; + loop { + match operation(attempt).await { + Ok(val) => return Ok(val), + Err(_) => { + // Use the number of failures so far (0-indexed) to compute delay, + // so the first retry uses `initial_backoff_ms`. + let is_last = attempt + 1 >= max; + if is_last { + return Err(RetryError::MaxRetriesExceeded); + } + let delay = BackoffCalculator::calculate_delay(config, attempt); + attempt += 1; // advance to the next attempt after computing delay + tokio::time::sleep(delay).await; + } + } + } + } + + /// Execute an operation that returns an HTTP Response with retries and backoff. + /// + /// Usage pattern: + /// - `operation(attempt)`: perform one attempt (0-based). Construct and send the request, + /// then return the `Response`. Do any per-attempt bookkeeping (e.g., load tracking, + /// circuit-breaker outcome recording) inside this closure. + /// - `should_retry(&response, attempt)`: decide if the given response should be retried + /// (e.g., based on HTTP status). Returning false short-circuits and returns the response. + /// - `on_backoff(delay, next_attempt)`: called before sleeping between attempts. + /// Use this to record metrics. + /// - `on_exhausted()`: called when the executor has exhausted all retry attempts. + /// + /// Example: + /// ```ignore + /// let resp = RetryExecutor::execute_response_with_retry( + /// &retry_cfg, + /// |attempt| async move { + /// let worker = select_cb_aware_worker()?; + /// let resp = send_request(worker).await; + /// worker.record_outcome(resp.status().is_success()); + /// resp + /// }, + /// |res, _| matches!(res.status(), StatusCode::REQUEST_TIMEOUT | StatusCode::TOO_MANY_REQUESTS | StatusCode::INTERNAL_SERVER_ERROR | StatusCode::BAD_GATEWAY | StatusCode::SERVICE_UNAVAILABLE | StatusCode::GATEWAY_TIMEOUT), + /// |delay, attempt| RouterMetrics::record_retry_backoff_duration(delay, attempt), + /// || RouterMetrics::record_retries_exhausted("/route"), + /// ).await; + /// ``` + pub async fn execute_response_with_retry( + config: &RetryConfig, + mut operation: Op, + should_retry: ShouldRetry, + on_backoff: OnBackoff, + mut on_exhausted: OnExhausted, + ) -> Response + where + Op: FnMut(u32) -> Fut, + Fut: std::future::Future, + ShouldRetry: Fn(&Response, u32) -> bool, + OnBackoff: Fn(Duration, u32), + OnExhausted: FnMut(), + { + let max = config.max_retries.max(1); + + let mut attempt: u32 = 0; + loop { + let response = operation(attempt).await; + let is_last = attempt + 1 >= max; + + if !should_retry(&response, attempt) { + return response; + } + + if is_last { + // Exhausted retries + on_exhausted(); + return response; + } + + // Backoff before next attempt + let next_attempt = attempt + 1; + // Compute delay based on the number of failures so far (0-indexed) + let delay = BackoffCalculator::calculate_delay(config, attempt); + debug!( + attempt = attempt, + next_attempt = next_attempt, + delay_ms = delay.as_millis() as u64, + "Retry backoff" + ); + on_backoff(delay, next_attempt); + tokio::time::sleep(delay).await; + + attempt = next_attempt; + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use axum::http::StatusCode; + use axum::response::IntoResponse; + use std::sync::atomic::{AtomicU32, Ordering}; + use std::sync::Arc; + + fn base_retry_config() -> RetryConfig { + RetryConfig { + max_retries: 3, + initial_backoff_ms: 1, + max_backoff_ms: 4, + backoff_multiplier: 2.0, + jitter_factor: 0.0, + } + } + + #[test] + fn test_backoff_no_jitter_progression_and_cap() { + let cfg = RetryConfig { + max_retries: 10, + initial_backoff_ms: 100, + max_backoff_ms: 250, + backoff_multiplier: 2.0, + jitter_factor: 0.0, + }; + // attempt=0 => 100ms + assert_eq!( + BackoffCalculator::calculate_delay(&cfg, 0), + Duration::from_millis(100) + ); + // attempt=1 => 200ms + assert_eq!( + BackoffCalculator::calculate_delay(&cfg, 1), + Duration::from_millis(200) + ); + // attempt=2 => 400ms -> capped to 250ms + assert_eq!( + BackoffCalculator::calculate_delay(&cfg, 2), + Duration::from_millis(250) + ); + // large attempt still capped + assert_eq!( + BackoffCalculator::calculate_delay(&cfg, 10), + Duration::from_millis(250) + ); + } + + #[test] + fn test_backoff_with_jitter_within_bounds() { + let cfg = RetryConfig { + max_retries: 5, + initial_backoff_ms: 100, + max_backoff_ms: 10_000, + backoff_multiplier: 2.0, + jitter_factor: 0.5, + }; + // attempt=2 => base 400ms, jitter in [0.5x, 1.5x] + let base = 400.0; + for _ in 0..50 { + let d = BackoffCalculator::calculate_delay(&cfg, 2).as_millis() as f32; + assert!(d >= base * 0.5 - 1.0 && d <= base * 1.5 + 1.0); + } + } + + #[tokio::test] + async fn test_execute_with_retry_success_after_failures() { + let cfg = base_retry_config(); + let remaining = Arc::new(AtomicU32::new(2)); + let calls = Arc::new(AtomicU32::new(0)); + + let res: Result = RetryExecutor::execute_with_retry(&cfg, { + let remaining = remaining.clone(); + let calls = calls.clone(); + move |_attempt| { + calls.fetch_add(1, Ordering::Relaxed); + let remaining = remaining.clone(); + async move { + if remaining + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |v| v.checked_sub(1)) + .is_ok() + { + Err(()) + } else { + Ok(42u32) + } + } + } + }) + .await; + + assert!(res.is_ok()); + assert_eq!(res.unwrap(), 42); + assert_eq!(calls.load(Ordering::Relaxed), 3); // 2 fails + 1 success + } + + #[tokio::test] + async fn test_execute_with_retry_exhausted() { + let cfg = base_retry_config(); + let calls = Arc::new(AtomicU32::new(0)); + let res: Result = RetryExecutor::execute_with_retry(&cfg, { + let calls = calls.clone(); + move |_attempt| { + calls.fetch_add(1, Ordering::Relaxed); + async move { Err(()) } + } + }) + .await; + + assert!(matches!(res, Err(RetryError::MaxRetriesExceeded))); + assert_eq!(calls.load(Ordering::Relaxed), cfg.max_retries); + } + + #[tokio::test] + async fn test_execute_response_with_retry_success_path_and_hooks() { + let cfg = base_retry_config(); + let remaining = Arc::new(AtomicU32::new(2)); + let calls = Arc::new(AtomicU32::new(0)); + let backoffs = Arc::new(AtomicU32::new(0)); + let exhausted = Arc::new(AtomicU32::new(0)); + + let response = RetryExecutor::execute_response_with_retry( + &cfg, + { + let remaining = remaining.clone(); + let calls = calls.clone(); + move |_attempt| { + calls.fetch_add(1, Ordering::Relaxed); + let remaining = remaining.clone(); + async move { + if remaining + .fetch_update(Ordering::AcqRel, Ordering::Acquire, |v| v.checked_sub(1)) + .is_ok() + { + (StatusCode::SERVICE_UNAVAILABLE, "fail").into_response() + } else { + (StatusCode::OK, "ok").into_response() + } + } + } + }, + |res, _attempt| !res.status().is_success(), // retry until success + { + let backoffs = backoffs.clone(); + move |_delay, _next_attempt| { + backoffs.fetch_add(1, Ordering::Relaxed); + } + }, + { + let exhausted = exhausted.clone(); + move || { + exhausted.fetch_add(1, Ordering::Relaxed); + } + }, + ) + .await; + + assert_eq!(response.status(), StatusCode::OK); + assert_eq!(calls.load(Ordering::Relaxed), 3); // 2 fails + 1 success + assert_eq!(backoffs.load(Ordering::Relaxed), 2); + assert_eq!(exhausted.load(Ordering::Relaxed), 0); + } + + #[tokio::test] + async fn test_execute_response_with_retry_non_retryable_short_circuit() { + let cfg = base_retry_config(); + let calls = Arc::new(AtomicU32::new(0)); + let backoffs = Arc::new(AtomicU32::new(0)); + let exhausted = Arc::new(AtomicU32::new(0)); + + let response = RetryExecutor::execute_response_with_retry( + &cfg, + { + let calls = calls.clone(); + move |_attempt| { + calls.fetch_add(1, Ordering::Relaxed); + async move { (StatusCode::BAD_REQUEST, "bad").into_response() } + } + }, + |_res, _attempt| false, // never retry + { + let backoffs = backoffs.clone(); + move |_delay, _next_attempt| { + backoffs.fetch_add(1, Ordering::Relaxed); + } + }, + { + let exhausted = exhausted.clone(); + move || { + exhausted.fetch_add(1, Ordering::Relaxed); + } + }, + ) + .await; + + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + assert_eq!(calls.load(Ordering::Relaxed), 1); + assert_eq!(backoffs.load(Ordering::Relaxed), 0); + assert_eq!(exhausted.load(Ordering::Relaxed), 0); + } + + #[tokio::test] + async fn test_execute_response_with_retry_exhausted_hooks() { + let cfg = base_retry_config(); + let calls = Arc::new(AtomicU32::new(0)); + let backoffs = Arc::new(AtomicU32::new(0)); + let exhausted = Arc::new(AtomicU32::new(0)); + + let response = RetryExecutor::execute_response_with_retry( + &cfg, + { + let calls = calls.clone(); + move |_attempt| { + calls.fetch_add(1, Ordering::Relaxed); + async move { (StatusCode::SERVICE_UNAVAILABLE, "fail").into_response() } + } + }, + |_res, _attempt| true, // keep retrying + { + let backoffs = backoffs.clone(); + move |_delay, _next_attempt| { + backoffs.fetch_add(1, Ordering::Relaxed); + } + }, + { + let exhausted = exhausted.clone(); + move || { + exhausted.fetch_add(1, Ordering::Relaxed); + } + }, + ) + .await; + + assert_eq!(response.status(), StatusCode::SERVICE_UNAVAILABLE); + assert_eq!(calls.load(Ordering::Relaxed), cfg.max_retries); + assert_eq!(backoffs.load(Ordering::Relaxed), cfg.max_retries - 1); + assert_eq!(exhausted.load(Ordering::Relaxed), 1); + } +} diff --git a/lmdeploy/router/src/core/token_bucket.rs b/lmdeploy/router/src/core/token_bucket.rs new file mode 100644 index 0000000000..5596f9e402 --- /dev/null +++ b/lmdeploy/router/src/core/token_bucket.rs @@ -0,0 +1,195 @@ +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::sync::{Mutex, Notify}; +use tracing::{debug, trace}; + +/// Token bucket for rate limiting +/// +/// This implementation provides: +/// - Smooth rate limiting with configurable refill rate +/// - Burst capacity handling +/// - Fair queuing for waiting requests +#[derive(Clone)] +pub struct TokenBucket { + inner: Arc>, + notify: Arc, + capacity: f64, + refill_rate: f64, // tokens per second +} + +struct TokenBucketInner { + tokens: f64, + last_refill: Instant, +} + +impl TokenBucket { + /// Create a new token bucket + /// + /// # Arguments + /// * `capacity` - Maximum number of tokens (burst capacity) + /// * `refill_rate` - Tokens added per second + pub fn new(capacity: usize, refill_rate: usize) -> Self { + let capacity = capacity as f64; + let refill_rate = refill_rate as f64; + + // Ensure refill_rate is not zero to prevent division by zero + let refill_rate = if refill_rate > 0.0 { + refill_rate + } else { + 1.0 // Default to 1 token per second if zero + }; + + Self { + inner: Arc::new(Mutex::new(TokenBucketInner { + tokens: capacity, // Start full + last_refill: Instant::now(), + })), + notify: Arc::new(Notify::new()), + capacity, + refill_rate, + } + } + + /// Try to acquire tokens immediately + pub async fn try_acquire(&self, tokens: f64) -> bool { + let mut inner = self.inner.lock().await; + + // Refill tokens based on elapsed time + let now = Instant::now(); + let elapsed = now.duration_since(inner.last_refill).as_secs_f64(); + let refill_amount = elapsed * self.refill_rate; + + inner.tokens = (inner.tokens + refill_amount).min(self.capacity); + inner.last_refill = now; + + trace!( + "Token bucket: {} tokens available, requesting {}", + inner.tokens, + tokens + ); + + if inner.tokens >= tokens { + inner.tokens -= tokens; + debug!( + "Token bucket: acquired {} tokens, {} remaining", + tokens, inner.tokens + ); + true + } else { + false + } + } + + /// Acquire tokens, waiting if necessary + pub async fn acquire(&self, tokens: f64) -> Result<(), tokio::time::error::Elapsed> { + // First try to acquire immediately + if self.try_acquire(tokens).await { + return Ok(()); + } + + // Calculate wait time + let wait_time = { + let inner = self.inner.lock().await; + let tokens_needed = (tokens - inner.tokens).max(0.0); + let wait_secs = (tokens_needed / self.refill_rate).max(0.0); + Duration::from_secs_f64(wait_secs) + }; + + debug!( + "Token bucket: waiting {:?} for {} tokens", + wait_time, tokens + ); + + // Wait for tokens to be available + tokio::time::timeout(wait_time, async { + loop { + // Check if we can acquire now + if self.try_acquire(tokens).await { + return; + } + + // Wait for notification or small interval + tokio::select! { + _ = self.notify.notified() => {}, + _ = tokio::time::sleep(Duration::from_millis(10)) => {}, + } + } + }) + .await?; + + Ok(()) + } + + /// Acquire tokens with custom timeout + pub async fn acquire_timeout( + &self, + tokens: f64, + timeout: Duration, + ) -> Result<(), tokio::time::error::Elapsed> { + tokio::time::timeout(timeout, self.acquire(tokens)).await? + } + + /// Return tokens to the bucket (for cancelled requests) + pub async fn return_tokens(&self, tokens: f64) { + let mut inner = self.inner.lock().await; + inner.tokens = (inner.tokens + tokens).min(self.capacity); + self.notify.notify_waiters(); + debug!( + "Token bucket: returned {} tokens, {} available", + tokens, inner.tokens + ); + } + + /// Get current available tokens (for monitoring) + pub async fn available_tokens(&self) -> f64 { + let mut inner = self.inner.lock().await; + + // Refill before checking + let now = Instant::now(); + let elapsed = now.duration_since(inner.last_refill).as_secs_f64(); + let refill_amount = elapsed * self.refill_rate; + + inner.tokens = (inner.tokens + refill_amount).min(self.capacity); + inner.last_refill = now; + + inner.tokens + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_token_bucket_basic() { + let bucket = TokenBucket::new(10, 5); // 10 capacity, 5 per second + + // Should succeed - bucket starts full + assert!(bucket.try_acquire(5.0).await); + assert!(bucket.try_acquire(5.0).await); + + // Should fail - no tokens left + assert!(!bucket.try_acquire(1.0).await); + + // Wait for refill + tokio::time::sleep(Duration::from_millis(300)).await; + + // Should have ~1.5 tokens now + assert!(bucket.try_acquire(1.0).await); + } + + #[tokio::test] + async fn test_token_bucket_refill() { + let bucket = TokenBucket::new(10, 10); // 10 capacity, 10 per second + + // Use all tokens + assert!(bucket.try_acquire(10.0).await); + + // Wait for partial refill + tokio::time::sleep(Duration::from_millis(500)).await; + + // Should have ~5 tokens + let available = bucket.available_tokens().await; + assert!((4.0..=6.0).contains(&available)); + } +} diff --git a/lmdeploy/router/src/core/worker.rs b/lmdeploy/router/src/core/worker.rs new file mode 100644 index 0000000000..80df98dcd0 --- /dev/null +++ b/lmdeploy/router/src/core/worker.rs @@ -0,0 +1,1248 @@ +use super::{CircuitBreaker, CircuitBreakerConfig, WorkerError, WorkerResult}; +use crate::metrics::RouterMetrics; +use async_trait::async_trait; +use futures; +use serde_json; +use std::fmt; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::sync::{Arc, LazyLock}; + +// Shared HTTP client for worker operations (health checks, server info, etc.) +static WORKER_CLIENT: LazyLock = LazyLock::new(|| { + reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(30)) // Default timeout, overridden per request + .build() + .expect("Failed to create worker HTTP client") +}); + +/// Core worker abstraction that represents a backend service +#[async_trait] +pub trait Worker: Send + Sync + fmt::Debug { + /// Get the worker's URL + fn url(&self) -> &str; + + /// Get the worker's type (Regular, Prefill, or Decode) + fn worker_type(&self) -> WorkerType; + + /// Get the worker's connection mode + fn connection_mode(&self) -> ConnectionMode; + + /// Check if the worker is currently healthy + fn is_healthy(&self) -> bool; + + /// Set the worker's health status + fn set_healthy(&self, healthy: bool); + + /// Perform an async health check on the worker + async fn check_health_async(&self) -> WorkerResult<()>; + + /// Synchronous health check wrapper (for compatibility) + fn check_health(&self) -> WorkerResult<()> { + // Use a small runtime for synchronous contexts + tokio::runtime::Builder::new_current_thread() + .enable_all() + .build() + .map_err(|e| WorkerError::HealthCheckFailed { + url: self.url().to_string(), + reason: format!("Failed to create runtime: {}", e), + })? + .block_on(self.check_health_async()) + } + + /// Get the current load (number of active requests) + fn load(&self) -> usize; + + /// Increment the load counter + fn increment_load(&self); + + /// Decrement the load counter + fn decrement_load(&self); + + /// Reset the load counter to 0 (for sync/recovery) + fn reset_load(&self) { + // Default implementation - does nothing + // Workers that track load should override this + } + + /// Get the number of processed requests + fn processed_requests(&self) -> usize; + + /// Increment the processed requests counter + fn increment_processed(&self); + + /// Get worker-specific metadata + fn metadata(&self) -> &WorkerMetadata; + + /// Get the circuit breaker for this worker + fn circuit_breaker(&self) -> &CircuitBreaker; + + /// Check if the worker is available (healthy + circuit closed/half-open) + fn is_available(&self) -> bool { + self.is_healthy() && self.circuit_breaker().can_execute() + } + + /// Record the outcome of a request to this worker + fn record_outcome(&self, success: bool) { + // Record outcome-level metric with worker label + let outcome_str = if success { "success" } else { "failure" }; + RouterMetrics::record_cb_outcome(self.url(), outcome_str); + + // Record into circuit breaker and infer state change for metrics + let before = self.circuit_breaker().state(); + self.circuit_breaker().record_outcome(success); + let after = self.circuit_breaker().state(); + + if before != after { + let from = match before { + crate::core::CircuitState::Closed => "closed", + crate::core::CircuitState::Open => "open", + crate::core::CircuitState::HalfOpen => "half_open", + }; + let to = match after { + crate::core::CircuitState::Closed => "closed", + crate::core::CircuitState::Open => "open", + crate::core::CircuitState::HalfOpen => "half_open", + }; + RouterMetrics::record_cb_state_transition(self.url(), from, to); + } + + let state_code = match self.circuit_breaker().state() { + crate::core::CircuitState::Closed => 0u8, + crate::core::CircuitState::Open => 1u8, + crate::core::CircuitState::HalfOpen => 2u8, + }; + RouterMetrics::set_cb_state(self.url(), state_code); + } + + /// Get the actual endpoint URL for requests + fn endpoint_url(&self, route: &str) -> String { + format!("{}{}", self.url(), route) + } + + /// Check if this worker can handle a specific request + fn can_handle(&self, _req: &serde_json::Value) -> bool { + true + } + + // === Multi-router support === + + // TODO: - Enhanced Worker Discovery + // The Worker trait should handle async discovery of metadata from the worker itself + // rather than having service discovery or other components query /get_server_info. + // This keeps service discovery decoupled from worker-specific APIs. + // + // Proposed additions: + // - async fn discover_metadata(&mut self) -> Result<(), Error> + // Query /get_server_info and populate metadata labels with model_id, priority, cost, etc. + // - async fn validate_configuration(&self) -> Result<(), Error> + // Ensure worker has required configuration for its mode + // - Make worker creation async to allow metadata discovery during initialization + // + // This way service discovery just calls router.add_worker() and the worker + // handles its own metadata discovery internally. + + /// Get the model ID this worker serves + fn model_id(&self) -> &str { + self.metadata() + .labels + .get("model_id") + .map(|s| s.as_str()) + .unwrap_or("unknown") + } + + /// Get the priority of this worker (higher value = higher priority) + fn priority(&self) -> u32 { + self.metadata() + .labels + .get("priority") + .and_then(|s| s.parse().ok()) + .unwrap_or(50) // Default priority is 50 (mid-range) + } + + /// Get the cost factor of this worker (1.0 = baseline) + fn cost(&self) -> f32 { + self.metadata() + .labels + .get("cost") + .and_then(|s| s.parse().ok()) + .unwrap_or(1.0) + } +} + +/// Connection mode for worker communication +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum ConnectionMode { + /// HTTP/REST connection + Http, +} + +impl fmt::Display for ConnectionMode { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + ConnectionMode::Http => write!(f, "HTTP"), + } + } +} + +/// Worker type classification +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum WorkerType { + /// Regular worker for standard routing + Regular, + /// Prefill worker for PD disaggregated mode + Prefill, + /// Decode worker for PD disaggregated mode + Decode, +} + +impl fmt::Display for WorkerType { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + WorkerType::Regular => write!(f, "Regular"), + WorkerType::Prefill => write!(f, "Prefill"), + WorkerType::Decode => write!(f, "Decode"), + } + } +} + +/// Health check configuration +#[derive(Debug, Clone)] +pub struct HealthConfig { + /// Timeout for health checks in seconds + pub timeout_secs: u64, + /// Interval between health checks in seconds + pub check_interval_secs: u64, + /// Health check endpoint path + pub endpoint: String, + /// Number of consecutive failures before marking unhealthy + pub failure_threshold: u32, + /// Number of consecutive successes before marking healthy + pub success_threshold: u32, +} + +impl Default for HealthConfig { + fn default() -> Self { + Self { + timeout_secs: 5, + check_interval_secs: 30, + endpoint: "/health".to_string(), + failure_threshold: 3, + success_threshold: 2, + } + } +} + +/// Metadata associated with a worker +#[derive(Debug, Clone)] +pub struct WorkerMetadata { + /// Worker URL + pub url: String, + /// Worker type + pub worker_type: WorkerType, + /// Connection mode + pub connection_mode: ConnectionMode, + /// Additional labels/tags + pub labels: std::collections::HashMap, + /// Health check configuration + pub health_config: HealthConfig, +} + +/// Basic worker implementation +#[derive(Clone)] +pub struct BasicWorker { + metadata: WorkerMetadata, + load_counter: Arc, + processed_counter: Arc, + healthy: Arc, + consecutive_failures: Arc, + consecutive_successes: Arc, + circuit_breaker: CircuitBreaker, +} + +impl fmt::Debug for BasicWorker { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("BasicWorker") + .field("metadata", &self.metadata) + .field("healthy", &self.healthy.load(Ordering::Relaxed)) + .field("circuit_breaker", &self.circuit_breaker) + .finish() + } +} + +impl BasicWorker { + pub fn new(url: String, worker_type: WorkerType) -> Self { + Self::with_connection_mode(url, worker_type, ConnectionMode::Http) + } + + pub fn with_connection_mode( + url: String, + worker_type: WorkerType, + connection_mode: ConnectionMode, + ) -> Self { + let metadata = WorkerMetadata { + url: url.clone(), + worker_type, + connection_mode, + labels: std::collections::HashMap::new(), + health_config: HealthConfig::default(), + }; + + Self { + metadata, + load_counter: Arc::new(AtomicUsize::new(0)), + processed_counter: Arc::new(AtomicUsize::new(0)), + healthy: Arc::new(AtomicBool::new(true)), + consecutive_failures: Arc::new(AtomicUsize::new(0)), + consecutive_successes: Arc::new(AtomicUsize::new(0)), + circuit_breaker: CircuitBreaker::new(), + } + } + + pub fn with_labels(mut self, labels: std::collections::HashMap) -> Self { + self.metadata.labels = labels; + self + } + + pub fn with_health_config(mut self, config: HealthConfig) -> Self { + self.metadata.health_config = config; + self + } + + pub fn with_circuit_breaker_config(mut self, config: CircuitBreakerConfig) -> Self { + self.circuit_breaker = CircuitBreaker::with_config(config); + self + } +} + +#[async_trait] +impl Worker for BasicWorker { + fn url(&self) -> &str { + &self.metadata.url + } + + fn worker_type(&self) -> WorkerType { + self.metadata.worker_type.clone() + } + + fn connection_mode(&self) -> ConnectionMode { + self.metadata.connection_mode.clone() + } + + fn is_healthy(&self) -> bool { + self.healthy.load(Ordering::Acquire) + } + + fn set_healthy(&self, healthy: bool) { + self.healthy.store(healthy, Ordering::Release); + RouterMetrics::set_worker_health(self.url(), healthy); + } + + async fn check_health_async(&self) -> WorkerResult<()> { + use std::time::Duration; + + let health_result = match &self.metadata.connection_mode { + ConnectionMode::Http => { + // Perform HTTP health check + let url = self.url(); + let health_url = format!("{}{}", url, self.metadata.health_config.endpoint); + let timeout = Duration::from_secs(self.metadata.health_config.timeout_secs); + + // Use the shared client with a custom timeout for this request + match WORKER_CLIENT.get(&health_url).timeout(timeout).send().await { + Ok(response) => response.status().is_success(), + Err(_) => false, + } + } + }; + + if health_result { + // Health check succeeded + self.consecutive_failures.store(0, Ordering::Release); + let successes = self.consecutive_successes.fetch_add(1, Ordering::AcqRel) + 1; + + // Mark healthy if we've reached the success threshold + if !self.is_healthy() + && successes >= self.metadata.health_config.success_threshold as usize + { + self.set_healthy(true); + self.consecutive_successes.store(0, Ordering::Release); + } + Ok(()) + } else { + // Health check failed + self.consecutive_successes.store(0, Ordering::Release); + let failures = self.consecutive_failures.fetch_add(1, Ordering::AcqRel) + 1; + + // Mark unhealthy if we've reached the failure threshold + if self.is_healthy() + && failures >= self.metadata.health_config.failure_threshold as usize + { + self.set_healthy(false); + self.consecutive_failures.store(0, Ordering::Release); + } + + Err(WorkerError::HealthCheckFailed { + url: self.metadata.url.clone(), + reason: format!("Health check failed (consecutive failures: {})", failures), + }) + } + } + + fn load(&self) -> usize { + self.load_counter.load(Ordering::Relaxed) + } + + fn increment_load(&self) { + self.load_counter.fetch_add(1, Ordering::Relaxed); + } + + fn decrement_load(&self) { + self.load_counter + .fetch_update(Ordering::Relaxed, Ordering::Relaxed, |current| { + current.checked_sub(1) + }) + .ok(); + } + + fn reset_load(&self) { + self.load_counter.store(0, Ordering::Relaxed); + } + + fn processed_requests(&self) -> usize { + self.processed_counter.load(Ordering::Relaxed) + } + + fn increment_processed(&self) { + self.processed_counter.fetch_add(1, Ordering::Relaxed); + } + + fn metadata(&self) -> &WorkerMetadata { + &self.metadata + } + + fn circuit_breaker(&self) -> &CircuitBreaker { + &self.circuit_breaker + } +} + +/// Worker factory for creating workers of different types +pub struct WorkerFactory; + +impl WorkerFactory { + /// Create a regular worker + pub fn create_regular(url: String) -> Box { + Box::new(BasicWorker::new(url, WorkerType::Regular)) + } + + /// Create a regular worker with custom circuit breaker configuration + pub fn create_regular_with_config( + url: String, + circuit_breaker_config: CircuitBreakerConfig, + ) -> Box { + Box::new( + BasicWorker::new(url, WorkerType::Regular) + .with_circuit_breaker_config(circuit_breaker_config), + ) + } + + /// Create a prefill worker + pub fn create_prefill(url: String) -> Box { + Box::new(BasicWorker::new(url, WorkerType::Prefill)) + } + + /// Create a prefill worker with custom circuit breaker configuration + pub fn create_prefill_with_config( + url: String, + circuit_breaker_config: CircuitBreakerConfig, + ) -> Box { + Box::new( + BasicWorker::new(url, WorkerType::Prefill) + .with_circuit_breaker_config(circuit_breaker_config), + ) + } + + /// Create a decode worker + pub fn create_decode(url: String) -> Box { + Box::new(BasicWorker::new(url, WorkerType::Decode)) + } + + /// Create a decode worker with custom circuit breaker configuration + pub fn create_decode_with_config( + url: String, + circuit_breaker_config: CircuitBreakerConfig, + ) -> Box { + Box::new( + BasicWorker::new(url, WorkerType::Decode) + .with_circuit_breaker_config(circuit_breaker_config), + ) + } + + /// Create workers from URLs with automatic type detection + #[allow(clippy::type_complexity)] + pub fn create_from_urls( + regular_urls: Vec, + prefill_urls: Vec, + decode_urls: Vec, + ) -> ( + Vec>, + Vec>, + Vec>, + ) { + let regular_workers: Vec> = + regular_urls.into_iter().map(Self::create_regular).collect(); + + let prefill_workers: Vec> = + prefill_urls.into_iter().map(Self::create_prefill).collect(); + + let decode_workers: Vec> = + decode_urls.into_iter().map(Self::create_decode).collect(); + + (regular_workers, prefill_workers, decode_workers) + } + + /// Create a regular worker with custom labels (for multi-router support) + pub fn create_regular_with_labels( + url: String, + labels: std::collections::HashMap, + circuit_breaker_config: CircuitBreakerConfig, + ) -> Box { + let mut worker = BasicWorker::new(url.clone(), WorkerType::Regular) + .with_circuit_breaker_config(circuit_breaker_config); + + // Add labels to metadata + worker.metadata.labels = labels; + + Box::new(worker) + } + + /// Create a prefill worker with labels + pub fn create_prefill_with_labels( + url: String, + labels: std::collections::HashMap, + circuit_breaker_config: CircuitBreakerConfig, + ) -> Box { + let mut worker = BasicWorker::new(url.clone(), WorkerType::Prefill) + .with_circuit_breaker_config(circuit_breaker_config); + + // Add labels to metadata + worker.metadata.labels = labels; + + Box::new(worker) + } + + /// Create a decode worker with labels + pub fn create_decode_with_labels( + url: String, + labels: std::collections::HashMap, + circuit_breaker_config: CircuitBreakerConfig, + ) -> Box { + let mut worker = BasicWorker::new(url.clone(), WorkerType::Decode) + .with_circuit_breaker_config(circuit_breaker_config); + + // Add labels to metadata + worker.metadata.labels = labels; + + Box::new(worker) + } + + /// Create regular workers (one per URL) + pub fn create_workers(urls: Vec) -> Vec> { + urls.into_iter() + .map(|url| Self::create_regular(url)) + .collect() + } +} + +/// Helper trait for collections of workers +pub trait WorkerCollection { + fn healthy_workers(&self) -> Vec<&dyn Worker>; + fn total_load(&self) -> usize; + fn find_worker(&self, url: &str) -> Option<&dyn Worker>; + fn find_worker_mut(&mut self, url: &str) -> Option<&mut Box>; +} + +impl WorkerCollection for Vec> { + fn healthy_workers(&self) -> Vec<&dyn Worker> { + self.iter() + .filter(|w| w.is_healthy()) + .map(|w| w.as_ref()) + .collect() + } + + fn total_load(&self) -> usize { + self.iter().map(|w| w.load()).sum() + } + + fn find_worker(&self, url: &str) -> Option<&dyn Worker> { + self.iter().find(|w| w.url() == url).map(|w| w.as_ref()) + } + + fn find_worker_mut(&mut self, url: &str) -> Option<&mut Box> { + self.iter_mut().find(|w| w.url() == url) + } +} + +/// Convert a list of worker URLs to worker trait objects +pub fn urls_to_workers(urls: Vec) -> Vec> { + urls.into_iter() + .map(WorkerFactory::create_regular) + .collect() +} + +/// Convert worker trait objects back to URLs +pub fn workers_to_urls(workers: &[Box]) -> Vec { + workers.iter().map(|w| w.url().to_string()).collect() +} + +/// RAII guard for worker load management +pub struct WorkerLoadGuard<'a> { + workers: Vec<&'a dyn Worker>, +} + +impl<'a> WorkerLoadGuard<'a> { + /// Create a new load guard for a single worker + pub fn new(worker: &'a dyn Worker) -> Self { + worker.increment_load(); + Self { + workers: vec![worker], + } + } + + /// Create a new load guard for multiple workers + pub fn new_multi(workers: Vec<&'a dyn Worker>) -> Self { + // Increment load counters for all workers + for worker in &workers { + worker.increment_load(); + } + Self { workers } + } +} + +impl<'a> Drop for WorkerLoadGuard<'a> { + fn drop(&mut self) { + // Decrement load counters for all workers + for worker in &self.workers { + worker.decrement_load(); + } + } +} + +/// Health checker handle with graceful shutdown +pub struct HealthChecker { + handle: tokio::task::JoinHandle<()>, + shutdown: Arc, +} + +impl fmt::Debug for HealthChecker { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + f.debug_struct("HealthChecker") + .field("shutdown", &self.shutdown.load(Ordering::Relaxed)) + .finish() + } +} + +impl HealthChecker { + /// Create a new HealthChecker + pub fn new(handle: tokio::task::JoinHandle<()>, shutdown: Arc) -> Self { + Self { handle, shutdown } + } + + /// Shutdown the health checker gracefully + pub async fn shutdown(self) { + self.shutdown.store(true, Ordering::Release); + let _ = self.handle.await; + } +} + +/// Start an async background health checker for a collection of workers +pub fn start_health_checker( + workers: std::sync::Arc>>>, + check_interval_secs: u64, +) -> HealthChecker { + let shutdown = Arc::new(AtomicBool::new(false)); + let shutdown_clone = shutdown.clone(); + + let handle = tokio::spawn(async move { + let mut interval = + tokio::time::interval(tokio::time::Duration::from_secs(check_interval_secs)); + + // Counter for periodic load reset (every 10 health check cycles) + let mut check_count = 0u64; + const LOAD_RESET_INTERVAL: u64 = 10; + + loop { + interval.tick().await; + + // Check for shutdown signal + if shutdown_clone.load(Ordering::Acquire) { + tracing::debug!("Health checker shutting down"); + break; + } + + check_count += 1; + + // Check health of all workers + let workers_to_check = match workers.read() { + Ok(guard) => guard.clone(), + Err(poisoned) => { + tracing::error!("Worker lock poisoned: {}", poisoned); + continue; + } + }; + + // Periodically reset load counters to prevent drift + // Only do this when we believe all workers should be idle + if check_count.is_multiple_of(LOAD_RESET_INTERVAL) { + let max_load = workers_to_check.iter().map(|w| w.load()).max().unwrap_or(0); + // Only reset if load appears to be very low (likely drift) + if max_load <= 2 { + tracing::debug!( + "Resetting load counters to prevent drift (max_load: {})", + max_load + ); + for worker in &workers_to_check { + worker.reset_load(); + } + } + } + + // Perform health checks concurrently + let health_checks = workers_to_check.iter().map(|worker| { + let worker_url = worker.url().to_string(); + let was_healthy = worker.is_healthy(); + + async move { + match worker.check_health_async().await { + Ok(_) => { + if !was_healthy { + tracing::info!("Worker {} is now healthy", worker_url); + } + } + Err(e) => { + if was_healthy { + tracing::warn!("Worker {} health check failed: {}", worker_url, e); + } else { + // Worker was already unhealthy, log at debug level + tracing::debug!("Worker {} remains unhealthy: {}", worker_url, e); + } + } + } + } + }); + + // Execute all health checks concurrently + futures::future::join_all(health_checks).await; + } + }); + + HealthChecker { handle, shutdown } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::time::Duration; + + // Test WorkerType + #[test] + fn test_worker_type_display() { + assert_eq!(WorkerType::Regular.to_string(), "Regular"); + assert_eq!(WorkerType::Prefill.to_string(), "Prefill"); + assert_eq!(WorkerType::Decode.to_string(), "Decode"); + } + + #[test] + fn test_worker_type_equality() { + assert_eq!(WorkerType::Regular, WorkerType::Regular); + assert_ne!(WorkerType::Regular, WorkerType::Decode); + assert_eq!(WorkerType::Prefill, WorkerType::Prefill); + } + + #[test] + fn test_worker_type_clone() { + let original = WorkerType::Prefill; + let cloned = original.clone(); + assert_eq!(original, cloned); + } + + // Test HealthConfig + #[test] + fn test_health_config_default() { + let config = HealthConfig::default(); + assert_eq!(config.timeout_secs, 5); + assert_eq!(config.check_interval_secs, 30); + assert_eq!(config.endpoint, "/health"); + assert_eq!(config.failure_threshold, 3); + assert_eq!(config.success_threshold, 2); + } + + #[test] + fn test_health_config_custom() { + let config = HealthConfig { + timeout_secs: 10, + check_interval_secs: 60, + endpoint: "/healthz".to_string(), + failure_threshold: 5, + success_threshold: 3, + }; + assert_eq!(config.timeout_secs, 10); + assert_eq!(config.check_interval_secs, 60); + assert_eq!(config.endpoint, "/healthz"); + assert_eq!(config.failure_threshold, 5); + assert_eq!(config.success_threshold, 3); + } + + // Test BasicWorker + #[test] + fn test_basic_worker_creation() { + let worker = BasicWorker::new("http://test:8080".to_string(), WorkerType::Regular); + assert_eq!(worker.url(), "http://test:8080"); + assert_eq!(worker.worker_type(), WorkerType::Regular); + assert!(worker.is_healthy()); + assert_eq!(worker.load(), 0); + assert_eq!(worker.processed_requests(), 0); + } + + #[test] + fn test_worker_with_labels() { + let mut labels = std::collections::HashMap::new(); + labels.insert("env".to_string(), "prod".to_string()); + labels.insert("zone".to_string(), "us-west".to_string()); + + let worker = BasicWorker::new("http://test:8080".to_string(), WorkerType::Regular) + .with_labels(labels.clone()); + + assert_eq!(worker.metadata().labels, labels); + } + + #[test] + fn test_worker_with_health_config() { + let custom_config = HealthConfig { + timeout_secs: 15, + check_interval_secs: 45, + endpoint: "/custom-health".to_string(), + failure_threshold: 4, + success_threshold: 2, + }; + + let worker = BasicWorker::new("http://test:8080".to_string(), WorkerType::Regular) + .with_health_config(custom_config.clone()); + + assert_eq!(worker.metadata().health_config.timeout_secs, 15); + assert_eq!(worker.metadata().health_config.check_interval_secs, 45); + assert_eq!(worker.metadata().health_config.endpoint, "/custom-health"); + } + + // Test Worker trait implementation + #[test] + fn test_worker_url() { + let worker = BasicWorker::new("http://worker1:8080".to_string(), WorkerType::Regular); + assert_eq!(worker.url(), "http://worker1:8080"); + } + + #[test] + fn test_worker_type_getter() { + let regular = BasicWorker::new("http://test:8080".to_string(), WorkerType::Regular); + assert_eq!(regular.worker_type(), WorkerType::Regular); + + let prefill = BasicWorker::new("http://test:8080".to_string(), WorkerType::Prefill); + assert_eq!(prefill.worker_type(), WorkerType::Prefill); + + let decode = BasicWorker::new("http://test:8080".to_string(), WorkerType::Decode); + assert_eq!(decode.worker_type(), WorkerType::Decode); + } + + #[test] + fn test_health_status() { + let worker = BasicWorker::new("http://test:8080".to_string(), WorkerType::Regular); + + // Initial state is healthy + assert!(worker.is_healthy()); + + // Set unhealthy + worker.set_healthy(false); + assert!(!worker.is_healthy()); + + // Set healthy again + worker.set_healthy(true); + assert!(worker.is_healthy()); + } + + #[test] + fn test_load_counter_operations() { + let worker = BasicWorker::new("http://test:8080".to_string(), WorkerType::Regular); + + // Initial load is 0 + assert_eq!(worker.load(), 0); + + // Increment once + worker.increment_load(); + assert_eq!(worker.load(), 1); + + // Increment twice more + worker.increment_load(); + worker.increment_load(); + assert_eq!(worker.load(), 3); + + // Decrement once + worker.decrement_load(); + assert_eq!(worker.load(), 2); + + // Decrement to 0 + worker.decrement_load(); + worker.decrement_load(); + assert_eq!(worker.load(), 0); + + // Decrement below 0 should stay at 0 + worker.decrement_load(); + assert_eq!(worker.load(), 0); + } + + #[test] + fn test_processed_counter() { + let worker = BasicWorker::new("http://test:8080".to_string(), WorkerType::Regular); + + // Initial count is 0 + assert_eq!(worker.processed_requests(), 0); + + // Increment multiple times + for i in 1..=100 { + worker.increment_processed(); + assert_eq!(worker.processed_requests(), i); + } + } + + // Test concurrent operations + #[tokio::test] + async fn test_concurrent_load_increments() { + let worker = Arc::new(BasicWorker::new( + "http://test:8080".to_string(), + WorkerType::Regular, + )); + + let mut handles = vec![]; + + // Spawn 100 tasks incrementing load + for _ in 0..100 { + let worker_clone = Arc::clone(&worker); + let handle = tokio::spawn(async move { + worker_clone.increment_load(); + }); + handles.push(handle); + } + + // Wait for all tasks + for handle in handles { + handle.await.unwrap(); + } + + // Final count should be 100 + assert_eq!(worker.load(), 100); + } + + #[tokio::test] + async fn test_concurrent_load_decrements() { + let worker = Arc::new(BasicWorker::new( + "http://test:8080".to_string(), + WorkerType::Regular, + )); + + // Set initial load to 100 + for _ in 0..100 { + worker.increment_load(); + } + assert_eq!(worker.load(), 100); + + let mut handles = vec![]; + + // Spawn 100 tasks decrementing load + for _ in 0..100 { + let worker_clone = Arc::clone(&worker); + let handle = tokio::spawn(async move { + worker_clone.decrement_load(); + }); + handles.push(handle); + } + + // Wait for all tasks + for handle in handles { + handle.await.unwrap(); + } + + // Final count should be 0 + assert_eq!(worker.load(), 0); + } + + #[tokio::test] + async fn test_concurrent_health_updates() { + let worker = Arc::new(BasicWorker::new( + "http://test:8080".to_string(), + WorkerType::Regular, + )); + + let mut handles = vec![]; + + // Spawn threads randomly setting health status + for i in 0..100 { + let worker_clone = Arc::clone(&worker); + let handle = tokio::spawn(async move { + worker_clone.set_healthy(i % 2 == 0); + tokio::time::sleep(Duration::from_micros(10)).await; + }); + handles.push(handle); + } + + // Wait for all tasks + for handle in handles { + handle.await.unwrap(); + } + } + + // Test WorkerFactory + #[test] + fn test_create_regular_worker() { + let worker = WorkerFactory::create_regular("http://regular:8080".to_string()); + assert_eq!(worker.url(), "http://regular:8080"); + assert_eq!(worker.worker_type(), WorkerType::Regular); + } + + #[test] + fn test_create_prefill_worker() { + let worker = WorkerFactory::create_prefill("http://prefill:8080".to_string()); + assert_eq!(worker.url(), "http://prefill:8080"); + assert_eq!(worker.worker_type(), WorkerType::Prefill); + } + + #[test] + fn test_create_decode_worker() { + let worker = WorkerFactory::create_decode("http://decode:8080".to_string()); + assert_eq!(worker.url(), "http://decode:8080"); + assert_eq!(worker.worker_type(), WorkerType::Decode); + } + + #[test] + fn test_create_from_urls() { + let regular_urls = vec![ + "http://regular1:8080".to_string(), + "http://regular2:8080".to_string(), + ]; + let prefill_urls = vec![ + "http://prefill1:8080".to_string(), + "http://prefill2:8080".to_string(), + ]; + let decode_urls = vec![ + "http://decode1:8080".to_string(), + "http://decode2:8080".to_string(), + ]; + + let (regular, prefill, decode) = + WorkerFactory::create_from_urls(regular_urls, prefill_urls, decode_urls); + + assert_eq!(regular.len(), 2); + assert_eq!(prefill.len(), 2); + assert_eq!(decode.len(), 2); + + assert_eq!(regular[0].url(), "http://regular1:8080"); + assert_eq!(prefill[0].url(), "http://prefill1:8080"); + assert_eq!(decode[0].url(), "http://decode1:8080"); + } + + // Test WorkerCollection trait + #[test] + fn test_healthy_workers_filter() { + let workers: Vec> = vec![ + WorkerFactory::create_regular("http://w1:8080".to_string()), + WorkerFactory::create_regular("http://w2:8080".to_string()), + WorkerFactory::create_regular("http://w3:8080".to_string()), + ]; + + // Set some workers unhealthy + workers[0].set_healthy(false); + workers[2].set_healthy(false); + + let healthy = workers.healthy_workers(); + assert_eq!(healthy.len(), 1); + assert_eq!(healthy[0].url(), "http://w2:8080"); + } + + #[test] + fn test_total_load_calculation() { + let workers: Vec> = vec![ + WorkerFactory::create_regular("http://w1:8080".to_string()), + WorkerFactory::create_regular("http://w2:8080".to_string()), + WorkerFactory::create_regular("http://w3:8080".to_string()), + ]; + + // Set different loads + workers[0].increment_load(); + workers[0].increment_load(); // load = 2 + + workers[1].increment_load(); + workers[1].increment_load(); + workers[1].increment_load(); // load = 3 + + workers[2].increment_load(); // load = 1 + + assert_eq!(workers.total_load(), 6); + } + + #[test] + fn test_find_worker() { + let workers: Vec> = vec![ + WorkerFactory::create_regular("http://w1:8080".to_string()), + WorkerFactory::create_regular("http://w2:8080".to_string()), + WorkerFactory::create_regular("http://w3:8080".to_string()), + ]; + + // Found case + let found = workers.find_worker("http://w2:8080"); + assert!(found.is_some()); + assert_eq!(found.unwrap().url(), "http://w2:8080"); + + // Not found case + let not_found = workers.find_worker("http://w4:8080"); + assert!(not_found.is_none()); + } + + #[test] + fn test_find_worker_mut() { + let mut workers: Vec> = vec![ + WorkerFactory::create_regular("http://w1:8080".to_string()), + WorkerFactory::create_regular("http://w2:8080".to_string()), + ]; + + // Find and modify + if let Some(worker) = workers.find_worker_mut("http://w1:8080") { + worker.set_healthy(false); + } + + // Verify modification + assert!(!workers[0].is_healthy()); + assert!(workers[1].is_healthy()); + } + + // Test WorkerLoadGuard + #[test] + fn test_load_guard_single_worker() { + let worker = BasicWorker::new("http://test:8080".to_string(), WorkerType::Regular); + assert_eq!(worker.load(), 0); + + { + let _guard = WorkerLoadGuard::new(&worker); + assert_eq!(worker.load(), 1); + } + + // Guard dropped, load decremented + assert_eq!(worker.load(), 0); + } + + #[test] + fn test_load_guard_multiple_workers() { + let workers: Vec> = vec![ + WorkerFactory::create_regular("http://w1:8080".to_string()), + WorkerFactory::create_regular("http://w2:8080".to_string()), + WorkerFactory::create_regular("http://w3:8080".to_string()), + ]; + + let worker_refs: Vec<&dyn Worker> = workers.iter().map(|w| w.as_ref()).collect(); + + { + let _guard = WorkerLoadGuard::new_multi(worker_refs); + // All loads incremented + assert_eq!(workers[0].load(), 1); + assert_eq!(workers[1].load(), 1); + assert_eq!(workers[2].load(), 1); + } + + // All loads decremented + assert_eq!(workers[0].load(), 0); + assert_eq!(workers[1].load(), 0); + assert_eq!(workers[2].load(), 0); + } + + #[test] + fn test_load_guard_panic_safety() { + let worker = Arc::new(BasicWorker::new( + "http://test:8080".to_string(), + WorkerType::Regular, + )); + assert_eq!(worker.load(), 0); + + // Clone for use inside catch_unwind + let worker_clone = Arc::clone(&worker); + + // Use AssertUnwindSafe wrapper for the test + // This is safe because we're only testing the load counter behavior. + use std::panic::AssertUnwindSafe; + + // This will panic, but the guard should still clean up + let result = std::panic::catch_unwind(AssertUnwindSafe(|| { + let _guard = WorkerLoadGuard::new(worker_clone.as_ref()); + assert_eq!(worker_clone.load(), 1); + panic!("Test panic"); + })); + + // Verify panic occurred + assert!(result.is_err()); + + // Load should be decremented even after panic + assert_eq!(worker.load(), 0); + } + + // Test helper functions + #[test] + fn test_urls_to_workers() { + let urls = vec!["http://w1:8080".to_string(), "http://w2:8080".to_string()]; + + let workers = urls_to_workers(urls); + assert_eq!(workers.len(), 2); + assert_eq!(workers[0].url(), "http://w1:8080"); + assert_eq!(workers[1].url(), "http://w2:8080"); + assert_eq!(workers[0].worker_type(), WorkerType::Regular); + } + + #[test] + fn test_workers_to_urls() { + let workers: Vec> = vec![ + WorkerFactory::create_regular("http://w1:8080".to_string()), + WorkerFactory::create_regular("http://w2:8080".to_string()), + ]; + + let urls = workers_to_urls(&workers); + assert_eq!(urls, vec!["http://w1:8080", "http://w2:8080"]); + } + + // Test synchronous health check wrapper + #[test] + fn test_check_health_sync_wrapper() { + // We can't easily test the actual HTTP call without mocking, + // but we can verify the sync wrapper works + let worker = BasicWorker::new("http://test:8080".to_string(), WorkerType::Regular); + + // This will fail because there's no server at this URL, + // but it tests that the sync wrapper doesn't panic + let result = worker.check_health(); + assert!(result.is_err()); + } + + // Performance test for load counter + #[test] + fn test_load_counter_performance() { + use std::time::Instant; + + let worker = BasicWorker::new("http://test:8080".to_string(), WorkerType::Regular); + let iterations = 1_000_000; + + let start = Instant::now(); + for _ in 0..iterations { + worker.increment_load(); + } + let duration = start.elapsed(); + + let ops_per_sec = iterations as f64 / duration.as_secs_f64(); + println!("Load counter operations per second: {:.0}", ops_per_sec); + + // Should be well over 1M ops/sec + assert!(ops_per_sec > 1_000_000.0); + } +} diff --git a/lmdeploy/router/src/core/worker_registry.rs b/lmdeploy/router/src/core/worker_registry.rs new file mode 100644 index 0000000000..ae6e6366c8 --- /dev/null +++ b/lmdeploy/router/src/core/worker_registry.rs @@ -0,0 +1,526 @@ +//! Worker Registry for multi-router support +//! +//! Provides centralized registry for workers with model-based indexing + +use crate::core::{ConnectionMode, Worker, WorkerType}; +use dashmap::DashMap; +use std::sync::{Arc, RwLock}; +use uuid::Uuid; + +/// Unique identifier for a worker +#[derive(Debug, Clone, Hash, Eq, PartialEq)] +pub struct WorkerId(String); + +impl WorkerId { + /// Create a new worker ID + pub fn new() -> Self { + Self(Uuid::new_v4().to_string()) + } + + /// Create a worker ID from a string + pub fn from_string(s: String) -> Self { + Self(s) + } + + /// Get the ID as a string + pub fn as_str(&self) -> &str { + &self.0 + } +} + +impl Default for WorkerId { + fn default() -> Self { + Self::new() + } +} + +/// Type alias for the model index to reduce complexity +type ModelIndex = Arc>>>>>; + +/// Worker registry with model-based indexing +#[derive(Debug)] +pub struct WorkerRegistry { + /// All workers indexed by ID + workers: Arc>>, + + /// Workers indexed by model ID (stores WorkerId for reference) + model_workers: Arc>>, + + /// Optimized model index for O(1) lookups (stores Arc directly) + model_index: ModelIndex, + + /// Workers indexed by worker type + type_workers: Arc>>, + + /// Workers indexed by connection mode + connection_workers: Arc>>, + + /// URL to worker ID mapping (for backward compatibility) + url_to_id: Arc>, +} + +impl WorkerRegistry { + /// Create a new worker registry + pub fn new() -> Self { + Self { + workers: Arc::new(DashMap::new()), + model_workers: Arc::new(DashMap::new()), + model_index: Arc::new(DashMap::new()), + type_workers: Arc::new(DashMap::new()), + connection_workers: Arc::new(DashMap::new()), + url_to_id: Arc::new(DashMap::new()), + } + } + + /// Register a new worker + pub fn register(&self, worker: Arc) -> WorkerId { + let worker_id = if let Some(existing_id) = self.url_to_id.get(worker.url()) { + // Worker with this URL already exists, update it + existing_id.clone() + } else { + WorkerId::new() + }; + + // Store worker + self.workers.insert(worker_id.clone(), worker.clone()); + + // Update URL mapping + self.url_to_id + .insert(worker.url().to_string(), worker_id.clone()); + + // Update model index (both ID-based and optimized) + let model_id = worker.model_id().to_string(); + self.model_workers + .entry(model_id.clone()) + .or_default() + .push(worker_id.clone()); + + // Update optimized model index for O(1) lookups + self.model_index + .entry(model_id) + .or_insert_with(|| Arc::new(RwLock::new(Vec::new()))) + .write() + .expect("RwLock for model_index is poisoned") + .push(worker.clone()); + + // Update type index + self.type_workers + .entry(worker.worker_type()) + .or_default() + .push(worker_id.clone()); + + // Update connection mode index + self.connection_workers + .entry(worker.connection_mode()) + .or_default() + .push(worker_id.clone()); + + worker_id + } + + /// Remove a worker by ID + pub fn remove(&self, worker_id: &WorkerId) -> Option> { + if let Some((_, worker)) = self.workers.remove(worker_id) { + // Remove from URL mapping + self.url_to_id.remove(worker.url()); + + // Remove from model index (both ID-based and optimized) + if let Some(mut model_workers) = self.model_workers.get_mut(worker.model_id()) { + model_workers.retain(|id| id != worker_id); + } + + // Remove from optimized model index + if let Some(model_index_entry) = self.model_index.get(worker.model_id()) { + let worker_url = worker.url(); + model_index_entry + .write() + .expect("RwLock for model_index is poisoned") + .retain(|w| w.url() != worker_url); + } + + // Remove from type index + if let Some(mut type_workers) = self.type_workers.get_mut(&worker.worker_type()) { + type_workers.retain(|id| id != worker_id); + } + + // Remove from connection mode index + if let Some(mut conn_workers) = + self.connection_workers.get_mut(&worker.connection_mode()) + { + conn_workers.retain(|id| id != worker_id); + } + + Some(worker) + } else { + None + } + } + + /// Remove a worker by URL + pub fn remove_by_url(&self, url: &str) -> Option> { + if let Some((_, worker_id)) = self.url_to_id.remove(url) { + self.remove(&worker_id) + } else { + None + } + } + + /// Get a worker by ID + pub fn get(&self, worker_id: &WorkerId) -> Option> { + self.workers.get(worker_id).map(|entry| entry.clone()) + } + + /// Get a worker by URL + pub fn get_by_url(&self, url: &str) -> Option> { + self.url_to_id.get(url).and_then(|id| self.get(&id)) + } + + /// Get all workers for a model + pub fn get_by_model(&self, model_id: &str) -> Vec> { + self.model_workers + .get(model_id) + .map(|ids| ids.iter().filter_map(|id| self.get(id)).collect()) + .unwrap_or_default() + } + + /// Get all workers for a model (O(1) optimized version) + /// This method uses the pre-indexed model_index for fast lookups + pub fn get_by_model_fast(&self, model_id: &str) -> Vec> { + self.model_index + .get(model_id) + .map(|workers| { + workers + .read() + .expect("RwLock for model_index is poisoned") + .clone() + }) + .unwrap_or_default() + } + + /// Get all workers by worker type + pub fn get_by_type(&self, worker_type: &WorkerType) -> Vec> { + self.type_workers + .get(worker_type) + .map(|ids| ids.iter().filter_map(|id| self.get(id)).collect()) + .unwrap_or_default() + } + + /// Get all prefill workers + pub fn get_prefill_workers(&self) -> Vec> { + self.workers + .iter() + .filter_map(|entry| { + let worker = entry.value(); + match worker.worker_type() { + WorkerType::Prefill => Some(worker.clone()), + _ => None, + } + }) + .collect() + } + + /// Get all decode workers + pub fn get_decode_workers(&self) -> Vec> { + self.get_by_type(&WorkerType::Decode) + } + + /// Get all workers by connection mode + pub fn get_by_connection(&self, connection_mode: &ConnectionMode) -> Vec> { + self.connection_workers + .get(connection_mode) + .map(|ids| ids.iter().filter_map(|id| self.get(id)).collect()) + .unwrap_or_default() + } + + /// Get all workers + pub fn get_all(&self) -> Vec> { + self.workers + .iter() + .map(|entry| entry.value().clone()) + .collect() + } + + /// Get all workers with their IDs + pub fn get_all_with_ids(&self) -> Vec<(WorkerId, Arc)> { + self.workers + .iter() + .map(|entry| (entry.key().clone(), entry.value().clone())) + .collect() + } + + /// Get all worker URLs + pub fn get_all_urls(&self) -> Vec { + self.workers + .iter() + .map(|entry| entry.value().url().to_string()) + .collect() + } + + /// Get all model IDs with workers + pub fn get_models(&self) -> Vec { + self.model_workers + .iter() + .filter(|entry| !entry.value().is_empty()) + .map(|entry| entry.key().clone()) + .collect() + } + + /// Get workers filtered by multiple criteria + /// + /// This method allows flexible filtering of workers based on: + /// - model_id: Filter by specific model + /// - worker_type: Filter by worker type (Regular, Prefill, Decode) + /// - connection_mode: Filter by connection mode + /// - healthy_only: Only return healthy workers + pub fn get_workers_filtered( + &self, + model_id: Option<&str>, + worker_type: Option, + connection_mode: Option, + healthy_only: bool, + ) -> Vec> { + // Start with the most efficient collection based on filters + // Use model index when possible as it's O(1) lookup + let workers = if let Some(model) = model_id { + self.get_by_model_fast(model) + } else { + self.get_all() + }; + + // Apply remaining filters + workers + .into_iter() + .filter(|w| { + // Check worker_type if specified + if let Some(ref wtype) = worker_type { + if w.worker_type() != *wtype { + return false; + } + } + + // Check connection_mode if specified + if let Some(ref conn) = connection_mode { + if w.connection_mode() != *conn { + return false; + } + } + + // Check health if required + if healthy_only && !w.is_healthy() { + return false; + } + + true + }) + .collect() + } + + /// Get worker statistics + pub fn stats(&self) -> WorkerRegistryStats { + let total_workers = self.workers.len(); + let total_models = self.get_models().len(); + + let mut healthy_count = 0; + let mut total_load = 0; + let mut regular_count = 0; + let mut prefill_count = 0; + let mut decode_count = 0; + + for worker in self.get_all() { + if worker.is_healthy() { + healthy_count += 1; + } + total_load += worker.load(); + + match worker.worker_type() { + WorkerType::Regular => regular_count += 1, + WorkerType::Prefill => prefill_count += 1, + WorkerType::Decode => decode_count += 1, + } + } + + WorkerRegistryStats { + total_workers, + total_models, + healthy_workers: healthy_count, + total_load, + regular_workers: regular_count, + prefill_workers: prefill_count, + decode_workers: decode_count, + } + } + + /// Start a health checker for all workers in the registry + /// This should be called once after the registry is populated with workers + pub fn start_health_checker(&self, check_interval_secs: u64) -> crate::core::HealthChecker { + use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::Arc; + + let shutdown = Arc::new(AtomicBool::new(false)); + let shutdown_clone = shutdown.clone(); + let workers_ref = self.workers.clone(); + + let handle = tokio::spawn(async move { + let mut interval = + tokio::time::interval(tokio::time::Duration::from_secs(check_interval_secs)); + + // Counter for periodic load reset (every 10 health check cycles) + let mut check_count = 0u64; + const LOAD_RESET_INTERVAL: u64 = 10; + + loop { + interval.tick().await; + + // Check for shutdown signal + if shutdown_clone.load(Ordering::Acquire) { + tracing::debug!("Registry health checker shutting down"); + break; + } + + // Get all workers from registry + let workers: Vec> = workers_ref + .iter() + .map(|entry| entry.value().clone()) + .collect(); + + // Perform health checks + for worker in &workers { + let _ = worker.check_health_async().await; // Use async version directly + } + + // Reset loads periodically + check_count += 1; + if check_count.is_multiple_of(LOAD_RESET_INTERVAL) { + tracing::debug!("Resetting worker loads (cycle {})", check_count); + for worker in &workers { + worker.reset_load(); + } + } + } + }); + + crate::core::HealthChecker::new(handle, shutdown) + } +} + +impl Default for WorkerRegistry { + fn default() -> Self { + Self::new() + } +} + +/// Statistics for the worker registry +#[derive(Debug, Clone)] +pub struct WorkerRegistryStats { + pub total_workers: usize, + pub total_models: usize, + pub healthy_workers: usize, + pub total_load: usize, + pub regular_workers: usize, + pub prefill_workers: usize, + pub decode_workers: usize, +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::{CircuitBreakerConfig, WorkerFactory}; + use std::collections::HashMap; + + #[test] + fn test_worker_registry() { + let registry = WorkerRegistry::new(); + + // Create a worker with labels + let mut labels = HashMap::new(); + labels.insert("model_id".to_string(), "llama-3-8b".to_string()); + labels.insert("priority".to_string(), "50".to_string()); + labels.insert("cost".to_string(), "0.8".to_string()); + + let worker = WorkerFactory::create_regular_with_labels( + "http://worker1:8080".to_string(), + labels, + CircuitBreakerConfig::default(), + ); + + // Register worker (WorkerFactory returns Box, convert to Arc) + let worker_id = registry.register(Arc::from(worker)); + + // Verify registration + assert!(registry.get(&worker_id).is_some()); + assert!(registry.get_by_url("http://worker1:8080").is_some()); + assert_eq!(registry.get_by_model("llama-3-8b").len(), 1); + assert_eq!(registry.get_by_type(&WorkerType::Regular).len(), 1); + assert_eq!(registry.get_by_connection(&ConnectionMode::Http).len(), 1); + + // Test stats + let stats = registry.stats(); + assert_eq!(stats.total_workers, 1); + assert_eq!(stats.total_models, 1); + + // Remove worker + registry.remove(&worker_id); + assert!(registry.get(&worker_id).is_none()); + } + + #[test] + fn test_model_index_fast_lookup() { + let registry = WorkerRegistry::new(); + + // Create workers for different models + let mut labels1 = HashMap::new(); + labels1.insert("model_id".to_string(), "llama-3".to_string()); + let worker1 = WorkerFactory::create_regular_with_labels( + "http://worker1:8080".to_string(), + labels1, + CircuitBreakerConfig::default(), + ); + + let mut labels2 = HashMap::new(); + labels2.insert("model_id".to_string(), "llama-3".to_string()); + let worker2 = WorkerFactory::create_regular_with_labels( + "http://worker2:8080".to_string(), + labels2, + CircuitBreakerConfig::default(), + ); + + let mut labels3 = HashMap::new(); + labels3.insert("model_id".to_string(), "gpt-4".to_string()); + let worker3 = WorkerFactory::create_regular_with_labels( + "http://worker3:8080".to_string(), + labels3, + CircuitBreakerConfig::default(), + ); + + // Register workers + registry.register(Arc::from(worker1)); + registry.register(Arc::from(worker2)); + registry.register(Arc::from(worker3)); + + // Test get_by_model_fast for llama-3 + let llama_workers = registry.get_by_model_fast("llama-3"); + assert_eq!(llama_workers.len(), 2); + let urls: Vec = llama_workers.iter().map(|w| w.url().to_string()).collect(); + assert!(urls.contains(&"http://worker1:8080".to_string())); + assert!(urls.contains(&"http://worker2:8080".to_string())); + + // Test get_by_model_fast for gpt-4 + let gpt_workers = registry.get_by_model_fast("gpt-4"); + assert_eq!(gpt_workers.len(), 1); + assert_eq!(gpt_workers[0].url(), "http://worker3:8080"); + + // Test get_by_model_fast for non-existent model + let unknown_workers = registry.get_by_model_fast("unknown-model"); + assert_eq!(unknown_workers.len(), 0); + + // Test that both get_by_model and get_by_model_fast return same results + let llama_workers_slow = registry.get_by_model("llama-3"); + assert_eq!(llama_workers.len(), llama_workers_slow.len()); + + // Test removal updates the model index + registry.remove_by_url("http://worker1:8080"); + let llama_workers_after = registry.get_by_model_fast("llama-3"); + assert_eq!(llama_workers_after.len(), 1); + assert_eq!(llama_workers_after[0].url(), "http://worker2:8080"); + } +} diff --git a/lmdeploy/router/src/data_connector/mod.rs b/lmdeploy/router/src/data_connector/mod.rs new file mode 100644 index 0000000000..1b4bf2073f --- /dev/null +++ b/lmdeploy/router/src/data_connector/mod.rs @@ -0,0 +1,11 @@ +// Data connector module for response storage +pub mod response_memory_store; +pub mod response_noop_store; +pub mod responses; + +pub use response_memory_store::MemoryResponseStorage; +pub use response_noop_store::NoOpResponseStorage; +pub use responses::{ + ResponseChain, ResponseId, ResponseStorage, ResponseStorageError, SharedResponseStorage, + StoredResponse, +}; diff --git a/lmdeploy/router/src/data_connector/response_memory_store.rs b/lmdeploy/router/src/data_connector/response_memory_store.rs new file mode 100644 index 0000000000..b2854c52b9 --- /dev/null +++ b/lmdeploy/router/src/data_connector/response_memory_store.rs @@ -0,0 +1,325 @@ +use async_trait::async_trait; +use parking_lot::RwLock; +use std::collections::HashMap; +use std::sync::Arc; + +use super::responses::{ResponseChain, ResponseId, ResponseStorage, Result, StoredResponse}; + +/// Internal store structure holding both maps together +#[derive(Default)] +struct InnerStore { + /// All stored responses indexed by ID + responses: HashMap, + /// Index of response IDs by user + user_index: HashMap>, +} + +/// In-memory implementation of response storage +pub struct MemoryResponseStorage { + /// Single lock wrapping both maps to prevent deadlocks and ensure atomic updates + store: Arc>, +} + +impl MemoryResponseStorage { + pub fn new() -> Self { + Self { + store: Arc::new(RwLock::new(InnerStore::default())), + } + } + + /// Get statistics about the store + pub fn stats(&self) -> MemoryStoreStats { + let store = self.store.read(); + MemoryStoreStats { + response_count: store.responses.len(), + user_count: store.user_index.len(), + } + } + + /// Clear all data (useful for testing) + pub fn clear(&self) { + let mut store = self.store.write(); + store.responses.clear(); + store.user_index.clear(); + } +} + +impl Default for MemoryResponseStorage { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl ResponseStorage for MemoryResponseStorage { + async fn store_response(&self, mut response: StoredResponse) -> Result { + // Generate ID if not set + if response.id.0.is_empty() { + response.id = ResponseId::new(); + } + + let response_id = response.id.clone(); + + // Single lock acquisition for atomic update + let mut store = self.store.write(); + + // Update user index if user is specified + if let Some(ref user) = response.user { + store + .user_index + .entry(user.clone()) + .or_default() + .push(response_id.clone()); + } + + // Store the response + store.responses.insert(response_id.clone(), response); + + Ok(response_id) + } + + async fn get_response(&self, response_id: &ResponseId) -> Result> { + let store = self.store.read(); + Ok(store.responses.get(response_id).cloned()) + } + + async fn delete_response(&self, response_id: &ResponseId) -> Result<()> { + let mut store = self.store.write(); + + // Remove the response and update user index if needed + if let Some(response) = store.responses.remove(response_id) { + if let Some(ref user) = response.user { + if let Some(user_responses) = store.user_index.get_mut(user) { + user_responses.retain(|id| id != response_id); + } + } + } + + Ok(()) + } + + async fn get_response_chain( + &self, + response_id: &ResponseId, + max_depth: Option, + ) -> Result { + let mut chain = ResponseChain::new(); + let max_depth = max_depth.unwrap_or(100); // Default max depth to prevent infinite loops + + // Collect all response IDs first + let mut response_ids = Vec::new(); + let mut current_id = Some(response_id.clone()); + let mut depth = 0; + + // Single lock acquisition to collect the chain + { + let store = self.store.read(); + while let Some(id) = current_id { + if depth >= max_depth { + break; + } + + if let Some(response) = store.responses.get(&id) { + response_ids.push(id); + current_id = response.previous_response_id.clone(); + depth += 1; + } else { + break; + } + } + } + + // Reverse to get chronological order (oldest first) + response_ids.reverse(); + + // Now collect the actual responses + let store = self.store.read(); + for id in response_ids { + if let Some(response) = store.responses.get(&id) { + chain.add_response(response.clone()); + } + } + + Ok(chain) + } + + async fn list_user_responses( + &self, + user: &str, + limit: Option, + ) -> Result> { + let store = self.store.read(); + + if let Some(user_response_ids) = store.user_index.get(user) { + // Collect responses with their timestamps for sorting + let mut responses_with_time: Vec<_> = user_response_ids + .iter() + .filter_map(|id| store.responses.get(id).map(|r| (r.created_at, id))) + .collect(); + + // Sort by creation time (newest first) + responses_with_time.sort_by_key(|b| std::cmp::Reverse(b.0)); + + // Apply limit and collect the actual responses + let limit = limit.unwrap_or(responses_with_time.len()); + let user_responses: Vec = responses_with_time + .into_iter() + .take(limit) + .filter_map(|(_, id)| store.responses.get(id).cloned()) + .collect(); + + Ok(user_responses) + } else { + Ok(Vec::new()) + } + } + + async fn delete_user_responses(&self, user: &str) -> Result { + let mut store = self.store.write(); + + if let Some(user_response_ids) = store.user_index.remove(user) { + let count = user_response_ids.len(); + for id in user_response_ids { + store.responses.remove(&id); + } + Ok(count) + } else { + Ok(0) + } + } +} + +/// Statistics for the memory store +#[derive(Debug, Clone)] +pub struct MemoryStoreStats { + pub response_count: usize, + pub user_count: usize, +} + +#[cfg(test)] +mod tests { + use super::*; + + #[tokio::test] + async fn test_memory_store_basic() { + let store = MemoryResponseStorage::new(); + + // Store a response + let response = StoredResponse::new("Hello".to_string(), "Hi there!".to_string(), None); + let response_id = store.store_response(response).await.unwrap(); + + // Retrieve it + let retrieved = store.get_response(&response_id).await.unwrap(); + assert!(retrieved.is_some()); + assert_eq!(retrieved.unwrap().input, "Hello"); + + // Delete it + store.delete_response(&response_id).await.unwrap(); + let deleted = store.get_response(&response_id).await.unwrap(); + assert!(deleted.is_none()); + } + + #[tokio::test] + async fn test_response_chain() { + let store = MemoryResponseStorage::new(); + + // Create a chain of responses + let response1 = + StoredResponse::new("First".to_string(), "First response".to_string(), None); + let id1 = store.store_response(response1).await.unwrap(); + + let response2 = StoredResponse::new( + "Second".to_string(), + "Second response".to_string(), + Some(id1.clone()), + ); + let id2 = store.store_response(response2).await.unwrap(); + + let response3 = StoredResponse::new( + "Third".to_string(), + "Third response".to_string(), + Some(id2.clone()), + ); + let id3 = store.store_response(response3).await.unwrap(); + + // Get the chain + let chain = store.get_response_chain(&id3, None).await.unwrap(); + assert_eq!(chain.responses.len(), 3); + assert_eq!(chain.responses[0].input, "First"); + assert_eq!(chain.responses[1].input, "Second"); + assert_eq!(chain.responses[2].input, "Third"); + + // Test with max_depth + let limited_chain = store.get_response_chain(&id3, Some(2)).await.unwrap(); + assert_eq!(limited_chain.responses.len(), 2); + assert_eq!(limited_chain.responses[0].input, "Second"); + assert_eq!(limited_chain.responses[1].input, "Third"); + } + + #[tokio::test] + async fn test_user_responses() { + let store = MemoryResponseStorage::new(); + + // Store responses for different users + let mut response1 = StoredResponse::new( + "User1 message".to_string(), + "Response to user1".to_string(), + None, + ); + response1.user = Some("user1".to_string()); + store.store_response(response1).await.unwrap(); + + let mut response2 = StoredResponse::new( + "Another user1 message".to_string(), + "Another response to user1".to_string(), + None, + ); + response2.user = Some("user1".to_string()); + store.store_response(response2).await.unwrap(); + + let mut response3 = StoredResponse::new( + "User2 message".to_string(), + "Response to user2".to_string(), + None, + ); + response3.user = Some("user2".to_string()); + store.store_response(response3).await.unwrap(); + + // List user1's responses + let user1_responses = store.list_user_responses("user1", None).await.unwrap(); + assert_eq!(user1_responses.len(), 2); + + // List user2's responses + let user2_responses = store.list_user_responses("user2", None).await.unwrap(); + assert_eq!(user2_responses.len(), 1); + + // Delete user1's responses + let deleted_count = store.delete_user_responses("user1").await.unwrap(); + assert_eq!(deleted_count, 2); + + // Verify they're gone + let user1_responses_after = store.list_user_responses("user1", None).await.unwrap(); + assert_eq!(user1_responses_after.len(), 0); + + // User2's responses should still be there + let user2_responses_after = store.list_user_responses("user2", None).await.unwrap(); + assert_eq!(user2_responses_after.len(), 1); + } + + #[tokio::test] + async fn test_memory_store_stats() { + let store = MemoryResponseStorage::new(); + + let mut response1 = StoredResponse::new("Test1".to_string(), "Reply1".to_string(), None); + response1.user = Some("user1".to_string()); + store.store_response(response1).await.unwrap(); + + let mut response2 = StoredResponse::new("Test2".to_string(), "Reply2".to_string(), None); + response2.user = Some("user2".to_string()); + store.store_response(response2).await.unwrap(); + + let stats = store.stats(); + assert_eq!(stats.response_count, 2); + assert_eq!(stats.user_count, 2); + } +} diff --git a/lmdeploy/router/src/data_connector/response_noop_store.rs b/lmdeploy/router/src/data_connector/response_noop_store.rs new file mode 100644 index 0000000000..968a329d72 --- /dev/null +++ b/lmdeploy/router/src/data_connector/response_noop_store.rs @@ -0,0 +1,53 @@ +use async_trait::async_trait; + +use super::responses::{ResponseChain, ResponseId, ResponseStorage, Result, StoredResponse}; + +/// No-op implementation of response storage (does nothing) +pub struct NoOpResponseStorage; + +impl NoOpResponseStorage { + pub fn new() -> Self { + Self + } +} + +impl Default for NoOpResponseStorage { + fn default() -> Self { + Self::new() + } +} + +#[async_trait] +impl ResponseStorage for NoOpResponseStorage { + async fn store_response(&self, response: StoredResponse) -> Result { + Ok(response.id) + } + + async fn get_response(&self, _response_id: &ResponseId) -> Result> { + Ok(None) + } + + async fn delete_response(&self, _response_id: &ResponseId) -> Result<()> { + Ok(()) + } + + async fn get_response_chain( + &self, + _response_id: &ResponseId, + _max_depth: Option, + ) -> Result { + Ok(ResponseChain::new()) + } + + async fn list_user_responses( + &self, + _user: &str, + _limit: Option, + ) -> Result> { + Ok(Vec::new()) + } + + async fn delete_user_responses(&self, _user: &str) -> Result { + Ok(0) + } +} diff --git a/lmdeploy/router/src/data_connector/responses.rs b/lmdeploy/router/src/data_connector/responses.rs new file mode 100644 index 0000000000..49693e9848 --- /dev/null +++ b/lmdeploy/router/src/data_connector/responses.rs @@ -0,0 +1,177 @@ +use async_trait::async_trait; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; +use std::sync::Arc; + +/// Response identifier +#[derive(Debug, Clone, Hash, Eq, PartialEq, Serialize, Deserialize)] +pub struct ResponseId(pub String); + +impl ResponseId { + pub fn new() -> Self { + Self(ulid::Ulid::new().to_string()) + } + + pub fn from_string(s: String) -> Self { + Self(s) + } +} + +impl Default for ResponseId { + fn default() -> Self { + Self::new() + } +} + +/// Stored response data +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct StoredResponse { + /// Unique response ID + pub id: ResponseId, + + /// ID of the previous response in the chain (if any) + pub previous_response_id: Option, + + /// The user input for this response + pub input: String, + + /// System instructions used + pub instructions: Option, + + /// The model's output + pub output: String, + + /// Tool calls made by the model (if any) + pub tool_calls: Vec, + + /// Custom metadata + pub metadata: HashMap, + + /// When this response was created + pub created_at: chrono::DateTime, + + /// User identifier (optional) + pub user: Option, + + /// Model used for generation + pub model: Option, +} + +impl StoredResponse { + pub fn new(input: String, output: String, previous_response_id: Option) -> Self { + Self { + id: ResponseId::new(), + previous_response_id, + input, + instructions: None, + output, + tool_calls: Vec::new(), + metadata: HashMap::new(), + created_at: chrono::Utc::now(), + user: None, + model: None, + } + } +} + +/// Response chain - a sequence of related responses +#[derive(Debug, Clone, Serialize, Deserialize)] +pub struct ResponseChain { + /// The responses in chronological order + pub responses: Vec, + + /// Metadata about the chain + pub metadata: HashMap, +} + +impl Default for ResponseChain { + fn default() -> Self { + Self::new() + } +} + +impl ResponseChain { + pub fn new() -> Self { + Self { + responses: Vec::new(), + metadata: HashMap::new(), + } + } + + /// Get the ID of the most recent response in the chain + pub fn latest_response_id(&self) -> Option<&ResponseId> { + self.responses.last().map(|r| &r.id) + } + + /// Add a response to the chain + pub fn add_response(&mut self, response: StoredResponse) { + self.responses.push(response); + } + + /// Build context from the chain for the next request + pub fn build_context(&self, max_responses: Option) -> Vec<(String, String)> { + let responses = if let Some(max) = max_responses { + let start = self.responses.len().saturating_sub(max); + &self.responses[start..] + } else { + &self.responses[..] + }; + + responses + .iter() + .map(|r| (r.input.clone(), r.output.clone())) + .collect() + } +} + +/// Error type for response storage operations +#[derive(Debug, thiserror::Error)] +pub enum ResponseStorageError { + #[error("Response not found: {0}")] + ResponseNotFound(String), + + #[error("Invalid chain: {0}")] + InvalidChain(String), + + #[error("Storage error: {0}")] + StorageError(String), + + #[error("Serialization error: {0}")] + SerializationError(#[from] serde_json::Error), +} + +pub type Result = std::result::Result; + +/// Trait for response storage +#[async_trait] +pub trait ResponseStorage: Send + Sync { + /// Store a new response + async fn store_response(&self, response: StoredResponse) -> Result; + + /// Get a response by ID + async fn get_response(&self, response_id: &ResponseId) -> Result>; + + /// Delete a response + async fn delete_response(&self, response_id: &ResponseId) -> Result<()>; + + /// Get the chain of responses leading to a given response + /// Returns responses in chronological order (oldest first) + async fn get_response_chain( + &self, + response_id: &ResponseId, + max_depth: Option, + ) -> Result; + + /// List recent responses for a user + async fn list_user_responses( + &self, + user: &str, + limit: Option, + ) -> Result>; + + /// Delete all responses for a user + async fn delete_user_responses(&self, user: &str) -> Result; +} + +/// Type alias for shared storage +pub type SharedResponseStorage = Arc; diff --git a/lmdeploy/router/src/handler.rs b/lmdeploy/router/src/handler.rs new file mode 100644 index 0000000000..309ae0ea40 --- /dev/null +++ b/lmdeploy/router/src/handler.rs @@ -0,0 +1,84 @@ +use crate::ChatCompletionRequest; +use crate::routes::RoutingTreeBuilder; +use crate::routes::interface::RouteHandle; +use crate::types::{AppError, ConfigurationError, Request, Response}; +use bytes::Bytes; +use http_body_util::{BodyExt, Full}; +use hyper::service::Service; +use std::convert::Infallible; +use std::pin::Pin; +use std::sync::Arc; + +#[derive(Clone)] +pub struct RequestHandler { + core: Arc, +} +struct Core { + root: Box, +} + +impl RequestHandler { + pub fn new(config: &std::path::PathBuf) -> Result { + let builder = RoutingTreeBuilder::from_file(config)?; + let root = builder.build_routing_tree()?; + + Ok(Self { + core: Arc::new(Core { root }), + }) + } + + async fn process_request(&self, request: Request) -> Response { + log::debug!("Got request"); + let result = match request.uri().path() { + "/v1/chat/completions" => self.handle_chat_completion(request).await, + nonexistent_route => Err(AppError::NotFound(format!( + "{} does not exist", + nonexistent_route + ))), + }; + + match result { + Ok(response) => { + log::debug!("Success"); + response + } + Err(app_error) => { + log::error!("Error"); + app_error.to_response() + } + } + } + + async fn handle_chat_completion(&self, request: Request) -> Result { + let body = request + .collect() + .await + .map_err(|e| AppError::InternalError(e.to_string()))? + .to_bytes(); + let chat_request: ChatCompletionRequest = serde_json::from_slice(&body).map_err(|e| { + AppError::BadRequest(format!("Invalid JSON in request body: {}", e.to_string())) + })?; + + let response = self.core.root.route(&chat_request).await?; + + Ok(hyper::Response::builder() + .body(Full::new(Bytes::from( + serde_json::to_string(&response).unwrap(), + ))) + .unwrap()) + } +} + +impl Service for RequestHandler { + // TODO Support streaming responses + type Response = Response; + // All errors are converted to proper HTTP responses instead of simply terminating connections + type Error = Infallible; + type Future = Pin> + Send>>; + + fn call(&self, request: Request) -> Self::Future { + let handler = self.clone(); + + Box::pin(async move { Ok(handler.process_request(request).await) }) + } +} diff --git a/lmdeploy/router/src/lib.rs b/lmdeploy/router/src/lib.rs new file mode 100644 index 0000000000..1e1bbfd82f --- /dev/null +++ b/lmdeploy/router/src/lib.rs @@ -0,0 +1,567 @@ +pub mod config; +pub mod logging; +#[cfg(feature = "python")] +use pyo3::prelude::*; +#[cfg(feature = "python")] +use std::collections::HashMap; + +pub mod core; +pub mod data_connector; +pub mod metrics; +pub mod middleware; +pub mod otel_http; +pub mod otel_trace; +pub mod policies; +pub mod protocols; +pub mod routers; +pub mod server; +pub mod service_discovery; +pub mod tokenizer; +pub mod tree; +#[cfg(feature = "python")] +use crate::metrics::PrometheusConfig; + +#[cfg(feature = "python")] +#[pyclass(eq)] +#[derive(Clone, PartialEq, Debug)] +pub enum PolicyType { + Random, + RoundRobin, + CacheAware, + PowerOfTwo, + ConsistentHash, + RendezvousHash, +} + +#[cfg(feature = "python")] +#[pyclass] +#[derive(Debug, Clone, PartialEq)] +struct Router { + host: String, + port: u16, + worker_urls: Vec, + policy: PolicyType, + worker_startup_timeout_secs: u64, + worker_startup_check_interval: u64, + cache_threshold: f32, + balance_abs_threshold: usize, + balance_rel_threshold: f32, + eviction_interval_secs: u64, + max_tree_size: usize, + max_payload_size: usize, + api_key: Option, + api_key_validation_urls: Vec, + log_dir: Option, + + log_level: Option, + service_discovery: bool, + selector: HashMap, + service_discovery_port: u16, + service_discovery_namespace: Option, + prefill_selector: HashMap, + decode_selector: HashMap, + prometheus_port: Option, + prometheus_host: Option, + request_timeout_secs: u64, + request_id_headers: Option>, + prefill_urls: Option>, + decode_urls: Option>, + prefill_policy: Option, + decode_policy: Option, + max_concurrent_requests: usize, + cors_allowed_origins: Vec, + // Retry configuration + retry_max_retries: u32, + retry_initial_backoff_ms: u64, + retry_max_backoff_ms: u64, + retry_backoff_multiplier: f32, + retry_jitter_factor: f32, + disable_retries: bool, + // Circuit breaker configuration + cb_failure_threshold: u32, + cb_success_threshold: u32, + cb_timeout_duration_secs: u64, + cb_window_duration_secs: u64, + disable_circuit_breaker: bool, + // Health check configuration + health_failure_threshold: u32, + health_success_threshold: u32, + health_check_timeout_secs: u64, + health_check_interval_secs: u64, + health_check_endpoint: String, + // IGW (Inference Gateway) configuration + enable_igw: bool, + queue_size: usize, + queue_timeout_secs: u64, + rate_limit_tokens_per_second: Option, + // OpenTelemetry tracing + enable_trace: bool, + otlp_traces_endpoint: Option, + // KV connector for PD disaggregation ("nixl" or "mooncake") + kv_connector: String, + // lmdeploy PD disaggregation + lmdeploy_pd_disaggregation: bool, + lmdeploy_migration_protocol: String, + lmdeploy_rdma_link_type: String, + lmdeploy_with_gdr: bool, + lmdeploy_dummy_prefill: bool, +} + +#[cfg(feature = "python")] +impl Router { + /// Convert PyO3 Router to RouterConfig + pub fn to_router_config(&self) -> config::ConfigResult { + use config::{ + DiscoveryConfig, MetricsConfig, PolicyConfig as ConfigPolicyConfig, RoutingMode, + }; + + // Convert policy helper function + let convert_policy = |policy: &PolicyType| -> ConfigPolicyConfig { + match policy { + PolicyType::Random => ConfigPolicyConfig::Random, + PolicyType::RoundRobin => ConfigPolicyConfig::RoundRobin, + PolicyType::CacheAware => ConfigPolicyConfig::CacheAware { + cache_threshold: self.cache_threshold, + balance_abs_threshold: self.balance_abs_threshold, + balance_rel_threshold: self.balance_rel_threshold, + eviction_interval_secs: self.eviction_interval_secs, + max_tree_size: self.max_tree_size, + }, + PolicyType::PowerOfTwo => ConfigPolicyConfig::PowerOfTwo { + load_check_interval_secs: 5, // Default value + }, + PolicyType::ConsistentHash => ConfigPolicyConfig::ConsistentHash { + virtual_nodes: 160, // Default value + }, + PolicyType::RendezvousHash => ConfigPolicyConfig::RendezvousHash, + } + }; + + // Determine routing mode + let mode = if self.enable_igw { + // IGW mode - routing mode is not used in IGW, but we need to provide a placeholder + RoutingMode::Regular { + worker_urls: vec![], + } + } else if self.lmdeploy_pd_disaggregation { + let migration_protocol = match self + .lmdeploy_migration_protocol + .to_ascii_lowercase() + .as_str() + { + "rdma" => config::LMDeployMigrationProtocol::Rdma, + "nvlink" => config::LMDeployMigrationProtocol::Nvlink, + other => { + return Err(config::ConfigError::ValidationFailed { + reason: format!( + "Invalid lmdeploy_migration_protocol '{}': expected 'rdma' or 'nvlink'", + other + ), + }); + } + }; + let rdma_config = if migration_protocol == config::LMDeployMigrationProtocol::Rdma { + let link_type = match self.lmdeploy_rdma_link_type.to_ascii_lowercase().as_str() { + "ib" => config::LMDeployRdmaLinkType::Ib, + "roce" => config::LMDeployRdmaLinkType::Roce, + other => { + return Err(config::ConfigError::ValidationFailed { + reason: format!( + "Invalid lmdeploy_rdma_link_type '{}': expected 'ib' or 'roce'", + other + ), + }); + } + }; + Some(config::LMDeployRdmaConfig { + with_gdr: self.lmdeploy_with_gdr, + link_type, + }) + } else { + None + }; + RoutingMode::LMDeployPrefillDecode { + prefill_urls: self.prefill_urls.clone().unwrap_or_default(), + decode_urls: self.decode_urls.clone().unwrap_or_default(), + prefill_policy: self.prefill_policy.as_ref().map(convert_policy), + decode_policy: self.decode_policy.as_ref().map(convert_policy), + migration_protocol, + rdma_config, + dummy_prefill: self.lmdeploy_dummy_prefill, + } + } else { + RoutingMode::Regular { + worker_urls: self.worker_urls.clone(), + } + }; + + // Convert main policy + let policy = convert_policy(&self.policy); + + // Service discovery configuration + let discovery = if self.service_discovery { + Some(DiscoveryConfig { + enabled: true, + namespace: self.service_discovery_namespace.clone(), + port: self.service_discovery_port, + check_interval_secs: 60, + selector: self.selector.clone(), + prefill_selector: self.prefill_selector.clone(), + decode_selector: self.decode_selector.clone(), + }) + } else { + None + }; + + // Metrics configuration + let metrics = match (self.prometheus_port, self.prometheus_host.as_ref()) { + (Some(port), Some(host)) => Some(MetricsConfig { + port, + host: host.clone(), + }), + _ => None, + }; + + Ok(config::RouterConfig { + mode, + policy, + host: self.host.clone(), + port: self.port, + connection_mode: config::ConnectionMode::Http, + max_payload_size: self.max_payload_size, + request_timeout_secs: self.request_timeout_secs, + worker_startup_timeout_secs: self.worker_startup_timeout_secs, + worker_startup_check_interval_secs: self.worker_startup_check_interval, + api_key: self.api_key.clone(), + api_key_validation_urls: self.api_key_validation_urls.clone(), + discovery, + metrics, + log_dir: self.log_dir.clone(), + log_level: self.log_level.clone(), + request_id_headers: self.request_id_headers.clone(), + max_concurrent_requests: self.max_concurrent_requests, + queue_size: self.queue_size, + queue_timeout_secs: self.queue_timeout_secs, + rate_limit_tokens_per_second: self.rate_limit_tokens_per_second, + cors_allowed_origins: self.cors_allowed_origins.clone(), + retry: config::RetryConfig { + max_retries: self.retry_max_retries, + initial_backoff_ms: self.retry_initial_backoff_ms, + max_backoff_ms: self.retry_max_backoff_ms, + backoff_multiplier: self.retry_backoff_multiplier, + jitter_factor: self.retry_jitter_factor, + }, + circuit_breaker: config::CircuitBreakerConfig { + failure_threshold: self.cb_failure_threshold, + success_threshold: self.cb_success_threshold, + timeout_duration_secs: self.cb_timeout_duration_secs, + window_duration_secs: self.cb_window_duration_secs, + }, + disable_retries: self.disable_retries, + disable_circuit_breaker: self.disable_circuit_breaker, + health_check: config::HealthCheckConfig { + failure_threshold: self.health_failure_threshold, + success_threshold: self.health_success_threshold, + timeout_secs: self.health_check_timeout_secs, + check_interval_secs: self.health_check_interval_secs, + endpoint: self.health_check_endpoint.clone(), + }, + enable_igw: self.enable_igw, + history_backend: config::HistoryBackend::Memory, + enable_profiling: false, // Profiling disabled in Python binding by default + profile_timeout_secs: 10, // Default profiling timeout + kv_connector: match self.kv_connector.to_ascii_lowercase().as_str() { + "nixl" => config::KvConnector::Nixl, + "mooncake" => config::KvConnector::Mooncake, + "moriio" => config::KvConnector::MoriIO, + other => { + return Err(config::ConfigError::ValidationFailed { + reason: format!( + "Invalid kv_connector '{}': expected 'nixl', 'mooncake', or 'moriio'", + other + ), + }); + } + }, + }) + } +} + +#[cfg(feature = "python")] +#[pymethods] +impl Router { + #[new] + #[pyo3(signature = ( + worker_urls, + policy = PolicyType::RoundRobin, + host = String::from("127.0.0.1"), + port = 3001, + worker_startup_timeout_secs = 600, + worker_startup_check_interval = 30, + cache_threshold = 0.3, + balance_abs_threshold = 64, + balance_rel_threshold = 1.5, + eviction_interval_secs = 120, + max_tree_size = 2usize.pow(26), + max_payload_size = 512 * 1024 * 1024, // 512MB default for large batches + api_key = None, + api_key_validation_urls = vec![], + log_dir = None, + log_level = None, + service_discovery = false, + selector = HashMap::new(), + service_discovery_port = 80, + service_discovery_namespace = None, + prefill_selector = HashMap::new(), + decode_selector = HashMap::new(), + prometheus_port = None, + prometheus_host = None, + request_timeout_secs = 1800, // Add configurable request timeout + request_id_headers = None, // Custom request ID headers + prefill_urls = None, + decode_urls = None, + prefill_policy = None, + decode_policy = None, + max_concurrent_requests = 32768, + cors_allowed_origins = vec![], + // Retry defaults + retry_max_retries = 5, + retry_initial_backoff_ms = 50, + retry_max_backoff_ms = 30_000, + retry_backoff_multiplier = 1.5, + retry_jitter_factor = 0.2, + disable_retries = false, + // Circuit breaker defaults + cb_failure_threshold = 10, + cb_success_threshold = 3, + cb_timeout_duration_secs = 60, + cb_window_duration_secs = 120, + disable_circuit_breaker = false, + // Health check defaults + health_failure_threshold = 3, + health_success_threshold = 2, + health_check_timeout_secs = 5, + health_check_interval_secs = 60, + health_check_endpoint = String::from("/health"), + // IGW defaults + enable_igw = false, + queue_size = 100, + queue_timeout_secs = 60, + rate_limit_tokens_per_second = None, + // Tracing defaults + enable_trace = false, + otlp_traces_endpoint = None, + // KV connector default (PD disaggregation) + kv_connector = String::from("nixl"), + // lmdeploy PD disaggregation defaults + lmdeploy_pd_disaggregation = false, + lmdeploy_migration_protocol = String::from("rdma"), + lmdeploy_rdma_link_type = String::from("roce"), + lmdeploy_with_gdr = true, + lmdeploy_dummy_prefill = false, + ))] + #[allow(clippy::too_many_arguments)] + fn new( + worker_urls: Vec, + policy: PolicyType, + host: String, + port: u16, + worker_startup_timeout_secs: u64, + worker_startup_check_interval: u64, + cache_threshold: f32, + balance_abs_threshold: usize, + balance_rel_threshold: f32, + eviction_interval_secs: u64, + max_tree_size: usize, + max_payload_size: usize, + api_key: Option, + api_key_validation_urls: Vec, + log_dir: Option, + log_level: Option, + service_discovery: bool, + selector: HashMap, + service_discovery_port: u16, + service_discovery_namespace: Option, + prefill_selector: HashMap, + decode_selector: HashMap, + prometheus_port: Option, + prometheus_host: Option, + request_timeout_secs: u64, + request_id_headers: Option>, + prefill_urls: Option>, + decode_urls: Option>, + prefill_policy: Option, + decode_policy: Option, + max_concurrent_requests: usize, + cors_allowed_origins: Vec, + retry_max_retries: u32, + retry_initial_backoff_ms: u64, + retry_max_backoff_ms: u64, + retry_backoff_multiplier: f32, + retry_jitter_factor: f32, + disable_retries: bool, + cb_failure_threshold: u32, + cb_success_threshold: u32, + cb_timeout_duration_secs: u64, + cb_window_duration_secs: u64, + disable_circuit_breaker: bool, + health_failure_threshold: u32, + health_success_threshold: u32, + health_check_timeout_secs: u64, + health_check_interval_secs: u64, + health_check_endpoint: String, + enable_igw: bool, + queue_size: usize, + queue_timeout_secs: u64, + rate_limit_tokens_per_second: Option, + enable_trace: bool, + otlp_traces_endpoint: Option, + kv_connector: String, + lmdeploy_pd_disaggregation: bool, + lmdeploy_migration_protocol: String, + lmdeploy_rdma_link_type: String, + lmdeploy_with_gdr: bool, + lmdeploy_dummy_prefill: bool, + ) -> PyResult { + Ok(Router { + host, + port, + worker_urls, + policy, + worker_startup_timeout_secs, + worker_startup_check_interval, + cache_threshold, + balance_abs_threshold, + balance_rel_threshold, + eviction_interval_secs, + max_tree_size, + max_payload_size, + api_key, + api_key_validation_urls, + log_dir, + log_level, + service_discovery, + selector, + service_discovery_port, + service_discovery_namespace, + prefill_selector, + decode_selector, + prometheus_port, + prometheus_host, + request_timeout_secs, + request_id_headers, + prefill_urls, + decode_urls, + prefill_policy, + decode_policy, + max_concurrent_requests, + cors_allowed_origins, + retry_max_retries, + retry_initial_backoff_ms, + retry_max_backoff_ms, + retry_backoff_multiplier, + retry_jitter_factor, + disable_retries, + cb_failure_threshold, + cb_success_threshold, + cb_timeout_duration_secs, + cb_window_duration_secs, + disable_circuit_breaker, + health_failure_threshold, + health_success_threshold, + health_check_timeout_secs, + health_check_interval_secs, + health_check_endpoint, + enable_igw, + queue_size, + queue_timeout_secs, + rate_limit_tokens_per_second, + enable_trace, + otlp_traces_endpoint, + kv_connector, + lmdeploy_pd_disaggregation, + lmdeploy_migration_protocol, + lmdeploy_rdma_link_type, + lmdeploy_with_gdr, + lmdeploy_dummy_prefill, + }) + } + + fn start(&self) -> PyResult<()> { + // Convert to RouterConfig and validate + let router_config = self.to_router_config().map_err(|e| { + pyo3::exceptions::PyValueError::new_err(format!("Configuration error: {}", e)) + })?; + + // Validate the configuration + router_config.validate().map_err(|e| { + pyo3::exceptions::PyValueError::new_err(format!( + "Configuration validation failed: {}", + e + )) + })?; + + // Create service discovery config if enabled + let service_discovery_config = if self.service_discovery { + Some(service_discovery::ServiceDiscoveryConfig { + enabled: true, + selector: self.selector.clone(), + check_interval: std::time::Duration::from_secs(60), + port: self.service_discovery_port, + namespace: self.service_discovery_namespace.clone(), + pd_mode: self.lmdeploy_pd_disaggregation, + prefill_selector: self.prefill_selector.clone(), + decode_selector: self.decode_selector.clone(), + }) + } else { + None + }; + + // Create Prometheus config if enabled + let prometheus_config = Some(PrometheusConfig { + port: self.prometheus_port.unwrap_or(29000), + host: self + .prometheus_host + .clone() + .unwrap_or_else(|| "127.0.0.1".to_string()), + }); + + // Use tokio runtime instead of actix-web System for better compatibility + let runtime = tokio::runtime::Runtime::new() + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string()))?; + + // Block on the async startup function + runtime.block_on(async move { + server::startup(server::ServerConfig { + host: self.host.clone(), + port: self.port, + router_config, + max_payload_size: self.max_payload_size, + log_dir: self.log_dir.clone(), + log_level: self.log_level.clone(), + service_discovery_config, + prometheus_config, + request_timeout_secs: self.request_timeout_secs, + request_id_headers: self.request_id_headers.clone(), + trace_config: if self.enable_trace { + Some(config::TraceConfig { + otlp_traces_endpoint: self.otlp_traces_endpoint.clone(), + ..Default::default() + }) + } else { + None + }, + }) + .await + .map_err(|e| pyo3::exceptions::PyRuntimeError::new_err(e.to_string())) + }) + } +} + +#[cfg(feature = "python")] +#[pymodule] +fn lmdeploy_router_rs(m: &Bound<'_, PyModule>) -> PyResult<()> { + m.add_class::()?; + m.add_class::()?; + Ok(()) +} diff --git a/lmdeploy/router/src/logger.rs b/lmdeploy/router/src/logger.rs new file mode 100644 index 0000000000..0c7a9a24ac --- /dev/null +++ b/lmdeploy/router/src/logger.rs @@ -0,0 +1,56 @@ +use chrono::{Datelike, Timelike}; +use log::{Level, Metadata, Record}; +use nix; +pub struct Logger; + +impl Logger { + fn thread_id() -> i32 { + // TODO Use gettid() on Linux + nix::unistd::getpid().into() + } + + fn to_char(level: Level) -> char { + match level { + Level::Error => 'E', + Level::Warn => 'W', + Level::Info => 'I', + Level::Debug => 'D', + Level::Trace => 'T', + } + } + + fn filename(filename: &str) -> &str { + filename.rsplit_once('/').unwrap().1 + } +} + +impl log::Log for Logger { + fn enabled(&self, metadata: &Metadata) -> bool { + metadata.level() <= Level::Info + } + + fn log(&self, record: &Record) { + if !self.enabled(record.metadata()) { + return; + } + + let now = chrono::Local::now(); + + println!( + "{}{:02}{:02} {:02}:{:02}:{:02}.{:06} {} {}:{}] {}", + Logger::to_char(record.level()), + now.month(), + now.day(), + now.hour(), + now.minute(), + now.second(), + now.timestamp_subsec_micros(), + Logger::thread_id(), + Logger::filename(record.file().unwrap_or("")), + record.line().unwrap_or(0), + record.args() + ); + } + + fn flush(&self) {} +} diff --git a/lmdeploy/router/src/logging.rs b/lmdeploy/router/src/logging.rs new file mode 100644 index 0000000000..656ee76e9b --- /dev/null +++ b/lmdeploy/router/src/logging.rs @@ -0,0 +1,201 @@ +use std::path::PathBuf; +use tracing::Level; +use tracing_appender::non_blocking::WorkerGuard; +use tracing_appender::rolling::{RollingFileAppender, Rotation}; +use tracing_subscriber::fmt::time::ChronoUtc; +use tracing_subscriber::layer::SubscriberExt; +use tracing_subscriber::util::SubscriberInitExt; +use tracing_subscriber::{EnvFilter, Layer}; + +use crate::config::TraceConfig; + +/// Configuration for the logging system +#[derive(Debug, Clone)] +pub struct LoggingConfig { + /// Log level for the application (default: INFO) + pub level: Level, + /// Whether to use json format for logs (default: false) + pub json_format: bool, + /// Path to store log files. If None, logs will only go to stdout/stderr + pub log_dir: Option, + /// Whether to colorize logs when output is a terminal (default: true) + pub colorize: bool, + /// Log file name to use if log_dir is specified (default: "lmdeploy-router") + pub log_file_name: String, + /// Custom log targets to filter (default: "lmdeploy_router_rs") + pub log_targets: Option>, +} + +impl Default for LoggingConfig { + fn default() -> Self { + Self { + level: Level::INFO, + json_format: false, + log_dir: None, + colorize: true, + log_file_name: "lmdeploy-router".to_string(), + log_targets: Some(vec!["lmdeploy_router_rs".to_string()]), + } + } +} + +/// Guard that keeps the file appender worker thread alive +/// +/// This must be kept in scope for the duration of the program +/// to ensure logs are properly written to files +#[allow(dead_code)] +pub struct LogGuard { + _file_guard: Option, +} + +/// Initialize the logging system with the given configuration +/// +/// # Arguments +/// * `config` - Configuration for the logging system +/// +/// # Returns +/// A LogGuard that must be kept alive for the duration of the program +/// +/// # Panics +/// Will not panic, as initialization errors are handled gracefully +pub fn init_logging(config: LoggingConfig, otel_layer_config: Option) -> LogGuard { + // Note: Do NOT call LogTracer::init() here. The tracing_subscriber + // try_init() method handles LogTracer initialization internally (when the + // "tracing-log" feature is enabled). Calling it beforehand causes + // try_init() to fail because log::set_logger() can only be called once. + + // Convert log level to filter string + let level_filter = match config.level { + Level::TRACE => "trace", + Level::DEBUG => "debug", + Level::INFO => "info", + Level::WARN => "warn", + Level::ERROR => "error", + }; + + // Create env filter + let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| { + // Format: =,=,... + let filter_string = if let Some(targets) = &config.log_targets { + targets + .iter() + .enumerate() + .map(|(i, target)| { + if i > 0 { + format!(",{}={}", target, level_filter) + } else { + format!("{}={}", target, level_filter) + } + }) + .collect::() + } else { + format!("lmdeploy_router_rs={}", level_filter) + }; + + EnvFilter::new(filter_string) + }); + + // When OTel tracing is enabled, ensure the OTel span target passes through + // the global EnvFilter regardless of log level. Without this, --log-level warn + // silently suppresses all info-level OTel spans. + let env_filter = if otel_layer_config.is_some() { + env_filter.add_directive( + "otel_trace=trace" + .parse() + .expect("valid EnvFilter directive"), + ) + } else { + env_filter + }; + + // Setup stdout/stderr layer + let mut layers = Vec::new(); + + // Standard timestamp format: YYYY-MM-DD HH:MM:SS + let time_format = "%Y-%m-%d %H:%M:%S".to_string(); + + // Configure the console stdout layer + let stdout_layer = tracing_subscriber::fmt::layer() + .with_ansi(config.colorize) + .with_file(true) + .with_line_number(true) + .with_timer(ChronoUtc::new(time_format.clone())); + + let stdout_layer = if config.json_format { + stdout_layer.json().flatten_event(true).boxed() + } else { + stdout_layer.boxed() + }; + + layers.push(stdout_layer); + + // Create a file appender if log_dir is specified + let mut file_guard = None; + + if let Some(log_dir) = &config.log_dir { + let file_name = config.log_file_name.clone(); + let log_dir = PathBuf::from(log_dir); + + // Create log directory if it doesn't exist + if !log_dir.exists() { + if let Err(e) = std::fs::create_dir_all(&log_dir) { + eprintln!("Failed to create log directory: {}", e); + return LogGuard { _file_guard: None }; + } + } + + let file_appender = RollingFileAppender::new(Rotation::DAILY, log_dir, file_name); + + let (non_blocking, guard) = tracing_appender::non_blocking(file_appender); + file_guard = Some(guard); + + let file_layer = tracing_subscriber::fmt::layer() + .with_ansi(false) // Never use ANSI colors in log files + .with_file(true) + .with_line_number(true) + .with_timer(ChronoUtc::new(time_format)) + .with_writer(non_blocking); + + let file_layer = if config.json_format { + file_layer.json().flatten_event(true).boxed() + } else { + file_layer.boxed() + }; + + layers.push(file_layer); + } + + let mut prepared_otel = None; + if let Some(trace_config) = otel_layer_config { + match crate::otel_trace::prepare_otel(&trace_config) { + Ok(prepared) => { + layers.push(prepared.layer()); + prepared_otel = Some(prepared); + } + Err(e) => { + eprintln!("Failed to initialize OpenTelemetry: {e}"); + } + } + } + + // Initialize the subscriber with all layers + // Use try_init to handle errors gracefully in case another subscriber is already set + let subscriber_installed = tracing_subscriber::registry() + .with(env_filter) + .with(layers) + .try_init() + .is_ok(); + + // Only activate OTel (set global propagator + provider, flip ENABLED) + // after both the layer and subscriber are live + if let Some(prepared) = prepared_otel { + if subscriber_installed { + crate::otel_trace::activate_otel(prepared); + } + } + + // Return the guard to keep the file appender worker thread alive + LogGuard { + _file_guard: file_guard, + } +} diff --git a/lmdeploy/router/src/main.rs b/lmdeploy/router/src/main.rs new file mode 100644 index 0000000000..cd804f1623 --- /dev/null +++ b/lmdeploy/router/src/main.rs @@ -0,0 +1,671 @@ +use clap::{ArgAction, Parser, ValueEnum}; +use lmdeploy_router_rs::config::{ + CircuitBreakerConfig, ConfigResult, ConnectionMode, DiscoveryConfig, HealthCheckConfig, + HistoryBackend, KvConnector, LMDeployMigrationProtocol, LMDeployRdmaConfig, + LMDeployRdmaLinkType, MetricsConfig, PolicyConfig, RetryConfig, RouterConfig, RoutingMode, + TraceConfig, +}; +use lmdeploy_router_rs::metrics::PrometheusConfig; +use lmdeploy_router_rs::server::{self, ServerConfig}; +use lmdeploy_router_rs::service_discovery::ServiceDiscoveryConfig; +use std::collections::HashMap; + +fn parse_prefill_args() -> Vec { + let args: Vec = std::env::args().collect(); + let mut prefill_entries = Vec::new(); + let mut i = 0; + + while i < args.len() { + if args[i] == "--prefill" && i + 1 < args.len() { + prefill_entries.push(args[i + 1].clone()); + i += 2; // Skip --prefill and URL + } else { + i += 1; + } + } + + prefill_entries +} + +#[derive(Copy, Clone, Debug, Eq, PartialEq, ValueEnum)] +pub enum Backend { + #[value(name = "lmdeploy")] + LMDeploy, + #[value(name = "openai")] + Openai, +} + +impl std::fmt::Display for Backend { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + let s = match self { + Backend::LMDeploy => "lmdeploy", + Backend::Openai => "openai", + }; + write!(f, "{}", s) + } +} + +#[derive(Parser, Debug)] +#[command(name = "lmdeploy-router")] +#[command(version)] +#[command(about = "LMDeploy Router - High-performance request distribution across worker nodes")] +#[command(long_about = r#" +LMDeploy Router - High-performance request distribution across worker nodes + +Usage: +This launcher enables starting a router with individual worker instances. It is useful for +multi-node setups or when you want to start workers and router separately. + +Examples: + # Regular mode + lmdeploy-router --worker-urls http://worker1:8000 http://worker2:8000 + + # LMDeploy dynamic registration (start the router before the API server) + lmdeploy-router --host 0.0.0.0 --port 30000 + lmdeploy serve api_server MODEL --proxy-url http://router:30000 + + # LMDeploy PD dynamic registration (no static URLs required) + lmdeploy-router --host 0.0.0.0 --port 30000 \ + --lmdeploy-pd-disaggregation --lmdeploy-migration-protocol rdma \ + --lmdeploy-rdma-link-type roce + lmdeploy serve api_server MODEL --role Prefill --proxy-url http://router:30000 + lmdeploy serve api_server MODEL --role Decode --proxy-url http://router:30000 + +"#)] +struct CliArgs { + /// Host address to bind the router server + #[arg(long, default_value = "127.0.0.1")] + host: String, + + /// Port number to bind the router server + #[arg(long, default_value_t = 30000)] + port: u16, + + /// Optional worker URLs. May be omitted when LMDeploy servers register via --proxy-url. + #[arg(long, num_args = 0..)] + worker_urls: Vec, + + /// Load balancing policy to use + #[arg(long, default_value = "cache_aware", value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "consistent_hash", "rendezvous_hash"])] + policy: String, + + /// Decode server URL (can be specified multiple times) + #[arg(long, action = ArgAction::Append)] + decode: Vec, + + /// Specific policy for prefill nodes in PD mode + #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "consistent_hash", "rendezvous_hash"])] + prefill_policy: Option, + + /// Specific policy for decode nodes in PD mode + #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "consistent_hash", "rendezvous_hash"])] + decode_policy: Option, + + /// Timeout in seconds for worker startup + #[arg(long, default_value_t = 600)] + worker_startup_timeout_secs: u64, + + /// Interval in seconds between checks for worker startup + #[arg(long, default_value_t = 30)] + worker_startup_check_interval: u64, + + /// Cache threshold (0.0-1.0) for cache-aware routing + #[arg(long, default_value_t = 0.3)] + cache_threshold: f32, + + /// Absolute threshold for load balancing + #[arg(long, default_value_t = 64)] + balance_abs_threshold: usize, + + /// Relative threshold for load balancing + #[arg(long, default_value_t = 1.5)] + balance_rel_threshold: f32, + + /// Interval in seconds between cache eviction operations + #[arg(long, default_value_t = 120)] + eviction_interval: u64, + + /// Maximum size of the approximation tree for cache-aware routing + #[arg(long, default_value_t = 67108864)] // 2^26 + max_tree_size: usize, + + /// Maximum payload size in bytes + #[arg(long, default_value_t = 536870912)] // 512MB + max_payload_size: usize, + + /// API key for worker authorization + #[arg(long)] + api_key: Option, + + /// API key validation URLs (defaults to env file) + #[arg(long, num_args = 0..)] + api_key_validation_urls: Vec, + + /// Backend to route requests to (lmdeploy, openai) + #[arg(long, value_enum, default_value_t = Backend::LMDeploy, alias = "runtime")] + backend: Backend, + + /// Directory to store log files + #[arg(long)] + log_dir: Option, + + /// Set the logging level + #[arg(long, default_value = "info", value_parser = ["debug", "info", "warn", "error"])] + log_level: String, + + /// Enable OpenTelemetry tracing + #[arg( + long, + default_value_t = false, + help_heading = "Tracing (OpenTelemetry)" + )] + enable_trace: bool, + + /// OTLP collector endpoint (format: host:port). If omitted, respects OTEL_EXPORTER_OTLP_ENDPOINT. + #[arg(long, help_heading = "Tracing (OpenTelemetry)")] + otlp_traces_endpoint: Option, + + /// Parent-based sampling ratio for OpenTelemetry traces. + #[arg(long, default_value_t = 1.0, help_heading = "Tracing (OpenTelemetry)")] + otel_sampling_ratio: f64, + + /// Exact HTTP paths to exclude from OpenTelemetry server spans. + #[arg(long, num_args = 0.., help_heading = "Tracing (OpenTelemetry)")] + otel_excluded_paths: Vec, + + /// Enable Kubernetes service discovery + #[arg(long, default_value_t = false)] + service_discovery: bool, + + /// Label selector for Kubernetes service discovery (format: key1=value1 key2=value2) + #[arg(long, num_args = 0..)] + selector: Vec, + + /// Port to use for discovered worker pods + #[arg(long, default_value_t = 80)] + service_discovery_port: u16, + + /// Kubernetes namespace to watch for pods + #[arg(long)] + service_discovery_namespace: Option, + + /// Label selector for prefill server pods in PD mode + #[arg(long, num_args = 0..)] + prefill_selector: Vec, + + /// Label selector for decode server pods in PD mode + #[arg(long, num_args = 0..)] + decode_selector: Vec, + + /// Port to expose Prometheus metrics + #[arg(long, default_value_t = 29000)] + prometheus_port: u16, + + /// Host address to bind the Prometheus metrics server + #[arg(long, default_value = "127.0.0.1")] + prometheus_host: String, + + /// Custom HTTP headers to check for request IDs + #[arg(long, num_args = 0..)] + request_id_headers: Vec, + + /// Request timeout in seconds + #[arg(long, default_value_t = 1800)] + request_timeout_secs: u64, + + /// Maximum number of concurrent requests allowed + #[arg(long, default_value_t = 32768)] + max_concurrent_requests: usize, + + /// CORS allowed origins + #[arg(long, num_args = 0..)] + cors_allowed_origins: Vec, + + // Retry configuration + /// Maximum number of retries + #[arg(long, default_value_t = 5)] + retry_max_retries: u32, + + /// Initial backoff in milliseconds for retries + #[arg(long, default_value_t = 50)] + retry_initial_backoff_ms: u64, + + /// Maximum backoff in milliseconds for retries + #[arg(long, default_value_t = 30000)] + retry_max_backoff_ms: u64, + + /// Backoff multiplier for exponential backoff + #[arg(long, default_value_t = 1.5)] + retry_backoff_multiplier: f32, + + /// Jitter factor for retry backoff + #[arg(long, default_value_t = 0.2)] + retry_jitter_factor: f32, + + /// Disable retries + #[arg(long, default_value_t = false)] + disable_retries: bool, + + // Circuit breaker configuration + /// Number of failures before circuit breaker opens + #[arg(long, default_value_t = 10)] + cb_failure_threshold: u32, + + /// Number of successes before circuit breaker closes + #[arg(long, default_value_t = 3)] + cb_success_threshold: u32, + + /// Timeout duration in seconds for circuit breaker + #[arg(long, default_value_t = 60)] + cb_timeout_duration_secs: u64, + + /// Window duration in seconds for circuit breaker + #[arg(long, default_value_t = 120)] + cb_window_duration_secs: u64, + + /// Disable circuit breaker + #[arg(long, default_value_t = false)] + disable_circuit_breaker: bool, + + // Health check configuration + /// Number of consecutive health check failures before marking worker unhealthy + #[arg(long, default_value_t = 3)] + health_failure_threshold: u32, + + /// Number of consecutive health check successes before marking worker healthy + #[arg(long, default_value_t = 2)] + health_success_threshold: u32, + + /// Timeout in seconds for health check requests + #[arg(long, default_value_t = 5)] + health_check_timeout_secs: u64, + + /// Interval in seconds between runtime health checks + #[arg(long, default_value_t = 60)] + health_check_interval_secs: u64, + + /// Health check endpoint path + #[arg(long, default_value = "/health")] + health_check_endpoint: String, + + // IGW (Inference Gateway) configuration + /// Enable Inference Gateway mode + #[arg(long, default_value_t = false)] + enable_igw: bool, + + /// History backend configuration (memory or none) + #[arg(long, default_value = "memory", value_parser = ["memory", "none"])] + history_backend: String, + + /// Enable profiling calls to backend workers + #[arg(long, default_value_t = false)] + profile: bool, + + /// KV connector type for PD disaggregation (nixl or mooncake) + #[arg(long, value_enum, default_value_t = KvConnector::Nixl)] + kv_connector: KvConnector, + + /// Enable lmdeploy PD disaggregation mode + #[arg(long = "lmdeploy-pd-disaggregation", default_value_t = false)] + lmdeploy_pd_disaggregation: bool, + + /// Migration protocol for lmdeploy PD + #[arg(long = "lmdeploy-migration-protocol", value_enum, default_value_t = LMDeployMigrationProtocol::Rdma)] + lmdeploy_migration_protocol: LMDeployMigrationProtocol, + + /// RDMA link type for lmdeploy PD + #[arg(long = "lmdeploy-rdma-link-type", value_enum, default_value_t = LMDeployRdmaLinkType::Roce)] + lmdeploy_rdma_link_type: LMDeployRdmaLinkType, + + /// Disable GPU Direct RDMA for lmdeploy PD + #[arg(long = "lmdeploy-disable-gdr", default_value_t = false)] + lmdeploy_disable_gdr: bool, + + /// Use dummy prefill for lmdeploy PD + #[arg(long = "lmdeploy-dummy-prefill", default_value_t = false)] + lmdeploy_dummy_prefill: bool, +} + +impl CliArgs { + /// Parse selector strings into HashMap + fn parse_selector(selector_list: &[String]) -> HashMap { + let mut map = HashMap::new(); + for item in selector_list { + if let Some(eq_pos) = item.find('=') { + let key = item[..eq_pos].to_string(); + let value = item[eq_pos + 1..].to_string(); + map.insert(key, value); + } + } + map + } + + /// Convert policy string to PolicyConfig + fn parse_policy(&self, policy_str: &str) -> PolicyConfig { + match policy_str { + "random" => PolicyConfig::Random, + "round_robin" => PolicyConfig::RoundRobin, + "cache_aware" => PolicyConfig::CacheAware { + cache_threshold: self.cache_threshold, + balance_abs_threshold: self.balance_abs_threshold, + balance_rel_threshold: self.balance_rel_threshold, + eviction_interval_secs: self.eviction_interval, + max_tree_size: self.max_tree_size, + }, + "power_of_two" => PolicyConfig::PowerOfTwo { + load_check_interval_secs: 5, // Default value + }, + "consistent_hash" => PolicyConfig::ConsistentHash { + virtual_nodes: 160, // Default value + }, + "rendezvous_hash" => PolicyConfig::RendezvousHash, + _ => PolicyConfig::RoundRobin, // Fallback + } + } + + /// Convert CLI arguments to RouterConfig + fn to_router_config(&self, prefill_urls: Vec) -> ConfigResult { + // Determine routing mode + let mode = if self.enable_igw { + // IGW mode - routing mode is not used in IGW, but we need to provide a placeholder + RoutingMode::Regular { + worker_urls: vec![], + } + } else if matches!(self.backend, Backend::Openai) { + // OpenAI backend mode - use worker_urls as base(s) + RoutingMode::OpenAI { + worker_urls: self.worker_urls.clone(), + } + } else if self.lmdeploy_pd_disaggregation { + // LMDeploy PD disaggregation mode + let decode_str_urls = self.decode.clone(); + + eprintln!("ℹ️ INFO: Using LMDeploy PD disaggregation mode."); + eprintln!(" Prefill URLs: {:?}", prefill_urls); + eprintln!(" Decode URLs: {:?}", decode_str_urls); + eprintln!( + " Migration protocol: {:?}", + self.lmdeploy_migration_protocol + ); + if self.lmdeploy_migration_protocol == LMDeployMigrationProtocol::Rdma { + eprintln!(" RDMA link type: {:?}", self.lmdeploy_rdma_link_type); + eprintln!(" GPU Direct RDMA: {}", !self.lmdeploy_disable_gdr); + } + eprintln!(" Dummy prefill: {}", self.lmdeploy_dummy_prefill); + + let rdma_config = if self.lmdeploy_migration_protocol == LMDeployMigrationProtocol::Rdma + { + Some(LMDeployRdmaConfig { + with_gdr: !self.lmdeploy_disable_gdr, + link_type: self.lmdeploy_rdma_link_type, + }) + } else { + None + }; + + RoutingMode::LMDeployPrefillDecode { + prefill_urls, + decode_urls: decode_str_urls, + prefill_policy: self.prefill_policy.as_ref().map(|p| self.parse_policy(p)), + decode_policy: self.decode_policy.as_ref().map(|p| self.parse_policy(p)), + migration_protocol: self.lmdeploy_migration_protocol, + rdma_config, + dummy_prefill: self.lmdeploy_dummy_prefill, + } + } else { + // Regular mode + RoutingMode::Regular { + worker_urls: self.worker_urls.clone(), + } + }; + + // Main policy + let policy = self.parse_policy(&self.policy); + + // Service discovery configuration + let discovery = if self.service_discovery { + Some(DiscoveryConfig { + enabled: true, + namespace: self.service_discovery_namespace.clone(), + port: self.service_discovery_port, + check_interval_secs: 60, + selector: Self::parse_selector(&self.selector), + prefill_selector: Self::parse_selector(&self.prefill_selector), + decode_selector: Self::parse_selector(&self.decode_selector), + }) + } else { + None + }; + + // Metrics configuration + let metrics = Some(MetricsConfig { + port: self.prometheus_port, + host: self.prometheus_host.clone(), + }); + + let connection_mode = ConnectionMode::Http; + + let api_key_validation_urls = if !self.api_key_validation_urls.is_empty() { + self.api_key_validation_urls.clone() + } else if let Ok(raw_urls) = std::env::var("API_KEY_VALIDATION_URLS") { + raw_urls + .split(',') + .map(|url| url.trim().to_string()) + .filter(|url| !url.is_empty()) + .collect() + } else { + Vec::new() + }; + + // Build RouterConfig + Ok(RouterConfig { + mode, + policy, + connection_mode, + host: self.host.clone(), + port: self.port, + max_payload_size: self.max_payload_size, + request_timeout_secs: self.request_timeout_secs, + worker_startup_timeout_secs: self.worker_startup_timeout_secs, + worker_startup_check_interval_secs: self.worker_startup_check_interval, + api_key: self.api_key.clone(), + api_key_validation_urls, + discovery, + metrics, + log_dir: self.log_dir.clone(), + log_level: Some(self.log_level.clone()), + request_id_headers: if self.request_id_headers.is_empty() { + None + } else { + Some(self.request_id_headers.clone()) + }, + max_concurrent_requests: self.max_concurrent_requests, + queue_size: 100, // Default queue size + queue_timeout_secs: 60, // Default timeout + cors_allowed_origins: self.cors_allowed_origins.clone(), + retry: RetryConfig { + max_retries: self.retry_max_retries, + initial_backoff_ms: self.retry_initial_backoff_ms, + max_backoff_ms: self.retry_max_backoff_ms, + backoff_multiplier: self.retry_backoff_multiplier, + jitter_factor: self.retry_jitter_factor, + }, + circuit_breaker: CircuitBreakerConfig { + failure_threshold: self.cb_failure_threshold, + success_threshold: self.cb_success_threshold, + timeout_duration_secs: self.cb_timeout_duration_secs, + window_duration_secs: self.cb_window_duration_secs, + }, + disable_retries: self.disable_retries, + disable_circuit_breaker: self.disable_circuit_breaker, + health_check: HealthCheckConfig { + failure_threshold: self.health_failure_threshold, + success_threshold: self.health_success_threshold, + timeout_secs: self.health_check_timeout_secs, + check_interval_secs: self.health_check_interval_secs, + endpoint: self.health_check_endpoint.clone(), + }, + enable_igw: self.enable_igw, + rate_limit_tokens_per_second: None, + history_backend: match self.history_backend.as_str() { + "none" => HistoryBackend::None, + _ => HistoryBackend::Memory, + }, + enable_profiling: self.profile, + profile_timeout_secs: 10, // Default profiling timeout + kv_connector: self.kv_connector, + }) + } + + /// Create ServerConfig from CLI args and RouterConfig + fn to_server_config(&self, router_config: RouterConfig) -> ServerConfig { + // Create service discovery config if enabled + let service_discovery_config = if self.service_discovery { + Some(ServiceDiscoveryConfig { + enabled: true, + selector: Self::parse_selector(&self.selector), + check_interval: std::time::Duration::from_secs(60), + port: self.service_discovery_port, + namespace: self.service_discovery_namespace.clone(), + // HTTP service discovery supports both PD router implementations. + pd_mode: self.lmdeploy_pd_disaggregation, + prefill_selector: Self::parse_selector(&self.prefill_selector), + decode_selector: Self::parse_selector(&self.decode_selector), + }) + } else { + None + }; + + // Create Prometheus config + let prometheus_config = Some(PrometheusConfig { + port: self.prometheus_port, + host: self.prometheus_host.clone(), + }); + + ServerConfig { + host: self.host.clone(), + port: self.port, + router_config, + max_payload_size: self.max_payload_size, + log_dir: self.log_dir.clone(), + log_level: Some(self.log_level.clone()), + service_discovery_config, + prometheus_config, + request_timeout_secs: self.request_timeout_secs, + request_id_headers: if self.request_id_headers.is_empty() { + None + } else { + Some(self.request_id_headers.clone()) + }, + trace_config: if self.enable_trace { + Some(TraceConfig { + otlp_traces_endpoint: self.otlp_traces_endpoint.clone(), + sampling_ratio: self.otel_sampling_ratio, + excluded_paths: if self.otel_excluded_paths.is_empty() { + TraceConfig::default_excluded_paths() + } else { + self.otel_excluded_paths.clone() + }, + }) + } else { + None + }, + } + } +} + +fn main() -> Result<(), Box> { + dotenvy::dotenv().ok(); + println!("DEBUG: Main function started"); + + // Parse prefill arguments manually before clap parsing + println!("DEBUG: Parsing prefill arguments"); + let prefill_urls = parse_prefill_args(); + println!("DEBUG: Prefill URLs parsed: {:?}", prefill_urls); + + // Filter out prefill arguments and their values before passing to clap + println!("DEBUG: Filtering CLI arguments"); + let mut filtered_args: Vec = Vec::new(); + let raw_args: Vec = std::env::args().collect(); + println!("DEBUG: Raw args: {:?}", raw_args); + let mut i = 0; + + while i < raw_args.len() { + if raw_args[i] == "--prefill" && i + 1 < raw_args.len() { + // Skip --prefill and its URL + i += 2; + } else { + filtered_args.push(raw_args[i].clone()); + i += 1; + } + } + + // Parse CLI arguments with clap using filtered args + println!("DEBUG: Parsing CLI arguments with clap"); + println!("DEBUG: Filtered args: {:?}", filtered_args); + let cli_args = CliArgs::parse_from(filtered_args); + println!("DEBUG: CLI args parsed successfully"); + // Print startup info + println!("LMDeploy Router starting..."); + println!("Host: {}:{}", cli_args.host, cli_args.port); + let mode_str = if cli_args.enable_igw { + "IGW (Inference Gateway)".to_string() + } else if matches!(cli_args.backend, Backend::Openai) { + "OpenAI Backend".to_string() + } else if cli_args.lmdeploy_pd_disaggregation { + "LMDeploy PD Disaggregated".to_string() + } else { + format!("Regular ({})", cli_args.backend) + }; + println!("Mode: {}", mode_str); + + if !cli_args.enable_igw { + println!("Policy: {}", cli_args.policy); + + if cli_args.lmdeploy_pd_disaggregation { + println!("LMDeploy Prefill nodes: {:?}", prefill_urls); + println!("LMDeploy Decode nodes: {:?}", cli_args.decode); + } + } + + // Convert to RouterConfig + println!("DEBUG: Converting to RouterConfig"); + let router_config = cli_args.to_router_config(prefill_urls)?; + println!("DEBUG: RouterConfig created successfully"); + + // Validate configuration + println!("DEBUG: Validating configuration"); + router_config.validate()?; + println!("DEBUG: Configuration validated successfully"); + + // Create ServerConfig + println!("DEBUG: Creating ServerConfig"); + println!( + "DEBUG: CLI host: {}, port: {}", + cli_args.host, cli_args.port + ); + let server_config = cli_args.to_server_config(router_config); + println!( + "DEBUG: ServerConfig created successfully - host: {}, port: {}", + server_config.host, server_config.port + ); + + // Create a new runtime for the server (like Python binding does) + println!("DEBUG: Creating Tokio runtime"); + let runtime = tokio::runtime::Runtime::new()?; + println!("DEBUG: Tokio runtime created successfully"); + + // Block on the async startup function + println!("DEBUG: Starting server startup function"); + runtime.block_on(async move { + let result = server::startup(server_config).await; + // Shut down OTel while the Tokio runtime is still alive so the + // BatchSpanProcessor can flush its final batch. + if lmdeploy_router_rs::otel_trace::is_otel_enabled() { + lmdeploy_router_rs::otel_trace::shutdown_otel(); + } + result + })?; + + Ok(()) +} diff --git a/lmdeploy/router/src/metrics.rs b/lmdeploy/router/src/metrics.rs new file mode 100644 index 0000000000..d5e0239189 --- /dev/null +++ b/lmdeploy/router/src/metrics.rs @@ -0,0 +1,1058 @@ +use metrics::{counter, describe_counter, describe_gauge, describe_histogram, gauge, histogram}; +use metrics_exporter_prometheus::{Matcher, PrometheusBuilder}; +use std::net::{IpAddr, Ipv4Addr, SocketAddr}; +use std::time::Duration; + +#[derive(Debug, Clone)] +pub struct PrometheusConfig { + pub port: u16, + pub host: String, +} + +impl Default for PrometheusConfig { + fn default() -> Self { + Self { + port: 29000, + host: "0.0.0.0".to_string(), + } + } +} + +pub fn init_metrics() { + // Request metrics + describe_counter!( + "lmdeploy_router_requests_total", + "Total number of requests by route and method" + ); + describe_histogram!( + "lmdeploy_router_request_duration_seconds", + "Request duration in seconds by route" + ); + describe_counter!( + "lmdeploy_router_request_errors_total", + "Total number of request errors by route and error type" + ); + describe_counter!( + "lmdeploy_router_retries_total", + "Total number of request retries by route" + ); + describe_histogram!( + "lmdeploy_router_retry_backoff_duration_seconds", + "Backoff duration in seconds by attempt index" + ); + describe_counter!( + "lmdeploy_router_retries_exhausted_total", + "Total number of requests that exhausted retries by route" + ); + + // Circuit breaker metrics + describe_gauge!( + "lmdeploy_router_cb_state", + "Circuit breaker state per worker (0=closed, 1=open, 2=half_open)" + ); + describe_counter!( + "lmdeploy_router_cb_state_transitions_total", + "Total number of circuit breaker state transitions by worker" + ); + describe_counter!( + "lmdeploy_router_cb_outcomes_total", + "Total number of circuit breaker outcomes by worker and outcome type (success/failure)" + ); + + // Worker metrics + describe_gauge!( + "lmdeploy_router_active_workers", + "Number of currently active workers" + ); + describe_gauge!( + "lmdeploy_router_worker_health", + "Worker health status (1=healthy, 0=unhealthy)" + ); + describe_gauge!("lmdeploy_router_worker_load", "Current load on each worker"); + describe_counter!( + "lmdeploy_router_processed_requests_total", + "Total requests processed by each worker" + ); + + // Policy metrics + describe_counter!( + "lmdeploy_router_policy_decisions_total", + "Total routing policy decisions by policy and worker" + ); + describe_counter!("lmdeploy_router_cache_hits_total", "Total cache hits"); + describe_counter!("lmdeploy_router_cache_misses_total", "Total cache misses"); + describe_gauge!( + "lmdeploy_router_tree_size", + "Current tree size for cache-aware routing" + ); + describe_counter!( + "lmdeploy_router_load_balancing_events_total", + "Total load balancing trigger events" + ); + describe_gauge!("lmdeploy_router_max_load", "Maximum worker load"); + describe_gauge!("lmdeploy_router_min_load", "Minimum worker load"); + + // PD-specific metrics + describe_counter!( + "lmdeploy_router_pd_requests_total", + "Total PD requests by route" + ); + describe_counter!( + "lmdeploy_router_pd_prefill_requests_total", + "Total prefill requests per worker" + ); + describe_counter!( + "lmdeploy_router_pd_decode_requests_total", + "Total decode requests per worker" + ); + describe_counter!( + "lmdeploy_router_pd_errors_total", + "Total PD errors by error type" + ); + describe_counter!( + "lmdeploy_router_pd_prefill_errors_total", + "Total prefill server errors" + ); + describe_counter!( + "lmdeploy_router_pd_decode_errors_total", + "Total decode server errors" + ); + describe_counter!( + "lmdeploy_router_pd_stream_errors_total", + "Total streaming errors per worker" + ); + describe_histogram!( + "lmdeploy_router_pd_request_duration_seconds", + "PD request duration by route" + ); + + // Service discovery metrics + describe_counter!( + "lmdeploy_router_discovery_updates_total", + "Total service discovery update events" + ); + describe_gauge!( + "lmdeploy_router_discovery_workers_added", + "Number of workers added in last discovery update" + ); + describe_gauge!( + "lmdeploy_router_discovery_workers_removed", + "Number of workers removed in last discovery update" + ); + + // Generate request specific metrics + describe_histogram!( + "lmdeploy_router_generate_duration_seconds", + "Generate request duration" + ); + + // Embedding request specific metrics + describe_counter!( + "lmdeploy_router_embeddings_total", + "Total embedding requests" + ); + describe_histogram!( + "lmdeploy_router_embeddings_duration_seconds", + "Embedding request duration" + ); + describe_counter!( + "lmdeploy_router_embeddings_errors_total", + "Embedding request errors" + ); + describe_gauge!( + "lmdeploy_router_embeddings_queue_size", + "Embedding queue size" + ); + + // Running requests gauge for cache-aware policy + describe_gauge!( + "lmdeploy_router_running_requests", + "Number of running requests per worker" + ); + + // Tokenizer metrics + describe_histogram!( + "lmdeploy_router_tokenizer_encode_duration_seconds", + "Time to encode text to tokens" + ); + describe_histogram!( + "lmdeploy_router_tokenizer_decode_duration_seconds", + "Time to decode tokens to text" + ); + describe_histogram!( + "lmdeploy_router_tokenizer_encode_batch_duration_seconds", + "Time to encode a batch of texts" + ); + describe_counter!( + "lmdeploy_router_tokenizer_encode_requests_total", + "Total number of encode requests by tokenizer type" + ); + describe_counter!( + "lmdeploy_router_tokenizer_decode_requests_total", + "Total number of decode requests by tokenizer type" + ); + describe_counter!( + "lmdeploy_router_tokenizer_encode_errors_total", + "Total number of encode errors by error type" + ); + describe_counter!( + "lmdeploy_router_tokenizer_decode_errors_total", + "Total number of decode errors by error type" + ); + describe_histogram!( + "lmdeploy_router_tokenizer_tokens_per_encode", + "Number of tokens produced per encode operation" + ); + describe_histogram!( + "lmdeploy_router_tokenizer_chars_per_encode", + "Number of characters in input text per encode" + ); + describe_histogram!( + "lmdeploy_router_tokenizer_tokens_per_decode", + "Number of tokens decoded per operation" + ); + describe_gauge!( + "lmdeploy_router_tokenizer_vocab_size", + "Vocabulary size of the loaded tokenizer" + ); + + // Stop sequence detection metrics + describe_counter!( + "lmdeploy_router_tokenizer_stop_sequences_detected_total", + "Total stop sequences detected by type" + ); + describe_counter!( + "lmdeploy_router_tokenizer_partial_matches_total", + "Total partial stop sequence matches (jailed text)" + ); + describe_histogram!( + "lmdeploy_router_tokenizer_stop_detection_duration_seconds", + "Time to check for stop sequences per token" + ); + + // Streaming decode metrics + describe_counter!( + "lmdeploy_router_tokenizer_stream_tokens_total", + "Total tokens processed in streaming decode" + ); + describe_counter!( + "lmdeploy_router_tokenizer_stream_incomplete_utf8_total", + "Total incomplete UTF-8 sequences detected" + ); + describe_histogram!( + "lmdeploy_router_tokenizer_stream_step_duration_seconds", + "Time per streaming decode step" + ); + + // Factory metrics + describe_counter!( + "lmdeploy_router_tokenizer_factory_loads_total", + "Total tokenizer loads by file type" + ); + describe_counter!( + "lmdeploy_router_tokenizer_factory_errors_total", + "Total tokenizer loading errors by type" + ); + describe_histogram!( + "lmdeploy_router_tokenizer_factory_load_duration_seconds", + "Time to load and initialize tokenizer" + ); +} + +pub fn start_prometheus(config: PrometheusConfig) { + // Initialize metric descriptions + init_metrics(); + + let duration_matcher = Matcher::Suffix(String::from("duration_seconds")); + let duration_bucket = [ + 0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 15.0, 30.0, 45.0, + 60.0, 90.0, 120.0, 180.0, 240.0, + ]; + + let ip_addr: IpAddr = config + .host + .parse() + .unwrap_or(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0))); + let socket_addr = SocketAddr::new(ip_addr, config.port); + + PrometheusBuilder::new() + .with_http_listener(socket_addr) + .upkeep_timeout(Duration::from_secs(5 * 60)) + .set_buckets_for_metric(duration_matcher, &duration_bucket) + .expect("failed to set duration bucket") + .install() + .expect("failed to install Prometheus metrics exporter"); +} + +pub struct RouterMetrics; + +pub struct TokenizerMetrics; + +impl RouterMetrics { + // Request metrics + pub fn record_request(route: &str) { + counter!("lmdeploy_router_requests_total", + "route" => route.to_string() + ) + .increment(1); + } + + pub fn record_request_duration(route: &str, duration: Duration) { + histogram!("lmdeploy_router_request_duration_seconds", + "route" => route.to_string() + ) + .record(duration.as_secs_f64()); + } + + pub fn record_request_error(route: &str, error_type: &str) { + counter!("lmdeploy_router_request_errors_total", + "route" => route.to_string(), + "error_type" => error_type.to_string() + ) + .increment(1); + } + + pub fn record_retry(route: &str) { + counter!("lmdeploy_router_retries_total", + "route" => route.to_string() + ) + .increment(1); + } + + pub fn record_retry_backoff_duration(duration: Duration, attempt: u32) { + histogram!("lmdeploy_router_retry_backoff_duration_seconds", + "attempt" => attempt.to_string() + ) + .record(duration.as_secs_f64()); + } + + pub fn record_retries_exhausted(route: &str) { + counter!("lmdeploy_router_retries_exhausted_total", + "route" => route.to_string() + ) + .increment(1); + } + + // Worker metrics + pub fn set_active_workers(count: usize) { + gauge!("lmdeploy_router_active_workers").set(count as f64); + } + + pub fn set_worker_health(worker_url: &str, healthy: bool) { + gauge!("lmdeploy_router_worker_health", + "worker" => worker_url.to_string() + ) + .set(if healthy { 1.0 } else { 0.0 }); + } + + pub fn set_worker_load(worker_url: &str, load: usize) { + gauge!("lmdeploy_router_worker_load", + "worker" => worker_url.to_string() + ) + .set(load as f64); + } + + pub fn record_processed_request(worker_url: &str) { + counter!("lmdeploy_router_processed_requests_total", + "worker" => worker_url.to_string() + ) + .increment(1); + } + + // Policy metrics + pub fn record_policy_decision(policy: &str, worker: &str) { + counter!("lmdeploy_router_policy_decisions_total", + "policy" => policy.to_string(), + "worker" => worker.to_string() + ) + .increment(1); + } + + pub fn record_cache_hit() { + counter!("lmdeploy_router_cache_hits_total").increment(1); + } + + pub fn record_cache_miss() { + counter!("lmdeploy_router_cache_misses_total").increment(1); + } + + pub fn set_tree_size(worker: &str, size: usize) { + gauge!("lmdeploy_router_tree_size", + "worker" => worker.to_string() + ) + .set(size as f64); + } + + pub fn record_load_balancing_event() { + counter!("lmdeploy_router_load_balancing_events_total").increment(1); + } + + pub fn set_load_range(max_load: usize, min_load: usize) { + gauge!("lmdeploy_router_max_load").set(max_load as f64); + gauge!("lmdeploy_router_min_load").set(min_load as f64); + } + + // PD-specific metrics + pub fn record_pd_request(route: &str) { + counter!("lmdeploy_router_pd_requests_total", + "route" => route.to_string() + ) + .increment(1); + } + + pub fn record_pd_request_duration(route: &str, duration: Duration) { + histogram!("lmdeploy_router_pd_request_duration_seconds", + "route" => route.to_string() + ) + .record(duration.as_secs_f64()); + } + + pub fn record_pd_prefill_request(worker: &str) { + counter!("lmdeploy_router_pd_prefill_requests_total", + "worker" => worker.to_string() + ) + .increment(1); + } + + pub fn record_pd_decode_request(worker: &str) { + counter!("lmdeploy_router_pd_decode_requests_total", + "worker" => worker.to_string() + ) + .increment(1); + } + + pub fn record_pd_error(error_type: &str) { + counter!("lmdeploy_router_pd_errors_total", + "error_type" => error_type.to_string() + ) + .increment(1); + } + + pub fn record_pd_prefill_error(worker: &str) { + counter!("lmdeploy_router_pd_prefill_errors_total", + "worker" => worker.to_string() + ) + .increment(1); + } + + pub fn record_pd_decode_error(worker: &str) { + counter!("lmdeploy_router_pd_decode_errors_total", + "worker" => worker.to_string() + ) + .increment(1); + } + + pub fn record_pd_stream_error(worker: &str) { + counter!("lmdeploy_router_pd_stream_errors_total", + "worker" => worker.to_string() + ) + .increment(1); + } + + // Service discovery metrics + pub fn record_discovery_update(added: usize, removed: usize) { + counter!("lmdeploy_router_discovery_updates_total").increment(1); + gauge!("lmdeploy_router_discovery_workers_added").set(added as f64); + gauge!("lmdeploy_router_discovery_workers_removed").set(removed as f64); + } + + // Generate request metrics + pub fn record_generate_duration(duration: Duration) { + histogram!("lmdeploy_router_generate_duration_seconds").record(duration.as_secs_f64()); + } + + // Embeddings metrics + pub fn record_embeddings_request() { + counter!("lmdeploy_router_embeddings_total").increment(1); + } + + pub fn record_embeddings_duration(duration: Duration) { + histogram!("lmdeploy_router_embeddings_duration_seconds").record(duration.as_secs_f64()); + } + + pub fn record_embeddings_error(error_type: &str) { + counter!( + "lmdeploy_router_embeddings_errors_total", + "error_type" => error_type.to_string() + ) + .increment(1); + } + + pub fn set_embeddings_queue_size(size: usize) { + gauge!("lmdeploy_router_embeddings_queue_size").set(size as f64); + } + + // Running requests for cache-aware policy + pub fn set_running_requests(worker: &str, count: usize) { + gauge!("lmdeploy_router_running_requests", + "worker" => worker.to_string() + ) + .set(count as f64); + } + + // Circuit breaker metrics + pub fn set_cb_state(worker: &str, state_code: u8) { + gauge!("lmdeploy_router_cb_state", + "worker" => worker.to_string() + ) + .set(state_code as f64); + } + + pub fn record_cb_state_transition(worker: &str, from: &str, to: &str) { + counter!("lmdeploy_router_cb_state_transitions_total", + "worker" => worker.to_string(), + "from" => from.to_string(), + "to" => to.to_string() + ) + .increment(1); + } + + pub fn record_cb_outcome(worker: &str, outcome: &str) { + counter!("lmdeploy_router_cb_outcomes_total", + "worker" => worker.to_string(), + "outcome" => outcome.to_string() + ) + .increment(1); + } +} + +impl TokenizerMetrics { + // Encoding metrics + pub fn record_encode_request(tokenizer_type: &str) { + counter!("lmdeploy_router_tokenizer_encode_requests_total", + "tokenizer_type" => tokenizer_type.to_string() + ) + .increment(1); + } + + pub fn record_encode_duration(duration: Duration) { + histogram!("lmdeploy_router_tokenizer_encode_duration_seconds") + .record(duration.as_secs_f64()); + } + + pub fn record_encode_error(error_type: &str) { + counter!("lmdeploy_router_tokenizer_encode_errors_total", + "error_type" => error_type.to_string() + ) + .increment(1); + } + + pub fn record_tokens_per_encode(token_count: usize) { + histogram!("lmdeploy_router_tokenizer_tokens_per_encode").record(token_count as f64); + } + + pub fn record_chars_per_encode(char_count: usize) { + histogram!("lmdeploy_router_tokenizer_chars_per_encode").record(char_count as f64); + } + + // Decoding metrics + pub fn record_decode_request(tokenizer_type: &str) { + counter!("lmdeploy_router_tokenizer_decode_requests_total", + "tokenizer_type" => tokenizer_type.to_string() + ) + .increment(1); + } + + pub fn record_decode_duration(duration: Duration) { + histogram!("lmdeploy_router_tokenizer_decode_duration_seconds") + .record(duration.as_secs_f64()); + } + + pub fn record_decode_error(error_type: &str) { + counter!("lmdeploy_router_tokenizer_decode_errors_total", + "error_type" => error_type.to_string() + ) + .increment(1); + } + + pub fn record_tokens_per_decode(token_count: usize) { + histogram!("lmdeploy_router_tokenizer_tokens_per_decode").record(token_count as f64); + } + + // Batch encoding metrics + pub fn record_encode_batch_duration(duration: Duration, batch_size: usize) { + histogram!("lmdeploy_router_tokenizer_encode_batch_duration_seconds", + "batch_size" => batch_size.to_string() + ) + .record(duration.as_secs_f64()); + } + + // Stop sequence detection metrics + pub fn record_stop_sequence_detected(stop_type: &str) { + counter!("lmdeploy_router_tokenizer_stop_sequences_detected_total", + "type" => stop_type.to_string() + ) + .increment(1); + } + + pub fn record_partial_match() { + counter!("lmdeploy_router_tokenizer_partial_matches_total").increment(1); + } + + pub fn record_stop_detection_duration(duration: Duration) { + histogram!("lmdeploy_router_tokenizer_stop_detection_duration_seconds") + .record(duration.as_secs_f64()); + } + + // Streaming decode metrics + pub fn record_stream_token() { + counter!("lmdeploy_router_tokenizer_stream_tokens_total").increment(1); + } + + pub fn record_incomplete_utf8() { + counter!("lmdeploy_router_tokenizer_stream_incomplete_utf8_total").increment(1); + } + + pub fn record_stream_step_duration(duration: Duration) { + histogram!("lmdeploy_router_tokenizer_stream_step_duration_seconds") + .record(duration.as_secs_f64()); + } + + // Factory metrics + pub fn record_factory_load(file_type: &str) { + counter!("lmdeploy_router_tokenizer_factory_loads_total", + "file_type" => file_type.to_string() + ) + .increment(1); + } + + pub fn record_factory_error(error_type: &str) { + counter!("lmdeploy_router_tokenizer_factory_errors_total", + "error_type" => error_type.to_string() + ) + .increment(1); + } + + pub fn record_factory_load_duration(duration: Duration) { + histogram!("lmdeploy_router_tokenizer_factory_load_duration_seconds") + .record(duration.as_secs_f64()); + } + + // Vocabulary metrics + pub fn set_vocab_size(tokenizer_type: &str, size: usize) { + gauge!("lmdeploy_router_tokenizer_vocab_size", + "tokenizer_type" => tokenizer_type.to_string() + ) + .set(size as f64); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::net::TcpListener; + + // ============= PrometheusConfig Tests ============= + + #[test] + fn test_prometheus_config_default() { + let config = PrometheusConfig::default(); + assert_eq!(config.port, 29000); + assert_eq!(config.host, "0.0.0.0"); + } + + #[test] + fn test_prometheus_config_custom() { + let config = PrometheusConfig { + port: 8080, + host: "127.0.0.1".to_string(), + }; + assert_eq!(config.port, 8080); + assert_eq!(config.host, "127.0.0.1"); + } + + #[test] + fn test_prometheus_config_clone() { + let config = PrometheusConfig { + port: 9090, + host: "192.168.1.1".to_string(), + }; + let cloned = config.clone(); + assert_eq!(cloned.port, config.port); + assert_eq!(cloned.host, config.host); + } + + // ============= IP Address Parsing Tests ============= + + #[test] + fn test_valid_ipv4_parsing() { + let test_cases = vec!["127.0.0.1", "192.168.1.1", "0.0.0.0"]; + + for ip_str in test_cases { + let config = PrometheusConfig { + port: 29000, + host: ip_str.to_string(), + }; + + let ip_addr: IpAddr = config.host.parse().unwrap(); + assert!(matches!(ip_addr, IpAddr::V4(_))); + } + } + + #[test] + fn test_valid_ipv6_parsing() { + let test_cases = vec!["::1", "2001:db8::1", "::"]; + + for ip_str in test_cases { + let config = PrometheusConfig { + port: 29000, + host: ip_str.to_string(), + }; + + let ip_addr: IpAddr = config.host.parse().unwrap(); + assert!(matches!(ip_addr, IpAddr::V6(_))); + } + } + + #[test] + fn test_invalid_ip_parsing() { + let test_cases = vec!["invalid", "256.256.256.256", "hostname"]; + + for ip_str in test_cases { + let config = PrometheusConfig { + port: 29000, + host: ip_str.to_string(), + }; + + let ip_addr: IpAddr = config + .host + .parse() + .unwrap_or(IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0))); + + // Should fall back to 0.0.0.0 + assert_eq!(ip_addr, IpAddr::V4(Ipv4Addr::new(0, 0, 0, 0))); + } + } + + // ============= Socket Address Creation Tests ============= + + #[test] + fn test_socket_addr_creation() { + let test_cases = vec![("127.0.0.1", 8080), ("0.0.0.0", 29000), ("::1", 9090)]; + + for (host, port) in test_cases { + let config = PrometheusConfig { + port, + host: host.to_string(), + }; + + let ip_addr: IpAddr = config.host.parse().unwrap(); + let socket_addr = SocketAddr::new(ip_addr, config.port); + + assert_eq!(socket_addr.port(), port); + assert_eq!(socket_addr.ip().to_string(), host); + } + } + + #[test] + fn test_socket_addr_with_different_ports() { + let ports = vec![0, 80, 8080, 65535]; + + for port in ports { + let config = PrometheusConfig { + port, + host: "127.0.0.1".to_string(), + }; + + let ip_addr: IpAddr = config.host.parse().unwrap(); + let socket_addr = SocketAddr::new(ip_addr, config.port); + + assert_eq!(socket_addr.port(), port); + } + } + + // ============= Duration Bucket Tests ============= + + #[test] + fn test_duration_bucket_coverage() { + let test_cases: [(f64, &str); 7] = [ + (0.0005, "sub-millisecond"), + (0.005, "5ms"), + (0.05, "50ms"), + (1.0, "1s"), + (10.0, "10s"), + (60.0, "1m"), + (240.0, "4m"), + ]; + + let buckets: [f64; 20] = [ + 0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 15.0, 30.0, 45.0, + 60.0, 90.0, 120.0, 180.0, 240.0, + ]; + + for (duration, label) in test_cases { + let bucket_found = buckets + .iter() + .any(|&b| (b - duration).abs() < 0.0001 || b > duration); + assert!(bucket_found, "No bucket found for {} ({})", duration, label); + } + } + + // ============= Matcher Configuration Tests ============= + + #[test] + fn test_duration_suffix_matcher() { + let matcher = Matcher::Suffix(String::from("duration_seconds")); + + // Test matching behavior + let _matching_metrics = [ + "request_duration_seconds", + "response_duration_seconds", + "lmdeploy_router_request_duration_seconds", + ]; + + let _non_matching_metrics = ["duration_total", "duration_seconds_total", "other_metric"]; + + // Note: We can't directly test Matcher matching without the internals, + // but we can verify the matcher is created correctly + match matcher { + Matcher::Suffix(suffix) => assert_eq!(suffix, "duration_seconds"), + _ => panic!("Expected Suffix matcher"), + } + } + + // ============= Builder Configuration Tests ============= + + #[test] + fn test_prometheus_builder_configuration() { + // This test verifies the builder configuration without actually starting Prometheus + let _config = PrometheusConfig::default(); + + let duration_matcher = Matcher::Suffix(String::from("duration_seconds")); + let duration_bucket = [ + 0.001, 0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 15.0, 30.0, 45.0, + 60.0, 90.0, 120.0, 180.0, 240.0, + ]; + + // Verify bucket configuration + assert_eq!(duration_bucket.len(), 20); + + // Verify matcher is suffix type + match duration_matcher { + Matcher::Suffix(s) => assert_eq!(s, "duration_seconds"), + _ => panic!("Expected Suffix matcher"), + } + } + + // ============= Upkeep Timeout Tests ============= + + #[test] + fn test_upkeep_timeout_duration() { + let timeout = Duration::from_secs(5 * 60); + assert_eq!(timeout.as_secs(), 300); + } + + // ============= Custom Bucket Tests ============= + + #[test] + fn test_custom_buckets_for_different_metrics() { + // Test that we can create different bucket configurations + let request_buckets = [0.001, 0.01, 0.1, 1.0, 10.0]; + let generate_buckets = [0.1, 0.5, 1.0, 5.0, 30.0, 60.0]; + + assert_eq!(request_buckets.len(), 5); + assert_eq!(generate_buckets.len(), 6); + + // Verify each set is sorted + for i in 1..request_buckets.len() { + assert!(request_buckets[i] > request_buckets[i - 1]); + } + + for i in 1..generate_buckets.len() { + assert!(generate_buckets[i] > generate_buckets[i - 1]); + } + } + + // ============= RouterMetrics Tests ============= + + #[test] + fn test_metrics_static_methods() { + // Test that all static methods can be called without panic + RouterMetrics::record_request("/generate"); + RouterMetrics::record_request_duration("/generate", Duration::from_millis(100)); + RouterMetrics::record_request_error("/generate", "timeout"); + RouterMetrics::record_retry("/generate"); + + RouterMetrics::set_active_workers(5); + RouterMetrics::set_worker_health("http://worker1", true); + RouterMetrics::set_worker_load("http://worker1", 10); + RouterMetrics::record_processed_request("http://worker1"); + + RouterMetrics::record_policy_decision("random", "http://worker1"); + RouterMetrics::record_cache_hit(); + RouterMetrics::record_cache_miss(); + RouterMetrics::set_tree_size("http://worker1", 1000); + RouterMetrics::record_load_balancing_event(); + RouterMetrics::set_load_range(20, 5); + + RouterMetrics::record_pd_request("/v1/chat/completions"); + RouterMetrics::record_pd_request_duration("/v1/chat/completions", Duration::from_secs(1)); + RouterMetrics::record_pd_prefill_request("http://prefill1"); + RouterMetrics::record_pd_decode_request("http://decode1"); + RouterMetrics::record_pd_error("invalid_request"); + RouterMetrics::record_pd_prefill_error("http://prefill1"); + RouterMetrics::record_pd_decode_error("http://decode1"); + RouterMetrics::record_pd_stream_error("http://decode1"); + + RouterMetrics::record_discovery_update(3, 1); + RouterMetrics::record_generate_duration(Duration::from_secs(2)); + RouterMetrics::set_running_requests("http://worker1", 15); + } + + #[test] + fn test_tokenizer_metrics_static_methods() { + // Test that all tokenizer metric methods can be called without panic + + // Encoding metrics + TokenizerMetrics::record_encode_request("huggingface"); + TokenizerMetrics::record_encode_duration(Duration::from_millis(10)); + TokenizerMetrics::record_encode_error("invalid_input"); + TokenizerMetrics::record_tokens_per_encode(100); + TokenizerMetrics::record_chars_per_encode(500); + + // Decoding metrics + TokenizerMetrics::record_decode_request("huggingface"); + TokenizerMetrics::record_decode_duration(Duration::from_millis(5)); + TokenizerMetrics::record_decode_error("invalid_tokens"); + TokenizerMetrics::record_tokens_per_decode(50); + + // Batch encoding + TokenizerMetrics::record_encode_batch_duration(Duration::from_millis(100), 10); + + // Stop sequence detection + TokenizerMetrics::record_stop_sequence_detected("token"); + TokenizerMetrics::record_stop_sequence_detected("string"); + TokenizerMetrics::record_partial_match(); + TokenizerMetrics::record_stop_detection_duration(Duration::from_micros(100)); + + // Streaming decode + TokenizerMetrics::record_stream_token(); + TokenizerMetrics::record_incomplete_utf8(); + TokenizerMetrics::record_stream_step_duration(Duration::from_micros(50)); + + // Factory metrics + TokenizerMetrics::record_factory_load("json"); + TokenizerMetrics::record_factory_error("unsupported_format"); + TokenizerMetrics::record_factory_load_duration(Duration::from_millis(200)); + + // Vocabulary metrics + TokenizerMetrics::set_vocab_size("huggingface", 50000); + } + + // ============= Port Availability Tests ============= + + #[test] + fn test_port_already_in_use() { + // Skip this test if we can't bind to the port + let port = 29123; // Use a different port to avoid conflicts + + if let Ok(_listener) = TcpListener::bind(("127.0.0.1", port)) { + // Port is available, we can test + let config = PrometheusConfig { + port, + host: "127.0.0.1".to_string(), + }; + + // Just verify config is created correctly + assert_eq!(config.port, port); + } + } + + // ============= Integration Test Helpers ============= + + #[test] + fn test_metrics_endpoint_accessibility() { + // This would be an integration test in practice + // Here we just verify the configuration + let config = PrometheusConfig { + port: 29000, + host: "127.0.0.1".to_string(), + }; + + let ip_addr: IpAddr = config.host.parse().unwrap(); + let socket_addr = SocketAddr::new(ip_addr, config.port); + + assert_eq!(socket_addr.to_string(), "127.0.0.1:29000"); + } + + #[test] + fn test_concurrent_metric_updates() { + // Test that metric updates can be called concurrently + use std::sync::atomic::{AtomicBool, Ordering}; + use std::sync::Arc; + use std::thread; + + let done = Arc::new(AtomicBool::new(false)); + let mut handles = vec![]; + + for i in 0..3 { + let done_clone = done.clone(); + let handle = thread::spawn(move || { + let worker = format!("http://worker{}", i); + while !done_clone.load(Ordering::Relaxed) { + RouterMetrics::set_worker_load(&worker, i * 10); + RouterMetrics::record_processed_request(&worker); + thread::sleep(Duration::from_millis(1)); + } + }); + handles.push(handle); + } + + // Let threads run briefly + thread::sleep(Duration::from_millis(10)); + done.store(true, Ordering::Relaxed); + + // Wait for all threads + for handle in handles { + handle.join().unwrap(); + } + } + + // ============= Edge Cases Tests ============= + + #[test] + fn test_empty_string_metrics() { + // Test that empty strings don't cause issues + RouterMetrics::record_request(""); + RouterMetrics::set_worker_health("", true); + RouterMetrics::record_policy_decision("", ""); + } + + #[test] + fn test_very_long_metric_labels() { + let long_label = "a".repeat(1000); + + RouterMetrics::record_request(&long_label); + RouterMetrics::set_worker_health(&long_label, false); + } + + #[test] + fn test_special_characters_in_labels() { + let special_labels = [ + "test/with/slashes", + "test-with-dashes", + "test_with_underscores", + "test.with.dots", + "test:with:colons", + ]; + + for label in special_labels { + RouterMetrics::record_request(label); + RouterMetrics::set_worker_health(label, true); + } + } + + #[test] + fn test_extreme_metric_values() { + // Test extreme values + RouterMetrics::set_active_workers(0); + RouterMetrics::set_active_workers(usize::MAX); + + RouterMetrics::set_worker_load("worker", 0); + RouterMetrics::set_worker_load("worker", usize::MAX); + + RouterMetrics::record_request_duration("route", Duration::from_nanos(1)); + // 24 hours + RouterMetrics::record_request_duration("route", Duration::from_secs(86400)); + } +} diff --git a/lmdeploy/router/src/middleware.rs b/lmdeploy/router/src/middleware.rs new file mode 100644 index 0000000000..b654df612e --- /dev/null +++ b/lmdeploy/router/src/middleware.rs @@ -0,0 +1,587 @@ +use axum::{ + extract::Request, extract::State, http::HeaderValue, http::StatusCode, middleware::Next, + response::IntoResponse, response::Response, +}; +use rand::Rng; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; +use std::time::Duration; +use std::time::Instant; +use tokio::sync::{mpsc, oneshot}; +use tower::{Layer, Service}; +use tower_http::classify::ServerErrorsFailureClass; +use tower_http::trace::{MakeSpan, OnFailure, OnRequest, OnResponse, TraceLayer}; +use tracing::{debug, error, field::Empty, info, info_span, warn, Span}; +use tracing_opentelemetry::OpenTelemetrySpanExt; + +pub use crate::core::token_bucket::TokenBucket; + +use crate::metrics::RouterMetrics; +use crate::server::AppState; + +/// Generate OpenAI-compatible request ID based on endpoint +fn generate_request_id(path: &str) -> String { + let prefix = if path.contains("/chat/completions") { + "chatcmpl-" + } else if path.contains("/completions") { + "cmpl-" + } else if path.contains("/generate") { + "gnt-" + } else { + "req-" + }; + + // Generate a random string similar to OpenAI's format + let chars = "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"; + let mut rng = rand::rng(); + let random_part: String = (0..24) + .map(|_| { + let idx = rng.random_range(0..chars.len()); + chars.chars().nth(idx).unwrap() + }) + .collect(); + + format!("{}{}", prefix, random_part) +} + +/// Extension type for storing request ID +#[derive(Clone, Debug)] +pub struct RequestId(pub String); + +/// Tower Layer for request ID middleware +#[derive(Clone)] +pub struct RequestIdLayer { + headers: Arc>, +} + +impl RequestIdLayer { + pub fn new(headers: Vec) -> Self { + Self { + headers: Arc::new(headers), + } + } +} + +impl Layer for RequestIdLayer { + type Service = RequestIdMiddleware; + + fn layer(&self, inner: S) -> Self::Service { + RequestIdMiddleware { + inner, + headers: self.headers.clone(), + } + } +} + +/// Tower Service for request ID middleware +#[derive(Clone)] +pub struct RequestIdMiddleware { + inner: S, + headers: Arc>, +} + +impl Service for RequestIdMiddleware +where + S: Service + Send + 'static, + S::Future: Send + 'static, +{ + type Response = S::Response; + type Error = S::Error; + type Future = std::pin::Pin< + Box> + Send>, + >; + + fn poll_ready( + &mut self, + cx: &mut std::task::Context<'_>, + ) -> std::task::Poll> { + self.inner.poll_ready(cx) + } + + fn call(&mut self, mut req: Request) -> Self::Future { + let headers = self.headers.clone(); + + // Extract request ID from headers or generate new one + let mut request_id = None; + + for header_name in headers.iter() { + if let Some(header_value) = req.headers().get(header_name) { + if let Ok(value) = header_value.to_str() { + request_id = Some(value.to_string()); + break; + } + } + } + + let request_id = request_id.unwrap_or_else(|| generate_request_id(req.uri().path())); + + // Insert request ID into request extensions for other middleware/handlers to use + req.extensions_mut().insert(RequestId(request_id.clone())); + + // Call the inner service + let future = self.inner.call(req); + + Box::pin(async move { + let mut response = future.await?; + + // Add request ID to response headers + response.headers_mut().insert( + "x-request-id", + HeaderValue::from_str(&request_id) + .unwrap_or_else(|_| HeaderValue::from_static("invalid-request-id")), + ); + + Ok(response) + }) + } +} + +// ============= Logging Middleware ============= + +/// Custom span maker that includes request ID +#[derive(Clone, Debug)] +pub struct RequestSpan; + +impl MakeSpan for RequestSpan { + fn make_span(&mut self, request: &Request) -> Span { + let request_path = request.uri().path(); + if !crate::otel_trace::is_otel_enabled() + || crate::otel_trace::is_excluded_http_path(request_path) + { + return Span::none(); + } + + // Don't try to extract request ID here - it won't be available yet + // The RequestIdLayer runs after TraceLayer creates the span + + // Use MatchedPath (route template) when available for low-cardinality + // span names, falling back to the raw URI path otherwise. + let route = request + .extensions() + .get::() + .map(axum::extract::MatchedPath::as_str); + let path_for_name = route.unwrap_or(request_path); + + let span = info_span!( + target: "otel_trace", + "http_request", + otel.kind = "server", + otel.name = %format_args!("{} {}", request.method(), path_for_name), + http.request.method = %request.method(), + http.route = Empty, + url.path = %request_path, + http.response.status_code = Empty, + request_id = Empty, // Will be set later + latency_ms = Empty, + error = Empty, + ); + + if let Some(route) = route { + span.record("http.route", route); + } + + // Extract W3C TraceContext (traceparent/tracestate) from incoming request + // headers and set it as the parent context on this span. This links the + // router's span as a child of the upstream caller's trace. + if !span.is_disabled() { + let parent_cx = opentelemetry::global::get_text_map_propagator(|propagator| { + propagator.extract(&opentelemetry_http::HeaderExtractor(request.headers())) + }); + span.set_parent(parent_cx); + } + + span + } +} + +/// Custom on_request handler +#[derive(Clone, Debug)] +pub struct RequestLogger; + +impl OnRequest for RequestLogger { + fn on_request(&mut self, request: &Request, span: &Span) { + let _enter = span.enter(); + + // Try to get the request ID from extensions + // This will work if RequestIdLayer has already run + if let Some(request_id) = request.extensions().get::() { + span.record("request_id", request_id.0.as_str()); + } + + // Log the request start + debug!( + target: "lmdeploy_router_rs::request", + "started processing request" + ); + } +} + +/// Custom on_response handler +#[derive(Clone, Debug)] +pub struct ResponseLogger { + _start_time: Instant, +} + +impl Default for ResponseLogger { + fn default() -> Self { + Self { + _start_time: Instant::now(), + } + } +} + +impl OnResponse for ResponseLogger { + fn on_response(self, response: &Response, latency: std::time::Duration, span: &Span) { + let status = response.status(); + + if !span.is_disabled() { + // Record OTel semantic convention attributes only when a real span exists. + span.record("http.response.status_code", status.as_u16()); + span.record("latency_ms", latency.as_secs_f64() * 1000.0); + + if status.is_server_error() { + span.set_status(opentelemetry::trace::Status::error(format!( + "HTTP {}", + status.as_u16() + ))); + } + } + + // Log the response completion + let _enter = span.enter(); + if status.is_server_error() { + error!( + target: "lmdeploy_router_rs::response", + "request failed with server error" + ); + } else if status.is_client_error() { + warn!( + target: "lmdeploy_router_rs::response", + "request failed with client error" + ); + } else { + debug!( + target: "lmdeploy_router_rs::response", + "finished processing request" + ); + } + } +} + +/// Custom on_failure handler for pre-response errors (e.g. connection drops). +#[derive(Clone, Debug, Default)] +pub struct FailureLogger; + +impl OnFailure for FailureLogger { + fn on_failure( + &mut self, + failure: ServerErrorsFailureClass, + latency: std::time::Duration, + span: &Span, + ) { + let message = match &failure { + ServerErrorsFailureClass::StatusCode(code) => { + format!("response failed with status {code}") + } + ServerErrorsFailureClass::Error(err) => { + format!("request failed: {err}") + } + }; + + let _enter = span.enter(); + error!( + target: "lmdeploy_router_rs::response", + "{message}" + ); + + if !span.is_disabled() { + span.record("latency_ms", latency.as_secs_f64() * 1000.0); + span.record("error", message.as_str()); + span.set_status(opentelemetry::trace::Status::error(message)); + } + } +} + +/// Create a configured TraceLayer for HTTP logging +/// Note: Actual request/response logging with request IDs is done in RequestIdService +pub fn create_logging_layer() -> TraceLayer< + tower_http::classify::SharedClassifier, + RequestSpan, + RequestLogger, + ResponseLogger, + (), + (), + FailureLogger, +> { + TraceLayer::new_for_http() + .make_span_with(RequestSpan) + .on_request(RequestLogger) + .on_response(ResponseLogger::default()) + .on_body_chunk(()) + .on_eos(()) + .on_failure(FailureLogger) +} + +/// Structured logging data for requests +#[derive(Debug, serde::Serialize)] +pub struct RequestLogEntry { + pub timestamp: String, + pub request_id: String, + pub method: String, + pub uri: String, + pub status: u16, + pub latency_ms: u64, + pub user_agent: Option, + pub remote_addr: Option, + pub error: Option, +} + +/// Log a request with structured data +pub fn log_request(entry: RequestLogEntry) { + if entry.status >= 500 { + tracing::error!( + target: "lmdeploy_router_rs::http", + request_id = %entry.request_id, + method = %entry.method, + uri = %entry.uri, + status = entry.status, + latency_ms = entry.latency_ms, + user_agent = ?entry.user_agent, + remote_addr = ?entry.remote_addr, + error = ?entry.error, + "HTTP request failed" + ); + } else if entry.status >= 400 { + tracing::warn!( + target: "lmdeploy_router_rs::http", + request_id = %entry.request_id, + method = %entry.method, + uri = %entry.uri, + status = entry.status, + latency_ms = entry.latency_ms, + user_agent = ?entry.user_agent, + remote_addr = ?entry.remote_addr, + "HTTP request client error" + ); + } else { + tracing::info!( + target: "lmdeploy_router_rs::http", + request_id = %entry.request_id, + method = %entry.method, + uri = %entry.uri, + status = entry.status, + latency_ms = entry.latency_ms, + user_agent = ?entry.user_agent, + remote_addr = ?entry.remote_addr, + "HTTP request completed" + ); + } +} + +// ============ Concurrency Limiting with Queue Support ============ + +/// Request queue entry +pub struct QueuedRequest { + /// Time when the request was queued + queued_at: Instant, + /// Channel to send the permit back when acquired + permit_tx: oneshot::Sender>, +} + +/// Queue metrics for monitoring +#[derive(Debug, Default)] +pub struct QueueMetrics { + pub total_queued: std::sync::atomic::AtomicU64, + pub current_queued: std::sync::atomic::AtomicU64, + pub total_timeout: std::sync::atomic::AtomicU64, + pub total_rejected: std::sync::atomic::AtomicU64, +} + +/// Queue processor that handles queued requests +pub struct QueueProcessor { + token_bucket: Arc, + queue_rx: mpsc::Receiver, + queue_timeout: Duration, +} + +impl QueueProcessor { + pub fn new( + token_bucket: Arc, + queue_rx: mpsc::Receiver, + queue_timeout: Duration, + ) -> Self { + Self { + token_bucket, + queue_rx, + queue_timeout, + } + } + + pub async fn run(mut self) { + info!("Starting concurrency queue processor"); + + // Process requests in a single task to reduce overhead + while let Some(queued) = self.queue_rx.recv().await { + // Check timeout immediately + let elapsed = queued.queued_at.elapsed(); + if elapsed >= self.queue_timeout { + warn!("Request already timed out in queue"); + let _ = queued.permit_tx.send(Err(StatusCode::REQUEST_TIMEOUT)); + continue; + } + + let remaining_timeout = self.queue_timeout - elapsed; + + // Try to acquire token for this request + if self.token_bucket.try_acquire(1.0).await { + // Got token immediately + debug!("Queue: acquired token immediately for queued request"); + let _ = queued.permit_tx.send(Ok(())); + } else { + // Need to wait for token + let token_bucket = self.token_bucket.clone(); + + // Spawn task only when we actually need to wait + tokio::spawn(async move { + if token_bucket + .acquire_timeout(1.0, remaining_timeout) + .await + .is_ok() + { + debug!("Queue: acquired token after waiting"); + let _ = queued.permit_tx.send(Ok(())); + } else { + warn!("Queue: request timed out waiting for token"); + let _ = queued.permit_tx.send(Err(StatusCode::REQUEST_TIMEOUT)); + } + }); + } + } + + warn!("Concurrency queue processor shutting down"); + } +} + +/// State for the concurrency limiter +pub struct ConcurrencyLimiter { + pub queue_tx: Option>, +} + +impl ConcurrencyLimiter { + /// Create new concurrency limiter with optional queue + pub fn new( + token_bucket: Arc, + queue_size: usize, + queue_timeout: Duration, + ) -> (Self, Option) { + if queue_size > 0 { + let (queue_tx, queue_rx) = mpsc::channel(queue_size); + let processor = QueueProcessor::new(token_bucket, queue_rx, queue_timeout); + + ( + Self { + queue_tx: Some(queue_tx), + }, + Some(processor), + ) + } else { + (Self { queue_tx: None }, None) + } + } +} + +/// Middleware function for concurrency limiting with optional queuing +pub async fn concurrency_limit_middleware( + State(app_state): State>, + request: Request, + next: Next, +) -> Response { + // Static counter for embeddings queue size + static EMBEDDINGS_QUEUE_SIZE: AtomicU64 = AtomicU64::new(0); + + // Identify if this is an embeddings request based on path + let is_embeddings = request.uri().path().contains("/v1/embeddings"); + let token_bucket = app_state.context.rate_limiter.clone(); + + // Try to acquire token immediately + if token_bucket.try_acquire(1.0).await { + debug!("Acquired token immediately"); + let response = next.run(request).await; + + // Return the token to the bucket + token_bucket.return_tokens(1.0).await; + + response + } else { + // No tokens available, try to queue if enabled + if let Some(queue_tx) = &app_state.concurrency_queue_tx { + debug!("No tokens available, attempting to queue request"); + + // Create a channel for the token response + let (permit_tx, permit_rx) = oneshot::channel(); + + let queued = QueuedRequest { + queued_at: Instant::now(), + permit_tx, + }; + + // Try to send to queue + match queue_tx.try_send(queued) { + Ok(_) => { + // On successful enqueue, update embeddings queue gauge if applicable + if is_embeddings { + let new_val = EMBEDDINGS_QUEUE_SIZE.fetch_add(1, Ordering::Relaxed) + 1; + RouterMetrics::set_embeddings_queue_size(new_val as usize); + } + + // Wait for token from queue processor + match permit_rx.await { + Ok(Ok(())) => { + debug!("Acquired token from queue"); + // Dequeue for embeddings + if is_embeddings { + let new_val = + EMBEDDINGS_QUEUE_SIZE.fetch_sub(1, Ordering::Relaxed) - 1; + RouterMetrics::set_embeddings_queue_size(new_val as usize); + } + + let response = next.run(request).await; + + // Return the token to the bucket + token_bucket.return_tokens(1.0).await; + + response + } + Ok(Err(status)) => { + warn!("Queue returned error status: {}", status); + // Dequeue for embeddings on error + if is_embeddings { + let new_val = + EMBEDDINGS_QUEUE_SIZE.fetch_sub(1, Ordering::Relaxed) - 1; + RouterMetrics::set_embeddings_queue_size(new_val as usize); + } + status.into_response() + } + Err(_) => { + error!("Queue response channel closed"); + // Dequeue for embeddings on channel error + if is_embeddings { + let new_val = + EMBEDDINGS_QUEUE_SIZE.fetch_sub(1, Ordering::Relaxed) - 1; + RouterMetrics::set_embeddings_queue_size(new_val as usize); + } + StatusCode::INTERNAL_SERVER_ERROR.into_response() + } + } + } + Err(_) => { + warn!("Request queue is full, returning 429"); + StatusCode::TOO_MANY_REQUESTS.into_response() + } + } + } else { + warn!("No tokens available and queuing is disabled, returning 429"); + StatusCode::TOO_MANY_REQUESTS.into_response() + } + } +} diff --git a/lmdeploy/router/src/otel_http.rs b/lmdeploy/router/src/otel_http.rs new file mode 100644 index 0000000000..0704374dff --- /dev/null +++ b/lmdeploy/router/src/otel_http.rs @@ -0,0 +1,152 @@ +use axum::http::HeaderMap; +use tracing::{field::Empty, info_span, Instrument, Span}; +use tracing_opentelemetry::OpenTelemetrySpanExt; +use url::Url; + +pub const TRACE_HEADER_NAMES: &[&str] = &["traceparent", "tracestate", "baggage"]; + +#[derive(Clone, Copy, Debug)] +pub struct ClientRequestOptions<'a> { + pub method: &'a str, + pub url: &'a str, + pub route: Option<&'a str>, + pub request_phase: Option<&'a str>, +} + +pub struct PreparedClientRequest { + request: reqwest::RequestBuilder, + span: Span, +} + +impl PreparedClientRequest { + pub async fn send(self) -> Result { + let Self { request, span } = self; + + let response = if span.is_none() { + request.send().await + } else { + request.send().instrument(span.clone()).await + }; + + match response { + Ok(response) => { + if !span.is_disabled() { + let status = response.status(); + span.record("http.response.status_code", status.as_u16()); + + if status.is_server_error() { + span.set_status(opentelemetry::trace::Status::error(format!( + "HTTP {}", + status.as_u16() + ))); + } + } + + Ok(response) + } + Err(error) => { + if !span.is_disabled() { + let message = error.to_string(); + span.record("error", message.as_str()); + span.set_status(opentelemetry::trace::Status::error(message)); + } + + Err(error) + } + } + } +} + +pub fn prepare_client_request( + mut request: reqwest::RequestBuilder, + incoming_headers: Option<&HeaderMap>, + options: ClientRequestOptions<'_>, +) -> PreparedClientRequest { + let span = if crate::otel_trace::is_otel_enabled() { + let span = info_span!( + target: "otel_trace", + "http_client_request", + otel.kind = "client", + otel.name = %format_args!( + "{} {}", + options.method, + options.route.unwrap_or(options.url) + ), + http.request.method = %options.method, + http.route = Empty, + url.full = %options.url, + server.address = Empty, + server.port = Empty, + inference.request_phase = Empty, + http.response.status_code = Empty, + error = Empty, + ); + + if let Some(route) = options.route { + span.record("http.route", route); + } + if let Some(request_phase) = options.request_phase { + span.record("inference.request_phase", request_phase); + } + record_server_address(&span, options.url); + + { + let _enter = span.enter(); + request = propagate_trace_headers(request, incoming_headers); + } + span + } else { + request = propagate_trace_headers(request, incoming_headers); + Span::none() + }; + + PreparedClientRequest { request, span } +} + +pub async fn send_client_request( + request: reqwest::RequestBuilder, + incoming_headers: Option<&HeaderMap>, + options: ClientRequestOptions<'_>, +) -> Result { + prepare_client_request(request, incoming_headers, options) + .send() + .await +} + +pub fn propagate_trace_headers( + mut request: reqwest::RequestBuilder, + headers: Option<&HeaderMap>, +) -> reqwest::RequestBuilder { + if crate::otel_trace::is_otel_enabled() { + let mut trace_headers = HeaderMap::new(); + crate::otel_trace::inject_trace_context_http(&mut trace_headers); + for (name, value) in trace_headers.iter() { + request = request.header(name, value); + } + return request; + } + + if let Some(headers) = headers { + for &name in TRACE_HEADER_NAMES { + if let Some(value) = headers.get(name) { + request = request.header(name, value); + } + } + } + + request +} + +fn record_server_address(span: &Span, url: &str) { + let Ok(parsed) = Url::parse(url) else { + return; + }; + + if let Some(host) = parsed.host_str() { + span.record("server.address", host); + } + + if let Some(port) = parsed.port_or_known_default() { + span.record("server.port", port); + } +} diff --git a/lmdeploy/router/src/otel_trace.rs b/lmdeploy/router/src/otel_trace.rs new file mode 100644 index 0000000000..8a1253d053 --- /dev/null +++ b/lmdeploy/router/src/otel_trace.rs @@ -0,0 +1,280 @@ +//! OpenTelemetry tracing integration. + +use std::{ + sync::{ + atomic::{AtomicBool, Ordering}, + Mutex, OnceLock, + }, + time::Duration, +}; + +use anyhow::Result; +use axum::http::{HeaderMap, HeaderName, HeaderValue}; +use opentelemetry::{ + global, propagation::TextMapCompositePropagator, trace::TracerProvider as _, KeyValue, +}; +use opentelemetry_otlp::WithExportConfig; +use opentelemetry_sdk::{ + propagation::{BaggagePropagator, TraceContextPropagator}, + runtime, + trace::{BatchConfigBuilder, BatchSpanProcessor, Sampler, Tracer as SdkTracer, TracerProvider}, + Resource, +}; +use tracing::{Metadata, Subscriber}; +use tracing_opentelemetry::OpenTelemetrySpanExt; +use tracing_subscriber::{ + layer::{Context, Filter}, + Layer, +}; + +use crate::config::TraceConfig; + +/// Whether OpenTelemetry tracing is enabled. +/// +/// This flag guards runtime code paths (make_span parent extraction, +/// inject_trace_context_http). Only set to true after both the OTel layer +/// and subscriber are installed via `activate_otel`. +static ENABLED: AtomicBool = AtomicBool::new(false); +static PROVIDER: Mutex> = Mutex::new(None); +static EXCLUDED_PATHS: OnceLock> = OnceLock::new(); + +const OTEL_SPAN_TARGET: &str = "otel_trace"; + +/// Filter that only allows specific module targets to be exported to OTEL. +#[derive(Clone, Copy, Default)] +pub(crate) struct CustomOtelFilter; + +impl CustomOtelFilter { + #[inline] + pub const fn new() -> Self { + Self + } + + #[inline] + fn is_allowed(target: &str) -> bool { + target.starts_with(OTEL_SPAN_TARGET) + } +} + +impl Filter for CustomOtelFilter +where + S: Subscriber, +{ + #[inline] + fn enabled(&self, meta: &Metadata<'_>, _cx: &Context<'_, S>) -> bool { + Self::is_allowed(meta.target()) + } + + #[inline] + fn callsite_enabled(&self, meta: &'static Metadata<'static>) -> tracing::subscriber::Interest { + if Self::is_allowed(meta.target()) { + tracing::subscriber::Interest::always() + } else { + tracing::subscriber::Interest::never() + } + } +} + +/// Prepared OTel state ready to be activated. +/// +/// Built by `prepare_otel()`, this holds the provider and tracer without +/// mutating any global state. Call `layer()` to create the tracing layer, +/// then `activate_otel()` after the subscriber is installed. +pub struct PreparedOtel { + provider: TracerProvider, + tracer: SdkTracer, + excluded_paths: Box<[String]>, +} + +impl PreparedOtel { + /// Build from pre-existing parts (useful for tests with in-memory exporters). + pub fn from_parts(provider: TracerProvider, tracer: SdkTracer) -> Self { + Self::from_parts_with_excluded_paths(provider, tracer, Vec::new()) + } + + /// Build from pre-existing parts with an explicit excluded-path list. + pub fn from_parts_with_excluded_paths( + provider: TracerProvider, + tracer: SdkTracer, + excluded_paths: Vec, + ) -> Self { + Self { + provider, + tracer, + excluded_paths: excluded_paths.into_boxed_slice(), + } + } + + /// Create the tracing-opentelemetry layer for this prepared state. + pub fn layer(&self) -> Box + Send + Sync + 'static> + where + S: Subscriber + for<'a> tracing_subscriber::registry::LookupSpan<'a> + Send + Sync, + { + let layer = tracing_opentelemetry::layer() + .with_tracer(self.tracer.clone()) + .with_filter(CustomOtelFilter::new()); + Box::new(layer) + } +} + +/// Cross-platform hostname with PID for unique instance identification. +fn service_instance_id() -> String { + let hostname = ["HOSTNAME", "HOST", "COMPUTERNAME"] + .iter() + .find_map(|key| std::env::var(key).ok().filter(|value| !value.is_empty())); + match hostname { + Some(hostname) => format!("{hostname}-{}", std::process::id()), + None => format!("pid-{}", std::process::id()), + } +} + +fn parent_based_sampler(sampling_ratio: f64) -> Sampler { + if sampling_ratio <= 0.0 { + Sampler::ParentBased(Box::new(Sampler::AlwaysOff)) + } else if sampling_ratio >= 1.0 { + Sampler::ParentBased(Box::new(Sampler::AlwaysOn)) + } else { + Sampler::ParentBased(Box::new(Sampler::TraceIdRatioBased(sampling_ratio))) + } +} + +/// Build OTel provider and tracer without mutating global state. +pub fn prepare_otel(trace_config: &TraceConfig) -> Result { + anyhow::ensure!( + (0.0..=1.0).contains(&trace_config.sampling_ratio), + "OTel sampling_ratio must be within [0.0, 1.0], got {}", + trace_config.sampling_ratio + ); + + let mut exporter_builder = opentelemetry_otlp::SpanExporter::builder() + .with_tonic() + .with_protocol(opentelemetry_otlp::Protocol::Grpc); + + // Only set the endpoint explicitly when configured; otherwise the SDK + // respects OTEL_EXPORTER_OTLP_ENDPOINT (defaulting to localhost:4317). + if let Some(ep) = trace_config.otlp_traces_endpoint.as_deref() { + let ep = if !ep.starts_with("http://") && !ep.starts_with("https://") { + format!("http://{ep}") + } else { + ep.to_string() + }; + exporter_builder = exporter_builder.with_endpoint(&ep); + } + + let exporter = exporter_builder + .build() + .map_err(|e| anyhow::anyhow!("Failed to create OTLP exporter: {e}"))?; + + let batch_config = BatchConfigBuilder::default() + .with_scheduled_delay(Duration::from_millis(500)) + .with_max_export_batch_size(64) + .build(); + + let span_processor = BatchSpanProcessor::builder(exporter, runtime::Tokio) + .with_batch_config(batch_config) + .build(); + + let resource_attrs = vec![ + KeyValue::new("service.name", "lmdeploy-router"), + KeyValue::new("service.version", env!("CARGO_PKG_VERSION")), + KeyValue::new("service.instance.id", service_instance_id()), + ]; + let resource = Resource::default().merge(&Resource::new(resource_attrs)); + + let provider = TracerProvider::builder() + .with_sampler(parent_based_sampler(trace_config.sampling_ratio)) + .with_span_processor(span_processor) + .with_resource(resource) + .build(); + + let tracer = provider.tracer("lmdeploy-router"); + + Ok(PreparedOtel { + provider, + tracer, + excluded_paths: trace_config.excluded_paths.clone().into_boxed_slice(), + }) +} + +/// Set global propagator and provider, flip ENABLED. +/// +/// Must be called after the subscriber with the OTel layer is installed. +pub fn activate_otel(prepared: PreparedOtel) { + let PreparedOtel { + provider, + excluded_paths, + .. + } = prepared; + + global::set_text_map_propagator(TextMapCompositePropagator::new(vec![ + Box::new(TraceContextPropagator::new()), + Box::new(BaggagePropagator::new()), + ])); + + let _ = EXCLUDED_PATHS.set(excluded_paths); + let _ = global::set_tracer_provider(provider.clone()); + *PROVIDER.lock().unwrap() = Some(provider); + ENABLED.store(true, Ordering::Release); + tracing::info!("[tracing] OpenTelemetry tracing enabled"); +} + +/// Mark OpenTelemetry tracing as enabled (for tests that manage their own provider). +pub fn mark_otel_enabled() { + ENABLED.store(true, Ordering::Release); +} + +/// Check if OpenTelemetry tracing is enabled. +#[inline] +pub fn is_otel_enabled() -> bool { + ENABLED.load(Ordering::Acquire) +} + +#[inline] +pub fn is_excluded_http_path(path: &str) -> bool { + EXCLUDED_PATHS + .get() + .is_some_and(|paths| paths.iter().any(|excluded| excluded == path)) +} + +pub fn shutdown_otel() { + if ENABLED.load(Ordering::Acquire) { + tracing::info!("[tracing] OpenTelemetry shutting down"); + if let Some(provider) = PROVIDER.lock().unwrap().take() { + let _ = provider.shutdown(); + } + ENABLED.store(false, Ordering::Release); + } +} + +/// Inject W3C trace context headers into an HTTP request's HeaderMap. +/// +/// When OTel is enabled, this injects the current span's trace context +/// (traceparent, tracestate) into outgoing requests, making the router's +/// span the parent of backend spans. +/// +/// When OTel is disabled, this is a no-op. +#[inline] +pub fn inject_trace_context_http(headers: &mut HeaderMap) { + if !is_otel_enabled() { + return; + } + + let context = tracing::Span::current().context(); + + struct HeaderInjector<'a>(&'a mut HeaderMap); + + impl opentelemetry::propagation::Injector for HeaderInjector<'_> { + #[inline] + fn set(&mut self, key: &str, value: String) { + if let Ok(header_name) = HeaderName::from_bytes(key.as_bytes()) { + if let Ok(header_value) = HeaderValue::from_str(&value) { + self.0.insert(header_name, header_value); + } + } + } + } + + global::get_text_map_propagator(|propagator| { + propagator.inject_context(&context, &mut HeaderInjector(headers)); + }); +} diff --git a/lmdeploy/router/src/policies/cache_aware.rs b/lmdeploy/router/src/policies/cache_aware.rs new file mode 100644 index 0000000000..dc209a193c --- /dev/null +++ b/lmdeploy/router/src/policies/cache_aware.rs @@ -0,0 +1,577 @@ +/* + Cache-Aware Load Balancing Router + + This router combines two strategies to optimize both cache utilization and request distribution: + + 1. Cache-Aware Routing (Approximate Tree) + 2. Load Balancing (Shortest Queue with Balance Thresholds) + + The router dynamically switches between these strategies based on load conditions: + - Uses load balancing when the system is imbalanced + - Uses cache-aware routing when the system is balanced + + A system is considered imbalanced if both conditions are met: + 1. (max - min) > abs_threshold + 2. max > rel_threshold * min + + Strategy Details: + + 1. Cache-Aware Routing (Approximate Tree) + ------------------------------------------- + This strategy maintains an approximate radix tree for each worker based on request history, + eliminating the need for direct cache state queries. The tree stores raw text characters + instead of token IDs to avoid tokenization overhead. + + Process: + a. For each request, find the worker with the highest prefix match + b. If match rate > cache_threshold: + Route to the worker with highest match (likely has relevant data cached) + c. If match rate ≤ cache_threshold: + Route to the worker with smallest tree size (most available cache capacity) + d. Background maintenance: + Periodically evict least recently used leaf nodes to prevent memory overflow + + 2. Load Balancing (Shortest Queue) + ------------------------------------------- + This strategy tracks pending request counts per worker and routes new requests + to the least busy worker when the system is detected to be imbalanced. + + Configuration Parameters: + ------------------------ + 1. cache_threshold: (float, 0.0 to 1.0) + Minimum prefix match ratio to use highest-match routing. + Below this threshold, routes to worker with most available cache space. + + 2. balance_abs_threshold: (integer) + Absolute difference threshold for load imbalance detection. + System is potentially imbalanced if (max_load - min_load) > abs_threshold + + 3. balance_rel_threshold: (float) + Relative ratio threshold for load imbalance detection. + System is potentially imbalanced if max_load > min_load * rel_threshold + Used in conjunction with abs_threshold to determine final imbalance state. + + 4. eviction_interval_secs: (integer) + Interval between LRU eviction cycles for the approximate trees. + + 5. max_tree_size: (integer) + Maximum nodes per tree. When exceeded, LRU leaf nodes are evicted + during the next eviction cycle. +*/ + +use super::{get_healthy_worker_indices, CacheAwareConfig, LoadBalancingPolicy, RequestHeaders}; +use crate::core::Worker; +use crate::metrics::RouterMetrics; +use crate::policies::normalize_model_key; +use crate::tree::Tree; +use dashmap::DashMap; +use rand::Rng; +use std::collections::HashMap; +use std::sync::Arc; +use std::thread; +use std::time::Duration; +use tracing::{debug, info}; + +/// Cache-aware routing policy +/// +/// Routes requests based on cache affinity when load is balanced, +/// switches to shortest-queue routing when load is imbalanced. +/// Maintains separate trees per model for multi-model support. +#[derive(Debug)] +pub struct CacheAwarePolicy { + config: CacheAwareConfig, + trees: Arc>>, // model_id -> Arc + eviction_handle: Option>, +} + +impl CacheAwarePolicy { + pub fn new() -> Self { + Self::with_config(CacheAwareConfig::default()) + } + + pub fn with_config(config: CacheAwareConfig) -> Self { + let trees = Arc::new(DashMap::>::new()); + + // Start background eviction thread if configured + let eviction_handle = if config.eviction_interval_secs > 0 { + let trees_clone = Arc::clone(&trees); + let max_tree_size = config.max_tree_size; + let interval = config.eviction_interval_secs; + + Some(thread::spawn(move || loop { + thread::sleep(Duration::from_secs(interval)); + + // Evict for all model trees + for entry in trees_clone.iter() { + let model_id = entry.key(); + let tree = entry.value(); + tree.evict_tenant_by_size(max_tree_size); + debug!( + "Cache eviction completed for model {}, max_size: {}", + model_id, max_tree_size + ); + } + })) + } else { + None + }; + + Self { + config, + trees, + eviction_handle, + } + } + + /// Add a single worker to the tree (incremental update) + pub fn add_worker(&self, worker: &dyn Worker) { + let tree_key = normalize_model_key(worker.model_id()); + let tree = self + .trees + .entry(tree_key.to_string()) + .or_insert_with(|| Arc::new(Tree::new())); + tree.insert("", worker.url()); + } + + /// Add a worker by URL and model (for backward compatibility) + pub fn add_worker_by_url(&self, url: &str, model_id: &str) { + let tree = self + .trees + .entry(model_id.to_string()) + .or_insert_with(|| Arc::new(Tree::new())); + tree.insert("", url); + } + + /// Remove a worker from the tree + pub fn remove_worker(&self, worker: &dyn Worker) { + let tree_key = normalize_model_key(worker.model_id()); + if let Some(tree) = self.trees.get(tree_key) { + tree.remove_tenant(worker.url()); + } + } + + /// Remove a worker by URL (removes from all model trees for backward compatibility) + pub fn remove_worker_by_url(&self, url: &str) { + // Remove from all trees since we don't know which model it belongs to + for tree_ref in self.trees.iter() { + tree_ref.value().remove_tenant(url); + } + } + + /// Run cache eviction to prevent unbounded growth + pub fn evict_cache(&self, max_size: usize) { + for tree_ref in self.trees.iter() { + let model_id = tree_ref.key(); + let tree = tree_ref.value(); + tree.evict_tenant_by_size(max_size); + debug!( + "Cache eviction for model {}, max_size: {}", + model_id, max_size + ); + } + } + + fn select_worker_min_load( + &self, + workers: &[Arc], + request_text: Option<&str>, + healthy_indices: &[usize], + model_id: &str, + max_load: usize, + min_load: usize, + ) -> Option { + // Log load balancing trigger (only compute worker loads if debug enabled) + if tracing::enabled!(tracing::Level::DEBUG) { + let worker_loads: Vec<(&str, usize)> = + workers.iter().map(|w| (w.url(), w.load())).collect(); + debug!( + "Load balancing triggered | max: {} | min: {} | workers: {:?}", + max_load, min_load, worker_loads + ); + } + + RouterMetrics::record_load_balancing_event(); + RouterMetrics::set_load_range(max_load, min_load); + + // Use shortest queue when imbalanced + let min_load_idx = healthy_indices + .iter() + .min_by_key(|&&idx| workers[idx].load()) + .copied()?; + + // Even in imbalanced mode, update the tree to maintain cache state + if let Some(text) = request_text { + // Get the tree reference without locking the entire HashMap + // DashMap only locks the specific shard containing this key + let tree = self.trees.get(model_id).map(|entry| entry.value().clone()); + + if let Some(tree) = tree { + let worker_url = workers[min_load_idx].url(); + tree.insert(text, worker_url); + } else { + debug!( + "Warning: No tree found for model '{}', skipping cache update", + model_id + ); + } + } + + // Increment processed counter + workers[min_load_idx].increment_processed(); + RouterMetrics::record_processed_request(workers[min_load_idx].url()); + RouterMetrics::record_policy_decision(self.name(), workers[min_load_idx].url()); + + Some(min_load_idx) + } +} + +impl LoadBalancingPolicy for CacheAwarePolicy { + fn select_worker_with_headers( + &self, + workers: &[Arc], + request_text: Option<&str>, + _headers: Option<&RequestHeaders>, + ) -> Option { + let healthy_indices = get_healthy_worker_indices(workers); + + if healthy_indices.is_empty() { + return None; + } + + // Determine the model for this set of workers (router pre-filters by model) + // All workers should be from the same model + let model_id = normalize_model_key(workers[healthy_indices[0]].model_id()); + + // Get current load statistics - compute min/max in single pass without allocation + let (min_load, max_load) = workers.iter().fold((usize::MAX, 0usize), |(min, max), w| { + let load = w.load(); + (min.min(load), max.max(load)) + }); + let min_load = if min_load == usize::MAX { 0 } else { min_load }; + + // Check if load is imbalanced + let is_imbalanced = max_load.saturating_sub(min_load) > self.config.balance_abs_threshold + && (max_load as f32) > (min_load as f32 * self.config.balance_rel_threshold); + + debug!( + "Load status for model: max_load={}, min_load={}, is_imbalanced={}", + max_load, min_load, is_imbalanced + ); + + if is_imbalanced { + return self.select_worker_min_load( + workers, + request_text, + &healthy_indices, + model_id, + max_load, + min_load, + ); + } + + // Use cache-aware routing when balanced + let text = request_text.unwrap_or(""); + + // Get the tree reference without locking the entire HashMap + // DashMap only locks the specific shard containing this key + let tree = self.trees.get(model_id).map(|entry| entry.value().clone()); + + let keys: Vec<_> = self.trees.iter().map(|entry| entry.key().clone()).collect(); + debug!("Available tree keys: {:?}", keys); + + let Some(tree) = tree else { + // No tree for this model, log warning and use random selection + debug!( + "Warning: No tree found for model '{}', using random worker selection", + model_id + ); + // Return a random healthy worker + let mut rng = rand::rng(); + let random_idx = rng.random_range(0..healthy_indices.len()); + let selected_idx = healthy_indices[random_idx]; + + workers[selected_idx].increment_processed(); + RouterMetrics::record_processed_request(workers[selected_idx].url()); + RouterMetrics::record_policy_decision(self.name(), workers[selected_idx].url()); + + return Some(selected_idx); + }; + debug!("Using cache-aware routing for model '{}'", model_id); + // Now we work with the tree without holding the HashMap lock + // Use prefix_match_with_counts to avoid redundant chars().count() calls + let result = tree.prefix_match_with_counts(text); + let match_rate = if result.input_char_count == 0 { + 0.0 + } else { + result.matched_char_count as f32 / result.input_char_count as f32 + }; + + debug!( + "Cache match for model '{}': matched_chars={}, input_chars={}, match_rate={:.2}", + model_id, result.matched_char_count, result.input_char_count, match_rate + ); + // Select worker without String allocation + let selected_idx = if match_rate > self.config.cache_threshold { + // Cache hit path: find worker by URL (compare &str directly, no allocation) + let tenant_url: &str = &result.tenant; + workers + .iter() + .position(|w| w.url() == tenant_url) + .filter(|&idx| workers[idx].is_healthy()) + } else { + // Low cache match: use worker with minimum load + healthy_indices + .iter() + .min_by_key(|&&idx| workers[idx].load()) + .copied() + }; + + if let Some(idx) = selected_idx { + // Update the tree with this request (use worker URL directly, no allocation) + tree.insert(text, workers[idx].url()); + + // Increment processed counter + workers[idx].increment_processed(); + RouterMetrics::record_processed_request(workers[idx].url()); + RouterMetrics::record_policy_decision(self.name(), workers[idx].url()); + + return Some(idx); + } + + // Selected worker no longer exists or unhealthy, remove stale tenant from tree + if match_rate > self.config.cache_threshold { + let tenant_url: &str = &result.tenant; + tree.remove_tenant(tenant_url); + debug!("Removed stale worker {} from cache tree", tenant_url); + } + + // Fallback to first healthy worker + if let Some(idx) = healthy_indices.first().copied() { + workers[idx].increment_processed(); + RouterMetrics::record_processed_request(workers[idx].url()); + RouterMetrics::record_policy_decision(self.name(), workers[idx].url()); + + Some(idx) + } else { + None + } + } + + fn name(&self) -> &'static str { + "cache_aware" + } + + fn needs_request_text(&self) -> bool { + true // Cache-aware policy needs request text for cache affinity + } + + fn on_request_complete(&self, worker_url: &str, success: bool) { + // Could track success rates per worker for more intelligent routing + if !success { + // Optionally reduce affinity for failed requests + tracing::debug!( + "Request to {} completed with success={}", + worker_url, + success + ); + } + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn select_worker_pair_with_headers( + &self, + prefill_workers: &[Arc], + decode_workers: &[Arc], + request_text: Option<&str>, + headers: Option<&RequestHeaders>, + ) -> Option<(usize, usize)> { + // DEPRECATED: This method is no longer used when separate policies are configured. + // The PD router now uses separate policies for prefill and decode selection. + // This implementation remains for backward compatibility when a single policy is used. + + // In PD mode with single policy: + // - Prefill: Use cache-aware routing for better cache utilization + // - Decode: Use least-load routing for better load distribution + + // Select prefill worker using cache-aware logic + let prefill_idx = + self.select_worker_with_headers(prefill_workers, request_text, headers)?; + + // Select decode worker using least-load logic + let healthy_decode = get_healthy_worker_indices(decode_workers); + if healthy_decode.is_empty() { + return None; + } + + let decode_idx = healthy_decode + .iter() + .min_by_key(|&&idx| decode_workers[idx].load()) + .copied()?; + + Some((prefill_idx, decode_idx)) + } + + fn requires_initialization(&self) -> bool { + true // Cache-aware policy requires init_workers() to set up trees + } + + fn init_workers(&self, workers: &[Arc]) { + // Group workers by model + info!( + "Initializing workers for cache-aware policy: {}", + workers + .iter() + .map(|w| w.url()) + .collect::>() + .join(", ") + ); + let mut model_workers: HashMap>> = HashMap::new(); + for worker in workers { + let tree_key = normalize_model_key(worker.model_id()); + model_workers + .entry(tree_key.to_string()) + .or_default() + .push(worker); + } + + // Initialize tree for each model + for (tree_key, model_workers) in model_workers { + info!( + "Creating tree for model key: '{}' with {} workers", + tree_key, + model_workers.len() + ); + let tree = self + .trees + .entry(tree_key) + .or_insert_with(|| Arc::new(Tree::new())) + .clone(); + for worker in model_workers { + tree.insert("", worker.url()); + } + } + } +} + +impl Default for CacheAwarePolicy { + fn default() -> Self { + Self::new() + } +} + +impl Drop for CacheAwarePolicy { + fn drop(&mut self) { + // Note: We can't properly stop the eviction thread since it's in an infinite loop + // In a production system, we'd use a channel or atomic flag to signal shutdown + if let Some(handle) = self.eviction_handle.take() { + // The thread will continue running until the program exits + // This is acceptable for now since the router typically runs for the lifetime of the program + drop(handle); + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::{BasicWorker, WorkerType}; + + #[test] + fn test_cache_aware_with_balanced_load() { + // Create policy without eviction thread for testing + let config = CacheAwareConfig { + eviction_interval_secs: 0, // Disable eviction thread + ..Default::default() + }; + let policy = CacheAwarePolicy::with_config(config); + let workers: Vec> = vec![ + Arc::new(BasicWorker::new( + "http://w1:8000".to_string(), + WorkerType::Regular, + )), + Arc::new(BasicWorker::new( + "http://w2:8000".to_string(), + WorkerType::Regular, + )), + ]; + + // Initialize the policy with workers + policy.init_workers(&workers); + + // First request should be distributed + let idx1 = policy.select_worker(&workers, Some("hello world")).unwrap(); + + // Same request should go to same worker (cache hit) + let idx2 = policy.select_worker(&workers, Some("hello world")).unwrap(); + assert_eq!(idx1, idx2); + + // Similar request should also go to same worker + let idx3 = policy.select_worker(&workers, Some("hello")).unwrap(); + assert_eq!(idx1, idx3); + } + + #[test] + fn test_cache_aware_with_imbalanced_load() { + let policy = CacheAwarePolicy::with_config(CacheAwareConfig { + cache_threshold: 0.5, + balance_abs_threshold: 5, + balance_rel_threshold: 2.0, + eviction_interval_secs: 0, // Disable eviction thread + max_tree_size: 10000, + }); + + let worker1 = BasicWorker::new("http://w1:8000".to_string(), WorkerType::Regular); + let worker2 = BasicWorker::new("http://w2:8000".to_string(), WorkerType::Regular); + + // Create significant load imbalance + for _ in 0..20 { + worker1.increment_load(); + } + // worker2 has load 0 + + let workers: Vec> = vec![Arc::new(worker1), Arc::new(worker2)]; + policy.init_workers(&workers); + + // Should select worker2 (lower load) despite cache affinity + for _ in 0..5 { + let idx = policy.select_worker(&workers, Some("test")).unwrap(); + assert_eq!(idx, 1); // Should always pick worker2 + } + } + + #[test] + fn test_cache_aware_worker_removal() { + let config = CacheAwareConfig { + eviction_interval_secs: 0, // Disable eviction thread + ..Default::default() + }; + let policy = CacheAwarePolicy::with_config(config); + let workers: Vec> = vec![ + Arc::new(BasicWorker::new( + "http://w1:8000".to_string(), + WorkerType::Regular, + )), + Arc::new(BasicWorker::new( + "http://w2:8000".to_string(), + WorkerType::Regular, + )), + ]; + + policy.init_workers(&workers); + + // Route some requests + policy.select_worker(&workers, Some("test1")); + policy.select_worker(&workers, Some("test2")); + + // Remove a worker + policy.remove_worker_by_url("http://w1:8000"); + workers[0].set_healthy(false); + + // All requests should now go to worker2 + let idx = policy.select_worker(&workers, Some("test1")).unwrap(); + assert_eq!(idx, 1); + } +} diff --git a/lmdeploy/router/src/policies/consistent_hash.rs b/lmdeploy/router/src/policies/consistent_hash.rs new file mode 100644 index 0000000000..150d0bc2af --- /dev/null +++ b/lmdeploy/router/src/policies/consistent_hash.rs @@ -0,0 +1,528 @@ +//! Consistent hashing load balancing policy +//! +//! This policy implements consistent hashing to route requests to workers based on +//! session ID or user ID, ensuring that requests from the same user/session are +//! consistently routed to the same worker for better cache locality. + +use std::collections::BTreeMap; +use std::sync::{Arc, RwLock}; + +use tracing::debug; +use tracing::info; + +use super::get_healthy_worker_indices; +use super::hash_key; +use super::LoadBalancingPolicy; +use super::RequestHeaders; +use crate::core::Worker; +use crate::metrics::RouterMetrics; + +/// Number of virtual nodes per physical worker (for better load distribution) +pub const VIRTUAL_NODES_PER_WORKER: u32 = 160; + +/// Consistent hashing policy +/// +/// Routes requests based on session ID or user ID using consistent hashing, +/// ensuring that requests from the same user/session consistently go to the same worker. +#[derive(Debug)] +pub struct ConsistentHashPolicy { + /// Hash ring mapping hash values to worker URLs + hash_ring: RwLock>, + /// Current set of workers (for detecting changes) + current_workers: RwLock>, +} + +impl ConsistentHashPolicy { + pub fn new() -> Self { + Self { + hash_ring: RwLock::new(BTreeMap::new()), + current_workers: RwLock::new(Vec::new()), + } + } + + /// MurmurHash64A implementation from Facebook's mcrouter/lib/fbi/hash.c + fn murmur_hash_64a(key: &[u8], seed: u32) -> u64 { + const M: u64 = 0xc6a4a7935bd1e995; + const R: i32 = 47; + + let mut h = (seed as u64) ^ ((key.len() as u64).wrapping_mul(M)); + + // Process 8-byte chunks + let (chunks, remainder) = key.as_chunks::<8>(); + + for chunk in chunks { + let mut k = u64::from_le_bytes([ + chunk[0], chunk[1], chunk[2], chunk[3], chunk[4], chunk[5], chunk[6], chunk[7], + ]); + + k = k.wrapping_mul(M); + k ^= k >> R; + k = k.wrapping_mul(M); + + h ^= k; + h = h.wrapping_mul(M); + } + + // Process remaining bytes + match remainder.len() { + 7 => { + h ^= (remainder[6] as u64) << 48; + h ^= (remainder[5] as u64) << 40; + h ^= (remainder[4] as u64) << 32; + h ^= (remainder[3] as u64) << 24; + h ^= (remainder[2] as u64) << 16; + h ^= (remainder[1] as u64) << 8; + h ^= remainder[0] as u64; + h = h.wrapping_mul(M); + } + 6 => { + h ^= (remainder[5] as u64) << 40; + h ^= (remainder[4] as u64) << 32; + h ^= (remainder[3] as u64) << 24; + h ^= (remainder[2] as u64) << 16; + h ^= (remainder[1] as u64) << 8; + h ^= remainder[0] as u64; + h = h.wrapping_mul(M); + } + 5 => { + h ^= (remainder[4] as u64) << 32; + h ^= (remainder[3] as u64) << 24; + h ^= (remainder[2] as u64) << 16; + h ^= (remainder[1] as u64) << 8; + h ^= remainder[0] as u64; + h = h.wrapping_mul(M); + } + 4 => { + h ^= (remainder[3] as u64) << 24; + h ^= (remainder[2] as u64) << 16; + h ^= (remainder[1] as u64) << 8; + h ^= remainder[0] as u64; + h = h.wrapping_mul(M); + } + 3 => { + h ^= (remainder[2] as u64) << 16; + h ^= (remainder[1] as u64) << 8; + h ^= remainder[0] as u64; + h = h.wrapping_mul(M); + } + 2 => { + h ^= (remainder[1] as u64) << 8; + h ^= remainder[0] as u64; + h = h.wrapping_mul(M); + } + 1 => { + h ^= remainder[0] as u64; + h = h.wrapping_mul(M); + } + _ => {} + } + + h ^= h >> R; + h = h.wrapping_mul(M); + h ^= h >> R; + + h + } + + /// Optimized MurmurHash64A for rehashing uint64_t keys (from hash.c) + fn murmur_rehash_64a(k: u64) -> u64 { + const M: u64 = 0xc6a4a7935bd1e995; + const R: i32 = 47; + const SEED: u64 = 4193360111; // From hash.c + + let mut h = SEED ^ (8u64.wrapping_mul(M)); // sizeof(uint64_t) * M + + let mut k = k; + k = k.wrapping_mul(M); + k ^= k >> R; + k = k.wrapping_mul(M); + + h ^= k; + h = h.wrapping_mul(M); + + h ^= h >> R; + h = h.wrapping_mul(M); + h ^= h >> R; + + h + } + + /// FurcHash implementation from Facebook's mcrouter/lib/fbi/hash.c + /// This is the actual consistent hash function used by Facebook + fn furc_hash(key: &str, m: u32) -> u32 { + const MAX_TRIES: u32 = 32; + const FURC_SHIFT: u32 = 23; + const FURC_CACHE_SIZE: usize = 1024; + + if m <= 1 { + return 0; + } + + let key_bytes = key.as_bytes(); + let mut hash_cache: Vec = Vec::with_capacity(FURC_CACHE_SIZE); + let mut old_ord: i32 = -1; + + // Calculate d = ceil(log2(m)) + let mut d = 0; + while m > (1u32 << d) { + d += 1; + } + + let mut a = d; + for _try in 0..MAX_TRIES { + // Generate bits until we find a 1 bit or run out of depth + while !Self::furc_get_bit(key_bytes, a, &mut hash_cache, &mut old_ord) { + if d == 0 { + return 0; + } + d -= 1; + a = d; + } + + a += FURC_SHIFT; + let mut num = 1u32; + + // Build the number bit by bit + for _ in 0..d.saturating_sub(1) { + num = (num << 1) + | Self::furc_get_bit(key_bytes, a, &mut hash_cache, &mut old_ord) as u32; + a += FURC_SHIFT; + } + + if num < m { + return num; + } + } + + // Give up and return 0, which is always a legal value + 0 + } + + /// Bit generator function from furc_hash (from hash.c) + fn furc_get_bit(key: &[u8], idx: u32, hash_cache: &mut Vec, old_ord: &mut i32) -> bool { + const SEED: u32 = 4193360111; + + let ord = (idx >> 6) as i32; // idx / 64 + + // Extend hash cache if needed + if *old_ord < ord { + for n in (*old_ord + 1)..=ord { + let hash_value = if n == 0 { + Self::murmur_hash_64a(key, SEED) + } else { + Self::murmur_rehash_64a(hash_cache[(n - 1) as usize]) + }; + + // Ensure cache has enough capacity + if hash_cache.len() <= n as usize { + hash_cache.resize((n as usize) + 1, 0); + } + hash_cache[n as usize] = hash_value; + } + *old_ord = ord; + } + + // Extract the bit + let hash_val = hash_cache[ord as usize]; + let bit_pos = idx & 0x3f; // idx % 64 + (hash_val >> bit_pos) & 0x1 != 0 + } + + /// Facebook-style hash function using furc_hash for consistent hashing + pub fn fbi_hash(key: &str) -> u64 { + // Use furc_hash with a large modulus to get good distribution + // Then expand to u64 for our hash ring + const LARGE_MODULUS: u32 = (1u32 << 23) - 1; // Max furc_hash modulus + + let furc_result = Self::furc_hash(key, LARGE_MODULUS); + + // Expand the 23-bit furc result to 64-bit by using MurmurHash + // This gives us the full 64-bit space while preserving consistency + Self::murmur_hash_64a( + &furc_result.to_le_bytes(), + 4193360111, // Same seed as furc_hash + ) + } + + /// Update the hash ring when workers change + fn update_hash_ring(&self, workers: &[Arc]) { + let worker_urls: Vec = workers.iter().map(|w| w.url().to_string()).collect(); + + // Check if workers have changed + { + let current = self.current_workers.read().unwrap(); + if *current == worker_urls { + return; // No change needed + } + } + + // Rebuild hash ring + let mut new_ring = BTreeMap::new(); + + for worker_url in &worker_urls { + // Create virtual nodes for better distribution + for i in 0..VIRTUAL_NODES_PER_WORKER { + let virtual_key = format!("{}:{}", worker_url, i); + let hash_value = Self::fbi_hash(&virtual_key); + new_ring.insert(hash_value, worker_url.clone()); + } + } + + // Update both the ring and current workers + { + let mut ring = self.hash_ring.write().unwrap(); + *ring = new_ring; + } + { + let mut current = self.current_workers.write().unwrap(); + *current = worker_urls; + } + + info!( + "Updated consistent hash ring with {} workers and {} virtual nodes", + workers.len(), + workers.len() as u32 * VIRTUAL_NODES_PER_WORKER + ); + } + + /// Find the worker for a given hash key using consistent hashing + fn find_worker_by_hash(&self, hash_key: &str) -> Option { + let hash_value = Self::fbi_hash(hash_key); + + let ring = self.hash_ring.read().unwrap(); + if ring.is_empty() { + return None; + } + + // Find the first worker with hash >= our hash value + // If none found, wrap around to the first worker (smallest hash) + let selected_worker = ring + .range(hash_value..) + .next() + .or_else(|| ring.iter().next()) + .map(|(_, worker_url)| worker_url.clone()); + + if let Some(ref worker) = selected_worker { + debug!( + "Consistent hash: key='{}' hash={:016x} -> worker='{}'", + hash_key, hash_value, worker + ); + } + + selected_worker + } +} + +impl LoadBalancingPolicy for ConsistentHashPolicy { + fn select_worker_with_headers( + &self, + workers: &[Arc], + request_text: Option<&str>, + headers: Option<&RequestHeaders>, + ) -> Option { + let healthy_indices = get_healthy_worker_indices(workers); + + if healthy_indices.is_empty() { + return None; + } + + // Update hash ring if needed + self.update_hash_ring(workers); + + // Extract hash key with priority: headers > body > fallback + let hash_key = hash_key::extract_hash_key(request_text, headers); + + // DEBUG: Log the request text and extracted hash key + if let Some(text) = request_text { + debug!("CONSISTENT_HASH_DEBUG: Request text length: {}", text.len()); + } + if let Some(hdrs) = headers { + debug!( + "CONSISTENT_HASH_DEBUG: Headers available: {:?}", + hdrs.keys().collect::>() + ); + } + info!("CONSISTENT_HASH_DEBUG: Extracted hash key: {}", hash_key); + + // Find target worker using consistent hashing + let target_worker_url = match self.find_worker_by_hash(&hash_key) { + Some(url) => { + info!( + "CONSISTENT_HASH_DEBUG: Hash key '{}' mapped to worker: {}", + hash_key, url + ); + url + } + None => { + // Fallback to first healthy worker if hash ring is empty + let fallback_idx = healthy_indices[0]; + let worker_url = workers[fallback_idx].url(); + info!( + "CONSISTENT_HASH_DEBUG: Hash ring empty, falling back to worker: {}", + worker_url + ); + RouterMetrics::record_processed_request(worker_url); + RouterMetrics::record_policy_decision(self.name(), worker_url); + return Some(fallback_idx); + } + }; + + // Find the worker index that matches our target + let selected_idx = workers.iter().position(|w| w.url() == target_worker_url); + + debug!( + "CONSISTENT_HASH_DEBUG: Target worker URL: {}", + target_worker_url + ); + + match selected_idx { + Some(idx) => { + // Verify the worker is healthy + if workers[idx].is_healthy() && workers[idx].circuit_breaker().can_execute() { + let worker_url = workers[idx].url(); + debug!( + "CONSISTENT_HASH_DEBUG: Selected worker at index {}: {}", + idx, worker_url + ); + info!( + "Consistent hash routing: key='{}' -> worker='{}' (index={})", + hash_key, worker_url, idx + ); + + // Increment processed counter + workers[idx].increment_processed(); + RouterMetrics::record_processed_request(worker_url); + RouterMetrics::record_policy_decision(self.name(), worker_url); + + Some(idx) + } else { + // Target worker is unhealthy, fall back to first healthy worker + debug!( + "Target worker '{}' is unhealthy, falling back to healthy worker", + workers[idx].url() + ); + let fallback_idx = healthy_indices[0]; + let worker_url = workers[fallback_idx].url(); + + workers[fallback_idx].increment_processed(); + RouterMetrics::record_processed_request(worker_url); + RouterMetrics::record_policy_decision(self.name(), worker_url); + + Some(fallback_idx) + } + } + None => { + // Worker not found in current set, fall back to first healthy worker + debug!( + "Target worker '{}' not found in current worker set, falling back", + target_worker_url + ); + let fallback_idx = healthy_indices[0]; + let worker_url = workers[fallback_idx].url(); + + workers[fallback_idx].increment_processed(); + RouterMetrics::record_processed_request(worker_url); + RouterMetrics::record_policy_decision(self.name(), worker_url); + + Some(fallback_idx) + } + } + } + + fn name(&self) -> &'static str { + "consistent_hash" + } + + fn needs_request_text(&self) -> bool { + true // We need request text to extract session_id/user_id as fallback + } + + fn needs_headers(&self) -> bool { + true // We prefer HTTP headers for routing (x-session-id, etc.) + } + + fn reset(&self) { + // Clear the hash ring and force rebuild on next request + { + let mut ring = self.hash_ring.write().unwrap(); + ring.clear(); + } + { + let mut current = self.current_workers.write().unwrap(); + current.clear(); + } + info!("Consistent hash policy reset - hash ring cleared"); + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } + + fn select_worker_pair_with_headers( + &self, + prefill_workers: &[Arc], + decode_workers: &[Arc], + request_text: Option<&str>, + headers: Option<&RequestHeaders>, + ) -> Option<(usize, usize)> { + // For PD mode, use consistent hashing for both prefill and decode + // This ensures that the same user's prefill and decode requests + // go to workers that can share state efficiently + + let prefill_idx = + self.select_worker_with_headers(prefill_workers, request_text, headers)?; + let decode_idx = self.select_worker_with_headers(decode_workers, request_text, headers)?; + + Some((prefill_idx, decode_idx)) + } +} + +impl Default for ConsistentHashPolicy { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::BasicWorker; + use crate::core::WorkerType; + + #[test] + fn test_fbi_hash_consistency() { + let key = "test_session_123"; + let hash1 = ConsistentHashPolicy::fbi_hash(key); + let hash2 = ConsistentHashPolicy::fbi_hash(key); + assert_eq!(hash1, hash2, "Hash function should be deterministic"); + } + + #[test] + fn test_consistent_hash_selection() { + let policy = ConsistentHashPolicy::new(); + let workers: Vec> = vec![ + Arc::new(BasicWorker::new( + "http://worker1:8000".to_string(), + WorkerType::Regular, + )), + Arc::new(BasicWorker::new( + "http://worker2:8000".to_string(), + WorkerType::Regular, + )), + Arc::new(BasicWorker::new( + "http://worker3:8000".to_string(), + WorkerType::Regular, + )), + ]; + + // Same session should always go to same worker + let request = r#"{"session_id": "consistent_test"}"#; + let idx1 = policy.select_worker(&workers, Some(request)); + let idx2 = policy.select_worker(&workers, Some(request)); + let idx3 = policy.select_worker(&workers, Some(request)); + + assert_eq!(idx1, idx2); + assert_eq!(idx2, idx3); + assert!(idx1.is_some()); + } +} diff --git a/lmdeploy/router/src/policies/factory.rs b/lmdeploy/router/src/policies/factory.rs new file mode 100644 index 0000000000..08ff2d80db --- /dev/null +++ b/lmdeploy/router/src/policies/factory.rs @@ -0,0 +1,115 @@ +//! Factory for creating load balancing policies + +use super::{ + CacheAwareConfig, CacheAwarePolicy, ConsistentHashPolicy, LoadBalancingPolicy, + PowerOfTwoPolicy, RandomPolicy, RendezvousHashPolicy, RoundRobinPolicy, +}; +use crate::config::PolicyConfig; +use std::sync::Arc; + +/// Factory for creating policy instances +pub struct PolicyFactory; + +impl PolicyFactory { + /// Create a policy from configuration + pub fn create_from_config(config: &PolicyConfig) -> Arc { + match config { + PolicyConfig::Random => Arc::new(RandomPolicy::new()), + PolicyConfig::RoundRobin => Arc::new(RoundRobinPolicy::new()), + PolicyConfig::PowerOfTwo { .. } => Arc::new(PowerOfTwoPolicy::new()), + PolicyConfig::CacheAware { + cache_threshold, + balance_abs_threshold, + balance_rel_threshold, + eviction_interval_secs, + max_tree_size, + } => { + let config = CacheAwareConfig { + cache_threshold: *cache_threshold, + balance_abs_threshold: *balance_abs_threshold, + balance_rel_threshold: *balance_rel_threshold, + eviction_interval_secs: *eviction_interval_secs, + max_tree_size: *max_tree_size, + }; + Arc::new(CacheAwarePolicy::with_config(config)) + } + PolicyConfig::ConsistentHash { virtual_nodes: _ } => { + // Note: virtual_nodes parameter is available but not currently used + // The consistent hash policy uses a hardcoded value for now + Arc::new(ConsistentHashPolicy::new()) + } + PolicyConfig::RendezvousHash => Arc::new(RendezvousHashPolicy::new()), + } + } + + /// Create a policy by name (for dynamic loading) + pub fn create_by_name(name: &str) -> Option> { + match name.to_lowercase().as_str() { + "random" => Some(Arc::new(RandomPolicy::new())), + "round_robin" | "roundrobin" => Some(Arc::new(RoundRobinPolicy::new())), + "power_of_two" | "poweroftwo" => Some(Arc::new(PowerOfTwoPolicy::new())), + "cache_aware" | "cacheaware" => Some(Arc::new(CacheAwarePolicy::new())), + "consistent_hash" | "consistenthash" => Some(Arc::new(ConsistentHashPolicy::new())), + "rendezvous_hash" | "rendezvoushash" => Some(Arc::new(RendezvousHashPolicy::new())), + _ => None, + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_create_from_config() { + // Test Random + let policy = PolicyFactory::create_from_config(&PolicyConfig::Random); + assert_eq!(policy.name(), "random"); + + // Test RoundRobin + let policy = PolicyFactory::create_from_config(&PolicyConfig::RoundRobin); + assert_eq!(policy.name(), "round_robin"); + + // Test PowerOfTwo + let policy = PolicyFactory::create_from_config(&PolicyConfig::PowerOfTwo { + load_check_interval_secs: 60, + }); + assert_eq!(policy.name(), "power_of_two"); + + // Test CacheAware + let policy = PolicyFactory::create_from_config(&PolicyConfig::CacheAware { + cache_threshold: 0.7, + balance_abs_threshold: 10, + balance_rel_threshold: 1.5, + eviction_interval_secs: 30, + max_tree_size: 1000, + }); + assert_eq!(policy.name(), "cache_aware"); + + // Test ConsistentHash + let policy = + PolicyFactory::create_from_config(&PolicyConfig::ConsistentHash { virtual_nodes: 160 }); + assert_eq!(policy.name(), "consistent_hash"); + + // Test RendezvousHash + let policy = PolicyFactory::create_from_config(&PolicyConfig::RendezvousHash); + assert_eq!(policy.name(), "rendezvous_hash"); + } + + #[test] + fn test_create_by_name() { + assert!(PolicyFactory::create_by_name("random").is_some()); + assert!(PolicyFactory::create_by_name("RANDOM").is_some()); + assert!(PolicyFactory::create_by_name("round_robin").is_some()); + assert!(PolicyFactory::create_by_name("RoundRobin").is_some()); + assert!(PolicyFactory::create_by_name("power_of_two").is_some()); + assert!(PolicyFactory::create_by_name("PowerOfTwo").is_some()); + assert!(PolicyFactory::create_by_name("cache_aware").is_some()); + assert!(PolicyFactory::create_by_name("CacheAware").is_some()); + assert!(PolicyFactory::create_by_name("consistent_hash").is_some()); + assert!(PolicyFactory::create_by_name("ConsistentHash").is_some()); + assert!(PolicyFactory::create_by_name("rendezvous_hash").is_some()); + assert!(PolicyFactory::create_by_name("RendezvousHash").is_some()); + assert!(PolicyFactory::create_by_name("unknown").is_none()); + } +} diff --git a/lmdeploy/router/src/policies/hash_key.rs b/lmdeploy/router/src/policies/hash_key.rs new file mode 100644 index 0000000000..acfe0a8d77 --- /dev/null +++ b/lmdeploy/router/src/policies/hash_key.rs @@ -0,0 +1,311 @@ +//! Shared hash key extraction for consistent hashing policies +//! +//! Extracts routing keys from HTTP headers and request bodies. +//! Used by both ConsistentHashPolicy and RendezvousHashPolicy. + +use super::RequestHeaders; +use crate::policies::ConsistentHashPolicy; +use tracing::debug; + +/// HTTP header names to check for session ID (case-insensitive, checked in order) +pub(crate) const SESSION_HEADER_NAMES: &[&str] = &[ + "x-session-id", + "x-user-id", + "x-tenant-id", + "x-correlation-id", // per-session — check before per-request + "x-request-id", + "x-trace-id", +]; + +/// Extract hash key with priority: HTTP headers > body fields > request content hash +/// +/// Priority order: +/// 1. HTTP Headers: x-session-id, x-user-id, x-tenant-id, x-correlation-id, x-request-id, x-trace-id +/// 2. Body: user field (OpenAI format) +/// 3. Body: session_id (LMDeploy /generate format) +/// 4. Body: user_id (legacy) +/// 5. Fallback: hash of request body (long) or raw text (short) +pub(crate) fn extract_hash_key( + request_text: Option<&str>, + headers: Option<&RequestHeaders>, +) -> String { + // 1. First priority: HTTP headers + if let Some(hdrs) = headers { + if let Some(key) = extract_hash_key_from_headers(hdrs) { + return key; + } + } + + // 2. Second priority: Body fields + if let Some(key) = extract_hash_key_from_body(request_text) { + return key; + } + + // 3. Final fallback: hash of request body + let text = request_text.unwrap_or(""); + if text.len() > 100 { + format!("request_hash:{:016x}", ConsistentHashPolicy::fbi_hash(text)) + } else { + format!("request:{}", text) + } +} + +/// Extract hash key from HTTP headers +pub(crate) fn extract_hash_key_from_headers(headers: &RequestHeaders) -> Option { + for header_name in SESSION_HEADER_NAMES { + if let Some(value) = headers.get(*header_name) { + if !value.is_empty() { + debug!( + "Hash key extraction: found session key in header '{}': {}", + header_name, value + ); + return Some(format!("header:{}:{}", header_name, value)); + } + } + } + None +} + +/// Extract hash key from request body fields +/// +/// Priority: user > session_id > user_id +pub(crate) fn extract_hash_key_from_body(request_text: Option<&str>) -> Option { + let text = request_text.unwrap_or(""); + if text.is_empty() { + return None; + } + + // 1. Try to extract direct user field (from OpenAI ChatCompletion/Completion requests) + if let Some(user) = extract_field_value(text, "user") { + debug!("Hash key extraction: found user field: {}", user); + return Some(format!("user:{}", user)); + } + + // 2. Extract the LMDeploy top-level session_id field + if let Some(session_id) = extract_field_value(text, "session_id") { + return Some(format!("session:{}", session_id)); + } + + // 3. Fallback: try legacy user_id field + if let Some(user_id) = extract_field_value(text, "user_id") { + return Some(format!("user:{}", user_id)); + } + + None +} + +/// Extract field value from JSON-like text (simple parser) +/// +/// Supports double-quoted, single-quoted, and unquoted values. +pub(crate) fn extract_field_value(text: &str, field_name: &str) -> Option { + let patterns = [ + format!("\"{}\"", field_name), + format!("'{}'", field_name), + field_name.to_string(), + ]; + + for pattern in &patterns { + if let Some(field_pos) = text.find(pattern) { + let after_field = &text[field_pos + pattern.len()..]; + + // Skip whitespace and look for colon + let mut colon_pos = None; + for (i, ch) in after_field.char_indices() { + if ch == ':' { + colon_pos = Some(i); + break; + } else if !ch.is_whitespace() { + break; + } + } + + if let Some(colon_idx) = colon_pos { + let after_colon = &after_field[colon_idx + 1..]; + let trimmed = after_colon.trim_start(); + + // Extract quoted string (double or single quotes) + if trimmed.starts_with('"') { + if let Some(stripped) = trimmed.strip_prefix('"') { + if let Some(end_quote) = stripped.find('"') { + return Some(stripped[..end_quote].to_string()); + } + } + } else if trimmed.starts_with('\'') { + if let Some(stripped) = trimmed.strip_prefix('\'') { + if let Some(end_quote) = stripped.find('\'') { + return Some(stripped[..end_quote].to_string()); + } + } + } else { + // Unquoted value - extract until delimiter + let end_pos = trimmed + .find(&[',', ' ', '}', ']', '\n', '\r', '\t'][..]) + .unwrap_or(trimmed.len()); + if end_pos > 0 { + return Some(trimmed[..end_pos].to_string()); + } + } + } + } + } + + None +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashMap; + + // === extract_field_value tests === + + #[test] + fn test_extract_field_value_double_quoted() { + let text = r#"{"session_id": "abc123", "prompt": "hello"}"#; + assert_eq!( + extract_field_value(text, "session_id"), + Some("abc123".to_string()) + ); + } + + #[test] + fn test_extract_field_value_single_quoted() { + let text = r#"{'session_id': 'def456', 'prompt': 'world'}"#; + assert_eq!( + extract_field_value(text, "session_id"), + Some("def456".to_string()) + ); + } + + #[test] + fn test_extract_field_value_unquoted() { + let text = r#"{"count": 42, "name": "test"}"#; + assert_eq!(extract_field_value(text, "count"), Some("42".to_string())); + } + + #[test] + fn test_extract_field_value_missing() { + let text = r#"{"other": "val"}"#; + assert_eq!(extract_field_value(text, "session_id"), None); + } + + #[test] + fn test_extract_field_value_no_space_after_colon() { + let text = r#"{"session_id":"compact_value"}"#; + assert_eq!( + extract_field_value(text, "session_id"), + Some("compact_value".to_string()) + ); + } + + #[test] + fn test_extract_field_value_multiple_fields() { + let text = r#"{"user": "bob", "prompt": "hi", "session_id": "sess1"}"#; + assert_eq!(extract_field_value(text, "user"), Some("bob".to_string())); + assert_eq!( + extract_field_value(text, "session_id"), + Some("sess1".to_string()) + ); + } + + // === extract_hash_key_from_headers tests === + + #[test] + fn test_header_extraction_priority() { + let mut headers = HashMap::new(); + headers.insert("x-request-id".to_string(), "req-1".to_string()); + headers.insert("x-session-id".to_string(), "sess-1".to_string()); + + // x-session-id has higher priority than x-request-id + let key = extract_hash_key_from_headers(&headers).unwrap(); + assert_eq!(key, "header:x-session-id:sess-1"); + } + + #[test] + fn test_header_extraction_skips_empty() { + let mut headers = HashMap::new(); + headers.insert("x-session-id".to_string(), "".to_string()); + headers.insert("x-user-id".to_string(), "user-1".to_string()); + + let key = extract_hash_key_from_headers(&headers).unwrap(); + assert_eq!(key, "header:x-user-id:user-1"); + } + + #[test] + fn test_header_extraction_no_match() { + let mut headers = HashMap::new(); + headers.insert("x-custom-header".to_string(), "val".to_string()); + assert_eq!(extract_hash_key_from_headers(&headers), None); + } + + // === extract_hash_key_from_body tests === + + #[test] + fn test_body_extraction_user_field() { + let text = r#"{"user": "alice", "prompt": "hi"}"#; + let key = extract_hash_key_from_body(Some(text)).unwrap(); + assert_eq!(key, "user:alice"); + } + + #[test] + fn test_body_extraction_session_id() { + let text = r#"{"session_id": "session123", "prompt": "hi"}"#; + let key = extract_hash_key_from_body(Some(text)).unwrap(); + assert_eq!(key, "session:session123"); + } + + #[test] + fn test_body_extraction_legacy_user_id() { + let text = r#"{"user_id": "uid456", "prompt": "hi"}"#; + let key = extract_hash_key_from_body(Some(text)).unwrap(); + assert_eq!(key, "user:uid456"); + } + + #[test] + fn test_body_extraction_empty() { + assert_eq!(extract_hash_key_from_body(None), None); + assert_eq!(extract_hash_key_from_body(Some("")), None); + } + + #[test] + fn test_body_extraction_no_known_fields() { + let text = r#"{"prompt": "hello", "model": "llama"}"#; + assert_eq!(extract_hash_key_from_body(Some(text)), None); + } + + // === extract_hash_key (top-level) tests === + + #[test] + fn test_hash_key_headers_over_body() { + let mut headers = HashMap::new(); + headers.insert("x-session-id".to_string(), "from-header".to_string()); + let body = r#"{"session_id": "from-body"}"#; + + let key = extract_hash_key(Some(body), Some(&headers)); + assert_eq!(key, "header:x-session-id:from-header"); + } + + #[test] + fn test_hash_key_fallback_short_text() { + let key = extract_hash_key(Some("short"), None); + assert_eq!(key, "request:short"); + } + + #[test] + fn test_hash_key_fallback_long_text() { + let long_text = "x".repeat(200); + let key = extract_hash_key(Some(&long_text), None); + assert!(key.starts_with("request_hash:")); + assert_eq!(key.len(), "request_hash:".len() + 16); // 16 hex chars + + // Same long text should produce same hash + let key2 = extract_hash_key(Some(&long_text), None); + assert_eq!(key, key2); + } + + #[test] + fn test_hash_key_fallback_none() { + let key = extract_hash_key(None, None); + assert_eq!(key, "request:"); + } +} diff --git a/lmdeploy/router/src/policies/mod.rs b/lmdeploy/router/src/policies/mod.rs new file mode 100644 index 0000000000..1148f57f82 --- /dev/null +++ b/lmdeploy/router/src/policies/mod.rs @@ -0,0 +1,219 @@ +//! Load balancing policies for the LMDeploy router +//! +//! This module provides a unified abstraction for routing policies that work +//! across both regular and prefill-decode (PD) routing modes. + +use crate::core::Worker; +use std::collections::HashMap; +use std::fmt::Debug; +use std::sync::Arc; + +mod cache_aware; +mod consistent_hash; +mod factory; +mod hash_key; +mod power_of_two; +mod random; +mod registry; +mod rendezvous_hash; +mod round_robin; + +pub use cache_aware::CacheAwarePolicy; +pub use consistent_hash::ConsistentHashPolicy; +pub use consistent_hash::VIRTUAL_NODES_PER_WORKER; +pub use factory::PolicyFactory; +pub use power_of_two::PowerOfTwoPolicy; +pub use random::RandomPolicy; +pub use registry::PolicyRegistry; +pub use rendezvous_hash::RendezvousHashPolicy; +pub use round_robin::RoundRobinPolicy; + +/// HTTP headers passed to policies for routing decisions +/// Key is lowercase header name, value is header value +pub type RequestHeaders = HashMap; + +/// Core trait for load balancing policies +/// +/// This trait provides a unified interface for implementing routing algorithms +/// that can work with both regular single-worker selection and PD dual-worker selection. +pub trait LoadBalancingPolicy: Send + Sync + Debug { + /// Select a single worker from the available workers + /// + /// This is used for regular routing mode where requests go to a single worker. + /// Now uses Arc for better performance and to avoid unnecessary cloning. + fn select_worker( + &self, + workers: &[Arc], + request_text: Option<&str>, + ) -> Option { + // Default implementation delegates to select_worker_with_headers without headers + self.select_worker_with_headers(workers, request_text, None) + } + + /// Select a single worker with optional HTTP headers for routing decisions + /// + /// Policies like consistent_hash can use headers (e.g., X-Session-ID) for routing. + /// Default implementation ignores headers and falls back to basic selection. + fn select_worker_with_headers( + &self, + workers: &[Arc], + request_text: Option<&str>, + headers: Option<&RequestHeaders>, + ) -> Option; + + /// Select a pair of workers (prefill and decode) for PD routing + /// + /// Returns indices of (prefill_worker, decode_worker) from their respective arrays. + /// Default implementation uses select_worker for each array independently. + fn select_worker_pair( + &self, + prefill_workers: &[Arc], + decode_workers: &[Arc], + request_text: Option<&str>, + ) -> Option<(usize, usize)> { + self.select_worker_pair_with_headers(prefill_workers, decode_workers, request_text, None) + } + + /// Select a pair of workers with optional HTTP headers for routing decisions + fn select_worker_pair_with_headers( + &self, + prefill_workers: &[Arc], + decode_workers: &[Arc], + request_text: Option<&str>, + headers: Option<&RequestHeaders>, + ) -> Option<(usize, usize)> { + // Default implementation: independently select from each pool + let prefill_idx = + self.select_worker_with_headers(prefill_workers, request_text, headers)?; + let decode_idx = self.select_worker_with_headers(decode_workers, request_text, headers)?; + Some((prefill_idx, decode_idx)) + } + + /// Update policy state after request completion + /// + /// This is called when a request completes (successfully or not) to allow + /// policies to update their internal state. + fn on_request_complete(&self, _worker_url: &str, _success: bool) { + // Default: no-op for stateless policies + } + + /// Get policy name for metrics and debugging + fn name(&self) -> &'static str; + + /// Check if this policy needs request text for routing decisions + fn needs_request_text(&self) -> bool { + false // Default: most policies don't need request text + } + + /// Check if this policy needs HTTP headers for routing decisions + fn needs_headers(&self) -> bool { + false // Default: most policies don't need headers + } + + /// Update worker load information + /// + /// This is called periodically with current load information for load-aware policies. + fn update_loads(&self, _loads: &std::collections::HashMap) { + // Default: no-op for policies that don't use load information + } + + /// Reset any internal state + /// + /// This is useful for policies that maintain state (e.g., round-robin counters). + fn reset(&self) { + // Default: no-op for stateless policies + } + + /// Get as Any for downcasting + fn as_any(&self) -> &dyn std::any::Any; + + /// Returns true if this policy requires init_workers() to be called before use. + /// Override to return true for stateful policies like cache-aware routing. + fn requires_initialization(&self) -> bool { + false // Default: most policies don't need initialization + } + + /// Initialize the policy with a set of workers. + /// Override for stateful policies that need to set up internal data structures. + /// Default is no-op for stateless policies like round-robin. + fn init_workers(&self, _workers: &[Arc]) { + // Default: no-op for policies that don't need initialization + } +} + +/// Configuration for cache-aware policy +#[derive(Debug, Clone)] +pub struct CacheAwareConfig { + pub cache_threshold: f32, + pub balance_abs_threshold: usize, + pub balance_rel_threshold: f32, + pub eviction_interval_secs: u64, + pub max_tree_size: usize, +} + +impl Default for CacheAwareConfig { + fn default() -> Self { + Self { + cache_threshold: 0.5, + balance_abs_threshold: 32, + balance_rel_threshold: 1.1, + eviction_interval_secs: 30, + max_tree_size: 10000, + } + } +} + +/// Helper function to filter healthy workers and return their indices +pub(crate) fn get_healthy_worker_indices(workers: &[Arc]) -> Vec { + workers + .iter() + .enumerate() + .filter(|(_, w)| w.is_healthy() && w.circuit_breaker().can_execute()) + .map(|(idx, _)| idx) + .collect() +} + +/// Helper function to normalize model_id to a key for policy lookups. +/// +/// Use "default" for unknown/empty model_ids for backward compatibility +#[inline] +pub(crate) fn normalize_model_key(model_id: &str) -> &str { + if model_id.is_empty() || model_id == "unknown" { + "default" + } else { + model_id + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::{BasicWorker, WorkerType}; + + #[test] + fn test_get_healthy_worker_indices() { + let workers: Vec> = vec![ + Arc::new(BasicWorker::new( + "http://w1:8000".to_string(), + WorkerType::Regular, + )), + Arc::new(BasicWorker::new( + "http://w2:8000".to_string(), + WorkerType::Regular, + )), + Arc::new(BasicWorker::new( + "http://w3:8000".to_string(), + WorkerType::Regular, + )), + ]; + + // All healthy initially + let indices = get_healthy_worker_indices(&workers); + assert_eq!(indices, vec![0, 1, 2]); + + // Mark one unhealthy + workers[1].set_healthy(false); + let indices = get_healthy_worker_indices(&workers); + assert_eq!(indices, vec![0, 2]); + } +} diff --git a/lmdeploy/router/src/policies/power_of_two.rs b/lmdeploy/router/src/policies/power_of_two.rs new file mode 100644 index 0000000000..4d06a0df96 --- /dev/null +++ b/lmdeploy/router/src/policies/power_of_two.rs @@ -0,0 +1,202 @@ +//! Power-of-two choices load balancing policy + +use super::{get_healthy_worker_indices, LoadBalancingPolicy, RequestHeaders}; +use crate::core::Worker; +use crate::metrics::RouterMetrics; +use rand::Rng; +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; +use tracing::info; + +/// Power-of-two choices policy +/// +/// Randomly selects two workers and routes to the one with lower load. +/// This provides good load distribution with minimal coordination overhead. +#[derive(Debug)] +pub struct PowerOfTwoPolicy { + /// Cached load information from external monitoring + cached_loads: RwLock>, +} + +impl PowerOfTwoPolicy { + pub fn new() -> Self { + Self { + cached_loads: RwLock::new(HashMap::new()), + } + } + + fn get_worker_load(&self, worker: &dyn Worker) -> isize { + // First check cached loads (from external monitoring) + if let Ok(loads) = self.cached_loads.read() { + if let Some(&load) = loads.get(worker.url()) { + return load; + } + } + + // Fall back to local load counter + worker.load() as isize + } +} + +impl LoadBalancingPolicy for PowerOfTwoPolicy { + fn select_worker_with_headers( + &self, + workers: &[Arc], + _request_text: Option<&str>, + _headers: Option<&RequestHeaders>, + ) -> Option { + let healthy_indices = get_healthy_worker_indices(workers); + + if healthy_indices.is_empty() { + return None; + } + + if healthy_indices.len() == 1 { + return Some(healthy_indices[0]); + } + + // Select two random workers + let mut rng = rand::rng(); + let idx1 = rng.random_range(0..healthy_indices.len()); + let mut idx2 = rng.random_range(0..healthy_indices.len()); + + // Ensure we pick two different workers + while idx2 == idx1 { + idx2 = rng.random_range(0..healthy_indices.len()); + } + + let worker_idx1 = healthy_indices[idx1]; + let worker_idx2 = healthy_indices[idx2]; + + // Compare loads and select the less loaded one + let load1 = self.get_worker_load(workers[worker_idx1].as_ref()); + let load2 = self.get_worker_load(workers[worker_idx2].as_ref()); + + // Log selection for debugging + let selected_idx = if load1 <= load2 { + worker_idx1 + } else { + worker_idx2 + }; + + info!( + "Power-of-two selection: {}={} vs {}={} -> selected {}", + workers[worker_idx1].url(), + load1, + workers[worker_idx2].url(), + load2, + workers[selected_idx].url() + ); + + // Increment processed counter + workers[selected_idx].increment_processed(); + RouterMetrics::record_processed_request(workers[selected_idx].url()); + RouterMetrics::record_policy_decision(self.name(), workers[selected_idx].url()); + + Some(selected_idx) + } + + fn name(&self) -> &'static str { + "power_of_two" + } + + fn update_loads(&self, loads: &HashMap) { + if let Ok(mut cached) = self.cached_loads.write() { + *cached = loads.clone(); + } + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +impl Default for PowerOfTwoPolicy { + fn default() -> Self { + Self::new() + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::{BasicWorker, WorkerType}; + + #[test] + fn test_power_of_two_selection() { + let policy = PowerOfTwoPolicy::new(); + let worker1 = BasicWorker::new("http://w1:8000".to_string(), WorkerType::Regular); + let worker2 = BasicWorker::new("http://w2:8000".to_string(), WorkerType::Regular); + let worker3 = BasicWorker::new("http://w3:8000".to_string(), WorkerType::Regular); + + // Set different loads + for _ in 0..10 { + worker1.increment_load(); + } + for _ in 0..5 { + worker2.increment_load(); + } + // worker3 has load 0 + + let workers: Vec> = + vec![Arc::new(worker1), Arc::new(worker2), Arc::new(worker3)]; + + // Run multiple selections + let mut selected_counts = [0; 3]; + for _ in 0..100 { + if let Some(idx) = policy.select_worker(&workers, None) { + selected_counts[idx] += 1; + } + } + + // Worker with lowest load (worker3) should be selected most often + assert!(selected_counts[2] > selected_counts[1]); + assert!(selected_counts[1] > selected_counts[0]); + } + + #[test] + fn test_power_of_two_with_cached_loads() { + let policy = PowerOfTwoPolicy::new(); + let workers: Vec> = vec![ + Arc::new(BasicWorker::new( + "http://w1:8000".to_string(), + WorkerType::Regular, + )), + Arc::new(BasicWorker::new( + "http://w2:8000".to_string(), + WorkerType::Regular, + )), + ]; + + // Update cached loads + let mut loads = HashMap::new(); + loads.insert("http://w1:8000".to_string(), 100); + loads.insert("http://w2:8000".to_string(), 10); + policy.update_loads(&loads); + + // Should prefer worker2 with lower cached load + let mut w2_selected = 0; + for _ in 0..50 { + if let Some(idx) = policy.select_worker(&workers, None) { + if idx == 1 { + w2_selected += 1; + } + } + } + + // Worker2 should be selected significantly more often + assert!(w2_selected > 35); // Should win most of the time + } + + #[test] + fn test_power_of_two_single_worker() { + let policy = PowerOfTwoPolicy::new(); + let workers: Vec> = vec![Arc::new(BasicWorker::new( + "http://w1:8000".to_string(), + WorkerType::Regular, + ))]; + + // With single worker, should always select it + assert_eq!(policy.select_worker(&workers, None), Some(0)); + } +} diff --git a/lmdeploy/router/src/policies/random.rs b/lmdeploy/router/src/policies/random.rs new file mode 100644 index 0000000000..e38089d071 --- /dev/null +++ b/lmdeploy/router/src/policies/random.rs @@ -0,0 +1,123 @@ +//! Random load balancing policy + +use super::{get_healthy_worker_indices, LoadBalancingPolicy, RequestHeaders}; +use crate::core::Worker; +use crate::metrics::RouterMetrics; +use rand::Rng; +use std::sync::Arc; + +/// Random selection policy +/// +/// Selects workers randomly with uniform distribution among healthy workers. +#[derive(Debug, Default)] +pub struct RandomPolicy; + +impl RandomPolicy { + pub fn new() -> Self { + Self + } +} + +impl LoadBalancingPolicy for RandomPolicy { + fn select_worker_with_headers( + &self, + workers: &[Arc], + _request_text: Option<&str>, + _headers: Option<&RequestHeaders>, + ) -> Option { + let healthy_indices = get_healthy_worker_indices(workers); + + if healthy_indices.is_empty() { + return None; + } + + let mut rng = rand::rng(); + let random_idx = rng.random_range(0..healthy_indices.len()); + let worker = workers[healthy_indices[random_idx]].url(); + + RouterMetrics::record_processed_request(worker); + RouterMetrics::record_policy_decision(self.name(), worker); + Some(healthy_indices[random_idx]) + } + + fn name(&self) -> &'static str { + "random" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::{BasicWorker, WorkerType}; + use std::collections::HashMap; + + #[test] + fn test_random_selection() { + let policy = RandomPolicy::new(); + let workers: Vec> = vec![ + Arc::new(BasicWorker::new( + "http://w1:8000".to_string(), + WorkerType::Regular, + )), + Arc::new(BasicWorker::new( + "http://w2:8000".to_string(), + WorkerType::Regular, + )), + Arc::new(BasicWorker::new( + "http://w3:8000".to_string(), + WorkerType::Regular, + )), + ]; + + // Test multiple selections to ensure randomness + let mut counts = HashMap::new(); + for _ in 0..100 { + if let Some(idx) = policy.select_worker(&workers, None) { + *counts.entry(idx).or_insert(0) += 1; + } + } + + // All workers should be selected at least once + assert_eq!(counts.len(), 3); + assert!(counts.values().all(|&count| count > 0)); + } + + #[test] + fn test_random_with_unhealthy_workers() { + let policy = RandomPolicy::new(); + let workers: Vec> = vec![ + Arc::new(BasicWorker::new( + "http://w1:8000".to_string(), + WorkerType::Regular, + )), + Arc::new(BasicWorker::new( + "http://w2:8000".to_string(), + WorkerType::Regular, + )), + ]; + + // Mark first worker as unhealthy + workers[0].set_healthy(false); + + // Should always select the healthy worker (index 1) + for _ in 0..10 { + assert_eq!(policy.select_worker(&workers, None), Some(1)); + } + } + + #[test] + fn test_random_no_healthy_workers() { + let policy = RandomPolicy::new(); + let workers: Vec> = vec![Arc::new(BasicWorker::new( + "http://w1:8000".to_string(), + WorkerType::Regular, + ))]; + + workers[0].set_healthy(false); + assert_eq!(policy.select_worker(&workers, None), None); + } +} diff --git a/lmdeploy/router/src/policies/registry.rs b/lmdeploy/router/src/policies/registry.rs new file mode 100644 index 0000000000..fb75901153 --- /dev/null +++ b/lmdeploy/router/src/policies/registry.rs @@ -0,0 +1,336 @@ +/// Policy Registry for managing model-to-policy mappings +/// +/// This registry manages the dynamic assignment of load balancing policies to models. +/// When the first worker of a new model is added, it determines the policy for that model. +/// All subsequent workers of the same model use the established policy. +/// When the last worker of a model is removed, the policy mapping is cleaned up. +use super::{ + CacheAwareConfig, CacheAwarePolicy, ConsistentHashPolicy, LoadBalancingPolicy, + PowerOfTwoPolicy, RandomPolicy, RendezvousHashPolicy, RoundRobinPolicy, +}; +use crate::config::types::PolicyConfig; +use std::collections::HashMap; +use std::sync::{Arc, RwLock}; +use tracing::{debug, info, warn}; + +/// Registry for managing model-to-policy mappings +#[derive(Clone)] +pub struct PolicyRegistry { + /// Model ID -> Policy instance mapping + model_policies: Arc>>>, + + /// Model ID -> Worker count for cleanup tracking + model_worker_counts: Arc>>, + + /// Default policy instance (cached) + default_policy: Arc, + + /// Prefill policy for PD mode + prefill_policy: Arc>>>, + + /// Decode policy for PD mode + decode_policy: Arc>>>, +} + +impl PolicyRegistry { + /// Create a new PolicyRegistry with a default policy + pub fn new(default_policy_config: PolicyConfig) -> Self { + let default_policy = Self::create_policy_from_config(&default_policy_config); + + Self { + model_policies: Arc::new(RwLock::new(HashMap::new())), + model_worker_counts: Arc::new(RwLock::new(HashMap::new())), + default_policy, + prefill_policy: Arc::new(RwLock::new(None)), + decode_policy: Arc::new(RwLock::new(None)), + } + } + + /// Called when a worker is added + /// Returns the policy that should be used for this worker's model + pub fn on_worker_added( + &self, + model_id: &str, + policy_hint: Option<&str>, + ) -> Arc { + // Increment worker count + { + let mut counts = self.model_worker_counts.write().unwrap(); + *counts.entry(model_id.to_string()).or_insert(0) += 1; + debug!( + "Worker added for model {}, count: {}", + model_id, + counts.get(model_id).unwrap() + ); + } + + // Check if model already has a policy + { + let policies = self.model_policies.read().unwrap(); + if let Some(existing_policy) = policies.get(model_id) { + debug!( + "Model {} already has policy: {}", + model_id, + existing_policy.name() + ); + return Arc::clone(existing_policy); + } + } + + // New model - determine policy + let policy = self.determine_policy_for_model(model_id, policy_hint); + + info!( + "Assigning policy {} to new model {}", + policy.name(), + model_id + ); + + // Store policy for this model + { + let mut policies = self.model_policies.write().unwrap(); + policies.insert(model_id.to_string(), Arc::clone(&policy)); + } + + policy + } + + /// Called when a worker is removed + pub fn on_worker_removed(&self, model_id: &str) { + let should_cleanup = { + let mut counts = self.model_worker_counts.write().unwrap(); + if let Some(count) = counts.get_mut(model_id) { + *count = count.saturating_sub(1); + debug!("Worker removed for model {}, count: {}", model_id, *count); + if *count == 0 { + counts.remove(model_id); + true + } else { + false + } + } else { + warn!( + "Attempted to remove worker for model {} with no registered workers", + model_id + ); + false + } + }; + + // Clean up policy if this was the last worker + if should_cleanup { + let mut policies = self.model_policies.write().unwrap(); + if let Some(policy) = policies.remove(model_id) { + info!( + "Removed policy {} for model {} (last worker removed)", + policy.name(), + model_id + ); + // Policy will be dropped here, cleaning up any resources + drop(policy); + } + } + } + + /// Get the policy for a model + pub fn get_policy(&self, model_id: &str) -> Option> { + self.model_policies.read().unwrap().get(model_id).cloned() + } + + /// Get the default policy + pub fn get_default_policy(&self) -> Arc { + Arc::clone(&self.default_policy) + } + + /// Get policy for a model, or default if not found + pub fn get_policy_or_default(&self, model_id: &str) -> Arc { + self.get_policy(model_id) + .unwrap_or_else(|| self.get_default_policy()) + } + + /// Determine policy for a new model + fn determine_policy_for_model( + &self, + model_id: &str, + policy_hint: Option<&str>, + ) -> Arc { + // 1. Check policy hint from worker + if let Some(policy_type) = policy_hint { + debug!("Using policy hint '{}' for model {}", policy_type, model_id); + return self.create_policy_from_type(policy_type); + } + + // 2. Use default policy + debug!("Using default policy for model {}", model_id); + Arc::clone(&self.default_policy) + } + + /// Create a policy from a type string + fn create_policy_from_type(&self, policy_type: &str) -> Arc { + match policy_type { + "round_robin" => Arc::new(RoundRobinPolicy::new()), + "random" => Arc::new(RandomPolicy::new()), + "cache_aware" => Arc::new(CacheAwarePolicy::new()), + "power_of_two" => Arc::new(PowerOfTwoPolicy::new()), + "rendezvous_hash" => Arc::new(RendezvousHashPolicy::new()), + _ => { + warn!("Unknown policy type '{}', using default", policy_type); + Arc::clone(&self.default_policy) + } + } + } + + /// Create a policy from a PolicyConfig + fn create_policy_from_config(config: &PolicyConfig) -> Arc { + match config { + PolicyConfig::RoundRobin => Arc::new(RoundRobinPolicy::new()), + PolicyConfig::Random => Arc::new(RandomPolicy::new()), + PolicyConfig::CacheAware { + cache_threshold, + balance_abs_threshold, + balance_rel_threshold, + eviction_interval_secs, + max_tree_size, + } => { + let cache_config = CacheAwareConfig { + cache_threshold: *cache_threshold, + balance_abs_threshold: *balance_abs_threshold, + balance_rel_threshold: *balance_rel_threshold, + eviction_interval_secs: *eviction_interval_secs, + max_tree_size: *max_tree_size, + }; + Arc::new(CacheAwarePolicy::with_config(cache_config)) + } + PolicyConfig::PowerOfTwo { .. } => Arc::new(PowerOfTwoPolicy::new()), + PolicyConfig::ConsistentHash { .. } => Arc::new(ConsistentHashPolicy::new()), + PolicyConfig::RendezvousHash => Arc::new(RendezvousHashPolicy::new()), + } + } + + /// Get current model->policy mappings (for debugging/monitoring) + pub fn get_all_mappings(&self) -> HashMap { + let policies = self.model_policies.read().unwrap(); + policies + .iter() + .map(|(model, policy)| (model.clone(), policy.name().to_string())) + .collect() + } + + /// Get worker counts per model + pub fn get_worker_counts(&self) -> HashMap { + self.model_worker_counts.read().unwrap().clone() + } + + /// Clear all policies (useful for testing) + pub fn clear(&self) { + let mut policies = self.model_policies.write().unwrap(); + policies.clear(); + let mut counts = self.model_worker_counts.write().unwrap(); + counts.clear(); + } + + /// Set the prefill policy for PD mode + pub fn set_prefill_policy(&self, policy: Arc) { + let mut prefill_policy = self.prefill_policy.write().unwrap(); + *prefill_policy = Some(policy); + } + + /// Set the decode policy for PD mode + pub fn set_decode_policy(&self, policy: Arc) { + let mut decode_policy = self.decode_policy.write().unwrap(); + *decode_policy = Some(policy); + } + + /// Get the prefill policy for PD mode, or default if not set + pub fn get_prefill_policy(&self) -> Arc { + let prefill_policy = self.prefill_policy.read().unwrap(); + prefill_policy + .as_ref() + .map(Arc::clone) + .unwrap_or_else(|| self.get_default_policy()) + } + + /// Get the decode policy for PD mode, or default if not set + pub fn get_decode_policy(&self) -> Arc { + let decode_policy = self.decode_policy.read().unwrap(); + decode_policy + .as_ref() + .map(Arc::clone) + .unwrap_or_else(|| self.get_default_policy()) + } +} + +impl std::fmt::Debug for PolicyRegistry { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("PolicyRegistry") + .field("model_policies", &self.model_policies) + .field("model_worker_counts", &self.model_worker_counts) + .field("default_policy", &self.default_policy.name()) + .finish() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_policy_registry_basic() { + let registry = PolicyRegistry::new(PolicyConfig::RoundRobin); + + // First worker of a model sets the policy + let policy1 = registry.on_worker_added("llama-3", Some("cache_aware")); + assert_eq!(policy1.name(), "cache_aware"); + + // Second worker of same model uses existing policy + let policy2 = registry.on_worker_added("llama-3", Some("round_robin")); + assert_eq!(policy2.name(), "cache_aware"); // Ignores hint, uses existing + + // Different model can have different policy + let policy3 = registry.on_worker_added("gpt-4", Some("random")); + assert_eq!(policy3.name(), "random"); + + // Check mappings + let mappings = registry.get_all_mappings(); + assert_eq!(mappings.get("llama-3").unwrap(), "cache_aware"); + assert_eq!(mappings.get("gpt-4").unwrap(), "random"); + + // Check worker counts + let counts = registry.get_worker_counts(); + assert_eq!(*counts.get("llama-3").unwrap(), 2); + assert_eq!(*counts.get("gpt-4").unwrap(), 1); + } + + #[test] + fn test_policy_registry_cleanup() { + let registry = PolicyRegistry::new(PolicyConfig::RoundRobin); + + // Add workers + registry.on_worker_added("llama-3", Some("cache_aware")); + registry.on_worker_added("llama-3", None); + assert_eq!(registry.get_worker_counts().get("llama-3"), Some(&2)); + + // Remove one worker - policy should remain + registry.on_worker_removed("llama-3"); + assert!(registry.get_policy("llama-3").is_some()); + assert_eq!(registry.get_worker_counts().get("llama-3"), Some(&1)); + + // Remove last worker - policy should be cleaned up + registry.on_worker_removed("llama-3"); + assert!(registry.get_policy("llama-3").is_none()); + assert_eq!(registry.get_worker_counts().get("llama-3"), None); + } + + #[test] + fn test_default_policy() { + let registry = PolicyRegistry::new(PolicyConfig::RoundRobin); + + // No hint, no template - uses default + let policy = registry.on_worker_added("unknown-model", None); + assert_eq!(policy.name(), "round_robin"); + + // Get default directly + let default = registry.get_default_policy(); + assert_eq!(default.name(), "round_robin"); + } +} diff --git a/lmdeploy/router/src/policies/rendezvous_hash.rs b/lmdeploy/router/src/policies/rendezvous_hash.rs new file mode 100644 index 0000000000..e6599aee00 --- /dev/null +++ b/lmdeploy/router/src/policies/rendezvous_hash.rs @@ -0,0 +1,229 @@ +//! Rendezvous hashing (Highest Random Weight) load balancing policy +//! +//! Routes requests to workers by scoring each worker with a hash of the +//! session key combined with the worker URL, then picking the highest score. +//! This provides better balance than ring-based consistent hashing with +//! zero tuning knobs, while preserving session stickiness and minimal +//! redistribution (~1/n) on worker add/remove. + +use std::sync::Arc; + +use tracing::info; + +use super::get_healthy_worker_indices; +use super::hash_key; +use super::LoadBalancingPolicy; +use super::RequestHeaders; +use crate::core::Worker; +use crate::metrics::RouterMetrics; +use crate::policies::ConsistentHashPolicy; + +/// Rendezvous hashing (Highest Random Weight) policy +/// +/// For each request, scores every healthy worker with +/// `fbi_hash("{hash_key}:{worker_url}")` and picks the one with the +/// highest score. This is deterministic — the same session always maps +/// to the same worker — and gives near-perfect balance across workers. +#[derive(Debug)] +pub struct RendezvousHashPolicy; + +impl RendezvousHashPolicy { + pub fn new() -> Self { + Self + } +} + +impl Default for RendezvousHashPolicy { + fn default() -> Self { + Self::new() + } +} + +impl LoadBalancingPolicy for RendezvousHashPolicy { + fn select_worker_with_headers( + &self, + workers: &[Arc], + request_text: Option<&str>, + headers: Option<&RequestHeaders>, + ) -> Option { + let healthy_indices = get_healthy_worker_indices(workers); + + if healthy_indices.is_empty() { + return None; + } + + let hash_key = hash_key::extract_hash_key(request_text, headers); + + // Score each healthy worker and pick the highest + let selected_idx = healthy_indices + .iter() + .max_by_key(|&&idx| { + let score_key = format!("{}:{}", hash_key, workers[idx].url()); + ConsistentHashPolicy::fbi_hash(&score_key) + }) + .copied() + .unwrap(); + + let worker_url = workers[selected_idx].url(); + info!( + "Rendezvous hash routing: key='{}' -> worker='{}' (index={})", + hash_key, worker_url, selected_idx + ); + + workers[selected_idx].increment_processed(); + RouterMetrics::record_processed_request(worker_url); + RouterMetrics::record_policy_decision(self.name(), worker_url); + + Some(selected_idx) + } + + fn name(&self) -> &'static str { + "rendezvous_hash" + } + + fn needs_request_text(&self) -> bool { + true + } + + fn needs_headers(&self) -> bool { + true + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::BasicWorker; + use crate::core::WorkerType; + use std::collections::HashSet; + + fn make_workers(urls: &[&str]) -> Vec> { + urls.iter() + .map(|url| { + Arc::new(BasicWorker::new(url.to_string(), WorkerType::Regular)) as Arc + }) + .collect() + } + + #[test] + fn test_rendezvous_hash_consistency() { + let policy = RendezvousHashPolicy::new(); + let workers = make_workers(&[ + "http://worker1:8000", + "http://worker2:8000", + "http://worker3:8000", + ]); + + let request = r#"{"session_id": "consistent_test"}"#; + let idx1 = policy.select_worker(&workers, Some(request)); + let idx2 = policy.select_worker(&workers, Some(request)); + let idx3 = policy.select_worker(&workers, Some(request)); + + assert_eq!(idx1, idx2); + assert_eq!(idx2, idx3); + assert!(idx1.is_some()); + } + + #[test] + fn test_rendezvous_distribution() { + let policy = RendezvousHashPolicy::new(); + let workers = make_workers(&[ + "http://worker1:8000", + "http://worker2:8000", + "http://worker3:8000", + ]); + + let mut selected_workers = HashSet::new(); + for i in 0..100 { + let request = format!(r#"{{"session_id": "session_{}"}}"#, i); + if let Some(idx) = policy.select_worker(&workers, Some(&request)) { + selected_workers.insert(idx); + } + } + + // All workers should be selected at least once with 100 different sessions + assert_eq!(selected_workers.len(), 3); + } + + #[test] + fn test_rendezvous_with_unhealthy_workers() { + let policy = RendezvousHashPolicy::new(); + let workers = make_workers(&["http://worker1:8000", "http://worker2:8000"]); + + workers[0].set_healthy(false); + + // Should always select the healthy worker + for i in 0..10 { + let request = format!(r#"{{"session_id": "session_{}"}}"#, i); + assert_eq!(policy.select_worker(&workers, Some(&request)), Some(1)); + } + } + + #[test] + fn test_rendezvous_no_healthy_workers() { + let policy = RendezvousHashPolicy::new(); + let workers = make_workers(&["http://worker1:8000"]); + + workers[0].set_healthy(false); + assert_eq!(policy.select_worker(&workers, None), None); + } + + #[test] + fn test_rendezvous_minimal_redistribution() { + let policy = RendezvousHashPolicy::new(); + let workers_3 = make_workers(&[ + "http://worker1:8000", + "http://worker2:8000", + "http://worker3:8000", + ]); + let workers_2 = make_workers(&["http://worker1:8000", "http://worker2:8000"]); + + let mut same_count = 0; + let total = 100; + + for i in 0..total { + let request = format!(r#"{{"session_id": "session_{}"}}"#, i); + let idx_3 = policy.select_worker(&workers_3, Some(&request)).unwrap(); + let idx_2 = policy.select_worker(&workers_2, Some(&request)).unwrap(); + + let url_3 = workers_3[idx_3].url(); + let url_2 = workers_2[idx_2].url(); + if url_3 != "http://worker3:8000" && url_3 == url_2 { + same_count += 1; + } + } + + let sessions_not_on_worker3 = (0..total) + .filter(|i| { + let request = format!(r#"{{"session_id": "session_{}"}}"#, i); + let idx = policy.select_worker(&workers_3, Some(&request)).unwrap(); + workers_3[idx].url() != "http://worker3:8000" + }) + .count(); + + // All sessions that were on worker1/worker2 should stay there + assert_eq!(same_count, sessions_not_on_worker3); + } + + #[test] + fn test_rendezvous_header_routing() { + let policy = RendezvousHashPolicy::new(); + let workers = make_workers(&[ + "http://worker1:8000", + "http://worker2:8000", + "http://worker3:8000", + ]); + + let mut headers = RequestHeaders::new(); + headers.insert("x-session-id".to_string(), "test-session-123".to_string()); + + let idx1 = policy.select_worker_with_headers(&workers, None, Some(&headers)); + let idx2 = policy.select_worker_with_headers(&workers, None, Some(&headers)); + assert_eq!(idx1, idx2); + assert!(idx1.is_some()); + } +} diff --git a/lmdeploy/router/src/policies/round_robin.rs b/lmdeploy/router/src/policies/round_robin.rs new file mode 100644 index 0000000000..1a65d60937 --- /dev/null +++ b/lmdeploy/router/src/policies/round_robin.rs @@ -0,0 +1,142 @@ +//! Round-robin load balancing policy + +use super::{get_healthy_worker_indices, LoadBalancingPolicy, RequestHeaders}; +use crate::core::Worker; +use crate::metrics::RouterMetrics; +use std::sync::atomic::{AtomicUsize, Ordering}; +use std::sync::Arc; + +/// Round-robin selection policy +/// +/// Selects workers in sequential order, cycling through all healthy workers. +#[derive(Debug, Default)] +pub struct RoundRobinPolicy { + counter: AtomicUsize, +} + +impl RoundRobinPolicy { + pub fn new() -> Self { + Self { + counter: AtomicUsize::new(0), + } + } +} + +impl LoadBalancingPolicy for RoundRobinPolicy { + fn select_worker_with_headers( + &self, + workers: &[Arc], + _request_text: Option<&str>, + _headers: Option<&RequestHeaders>, + ) -> Option { + let healthy_indices = get_healthy_worker_indices(workers); + + if healthy_indices.is_empty() { + return None; + } + + // Get and increment counter atomically + let count = self.counter.fetch_add(1, Ordering::Relaxed); + let selected_idx = count % healthy_indices.len(); + let worker = workers[healthy_indices[selected_idx]].url(); + + RouterMetrics::record_processed_request(worker); + RouterMetrics::record_policy_decision(self.name(), worker); + Some(healthy_indices[selected_idx]) + } + + fn name(&self) -> &'static str { + "round_robin" + } + + fn reset(&self) { + self.counter.store(0, Ordering::Relaxed); + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::{BasicWorker, WorkerType}; + + #[test] + fn test_round_robin_selection() { + let policy = RoundRobinPolicy::new(); + let workers: Vec> = vec![ + Arc::new(BasicWorker::new( + "http://w1:8000".to_string(), + WorkerType::Regular, + )), + Arc::new(BasicWorker::new( + "http://w2:8000".to_string(), + WorkerType::Regular, + )), + Arc::new(BasicWorker::new( + "http://w3:8000".to_string(), + WorkerType::Regular, + )), + ]; + + // Should select workers in order: 0, 1, 2, 0, 1, 2, ... + assert_eq!(policy.select_worker(&workers, None), Some(0)); + assert_eq!(policy.select_worker(&workers, None), Some(1)); + assert_eq!(policy.select_worker(&workers, None), Some(2)); + assert_eq!(policy.select_worker(&workers, None), Some(0)); + assert_eq!(policy.select_worker(&workers, None), Some(1)); + } + + #[test] + fn test_round_robin_with_unhealthy_workers() { + let policy = RoundRobinPolicy::new(); + let workers: Vec> = vec![ + Arc::new(BasicWorker::new( + "http://w1:8000".to_string(), + WorkerType::Regular, + )), + Arc::new(BasicWorker::new( + "http://w2:8000".to_string(), + WorkerType::Regular, + )), + Arc::new(BasicWorker::new( + "http://w3:8000".to_string(), + WorkerType::Regular, + )), + ]; + + // Mark middle worker as unhealthy + workers[1].set_healthy(false); + + // Should skip unhealthy worker: 0, 2, 0, 2, ... + assert_eq!(policy.select_worker(&workers, None), Some(0)); + assert_eq!(policy.select_worker(&workers, None), Some(2)); + assert_eq!(policy.select_worker(&workers, None), Some(0)); + assert_eq!(policy.select_worker(&workers, None), Some(2)); + } + + #[test] + fn test_round_robin_reset() { + let policy = RoundRobinPolicy::new(); + let workers: Vec> = vec![ + Arc::new(BasicWorker::new( + "http://w1:8000".to_string(), + WorkerType::Regular, + )), + Arc::new(BasicWorker::new( + "http://w2:8000".to_string(), + WorkerType::Regular, + )), + ]; + + // Advance the counter + assert_eq!(policy.select_worker(&workers, None), Some(0)); + assert_eq!(policy.select_worker(&workers, None), Some(1)); + + // Reset should start from beginning + policy.reset(); + assert_eq!(policy.select_worker(&workers, None), Some(0)); + } +} diff --git a/lmdeploy/router/src/protocols/mod.rs b/lmdeploy/router/src/protocols/mod.rs new file mode 100644 index 0000000000..7359a3d2e2 --- /dev/null +++ b/lmdeploy/router/src/protocols/mod.rs @@ -0,0 +1,6 @@ +// Protocol definitions and validation for various LLM APIs +// This module provides a structured approach to handling different API protocols + +pub mod spec; +pub mod validation; +pub mod worker_spec; diff --git a/lmdeploy/router/src/protocols/spec.rs b/lmdeploy/router/src/protocols/spec.rs new file mode 100644 index 0000000000..8ce1c6b273 --- /dev/null +++ b/lmdeploy/router/src/protocols/spec.rs @@ -0,0 +1,3141 @@ +use serde::{Deserialize, Serialize}; +use serde_json::Value; +use std::collections::HashMap; + +// # Protocol Specifications +// +// This module contains all protocol definitions for OpenAI-compatible APIs. +// +// ## Table of Contents +// +// 1. **OPENAI SPEC - Chat Completions API** +// - Message Types +// - Response Format Types +// - Tool/Function Types +// - Streaming Delta Types +// - Request/Response structures +// +// 2. **OPENAI SPEC - Completions API** +// - Request/Response structures +// - Streaming support +// +// 3. **OPENAI SPEC - Responses API** +// - Tool Definitions +// - Reasoning Configuration +// - Input/Output Items +// - Service Tier & Tool Choice +// - Request/Response structures +// +// 4. **OPENAI SPEC - Common** +// - Shared Request Components +// - Tool Choice Types +// - Usage Tracking +// - Logprobs Types +// - Error Response Types +// +// 5. **LMDeploy SPEC - GENERATE API** +// - Generate Parameters +// - Sampling Parameters +// - Request/Response structures +// +// 7. **OPENAI SPEC - Embeddings API** +// - Request structures +// +// 8. **COMMON** +// - GenerationRequest trait +// - StringOrArray & LoRAPath types +// - Helper functions + +// ================================================================== +// = OPENAI SPEC - Chat Completions API = +// ================================================================== + +// ============= Message Types ============= + +// Note: We implement Deserialize manually below to dispatch based on the "role" field. +// This fixes a bug where serde's untagged enum matching would incorrectly match Assistant +// messages as System messages (because System only requires role + content, serde tries +// it first and silently ignores reasoning and other Assistant-specific fields). +#[derive(Debug, Clone, Serialize)] +#[serde(untagged)] +pub enum ChatMessage { + System { + role: String, + content: String, + #[serde(skip_serializing_if = "Option::is_none")] + name: Option, + }, + User { + role: String, // "user" + content: UserMessageContent, + #[serde(skip_serializing_if = "Option::is_none")] + name: Option, + }, + Assistant { + role: String, // "assistant" + #[serde(skip_serializing_if = "Option::is_none")] + content: Option, + #[serde(skip_serializing_if = "Option::is_none")] + name: Option, + #[serde(skip_serializing_if = "Option::is_none")] + tool_calls: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + function_call: Option, + /// Reasoning content for reasoning models. The custom Deserialize impl + /// below also accepts the deprecated `reasoning_content` alias. + #[serde(skip_serializing_if = "Option::is_none")] + reasoning: Option, + }, + Tool { + role: String, // "tool" + content: Value, + tool_call_id: String, + }, + Function { + role: String, // "function" + content: String, + name: String, + }, +} + +impl<'de> Deserialize<'de> for ChatMessage { + fn deserialize(deserializer: D) -> Result + where + D: serde::Deserializer<'de>, + { + use serde::de::Error; + + let value = Value::deserialize(deserializer)?; + let role = value + .get("role") + .and_then(|r| r.as_str()) + .ok_or_else(|| D::Error::custom("missing role field"))?; + + match role { + "assistant" => Ok(ChatMessage::Assistant { + role: role.to_string(), + content: value.get("content").and_then(|c| { + if c.is_null() { + None + } else { + c.as_str().map(String::from) + } + }), + name: value.get("name").and_then(|n| { + if n.is_null() { + None + } else { + n.as_str().map(String::from) + } + }), + tool_calls: value.get("tool_calls").and_then(|tc| { + if tc.is_null() { + None + } else { + serde_json::from_value(tc.clone()).ok() + } + }), + function_call: value.get("function_call").and_then(|fc| { + if fc.is_null() { + None + } else { + serde_json::from_value(fc.clone()).ok() + } + }), + // `reasoning` is the canonical field. Prefer a usable canonical + // string when both keys are present; otherwise fall back to the + // deprecated `reasoning_content` alias, including when `reasoning` + // is null or not a string. Serialization remains canonical. + reasoning: value + .get("reasoning") + .and_then(Value::as_str) + .or_else(|| value.get("reasoning_content").and_then(Value::as_str)) + .map(String::from), + }), + "system" => Ok(ChatMessage::System { + role: role.to_string(), + content: value + .get("content") + .and_then(|c| c.as_str()) + .unwrap_or("") + .to_string(), + name: value.get("name").and_then(|n| { + if n.is_null() { + None + } else { + n.as_str().map(String::from) + } + }), + }), + "user" => { + let content = value + .get("content") + .map(|c| { + serde_json::from_value(c.clone()) + .unwrap_or(UserMessageContent::Text(String::new())) + }) + .unwrap_or(UserMessageContent::Text(String::new())); + Ok(ChatMessage::User { + role: role.to_string(), + content, + name: value.get("name").and_then(|n| { + if n.is_null() { + None + } else { + n.as_str().map(String::from) + } + }), + }) + } + "tool" => Ok(ChatMessage::Tool { + role: role.to_string(), + content: value + .get("content") + .cloned() + .unwrap_or_else(|| Value::String(String::new())), + tool_call_id: value + .get("tool_call_id") + .and_then(|t| t.as_str()) + .unwrap_or("") + .to_string(), + }), + "function" => Ok(ChatMessage::Function { + role: role.to_string(), + content: value + .get("content") + .and_then(|c| c.as_str()) + .unwrap_or("") + .to_string(), + name: value + .get("name") + .and_then(|n| n.as_str()) + .unwrap_or("") + .to_string(), + }), + _ => Err(D::Error::custom(format!("unknown role: {}", role))), + } + } +} + +#[derive(Debug, Clone, Deserialize, Serialize, Default)] +pub struct StructuredOutputsParams { + /// JSON schema for structured output (mutually exclusive with other constraints) + #[serde(skip_serializing_if = "Option::is_none")] + pub json: Option, + + /// Regex pattern for structured output (mutually exclusive with other constraints) + #[serde(skip_serializing_if = "Option::is_none")] + pub regex: Option, + + /// List of choices for structured output (mutually exclusive with other constraints) + #[serde(skip_serializing_if = "Option::is_none")] + pub choice: Option>, + + /// Grammar for structured output (mutually exclusive with other constraints) + #[serde(skip_serializing_if = "Option::is_none")] + pub grammar: Option, + + /// JSON object mode (mutually exclusive with other constraints) + #[serde(skip_serializing_if = "Option::is_none")] + pub json_object: Option, + + /// Structural tag (mutually exclusive with other constraints) + #[serde(skip_serializing_if = "Option::is_none")] + pub structural_tag: Option, + + /// Disable fallback to non-structured output + #[serde(skip_serializing_if = "Option::is_none")] + pub disable_fallback: Option, + + /// Disable any whitespace in structured output + #[serde(skip_serializing_if = "Option::is_none")] + pub disable_any_whitespace: Option, + + /// Disable additional properties in JSON schema + #[serde(skip_serializing_if = "Option::is_none")] + pub disable_additional_properties: Option, + + /// Whitespace pattern for structured output + #[serde(skip_serializing_if = "Option::is_none")] + pub whitespace_pattern: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(untagged)] +pub enum UserMessageContent { + Text(String), + Parts(Vec), +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(tag = "type")] +pub enum ContentPart { + #[serde(rename = "text")] + Text { text: String }, + #[serde(rename = "image_url")] + ImageUrl { image_url: ImageUrl }, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ImageUrl { + pub url: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub detail: Option, // "auto", "low", or "high" +} + +// ============= Response Format Types ============= + +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(tag = "type")] +pub enum ResponseFormat { + #[serde(rename = "text")] + Text, + #[serde(rename = "json_object")] + JsonObject, + #[serde(rename = "json_schema")] + JsonSchema { json_schema: JsonSchemaFormat }, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct JsonSchemaFormat { + pub name: String, + pub schema: Value, + #[serde(skip_serializing_if = "Option::is_none")] + pub strict: Option, +} + +// ============= Streaming Delta Types ============= + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ChatMessageDelta { + #[serde(skip_serializing_if = "Option::is_none")] + pub role: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub content: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_calls: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub function_call: Option, + /// Reasoning content delta for reasoning models + #[serde(skip_serializing_if = "Option::is_none")] + pub reasoning: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ToolCallDelta { + pub index: u32, + #[serde(skip_serializing_if = "Option::is_none")] + pub id: Option, + #[serde(skip_serializing_if = "Option::is_none")] + #[serde(rename = "type")] + pub tool_type: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub function: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct FunctionCallDelta { + #[serde(skip_serializing_if = "Option::is_none")] + pub name: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub arguments: Option, +} + +// ============= Request ============= + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ChatCompletionRequest { + /// ID of the model to use (optional, LMDeploy supports requests without model) + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model: Option, + + /// A list of messages comprising the conversation so far + pub messages: Vec, + + /// Pre-tokenized input for LMDeploy token-in-token-out requests. + /// This is used when `messages` is empty. + #[serde(skip_serializing_if = "Option::is_none")] + pub input_ids: Option>, + + /// What sampling temperature to use, between 0 and 2 + #[serde(skip_serializing_if = "Option::is_none")] + pub temperature: Option, + + /// An alternative to sampling with temperature + #[serde(skip_serializing_if = "Option::is_none")] + pub top_p: Option, + + /// How many chat completion choices to generate for each input message + #[serde(skip_serializing_if = "Option::is_none")] + pub n: Option, + + /// If set, partial message deltas will be sent + #[serde(default)] + pub stream: bool, + + /// Options for streaming response + #[serde(skip_serializing_if = "Option::is_none")] + pub stream_options: Option, + + /// Up to 4 sequences where the API will stop generating further tokens + #[serde(skip_serializing_if = "Option::is_none")] + pub stop: Option, + + /// The maximum number of tokens to generate + #[serde(skip_serializing_if = "Option::is_none")] + pub max_tokens: Option, + + /// An upper bound for the number of tokens that can be generated for a completion + #[serde(skip_serializing_if = "Option::is_none")] + pub max_completion_tokens: Option, + + /// Number between -2.0 and 2.0. Positive values penalize new tokens based on whether they appear in the text so far + #[serde(skip_serializing_if = "Option::is_none")] + pub presence_penalty: Option, + + /// Number between -2.0 and 2.0. Positive values penalize new tokens based on their existing frequency in the text so far + #[serde(skip_serializing_if = "Option::is_none")] + pub frequency_penalty: Option, + + /// Modify the likelihood of specified tokens appearing in the completion + #[serde(skip_serializing_if = "Option::is_none")] + pub logit_bias: Option>, + + /// A unique identifier representing your end-user + #[serde(skip_serializing_if = "Option::is_none")] + pub user: Option, + + /// If specified, our system will make a best effort to sample deterministically + #[serde(skip_serializing_if = "Option::is_none")] + pub seed: Option, + + /// Whether to return log probabilities of the output tokens + #[serde(default)] + pub logprobs: bool, + + /// An integer between 0 and 20 specifying the number of most likely tokens to return + #[serde(skip_serializing_if = "Option::is_none")] + pub top_logprobs: Option, + + /// An object specifying the format that the model must output + #[serde(skip_serializing_if = "Option::is_none")] + pub response_format: Option, + + /// A list of tools the model may call + #[serde(skip_serializing_if = "Option::is_none")] + pub tools: Option>, + + /// Controls which (if any) tool is called by the model + #[serde(skip_serializing_if = "Option::is_none")] + pub tool_choice: Option, + + /// Whether to enable parallel function calling during tool use + #[serde(skip_serializing_if = "Option::is_none")] + pub parallel_tool_calls: Option, + + /// Deprecated: use tools instead + #[serde(skip_serializing_if = "Option::is_none")] + pub functions: Option>, + + /// Deprecated: use tool_choice instead + #[serde(skip_serializing_if = "Option::is_none")] + pub function_call: Option, + + // ============= Generation Extensions ============= + /// Top-k sampling parameter (-1 to disable) + #[serde(skip_serializing_if = "Option::is_none")] + pub top_k: Option, + + /// Min-p nucleus sampling parameter + #[serde(skip_serializing_if = "Option::is_none")] + pub min_p: Option, + + /// Minimum number of tokens to generate + #[serde(skip_serializing_if = "Option::is_none")] + pub min_tokens: Option, + + /// Repetition penalty for reducing repetitive text + #[serde(skip_serializing_if = "Option::is_none")] + pub repetition_penalty: Option, + + /// Regex constraint for output generation + #[serde(skip_serializing_if = "Option::is_none")] + pub regex: Option, + + /// EBNF grammar constraint for structured output + #[serde(skip_serializing_if = "Option::is_none")] + pub ebnf: Option, + + /// Specific token IDs to use as stop conditions + #[serde(skip_serializing_if = "Option::is_none")] + pub stop_token_ids: Option>, + + /// Skip trimming stop tokens from output + #[serde(default)] + pub no_stop_trim: bool, + + /// Ignore end-of-sequence tokens during generation + #[serde(default)] + pub ignore_eos: bool, + + /// Add generation prompt to the chat template + #[serde(default = "default_true")] + pub add_generation_prompt: bool, + + /// Continue generating from final assistant message + #[serde(default)] + pub continue_final_message: bool, + + /// Skip special tokens during detokenization + #[serde(default = "default_true")] + pub skip_special_tokens: bool, + + // ============= Generation Extensions ============= + /// Separate reasoning content from final answer (O1-style models) + #[serde(default = "default_true")] + pub separate_reasoning: bool, + + /// Stream reasoning tokens during generation + #[serde(default = "default_true")] + pub stream_reasoning: bool, + + /// Chat template kwargs + #[serde(skip_serializing_if = "Option::is_none")] + pub chat_template_kwargs: Option>, + + /// Echo back the prompt in addition to the completion + #[serde(skip_serializing_if = "Option::is_none")] + pub echo: Option, + + /// Reasoning effort level for reasoning models (low, medium, high) + #[serde(skip_serializing_if = "Option::is_none")] + pub reasoning_effort: Option, + + /// Whether to include reasoning in the response + #[serde(default = "default_true")] + pub include_reasoning: bool, + /// Structured outputs parameters + #[serde(skip_serializing_if = "Option::is_none")] + pub structured_outputs: Option, + + /// Additional fields passed through transparently to the backend + #[serde(flatten)] + pub other: serde_json::Map, +} + +impl GenerationRequest for ChatCompletionRequest { + fn is_stream(&self) -> bool { + self.stream + } + + fn get_model(&self) -> Option<&str> { + self.model.as_deref() + } + + fn extract_text_for_routing(&self) -> String { + // LMDeploy fallback: when messages is empty, input_ids is the active input. + // Use the typed field directly so consistent_hash/cache_aware receive a + // deterministic routing key without reparsing flattened JSON. + if let Some(input_ids) = self.input_ids.as_ref().filter(|ids| !ids.is_empty()) { + return input_ids + .iter() + .map(i32::to_string) + .collect::>() + .join(" "); + } + + // Return empty string if no routing key found - random/roundrobin/power_of_two + // policies work without a key; consistent_hash/cache_aware will treat all as same bucket. + String::new() + } +} + +// ============= Regular Response ============= + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ChatCompletionResponse { + pub id: String, + pub object: String, // "chat.completion" + pub created: u64, + pub model: String, + pub choices: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub usage: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub system_fingerprint: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ChatChoice { + pub index: u32, + pub message: ChatMessage, + #[serde(skip_serializing_if = "Option::is_none")] + pub logprobs: Option, + pub finish_reason: Option, // "stop", "length", "tool_calls", "content_filter", "function_call" + /// Information about which stop condition was matched + #[serde(skip_serializing_if = "Option::is_none")] + pub matched_stop: Option, // Can be string or integer + /// Generated token IDs (LMDeploy extension, enabled by `return_token_ids`). + #[serde(skip_serializing_if = "Option::is_none")] + pub output_ids: Option>, + /// Preserve LMDeploy extensions such as `output_token_logprobs` and + /// `routed_experts`, as well as future backend-specific choice fields. + #[serde(flatten)] + pub other: serde_json::Map, +} + +// ============= Streaming Response ============= + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ChatCompletionStreamResponse { + pub id: String, + pub object: String, // "chat.completion.chunk" + pub created: u64, + pub model: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub system_fingerprint: Option, + pub choices: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub usage: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ChatStreamChoice { + pub index: u32, + pub delta: ChatMessageDelta, + #[serde(skip_serializing_if = "Option::is_none")] + pub logprobs: Option, + pub finish_reason: Option, + /// Generated token IDs for this stream chunk (LMDeploy extension). + #[serde(skip_serializing_if = "Option::is_none")] + pub output_ids: Option>, + /// Preserve LMDeploy extensions such as `output_token_logprobs` and + /// `routed_experts`, as well as future backend-specific choice fields. + #[serde(flatten)] + pub other: serde_json::Map, +} + +// ================================================================== +// = OPENAI SPEC - Completions API = +// ================================================================== +// Completions API request types (v1/completions) - DEPRECATED but still supported + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct CompletionRequest { + /// ID of the model to use (optional, LMDeploy supports requests without model) + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model: Option, + + /// The prompt(s) to generate completions for + /// Supports: str, list[str], list[int], list[list[int]] + pub prompt: PromptInput, + + /// The suffix that comes after a completion of inserted text + #[serde(skip_serializing_if = "Option::is_none")] + pub suffix: Option, + + /// The maximum number of tokens to generate + #[serde(skip_serializing_if = "Option::is_none")] + pub max_tokens: Option, + + /// What sampling temperature to use, between 0 and 2 + #[serde(skip_serializing_if = "Option::is_none")] + pub temperature: Option, + + /// An alternative to sampling with temperature (nucleus sampling) + #[serde(skip_serializing_if = "Option::is_none")] + pub top_p: Option, + + /// How many completions to generate for each prompt + #[serde(skip_serializing_if = "Option::is_none")] + pub n: Option, + + /// Whether to stream back partial progress + #[serde(default)] + pub stream: bool, + + /// Options for streaming response + #[serde(skip_serializing_if = "Option::is_none")] + pub stream_options: Option, + + /// Include the log probabilities on the logprobs most likely tokens + #[serde(skip_serializing_if = "Option::is_none")] + pub logprobs: Option, + + /// Echo back the prompt in addition to the completion + #[serde(default)] + pub echo: bool, + + /// Up to 4 sequences where the API will stop generating further tokens + #[serde(skip_serializing_if = "Option::is_none")] + pub stop: Option, + + /// Number between -2.0 and 2.0. Positive values penalize new tokens based on whether they appear in the text so far + #[serde(skip_serializing_if = "Option::is_none")] + pub presence_penalty: Option, + + /// Number between -2.0 and 2.0. Positive values penalize new tokens based on their existing frequency in the text so far + #[serde(skip_serializing_if = "Option::is_none")] + pub frequency_penalty: Option, + + /// Generates best_of completions server-side and returns the "best" + #[serde(skip_serializing_if = "Option::is_none")] + pub best_of: Option, + + /// Modify the likelihood of specified tokens appearing in the completion + #[serde(skip_serializing_if = "Option::is_none")] + pub logit_bias: Option>, + + /// A unique identifier representing your end-user + #[serde(skip_serializing_if = "Option::is_none")] + pub user: Option, + + /// If specified, our system will make a best effort to sample deterministically + #[serde(skip_serializing_if = "Option::is_none")] + pub seed: Option, + + // ============= Generation Extensions ============= + /// Top-k sampling parameter (-1 to disable) + #[serde(skip_serializing_if = "Option::is_none")] + pub top_k: Option, + + /// Min-p nucleus sampling parameter + #[serde(skip_serializing_if = "Option::is_none")] + pub min_p: Option, + + /// Minimum number of tokens to generate + #[serde(skip_serializing_if = "Option::is_none")] + pub min_tokens: Option, + + /// Repetition penalty for reducing repetitive text + #[serde(skip_serializing_if = "Option::is_none")] + pub repetition_penalty: Option, + + /// Regex constraint for output generation + #[serde(skip_serializing_if = "Option::is_none")] + pub regex: Option, + + /// EBNF grammar constraint for structured output + #[serde(skip_serializing_if = "Option::is_none")] + pub ebnf: Option, + + /// JSON schema constraint for structured output + #[serde(skip_serializing_if = "Option::is_none")] + pub json_schema: Option, + + /// Specific token IDs to use as stop conditions + #[serde(skip_serializing_if = "Option::is_none")] + pub stop_token_ids: Option>, + + /// Skip trimming stop tokens from output + #[serde(default)] + pub no_stop_trim: bool, + + /// Ignore end-of-sequence tokens during generation + #[serde(default)] + pub ignore_eos: bool, + + /// Skip special tokens during detokenization + #[serde(default = "default_true")] + pub skip_special_tokens: bool, + + /// Additional fields passed through transparently to the backend + #[serde(flatten)] + pub other: serde_json::Map, +} + +impl GenerationRequest for CompletionRequest { + fn is_stream(&self) -> bool { + self.stream + } + + fn get_model(&self) -> Option<&str> { + self.model.as_deref() + } + + fn extract_text_for_routing(&self) -> String { + if let Some(user) = self.user.as_deref() { + return format!("user:{user}"); + } + self.prompt.extract_text_for_routing() + } +} + +// ============= Regular Response ============= + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct CompletionResponse { + pub id: String, + pub object: String, // "text_completion" + pub created: u64, + pub model: String, + pub choices: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + pub usage: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub system_fingerprint: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct CompletionChoice { + pub text: String, + pub index: u32, + #[serde(skip_serializing_if = "Option::is_none")] + pub logprobs: Option, + pub finish_reason: Option, // "stop", "length", "content_filter", etc. + /// Information about which stop condition was matched + #[serde(skip_serializing_if = "Option::is_none")] + pub matched_stop: Option, // Can be string or integer +} + +// ============= Streaming Response ============= + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct CompletionStreamResponse { + pub id: String, + pub object: String, // "text_completion" + pub created: u64, + pub choices: Vec, + pub model: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub system_fingerprint: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct CompletionStreamChoice { + pub text: String, + pub index: u32, + #[serde(skip_serializing_if = "Option::is_none")] + pub logprobs: Option, + pub finish_reason: Option, +} + +// ================================================================== +// = OPENAI SPEC - Responses API = +// ================================================================== + +// ============= Tool Definitions ============= + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ResponseTool { + #[serde(rename = "type")] + pub r#type: ResponseToolType, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum ResponseToolType { + WebSearchPreview, + CodeInterpreter, +} + +// ============= Reasoning Configuration ============= + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ResponseReasoningParam { + #[serde(default = "default_reasoning_effort")] + pub effort: Option, +} + +fn default_reasoning_effort() -> Option { + Some(ReasoningEffort::Medium) +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum ReasoningEffort { + Low, + Medium, + High, +} + +// ============= Input/Output Items ============= + +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(tag = "type")] +#[serde(rename_all = "snake_case")] +pub enum ResponseInputOutputItem { + #[serde(rename = "message")] + Message { + id: String, + role: String, + content: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + status: Option, + }, + #[serde(rename = "reasoning")] + Reasoning { + id: String, + #[serde(skip_serializing_if = "Vec::is_empty")] + summary: Vec, + content: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + status: Option, + }, + #[serde(rename = "function_tool_call")] + FunctionToolCall { + id: String, + name: String, + arguments: String, + #[serde(skip_serializing_if = "Option::is_none")] + output: Option, + #[serde(skip_serializing_if = "Option::is_none")] + status: Option, + }, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(tag = "type")] +#[serde(rename_all = "snake_case")] +pub enum ResponseContentPart { + #[serde(rename = "output_text")] + OutputText { + text: String, + #[serde(skip_serializing_if = "Vec::is_empty")] + annotations: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + logprobs: Option, + }, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(tag = "type")] +#[serde(rename_all = "snake_case")] +pub enum ResponseReasoningContent { + #[serde(rename = "reasoning_text")] + ReasoningText { text: String }, +} + +// ============= Output Items for Response ============= + +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(tag = "type")] +#[serde(rename_all = "snake_case")] +pub enum ResponseOutputItem { + #[serde(rename = "message")] + Message { + id: String, + role: String, + content: Vec, + status: String, + }, + #[serde(rename = "reasoning")] + Reasoning { + id: String, + #[serde(skip_serializing_if = "Vec::is_empty")] + summary: Vec, + content: Vec, + #[serde(skip_serializing_if = "Option::is_none")] + status: Option, + }, + #[serde(rename = "function_tool_call")] + FunctionToolCall { + id: String, + name: String, + arguments: String, + #[serde(skip_serializing_if = "Option::is_none")] + output: Option, + status: String, + }, +} + +// ============= Service Tier ============= + +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(rename_all = "snake_case")] +#[derive(Default)] +pub enum ServiceTier { + #[default] + Auto, + Default, + Flex, + Scale, + Priority, +} + +// ============= Truncation ============= + +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(rename_all = "snake_case")] +#[derive(Default)] +pub enum Truncation { + Auto, + #[default] + Disabled, +} + +// ============= Response Status ============= + +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum ResponseStatus { + Queued, + InProgress, + Completed, + Failed, + Cancelled, +} + +// ============= Include Fields ============= + +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum IncludeField { + #[serde(rename = "code_interpreter_call.outputs")] + CodeInterpreterCallOutputs, + #[serde(rename = "computer_call_output.output.image_url")] + ComputerCallOutputImageUrl, + #[serde(rename = "file_search_call.results")] + FileSearchCallResults, + #[serde(rename = "message.input_image.image_url")] + MessageInputImageUrl, + #[serde(rename = "message.output_text.logprobs")] + MessageOutputTextLogprobs, + #[serde(rename = "reasoning.encrypted_content")] + ReasoningEncryptedContent, +} + +// ============= Usage Info ============= + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct UsageInfo { + pub prompt_tokens: u32, + pub completion_tokens: u32, + pub total_tokens: u32, + #[serde(skip_serializing_if = "Option::is_none")] + pub reasoning_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub prompt_tokens_details: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct PromptTokenUsageInfo { + pub cached_tokens: u32, +} + +// ============= Response Usage Format ============= + +/// OpenAI Responses API usage format (different from standard UsageInfo) +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ResponseUsage { + pub input_tokens: u32, + pub output_tokens: u32, + pub total_tokens: u32, + #[serde(skip_serializing_if = "Option::is_none")] + pub input_tokens_details: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub output_tokens_details: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct InputTokensDetails { + pub cached_tokens: u32, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct OutputTokensDetails { + pub reasoning_tokens: u32, +} + +impl UsageInfo { + /// Convert to OpenAI Responses API format + pub fn to_response_usage(&self) -> ResponseUsage { + ResponseUsage { + input_tokens: self.prompt_tokens, + output_tokens: self.completion_tokens, + total_tokens: self.total_tokens, + input_tokens_details: self.prompt_tokens_details.as_ref().map(|details| { + InputTokensDetails { + cached_tokens: details.cached_tokens, + } + }), + output_tokens_details: self.reasoning_tokens.map(|tokens| OutputTokensDetails { + reasoning_tokens: tokens, + }), + } + } +} + +impl From for ResponseUsage { + fn from(usage: UsageInfo) -> Self { + usage.to_response_usage() + } +} + +impl ResponseUsage { + /// Convert back to standard UsageInfo format + pub fn to_usage_info(&self) -> UsageInfo { + UsageInfo { + prompt_tokens: self.input_tokens, + completion_tokens: self.output_tokens, + total_tokens: self.total_tokens, + reasoning_tokens: self + .output_tokens_details + .as_ref() + .map(|details| details.reasoning_tokens), + prompt_tokens_details: self.input_tokens_details.as_ref().map(|details| { + PromptTokenUsageInfo { + cached_tokens: details.cached_tokens, + } + }), + } + } +} + +fn generate_request_id() -> String { + format!("resp_{}", uuid::Uuid::new_v4().simple()) +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ResponsesRequest { + // ============= Core OpenAI API fields ============= + /// Run the request in the background + #[serde(default)] + pub background: bool, + + /// Fields to include in the response + #[serde(skip_serializing_if = "Option::is_none")] + pub include: Option>, + + /// Input content - can be string or structured items + pub input: ResponseInput, + + /// System instructions for the model + #[serde(skip_serializing_if = "Option::is_none")] + pub instructions: Option, + + /// Maximum number of output tokens + #[serde(skip_serializing_if = "Option::is_none")] + pub max_output_tokens: Option, + + /// Maximum number of tool calls + #[serde(skip_serializing_if = "Option::is_none")] + pub max_tool_calls: Option, + + /// Additional metadata + #[serde(skip_serializing_if = "Option::is_none")] + pub metadata: Option>, + + /// Model to use (optional, LMDeploy supports requests without model) + #[serde(skip_serializing_if = "Option::is_none")] + pub model: Option, + + /// Whether to enable parallel tool calls + #[serde(default = "default_true")] + pub parallel_tool_calls: bool, + + /// ID of previous response to continue from + #[serde(skip_serializing_if = "Option::is_none")] + pub previous_response_id: Option, + + /// Reasoning configuration + #[serde(skip_serializing_if = "Option::is_none")] + pub reasoning: Option, + + /// Service tier + #[serde(default)] + pub service_tier: ServiceTier, + + /// Whether to store the response + #[serde(default = "default_true")] + pub store: bool, + + /// Whether to stream the response + #[serde(default)] + pub stream: bool, + + /// Temperature for sampling + #[serde(skip_serializing_if = "Option::is_none")] + pub temperature: Option, + + /// Tool choice behavior + #[serde(default)] + pub tool_choice: ToolChoice, + + /// Available tools + #[serde(default)] + pub tools: Vec, + + /// Number of top logprobs to return + #[serde(default)] + pub top_logprobs: u32, + + /// Top-p sampling parameter + #[serde(skip_serializing_if = "Option::is_none")] + pub top_p: Option, + + /// Truncation behavior + #[serde(default)] + pub truncation: Truncation, + + /// User identifier + #[serde(skip_serializing_if = "Option::is_none")] + pub user: Option, + + // ============= Generation Extensions ============= + /// Request ID + #[serde(default = "generate_request_id")] + pub request_id: String, + + /// Request priority + #[serde(default)] + pub priority: i32, + + /// Frequency penalty + #[serde(default)] + pub frequency_penalty: f32, + + /// Presence penalty + #[serde(default)] + pub presence_penalty: f32, + + /// Stop sequences + #[serde(skip_serializing_if = "Option::is_none")] + pub stop: Option, + + /// Top-k sampling parameter + #[serde(default = "default_top_k")] + pub top_k: i32, + + /// Min-p sampling parameter + #[serde(default)] + pub min_p: f32, + + /// Repetition penalty + #[serde(default = "default_repetition_penalty")] + pub repetition_penalty: f32, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(untagged)] +pub enum ResponseInput { + Text(String), + Items(Vec), +} + +fn default_top_k() -> i32 { + -1 +} + +fn default_repetition_penalty() -> f32 { + 1.0 +} + +impl ResponsesRequest { + /// Default sampling parameters + const DEFAULT_TEMPERATURE: f32 = 0.7; + const DEFAULT_TOP_P: f32 = 1.0; + + /// Convert to sampling parameters for generation + pub fn to_sampling_params( + &self, + default_max_tokens: u32, + default_params: Option>, + ) -> HashMap { + let mut params = HashMap::new(); + + // Use max_output_tokens if available + let max_tokens = if let Some(max_output) = self.max_output_tokens { + std::cmp::min(max_output, default_max_tokens) + } else { + default_max_tokens + }; + + // Avoid exceeding context length by minus 1 token + let max_tokens = max_tokens.saturating_sub(1); + + // Temperature + let temperature = self.temperature.unwrap_or_else(|| { + default_params + .as_ref() + .and_then(|p| p.get("temperature")) + .and_then(|v| v.as_f64()) + .map(|v| v as f32) + .unwrap_or(Self::DEFAULT_TEMPERATURE) + }); + + // Top-p + let top_p = self.top_p.unwrap_or_else(|| { + default_params + .as_ref() + .and_then(|p| p.get("top_p")) + .and_then(|v| v.as_f64()) + .map(|v| v as f32) + .unwrap_or(Self::DEFAULT_TOP_P) + }); + + params.insert( + "max_new_tokens".to_string(), + serde_json::Value::Number(serde_json::Number::from(max_tokens)), + ); + params.insert( + "temperature".to_string(), + serde_json::Value::Number(serde_json::Number::from_f64(temperature as f64).unwrap()), + ); + params.insert( + "top_p".to_string(), + serde_json::Value::Number(serde_json::Number::from_f64(top_p as f64).unwrap()), + ); + params.insert( + "frequency_penalty".to_string(), + serde_json::Value::Number( + serde_json::Number::from_f64(self.frequency_penalty as f64).unwrap(), + ), + ); + params.insert( + "presence_penalty".to_string(), + serde_json::Value::Number( + serde_json::Number::from_f64(self.presence_penalty as f64).unwrap(), + ), + ); + params.insert( + "top_k".to_string(), + serde_json::Value::Number(serde_json::Number::from(self.top_k)), + ); + params.insert( + "min_p".to_string(), + serde_json::Value::Number(serde_json::Number::from_f64(self.min_p as f64).unwrap()), + ); + params.insert( + "repetition_penalty".to_string(), + serde_json::Value::Number( + serde_json::Number::from_f64(self.repetition_penalty as f64).unwrap(), + ), + ); + + if let Some(ref stop) = self.stop { + match serde_json::to_value(stop) { + Ok(value) => params.insert("stop".to_string(), value), + Err(_) => params.insert("stop".to_string(), serde_json::Value::Null), + }; + } + + // Apply any additional default parameters + if let Some(default_params) = default_params { + for (key, value) in default_params { + params.entry(key).or_insert(value); + } + } + + params + } +} + +impl GenerationRequest for ResponsesRequest { + fn is_stream(&self) -> bool { + self.stream + } + + fn get_model(&self) -> Option<&str> { + self.model.as_deref() + } + + fn extract_text_for_routing(&self) -> String { + match &self.input { + ResponseInput::Text(text) => text.clone(), + ResponseInput::Items(items) => items + .iter() + .filter_map(|item| match item { + ResponseInputOutputItem::Message { content, .. } => { + let texts: Vec = content + .iter() + .map(|part| match part { + ResponseContentPart::OutputText { text, .. } => text.clone(), + }) + .collect(); + if texts.is_empty() { + None + } else { + Some(texts.join(" ")) + } + } + ResponseInputOutputItem::Reasoning { content, .. } => { + let texts: Vec = content + .iter() + .map(|part| match part { + ResponseReasoningContent::ReasoningText { text } => text.clone(), + }) + .collect(); + if texts.is_empty() { + None + } else { + Some(texts.join(" ")) + } + } + ResponseInputOutputItem::FunctionToolCall { arguments, .. } => { + Some(arguments.clone()) + } + }) + .collect::>() + .join(" "), + } + } +} + +fn generate_response_id() -> String { + format!("resp_{}", uuid::Uuid::new_v4().simple()) +} + +fn current_timestamp() -> i64 { + std::time::SystemTime::now() + .duration_since(std::time::UNIX_EPOCH) + .unwrap_or_else(|_| std::time::Duration::from_secs(0)) + .as_secs() as i64 +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ResponsesResponse { + /// Response ID + #[serde(default = "generate_response_id")] + pub id: String, + + /// Object type + #[serde(default = "default_object_type")] + pub object: String, + + /// Creation timestamp + #[serde(default = "current_timestamp")] + pub created_at: i64, + + /// Model name + pub model: String, + + /// Output items + #[serde(default)] + pub output: Vec, + + /// Response status + pub status: ResponseStatus, + + /// Usage information + #[serde(skip_serializing_if = "Option::is_none")] + pub usage: Option, + + /// Whether parallel tool calls are enabled + #[serde(default = "default_true")] + pub parallel_tool_calls: bool, + + /// Tool choice setting + #[serde(default = "default_tool_choice")] + pub tool_choice: String, + + /// Available tools + #[serde(default)] + pub tools: Vec, +} + +fn default_object_type() -> String { + "response".to_string() +} + +fn default_tool_choice() -> String { + "auto".to_string() +} + +impl ResponsesResponse { + /// Create a response from a request + #[allow(clippy::too_many_arguments)] + pub fn from_request( + request: &ResponsesRequest, + _sampling_params: &HashMap, + model_name: String, + created_time: i64, + output: Vec, + status: ResponseStatus, + usage: Option, + ) -> Self { + Self { + id: request.request_id.clone(), + object: "response".to_string(), + created_at: created_time, + model: model_name, + output, + status, + usage, + parallel_tool_calls: request.parallel_tool_calls, + tool_choice: match &request.tool_choice { + ToolChoice::Value(ToolChoiceValue::Auto) => "auto".to_string(), + ToolChoice::Value(ToolChoiceValue::Required) => "required".to_string(), + ToolChoice::Value(ToolChoiceValue::None) => "none".to_string(), + ToolChoice::Function { .. } => "function".to_string(), + }, + tools: request.tools.clone(), + } + } + + /// Create a new response with default values + pub fn new(request_id: String, model: String, status: ResponseStatus) -> Self { + Self { + id: request_id, + object: "response".to_string(), + created_at: current_timestamp(), + model, + output: Vec::new(), + status, + usage: None, + parallel_tool_calls: true, + tool_choice: "auto".to_string(), + tools: Vec::new(), + } + } + + /// Add an output item to the response + pub fn add_output(&mut self, item: ResponseOutputItem) { + self.output.push(item); + } + + /// Set the usage information + pub fn set_usage(&mut self, usage: UsageInfo) { + self.usage = Some(usage); + } + + /// Update the status + pub fn set_status(&mut self, status: ResponseStatus) { + self.status = status; + } + + /// Check if the response is complete + pub fn is_complete(&self) -> bool { + matches!(self.status, ResponseStatus::Completed) + } + + /// Check if the response is in progress + pub fn is_in_progress(&self) -> bool { + matches!(self.status, ResponseStatus::InProgress) + } + + /// Check if the response failed + pub fn is_failed(&self) -> bool { + matches!(self.status, ResponseStatus::Failed) + } + + /// Check if the response was cancelled + pub fn is_cancelled(&self) -> bool { + matches!(self.status, ResponseStatus::Cancelled) + } + + /// Check if the response is queued + pub fn is_queued(&self) -> bool { + matches!(self.status, ResponseStatus::Queued) + } + + /// Convert usage to OpenAI Responses API format + pub fn usage_in_response_format(&self) -> Option { + self.usage.as_ref().map(|usage| usage.to_response_usage()) + } + + /// Get the response as a JSON value with usage in response format + pub fn to_response_format(&self) -> serde_json::Value { + let mut response = serde_json::to_value(self).unwrap_or(serde_json::Value::Null); + + // Convert usage to response format if present + if let Some(usage) = &self.usage { + if let Ok(usage_value) = serde_json::to_value(usage.to_response_usage()) { + response["usage"] = usage_value; + } + } + + response + } +} + +// ============= Helper Functions ============= + +impl ResponseOutputItem { + /// Create a new message output item + pub fn new_message( + id: String, + role: String, + content: Vec, + status: String, + ) -> Self { + Self::Message { + id, + role, + content, + status, + } + } + + /// Create a new reasoning output item + pub fn new_reasoning( + id: String, + summary: Vec, + content: Vec, + status: Option, + ) -> Self { + Self::Reasoning { + id, + summary, + content, + status, + } + } + + /// Create a new function tool call output item + pub fn new_function_tool_call( + id: String, + name: String, + arguments: String, + output: Option, + status: String, + ) -> Self { + Self::FunctionToolCall { + id, + name, + arguments, + output, + status, + } + } +} + +impl ResponseContentPart { + /// Create a new text content part + pub fn new_text( + text: String, + annotations: Vec, + logprobs: Option, + ) -> Self { + Self::OutputText { + text, + annotations, + logprobs, + } + } +} + +impl ResponseReasoningContent { + /// Create a new reasoning text content + pub fn new_reasoning_text(text: String) -> Self { + Self::ReasoningText { text } + } +} + +impl UsageInfo { + /// Create a new usage info with token counts + pub fn new(prompt_tokens: u32, completion_tokens: u32, reasoning_tokens: Option) -> Self { + Self { + prompt_tokens, + completion_tokens, + total_tokens: prompt_tokens + completion_tokens, + reasoning_tokens, + prompt_tokens_details: None, + } + } + + /// Create usage info with cached token details + pub fn new_with_cached( + prompt_tokens: u32, + completion_tokens: u32, + reasoning_tokens: Option, + cached_tokens: u32, + ) -> Self { + Self { + prompt_tokens, + completion_tokens, + total_tokens: prompt_tokens + completion_tokens, + reasoning_tokens, + prompt_tokens_details: Some(PromptTokenUsageInfo { cached_tokens }), + } + } +} + +// ================================================================== +// = OPENAI SPEC - Common = +// ================================================================== + +// ============= Shared Request Components ============= + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct StreamOptions { + #[serde(skip_serializing_if = "Option::is_none")] + pub include_usage: Option, +} + +// ============= Tool Choice Types ============= + +/// Tool choice value for simple string options +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(rename_all = "snake_case")] +pub enum ToolChoiceValue { + Auto, + Required, + None, +} + +/// Tool choice for both Chat Completion and Responses APIs +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(untagged)] +pub enum ToolChoice { + Value(ToolChoiceValue), + Function { + #[serde(rename = "type")] + tool_type: String, // "function" + function: FunctionChoice, + }, +} + +impl Default for ToolChoice { + fn default() -> Self { + Self::Value(ToolChoiceValue::Auto) + } +} + +/// Function choice specification for ToolChoice::Function +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct FunctionChoice { + pub name: String, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct Tool { + #[serde(rename = "type")] + pub tool_type: String, // "function" + pub function: Function, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct Function { + pub name: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub description: Option, + pub parameters: Value, // JSON Schema +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ToolCall { + pub id: String, + #[serde(rename = "type")] + pub tool_type: String, // "function" + pub function: FunctionCallResponse, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(untagged)] +pub enum FunctionCall { + None, + Auto, + Function { name: String }, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct FunctionCallResponse { + pub name: String, + #[serde(default)] + pub arguments: Option, // JSON string +} + +// ============= Usage Tracking ============= + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct Usage { + pub prompt_tokens: u32, + pub completion_tokens: u32, + pub total_tokens: u32, + #[serde(skip_serializing_if = "Option::is_none")] + pub completion_tokens_details: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct CompletionTokensDetails { + pub reasoning_tokens: Option, +} + +// ============= Logprobs Types ============= + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct LogProbs { + pub tokens: Vec, + pub token_logprobs: Vec>, + pub top_logprobs: Vec>>, + pub text_offset: Vec, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ChatLogProbs { + pub content: Option>, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ChatLogProbsContent { + pub token: String, + pub logprob: f32, + pub bytes: Option>, + pub top_logprobs: Vec, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct TopLogProb { + pub token: String, + pub logprob: f32, + pub bytes: Option>, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ErrorResponse { + pub error: ErrorDetail, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ErrorDetail { + pub message: String, + #[serde(rename = "type")] + pub error_type: String, + #[serde(skip_serializing_if = "Option::is_none")] + pub param: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub code: Option, +} + +// ================================================================== +// = LMDeploy SPEC - GENERATE API = +// ================================================================== + +#[derive(Debug, Clone, Deserialize, Serialize)] +#[serde(untagged)] +pub enum InputIds { + Single(Vec), + Batch(Vec>), +} + +#[derive(Debug, Clone, Deserialize, Serialize, Default)] +pub struct GenerateParameters { + #[serde(skip_serializing_if = "Option::is_none")] + pub best_of: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub decoder_input_details: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub details: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub do_sample: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub max_new_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub repetition_penalty: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub return_full_text: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub seed: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub stop: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub temperature: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub top_k: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub top_p: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub truncate: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub typical_p: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub watermark: Option, +} + +#[derive(Debug, Clone, Deserialize, Serialize, Default)] +pub struct SamplingParams { + #[serde(skip_serializing_if = "Option::is_none")] + pub temperature: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub max_new_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub top_p: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub top_k: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub frequency_penalty: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub presence_penalty: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub repetition_penalty: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub stop: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub ignore_eos: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub skip_special_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub json_schema: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub regex: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub ebnf: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub min_p: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub min_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub stop_token_ids: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub no_stop_trim: Option, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct GenerateRequest { + /// The prompt to generate from (OpenAI style) + #[serde(skip_serializing_if = "Option::is_none")] + pub prompt: Option, + + /// Text input retained for clients that use compact prompt form + #[serde(skip_serializing_if = "Option::is_none")] + pub text: Option, + + /// Input IDs for tokenized input + #[serde(skip_serializing_if = "Option::is_none")] + pub input_ids: Option, + + /// Generation parameters + #[serde(default, skip_serializing_if = "Option::is_none")] + pub parameters: Option, + + /// Sampling parameters + #[serde(skip_serializing_if = "Option::is_none")] + pub sampling_params: Option, + + /// Whether to stream the response + #[serde(default)] + pub stream: bool, + + /// Whether to return logprobs + #[serde(default)] + pub return_logprob: bool, + + /// Request ID for tracking + #[serde(skip_serializing_if = "Option::is_none")] + pub rid: Option, + + /// Pass-through fields (e.g. max_tokens, temperature, stop, stop_token_ids, + /// session_id, repetition_penalty, top_k, top_p, min_p, ignore_eos, etc.) + /// for backends like lmdeploy whose /generate endpoint accepts these as + /// top-level fields. Without this flatten map, unknown fields would be + /// silently dropped during deserialization and lost when re-serializing + /// the request to the backend. + #[serde(flatten)] + pub other: serde_json::Map, +} + +/// LMDeploy `/generate` response metadata. +/// +/// LMDeploy evolves this object independently of the router, so known token +/// accounting fields are typed while additional fields are retained verbatim. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct GenerateResponseMeta { + #[serde(skip_serializing_if = "Option::is_none")] + pub prompt_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub completion_tokens: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub finish_reason: Option, + #[serde(skip_serializing_if = "Option::is_none")] + pub output_token_logprobs: Option>, + #[serde(skip_serializing_if = "Option::is_none")] + pub routed_experts: Option, + #[serde(flatten)] + pub other: serde_json::Map, +} + +/// LMDeploy `/generate` token-in-token-out response. +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct GenerateResponse { + pub text: String, + pub output_ids: Vec, + pub meta_info: GenerateResponseMeta, + #[serde(flatten)] + pub other: serde_json::Map, +} + +impl GenerationRequest for GenerateRequest { + fn is_stream(&self) -> bool { + self.stream + } + + fn get_model(&self) -> Option<&str> { + // Generate requests typically don't have a model field + None + } + + fn extract_text_for_routing(&self) -> String { + // Check fields in priority order: text, prompt, inputs + if let Some(ref text) = self.text { + return text.clone(); + } + + if let Some(ref prompt) = self.prompt { + return match prompt { + StringOrArray::String(s) => s.clone(), + StringOrArray::Array(v) => v.join(" "), + }; + } + + if let Some(ref input_ids) = self.input_ids { + return match input_ids { + InputIds::Single(ids) => ids + .iter() + .map(|&id| id.to_string()) + .collect::>() + .join(" "), + InputIds::Batch(batches) => batches + .iter() + .flat_map(|batch| batch.iter().map(|&id| id.to_string())) + .collect::>() + .join(" "), + }; + } + + // No text input found + String::new() + } +} + +// ================================================================== +// = OPENAI SPEC - Embeddings API = +// ================================================================== + +/// Embeddings request compatible with OpenAI API +/// We intentionally keep fields flexible to pass through to workers. +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct EmbeddingRequest { + /// ID of the model to use (optional; some backends omit it) + #[serde(default, skip_serializing_if = "Option::is_none")] + pub model: Option, + + /// Input can be a string, array of strings, tokens, or batch inputs + pub input: serde_json::Value, + + /// Optional encoding format (e.g., "float", "base64") + #[serde(skip_serializing_if = "Option::is_none")] + pub encoding_format: Option, + + /// Optional user identifier + #[serde(skip_serializing_if = "Option::is_none")] + pub user: Option, + + /// Optional number of dimensions for the embedding + #[serde(skip_serializing_if = "Option::is_none")] + pub dimensions: Option, + + /// Request ID for tracking + #[serde(skip_serializing_if = "Option::is_none")] + pub rid: Option, +} + +impl GenerationRequest for EmbeddingRequest { + fn is_stream(&self) -> bool { + // Embeddings are non-streaming + false + } + + fn get_model(&self) -> Option<&str> { + self.model.as_deref() + } + + fn extract_text_for_routing(&self) -> String { + // Best effort: extract text content for routing decisions + match &self.input { + serde_json::Value::String(s) => s.clone(), + serde_json::Value::Array(arr) => arr + .iter() + .filter_map(|v| v.as_str()) + .collect::>() + .join(" "), + _ => String::new(), + } + } +} + +// ================================================================== +// = COMMON = +// ================================================================== + +/// Helper function for serde default value +pub fn default_true() -> bool { + true +} + +/// Common trait for all generation requests across different APIs +pub trait GenerationRequest: Send + Sync { + /// Check if the request is for streaming + fn is_stream(&self) -> bool; + + /// Get the model name if specified + fn get_model(&self) -> Option<&str>; + + /// Extract text content for routing decisions + fn extract_text_for_routing(&self) -> String; +} + +/// Helper type for string or array of strings +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)] +#[serde(untagged)] +pub enum StringOrArray { + String(String), + Array(Vec), +} +impl StringOrArray { + /// Get the number of items in the StringOrArray + pub fn len(&self) -> usize { + match self { + StringOrArray::String(_) => 1, + StringOrArray::Array(arr) => arr.len(), + } + } + + /// Check if the StringOrArray is empty + pub fn is_empty(&self) -> bool { + match self { + StringOrArray::String(s) => s.is_empty(), + StringOrArray::Array(arr) => arr.is_empty(), + } + } + + /// Convert to a vector of strings + pub fn to_vec(&self) -> Vec { + match self { + StringOrArray::String(s) => vec![s.clone()], + StringOrArray::Array(arr) => arr.clone(), + } + } +} + +/// Prompt input type supporting both text and token IDs +/// Compatible with the OpenAI prompt field format +#[derive(Debug, Clone, Deserialize, Serialize, PartialEq)] +#[serde(untagged)] +pub enum PromptInput { + /// Batch of token ID sequences: list[list[int]] + /// This must come first due to serde untagged matching order + IntBatch(Vec>), + /// Array of strings: list[str] + StringArray(Vec), + /// Single token ID sequence: list[int] + IntArray(Vec), + /// Single string: str + String(String), +} + +impl PromptInput { + /// Get the number of items in the PromptInput + pub fn len(&self) -> usize { + match self { + PromptInput::String(_) => 1, + PromptInput::StringArray(arr) => arr.len(), + PromptInput::IntArray(_) => 1, + PromptInput::IntBatch(batch) => batch.len(), + } + } + + /// Check if the PromptInput is empty + pub fn is_empty(&self) -> bool { + match self { + PromptInput::String(s) => s.is_empty(), + PromptInput::StringArray(arr) => arr.is_empty(), + PromptInput::IntArray(arr) => arr.is_empty(), + PromptInput::IntBatch(batch) => batch.is_empty(), + } + } + + /// Extract text representation for routing decisions + /// For token IDs, converts to a string representation + pub fn extract_text_for_routing(&self) -> String { + match self { + PromptInput::String(s) => s.clone(), + PromptInput::StringArray(arr) => arr.join(" "), + PromptInput::IntArray(ids) => { + // Convert token IDs to string representation for routing + // Format: "token_ids:" to indicate this is a token-based prompt + format!("token_ids:{}", ids.len()) + } + PromptInput::IntBatch(batches) => { + // For batches, use total token count + let total_tokens: usize = batches.iter().map(|b| b.len()).sum(); + format!("token_ids_batch:{}:{}", batches.len(), total_tokens) + } + } + } + + /// Check if this prompt is token-based (not text) + pub fn is_token_based(&self) -> bool { + matches!(self, PromptInput::IntArray(_) | PromptInput::IntBatch(_)) + } + + /// Get the total token count if token-based, otherwise estimate from text + pub fn estimated_token_count(&self) -> usize { + match self { + PromptInput::String(s) => { + // Rough estimate: ~4 chars per token on average + s.len() / 4 + } + PromptInput::StringArray(arr) => { + // Sum estimated tokens for all strings + arr.iter().map(|s| s.len() / 4).sum() + } + PromptInput::IntArray(ids) => ids.len(), + PromptInput::IntBatch(batches) => batches.iter().map(|b| b.len()).sum(), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + // ================================================================== + // = EMBEDDINGS REQUEST TESTS = + // ================================================================== + + #[test] + fn test_embedding_request_serialization_string_input() { + let req = EmbeddingRequest { + model: Some("test-emb".to_string()), + input: serde_json::Value::String("hello".to_string()), + encoding_format: Some("float".to_string()), + user: Some("user-1".to_string()), + dimensions: Some(128), + rid: Some("rid-123".to_string()), + }; + + let serialized = serde_json::to_string(&req).unwrap(); + let deserialized: EmbeddingRequest = serde_json::from_str(&serialized).unwrap(); + + assert_eq!(deserialized.model, req.model); + assert_eq!(deserialized.input, req.input); + assert_eq!(deserialized.encoding_format, req.encoding_format); + assert_eq!(deserialized.user, req.user); + assert_eq!(deserialized.dimensions, req.dimensions); + assert_eq!(deserialized.rid, req.rid); + } + + #[test] + fn test_embedding_request_serialization_array_input() { + let req = EmbeddingRequest { + model: Some("test-emb".to_string()), + input: serde_json::json!(["a", "b", "c"]), + encoding_format: None, + user: None, + dimensions: None, + rid: None, + }; + + let serialized = serde_json::to_string(&req).unwrap(); + let de: EmbeddingRequest = serde_json::from_str(&serialized).unwrap(); + assert_eq!(de.model, req.model); + assert_eq!(de.input, req.input); + } + + #[test] + fn test_embedding_generation_request_trait_string() { + let req = EmbeddingRequest { + model: Some("emb-model".to_string()), + input: serde_json::Value::String("hello".to_string()), + encoding_format: None, + user: None, + dimensions: None, + rid: None, + }; + assert!(!req.is_stream()); + assert_eq!(req.get_model(), Some("emb-model")); + assert_eq!(req.extract_text_for_routing(), "hello"); + } + + #[test] + fn test_embedding_generation_request_trait_array() { + let req = EmbeddingRequest { + model: Some("emb-model".to_string()), + input: serde_json::json!(["hello", "world"]), + encoding_format: None, + user: None, + dimensions: None, + rid: None, + }; + assert_eq!(req.extract_text_for_routing(), "hello world"); + } + + #[test] + fn test_embedding_generation_request_trait_non_text() { + let req = EmbeddingRequest { + model: Some("emb-model".to_string()), + input: serde_json::json!({"tokens": [1, 2, 3]}), + encoding_format: None, + user: None, + dimensions: None, + rid: None, + }; + assert_eq!(req.extract_text_for_routing(), ""); + } + + #[test] + fn test_embedding_generation_request_trait_mixed_array_ignores_nested() { + let req = EmbeddingRequest { + model: Some("emb-model".to_string()), + input: serde_json::json!(["a", ["b", "c"], 123, {"k": "v"}]), + encoding_format: None, + user: None, + dimensions: None, + rid: None, + }; + // Only top-level string elements are extracted + assert_eq!(req.extract_text_for_routing(), "a"); + } + + // ================================================================== + // = OPTIONAL MODEL FIELD TESTS = + // ================================================================== + + #[test] + fn test_chat_completion_request_without_model() { + let json = r#"{ + "messages": [{"role": "user", "content": "Hello"}] + }"#; + + let request: ChatCompletionRequest = serde_json::from_str(json).unwrap(); + assert!(request.model.is_none()); + assert_eq!(request.get_model(), None); + } + + #[test] + fn test_chat_completion_request_with_model() { + let json = r#"{ + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}] + }"#; + + let request: ChatCompletionRequest = serde_json::from_str(json).unwrap(); + assert_eq!(request.model.as_deref(), Some("gpt-4")); + assert_eq!(request.get_model(), Some("gpt-4")); + } + + #[test] + fn test_chat_completion_request_with_null_model() { + let json = r#"{ + "model": null, + "messages": [{"role": "user", "content": "Hello"}] + }"#; + + let request: ChatCompletionRequest = serde_json::from_str(json).unwrap(); + assert!(request.model.is_none()); + assert_eq!(request.get_model(), None); + } + + #[test] + fn test_chat_completion_request_without_model_roundtrip() { + let json = r#"{ + "messages": [{"role": "user", "content": "Hello"}] + }"#; + + let request: ChatCompletionRequest = serde_json::from_str(json).unwrap(); + let serialized = serde_json::to_string(&request).unwrap(); + // model should not appear in serialized output when None + assert!(!serialized.contains("\"model\"")); + + let roundtrip: ChatCompletionRequest = serde_json::from_str(&serialized).unwrap(); + assert!(roundtrip.model.is_none()); + } + + #[test] + fn test_chat_completion_request_without_model_validates() { + use crate::protocols::validation::ValidatableRequest; + + let json = r#"{ + "messages": [{"role": "user", "content": "Hello"}] + }"#; + + let request: ChatCompletionRequest = serde_json::from_str(json).unwrap(); + assert!(request.model.is_none()); + assert!(request.validate().is_ok()); + } + + #[test] + fn test_completion_request_without_model() { + let json = r#"{ + "prompt": "Hello, world!" + }"#; + + let request: CompletionRequest = serde_json::from_str(json).unwrap(); + assert!(request.model.is_none()); + assert_eq!(request.get_model(), None); + } + + #[test] + fn test_completion_request_with_model() { + let json = r#"{ + "model": "text-davinci-003", + "prompt": "Hello, world!" + }"#; + + let request: CompletionRequest = serde_json::from_str(json).unwrap(); + assert_eq!(request.model.as_deref(), Some("text-davinci-003")); + assert_eq!(request.get_model(), Some("text-davinci-003")); + } + + #[test] + fn test_completion_request_without_model_roundtrip() { + let json = r#"{ + "prompt": "Hello, world!" + }"#; + + let request: CompletionRequest = serde_json::from_str(json).unwrap(); + let serialized = serde_json::to_string(&request).unwrap(); + assert!(!serialized.contains("\"model\"")); + + let roundtrip: CompletionRequest = serde_json::from_str(&serialized).unwrap(); + assert!(roundtrip.model.is_none()); + } + + #[test] + fn test_embedding_request_without_model() { + let json = r#"{ + "input": "Hello, world!" + }"#; + + let request: EmbeddingRequest = serde_json::from_str(json).unwrap(); + assert!(request.model.is_none()); + assert_eq!(request.get_model(), None); + } + + #[test] + fn test_embedding_request_with_model() { + let json = r#"{ + "model": "text-embedding-ada-002", + "input": "Hello, world!" + }"#; + + let request: EmbeddingRequest = serde_json::from_str(json).unwrap(); + assert_eq!(request.model.as_deref(), Some("text-embedding-ada-002")); + assert_eq!(request.get_model(), Some("text-embedding-ada-002")); + } + + #[test] + fn test_embedding_request_without_model_roundtrip() { + let json = r#"{ + "input": "Hello, world!" + }"#; + + let request: EmbeddingRequest = serde_json::from_str(json).unwrap(); + let serialized = serde_json::to_string(&request).unwrap(); + assert!(!serialized.contains("\"model\"")); + + let roundtrip: EmbeddingRequest = serde_json::from_str(&serialized).unwrap(); + assert!(roundtrip.model.is_none()); + } + + // ================================================================== + // = STRUCTURED OUTPUTS PARAMS TESTS = + // ================================================================== + + #[test] + fn test_structured_outputs_params_default() { + let params = StructuredOutputsParams::default(); + + assert!(params.json.is_none()); + assert!(params.regex.is_none()); + assert!(params.choice.is_none()); + assert!(params.grammar.is_none()); + assert!(params.json_object.is_none()); + assert!(params.structural_tag.is_none()); + assert!(params.disable_fallback.is_none()); + assert!(params.disable_any_whitespace.is_none()); + assert!(params.disable_additional_properties.is_none()); + assert!(params.whitespace_pattern.is_none()); + } + + #[test] + fn test_structured_outputs_params_with_json_schema() { + let schema = serde_json::json!({ + "type": "object", + "properties": { + "name": {"type": "string"}, + "age": {"type": "integer"} + }, + "required": ["name", "age"] + }); + + let params = StructuredOutputsParams { + json: Some(schema.clone()), + ..Default::default() + }; + + assert_eq!(params.json, Some(schema)); + assert!(params.regex.is_none()); + } + + #[test] + fn test_structured_outputs_params_with_regex() { + let params = StructuredOutputsParams { + regex: Some(r"^\d{3}-\d{2}-\d{4}$".to_string()), + ..Default::default() + }; + + assert!(params.json.is_none()); + assert_eq!(params.regex, Some(r"^\d{3}-\d{2}-\d{4}$".to_string())); + } + + #[test] + fn test_structured_outputs_params_with_choice() { + let choices = vec!["yes".to_string(), "no".to_string(), "maybe".to_string()]; + + let params = StructuredOutputsParams { + choice: Some(choices.clone()), + ..Default::default() + }; + + assert_eq!(params.choice, Some(choices)); + } + + #[test] + fn test_structured_outputs_params_with_grammar() { + let grammar = r#" + root ::= sentence + sentence ::= subject " " verb " " object + subject ::= "I" | "You" | "They" + verb ::= "eat" | "drink" | "see" + object ::= "apple" | "water" | "sky" + "# + .to_string(); + + let params = StructuredOutputsParams { + grammar: Some(grammar.clone()), + ..Default::default() + }; + + assert_eq!(params.grammar, Some(grammar)); + } + + #[test] + fn test_structured_outputs_params_with_json_object_mode() { + let params = StructuredOutputsParams { + json_object: Some(true), + ..Default::default() + }; + + assert_eq!(params.json_object, Some(true)); + } + + #[test] + fn test_structured_outputs_params_with_structural_tag() { + let params = StructuredOutputsParams { + structural_tag: Some("".to_string()), + ..Default::default() + }; + + assert_eq!(params.structural_tag, Some("".to_string())); + } + + #[test] + fn test_structured_outputs_params_with_all_flags() { + let params = StructuredOutputsParams { + disable_fallback: Some(true), + disable_any_whitespace: Some(true), + disable_additional_properties: Some(false), + whitespace_pattern: Some(r"\s*".to_string()), + ..Default::default() + }; + + assert_eq!(params.disable_fallback, Some(true)); + assert_eq!(params.disable_any_whitespace, Some(true)); + assert_eq!(params.disable_additional_properties, Some(false)); + assert_eq!(params.whitespace_pattern, Some(r"\s*".to_string())); + } + + #[test] + fn test_structured_outputs_params_serialization_empty() { + let params = StructuredOutputsParams::default(); + + let serialized = serde_json::to_string(¶ms).unwrap(); + // All fields should be skipped when None + assert_eq!(serialized, "{}"); + } + + #[test] + fn test_structured_outputs_params_deserialization_partial() { + let json = r#"{ + "regex": "^[a-z]+$", + "disable_fallback": true + }"#; + + let params: StructuredOutputsParams = serde_json::from_str(json).unwrap(); + + assert!(params.json.is_none()); + assert_eq!(params.regex, Some("^[a-z]+$".to_string())); + assert!(params.choice.is_none()); + assert!(params.grammar.is_none()); + assert!(params.json_object.is_none()); + assert!(params.structural_tag.is_none()); + assert_eq!(params.disable_fallback, Some(true)); + assert!(params.disable_any_whitespace.is_none()); + assert!(params.disable_additional_properties.is_none()); + assert!(params.whitespace_pattern.is_none()); + } + + #[test] + fn test_structured_outputs_params_complex_json_schema() { + let schema = serde_json::json!({ + "type": "object", + "properties": { + "user": { + "type": "object", + "properties": { + "name": {"type": "string", "minLength": 1}, + "email": {"type": "string", "format": "email"}, + "age": {"type": "integer", "minimum": 0, "maximum": 150} + }, + "required": ["name", "email"] + }, + "items": { + "type": "array", + "items": { + "type": "object", + "properties": { + "id": {"type": "integer"}, + "value": {"type": "number"} + } + } + } + }, + "required": ["user"] + }); + + let params = StructuredOutputsParams { + json: Some(schema.clone()), + disable_additional_properties: Some(true), + ..Default::default() + }; + + let serialized = serde_json::to_string(¶ms).unwrap(); + let deserialized: StructuredOutputsParams = serde_json::from_str(&serialized).unwrap(); + + assert_eq!(deserialized.json, Some(schema)); + assert_eq!(deserialized.disable_additional_properties, Some(true)); + } + + #[test] + fn test_structured_outputs_params_empty_choice_array() { + let params = StructuredOutputsParams { + choice: Some(vec![]), + ..Default::default() + }; + + let serialized = serde_json::to_string(¶ms).unwrap(); + let deserialized: StructuredOutputsParams = serde_json::from_str(&serialized).unwrap(); + + assert_eq!(deserialized.choice, Some(vec![])); + } + + #[test] + fn test_structured_outputs_params_in_chat_completion_request() { + let json = r#"{ + "model": "test-model", + "messages": [{"role": "user", "content": "Generate JSON"}], + "structured_outputs": { + "json": {"type": "object"}, + "disable_fallback": true + } + }"#; + + let request: ChatCompletionRequest = serde_json::from_str(json).unwrap(); + + assert!(request.structured_outputs.is_some()); + let structured_outputs = request.structured_outputs.unwrap(); + assert!(structured_outputs.json.is_some()); + assert_eq!(structured_outputs.disable_fallback, Some(true)); + } + + #[test] + fn test_structured_outputs_params_in_chat_completion_request_none() { + let json = r#"{ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}] + }"#; + + let request: ChatCompletionRequest = serde_json::from_str(json).unwrap(); + + assert!(request.structured_outputs.is_none()); + } + + // ================================================================== + // = CHAT MESSAGE DESERIALIZATION TESTS = + // ================================================================== + + #[test] + fn test_chat_message_assistant_with_reasoning() { + let json = r#"{ + "role": "assistant", + "content": "Hello there!", + "reasoning": "Let me think about how to greet the user..." + }"#; + + let message: ChatMessage = serde_json::from_str(json).unwrap(); + + match message { + ChatMessage::Assistant { + content, reasoning, .. + } => { + assert_eq!(content.as_ref().unwrap(), "Hello there!"); + assert_eq!( + reasoning.as_ref().unwrap(), + "Let me think about how to greet the user..." + ); + } + other => panic!( + "Expected Assistant message but got {:?}", + std::mem::discriminant(&other) + ), + } + } + + #[test] + fn test_chat_message_assistant_with_null_content_and_reasoning() { + let json = r#"{ + "role": "assistant", + "content": null, + "reasoning": "Deep thinking in progress..." + }"#; + + let message: ChatMessage = serde_json::from_str(json).unwrap(); + + match message { + ChatMessage::Assistant { + content, reasoning, .. + } => { + assert!(content.is_none()); + assert_eq!(reasoning.as_ref().unwrap(), "Deep thinking in progress..."); + } + _ => panic!("Expected Assistant message"), + } + } + + #[test] + fn test_chat_message_assistant_with_tool_calls() { + let json = r#"{ + "role": "assistant", + "content": null, + "tool_calls": [{ + "id": "call_abc123", + "type": "function", + "function": { + "name": "get_weather", + "arguments": "{\"location\": \"Boston\"}" + } + }] + }"#; + + let message: ChatMessage = serde_json::from_str(json).unwrap(); + + match message { + ChatMessage::Assistant { tool_calls, .. } => { + let calls = tool_calls.unwrap(); + assert_eq!(calls.len(), 1); + assert_eq!(calls[0].id, "call_abc123"); + assert_eq!(calls[0].function.name, "get_weather"); + } + _ => panic!("Expected Assistant message"), + } + } + + #[test] + fn test_chat_message_system() { + let json = r#"{ + "role": "system", + "content": "You are a helpful assistant." + }"#; + + let message: ChatMessage = serde_json::from_str(json).unwrap(); + + match message { + ChatMessage::System { content, .. } => { + assert_eq!(content, "You are a helpful assistant."); + } + _ => panic!("Expected System message"), + } + } + + #[test] + fn test_chat_message_user_text() { + let json = r#"{ + "role": "user", + "content": "Hello!" + }"#; + + let message: ChatMessage = serde_json::from_str(json).unwrap(); + + match message { + ChatMessage::User { content, .. } => match content { + UserMessageContent::Text(text) => assert_eq!(text, "Hello!"), + _ => panic!("Expected Text content"), + }, + _ => panic!("Expected User message"), + } + } + + #[test] + fn test_chat_message_tool() { + let json = r#"{ + "role": "tool", + "content": "Tool result here", + "tool_call_id": "call_123" + }"#; + + let message: ChatMessage = serde_json::from_str(json).unwrap(); + + match message { + ChatMessage::Tool { + content, + tool_call_id, + .. + } => { + assert_eq!(content, Value::String("Tool result here".to_string())); + assert_eq!(tool_call_id, "call_123"); + } + _ => panic!("Expected Tool message"), + } + } + + #[test] + fn test_chat_message_tool_preserves_array_content() { + let json = r#"{ + "role": "tool", + "content": [{"type": "text", "text": "1 files found: ['wi_example.docx']"}], + "tool_call_id": "call_456" + }"#; + + let message: ChatMessage = serde_json::from_str(json).unwrap(); + + match message { + ChatMessage::Tool { + content, + tool_call_id, + .. + } => { + let expected = serde_json::json!([ + {"type": "text", "text": "1 files found: ['wi_example.docx']"} + ]); + assert_eq!(content, expected); + assert_eq!(tool_call_id, "call_456"); + } + _ => panic!("Expected Tool message"), + } + } + + #[test] + fn test_chat_message_function() { + let json = r#"{ + "role": "function", + "content": "Function result", + "name": "my_function" + }"#; + + let message: ChatMessage = serde_json::from_str(json).unwrap(); + + match message { + ChatMessage::Function { content, name, .. } => { + assert_eq!(content, "Function result"); + assert_eq!(name, "my_function"); + } + _ => panic!("Expected Function message"), + } + } + + #[test] + fn test_chat_message_roundtrip_serialization() { + let original = ChatMessage::Assistant { + role: "assistant".to_string(), + content: Some("Hello!".to_string()), + name: None, + tool_calls: None, + function_call: None, + reasoning: Some("Thinking...".to_string()), + }; + + let serialized = serde_json::to_string(&original).unwrap(); + let deserialized: ChatMessage = serde_json::from_str(&serialized).unwrap(); + + match deserialized { + ChatMessage::Assistant { + content, reasoning, .. + } => { + assert_eq!(content.as_ref().unwrap(), "Hello!"); + assert_eq!(reasoning.as_ref().unwrap(), "Thinking..."); + } + _ => panic!("Expected Assistant message"), + } + } + + #[test] + fn test_chat_message_assistant_reasoning_content_alias_roundtrip() { + // Regression: assistant turns using the deprecated `reasoning_content` + // alias must survive forwarding and normalize to canonical `reasoning`. + let json = r#"{ + "role": "assistant", + "content": "The answer is 42.", + "reasoning_content": "First I considered X, then Y..." + }"#; + + let message: ChatMessage = serde_json::from_str(json).unwrap(); + match &message { + ChatMessage::Assistant { reasoning, .. } => { + assert_eq!( + reasoning.as_deref(), + Some("First I considered X, then Y...") + ); + } + _ => panic!("Expected Assistant message"), + } + + // Re-serialization emits only the canonical field. + let serialized = serde_json::to_value(&message).unwrap(); + assert_eq!( + serialized.get("reasoning").and_then(|v| v.as_str()), + Some("First I considered X, then Y..."), + ); + assert!(serialized.get("reasoning_content").is_none()); + + let reparsed: ChatMessage = serde_json::from_value(serialized).unwrap(); + match reparsed { + ChatMessage::Assistant { reasoning, .. } => { + assert_eq!( + reasoning.as_deref(), + Some("First I considered X, then Y...") + ); + } + _ => panic!("Expected Assistant message"), + } + } + + #[test] + fn test_chat_message_assistant_reasoning_alias_precedence() { + let cases = [ + ( + r#"{ + "role": "assistant", + "reasoning": "canonical", + "reasoning_content": "deprecated alias" + }"#, + "canonical", + ), + ( + r#"{ + "role": "assistant", + "reasoning": null, + "reasoning_content": "deprecated alias" + }"#, + "deprecated alias", + ), + ( + r#"{ + "role": "assistant", + "reasoning": 42, + "reasoning_content": "deprecated alias" + }"#, + "deprecated alias", + ), + ]; + + for (json, expected) in cases { + let message: ChatMessage = serde_json::from_str(json).unwrap(); + match message { + ChatMessage::Assistant { reasoning, .. } => { + assert_eq!(reasoning.as_deref(), Some(expected)); + } + _ => panic!("Expected Assistant message"), + } + } + } + + // ===== lmdeploy token-in-token-out routing tests ===== + + #[test] + fn test_chat_completion_extract_routing_with_input_ids() { + // Build a ChatCompletionRequest via JSON deserialization (lmdeploy-style: + // empty messages + input_ids fallback). + let body = serde_json::json!({ + "model": "test-model", + "messages": [], + "input_ids": [151644, 8948, 198, 2610], + }); + let req: ChatCompletionRequest = + serde_json::from_value(body).expect("deserialize chat request"); + assert_eq!(req.extract_text_for_routing(), "151644 8948 198 2610"); + } + + #[test] + fn test_chat_completion_extract_routing_ignores_wrapped_session_id() { + let body = serde_json::json!({ + "model": "test-model", + "messages": [], + "input_ids": [1, 2, 3], + "session_id": "sess-abc" + }); + let req: ChatCompletionRequest = + serde_json::from_value(body).expect("deserialize chat request"); + assert_eq!(req.extract_text_for_routing(), "1 2 3"); + } + + #[test] + fn test_chat_completion_extract_routing_empty_when_no_keys() { + let body = serde_json::json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "hello"}], + }); + let req: ChatCompletionRequest = + serde_json::from_value(body).expect("deserialize chat request"); + assert_eq!(req.extract_text_for_routing(), ""); + } + + #[test] + fn test_chat_completion_rejects_non_lmdeploy_input_ids_shapes() { + for input_ids in [ + serde_json::json!([[10, 20], [30, 40]]), + serde_json::json!([10, "invalid", 20]), + serde_json::json!("not-an-array"), + ] { + let body = serde_json::json!({ + "model": "test-model", + "messages": [], + "input_ids": input_ids, + }); + assert!(serde_json::from_value::(body).is_err()); + } + } + + #[test] + fn test_chat_completion_preserves_typed_input_ids_on_serialize() { + // Ensure input_ids survives round-trip serialization (forwarded to backend). + let body = serde_json::json!({ + "model": "test-model", + "messages": [], + "input_ids": [1, 2, 3], + "max_tokens": 16, + }); + let req: ChatCompletionRequest = + serde_json::from_value(body).expect("deserialize chat request"); + assert_eq!(req.input_ids, Some(vec![1, 2, 3])); + assert!(!req.other.contains_key("input_ids")); + let out = serde_json::to_value(&req).expect("serialize chat request"); + assert_eq!(out["input_ids"], serde_json::json!([1, 2, 3])); + } + + #[test] + fn test_chat_completion_preserves_all_lmdeploy_token_input_fields() { + let body = serde_json::json!({ + "model": "test-model", + "messages": [], + "input_ids": [1, 2, 3], + "image_data": ["data:image/png;base64,AAAA"], + "do_preprocess": false, + "return_token_ids": true, + }); + let req: ChatCompletionRequest = + serde_json::from_value(body.clone()).expect("deserialize chat request"); + let out = serde_json::to_value(&req).expect("serialize chat request"); + for field in [ + "input_ids", + "image_data", + "do_preprocess", + "return_token_ids", + ] { + assert_eq!(out[field], body[field], "field {field} must survive"); + } + } + + #[test] + fn test_generate_request_preserves_lmdeploy_token_input_fields() { + let body = serde_json::json!({ + "input_ids": [151644, 8948, 198], + "session_id": 42, + "max_tokens": 16, + "return_logprob": true, + "stream": false, + }); + let req: GenerateRequest = + serde_json::from_value(body.clone()).expect("deserialize generate request"); + assert_eq!(req.extract_text_for_routing(), "151644 8948 198"); + + let out = serde_json::to_value(&req).expect("serialize generate request"); + for field in ["input_ids", "session_id", "max_tokens", "return_logprob"] { + assert_eq!(out[field], body[field], "field {field} must survive"); + } + } + + #[test] + fn test_lmdeploy_chat_response_preserves_output_ids() { + let body = serde_json::json!({ + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 1, + "model": "test-model", + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + "output_ids": [100, 101], + "output_token_logprobs": [[-0.1, 100]], + "routed_experts": [[[1, 2]]] + }], + "usage": { + "prompt_tokens": 3, + "completion_tokens": 2, + "total_tokens": 5 + } + }); + let response: ChatCompletionResponse = + serde_json::from_value(body.clone()).expect("deserialize chat response"); + assert_eq!(response.choices[0].output_ids, Some(vec![100, 101])); + assert!(response.choices[0] + .other + .contains_key("output_token_logprobs")); + assert!(response.choices[0].other.contains_key("routed_experts")); + + let out = serde_json::to_value(response).expect("serialize chat response"); + assert_eq!( + out["choices"][0]["output_ids"], + body["choices"][0]["output_ids"] + ); + assert_eq!( + out["choices"][0]["output_token_logprobs"], + body["choices"][0]["output_token_logprobs"] + ); + } + + #[test] + fn test_lmdeploy_chat_stream_response_preserves_output_ids() { + let body = serde_json::json!({ + "id": "chatcmpl-1", + "object": "chat.completion.chunk", + "created": 1, + "model": "test-model", + "choices": [{ + "index": 0, + "delta": {"role": "assistant", "content": "ok"}, + "finish_reason": null, + "output_ids": [100, 101] + }] + }); + let response: ChatCompletionStreamResponse = + serde_json::from_value(body).expect("deserialize chat stream response"); + assert_eq!(response.choices[0].output_ids, Some(vec![100, 101])); + } + + #[test] + fn test_lmdeploy_generate_response_preserves_output_ids() { + let body = serde_json::json!({ + "text": "ok", + "output_ids": [100, 101], + "meta_info": { + "prompt_tokens": 3, + "completion_tokens": 2, + "finish_reason": {"type": "stop"}, + "output_token_logprobs": [[-0.1, 100]], + "first_token_latency": 0.01 + } + }); + let response: GenerateResponse = + serde_json::from_value(body.clone()).expect("deserialize generate response"); + assert_eq!(response.output_ids, vec![100, 101]); + assert_eq!(response.meta_info.prompt_tokens, Some(3)); + assert!(response.meta_info.other.contains_key("first_token_latency")); + + let out = serde_json::to_value(response).expect("serialize generate response"); + assert_eq!(out["output_ids"], body["output_ids"]); + assert_eq!(out["meta_info"]["first_token_latency"], 0.01); + } +} diff --git a/lmdeploy/router/src/protocols/validation.rs b/lmdeploy/router/src/protocols/validation.rs new file mode 100644 index 0000000000..ab2cd87bdc --- /dev/null +++ b/lmdeploy/router/src/protocols/validation.rs @@ -0,0 +1,1264 @@ +// Core validation infrastructure for API parameter validation + +use anyhow::Result; +use serde::{Deserialize, Serialize}; +use std::fmt::Display; + +// Import types from spec module +use crate::protocols::spec::{ + ChatCompletionRequest, ChatMessage, ResponseFormat, StringOrArray, UserMessageContent, +}; + +/// Validation constants for OpenAI API parameters +pub mod constants { + /// Temperature range: 0.0 to 2.0 (OpenAI spec) + pub const TEMPERATURE_RANGE: (f32, f32) = (0.0, 2.0); + + /// Top-p range: 0.0 to 1.0 (exclusive of 0.0) + pub const TOP_P_RANGE: (f32, f32) = (0.0, 1.0); + + /// Presence penalty range: -2.0 to 2.0 (OpenAI spec) + pub const PRESENCE_PENALTY_RANGE: (f32, f32) = (-2.0, 2.0); + + /// Frequency penalty range: -2.0 to 2.0 (OpenAI spec) + pub const FREQUENCY_PENALTY_RANGE: (f32, f32) = (-2.0, 2.0); + + /// Logprobs range for completions API: 0 to 5 + pub const LOGPROBS_RANGE: (u32, u32) = (0, 5); + + /// Top logprobs range for chat completions: 0 to 20 + pub const TOP_LOGPROBS_RANGE: (u32, u32) = (0, 20); + + /// Maximum number of stop sequences allowed + pub const MAX_STOP_SEQUENCES: usize = 4; + + /// Generation extension validation constants + pub mod generation_extensions { + /// Min-p range: 0.0 to 1.0 + pub const MIN_P_RANGE: (f32, f32) = (0.0, 1.0); + + /// Top-k minimum value: -1 to disable, otherwise positive + pub const TOP_K_MIN: i32 = -1; + + /// Repetition penalty range: 0.0 to 2.0 + /// 1.0 = no penalty, >1.0 = discourage repetition, <1.0 = encourage repetition + pub const REPETITION_PENALTY_RANGE: (f32, f32) = (0.0, 2.0); + } +} + +/// Core validation error types +#[derive(Debug, Clone, Serialize, Deserialize)] +pub enum ValidationError { + /// Parameter value out of valid range + OutOfRange { + parameter: String, + value: String, + min: String, + max: String, + }, + /// Invalid parameter value format or type + InvalidValue { + parameter: String, + value: String, + reason: String, + }, + /// Cross-parameter validation failure + ConflictingParameters { + parameter1: String, + parameter2: String, + reason: String, + }, + /// Required parameter missing + MissingRequired { parameter: String }, + /// Too many items in array parameter + TooManyItems { + parameter: String, + count: usize, + max: usize, + }, + /// Custom validation error + Custom(String), +} + +impl Display for ValidationError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + ValidationError::OutOfRange { + parameter, + value, + min, + max, + } => { + write!( + f, + "Parameter '{}' must be between {} and {}, got {}", + parameter, min, max, value + ) + } + ValidationError::InvalidValue { + parameter, + value, + reason, + } => { + write!( + f, + "Invalid value for parameter '{}': {} ({})", + parameter, value, reason + ) + } + ValidationError::ConflictingParameters { + parameter1, + parameter2, + reason, + } => { + write!( + f, + "Conflicting parameters '{}' and '{}': {}", + parameter1, parameter2, reason + ) + } + ValidationError::MissingRequired { parameter } => { + write!(f, "Required parameter '{}' is missing", parameter) + } + ValidationError::TooManyItems { + parameter, + count, + max, + } => { + write!( + f, + "Parameter '{}' has too many items: {} (maximum: {})", + parameter, count, max + ) + } + ValidationError::Custom(msg) => write!(f, "{}", msg), + } + } +} + +impl std::error::Error for ValidationError {} + +/// Core validation utility functions +pub mod utils { + use super::*; + + /// Validate that a numeric value is within the specified range (inclusive) + pub fn validate_range( + value: T, + range: &(T, T), + param_name: &str, + ) -> Result + where + T: PartialOrd + Display + Copy, + { + if value >= range.0 && value <= range.1 { + Ok(value) + } else { + Err(ValidationError::OutOfRange { + parameter: param_name.to_string(), + value: value.to_string(), + min: range.0.to_string(), + max: range.1.to_string(), + }) + } + } + + /// Validate that a positive number is actually positive + pub fn validate_positive(value: T, param_name: &str) -> Result + where + T: PartialOrd + Display + Copy + Default, + { + if value > T::default() { + Ok(value) + } else { + Err(ValidationError::InvalidValue { + parameter: param_name.to_string(), + value: value.to_string(), + reason: "must be positive".to_string(), + }) + } + } + + /// Validate that an array doesn't exceed maximum length + pub fn validate_max_items( + items: &[T], + max_count: usize, + param_name: &str, + ) -> Result<(), ValidationError> { + if items.len() <= max_count { + Ok(()) + } else { + Err(ValidationError::TooManyItems { + parameter: param_name.to_string(), + count: items.len(), + max: max_count, + }) + } + } + + /// Validate that a required parameter is present + pub fn validate_required<'a, T>( + value: &'a Option, + param_name: &str, + ) -> Result<&'a T, ValidationError> { + value + .as_ref() + .ok_or_else(|| ValidationError::MissingRequired { + parameter: param_name.to_string(), + }) + } + + /// Validate top_k parameter + pub fn validate_top_k(top_k: i32) -> Result { + if top_k == constants::generation_extensions::TOP_K_MIN || top_k > 0 { + Ok(top_k) + } else { + Err(ValidationError::InvalidValue { + parameter: "top_k".to_string(), + value: top_k.to_string(), + reason: "must be -1 (disabled) or positive".to_string(), + }) + } + } + + /// Generic validation function for sampling options + pub fn validate_sampling_options( + request: &T, + ) -> Result<(), ValidationError> { + // Validate temperature (0.0 to 2.0) + if let Some(temp) = request.get_temperature() { + validate_range(temp, &constants::TEMPERATURE_RANGE, "temperature")?; + } + + // Validate top_p (0.0 to 1.0) + if let Some(top_p) = request.get_top_p() { + validate_range(top_p, &constants::TOP_P_RANGE, "top_p")?; + } + + // Validate frequency_penalty (-2.0 to 2.0) + if let Some(freq_penalty) = request.get_frequency_penalty() { + validate_range( + freq_penalty, + &constants::FREQUENCY_PENALTY_RANGE, + "frequency_penalty", + )?; + } + + // Validate presence_penalty (-2.0 to 2.0) + if let Some(pres_penalty) = request.get_presence_penalty() { + validate_range( + pres_penalty, + &constants::PRESENCE_PENALTY_RANGE, + "presence_penalty", + )?; + } + + Ok(()) + } + + /// Generic validation function for stop conditions + pub fn validate_stop_conditions( + request: &T, + ) -> Result<(), ValidationError> { + if let Some(stop) = request.get_stop_sequences() { + match stop { + StringOrArray::String(s) => { + if s.is_empty() { + return Err(ValidationError::InvalidValue { + parameter: "stop".to_string(), + value: "empty string".to_string(), + reason: "stop sequences cannot be empty".to_string(), + }); + } + } + StringOrArray::Array(arr) => { + validate_max_items(arr, constants::MAX_STOP_SEQUENCES, "stop")?; + for (i, s) in arr.iter().enumerate() { + if s.is_empty() { + return Err(ValidationError::InvalidValue { + parameter: format!("stop[{}]", i), + value: "empty string".to_string(), + reason: "stop sequences cannot be empty".to_string(), + }); + } + } + } + } + } + + Ok(()) + } + + /// Generic validation function for token limits + pub fn validate_token_limits( + request: &T, + ) -> Result<(), ValidationError> { + // Validate max_tokens if provided + if let Some(max_tokens) = request.get_max_tokens() { + validate_positive(max_tokens, "max_tokens")?; + } + + // Validate min_tokens if provided + if let Some(min_tokens) = request.get_min_tokens() { + validate_positive(min_tokens, "min_tokens")?; + } + + Ok(()) + } + + /// Generic validation function for logprobs + pub fn validate_logprobs( + request: &T, + ) -> Result<(), ValidationError> { + // Validate logprobs (completions API - 0 to 5) + if let Some(logprobs) = request.get_logprobs() { + validate_range(logprobs, &constants::LOGPROBS_RANGE, "logprobs")?; + } + + // Validate top_logprobs (chat API - 0 to 20) + if let Some(top_logprobs) = request.get_top_logprobs() { + validate_range(top_logprobs, &constants::TOP_LOGPROBS_RANGE, "top_logprobs")?; + } + + Ok(()) + } + + /// Generic cross-parameter validation + pub fn validate_cross_parameters( + request: &T, + ) -> Result<(), ValidationError> { + // Check min_tokens <= max_tokens if both are specified + if let (Some(min_tokens), Some(max_tokens)) = + (request.get_min_tokens(), request.get_max_tokens()) + { + if min_tokens > max_tokens { + return Err(ValidationError::ConflictingParameters { + parameter1: "min_tokens".to_string(), + parameter2: "max_tokens".to_string(), + reason: format!( + "min_tokens ({}) cannot be greater than max_tokens ({})", + min_tokens, max_tokens + ), + }); + } + } + + Ok(()) + } + + /// Validate conflicting structured output constraints + pub fn validate_conflicting_parameters( + param1_name: &str, + param1_value: bool, + param2_name: &str, + param2_value: bool, + reason: &str, + ) -> Result<(), ValidationError> { + if param1_value && param2_value { + return Err(ValidationError::ConflictingParameters { + parameter1: param1_name.to_string(), + parameter2: param2_name.to_string(), + reason: reason.to_string(), + }); + } + Ok(()) + } + + /// Validate that only one option from a set is active + pub fn validate_mutually_exclusive_options( + options: &[(&str, bool)], + error_msg: &str, + ) -> Result<(), ValidationError> { + let active_count = options.iter().filter(|(_, is_active)| *is_active).count(); + if active_count > 1 { + return Err(ValidationError::Custom(error_msg.to_string())); + } + Ok(()) + } + + /// Generic validation for generation extensions + pub fn validate_generation_extensions( + request: &T, + ) -> Result<(), ValidationError> { + // Validate top_k (-1 to disable, or positive) + if let Some(top_k) = request.get_top_k() { + validate_top_k(top_k)?; + } + + // Validate min_p (0.0 to 1.0) + if let Some(min_p) = request.get_min_p() { + validate_range( + min_p, + &constants::generation_extensions::MIN_P_RANGE, + "min_p", + )?; + } + + // Validate repetition_penalty (0.0 to 2.0) + if let Some(rep_penalty) = request.get_repetition_penalty() { + validate_range( + rep_penalty, + &constants::generation_extensions::REPETITION_PENALTY_RANGE, + "repetition_penalty", + )?; + } + + Ok(()) + } + + /// Generic validation for n parameter (number of completions) + pub fn validate_completion_count( + request: &T, + ) -> Result<(), ValidationError> { + const N_RANGE: (u32, u32) = (1, 10); + + if let Some(n) = request.get_n() { + validate_range(n, &N_RANGE, "n")?; + } + + Ok(()) + } + + /// Validate that an array is not empty + pub fn validate_non_empty_array( + items: &[T], + param_name: &str, + ) -> Result<(), ValidationError> { + if items.is_empty() { + return Err(ValidationError::MissingRequired { + parameter: param_name.to_string(), + }); + } + Ok(()) + } + + /// Validate common request parameters that are shared across all API types + pub fn validate_common_request_params(request: &T) -> Result<(), ValidationError> + where + T: SamplingOptionsProvider + + StopConditionsProvider + + TokenLimitsProvider + + LogProbsProvider + + GenerationExtensionsProvider + + CompletionCountProvider + + ?Sized, + { + // Validate all standard parameters + validate_sampling_options(request)?; + validate_stop_conditions(request)?; + validate_token_limits(request)?; + validate_logprobs(request)?; + + // Validate generation extensions and completion count + validate_generation_extensions(request)?; + validate_completion_count(request)?; + + // Perform cross-parameter validation + validate_cross_parameters(request)?; + + Ok(()) + } +} + +/// Core validation traits for different parameter categories +pub trait SamplingOptionsProvider { + /// Get temperature parameter + fn get_temperature(&self) -> Option; + + /// Get top_p parameter + fn get_top_p(&self) -> Option; + + /// Get frequency penalty parameter + fn get_frequency_penalty(&self) -> Option; + + /// Get presence penalty parameter + fn get_presence_penalty(&self) -> Option; +} + +/// Trait for validating stop conditions +pub trait StopConditionsProvider { + /// Get stop sequences + fn get_stop_sequences(&self) -> Option<&StringOrArray>; +} + +/// Trait for validating token limits +pub trait TokenLimitsProvider { + /// Get maximum tokens parameter + fn get_max_tokens(&self) -> Option; + + /// Get minimum tokens parameter + fn get_min_tokens(&self) -> Option; +} + +/// Trait for validating logprobs parameters +pub trait LogProbsProvider { + /// Get logprobs parameter (completions API) + fn get_logprobs(&self) -> Option; + + /// Get top_logprobs parameter (chat API) + fn get_top_logprobs(&self) -> Option; +} + +/// Trait for generation extensions +pub trait GenerationExtensionsProvider { + /// Get top_k parameter + fn get_top_k(&self) -> Option { + None + } + + /// Get min_p parameter + fn get_min_p(&self) -> Option { + None + } + + /// Get repetition_penalty parameter + fn get_repetition_penalty(&self) -> Option { + None + } +} + +/// Trait for n parameter (number of completions) +pub trait CompletionCountProvider { + /// Get n parameter + fn get_n(&self) -> Option { + None + } +} + +/// Comprehensive validation trait that combines all validation aspects +pub trait ValidatableRequest: + SamplingOptionsProvider + + StopConditionsProvider + + TokenLimitsProvider + + LogProbsProvider + + GenerationExtensionsProvider + + CompletionCountProvider +{ + /// Perform comprehensive validation of the entire request + fn validate(&self) -> Result<(), ValidationError> { + // Use the common validation function + utils::validate_common_request_params(self) + } +} + +// ================================================================== +// = OPENAI CHAT COMPLETION VALIDATION = +// ================================================================== + +impl SamplingOptionsProvider for ChatCompletionRequest { + fn get_temperature(&self) -> Option { + self.temperature + } + fn get_top_p(&self) -> Option { + self.top_p + } + fn get_frequency_penalty(&self) -> Option { + self.frequency_penalty + } + fn get_presence_penalty(&self) -> Option { + self.presence_penalty + } +} + +impl StopConditionsProvider for ChatCompletionRequest { + fn get_stop_sequences(&self) -> Option<&StringOrArray> { + self.stop.as_ref() + } +} + +impl TokenLimitsProvider for ChatCompletionRequest { + fn get_max_tokens(&self) -> Option { + // Prefer max_completion_tokens over max_tokens if both are set + self.max_completion_tokens.or(self.max_tokens) + } + + fn get_min_tokens(&self) -> Option { + self.min_tokens + } +} + +impl LogProbsProvider for ChatCompletionRequest { + fn get_logprobs(&self) -> Option { + // For chat API, logprobs is a boolean, return 1 if true for validation purposes + if self.logprobs { + Some(1) + } else { + None + } + } + + fn get_top_logprobs(&self) -> Option { + self.top_logprobs + } +} + +impl GenerationExtensionsProvider for ChatCompletionRequest { + fn get_top_k(&self) -> Option { + self.top_k + } + + fn get_min_p(&self) -> Option { + self.min_p + } + + fn get_repetition_penalty(&self) -> Option { + self.repetition_penalty + } +} + +impl CompletionCountProvider for ChatCompletionRequest { + fn get_n(&self) -> Option { + self.n + } +} + +impl ChatCompletionRequest { + /// Validate message-specific requirements + pub fn validate_messages(&self) -> Result<(), ValidationError> { + // LMDeploy token-in-token-out requests use an empty messages array and + // carry the active prompt in `input_ids`. + if self.messages.is_empty() { + return match self.input_ids.as_ref() { + Some(ids) if !ids.is_empty() => Ok(()), + Some(_) => Err(ValidationError::InvalidValue { + parameter: "input_ids".to_string(), + value: "empty array".to_string(), + reason: "input_ids cannot be empty when messages is empty".to_string(), + }), + None => utils::validate_non_empty_array(&self.messages, "messages"), + }; + } + + // LMDeploy gives messages priority and rejects simultaneous token input. + if self.input_ids.is_some() { + return Err(ValidationError::InvalidValue { + parameter: "input_ids".to_string(), + value: "set with non-empty messages".to_string(), + reason: "input_ids cannot be used when messages is non-empty".to_string(), + }); + } + + // Validate message content is not empty + for (i, msg) in self.messages.iter().enumerate() { + if let ChatMessage::User { content, .. } = msg { + match content { + UserMessageContent::Text(text) if text.is_empty() => { + return Err(ValidationError::InvalidValue { + parameter: format!("messages[{}].content", i), + value: "empty".to_string(), + reason: "message content cannot be empty".to_string(), + }); + } + UserMessageContent::Parts(parts) if parts.is_empty() => { + return Err(ValidationError::InvalidValue { + parameter: format!("messages[{}].content", i), + value: "empty array".to_string(), + reason: "message content parts cannot be empty".to_string(), + }); + } + _ => {} + } + } + } + + Ok(()) + } + + /// Validate response format if specified + pub fn validate_response_format(&self) -> Result<(), ValidationError> { + if let Some(ResponseFormat::JsonSchema { json_schema }) = &self.response_format { + if json_schema.name.is_empty() { + return Err(ValidationError::InvalidValue { + parameter: "response_format.json_schema.name".to_string(), + value: "empty".to_string(), + reason: "JSON schema name cannot be empty".to_string(), + }); + } + } + Ok(()) + } + + /// Validate chat API specific logprobs requirements + pub fn validate_chat_logprobs(&self) -> Result<(), ValidationError> { + // In chat API, if logprobs=true, top_logprobs must be specified + if self.logprobs && self.top_logprobs.is_none() { + return Err(ValidationError::MissingRequired { + parameter: "top_logprobs".to_string(), + }); + } + + // If top_logprobs is specified, logprobs should be true + if self.top_logprobs.is_some() && !self.logprobs { + return Err(ValidationError::InvalidValue { + parameter: "logprobs".to_string(), + value: "false".to_string(), + reason: "must be true when top_logprobs is specified".to_string(), + }); + } + + Ok(()) + } + + /// Validate cross-parameter relationships specific to chat completions + pub fn validate_chat_cross_parameters(&self) -> Result<(), ValidationError> { + // Validate that both max_tokens and max_completion_tokens aren't set + utils::validate_conflicting_parameters( + "max_tokens", + self.max_tokens.is_some(), + "max_completion_tokens", + self.max_completion_tokens.is_some(), + "cannot specify both max_tokens and max_completion_tokens", + )?; + + // Validate that tools and functions aren't both specified (deprecated) + utils::validate_conflicting_parameters( + "tools", + self.tools.is_some(), + "functions", + self.functions.is_some(), + "functions is deprecated, use tools instead", + )?; + + // Validate structured output constraints don't conflict with JSON response format + let has_json_format = matches!( + self.response_format, + Some(ResponseFormat::JsonObject | ResponseFormat::JsonSchema { .. }) + ); + + utils::validate_conflicting_parameters( + "response_format", + has_json_format, + "regex", + self.regex.is_some(), + "cannot use regex constraint with JSON response format", + )?; + + utils::validate_conflicting_parameters( + "response_format", + has_json_format, + "ebnf", + self.ebnf.is_some(), + "cannot use EBNF constraint with JSON response format", + )?; + + // Only one structured output constraint should be active + let structured_constraints = [ + ("regex", self.regex.is_some()), + ("ebnf", self.ebnf.is_some()), + ( + "json_schema", + matches!( + self.response_format, + Some(ResponseFormat::JsonSchema { .. }) + ), + ), + ]; + + utils::validate_mutually_exclusive_options( + &structured_constraints, + "Only one structured output constraint (regex, ebnf, or json_schema) can be active at a time", + )?; + + Ok(()) + } +} + +impl ValidatableRequest for ChatCompletionRequest { + fn validate(&self) -> Result<(), ValidationError> { + // Call the common validation function from the validation module + utils::validate_common_request_params(self)?; + + // Then validate chat-specific parameters + self.validate_messages()?; + self.validate_response_format()?; + self.validate_chat_logprobs()?; + self.validate_chat_cross_parameters()?; + + Ok(()) + } +} + +#[cfg(test)] +mod tests { + use super::constants::*; + use super::utils::*; + use super::*; + use crate::protocols::spec::StringOrArray; + + // Mock request type for testing validation traits + #[derive(Debug, Default)] + struct MockRequest { + temperature: Option, + stop: Option, + max_tokens: Option, + min_tokens: Option, + } + + impl SamplingOptionsProvider for MockRequest { + fn get_temperature(&self) -> Option { + self.temperature + } + fn get_top_p(&self) -> Option { + None + } + fn get_frequency_penalty(&self) -> Option { + None + } + fn get_presence_penalty(&self) -> Option { + None + } + } + + impl StopConditionsProvider for MockRequest { + fn get_stop_sequences(&self) -> Option<&StringOrArray> { + self.stop.as_ref() + } + } + + impl TokenLimitsProvider for MockRequest { + fn get_max_tokens(&self) -> Option { + self.max_tokens + } + fn get_min_tokens(&self) -> Option { + self.min_tokens + } + } + + impl LogProbsProvider for MockRequest { + fn get_logprobs(&self) -> Option { + None + } + fn get_top_logprobs(&self) -> Option { + None + } + } + + impl GenerationExtensionsProvider for MockRequest {} + impl CompletionCountProvider for MockRequest {} + impl ValidatableRequest for MockRequest {} + + #[test] + fn test_range_validation() { + // Valid range + assert!(validate_range(1.5f32, &TEMPERATURE_RANGE, "temperature").is_ok()); + // Invalid range + assert!(validate_range(-0.1f32, &TEMPERATURE_RANGE, "temperature").is_err()); + assert!(validate_range(3.0f32, &TEMPERATURE_RANGE, "temperature").is_err()); + } + + #[test] + fn test_generation_top_k_validation() { + assert!(validate_top_k(-1).is_ok()); // Disabled + assert!(validate_top_k(50).is_ok()); // Valid positive + assert!(validate_top_k(0).is_err()); // Invalid + assert!(validate_top_k(-5).is_err()); // Invalid + } + + #[test] + fn test_stop_sequences_limits() { + let request = MockRequest { + stop: Some(StringOrArray::Array(vec![ + "stop1".to_string(), + "stop2".to_string(), + "stop3".to_string(), + "stop4".to_string(), + "stop5".to_string(), // Too many + ])), + ..Default::default() + }; + assert!(request.validate().is_err()); + } + + #[test] + fn test_token_limits_conflict() { + let request = MockRequest { + min_tokens: Some(100), + max_tokens: Some(50), // min > max + ..Default::default() + }; + assert!(request.validate().is_err()); + } + + #[test] + fn test_valid_request() { + let request = MockRequest { + temperature: Some(1.0), + stop: Some(StringOrArray::Array(vec!["stop".to_string()])), + max_tokens: Some(100), + min_tokens: Some(10), + }; + assert!(request.validate().is_ok()); + } + + // Chat completion specific tests + #[cfg(test)] + mod chat_tests { + use super::*; + use crate::protocols::spec::ReasoningEffort; + + fn create_valid_chat_request() -> ChatCompletionRequest { + ChatCompletionRequest { + model: Some("gpt-4".to_string()), + messages: vec![ChatMessage::User { + role: "user".to_string(), + content: UserMessageContent::Text("Hello".to_string()), + name: None, + }], + input_ids: None, + temperature: Some(1.0), + top_p: Some(0.9), + n: Some(1), + stream: false, + stream_options: None, + stop: None, + max_tokens: Some(100), + max_completion_tokens: None, + presence_penalty: Some(0.0), + frequency_penalty: Some(0.0), + logit_bias: None, + user: None, + seed: None, + logprobs: false, + top_logprobs: None, + response_format: None, + tools: None, + tool_choice: None, + parallel_tool_calls: None, + functions: None, + function_call: None, + // Generation extensions + top_k: None, + min_p: None, + min_tokens: None, + repetition_penalty: None, + regex: None, + ebnf: None, + stop_token_ids: None, + echo: None, + reasoning_effort: None, + include_reasoning: true, + no_stop_trim: false, + ignore_eos: false, + add_generation_prompt: true, + continue_final_message: false, + skip_special_tokens: true, + separate_reasoning: true, + stream_reasoning: true, + chat_template_kwargs: None, + structured_outputs: None, + other: Default::default(), + } + } + + #[test] + fn test_chat_validation_basics() { + // Valid request + assert!(create_valid_chat_request().validate().is_ok()); + + // Empty messages + let mut request = create_valid_chat_request(); + request.messages = vec![]; + assert!(request.validate().is_err()); + + // Invalid temperature + let mut request = create_valid_chat_request(); + request.temperature = Some(3.0); + assert!(request.validate().is_err()); + } + + #[test] + fn test_chat_validation_accepts_lmdeploy_token_input() { + let mut request = create_valid_chat_request(); + request.messages = vec![]; + request.input_ids = Some(vec![1, 2, 3]); + assert!(request.validate().is_ok()); + + request.input_ids = Some(vec![]); + assert!(request.validate().is_err()); + } + + #[test] + fn test_chat_validation_rejects_messages_with_lmdeploy_token_input() { + let mut request = create_valid_chat_request(); + request.input_ids = Some(vec![1, 2, 3]); + assert!(request.validate().is_err()); + } + + #[test] + fn test_chat_conflicts() { + let mut request = create_valid_chat_request(); + + // Conflicting max_tokens + request.max_tokens = Some(100); + request.max_completion_tokens = Some(200); + assert!(request.validate().is_err()); + + // Logprobs without top_logprobs + request.max_tokens = None; + request.logprobs = true; + request.top_logprobs = None; + assert!(request.validate().is_err()); + } + + #[test] + fn test_generation_extensions() { + let mut request = create_valid_chat_request(); + + // Valid generation extension parameters + request.top_k = Some(-1); + request.min_p = Some(0.1); + request.repetition_penalty = Some(1.2); + assert!(request.validate().is_ok()); + + // Invalid parameters + request.top_k = Some(0); // Invalid + assert!(request.validate().is_err()); + } + + #[test] + fn test_parameter_ranges() { + let mut request = create_valid_chat_request(); + + // Test temperature range (0.0 to 2.0) + request.temperature = Some(1.5); + assert!(request.validate().is_ok()); + request.temperature = Some(-0.1); + assert!(request.validate().is_err()); + request.temperature = Some(3.0); + assert!(request.validate().is_err()); + + // Test top_p range (0.0 to 1.0) + request.temperature = Some(1.0); // Reset + request.top_p = Some(0.9); + assert!(request.validate().is_ok()); + request.top_p = Some(-0.1); + assert!(request.validate().is_err()); + request.top_p = Some(1.5); + assert!(request.validate().is_err()); + + // Test frequency_penalty range (-2.0 to 2.0) + request.top_p = Some(0.9); // Reset + request.frequency_penalty = Some(1.5); + assert!(request.validate().is_ok()); + request.frequency_penalty = Some(-2.5); + assert!(request.validate().is_err()); + request.frequency_penalty = Some(3.0); + assert!(request.validate().is_err()); + + // Test presence_penalty range (-2.0 to 2.0) + request.frequency_penalty = Some(0.0); // Reset + request.presence_penalty = Some(-1.5); + assert!(request.validate().is_ok()); + request.presence_penalty = Some(-3.0); + assert!(request.validate().is_err()); + request.presence_penalty = Some(2.5); + assert!(request.validate().is_err()); + + // Test repetition_penalty range (0.0 to 2.0) + request.presence_penalty = Some(0.0); // Reset + request.repetition_penalty = Some(1.2); + assert!(request.validate().is_ok()); + request.repetition_penalty = Some(-0.1); + assert!(request.validate().is_err()); + request.repetition_penalty = Some(2.1); + assert!(request.validate().is_err()); + + // Test min_p range (0.0 to 1.0) + request.repetition_penalty = Some(1.0); // Reset + request.min_p = Some(0.5); + assert!(request.validate().is_ok()); + request.min_p = Some(-0.1); + assert!(request.validate().is_err()); + request.min_p = Some(1.5); + assert!(request.validate().is_err()); + } + + #[test] + fn test_structured_output_conflicts() { + let mut request = create_valid_chat_request(); + + // JSON response format with regex should conflict + request.response_format = Some(ResponseFormat::JsonObject); + request.regex = Some(".*".to_string()); + assert!(request.validate().is_err()); + + // JSON response format with EBNF should conflict + request.regex = None; + request.ebnf = Some("grammar".to_string()); + assert!(request.validate().is_err()); + + // Multiple structured constraints should conflict + request.response_format = None; + request.regex = Some(".*".to_string()); + request.ebnf = Some("grammar".to_string()); + assert!(request.validate().is_err()); + + // Only one constraint should work + request.ebnf = None; + request.regex = Some(".*".to_string()); + assert!(request.validate().is_ok()); + + request.regex = None; + request.ebnf = Some("grammar".to_string()); + assert!(request.validate().is_ok()); + + request.ebnf = None; + request.response_format = Some(ResponseFormat::JsonObject); + assert!(request.validate().is_ok()); + } + + #[test] + fn test_stop_sequences_validation() { + let mut request = create_valid_chat_request(); + + // Valid stop sequences + request.stop = Some(StringOrArray::Array(vec![ + "stop1".to_string(), + "stop2".to_string(), + ])); + assert!(request.validate().is_ok()); + + // Too many stop sequences (max 4) + request.stop = Some(StringOrArray::Array(vec![ + "stop1".to_string(), + "stop2".to_string(), + "stop3".to_string(), + "stop4".to_string(), + "stop5".to_string(), + ])); + assert!(request.validate().is_err()); + + // Empty stop sequence should fail + request.stop = Some(StringOrArray::String("".to_string())); + assert!(request.validate().is_err()); + + // Empty string in array should fail + request.stop = Some(StringOrArray::Array(vec![ + "stop1".to_string(), + "".to_string(), + ])); + assert!(request.validate().is_err()); + } + + #[test] + fn test_logprobs_validation() { + let mut request = create_valid_chat_request(); + + // Valid logprobs configuration + request.logprobs = true; + request.top_logprobs = Some(10); + assert!(request.validate().is_ok()); + + // logprobs=true without top_logprobs should fail + request.top_logprobs = None; + assert!(request.validate().is_err()); + + // top_logprobs without logprobs=true should fail + request.logprobs = false; + request.top_logprobs = Some(10); + assert!(request.validate().is_err()); + + // top_logprobs out of range (0-20) + request.logprobs = true; + request.top_logprobs = Some(25); + assert!(request.validate().is_err()); + } + + #[test] + fn test_n_parameter_validation() { + let mut request = create_valid_chat_request(); + + // Valid n values (1-10) + request.n = Some(1); + assert!(request.validate().is_ok()); + request.n = Some(5); + assert!(request.validate().is_ok()); + request.n = Some(10); + assert!(request.validate().is_ok()); + + // Invalid n values + request.n = Some(0); + assert!(request.validate().is_err()); + request.n = Some(15); + assert!(request.validate().is_err()); + } + + #[test] + fn test_min_max_tokens_validation() { + let mut request = create_valid_chat_request(); + + // Valid token limits + request.min_tokens = Some(10); + request.max_tokens = Some(100); + assert!(request.validate().is_ok()); + + // min_tokens > max_tokens should fail + request.min_tokens = Some(150); + request.max_tokens = Some(100); + assert!(request.validate().is_err()); + + // Should work with max_completion_tokens instead + request.max_tokens = None; + request.max_completion_tokens = Some(200); + request.min_tokens = Some(50); + assert!(request.validate().is_ok()); + + // min_tokens > max_completion_tokens should fail + request.min_tokens = Some(250); + assert!(request.validate().is_err()); + } + + #[test] + fn test_echo_field() { + let mut request = create_valid_chat_request(); + + // echo is None by default, should be valid + assert!(request.echo.is_none()); + assert!(request.validate().is_ok()); + + // echo = Some(true) should be valid + request.echo = Some(true); + assert!(request.validate().is_ok()); + + // echo = Some(false) should be valid + request.echo = Some(false); + assert!(request.validate().is_ok()); + } + + #[test] + fn test_reasoning_effort_field() { + let mut request = create_valid_chat_request(); + + // reasoning_effort is None by default, should be valid + assert!(request.reasoning_effort.is_none()); + assert!(request.validate().is_ok()); + + // Valid reasoning_effort values: Low, Medium, High + request.reasoning_effort = Some(ReasoningEffort::Low); + assert!(request.validate().is_ok()); + + request.reasoning_effort = Some(ReasoningEffort::Medium); + assert!(request.validate().is_ok()); + + request.reasoning_effort = Some(ReasoningEffort::High); + assert!(request.validate().is_ok()); + } + + #[test] + fn test_include_reasoning_field() { + let mut request = create_valid_chat_request(); + + // include_reasoning defaults to true + assert!(request.include_reasoning); + assert!(request.validate().is_ok()); + + // include_reasoning = true should be valid + request.include_reasoning = true; + assert!(request.validate().is_ok()); + + // include_reasoning = false should be valid + request.include_reasoning = false; + assert!(request.validate().is_ok()); + } + } +} diff --git a/lmdeploy/router/src/protocols/worker_spec.rs b/lmdeploy/router/src/protocols/worker_spec.rs new file mode 100644 index 0000000000..9718868ef3 --- /dev/null +++ b/lmdeploy/router/src/protocols/worker_spec.rs @@ -0,0 +1,151 @@ +//! Worker management API specifications +//! +//! Defines the request/response structures for worker management endpoints + +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; + +/// Worker configuration for API requests +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct WorkerConfigRequest { + /// Worker URL (required) + pub url: String, + + /// Model ID (optional, will query from server if not provided) + #[serde(skip_serializing_if = "Option::is_none")] + pub model_id: Option, + + /// Worker priority (optional, default: 50, higher = preferred) + #[serde(skip_serializing_if = "Option::is_none")] + pub priority: Option, + + /// Worker cost factor (optional, default: 1.0) + #[serde(skip_serializing_if = "Option::is_none")] + pub cost: Option, + + /// Worker type (optional: "regular", "prefill", "decode") + #[serde(skip_serializing_if = "Option::is_none")] + pub worker_type: Option, + + /// Additional labels (optional) + #[serde(default, skip_serializing_if = "HashMap::is_empty")] + pub labels: HashMap, +} + +/// Worker information for API responses +#[derive(Debug, Clone, Serialize)] +pub struct WorkerInfo { + /// Worker unique identifier + pub id: String, + + /// Worker URL + pub url: String, + + /// Model ID this worker serves + pub model_id: String, + + /// Worker priority + pub priority: u32, + + /// Worker cost factor + pub cost: f32, + + /// Worker type + pub worker_type: String, + + /// Whether the worker is healthy + pub is_healthy: bool, + + /// Current load on the worker + pub load: usize, + + /// Connection mode + pub connection_mode: String, + + /// Additional metadata + #[serde(skip_serializing_if = "HashMap::is_empty")] + pub metadata: HashMap, +} + +/// Worker list response +#[derive(Debug, Clone, Serialize)] +pub struct WorkerListResponse { + /// List of workers + pub workers: Vec, + + /// Total count + pub total: usize, + + /// Statistics + pub stats: WorkerStats, +} + +/// Worker statistics +#[derive(Debug, Clone, Serialize)] +pub struct WorkerStats { + pub total_workers: usize, + pub healthy_workers: usize, + pub total_models: usize, + pub total_load: usize, + pub by_type: WorkerTypeStats, +} + +/// Worker statistics by type +#[derive(Debug, Clone, Serialize)] +pub struct WorkerTypeStats { + pub regular: usize, + pub prefill: usize, + pub decode: usize, +} + +/// Worker update request +#[derive(Debug, Clone, Deserialize)] +pub struct WorkerUpdateRequest { + /// Update priority + #[serde(skip_serializing_if = "Option::is_none")] + pub priority: Option, + + /// Update cost + #[serde(skip_serializing_if = "Option::is_none")] + pub cost: Option, + + /// Update labels + #[serde(skip_serializing_if = "Option::is_none")] + pub labels: Option>, +} + +/// Generic API response +#[derive(Debug, Clone, Serialize)] +pub struct WorkerApiResponse { + pub success: bool, + pub message: String, + + #[serde(skip_serializing_if = "Option::is_none")] + pub worker: Option, +} + +/// Error response +#[derive(Debug, Clone, Serialize)] +pub struct WorkerErrorResponse { + pub error: String, + pub code: String, +} + +/// Server info response from /get_server_info endpoint +#[derive(Debug, Clone, Deserialize)] +pub struct ServerInfo { + #[serde(skip_serializing_if = "Option::is_none")] + pub model_id: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + pub model_path: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + pub priority: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + pub cost: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + pub worker_type: Option, +} diff --git a/lmdeploy/router/src/routers/factory.rs b/lmdeploy/router/src/routers/factory.rs new file mode 100644 index 0000000000..92e15aba44 --- /dev/null +++ b/lmdeploy/router/src/routers/factory.rs @@ -0,0 +1,137 @@ +//! Factory for creating router instances + +use super::{ + http::{lmdeploy_pd_router::LMDeployPDRouter, openai_router::OpenAIRouter, router::Router}, + RouterTrait, +}; +use crate::config::{LMDeployMigrationProtocol, LMDeployRdmaConfig, PolicyConfig, RoutingMode}; +use crate::policies::PolicyFactory; +use crate::server::AppContext; +use std::sync::Arc; + +/// Static inputs needed to build an LMDeploy PD router. +pub struct LMDeployPDRouterParams<'a> { + pub prefill_urls: &'a [String], + pub decode_urls: &'a [String], + pub migration_protocol: LMDeployMigrationProtocol, + pub rdma_config: Option, + pub dummy_prefill: bool, + pub prefill_policy_config: Option<&'a PolicyConfig>, + pub decode_policy_config: Option<&'a PolicyConfig>, + pub main_policy_config: &'a PolicyConfig, + pub ctx: &'a Arc, +} + +/// Factory for creating router instances based on configuration +pub struct RouterFactory; + +impl RouterFactory { + /// Create a router instance from application context + pub async fn create_router(ctx: &Arc) -> Result, String> { + match &ctx.router_config.mode { + RoutingMode::Regular { worker_urls } => { + Self::create_regular_router(worker_urls, ctx).await + } + RoutingMode::LMDeployPrefillDecode { + prefill_urls, + decode_urls, + prefill_policy, + decode_policy, + migration_protocol, + rdma_config, + dummy_prefill, + } => { + tracing::info!( + "Creating LMDeployPDRouter with prefill_urls: {:?}, decode_urls: {:?}, migration_protocol: {:?}", + prefill_urls, decode_urls, migration_protocol + ); + Self::create_lmdeploy_pd_router(LMDeployPDRouterParams { + prefill_urls, + decode_urls, + migration_protocol: *migration_protocol, + rdma_config: rdma_config.clone(), + dummy_prefill: *dummy_prefill, + prefill_policy_config: prefill_policy.as_ref(), + decode_policy_config: decode_policy.as_ref(), + main_policy_config: &ctx.router_config.policy, + ctx, + }) + .await + } + RoutingMode::OpenAI { worker_urls, .. } => { + Self::create_openai_router(worker_urls.clone(), ctx).await + } + } + } + + /// Create a regular router + pub async fn create_regular_router( + worker_urls: &[String], + ctx: &Arc, + ) -> Result, String> { + // Create regular router with context + let router = Router::new(worker_urls.to_vec(), ctx).await?; + + Ok(Box::new(router)) + } + + /// Create an LMD (lmdeploy) PD router with static URLs + pub async fn create_lmdeploy_pd_router( + params: LMDeployPDRouterParams<'_>, + ) -> Result, String> { + let LMDeployPDRouterParams { + prefill_urls, + decode_urls, + migration_protocol, + rdma_config, + dummy_prefill, + prefill_policy_config, + decode_policy_config, + main_policy_config, + ctx, + } = params; + let prefill_policy = + PolicyFactory::create_from_config(prefill_policy_config.unwrap_or(main_policy_config)); + let decode_policy = + PolicyFactory::create_from_config(decode_policy_config.unwrap_or(main_policy_config)); + + ctx.policy_registry.set_prefill_policy(prefill_policy); + ctx.policy_registry.set_decode_policy(decode_policy); + + tracing::info!( + "Creating LMDeployPDRouter with static URLs - prefill: {:?}, decode: {:?}", + prefill_urls, + decode_urls + ); + + let router = LMDeployPDRouter::new( + prefill_urls.to_vec(), + decode_urls.to_vec(), + migration_protocol, + rdma_config, + dummy_prefill, + ctx, + ) + .await?; + tracing::info!("LMDeployPDRouter instance created successfully"); + + Ok(Box::new(router)) + } + + /// Create an OpenAI router + async fn create_openai_router( + worker_urls: Vec, + ctx: &Arc, + ) -> Result, String> { + // Use the first worker URL as the OpenAI-compatible base + let base_url = worker_urls + .first() + .cloned() + .ok_or_else(|| "OpenAI mode requires at least one worker URL".to_string())?; + + let router = + OpenAIRouter::new(base_url, Some(ctx.router_config.circuit_breaker.clone())).await?; + + Ok(Box::new(router)) + } +} diff --git a/lmdeploy/router/src/routers/header_utils.rs b/lmdeploy/router/src/routers/header_utils.rs new file mode 100644 index 0000000000..bb6ea3db4e --- /dev/null +++ b/lmdeploy/router/src/routers/header_utils.rs @@ -0,0 +1,95 @@ +use axum::body::Body; +use axum::extract::Request; +use axum::http::HeaderMap; + +pub use crate::otel_http::TRACE_HEADER_NAMES; + +/// Copy request headers to a Vec of name-value string pairs +/// Used for forwarding headers to backend workers +pub fn copy_request_headers(req: &Request) -> Vec<(String, String)> { + req.headers() + .iter() + .filter_map(|(name, value)| { + // Convert header value to string, skipping non-UTF8 headers + value + .to_str() + .ok() + .map(|v| (name.to_string(), v.to_string())) + }) + .collect() +} + +/// Convert headers from reqwest Response to axum HeaderMap +/// Filters out hop-by-hop headers that shouldn't be forwarded +pub fn preserve_response_headers(reqwest_headers: &HeaderMap) -> HeaderMap { + let mut headers = HeaderMap::new(); + + for (name, value) in reqwest_headers.iter() { + // Skip hop-by-hop headers that shouldn't be forwarded + let name_str = name.as_str().to_lowercase(); + if should_forward_header(&name_str) { + // The original name and value are already valid, so we can just clone them + headers.insert(name.clone(), value.clone()); + } + } + + headers +} + +/// Determine if a header should be forwarded from backend to client +fn should_forward_header(name: &str) -> bool { + // List of headers that should NOT be forwarded (hop-by-hop headers) + !matches!( + name, + "connection" | + "keep-alive" | + "proxy-authenticate" | + "proxy-authorization" | + "te" | + "trailers" | + "transfer-encoding" | + "upgrade" | + "content-encoding" | // Let axum/hyper handle encoding + "host" // Should not forward the backend's host header + ) +} + +/// Propagate OpenTelemetry trace headers to a reqwest RequestBuilder +/// +/// When OTel is enabled: actively injects the current span's trace context, +/// making the router's span the parent of the backend request's span. +/// When OTel is disabled: passively forwards existing trace headers from +/// the incoming request. +pub fn propagate_trace_headers( + request: reqwest::RequestBuilder, + headers: Option<&HeaderMap>, +) -> reqwest::RequestBuilder { + crate::otel_http::propagate_trace_headers(request, headers) +} + +/// Propagate specific headers from incoming request to outgoing reqwest RequestBuilder +/// +/// This is a general-purpose helper for selectively forwarding headers by name. +/// Only headers whose names match the provided list (case-insensitive) are propagated. +/// +/// # Arguments +/// * `request` - The reqwest RequestBuilder to add headers to +/// * `headers` - Optional incoming headers to check +/// * `header_names` - List of header names to propagate (matched case-insensitively) +/// +/// # Returns +/// The RequestBuilder with matching headers added +pub fn propagate_headers( + mut request: reqwest::RequestBuilder, + headers: Option<&HeaderMap>, + header_names: &[&str], +) -> reqwest::RequestBuilder { + if let Some(h) = headers { + for &name in header_names { + if let Some(value) = h.get(name) { + request = request.header(name, value); + } + } + } + request +} diff --git a/lmdeploy/router/src/routers/http/lmdeploy_pd_router.rs b/lmdeploy/router/src/routers/http/lmdeploy_pd_router.rs new file mode 100644 index 0000000000..9210b62cec --- /dev/null +++ b/lmdeploy/router/src/routers/http/lmdeploy_pd_router.rs @@ -0,0 +1,1533 @@ +// lmdeploy PD (Prefill-Decode) Router Implementation +// Extends PdRouterBase to handle lmdeploy-specific two-stage processing. +// +// LMDeploy PD protocol differs from the legacy PD protocol: +// - Prefill request sets with_cache=true, preserve_cache=true, max_tokens=1 +// - Prefill response carries top-level id, cache_block_ids, remote_token_ids +// - P2P RDMA connection established via /distserve/p2p_initialize + /distserve/p2p_connect +// - Decode request carries migration_request field (not kv_transfer_params) +use super::pd_router::PdRouterBase; +use super::pd_types::{error_chain, PDRouterError}; +use crate::config::{LMDeployMigrationProtocol, LMDeployRdmaConfig, LMDeployRdmaLinkType}; +use crate::core::Worker; +use crate::metrics::RouterMetrics; +use crate::otel_http::{self, ClientRequestOptions}; +use crate::policies::PolicyRegistry; +use crate::routers::{header_utils, RouterTrait, WorkerManagement}; +use async_trait::async_trait; +use axum::{ + body::Body, + extract::Request, + http::{HeaderMap, Method, StatusCode}, + response::{IntoResponse, Response}, +}; +use dashmap::mapref::entry::Entry; +use dashmap::DashMap; +use futures_util::StreamExt; +use serde_json::{json, Value}; +use std::collections::HashMap; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio::sync::OnceCell; +use tracing::{debug, error, info, warn}; +use uuid::Uuid; + +type P2pKey = (String, String); +type P2pConnectionCell = Arc>; + +/// lmdeploy PD Router that extends PdRouterBase with lmdeploy-specific request handling +#[derive(Debug)] +pub struct LMDeployPDRouter { + /// Underlying PD router for worker management, health, models + pd_router: PdRouterBase, + /// HTTP client for P2P setup calls (/distserve/*) + http_client: reqwest::Client, + /// Policy registry for load balancing + policy_registry: Arc, + /// KV migration protocol (rdma/nvlink) + migration_protocol: LMDeployMigrationProtocol, + /// Optional RDMA configuration + rdma_config: Option, + /// Whether to use dummy prefill + dummy_prefill: bool, + /// Per-pair, single-flight P2P connection initialization. + p2p_pool: Arc>, +} + +/// Releases a preserved prefill session if the request exits before a successful +/// decode response has completed. The asynchronous cleanup is spawned from Drop +/// so cancellation and streaming-client disconnects are covered as well. +#[derive(Debug)] +struct PrefillCacheGuard { + client: reqwest::Client, + prefill_base_url: String, + session_id: i64, + completed: Arc, +} + +impl PrefillCacheGuard { + fn new(client: reqwest::Client, prefill_base_url: String, session_id: i64) -> Self { + Self { + client, + prefill_base_url, + session_id, + completed: Arc::new(AtomicBool::new(false)), + } + } + + fn mark_completed(&self) { + self.completed.store(true, Ordering::Release); + } +} + +impl Drop for PrefillCacheGuard { + fn drop(&mut self) { + if self.completed.load(Ordering::Acquire) { + return; + } + + let client = self.client.clone(); + let prefill_base_url = self.prefill_base_url.clone(); + let session_id = self.session_id; + if let Ok(handle) = tokio::runtime::Handle::try_current() { + handle.spawn(async move { + let url = format!("{}/distserve/free_cache", prefill_base_url); + let result = client + .post(&url) + .json(&json!({ + "remote_engine_id": prefill_base_url, + "remote_session_id": session_id, + })) + .send() + .await; + match result { + Ok(response) if response.status().is_success() => { + debug!("Released preserved prefill session {}", session_id); + } + Ok(response) => { + warn!( + "Failed to release prefill session {}: status {}", + session_id, + response.status() + ); + } + Err(error) => { + warn!( + "Failed to release prefill session {} at {}: {}", + session_id, url, error + ); + } + } + }); + } + } +} + +impl LMDeployPDRouter { + /// Create a new LMDeployPDRouter with static prefill/decode URLs + pub async fn new( + prefill_urls: Vec, + decode_urls: Vec, + migration_protocol: LMDeployMigrationProtocol, + rdma_config: Option, + dummy_prefill: bool, + ctx: &Arc, + ) -> Result { + info!( + "LMDeployPDRouter::new with {} prefill, {} decode, protocol={:?}, dummy_prefill={}", + prefill_urls.len(), + decode_urls.len(), + migration_protocol, + dummy_prefill + ); + + let pd_router = PdRouterBase::new(prefill_urls, decode_urls, ctx).await?; + + // Initialize policies with workers from registry + let prefill_workers = pd_router.worker_registry.get_prefill_workers(); + let decode_workers = pd_router.worker_registry.get_decode_workers(); + let prefill_policy = ctx.policy_registry.get_prefill_policy(); + let decode_policy = ctx.policy_registry.get_decode_policy(); + if prefill_policy.requires_initialization() { + info!("Initializing prefill policy with workers."); + prefill_policy.init_workers(&prefill_workers); + } + if decode_policy.requires_initialization() { + info!("Initializing decode policy with workers."); + decode_policy.init_workers(&decode_workers); + } + + let rdma_config = match migration_protocol { + LMDeployMigrationProtocol::Rdma => Some(rdma_config.unwrap_or_default()), + LMDeployMigrationProtocol::Nvlink => None, + }; + + Ok(Self { + pd_router, + http_client: reqwest::Client::builder() + .connect_timeout(Duration::from_secs(10)) + .timeout(Duration::from_secs(120)) + .build() + .map_err(|e| format!("Failed to build LMDeploy control-plane client: {e}"))?, + policy_registry: ctx.policy_registry.clone(), + migration_protocol, + rdma_config, + dummy_prefill, + p2p_pool: Arc::new(DashMap::new()), + }) + } + + /// LMDeploy uses `enum.Enum` with `enum.auto()`. Pydantic's JSON mode + /// therefore serializes the wire values as integers, not enum names. + fn protocol_wire_value(&self) -> u8 { + match self.migration_protocol { + LMDeployMigrationProtocol::Rdma => 2, + LMDeployMigrationProtocol::Nvlink => 3, + } + } + + fn rdma_config_json(&self) -> Value { + match &self.rdma_config { + Some(cfg) => json!({ + "with_gdr": cfg.with_gdr, + "link_type": match cfg.link_type { + LMDeployRdmaLinkType::Ib => 1, + LMDeployRdmaLinkType::Roce => 2, + }, + }), + None => Value::Null, + } + } + + /// Prepare prefill request: force max_tokens=1, stream=false, + /// with_cache=true, preserve_cache=true (lmdeploy PD-specific fields). + fn prepare_prefill_request(mut request: Value, _path: &str) -> Value { + request["max_tokens"] = json!(1); + if request.get("max_completion_tokens").is_some() { + request["max_completion_tokens"] = json!(1); + } + // lmdeploy uses min_new_tokens (not min_tokens) + if let Some(min_new_tokens) = request.get("min_new_tokens").and_then(|v| v.as_u64()) { + if min_new_tokens > 1 { + request["min_new_tokens"] = json!(1); + } + } + // Force non-streaming for prefill to get JSON response with cache metadata + request["stream"] = json!(false); + if let Some(obj) = request.as_object_mut() { + obj.remove("stream_options"); + // Never forward client-supplied control-plane state to prefill. + obj.remove("migration_request"); + } + // lmdeploy PD-specific cache fields + request["with_cache"] = json!(true); + request["preserve_cache"] = json!(true); + request + } + + /// Build and strictly validate migration_request from a prefill response. + fn build_migration_request( + &self, + prefill_base_url: &str, + prefill_response: &Value, + ) -> Result { + Self::build_migration_request_static( + prefill_base_url, + prefill_response, + self.protocol_wire_value(), + self.dummy_prefill, + ) + } + + /// Static helper to build migration_request from prefill response. + /// Extracted for unit testing without requiring a full LMDeployPDRouter instance. + fn build_migration_request_static( + prefill_base_url: &str, + prefill_response: &Value, + protocol: u8, + dummy_prefill: bool, + ) -> Result { + if dummy_prefill { + return Ok(json!({ + "protocol": protocol, + "remote_engine_id": "dummy:dummy", + "remote_session_id": 0, + "remote_token_id": 0, + "remote_block_ids": [], + "is_dummy_prefill": true, + })); + } + + let session_id = Self::extract_session_id(prefill_response) + .ok_or_else(|| "Prefill response is missing a numeric remote session id".to_string())?; + + let remote_block_ids: Vec = prefill_response + .get("cache_block_ids") + .and_then(|v| v.as_array()) + .ok_or_else(|| "Prefill response is missing cache_block_ids".to_string())? + .iter() + .map(|value| { + value + .as_i64() + .ok_or_else(|| "cache_block_ids must contain only integers".to_string()) + }) + .collect::>()?; + if remote_block_ids.is_empty() { + return Err("Prefill response contains no cache_block_ids".to_string()); + } + + let remote_token_id: i64 = prefill_response + .get("remote_token_ids") + .and_then(|v| v.as_array()) + .and_then(|arr| arr.last()) + .and_then(|x| x.as_i64()) + .ok_or_else(|| { + "Prefill response is missing a final integer remote_token_id".to_string() + })?; + + Ok(json!({ + "protocol": protocol, + "remote_engine_id": prefill_base_url, + "remote_session_id": session_id, + "remote_token_id": remote_token_id, + "remote_block_ids": remote_block_ids, + "is_dummy_prefill": dummy_prefill, + })) + } + + fn extract_session_id(prefill_response: &Value) -> Option { + let parse_session_id = |value: &Value| { + value + .as_i64() + .or_else(|| value.as_str().and_then(|text| text.parse::().ok())) + }; + + prefill_response + .get("remote_session_id") + .and_then(parse_session_id) + .or_else(|| prefill_response.get("id").and_then(parse_session_id)) + } + + /// Ensure P2P RDMA connection between prefill and decode is established. + /// Cached per (prefill_base_url, decode_base_url) pair. + async fn ensure_p2p_connection( + &self, + prefill_base_url: &str, + decode_base_url: &str, + ) -> Result<(), String> { + let key = (prefill_base_url.to_string(), decode_base_url.to_string()); + let cell = self + .p2p_pool + .entry(key.clone()) + .or_insert_with(|| Arc::new(OnceCell::new())) + .clone(); + + let result = cell + .get_or_try_init(|| async { + self.establish_p2p_connection(prefill_base_url, decode_base_url) + .await + }) + .await + .map(|_| ()); + + if result.is_err() { + // Remove only the cell that failed. Another request may already have + // replaced it with a new initialization attempt. + if let Entry::Occupied(entry) = self.p2p_pool.entry(key) { + if Arc::ptr_eq(entry.get(), &cell) { + entry.remove(); + } + } + } + result + } + + async fn establish_p2p_connection( + &self, + prefill_base_url: &str, + decode_base_url: &str, + ) -> Result<(), String> { + info!( + "Establishing P2P connection: prefill={}, decode={}, protocol={:?}", + prefill_base_url, decode_base_url, self.migration_protocol + ); + + let result = async { + let (prefill_engine_info, decode_engine_info) = tokio::try_join!( + self.fetch_engine_info(prefill_base_url), + self.fetch_engine_info(decode_base_url), + )?; + + let prefill_tp_size = Self::engine_tp_size(&prefill_engine_info, prefill_base_url)?; + let decode_tp_size = Self::engine_tp_size(&decode_engine_info, decode_base_url)?; + if prefill_tp_size != decode_tp_size { + return Err(format!( + "LMDeploy PD peers must use the same tp_size: prefill={}, decode={}", + prefill_tp_size, decode_tp_size + )); + } + + let prefill_endpoints = self + .p2p_initialize( + prefill_base_url, + prefill_engine_info.clone(), + decode_base_url, + decode_engine_info.clone(), + ) + .await?; + + let decode_endpoints = self + .p2p_initialize( + decode_base_url, + decode_engine_info, + prefill_base_url, + prefill_engine_info, + ) + .await?; + + self.p2p_connect(decode_base_url, prefill_base_url, &prefill_endpoints) + .await?; + self.p2p_connect(prefill_base_url, decode_base_url, &decode_endpoints) + .await?; + Ok(()) + } + .await; + + if result.is_err() { + self.drop_p2p_connection(prefill_base_url, decode_base_url) + .await; + } else { + info!( + "P2P connection established successfully: prefill={}, decode={}", + prefill_base_url, decode_base_url + ); + } + result + } + + /// GET {base_url}/distserve/engine_info -> DistServeEngineConfig JSON + async fn fetch_engine_info(&self, base_url: &str) -> Result { + let url = format!("{}/distserve/engine_info", base_url); + let resp = self + .http_client + .get(&url) + .send() + .await + .map_err(|e| format!("GET {} failed: {}", url, e))?; + let status = resp.status(); + if !status.is_success() { + let body = resp.text().await.unwrap_or_default(); + return Err(format!("GET {} returned status {}: {}", url, status, body)); + } + let value = resp + .json::() + .await + .map_err(|e| format!("Failed to parse engine_info from {}: {}", url, e))?; + Self::parse_engine_info_value(value) + } + + fn parse_engine_info_value(value: Value) -> Result { + let value = match value { + Value::String(encoded) => serde_json::from_str(&encoded) + .map_err(|e| format!("engine_info contained invalid nested JSON: {e}"))?, + value => value, + }; + if !value.is_object() { + return Err("engine_info must decode to a JSON object".to_string()); + } + Ok(value) + } + + fn engine_tp_size(engine_info: &Value, base_url: &str) -> Result { + engine_info + .get("tp_size") + .and_then(Value::as_u64) + .filter(|size| *size > 0) + .ok_or_else(|| format!("engine_info from {base_url} has no positive integer tp_size")) + } + + /// POST {local_url}/distserve/p2p_initialize with DistServeInitRequest + async fn p2p_initialize( + &self, + local_url: &str, + local_engine_config: Value, + remote_url: &str, + remote_engine_config: Value, + ) -> Result { + let url = format!("{}/distserve/p2p_initialize", local_url); + let body = json!({ + "local_engine_id": local_url, + "local_engine_config": local_engine_config, + "remote_engine_id": remote_url, + "remote_engine_config": remote_engine_config, + "protocol": self.protocol_wire_value(), + "rdma_config": self.rdma_config_json(), + "nvlink_config": if self.migration_protocol == LMDeployMigrationProtocol::Nvlink { + json!({}) + } else { + Value::Null + }, + }); + let resp = self + .http_client + .post(&url) + .json(&body) + .send() + .await + .map_err(|e| format!("POST {} failed: {}", url, e))?; + let status = resp.status(); + if !status.is_success() { + let body = resp.text().await.unwrap_or_default(); + return Err(format!("POST {} returned status {}: {}", url, status, body)); + } + let value = resp.json::().await.map_err(|e| { + format!( + "Failed to parse p2p_initialize response from {}: {}", + url, e + ) + })?; + Self::validate_connection_response(&value, true)?; + Ok(value) + } + + /// POST {local_url}/distserve/p2p_connect with DistServeConnectionRequest + async fn p2p_connect( + &self, + local_url: &str, + remote_url: &str, + remote_endpoints: &Value, + ) -> Result<(), String> { + let url = format!("{}/distserve/p2p_connect", local_url); + let remote_engine_endpoint_info = remote_endpoints + .get("engine_endpoint_info") + .cloned() + .ok_or_else(|| "p2p_initialize response is missing engine_endpoint_info".to_string())?; + let remote_kvtransfer_endpoint_info = remote_endpoints + .get("kvtransfer_endpoint_info") + .cloned() + .filter(Value::is_array) + .ok_or_else(|| { + "p2p_initialize response is missing kvtransfer_endpoint_info".to_string() + })?; + let body = json!({ + "protocol": self.protocol_wire_value(), + "remote_engine_id": remote_url, + "remote_engine_endpoint_info": remote_engine_endpoint_info, + "remote_kvtransfer_endpoint_info": remote_kvtransfer_endpoint_info, + }); + let resp = self + .http_client + .post(&url) + .json(&body) + .send() + .await + .map_err(|e| format!("POST {} failed: {}", url, e))?; + let status = resp.status(); + if !status.is_success() { + let body = resp.text().await.unwrap_or_default(); + return Err(format!("POST {} returned status {}: {}", url, status, body)); + } + let value = resp + .json::() + .await + .map_err(|e| format!("Failed to parse p2p_connect response from {}: {}", url, e))?; + Self::validate_connection_response(&value, false)?; + Ok(()) + } + + fn validate_connection_response(value: &Value, require_endpoints: bool) -> Result<(), String> { + if value.get("status").and_then(Value::as_u64) != Some(1) { + return Err(format!("LMDeploy connection operation failed: {value}")); + } + if require_endpoints + && (value.get("engine_endpoint_info").is_none() + || !value + .get("kvtransfer_endpoint_info") + .is_some_and(Value::is_array)) + { + return Err(format!( + "LMDeploy initialization response is missing endpoint information: {value}" + )); + } + Ok(()) + } + + async fn drop_p2p_connection(&self, prefill_base_url: &str, decode_base_url: &str) { + let requests = [ + (prefill_base_url, decode_base_url), + (decode_base_url, prefill_base_url), + ]; + for (local, remote) in requests { + let url = format!("{}/distserve/p2p_drop_connect", local); + if let Err(error) = self + .http_client + .post(&url) + .json(&json!({ + "engine_id": local, + "remote_engine_id": remote, + })) + .send() + .await + { + debug!("Best-effort P2P drop failed for {}: {}", url, error); + } + } + } + + async fn invalidate_p2p_connection(&self, prefill_base_url: &str, decode_base_url: &str) { + self.p2p_pool + .remove(&(prefill_base_url.to_string(), decode_base_url.to_string())); + self.drop_p2p_connection(prefill_base_url, decode_base_url) + .await; + } + + fn invalidate_worker_connections(&self, worker_url: &str) { + self.p2p_pool + .retain(|key, _| key.0 != worker_url && key.1 != worker_url); + } + + async fn invalidate_worker_connections_and_drop(&self, worker_url: &str) { + let keys: Vec = self + .p2p_pool + .iter() + .filter(|entry| { + let (prefill, decode) = entry.key(); + prefill == worker_url || decode == worker_url + }) + .map(|entry| entry.key().clone()) + .collect(); + for (prefill, decode) in keys { + self.invalidate_p2p_connection(&prefill, &decode).await; + } + } + + /// Process lmdeploy two-stage request: prefill -> decode + async fn process_lmd_two_stage_request( + &self, + original_request: Value, + prefill_worker: Option>, + decode_worker: Arc, + path: &str, + headers: Option<&HeaderMap>, + ) -> Result { + let start_time = Instant::now(); + let request_id = format!("lmd-{}", Uuid::new_v4()); + let decode_base_url = decode_worker.url().to_string(); + let decode_url = decode_worker.endpoint_url(path); + + debug!( + "LMD prefill={}, decode={}, path={}, dummy_prefill={}", + prefill_worker + .as_ref() + .map(|worker| worker.url()) + .unwrap_or("dummy:dummy"), + decode_worker.url(), + path, + self.dummy_prefill + ); + + let (prefill_base_url, migration_request, cache_guard) = if self.dummy_prefill { + let migration_request = self + .build_migration_request("dummy:dummy", &json!({})) + .map_err(|message| PDRouterError::NetworkError { message })?; + ("dummy:dummy".to_string(), migration_request, None) + } else { + let prefill_worker = prefill_worker.ok_or_else(|| PDRouterError::NetworkError { + message: "LMDeploy PD request has no prefill worker".to_string(), + })?; + let prefill_base_url = prefill_worker.url().to_string(); + + // Establish the transport before preserving a prefill cache. This avoids + // retaining KV blocks when the control-plane handshake itself fails. + if let Err(message) = self + .ensure_p2p_connection(&prefill_base_url, &decode_base_url) + .await + { + RouterMetrics::record_pd_prefill_error(&prefill_base_url); + return Err(PDRouterError::NetworkError { + message: format!( + "P2P connection setup failed for ({}, {}): {}", + prefill_base_url, decode_base_url, message + ), + }); + } + + prefill_worker.increment_load(); + + // Stage 1: Prefill + let prefill_request = Self::prepare_prefill_request(original_request.clone(), path); + let prefill_url = prefill_worker.endpoint_url(path); + + debug!("LMD Stage 1 - Prefill: {}", prefill_url); + + let prefill_request_builder = self + .pd_router + .client + .post(&prefill_url) + .header("Content-Type", "application/json") + .header( + "Authorization", + format!( + "Bearer {}", + std::env::var("OPENAI_API_KEY").unwrap_or_default() + ), + ) + .header("X-Request-Id", &request_id); + + let prefill_response = match otel_http::send_client_request( + prefill_request_builder.json(&prefill_request), + headers, + ClientRequestOptions { + method: "POST", + url: &prefill_url, + route: Some(path), + request_phase: Some("prefill"), + }, + ) + .await + { + Ok(response) => response, + Err(error) => { + prefill_worker.decrement_load(); + RouterMetrics::record_pd_prefill_error(&prefill_base_url); + return Err(PDRouterError::NetworkError { + message: format!( + "Prefill request failed to {}: {}", + prefill_url, + error_chain(&error) + ), + }); + } + }; + + let prefill_status = prefill_response.status(); + let prefill_bytes = match prefill_response.bytes().await { + Ok(bytes) => bytes, + Err(error) => { + prefill_worker.decrement_load(); + RouterMetrics::record_pd_prefill_error(&prefill_base_url); + return Err(PDRouterError::NetworkError { + message: format!( + "Failed to read prefill response from {}: {}", + prefill_url, + error_chain(&error) + ), + }); + } + }; + prefill_worker.decrement_load(); + + if !prefill_status.is_success() { + RouterMetrics::record_pd_prefill_error(&prefill_base_url); + return Err(PDRouterError::UpstreamResponse { + url: prefill_base_url, + status: prefill_status, + body: prefill_bytes.to_vec(), + }); + } + + let prefill_response_json: Value = + serde_json::from_slice(&prefill_bytes).map_err(|e| { + RouterMetrics::record_pd_prefill_error(&prefill_base_url); + PDRouterError::NetworkError { + message: format!("Failed to parse prefill response as JSON: {e}"), + } + })?; + + let session_id = Self::extract_session_id(&prefill_response_json).ok_or_else(|| { + RouterMetrics::record_pd_prefill_error(&prefill_base_url); + PDRouterError::NetworkError { + message: "Prefill response is missing a numeric id".to_string(), + } + })?; + let cache_guard = PrefillCacheGuard::new( + self.http_client.clone(), + prefill_base_url.clone(), + session_id, + ); + let migration_request = self + .build_migration_request(&prefill_base_url, &prefill_response_json) + .map_err(|message| { + RouterMetrics::record_pd_prefill_error(&prefill_base_url); + PDRouterError::NetworkError { message } + })?; + + (prefill_base_url, migration_request, Some(cache_guard)) + }; + + // Stage 2: Decode + decode_worker.increment_load(); + let mut decode_request = original_request.clone(); + if let Some(object) = decode_request.as_object_mut() { + object.remove("with_cache"); + object.remove("preserve_cache"); + object.remove("migration_request"); + object.insert("migration_request".to_string(), migration_request); + } + + debug!("LMD Stage 2 - Decode: {}", decode_url); + + let decode_request_builder = self + .pd_router + .client + .post(&decode_url) + .header("Content-Type", "application/json") + .header( + "Authorization", + format!( + "Bearer {}", + std::env::var("OPENAI_API_KEY").unwrap_or_default() + ), + ) + .header("X-Request-Id", &request_id); + + let decode_response = match otel_http::send_client_request( + decode_request_builder.json(&decode_request), + headers, + ClientRequestOptions { + method: "POST", + url: &decode_url, + route: Some(path), + request_phase: Some("decode"), + }, + ) + .await + { + Ok(resp) => resp, + Err(e) => { + decode_worker.decrement_load(); + if !self.dummy_prefill { + self.invalidate_p2p_connection(&prefill_base_url, &decode_base_url) + .await; + } + let full_error = error_chain(&e); + let duration = start_time.elapsed(); + RouterMetrics::record_pd_decode_error(&decode_base_url); + RouterMetrics::record_pd_request(path); + RouterMetrics::record_pd_request_duration(path, duration); + RouterMetrics::record_pd_prefill_request(&prefill_base_url); + return Err(PDRouterError::NetworkError { + message: format!("Decode request failed to {}: {}", decode_url, full_error), + }); + } + }; + + decode_worker.decrement_load(); + let status = decode_response.status(); + + if !status.is_success() && !self.dummy_prefill { + self.invalidate_p2p_connection(&prefill_base_url, &decode_base_url) + .await; + } + + let duration = start_time.elapsed(); + RouterMetrics::record_pd_request(path); + RouterMetrics::record_pd_request_duration(path, duration); + RouterMetrics::record_pd_prefill_request(&prefill_base_url); + RouterMetrics::record_pd_decode_request(&decode_base_url); + + if !status.is_success() { + RouterMetrics::record_pd_decode_error(&decode_base_url); + } + + // Determine streaming vs full body (no logprobs merge for lmdeploy v1) + let is_streaming = original_request + .get("stream") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + + if is_streaming { + let mut response_builder = axum::http::Response::builder().status(status); + let mut decode_headers = + header_utils::preserve_response_headers(decode_response.headers()); + decode_headers.remove(axum::http::header::CONTENT_LENGTH); + for (name, value) in decode_headers.iter() { + response_builder = response_builder.header(name, value); + } + let stream = Box::pin(decode_response.bytes_stream()); + let body = match (status.is_success(), cache_guard) { + (true, Some(guard)) => { + let guarded_stream = futures_util::stream::unfold( + (stream, guard), + |(mut stream, guard)| async move { + match stream.next().await { + Some(item) => { + // LMDeploy sends the prefill cache-free message over + // ZMQ immediately after migration and before yielding + // the first decode output. Once a successful body chunk + // arrives, HTTP fallback cleanup would be a double free. + if item.is_ok() { + guard.mark_completed(); + } + Some((item, (stream, guard))) + } + // If no successful body chunk was observed, dropping the + // guard retains the fallback cleanup behavior. + None => None, + } + }, + ); + axum::body::Body::from_stream(guarded_stream) + } + (_, guard) => { + drop(guard); + axum::body::Body::from_stream(stream) + } + }; + response_builder + .body(body) + .map_err(|e| PDRouterError::NetworkError { + message: format!("Failed to build streaming response: {}", e), + }) + } else { + let decode_headers = decode_response.headers().clone(); + let body = decode_response + .bytes() + .await + .map_err(|e| PDRouterError::NetworkError { + message: format!("Failed to read decode response: {}", e), + })?; + if status.is_success() { + if let Some(guard) = cache_guard.as_ref() { + guard.mark_completed(); + } + } + drop(cache_guard); + let mut response_builder = axum::http::Response::builder().status(status); + for (name, value) in decode_headers.iter() { + response_builder = response_builder.header(name, value); + } + response_builder + .body(axum::body::Body::from(body)) + .map_err(|e| PDRouterError::NetworkError { + message: format!("Failed to build response: {}", e), + }) + } + } + + /// Process lmdeploy request: select PD pair via policy, then run two-stage + async fn process_lmd_request( + &self, + headers: Option<&HeaderMap>, + request_json: Value, + path: &str, + _model_id: Option<&str>, + ) -> Response { + let prefill_workers = self.pd_router.worker_registry.get_prefill_workers(); + let decode_workers = self.pd_router.worker_registry.get_decode_workers(); + + if (!self.dummy_prefill && prefill_workers.is_empty()) || decode_workers.is_empty() { + RouterMetrics::record_pd_error("server_selection"); + return ( + StatusCode::SERVICE_UNAVAILABLE, + format!( + "No workers available: {} prefill, {} decode", + prefill_workers.len(), + decode_workers.len() + ), + ) + .into_response(); + } + + let request_text = serde_json::to_string(&request_json).ok(); + let request_str = request_text.as_deref(); + let request_headers: Option> = headers.map(|h| { + h.iter() + .filter_map(|(name, value)| { + value + .to_str() + .ok() + .map(|v| (name.as_str().to_lowercase(), v.to_string())) + }) + .collect() + }); + + let prefill_policy = self.policy_registry.get_prefill_policy(); + let decode_policy = self.policy_registry.get_decode_policy(); + + let prefill_worker = if self.dummy_prefill { + None + } else { + let prefill_idx = match prefill_policy.select_worker_with_headers( + &prefill_workers, + request_str, + request_headers.as_ref(), + ) { + Some(idx) => idx, + None => { + RouterMetrics::record_pd_error("server_selection"); + return ( + StatusCode::SERVICE_UNAVAILABLE, + "Prefill policy failed to select a worker".to_string(), + ) + .into_response(); + } + }; + Some(prefill_workers[prefill_idx].clone()) + }; + + let decode_idx = match decode_policy.select_worker_with_headers( + &decode_workers, + request_str, + request_headers.as_ref(), + ) { + Some(idx) => idx, + None => { + RouterMetrics::record_pd_error("server_selection"); + return ( + StatusCode::SERVICE_UNAVAILABLE, + "Decode policy failed to select a worker".to_string(), + ) + .into_response(); + } + }; + + let decode_worker = decode_workers[decode_idx].clone(); + + if let Some(worker) = prefill_worker.as_ref() { + info!( + "LMD routing: prefill={} [policy:{}], decode={} [policy:{}]", + worker.url(), + prefill_policy.name(), + decode_worker.url(), + decode_policy.name() + ); + } else { + info!( + "LMD routing: dummy prefill, decode={} [policy:{}]", + decode_worker.url(), + decode_policy.name() + ); + } + + match self + .process_lmd_two_stage_request( + request_json, + prefill_worker, + decode_worker, + path, + headers, + ) + .await + { + Ok(response) => response, + Err(PDRouterError::UpstreamResponse { status, body, .. }) => Response::builder() + .status(status) + .header(http::header::CONTENT_TYPE, "application/json") + .body(Body::from(body)) + .unwrap_or_else(|error| { + error!("Failed to build upstream response: {}", error); + StatusCode::INTERNAL_SERVER_ERROR.into_response() + }), + Err(e) => { + error!("LMD two-stage processing failed: {}", e); + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Request processing failed: {}", e), + ) + .into_response() + } + } + } + + pub async fn add_prefill_server(&self, url: String) -> Result { + self.pd_router.add_prefill_server(url).await + } + + pub async fn add_decode_server(&self, url: String) -> Result { + self.pd_router.add_decode_server(url).await + } + + pub fn register_prefill_server_unchecked(&self, url: String) -> Result { + self.pd_router.register_prefill_server_unchecked(url) + } + + pub fn register_decode_server_unchecked(&self, url: String) -> Result { + self.pd_router.register_decode_server_unchecked(url) + } + + pub async fn remove_prefill_server(&self, url: &str) -> Result { + self.invalidate_worker_connections_and_drop(url).await; + self.pd_router.remove_prefill_server(url).await + } + + pub async fn remove_decode_server(&self, url: &str) -> Result { + self.invalidate_worker_connections_and_drop(url).await; + self.pd_router.remove_decode_server(url).await + } +} + +#[async_trait] +impl RouterTrait for LMDeployPDRouter { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + async fn health(&self, req: Request) -> Response { + self.pd_router.health(req).await + } + + async fn health_generate(&self, _req: Request) -> Response { + // lmdeploy has no /health_generate endpoint + ( + StatusCode::NOT_IMPLEMENTED, + "health_generate not supported for lmdeploy PD router", + ) + .into_response() + } + + async fn get_server_info(&self, _req: Request) -> Response { + // lmdeploy has no /get_server_info endpoint + ( + StatusCode::NOT_IMPLEMENTED, + "get_server_info not supported for lmdeploy PD router", + ) + .into_response() + } + + async fn get_models(&self, req: Request) -> Response { + self.pd_router.get_models(req).await + } + + async fn get_model_info(&self, req: Request) -> Response { + self.pd_router.get_model_info(req).await + } + + async fn route_generate( + &self, + _headers: Option<&HeaderMap>, + _body: &crate::protocols::spec::GenerateRequest, + _model_id: Option<&str>, + ) -> Response { + ( + StatusCode::NOT_IMPLEMENTED, + "LMDeploy /generate does not expose the cache migration fields required for PD", + ) + .into_response() + } + + async fn route_chat( + &self, + headers: Option<&HeaderMap>, + body: &crate::protocols::spec::ChatCompletionRequest, + model_id: Option<&str>, + ) -> Response { + let request_json = match serde_json::to_value(body) { + Ok(json) => json, + Err(e) => { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Serialization error: {}", e), + ) + .into_response() + } + }; + self.process_lmd_request(headers, request_json, "/v1/chat/completions", model_id) + .await + } + + async fn route_completion( + &self, + headers: Option<&HeaderMap>, + body: &crate::protocols::spec::CompletionRequest, + model_id: Option<&str>, + ) -> Response { + let request_json = match serde_json::to_value(body) { + Ok(json) => json, + Err(e) => { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Serialization error: {}", e), + ) + .into_response() + } + }; + self.process_lmd_request(headers, request_json, "/v1/completions", model_id) + .await + } + + async fn route_responses( + &self, + _headers: Option<&HeaderMap>, + _body: &crate::protocols::spec::ResponsesRequest, + _model_id: Option<&str>, + ) -> Response { + ( + StatusCode::NOT_IMPLEMENTED, + "LMDeploy /v1/responses does not expose the cache migration fields required for PD", + ) + .into_response() + } + + async fn get_response(&self, _headers: Option<&HeaderMap>, _response_id: &str) -> Response { + ( + StatusCode::NOT_IMPLEMENTED, + "Responses retrieve not supported for lmdeploy PD router", + ) + .into_response() + } + + async fn cancel_response(&self, _headers: Option<&HeaderMap>, _response_id: &str) -> Response { + ( + StatusCode::NOT_IMPLEMENTED, + "Responses cancel not supported for lmdeploy PD router", + ) + .into_response() + } + + async fn route_embeddings( + &self, + _headers: Option<&HeaderMap>, + _body: &crate::protocols::spec::EmbeddingRequest, + _model_id: Option<&str>, + ) -> Response { + ( + StatusCode::NOT_IMPLEMENTED, + "Embeddings not supported for lmdeploy PD router", + ) + .into_response() + } + + async fn flush_cache(&self) -> Response { + // lmdeploy has POST /terminate but requires --allow-terminate-by-client + ( + StatusCode::NOT_IMPLEMENTED, + "flush_cache not supported for lmdeploy PD router", + ) + .into_response() + } + + async fn get_worker_loads(&self) -> Response { + self.pd_router.get_worker_loads().await + } + + fn router_type(&self) -> &'static str { + "lmdeploy-pd" + } + + fn is_pd_mode(&self) -> bool { + true + } + + fn readiness(&self) -> Response { + self.pd_router.readiness() + } + + async fn route_transparent( + &self, + headers: Option<&HeaderMap>, + path: &str, + method: &Method, + body: serde_json::Value, + ) -> Response { + if *method != Method::POST { + return ( + StatusCode::METHOD_NOT_ALLOWED, + "Only POST requests are supported for transparent proxy", + ) + .into_response(); + } + if !matches!(path, "/v1/chat/completions" | "/v1/completions") { + return ( + StatusCode::NOT_IMPLEMENTED, + format!("Endpoint {path} is not supported by the LMDeploy PD protocol"), + ) + .into_response(); + } + self.process_lmd_request(headers, body, path, None).await + } +} + +#[async_trait] +impl WorkerManagement for LMDeployPDRouter { + async fn add_worker(&self, worker_url: &str) -> Result { + self.pd_router.add_worker(worker_url).await + } + + fn remove_worker(&self, worker_url: &str) { + self.invalidate_worker_connections(worker_url); + self.pd_router.remove_worker(worker_url); + } + + fn get_worker_urls(&self) -> Vec { + self.pd_router.get_worker_urls() + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_prepare_prefill_request_sets_max_tokens_1() { + let request = json!({ + "model": "test", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 512, + "stream": true + }); + let result = LMDeployPDRouter::prepare_prefill_request(request, "/v1/chat/completions"); + assert_eq!(result["max_tokens"], 1); + assert_eq!(result["stream"], false); + assert_eq!(result["with_cache"], true); + assert_eq!(result["preserve_cache"], true); + } + + #[test] + fn test_prepare_prefill_request_sets_max_completion_tokens_1() { + let request = json!({ + "model": "test", + "max_tokens": 512, + "max_completion_tokens": 256 + }); + let result = LMDeployPDRouter::prepare_prefill_request(request, "/v1/chat/completions"); + assert_eq!(result["max_tokens"], 1); + assert_eq!(result["max_completion_tokens"], 1); + } + + #[test] + fn test_prepare_prefill_request_clamps_min_new_tokens() { + let request = json!({ + "model": "test", + "max_tokens": 512, + "min_new_tokens": 100 + }); + let result = LMDeployPDRouter::prepare_prefill_request(request, "/v1/chat/completions"); + assert_eq!(result["min_new_tokens"], 1); + } + + #[test] + fn test_prepare_prefill_request_removes_stream_options() { + let request = json!({ + "model": "test", + "max_tokens": 512, + "stream": true, + "stream_options": {"include_usage": true} + }); + let result = LMDeployPDRouter::prepare_prefill_request(request, "/v1/chat/completions"); + assert_eq!(result["stream"], false); + assert!(result.get("stream_options").is_none()); + } + + #[test] + fn test_protocol_config_serialization() { + assert_eq!( + serde_json::to_value(LMDeployMigrationProtocol::Rdma).unwrap(), + serde_json::json!("rdma") + ); + assert_eq!( + serde_json::to_value(LMDeployMigrationProtocol::Nvlink).unwrap(), + serde_json::json!("nvlink") + ); + } + + #[test] + fn test_migration_protocol_default_is_rdma() { + assert_eq!( + LMDeployMigrationProtocol::default(), + LMDeployMigrationProtocol::Rdma + ); + } + + #[test] + fn test_build_migration_request_extracts_fields() { + let prefill_response = json!({ + "id": 12345, + "cache_block_ids": [1, 2, 3], + "remote_token_ids": [10, 20, 30] + }); + let result = LMDeployPDRouter::build_migration_request_static( + "http://prefill:23333", + &prefill_response, + 2, + false, + ); + let req = result.unwrap(); + assert_eq!(req["protocol"], 2); + assert_eq!(req["remote_engine_id"], "http://prefill:23333"); + assert_eq!(req["remote_session_id"], 12345); + assert_eq!(req["remote_token_id"], 30); + assert_eq!(req["remote_block_ids"], json!([1, 2, 3])); + assert_eq!(req["is_dummy_prefill"], false); + } + + #[test] + fn test_build_migration_request_prefers_remote_session_id() { + let prefill_response = json!({ + "id": "chatcmpl-not-a-number", + "remote_session_id": 67890, + "cache_block_ids": [5], + "remote_token_ids": [42] + }); + let result = LMDeployPDRouter::build_migration_request_static( + "http://p:8000", + &prefill_response, + 2, + false, + ); + let req = result.unwrap(); + assert_eq!(req["remote_session_id"], 67890); + } + + #[test] + fn test_build_migration_request_parses_string_id() { + let prefill_response = json!({ + "id": "67890", + "cache_block_ids": [5], + "remote_token_ids": [42] + }); + let result = LMDeployPDRouter::build_migration_request_static( + "http://p:8000", + &prefill_response, + 2, + false, + ); + let req = result.unwrap(); + assert_eq!(req["remote_session_id"], 67890); + assert_eq!(req["remote_token_id"], 42); + assert_eq!(req["remote_block_ids"], json!([5])); + assert_eq!(req["is_dummy_prefill"], false); + } + + #[test] + fn test_build_migration_request_rejects_missing_id() { + let prefill_response = json!({ + "cache_block_ids": [1, 2], + "remote_token_ids": [10] + }); + let result = LMDeployPDRouter::build_migration_request_static( + "http://p:8000", + &prefill_response, + 2, + false, + ); + assert!(result.is_err()); + } + + #[test] + fn test_build_migration_request_rejects_non_numeric_id() { + let prefill_response = json!({ + "id": "not-a-number", + "cache_block_ids": [1], + "remote_token_ids": [10] + }); + let result = LMDeployPDRouter::build_migration_request_static( + "http://p:8000", + &prefill_response, + 2, + false, + ); + assert!(result.is_err()); + } + + #[test] + fn test_build_migration_request_rejects_missing_cache_block_ids() { + let prefill_response = json!({ + "id": 100, + "remote_token_ids": [7] + }); + let result = LMDeployPDRouter::build_migration_request_static( + "http://p:8000", + &prefill_response, + 2, + false, + ); + assert!(result.is_err()); + } + + #[test] + fn test_build_migration_request_rejects_empty_remote_token_ids() { + let prefill_response = json!({ + "id": 200, + "cache_block_ids": [1, 2], + "remote_token_ids": [] + }); + let result = LMDeployPDRouter::build_migration_request_static( + "http://p:8000", + &prefill_response, + 2, + false, + ); + assert!(result.is_err()); + } + + #[test] + fn test_build_dummy_migration_request_skips_prefill_metadata() { + let request = + LMDeployPDRouter::build_migration_request_static("unused", &json!({}), 2, true) + .unwrap(); + assert_eq!(request["protocol"], 2); + assert_eq!(request["remote_engine_id"], "dummy:dummy"); + assert_eq!(request["remote_session_id"], 0); + assert_eq!(request["remote_block_ids"], json!([])); + assert_eq!(request["is_dummy_prefill"], true); + } + + #[test] + fn test_engine_info_nested_json_is_decoded() { + let value = Value::String(r#"{"tp_size":8,"dp_size":1}"#.to_string()); + let decoded = LMDeployPDRouter::parse_engine_info_value(value).unwrap(); + assert_eq!(decoded["tp_size"], 8); + assert_eq!(decoded["dp_size"], 1); + } + + #[test] + fn test_connection_response_requires_success_and_endpoints() { + let response = json!({ + "status": 1, + "engine_endpoint_info": {"zmq_address": "tcp://host:1234"}, + "kvtransfer_endpoint_info": [], + }); + assert!(LMDeployPDRouter::validate_connection_response(&response, true).is_ok()); + assert!( + LMDeployPDRouter::validate_connection_response(&json!({"status": 2}), false).is_err() + ); + } + + #[test] + fn test_prepare_prefill_request_sets_with_cache_and_preserve_cache() { + let request = json!({ + "model": "test-model", + "messages": [{"role": "user", "content": "hi"}], + "max_tokens": 512, + }); + let result = LMDeployPDRouter::prepare_prefill_request(request, "/v1/chat/completions"); + assert_eq!(result["with_cache"], true); + assert_eq!(result["preserve_cache"], true); + assert_eq!(result["max_tokens"], 1); + assert_eq!(result["stream"], false); + } + + #[test] + fn test_prepare_prefill_request_preserves_token_io_fields() { + let request = json!({ + "model": "test-model", + "messages": [], + "input_ids": [151644, 8948, 198], + "return_token_ids": true, + "do_preprocess": false, + "max_tokens": 32, + }); + let result = LMDeployPDRouter::prepare_prefill_request(request, "/v1/chat/completions"); + assert_eq!(result["input_ids"], json!([151644, 8948, 198])); + assert_eq!(result["return_token_ids"], true); + assert_eq!(result["do_preprocess"], false); + assert_eq!(result["max_tokens"], 1); + } + + #[test] + fn test_prepare_prefill_request_handles_generate_path() { + let request = json!({ + "model": "test-model", + "prompt": "hello", + "max_tokens": 256, + "stream": true, + }); + let result = LMDeployPDRouter::prepare_prefill_request(request, "/generate"); + assert_eq!(result["max_tokens"], 1); + assert_eq!(result["stream"], false); + assert_eq!(result["with_cache"], true); + assert_eq!(result["preserve_cache"], true); + } +} diff --git a/lmdeploy/router/src/routers/http/mod.rs b/lmdeploy/router/src/routers/http/mod.rs new file mode 100644 index 0000000000..53380af321 --- /dev/null +++ b/lmdeploy/router/src/routers/http/mod.rs @@ -0,0 +1,7 @@ +//! HTTP router implementations + +pub mod lmdeploy_pd_router; +pub mod openai_router; +pub mod pd_router; +pub mod pd_types; +pub mod router; diff --git a/lmdeploy/router/src/routers/http/openai_router.rs b/lmdeploy/router/src/routers/http/openai_router.rs new file mode 100644 index 0000000000..f0a215fe21 --- /dev/null +++ b/lmdeploy/router/src/routers/http/openai_router.rs @@ -0,0 +1,429 @@ +//! OpenAI router implementation (reqwest-based) + +use crate::config::CircuitBreakerConfig; +use crate::core::{CircuitBreaker, CircuitBreakerConfig as CoreCircuitBreakerConfig}; +use crate::otel_http::{self, ClientRequestOptions}; +use crate::protocols::spec::{ChatCompletionRequest, CompletionRequest, GenerateRequest}; +use async_trait::async_trait; +use axum::{ + body::Body, + extract::Request, + http::{header::CONTENT_TYPE, HeaderMap, HeaderValue, StatusCode}, + response::{IntoResponse, Response}, +}; +use futures_util::StreamExt; +use std::{ + any::Any, + sync::atomic::{AtomicBool, Ordering}, +}; + +/// Router for OpenAI backend +#[derive(Debug)] +pub struct OpenAIRouter { + /// HTTP client for upstream OpenAI-compatible API + client: reqwest::Client, + /// Base URL for identification (no trailing slash) + base_url: String, + /// Circuit breaker + circuit_breaker: CircuitBreaker, + /// Health status + healthy: AtomicBool, +} + +impl OpenAIRouter { + /// Create a new OpenAI router + pub async fn new( + base_url: String, + circuit_breaker_config: Option, + ) -> Result { + let client = reqwest::Client::builder() + .timeout(std::time::Duration::from_secs(300)) + .build() + .map_err(|e| format!("Failed to create HTTP client: {}", e))?; + + let base_url = base_url.trim_end_matches('/').to_string(); + + // Convert circuit breaker config + let core_cb_config = circuit_breaker_config + .map(|cb| CoreCircuitBreakerConfig { + failure_threshold: cb.failure_threshold, + success_threshold: cb.success_threshold, + timeout_duration: std::time::Duration::from_secs(cb.timeout_duration_secs), + window_duration: std::time::Duration::from_secs(cb.window_duration_secs), + }) + .unwrap_or_default(); + + let circuit_breaker = CircuitBreaker::with_config(core_cb_config); + + Ok(Self { + client, + base_url, + circuit_breaker, + healthy: AtomicBool::new(true), + }) + } +} + +#[async_trait] +impl super::super::WorkerManagement for OpenAIRouter { + async fn add_worker(&self, _worker_url: &str) -> Result { + Err("Cannot add workers to OpenAI router".to_string()) + } + + fn remove_worker(&self, _worker_url: &str) { + // No-op for OpenAI router + } + + fn get_worker_urls(&self) -> Vec { + vec![self.base_url.clone()] + } +} + +#[async_trait] +impl super::super::RouterTrait for OpenAIRouter { + fn as_any(&self) -> &dyn Any { + self + } + + async fn health(&self, _req: Request) -> Response { + // Simple upstream probe: GET {base}/v1/models without auth + let url = format!("{}/v1/models", self.base_url); + match self + .client + .get(&url) + .timeout(std::time::Duration::from_secs(2)) + .send() + .await + { + Ok(resp) => { + let code = resp.status(); + // Treat success and auth-required as healthy (endpoint reachable) + if code.is_success() || code.as_u16() == 401 || code.as_u16() == 403 { + (StatusCode::OK, "OK").into_response() + } else { + ( + StatusCode::SERVICE_UNAVAILABLE, + format!("Upstream status: {}", code), + ) + .into_response() + } + } + Err(e) => ( + StatusCode::SERVICE_UNAVAILABLE, + format!("Upstream error: {}", e), + ) + .into_response(), + } + } + + async fn health_generate(&self, _req: Request) -> Response { + // For OpenAI, health_generate is the same as health + self.health(_req).await + } + + async fn get_server_info(&self, _req: Request) -> Response { + let info = serde_json::json!({ + "router_type": "openai", + "workers": 1, + "base_url": &self.base_url + }); + (StatusCode::OK, info.to_string()).into_response() + } + + async fn get_models(&self, req: Request) -> Response { + // Proxy to upstream /v1/models; forward Authorization header if provided + let headers = req.headers(); + + let url = format!("{}/v1/models", self.base_url); + let mut upstream = self.client.get(&url); + + if let Some(auth) = headers + .get("authorization") + .or_else(|| headers.get("Authorization")) + { + upstream = upstream.header("Authorization", auth); + } + + match otel_http::send_client_request( + upstream, + Some(headers), + ClientRequestOptions { + method: "GET", + url: &url, + route: Some("/v1/models"), + request_phase: None, + }, + ) + .await + { + Ok(res) => { + let status = StatusCode::from_u16(res.status().as_u16()) + .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); + let content_type = res.headers().get(CONTENT_TYPE).cloned(); + match res.bytes().await { + Ok(body) => { + let mut response = Response::new(axum::body::Body::from(body)); + *response.status_mut() = status; + if let Some(ct) = content_type { + response.headers_mut().insert(CONTENT_TYPE, ct); + } + response + } + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to read upstream response: {}", e), + ) + .into_response(), + } + } + Err(e) => ( + StatusCode::BAD_GATEWAY, + format!("Failed to contact upstream: {}", e), + ) + .into_response(), + } + } + + async fn get_model_info(&self, _req: Request) -> Response { + // Not directly supported without model param; return 501 + ( + StatusCode::NOT_IMPLEMENTED, + "get_model_info not implemented for OpenAI router", + ) + .into_response() + } + + async fn route_generate( + &self, + _headers: Option<&HeaderMap>, + _body: &GenerateRequest, + _model_id: Option<&str>, + ) -> Response { + // Generate endpoint is backend-specific and not supported by the OpenAI backend + ( + StatusCode::NOT_IMPLEMENTED, + "Generate endpoint not supported for OpenAI backend", + ) + .into_response() + } + + async fn route_chat( + &self, + headers: Option<&HeaderMap>, + body: &ChatCompletionRequest, + _model_id: Option<&str>, + ) -> Response { + if !self.circuit_breaker.can_execute() { + return (StatusCode::SERVICE_UNAVAILABLE, "Circuit breaker open").into_response(); + } + + // Serialize request body, removing backend-specific fields + let mut payload = match serde_json::to_value(body) { + Ok(v) => v, + Err(e) => { + return ( + StatusCode::BAD_REQUEST, + format!("Failed to serialize request: {}", e), + ) + .into_response(); + } + }; + if let Some(obj) = payload.as_object_mut() { + for key in [ + "top_k", + "min_p", + "min_tokens", + "regex", + "ebnf", + "stop_token_ids", + "no_stop_trim", + "ignore_eos", + "continue_final_message", + "skip_special_tokens", + "separate_reasoning", + "stream_reasoning", + "chat_template_kwargs", + "repetition_penalty", + ] { + obj.remove(key); + } + } + + let url = format!("{}/v1/chat/completions", self.base_url); + let mut req = self.client.post(&url).json(&payload); + + // Forward Authorization header if provided + if let Some(h) = headers { + if let Some(auth) = h.get("authorization").or_else(|| h.get("Authorization")) { + req = req.header("Authorization", auth); + } + } + + // Accept SSE when stream=true + if body.stream { + req = req.header("Accept", "text/event-stream"); + } + + let resp = match otel_http::send_client_request( + req, + headers, + ClientRequestOptions { + method: "POST", + url: &url, + route: Some("/v1/chat/completions"), + request_phase: Some("inference"), + }, + ) + .await + { + Ok(r) => r, + Err(e) => { + self.circuit_breaker.record_failure(); + return ( + StatusCode::SERVICE_UNAVAILABLE, + format!("Failed to contact upstream: {}", e), + ) + .into_response(); + } + }; + + let status = StatusCode::from_u16(resp.status().as_u16()) + .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); + + if !body.stream { + // Capture Content-Type before consuming response body + let content_type = resp.headers().get(CONTENT_TYPE).cloned(); + match resp.bytes().await { + Ok(body) => { + self.circuit_breaker.record_success(); + let mut response = Response::new(axum::body::Body::from(body)); + *response.status_mut() = status; + if let Some(ct) = content_type { + response.headers_mut().insert(CONTENT_TYPE, ct); + } + response + } + Err(e) => { + self.circuit_breaker.record_failure(); + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to read response: {}", e), + ) + .into_response() + } + } + } else { + // Stream SSE bytes to client + let stream = resp.bytes_stream(); + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + tokio::spawn(async move { + let mut s = stream; + while let Some(chunk) = s.next().await { + match chunk { + Ok(bytes) => { + if tx.send(Ok(bytes)).is_err() { + break; + } + } + Err(e) => { + let _ = tx.send(Err(format!("Stream error: {}", e))); + break; + } + } + } + }); + let mut response = Response::new(Body::from_stream( + tokio_stream::wrappers::UnboundedReceiverStream::new(rx), + )); + *response.status_mut() = status; + response + .headers_mut() + .insert(CONTENT_TYPE, HeaderValue::from_static("text/event-stream")); + response + } + } + + async fn route_completion( + &self, + _headers: Option<&HeaderMap>, + _body: &CompletionRequest, + _model_id: Option<&str>, + ) -> Response { + // Completion endpoint not implemented for OpenAI backend + ( + StatusCode::NOT_IMPLEMENTED, + "Completion endpoint not implemented for OpenAI backend", + ) + .into_response() + } + + async fn route_responses( + &self, + _headers: Option<&HeaderMap>, + _body: &crate::protocols::spec::ResponsesRequest, + _model_id: Option<&str>, + ) -> Response { + ( + StatusCode::NOT_IMPLEMENTED, + "Responses endpoint not implemented for OpenAI router", + ) + .into_response() + } + + async fn get_response(&self, _headers: Option<&HeaderMap>, _response_id: &str) -> Response { + ( + StatusCode::NOT_IMPLEMENTED, + "Responses retrieve endpoint not implemented for OpenAI router", + ) + .into_response() + } + + async fn cancel_response(&self, _headers: Option<&HeaderMap>, _response_id: &str) -> Response { + ( + StatusCode::NOT_IMPLEMENTED, + "Responses cancel endpoint not implemented for OpenAI router", + ) + .into_response() + } + + async fn flush_cache(&self) -> Response { + ( + StatusCode::NOT_IMPLEMENTED, + "flush_cache not supported for OpenAI router", + ) + .into_response() + } + + async fn get_worker_loads(&self) -> Response { + ( + StatusCode::NOT_IMPLEMENTED, + "get_worker_loads not supported for OpenAI router", + ) + .into_response() + } + + fn router_type(&self) -> &'static str { + "openai" + } + + fn readiness(&self) -> Response { + if self.healthy.load(Ordering::Acquire) && self.circuit_breaker.can_execute() { + (StatusCode::OK, "Ready").into_response() + } else { + (StatusCode::SERVICE_UNAVAILABLE, "Not ready").into_response() + } + } + + async fn route_embeddings( + &self, + _headers: Option<&HeaderMap>, + _body: &crate::protocols::spec::EmbeddingRequest, + _model_id: Option<&str>, + ) -> Response { + ( + StatusCode::NOT_IMPLEMENTED, + "Embeddings endpoint not implemented for OpenAI backend", + ) + .into_response() + } +} diff --git a/lmdeploy/router/src/routers/http/pd_router.rs b/lmdeploy/router/src/routers/http/pd_router.rs new file mode 100644 index 0000000000..12ec9358a1 --- /dev/null +++ b/lmdeploy/router/src/routers/http/pd_router.rs @@ -0,0 +1,1245 @@ +// PD (Prefill-Decode) Router Implementation +// This module handles routing for disaggregated prefill-decode systems +use super::pd_types::PDRouterError; +use crate::core::{ + BasicWorker, CircuitBreakerConfig, HealthConfig, Worker, WorkerFactory, WorkerRegistry, + WorkerType, +}; +use crate::policies::{LoadBalancingPolicy, PolicyRegistry}; +use crate::routers::header_utils; +use axum::{ + body::Body, + extract::Request, + http::{HeaderMap, StatusCode}, + response::{IntoResponse, Response}, + Json, +}; +use reqwest::Client; +use serde_json::{json, Value}; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::Duration; +use tracing::{debug, error, info, warn}; + +#[derive(Debug, Clone)] +pub struct PdRouterBase { + pub worker_registry: Arc, + pub policy_registry: Arc, + pub worker_startup_timeout_secs: u64, + pub worker_startup_check_interval_secs: u64, + pub worker_loads: Arc>>, + pub load_monitor_handle: Option>>, + pub client: Client, + pub circuit_breaker_config: CircuitBreakerConfig, +} + +impl PdRouterBase { + // Private helper method to perform health check on a new server + async fn wait_for_server_health(&self, url: &str) -> Result<(), PDRouterError> { + crate::routers::http::router::Router::wait_for_healthy_workers( + &[url.to_string()], + self.worker_startup_timeout_secs, + self.worker_startup_check_interval_secs, + ) + .await + .map_err(|_| PDRouterError::HealthCheckFailed { + url: url.to_string(), + }) + } + + // Generic helper for processing all workers with an endpoint + async fn process_workers( + &self, + worker_type_enum: WorkerType, + worker_type: &str, + endpoint: &str, + ) -> (Vec, Vec) { + let mut results = Vec::new(); + let mut errors = Vec::new(); + + // Get workers from registry based on type + let workers = self.worker_registry.get_by_type(&worker_type_enum); + let urls: Vec = workers.iter().map(|w| w.url().to_string()).collect(); + + // Process each worker + for worker_url in urls { + let url = format!("{}/{}", worker_url, endpoint); + match self.client.post(&url).send().await { + Ok(res) if res.status().is_success() => { + results.push(format!("{} {}: OK", worker_type, worker_url)); + } + Ok(res) => { + errors.push(format!( + "{} {} returned status: {}", + worker_type, + worker_url, + res.status() + )); + } + Err(e) => { + errors.push(format!("{} {} error: {}", worker_type, worker_url, e)); + } + } + } + + (results, errors) + } + + // Helper to get prefill worker URLs + fn get_prefill_worker_urls(&self) -> Vec { + self.worker_registry + .get_prefill_workers() + .iter() + .map(|w| w.url().to_string()) + .collect() + } + + // Helper to get decode worker URLs + fn get_decode_worker_urls(&self) -> Vec { + self.worker_registry + .get_decode_workers() + .iter() + .map(|w| w.url().to_string()) + .collect() + } + + /// Start profiling on a backend server + pub async fn start_profiling(&self, worker_url: &str) { + let url = format!("{}/start_profile", worker_url); + match self.client.post(&url).send().await { + Ok(res) if res.status().is_success() => { + info!("Started profiling on {}", worker_url); + } + Ok(res) => { + warn!( + "Failed to start profiling on {}: status {}", + worker_url, + res.status() + ); + } + Err(e) => { + warn!("Error starting profiling on {}: {}", worker_url, e); + } + } + } + + /// Stop profiling on a backend server + pub async fn stop_profiling(&self, worker_url: &str) { + let url = format!("{}/stop_profile", worker_url); + match self.client.post(&url).send().await { + Ok(res) if res.status().is_success() => { + info!("Stopped profiling on {}", worker_url); + } + Ok(res) => { + warn!( + "Failed to stop profiling on {}: status {}", + worker_url, + res.status() + ); + } + Err(e) => { + warn!("Error stopping profiling on {}: {}", worker_url, e); + } + } + } + + // Helper for proxying requests to the first prefill worker + async fn proxy_to_first_prefill_worker( + &self, + endpoint: &str, + headers: Option>, + ) -> Response { + let workers = self.worker_registry.get_prefill_workers(); + let first_worker_url = workers.first().map(|w| w.url().to_string()); + + if let Some(worker_url) = first_worker_url { + self.proxy_to_worker(worker_url, endpoint, headers).await + } else { + ( + StatusCode::SERVICE_UNAVAILABLE, + "No prefill servers available".to_string(), + ) + .into_response() + } + } + + // Generic helper for proxying to a specific worker + async fn proxy_to_worker( + &self, + worker_url: String, + endpoint: &str, + headers: Option>, + ) -> Response { + let url = format!("{}/{}", worker_url, endpoint); + let mut request_builder = self.client.get(&url); + + // Add headers if provided + if let Some(headers) = headers { + for (name, value) in headers { + request_builder = request_builder.header(name, value); + } + } + + match request_builder.send().await { + Ok(res) if res.status().is_success() => { + let response_headers = header_utils::preserve_response_headers(res.headers()); + + match res.bytes().await { + Ok(body) => { + let mut response = Response::new(axum::body::Body::from(body)); + *response.status_mut() = StatusCode::OK; + *response.headers_mut() = response_headers; + response + } + Err(e) => { + error!("Failed to read response body: {}", e); + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to read response body: {}", e), + ) + .into_response() + } + } + } + Ok(res) => { + let status = StatusCode::from_u16(res.status().as_u16()) + .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); + (status, format!("{} server returned status: ", res.status())).into_response() + } + Err(e) => { + error!("Failed to proxy request server: {}", e); + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to proxy request: {}", e), + ) + .into_response() + } + } + } + + pub async fn add_prefill_server(&self, url: String) -> Result { + // Wait for the new server to be healthy + self.wait_for_server_health(&url).await?; + + self.register_prefill_server_unchecked(url) + } + + /// Register a prefill server without probing it first. This is used by + /// LMDeploy's startup-time `/nodes/add` callback, before the API server is + /// able to answer `/health`. + pub fn register_prefill_server_unchecked(&self, url: String) -> Result { + if self.worker_registry.get_by_url(&url).is_some() { + return Err(PDRouterError::WorkerAlreadyExists { url: url.clone() }); + } + let worker_arc: Arc = Arc::from(WorkerFactory::create_prefill_with_config( + url.clone(), + self.circuit_breaker_config.clone(), + )); + self.register_and_notify(worker_arc); + info!("Added prefill server: {}", url); + + Ok(format!("Successfully added prefill server: {}", url)) + } + + pub async fn add_decode_server(&self, url: String) -> Result { + // Wait for the new server to be healthy + self.wait_for_server_health(&url).await?; + + self.register_decode_server_unchecked(url) + } + + /// Register a decode server without probing it first. See + /// `register_prefill_server_unchecked` for why startup-time registration + /// must not synchronously perform a health check. + pub fn register_decode_server_unchecked(&self, url: String) -> Result { + if self.worker_registry.get_by_url(&url).is_some() { + return Err(PDRouterError::WorkerAlreadyExists { url: url.clone() }); + } + let worker_arc: Arc = Arc::from(WorkerFactory::create_decode_with_config( + url.clone(), + self.circuit_breaker_config.clone(), + )); + self.register_and_notify(worker_arc); + info!("Added decode server: {}", url); + + Ok(format!("Successfully added decode server: {}", url)) + } + + /// Register a worker and notify the policy registry. + fn register_and_notify(&self, worker_arc: Arc) { + self.worker_registry.register(worker_arc.clone()); + let model_id = worker_arc.model_id(); + let policy = self.policy_registry.on_worker_added(model_id, None); + if policy.name() == "cache_aware" { + if let Some(cache_aware) = policy + .as_any() + .downcast_ref::() + { + let model_workers = self.worker_registry.get_by_model_fast(model_id); + cache_aware.init_workers(&model_workers); + } + } + } + + pub async fn remove_prefill_server(&self, url: &str) -> Result { + let model_id = match self.worker_registry.get_by_url(url) { + Some(worker) => worker.model_id().to_string(), + None => { + return Err(PDRouterError::WorkerNotFound { + url: url.to_string(), + }); + } + }; + + let removed = self.worker_registry.remove_by_url(url); + + if removed.is_some() { + self.policy_registry.on_worker_removed(&model_id); + if let Some(policy) = self.policy_registry.get_policy(&model_id) { + if policy.name() == "cache_aware" { + if let Some(cache_aware) = policy + .as_any() + .downcast_ref::() + { + cache_aware.remove_worker_by_url(url); + } + } + } + } + + if removed.is_some() { + info!("Removed prefill server: {}", url); + Ok(format!("Successfully removed prefill server: {}", url)) + } else { + Err(PDRouterError::WorkerNotFound { + url: url.to_string(), + }) + } + } + + pub async fn remove_decode_server(&self, url: &str) -> Result { + let model_id = match self.worker_registry.get_by_url(url) { + Some(worker) => worker.model_id().to_string(), + None => { + return Err(PDRouterError::WorkerNotFound { + url: url.to_string(), + }); + } + }; + + let removed = self.worker_registry.remove_by_url(url); + + if removed.is_some() { + // Notify PolicyRegistry about the removed worker + self.policy_registry.on_worker_removed(&model_id); + + // Get the policy for this model to update cache-aware if needed + if let Some(policy) = self.policy_registry.get_policy(&model_id) { + if policy.name() == "cache_aware" { + if let Some(cache_aware) = policy + .as_any() + .downcast_ref::() + { + cache_aware.remove_worker_by_url(url); + } + } + } + } + + if removed.is_some() { + info!("Removed decode server: {}", url); + Ok(format!("Successfully removed decode server: {}", url)) + } else { + Err(PDRouterError::WorkerNotFound { + url: url.to_string(), + }) + } + } + + #[allow(clippy::too_many_arguments)] + pub async fn new( + prefill_urls: Vec, + decode_urls: Vec, + ctx: &Arc, + ) -> Result { + // Convert config CircuitBreakerConfig to core CircuitBreakerConfig + let circuit_breaker_config = ctx.router_config.effective_circuit_breaker_config(); + let core_cb_config = CircuitBreakerConfig { + failure_threshold: circuit_breaker_config.failure_threshold, + success_threshold: circuit_breaker_config.success_threshold, + timeout_duration: Duration::from_secs(circuit_breaker_config.timeout_duration_secs), + window_duration: Duration::from_secs(circuit_breaker_config.window_duration_secs), + }; + + let mut prefill_workers_urls = vec![]; + let mut decode_workers_urls = vec![]; + let health_config = HealthConfig { + timeout_secs: ctx.router_config.health_check.timeout_secs, + check_interval_secs: ctx.router_config.health_check.check_interval_secs, + endpoint: ctx.router_config.health_check.endpoint.clone(), + failure_threshold: ctx.router_config.health_check.failure_threshold, + success_threshold: ctx.router_config.health_check.success_threshold, + }; + + // Register prefill workers in the registry + for url in prefill_urls { + prefill_workers_urls.push(url.clone()); + let worker: Arc = Arc::new( + BasicWorker::new(url, WorkerType::Prefill) + .with_circuit_breaker_config(core_cb_config.clone()) + .with_health_config(health_config.clone()), + ); + ctx.worker_registry.register(worker); + } + + // Register decode workers in the registry + for url in decode_urls { + decode_workers_urls.push(url.clone()); + let worker: Arc = Arc::new( + BasicWorker::new(url, WorkerType::Decode) + .with_circuit_breaker_config(core_cb_config.clone()) + .with_health_config(health_config.clone()), + ); + ctx.worker_registry.register(worker); + } + + // Get all workers from registry for health check + let all_workers = ctx.worker_registry.get_all(); + let all_urls: Vec = all_workers + .iter() + .map(|worker| worker.url().to_string()) + .collect(); + // At least one prefill and one decode are up + if !prefill_workers_urls.is_empty() { + crate::routers::http::router::Router::wait_for_healthy_workers( + &prefill_workers_urls, + ctx.router_config.worker_startup_timeout_secs, + ctx.router_config.worker_startup_check_interval_secs, + ) + .await?; + } + + if !decode_workers_urls.is_empty() { + crate::routers::http::router::Router::wait_for_healthy_workers( + &decode_workers_urls, + ctx.router_config.worker_startup_timeout_secs, + ctx.router_config.worker_startup_check_interval_secs, + ) + .await?; + } + + // Initialize cache-aware policies with workers from registry + // Note: We need to get workers by type and convert to Box for CacheAwarePolicy + // This is a temporary workaround until CacheAwarePolicy is updated to work with Arc + // TODO: Update CacheAwarePolicy to accept Arc instead of Box + + // Set up background load monitoring for power-of-two selection + let (tx, rx) = tokio::sync::watch::channel(HashMap::new()); + let worker_loads = Arc::new(rx); + + // Get policies from registry to check if we need load monitoring + let prefill_policy = ctx.policy_registry.get_prefill_policy(); + let decode_policy = ctx.policy_registry.get_decode_policy(); + + let load_monitor_handle = + if prefill_policy.name() == "power_of_two" || decode_policy.name() == "power_of_two" { + let monitor_urls = all_urls.clone(); + let monitor_interval = ctx.router_config.worker_startup_check_interval_secs; + let monitor_client = ctx.client.clone(); + let prefill_policy_clone = Arc::clone(&prefill_policy); + let decode_policy_clone = Arc::clone(&decode_policy); + + Some(Arc::new(tokio::spawn(async move { + Self::monitor_worker_loads_with_client( + monitor_urls, + tx, + monitor_interval, + monitor_client, + prefill_policy_clone, + decode_policy_clone, + ) + .await; + }))) + } else { + None + }; + + // Note: Health checking is now handled centrally by RouterManager + // Individual routers no longer need to manage health checkers + + Ok(PdRouterBase { + worker_registry: Arc::clone(&ctx.worker_registry), + policy_registry: Arc::clone(&ctx.policy_registry), + worker_startup_timeout_secs: ctx.router_config.worker_startup_timeout_secs, + worker_startup_check_interval_secs: ctx + .router_config + .worker_startup_check_interval_secs, + worker_loads, + load_monitor_handle, + client: ctx.client.clone(), + circuit_breaker_config: core_cb_config, + }) + } + + // Select a pair of prefill and decode servers considering circuit breaker state + async fn select_pd_pair( + &self, + request_text: Option<&str>, + model_id: Option<&str>, + ) -> Result<(Arc, Arc), String> { + // Get workers from registry - filter by model if provided + let prefill_workers = if let Some(model) = model_id { + // Get model-specific workers and filter for prefill type + self.worker_registry + .get_by_model_fast(model) + .into_iter() + .filter(|w| matches!(w.worker_type(), WorkerType::Prefill)) + .collect() + } else { + self.worker_registry.get_prefill_workers() + }; + + let decode_workers = if let Some(model) = model_id { + // Get model-specific workers and filter for decode type + self.worker_registry + .get_by_model_fast(model) + .into_iter() + .filter(|w| matches!(w.worker_type(), WorkerType::Decode)) + .collect() + } else { + self.worker_registry.get_decode_workers() + }; + + // Select workers using helper function + // Use separate policies for prefill and decode to avoid counter conflicts + let prefill_policy = self.policy_registry.get_prefill_policy(); + let decode_policy = self.policy_registry.get_decode_policy(); + + let prefill = Self::pick_worker_by_policy_arc( + &prefill_workers, + &*prefill_policy, + request_text, + "prefill", + )?; + + let decode = Self::pick_worker_by_policy_arc( + &decode_workers, + &*decode_policy, + request_text, + "decode", + )?; + + Ok((prefill, decode)) + } + + // Helper function to select a worker using the policy (Arc version) + fn pick_worker_by_policy_arc( + workers: &[Arc], + policy: &dyn LoadBalancingPolicy, + request_text: Option<&str>, + worker_type: &str, + ) -> Result, String> { + // Check if we have any workers + if workers.is_empty() { + return Err(format!( + "No {} workers available. Please check if {} servers are configured and healthy.", + worker_type, worker_type + )); + } + + // Filter available workers (healthy + circuit breaker not open) + let available_workers: Vec> = workers + .iter() + .filter(|w| w.is_available()) + .cloned() + .collect(); + + if available_workers.is_empty() { + return Err(format!( + "No available {} workers (all circuits open or unhealthy)", + worker_type + )); + } + + // Let policy select from available workers (no conversion needed now!) + let selected_idx = policy + .select_worker(&available_workers, request_text) + .ok_or_else(|| { + format!( + "Policy {} failed to select a {} worker", + policy.name(), + worker_type + ) + })?; + + // Return the selected Arc worker + Ok(available_workers[selected_idx].clone()) + } + + // Background task to monitor worker loads with shared client + async fn monitor_worker_loads_with_client( + worker_urls: Vec, + tx: tokio::sync::watch::Sender>, + interval_secs: u64, + client: Client, + prefill_policy: Arc, + decode_policy: Arc, + ) { + loop { + let mut loads = HashMap::new(); + + let futures: Vec<_> = worker_urls + .iter() + .map(|url| { + let client = client.clone(); + let url = url.clone(); + async move { + let load = get_worker_load(&client, &url).await.unwrap_or(0); + (url, load) + } + }) + .collect(); + + let results = futures_util::future::join_all(futures).await; + + for (url, load) in results { + loads.insert(url, load); + } + + debug!("Worker loads updated: {:?}", loads); + + // Update both policies with current loads + prefill_policy.update_loads(&loads); + decode_policy.update_loads(&loads); + + // Check if receiver is still active + if tx.send(loads).is_err() { + info!("Load monitor receiver dropped, shutting down monitor task"); + break; + } + + tokio::time::sleep(Duration::from_secs(interval_secs)).await; + } + } +} + +// Helper functions + +async fn get_worker_load(client: &Client, worker_url: &str) -> Option { + match client.get(format!("{}/get_load", worker_url)).send().await { + Ok(res) if res.status().is_success() => match res.bytes().await { + Ok(bytes) => match serde_json::from_slice::(&bytes) { + Ok(data) => data + .get("load") + .and_then(|v| v.as_i64()) + .map(|v| v as isize), + Err(e) => { + debug!("Failed to parse load response from {}: {}", worker_url, e); + None + } + }, + Err(e) => { + debug!("Failed to read load response from {}: {}", worker_url, e); + None + } + }, + Ok(res) => { + debug!( + "Worker {} returned non-success status: {}", + worker_url, + res.status() + ); + None + } + Err(e) => { + debug!("Failed to get load from {}: {}", worker_url, e); + None + } + } +} + +impl PdRouterBase { + pub async fn add_worker(&self, _worker_url: &str) -> Result { + // For PD router, we don't support adding workers via this generic method + Err( + "PD router requires specific add_prefill_server or add_decode_server methods" + .to_string(), + ) + } + + pub fn remove_worker(&self, worker_url: &str) { + // Remove from registry + if let Some(worker) = self.worker_registry.remove_by_url(worker_url) { + match worker.worker_type() { + WorkerType::Prefill => { + info!("Removed prefill worker: {}", worker_url); + } + WorkerType::Decode => { + info!("Removed decode worker: {}", worker_url); + } + _ => { + info!("Removed worker: {}", worker_url); + } + } + } + } + + pub fn get_worker_urls(&self) -> Vec { + self.worker_registry.get_all_urls() + } +} + +impl PdRouterBase { + pub async fn health(&self, _req: Request) -> Response { + // This is a server readiness check - checking if we have healthy workers + // Workers handle their own health checks in the background + let mut all_healthy = true; + let mut unhealthy_servers = Vec::new(); + + // Check all workers + for worker in self.worker_registry.get_all() { + if !worker.is_healthy() { + all_healthy = false; + let worker_type = match worker.worker_type() { + WorkerType::Prefill => "Prefill", + WorkerType::Decode => "Decode", + _ => "Worker", + }; + unhealthy_servers.push(format!("{}: {}", worker_type, worker.url())); + } + } + + if all_healthy { + (StatusCode::OK, "All servers healthy").into_response() + } else { + ( + StatusCode::SERVICE_UNAVAILABLE, + format!("Unhealthy servers: {:?}", unhealthy_servers), + ) + .into_response() + } + } + + pub async fn health_generate(&self, _req: Request) -> Response { + // Test model generation capability by selecting a random pair and testing them + // Note: This endpoint actually causes the model to generate tokens, so we only test one pair + + // Select a random worker pair using the policy + let (prefill, decode) = match self.select_pd_pair(None, None).await { + Ok(pair) => pair, + Err(e) => { + return ( + StatusCode::SERVICE_UNAVAILABLE, + format!("No healthy worker pair available: {}", e), + ) + .into_response(); + } + }; + + let (prefill_result, decode_result) = tokio::join!( + self.client + .get(format!("{}/health_generate", prefill.url())) + .send(), + self.client + .get(format!("{}/health_generate", decode.url())) + .send() + ); + + // Check results + let mut errors = Vec::new(); + + match prefill_result { + Ok(res) if res.status().is_success() => { + debug!( + "Health generate passed for prefill server: {}", + prefill.url() + ); + } + Ok(res) => { + errors.push(format!( + "Prefill {} returned status {}", + prefill.url(), + res.status() + )); + } + Err(e) => { + errors.push(format!("Prefill {} error: {}", prefill.url(), e)); + } + } + + match decode_result { + Ok(res) if res.status().is_success() => { + debug!("Health generate passed for decode server: {}", decode.url()); + } + Ok(res) => { + errors.push(format!( + "Decode {} returned status {}", + decode.url(), + res.status() + )); + } + Err(e) => { + errors.push(format!("Decode {} error: {}", decode.url(), e)); + } + } + + if errors.is_empty() { + ( + StatusCode::OK, + format!( + "Health generate passed on selected pair: prefill={}, decode={}", + prefill.url(), + decode.url() + ), + ) + .into_response() + } else { + ( + StatusCode::SERVICE_UNAVAILABLE, + format!("Health generate failed: {:?}", errors), + ) + .into_response() + } + } + + pub async fn get_server_info(&self, _req: Request) -> Response { + // Get info from the first decode server + // Note: We use decode workers for server info to match expected format + self.proxy_to_first_prefill_worker("get_server_info", None) + .await + } + + pub async fn get_models(&self, req: Request) -> Response { + // Extract headers first to avoid Send issues + let headers = header_utils::copy_request_headers(&req); + + // Proxy to first prefill worker + self.proxy_to_first_prefill_worker("v1/models", Some(headers)) + .await + } + + pub async fn get_model_info(&self, req: Request) -> Response { + // Extract headers first to avoid Send issues + let headers = header_utils::copy_request_headers(&req); + + // Proxy to first prefill worker + self.proxy_to_first_prefill_worker("get_model_info", Some(headers)) + .await + } + + pub async fn get_response(&self, _headers: Option<&HeaderMap>, _response_id: &str) -> Response { + ( + StatusCode::NOT_IMPLEMENTED, + "Responses retrieve endpoint not implemented for PD router", + ) + .into_response() + } + + pub async fn cancel_response( + &self, + _headers: Option<&HeaderMap>, + _response_id: &str, + ) -> Response { + ( + StatusCode::NOT_IMPLEMENTED, + "Responses cancel endpoint not implemented for PD router", + ) + .into_response() + } + + pub async fn route_embeddings( + &self, + _headers: Option<&HeaderMap>, + _body: &crate::protocols::spec::EmbeddingRequest, + _model_id: Option<&str>, + ) -> Response { + ( + StatusCode::NOT_IMPLEMENTED, + "Embeddings endpoint not implemented for PD router", + ) + .into_response() + } + + pub async fn flush_cache(&self) -> Response { + // Process both prefill and decode workers + let (prefill_results, prefill_errors) = self + .process_workers(WorkerType::Prefill, "Prefill", "flush_cache") + .await; + let (decode_results, decode_errors) = self + .process_workers(WorkerType::Decode, "Decode", "flush_cache") + .await; + + // Combine results and errors + let mut results = prefill_results; + results.extend(decode_results); + let mut errors = prefill_errors; + errors.extend(decode_errors); + + if errors.is_empty() { + ( + StatusCode::OK, + format!("Cache flushed successfully: {:?}", results), + ) + .into_response() + } else { + ( + StatusCode::PARTIAL_CONTENT, + format!( + "Partial success. Results: {:?}, Errors: {:?}", + results, errors + ), + ) + .into_response() + } + } + + pub async fn get_worker_loads(&self) -> Response { + let mut loads = HashMap::new(); + let mut errors = Vec::new(); + + // Process prefill workers + let prefill_urls = self.get_prefill_worker_urls(); + for worker_url in prefill_urls { + match get_worker_load(&self.client, &worker_url).await { + Some(load) => { + loads.insert(format!("prefill_{}", worker_url), load); + } + None => { + errors.push(format!("Failed to get load from prefill {}", worker_url)); + } + } + } + + // Process decode workers + let decode_urls = self.get_decode_worker_urls(); + for worker_url in decode_urls { + match get_worker_load(&self.client, &worker_url).await { + Some(load) => { + loads.insert(format!("decode_{}", worker_url), load); + } + None => { + errors.push(format!("Failed to get load from decode {}", worker_url)); + } + } + } + + let response_data = serde_json::json!({ + "loads": loads, + "errors": errors + }); + + (StatusCode::OK, Json(response_data)).into_response() + } + + pub fn readiness(&self) -> Response { + // PD router is ready if it has at least one healthy prefill AND one healthy decode worker + let prefill_workers = self.worker_registry.get_prefill_workers(); + let decode_workers = self.worker_registry.get_decode_workers(); + + let healthy_prefill_count = prefill_workers.iter().filter(|w| w.is_healthy()).count(); + + let healthy_decode_count = decode_workers.iter().filter(|w| w.is_healthy()).count(); + + let total_prefill = prefill_workers.len(); + let total_decode = decode_workers.len(); + + if healthy_prefill_count > 0 && healthy_decode_count > 0 { + Json(json!({ + "status": "ready", + "prefill": { + "healthy": healthy_prefill_count, + "total": total_prefill + }, + "decode": { + "healthy": healthy_decode_count, + "total": total_decode + } + })) + .into_response() + } else { + let mut reasons = Vec::new(); + if healthy_prefill_count == 0 { + reasons.push("no healthy prefill workers"); + } + if healthy_decode_count == 0 { + reasons.push("no healthy decode workers"); + } + + ( + StatusCode::SERVICE_UNAVAILABLE, + Json(serde_json::json!({ + "status": "not_ready", + "reason": reasons.join(", "), + "prefill": { + "healthy": healthy_prefill_count, + "total": total_prefill + }, + "decode": { + "healthy": healthy_decode_count, + "total": total_decode + } + })), + ) + .into_response() + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::{BasicWorker, Worker, WorkerType}; + + fn create_test_pd_router() -> PdRouterBase { + let worker_registry = Arc::new(WorkerRegistry::new()); + let policy_registry = + Arc::new(PolicyRegistry::new(crate::config::PolicyConfig::RoundRobin)); + + PdRouterBase { + worker_registry, + policy_registry, + worker_startup_timeout_secs: 5, + worker_startup_check_interval_secs: 1, + worker_loads: Arc::new(tokio::sync::watch::channel(HashMap::new()).1), + load_monitor_handle: None, + client: Client::new(), + circuit_breaker_config: CircuitBreakerConfig::default(), + } + } + + fn create_test_worker(url: String, worker_type: WorkerType, healthy: bool) -> Box { + let worker = BasicWorker::new(url, worker_type); + worker.set_healthy(healthy); + Box::new(worker) + } + + // ============= Worker Management Tests ============= + + #[tokio::test] + async fn test_add_prefill_server_already_exists() { + let router = create_test_pd_router(); + + // Add a worker first + let worker = create_test_worker( + "http://localhost:8000".to_string(), + WorkerType::Prefill, + true, + ); + router.worker_registry.register(Arc::from(worker)); + + // Try to add the same URL again - this would fail during health check in real scenario + // For unit test, we test the duplicate check logic + let exists = router + .worker_registry + .get_by_url("http://localhost:8000") + .is_some(); + assert!(exists); + } + + #[tokio::test] + async fn test_remove_prefill_server_success() { + let router = create_test_pd_router(); + + // Add servers first + let worker1 = create_test_worker("http://worker1".to_string(), WorkerType::Prefill, true); + let worker2 = create_test_worker("http://worker2".to_string(), WorkerType::Prefill, true); + + router.worker_registry.register(Arc::from(worker1)); + router.worker_registry.register(Arc::from(worker2)); + + // Remove one + let result = router.remove_prefill_server("http://worker1").await; + + assert!(result.is_ok()); + assert!(result.unwrap().contains("Successfully removed")); + + let workers = router.worker_registry.get_prefill_workers(); + assert_eq!(workers.len(), 1); + assert_eq!(workers[0].url(), "http://worker2"); + } + + #[tokio::test] + async fn test_remove_prefill_server_not_found() { + let router = create_test_pd_router(); + + let result = router.remove_prefill_server("http://nonexistent").await; + + assert!(result.is_err()); + match result.unwrap_err() { + PDRouterError::WorkerNotFound { url } => { + assert_eq!(url, "http://nonexistent"); + } + _ => panic!("Expected WorkerNotFound error"), + } + } + + #[tokio::test] + async fn test_remove_decode_server_success() { + let router = create_test_pd_router(); + + // Add server first + let worker = create_test_worker("http://decode1".to_string(), WorkerType::Decode, true); + router.worker_registry.register(Arc::from(worker)); + + let result = router.remove_decode_server("http://decode1").await; + + assert!(result.is_ok()); + assert!(result.unwrap().contains("Successfully removed")); + + let workers = router.worker_registry.get_decode_workers(); + assert_eq!(workers.len(), 0); + } + + // ============= Lock Error Handling Tests ============= + + #[test] + fn test_registry_operations() { + let router = create_test_pd_router(); + + // Test registry operations + let workers = router.worker_registry.get_all(); + assert_eq!(workers.len(), 0); + + // Add a worker + let worker = create_test_worker("http://test".to_string(), WorkerType::Prefill, true); + router.worker_registry.register(Arc::from(worker)); + + let workers = router.worker_registry.get_all(); + assert_eq!(workers.len(), 1); + + let prefill_workers = router.worker_registry.get_prefill_workers(); + assert_eq!(prefill_workers.len(), 1); + } + + // ============= Worker Selection Tests ============= + + #[tokio::test] + async fn test_select_healthy_prefill_worker() { + let router = create_test_pd_router(); + + // Add mix of healthy and unhealthy workers + let healthy_worker = + create_test_worker("http://healthy".to_string(), WorkerType::Prefill, true); + let unhealthy_worker = + create_test_worker("http://unhealthy".to_string(), WorkerType::Prefill, false); + let decode_worker = + create_test_worker("http://decode".to_string(), WorkerType::Decode, true); + + router.worker_registry.register(Arc::from(unhealthy_worker)); + router.worker_registry.register(Arc::from(healthy_worker)); + router.worker_registry.register(Arc::from(decode_worker)); + + let result = router.select_pd_pair(None, None).await; + + assert!(result.is_ok()); + let (prefill, _decode) = result.unwrap(); + + // Should select the healthy worker + assert_eq!(prefill.url(), "http://healthy"); + assert!(prefill.is_healthy()); + } + + #[tokio::test] + async fn test_empty_worker_lists() { + let router = create_test_pd_router(); + + let result = router.select_pd_pair(None, None).await; + + assert!(result.is_err()); + assert!(result.unwrap_err().contains("No prefill workers available")); + } + + // ============= Health Endpoints Tests ============= + + #[tokio::test] + async fn test_health_endpoints() { + let router = create_test_pd_router(); + + // Add healthy workers - create_test_worker returns Box, convert to Arc + let prefill_worker = create_test_worker( + "http://localhost:8000".to_string(), + WorkerType::Prefill, + true, + ); + let decode_worker = create_test_worker( + "http://localhost:8001".to_string(), + WorkerType::Decode, + true, + ); + + router.worker_registry.register(Arc::from(prefill_worker)); + router.worker_registry.register(Arc::from(decode_worker)); + + // Test health endpoint + let http_req = axum::http::Request::builder() + .body(axum::body::Body::empty()) + .unwrap(); + let response = router.health(http_req).await; + + assert_eq!(response.status(), 200); + + // Test readiness endpoint + let response = router.readiness(); + assert_eq!(response.status(), 200); + } + + // ============= Load Monitoring Tests ============= + + #[tokio::test] + async fn test_load_monitor_updates() { + let power_of_two_policy = Arc::new(crate::policies::PowerOfTwoPolicy::new()); + let mut router = create_test_pd_router(); + // Set power_of_two policies in the registry + router + .policy_registry + .set_prefill_policy(power_of_two_policy.clone()); + router + .policy_registry + .set_decode_policy(power_of_two_policy); + + // Create load channel + let (tx, rx) = tokio::sync::watch::channel(HashMap::new()); + router.worker_loads = Arc::new(rx); + + // Simulate load updates + let mut loads = HashMap::new(); + loads.insert("http://worker1".to_string(), 10); + loads.insert("http://worker2".to_string(), 5); + + let _ = tx.send(loads.clone()); + + // Router should receive updates + let received = router.worker_loads.borrow().clone(); + assert_eq!(received.get("http://worker1"), Some(&10)); + assert_eq!(received.get("http://worker2"), Some(&5)); + } + + // ============= Concurrent Operations Tests ============= + + #[tokio::test] + async fn test_concurrent_worker_operations() { + let router = Arc::new(create_test_pd_router()); + + let mut handles = vec![]; + + // Spawn tasks to add workers + for i in 0..5 { + let router_clone = Arc::clone(&router); + let url = format!("http://worker{}", i); + let handle = tokio::spawn(async move { + let worker = create_test_worker(url, WorkerType::Prefill, true); + router_clone.worker_registry.register(Arc::from(worker)); + }); + handles.push(handle); + } + + // Wait for all tasks + for handle in handles { + let _ = handle.await; + } + + // Check final state + let workers = router.worker_registry.get_prefill_workers(); + assert_eq!(workers.len(), 5); + } +} diff --git a/lmdeploy/router/src/routers/http/pd_types.rs b/lmdeploy/router/src/routers/http/pd_types.rs new file mode 100644 index 0000000000..d14c3f54b5 --- /dev/null +++ b/lmdeploy/router/src/routers/http/pd_types.rs @@ -0,0 +1,110 @@ +// Custom error type for PD router operations +#[derive(Debug, thiserror::Error)] +pub enum PDRouterError { + #[error("Worker already exists: {url}")] + WorkerAlreadyExists { url: String }, + + #[error("Worker not found: {url}")] + WorkerNotFound { url: String }, + + #[error("Lock acquisition failed: {operation}")] + LockError { operation: String }, + + #[error("Health check failed for worker: {url}")] + HealthCheckFailed { url: String }, + + #[error("Invalid worker configuration: {reason}")] + InvalidConfiguration { reason: String }, + + #[error("Network error: {message}")] + NetworkError { message: String }, + + #[error("Upstream server {url} returned status {status}")] + UpstreamResponse { + url: String, + status: reqwest::StatusCode, + body: Vec, + }, + + #[error("Timeout waiting for worker: {url}")] + Timeout { url: String }, +} + +/// Format a full error chain for debugging (walks source() recursively). +/// Produces output like: "outer error caused by: middle error caused by: root cause" +pub fn error_chain(err: &dyn std::error::Error) -> String { + let mut chain = vec![err.to_string()]; + let mut source = err.source(); + while let Some(s) = source { + chain.push(s.to_string()); + source = s.source(); + } + chain.join(" caused by: ") +} + +#[cfg(test)] +mod tests { + use super::*; + use std::fmt; + + // Simple custom error for testing error chains + #[derive(Debug)] + struct TestError { + msg: String, + source: Option>, + } + + impl fmt::Display for TestError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{}", self.msg) + } + } + + impl std::error::Error for TestError { + fn source(&self) -> Option<&(dyn std::error::Error + 'static)> { + self.source.as_deref() + } + } + + #[test] + fn test_error_chain_single_error() { + let err = TestError { + msg: "something broke".into(), + source: None, + }; + assert_eq!(error_chain(&err), "something broke"); + } + + #[test] + fn test_error_chain_nested_errors() { + let inner = TestError { + msg: "root cause".into(), + source: None, + }; + let outer = TestError { + msg: "outer error".into(), + source: Some(Box::new(inner)), + }; + assert_eq!(error_chain(&outer), "outer error caused by: root cause"); + } + + #[test] + fn test_error_chain_triple_nested() { + let root = TestError { + msg: "connection reset".into(), + source: None, + }; + let middle = TestError { + msg: "HTTP send failed".into(), + source: Some(Box::new(root)), + }; + let top = TestError { + msg: "prefill request failed".into(), + source: Some(Box::new(middle)), + }; + assert_eq!( + error_chain(&top), + "prefill request failed caused by: HTTP send failed caused by: connection reset" + ); + } +} diff --git a/lmdeploy/router/src/routers/http/router.rs b/lmdeploy/router/src/routers/http/router.rs new file mode 100644 index 0000000000..4e33e5a459 --- /dev/null +++ b/lmdeploy/router/src/routers/http/router.rs @@ -0,0 +1,1893 @@ +use crate::config::types::RetryConfig; +use crate::core::{ + is_retryable_status, BasicWorker, CircuitBreakerConfig, HealthConfig, RetryExecutor, Worker, + WorkerRegistry, WorkerType, +}; +use crate::metrics::RouterMetrics; +use crate::otel_http::{self, ClientRequestOptions}; +use crate::policies::{LoadBalancingPolicy, PolicyRegistry}; +use crate::protocols::spec::{ + ChatCompletionRequest, CompletionRequest, EmbeddingRequest, GenerateRequest, GenerationRequest, + ResponsesRequest, +}; +use crate::routers::header_utils; +use crate::routers::{RouterTrait, WorkerManagement}; +use axum::{ + body::Body, + extract::Request, + http::{ + header::CONTENT_LENGTH, header::CONTENT_TYPE, HeaderMap, HeaderValue, Method, StatusCode, + }, + response::{IntoResponse, Response}, + Json, +}; +use futures_util::StreamExt; +use reqwest::Client; +use std::collections::HashMap; +use std::sync::Arc; +use std::time::{Duration, Instant}; +use tokio_stream::wrappers::UnboundedReceiverStream; +use tracing::{debug, error, info, warn}; + +/// Regular router that uses injected load balancing policies +#[derive(Debug)] +pub struct Router { + worker_registry: Arc, + policy_registry: Arc, + client: Client, + worker_startup_timeout_secs: u64, + worker_startup_check_interval_secs: u64, + api_key: Option, + retry_config: RetryConfig, + circuit_breaker_config: CircuitBreakerConfig, + _worker_loads: Arc>>, + _load_monitor_handle: Option>>, +} + +impl Router { + /// Create a new router with injected policy and client + #[allow(clippy::too_many_arguments)] + pub async fn new( + worker_urls: Vec, + ctx: &Arc, + ) -> Result { + // Update active workers gauge + RouterMetrics::set_active_workers(worker_urls.len()); + + // Wait for workers to be healthy (skip if empty - for service discovery mode) + if !worker_urls.is_empty() { + Self::wait_for_healthy_workers( + &worker_urls, + ctx.router_config.worker_startup_timeout_secs, + ctx.router_config.worker_startup_check_interval_secs, + ) + .await?; + } + + // Convert config CircuitBreakerConfig to core CircuitBreakerConfig + let circuit_breaker_config = ctx.router_config.effective_circuit_breaker_config(); + let core_cb_config = CircuitBreakerConfig { + failure_threshold: circuit_breaker_config.failure_threshold, + success_threshold: circuit_breaker_config.success_threshold, + timeout_duration: Duration::from_secs(circuit_breaker_config.timeout_duration_secs), + window_duration: Duration::from_secs(circuit_breaker_config.window_duration_secs), + }; + + // Register workers in the registry + // In IGW mode, we need to fetch model info from workers + let health_config = HealthConfig { + timeout_secs: ctx.router_config.health_check.timeout_secs, + check_interval_secs: ctx.router_config.health_check.check_interval_secs, + endpoint: ctx.router_config.health_check.endpoint.clone(), + failure_threshold: ctx.router_config.health_check.failure_threshold, + success_threshold: ctx.router_config.health_check.success_threshold, + }; + for url in &worker_urls { + // TODO: In IGW mode, fetch model_id from worker's /get_model_info endpoint + // For now, create worker without model_id + let worker_arc: Arc = Arc::new( + BasicWorker::new(url.clone(), WorkerType::Regular) + .with_circuit_breaker_config(core_cb_config.clone()) + .with_health_config(health_config.clone()), + ); + ctx.worker_registry.register(worker_arc.clone()); + + // Notify PolicyRegistry about the new worker + let model_id = worker_arc.model_id(); + let policy = ctx.policy_registry.on_worker_added(model_id, None); + + // If this is a cache-aware policy and it's the first worker for this model, + // initialize it with the worker + if policy.name() == "cache_aware" { + if let Some(cache_aware) = policy + .as_any() + .downcast_ref::() + { + let worker_dyn: Arc = worker_arc.clone(); + cache_aware.init_workers(std::slice::from_ref(&worker_dyn)); + } + } + } + + // Setup load monitoring for PowerOfTwo policy + let (tx, rx) = tokio::sync::watch::channel(HashMap::new()); + let worker_loads = Arc::new(rx); + + // Check if default policy is power_of_two for load monitoring + let default_policy = ctx.policy_registry.get_default_policy(); + let load_monitor_handle = if default_policy.name() == "power_of_two" { + let monitor_urls = worker_urls.clone(); + let monitor_interval = ctx.router_config.worker_startup_check_interval_secs; + let policy_clone = default_policy.clone(); + let client_clone = ctx.client.clone(); + + Some(Arc::new(tokio::spawn(async move { + Self::monitor_worker_loads( + monitor_urls, + tx, + monitor_interval, + policy_clone, + client_clone, + ) + .await; + }))) + } else { + None + }; + + Ok(Router { + worker_registry: ctx.worker_registry.clone(), + policy_registry: ctx.policy_registry.clone(), + client: ctx.client.clone(), + worker_startup_timeout_secs: ctx.router_config.worker_startup_timeout_secs, + worker_startup_check_interval_secs: ctx + .router_config + .worker_startup_check_interval_secs, + api_key: ctx.router_config.api_key.clone(), + retry_config: ctx.router_config.effective_retry_config(), + circuit_breaker_config: core_cb_config, + _worker_loads: worker_loads, + _load_monitor_handle: load_monitor_handle, + }) + } + + /// Get the current list of worker URLs + pub fn get_worker_urls(&self) -> Vec { + self.worker_registry.get_all_urls() + } + + /// Get worker URLs for a specific model + pub fn get_worker_urls_for_model(&self, model_id: Option<&str>) -> Vec { + let workers = match model_id { + Some(model) => self.worker_registry.get_by_model_fast(model), + None => self.worker_registry.get_all(), + }; + workers.iter().map(|w| w.url().to_string()).collect() + } + + pub async fn wait_for_healthy_workers( + worker_urls: &[String], + worker_startup_timeout_secs: u64, + worker_startup_check_interval_secs: u64, + ) -> Result<(), String> { + if worker_urls.is_empty() { + return Err( + "Timeout waiting for workers to become healthy: no workers provided".to_string(), + ); + } + + // Perform health check asynchronously + Self::wait_for_healthy_workers_async( + worker_urls, + worker_startup_timeout_secs, + worker_startup_check_interval_secs, + ) + .await + } + + async fn wait_for_healthy_workers_async( + worker_urls: &[String], + worker_startup_timeout_secs: u64, + worker_startup_check_interval_secs: u64, + ) -> Result<(), String> { + let unique_hosts_vec: Vec = worker_urls.to_vec(); + + info!( + "Waiting for {} workers to become healthy (timeout: {}s)", + unique_hosts_vec.len(), + worker_startup_timeout_secs + ); + + let start_time = std::time::Instant::now(); + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(2)) + .build() + .map_err(|e| format!("Failed to create HTTP client: {}", e))?; + + loop { + if start_time.elapsed() > Duration::from_secs(worker_startup_timeout_secs) { + error!( + "Timeout {}s waiting for hosts {:?} to become healthy. Set --worker-startup-timeout-secs to a larger value", + worker_startup_timeout_secs, unique_hosts_vec + ); + return Err(format!( + "Timeout {}s waiting for hosts {:?} to become healthy. Set --worker-startup-timeout-secs to a larger value", + worker_startup_timeout_secs, unique_hosts_vec + )); + } + + let mut health_checks = Vec::new(); + for worker_url in &unique_hosts_vec { + let client_clone = client.clone(); + let url_clone = worker_url.clone(); + + let check_health = tokio::spawn(async move { + let health_url = format!("{}/health", url_clone); + match client_clone.get(&health_url).send().await { + Ok(res) => { + if res.status().is_success() { + None + } else { + Some((url_clone, format!("status: {}", res.status()))) + } + } + Err(_) => Some((url_clone, "not ready".to_string())), + } + }); + + health_checks.push(check_health); + } + + // Wait for all health checks to complete + let results = futures::future::join_all(health_checks).await; + + let mut unhealthy_hosts = Vec::new(); + let mut healthy_host_count = 0; + + for result in results { + match result { + Ok(None) => { + healthy_host_count += 1; + // Host is healthy + } + Ok(Some((url, reason))) => { + unhealthy_hosts.push((url, reason)); + } + Err(e) => { + unhealthy_hosts.push(("unknown".to_string(), format!("task error: {}", e))); + } + } + } + + if healthy_host_count > 0 { + info!( + "{} out of {} workers are healthy", + healthy_host_count, + unique_hosts_vec.len() + ); + return Ok(()); + } else { + debug!( + "Waiting for at least 1 of {} workers to become healthy ({} unhealthy: {:?})", + unique_hosts_vec.len(), + unhealthy_hosts.len(), + unhealthy_hosts + ); + tokio::time::sleep(Duration::from_secs(worker_startup_check_interval_secs)).await; + } + } + } + + fn select_first_worker(&self) -> Result { + let workers = self.worker_registry.get_all(); + if workers.is_empty() { + Err("No workers are available".to_string()) + } else { + Ok(workers[0].url().to_string()) + } + } + + #[allow(dead_code)] + fn select_first_worker_for_model(&self, model_id: Option<&str>) -> Result { + let workers = match model_id { + Some(model) => self.worker_registry.get_by_model_fast(model), + None => self.worker_registry.get_all(), + }; + if workers.is_empty() { + Err(format!( + "No workers are available for model: {:?}", + model_id + )) + } else { + Ok(workers[0].url().to_string()) + } + } + + pub async fn send_health_check(&self, worker_url: &str) -> Response { + let request_builder = self.client.get(format!("{}/health", worker_url)); + + let response = match request_builder.send().await { + Ok(res) => { + let status = StatusCode::from_u16(res.status().as_u16()) + .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); + + match res.bytes().await { + Ok(body) => (status, body).into_response(), + Err(e) => { + error!( + worker_url = %worker_url, + error = %e, + "Failed to read health response body" + ); + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to read response body: {}", e), + ) + .into_response() + } + } + } + Err(e) => { + error!( + worker_url = %worker_url, + error = %e, + "Failed to send health request to worker" + ); + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to send request to worker {}: {}", worker_url, e), + ) + .into_response() + } + }; + + // Don't record metrics for health checks + response + } + + // Helper method to proxy GET requests to the first available worker + async fn proxy_get_request(&self, req: Request, endpoint: &str) -> Response { + let incoming_headers = req.headers(); + let headers = header_utils::copy_request_headers(&req); + + match self.select_first_worker() { + Ok(worker_url) => { + let url = format!("{}/{}", worker_url, endpoint); + let route_name = format!("/{}", endpoint); + let mut request_builder = self.client.get(&url); + for (name, value) in headers { + let name_lc = name.to_lowercase(); + if name_lc != "content-type" + && name_lc != "content-length" + && !header_utils::TRACE_HEADER_NAMES.contains(&name_lc.as_str()) + { + request_builder = request_builder.header(name, value); + } + } + + match otel_http::send_client_request( + request_builder, + Some(incoming_headers), + ClientRequestOptions { + method: "GET", + url: &url, + route: Some(&route_name), + request_phase: None, + }, + ) + .await + { + Ok(res) => { + let status = StatusCode::from_u16(res.status().as_u16()) + .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); + + // Preserve headers from backend + let response_headers = + header_utils::preserve_response_headers(res.headers()); + + match res.bytes().await { + Ok(body) => { + let mut response = Response::new(axum::body::Body::from(body)); + *response.status_mut() = status; + *response.headers_mut() = response_headers; + response + } + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to read response: {}", e), + ) + .into_response(), + } + } + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Request failed: {}", e), + ) + .into_response(), + } + } + Err(e) => (StatusCode::SERVICE_UNAVAILABLE, e).into_response(), + } + } + + /// Convert axum HeaderMap to policy RequestHeaders (HashMap) + fn headers_to_request_headers( + headers: Option<&HeaderMap>, + ) -> Option { + headers.map(|h| { + h.iter() + .filter_map(|(name, value)| { + value + .to_str() + .ok() + .map(|v| (name.as_str().to_lowercase(), v.to_string())) + }) + .collect() + }) + } + + /// Select worker for a specific model considering circuit breaker state + fn select_worker_for_model( + &self, + model_id: Option<&str>, + text: Option<&str>, + headers: Option<&HeaderMap>, + ) -> Option> { + // Get workers for the specified model (O(1) lookup if model_id is provided) + let workers = match model_id { + Some(model) => self.worker_registry.get_by_model_fast(model), + None => self.worker_registry.get_all(), + }; + + let available: Vec> = workers + .iter() + .filter(|w| w.is_available()) + .cloned() + .collect(); + if available.is_empty() { + return None; + } + + // Get the appropriate policy for this model + let policy = match model_id { + Some(model) => self.policy_registry.get_policy_or_default(model), + None => self.policy_registry.get_default_policy(), + }; + + // Convert headers for policies that need them (e.g., consistent_hash) + let request_headers = Self::headers_to_request_headers(headers); + + let idx = policy.select_worker_with_headers(&available, text, request_headers.as_ref())?; + Some(available[idx].clone()) + } + + pub async fn route_typed_request( + &self, + headers: Option<&HeaderMap>, + typed_req: &T, + route: &str, + model_id: Option<&str>, + ) -> Response { + let start = Instant::now(); + let is_stream = typed_req.is_stream(); + let text = typed_req.extract_text_for_routing(); + + let response = RetryExecutor::execute_response_with_retry( + &self.retry_config, + // operation per attempt + |_: u32| async { + let worker = match self.select_worker_for_model(model_id, Some(&text), headers) { + Some(w) => w, + None => { + RouterMetrics::record_request_error(route, "no_available_workers"); + return ( + StatusCode::SERVICE_UNAVAILABLE, + "No available workers (all circuits open or unhealthy)", + ) + .into_response(); + } + }; + + // Optional load tracking for cache-aware policy + // Get the policy for this model to check if it's cache-aware + let policy = match model_id { + Some(model) => self.policy_registry.get_policy_or_default(model), + None => self.policy_registry.get_default_policy(), + }; + + let load_incremented = if policy.name() == "cache_aware" { + worker.increment_load(); + RouterMetrics::set_running_requests(worker.url(), worker.load()); + true + } else { + false + }; + + // Keep a clone for potential cleanup on retry + let worker_for_cleanup = if load_incremented { + Some(worker.clone()) + } else { + None + }; + + let response = self + .send_typed_request( + headers, + typed_req, + route, + worker.url(), + is_stream, + load_incremented, + ) + .await; + + // Client errors (4xx) are not worker failures - only server errors (5xx) + // should count against the circuit breaker. + let status = response.status(); + worker.record_outcome(status.is_success() || status.is_client_error()); + + // For retryable failures, we need to decrement load since send_typed_request + // won't have done it (it only decrements on success or non-retryable failures) + if is_retryable_status(response.status()) && load_incremented { + if let Some(cleanup_worker) = worker_for_cleanup { + cleanup_worker.decrement_load(); + RouterMetrics::set_running_requests( + cleanup_worker.url(), + cleanup_worker.load(), + ); + } + } + + response + }, + // should_retry predicate + |res, _attempt| is_retryable_status(res.status()), + // on_backoff hook + |delay, attempt| { + RouterMetrics::record_retry(route); + RouterMetrics::record_retry_backoff_duration(delay, attempt); + }, + // on_exhausted hook + || RouterMetrics::record_retries_exhausted(route), + ) + .await; + + if response.status().is_success() { + let duration = start.elapsed(); + RouterMetrics::record_request(route); + RouterMetrics::record_generate_duration(duration); + } else if !is_retryable_status(response.status()) { + RouterMetrics::record_request_error(route, "non_retryable_error"); + } + + response + } + + // Generic simple routing for GET/POST without JSON body + async fn route_simple_request( + &self, + headers: Option<&HeaderMap>, + endpoint: &str, + method: Method, + ) -> Response { + // TODO: currently the backend is using in-memory state management, so this implementation has to fan out to all workers. + // Eventually, we need to have router to manage the chat history with a proper database, will update this implementation accordingly. + let worker_urls = self.get_worker_urls(); + if worker_urls.is_empty() { + return (StatusCode::SERVICE_UNAVAILABLE, "No available workers").into_response(); + } + + let mut last_response: Option = None; + for worker_url in worker_urls { + let url = format!("{}/{}", worker_url, endpoint); + + let route_name = format!("/{}", endpoint); + let method_name = method.as_str().to_string(); + let mut request_builder = match method.clone() { + Method::GET => self.client.get(&url), + Method::POST => self.client.post(&url), + _ => { + return ( + StatusCode::METHOD_NOT_ALLOWED, + "Unsupported method for simple routing", + ) + .into_response() + } + }; + + if let Some(hdrs) = headers { + for (name, value) in hdrs { + let name_lc = name.as_str().to_lowercase(); + if name_lc != "content-type" + && name_lc != "content-length" + && !header_utils::TRACE_HEADER_NAMES.contains(&name_lc.as_str()) + { + request_builder = request_builder.header(name, value); + } + } + } + + match otel_http::send_client_request( + request_builder, + headers, + ClientRequestOptions { + method: &method_name, + url: &url, + route: Some(&route_name), + request_phase: None, + }, + ) + .await + { + Ok(res) => { + let status = StatusCode::from_u16(res.status().as_u16()) + .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); + let response_headers = header_utils::preserve_response_headers(res.headers()); + match res.bytes().await { + Ok(body) => { + let mut response = Response::new(axum::body::Body::from(body)); + *response.status_mut() = status; + *response.headers_mut() = response_headers; + if status.is_success() { + return response; + } + last_response = Some(response); + } + Err(e) => { + last_response = Some( + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to read response: {}", e), + ) + .into_response(), + ); + } + } + } + Err(e) => { + last_response = Some( + ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Request failed: {}", e), + ) + .into_response(), + ); + } + } + } + + last_response + .unwrap_or_else(|| (StatusCode::BAD_GATEWAY, "No worker response").into_response()) + } + + // Route a GET request with provided headers to a specific endpoint + async fn route_get_request(&self, headers: Option<&HeaderMap>, endpoint: &str) -> Response { + self.route_simple_request(headers, endpoint, Method::GET) + .await + } + + // Route a POST request with empty body to a specific endpoint + async fn route_post_empty_request( + &self, + headers: Option<&HeaderMap>, + endpoint: &str, + ) -> Response { + self.route_simple_request(headers, endpoint, Method::POST) + .await + } + + // Send typed request directly without conversion + async fn send_typed_request( + &self, + headers: Option<&HeaderMap>, + typed_req: &T, + route: &str, + worker_url: &str, + is_stream: bool, + load_incremented: bool, // Whether load was incremented for this request + ) -> Response { + let request_url = format!("{}{}", worker_url, route); + let mut request_builder = self.client.post(&request_url).json(typed_req); + + // Copy all headers from original request if provided, skipping + // Content-Type/Content-Length (.json() sets them) and trace headers + // (propagate_trace_headers below injects fresh context). + if let Some(headers) = headers { + for (name, value) in headers { + if *name != CONTENT_TYPE + && *name != CONTENT_LENGTH + && !header_utils::TRACE_HEADER_NAMES + .iter() + .any(|&th| name.as_str().eq_ignore_ascii_case(th)) + { + request_builder = request_builder.header(name, value); + } + } + } + + let res = match otel_http::send_client_request( + request_builder, + headers, + ClientRequestOptions { + method: "POST", + url: &request_url, + route: Some(route), + request_phase: Some("inference"), + }, + ) + .await + { + Ok(res) => res, + Err(e) => { + error!( + "Failed to send typed request worker_url={} route={} error={}", + worker_url, route, e + ); + + // Decrement load on error if it was incremented + if load_incremented { + if let Some(worker) = self.worker_registry.get_by_url(worker_url) { + worker.decrement_load(); + RouterMetrics::set_running_requests(worker_url, worker.load()); + } + } + + return ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Request failed: {}", e), + ) + .into_response(); + } + }; + + let status = StatusCode::from_u16(res.status().as_u16()) + .unwrap_or(StatusCode::INTERNAL_SERVER_ERROR); + + if !is_stream { + // For non-streaming requests, preserve headers + let response_headers = header_utils::preserve_response_headers(res.headers()); + + let response = match res.bytes().await { + Ok(body) => { + let mut response = Response::new(axum::body::Body::from(body)); + *response.status_mut() = status; + *response.headers_mut() = response_headers; + response + } + Err(e) => { + // IMPORTANT: Decrement load on error before returning + if load_incremented { + if let Some(worker) = self.worker_registry.get_by_url(worker_url) { + worker.decrement_load(); + RouterMetrics::set_running_requests(worker_url, worker.load()); + } + } + + let error_msg = format!("Failed to get response body: {}", e); + (StatusCode::INTERNAL_SERVER_ERROR, error_msg).into_response() + } + }; + + // Decrement load counter for non-streaming requests if it was incremented + if load_incremented { + if let Some(worker) = self.worker_registry.get_by_url(worker_url) { + worker.decrement_load(); + RouterMetrics::set_running_requests(worker_url, worker.load()); + } + } + + response + } else if load_incremented { + // For streaming with load tracking, we need to manually decrement when done + let registry = Arc::clone(&self.worker_registry); + let worker_url = worker_url.to_string(); + + // Preserve headers for streaming response + let mut response_headers = header_utils::preserve_response_headers(res.headers()); + // Ensure we set the correct content-type for SSE + response_headers.insert(CONTENT_TYPE, HeaderValue::from_static("text/event-stream")); + + let stream = res.bytes_stream(); + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + + // Spawn task to forward stream and detect completion + tokio::spawn(async move { + let mut stream = stream; + let mut decremented = false; + while let Some(chunk) = stream.next().await { + match chunk { + Ok(bytes) => { + // Check for stream end marker + if bytes + .as_ref() + .windows(12) + .any(|window| window == b"data: [DONE]") + { + if let Some(worker) = registry.get_by_url(&worker_url) { + worker.decrement_load(); + RouterMetrics::set_running_requests(&worker_url, worker.load()); + decremented = true; + } + } + if tx.send(Ok(bytes)).is_err() { + break; + } + } + Err(e) => { + let _ = tx.send(Err(format!("Stream error: {}", e))); + break; + } + } + } + if !decremented { + if let Some(worker) = registry.get_by_url(&worker_url) { + worker.decrement_load(); + RouterMetrics::set_running_requests(&worker_url, worker.load()); + } + } + }); + + let stream = UnboundedReceiverStream::new(rx); + let body = Body::from_stream(stream); + + let mut response = Response::new(body); + *response.status_mut() = status; + *response.headers_mut() = response_headers; + response + } else { + // For requests without load tracking, just stream + // Preserve headers for streaming response + let mut response_headers = header_utils::preserve_response_headers(res.headers()); + // Ensure we set the correct content-type for SSE + response_headers.insert(CONTENT_TYPE, HeaderValue::from_static("text/event-stream")); + + let stream = res.bytes_stream(); + let (tx, rx) = tokio::sync::mpsc::unbounded_channel(); + + // Spawn task to forward stream + tokio::spawn(async move { + let mut stream = stream; + while let Some(chunk) = stream.next().await { + match chunk { + Ok(bytes) => { + if tx.send(Ok(bytes)).is_err() { + break; + } + } + Err(e) => { + let _ = tx.send(Err(format!("Stream error: {}", e))); + break; + } + } + } + }); + + let stream = UnboundedReceiverStream::new(rx); + let body = Body::from_stream(stream); + + let mut response = Response::new(body); + *response.status_mut() = status; + *response.headers_mut() = response_headers; + response + } + } + + pub async fn add_worker(&self, worker_url: &str) -> Result { + let start_time = std::time::Instant::now(); + let client = reqwest::Client::builder() + .timeout(Duration::from_secs(self.worker_startup_timeout_secs)) + .build() + .map_err(|e| format!("Failed to create HTTP client: {}", e))?; + + loop { + if start_time.elapsed() > Duration::from_secs(self.worker_startup_timeout_secs) { + error!( + "Timeout {}s waiting for worker {} to become healthy. Set --worker-startup-timeout-secs to a larger value", + self.worker_startup_timeout_secs, worker_url + ); + return Err(format!( + "Timeout {}s waiting for worker {} to become healthy. Set --worker-startup-timeout-secs to a larger value", + self.worker_startup_timeout_secs, worker_url + )); + } + + match client.get(format!("{}/health", worker_url)).send().await { + Ok(res) => { + if res.status().is_success() { + return self.register_worker_unchecked(worker_url); + } else { + debug!( + "Worker {} health check pending - status: {}", + worker_url, + res.status() + ); + // if the url does not have http or https prefix, warn users + if !worker_url.starts_with("http://") && !worker_url.starts_with("https://") + { + warn!("The worker url {} does not have http or https prefix. Please add the prefix to the url.", worker_url); + } + + tokio::time::sleep(Duration::from_secs( + self.worker_startup_check_interval_secs, + )) + .await; + continue; + } + } + Err(e) => { + debug!("Worker {} health check pending - error: {}", worker_url, e); + + // if the url does not have http or https prefix, warn users + if !worker_url.starts_with("http://") && !worker_url.starts_with("https://") { + warn!("The worker url {} does not have http or https prefix. Please add the prefix to the url.", worker_url); + } + + tokio::time::sleep(Duration::from_secs( + self.worker_startup_check_interval_secs, + )) + .await; + continue; + } + } + } + } + + /// Register a worker without waiting for its HTTP server to become healthy. + /// + /// LMDeploy registers from FastAPI's startup callback, before uvicorn starts + /// accepting requests. Performing a synchronous health check from + /// `/nodes/add` would therefore deadlock registration until the router's + /// startup timeout expires. The central health checker takes over after the + /// worker has been registered. + pub fn register_worker_unchecked(&self, worker_url: &str) -> Result { + if self.worker_registry.get_by_url(worker_url).is_some() { + return Err(format!("Worker {} already exists", worker_url)); + } + + let worker_arc: Arc = Arc::new( + BasicWorker::new(worker_url.to_string(), WorkerType::Regular) + .with_circuit_breaker_config(self.circuit_breaker_config.clone()), + ); + self.worker_registry.register(worker_arc.clone()); + let model_id = worker_arc.model_id(); + let policy = self.policy_registry.on_worker_added(model_id, None); + if policy.name() == "cache_aware" { + if let Some(cache_aware) = policy + .as_any() + .downcast_ref::() + { + let model_workers = self.worker_registry.get_by_model_fast(model_id); + cache_aware.init_workers(&model_workers); + } + } + info!( + "Registered worker without startup health check: {}", + worker_url + ); + + RouterMetrics::set_active_workers(self.worker_registry.get_all().len()); + Ok(format!("Successfully added worker: {}", worker_url)) + } + + pub fn remove_worker(&self, worker_url: &str) { + // Get the worker first to extract model_id + let model_id = if let Some(worker) = self.worker_registry.get_by_url(worker_url) { + worker.model_id().to_string() + } else { + warn!("Worker {} not found, skipping removal", worker_url); + return; + }; + + if self.worker_registry.remove_by_url(worker_url).is_some() { + info!("Removed worker: {}", worker_url); + + // Notify PolicyRegistry about the removed worker + self.policy_registry.on_worker_removed(&model_id); + + RouterMetrics::set_active_workers(self.worker_registry.get_all().len()); + } + + // If the model is using cache aware policy, remove the worker from the tree + if let Some(policy) = self.policy_registry.get_policy(&model_id) { + if let Some(cache_aware) = policy + .as_any() + .downcast_ref::() + { + cache_aware.remove_worker_by_url(worker_url); + info!("Removed worker from cache-aware tree: {}", worker_url); + } + } + } + + async fn get_worker_load(&self, worker_url: &str) -> Option { + match self + .client + .get(format!("{}/get_load", worker_url)) + .send() + .await + { + Ok(res) if res.status().is_success() => match res.bytes().await { + Ok(bytes) => match serde_json::from_slice::(&bytes) { + Ok(data) => data + .get("load") + .and_then(|v| v.as_i64()) + .map(|v| v as isize), + Err(e) => { + debug!("Failed to parse load response from {}: {}", worker_url, e); + None + } + }, + Err(e) => { + debug!("Failed to read load response from {}: {}", worker_url, e); + None + } + }, + Ok(res) => { + debug!( + "Worker {} returned non-success status: {}", + worker_url, + res.status() + ); + None + } + Err(e) => { + debug!("Failed to get load from {}: {}", worker_url, e); + None + } + } + } + + // Background task to monitor worker loads + async fn monitor_worker_loads( + worker_urls: Vec, + tx: tokio::sync::watch::Sender>, + interval_secs: u64, + policy: Arc, + client: Client, + ) { + let mut interval = tokio::time::interval(Duration::from_secs(interval_secs)); + + loop { + interval.tick().await; + + let mut loads = HashMap::new(); + for url in &worker_urls { + if let Some(load) = Self::get_worker_load_static(&client, url).await { + loads.insert(url.clone(), load); + } + } + + if !loads.is_empty() { + // Update policy with new loads + policy.update_loads(&loads); + + // Send to watchers + if let Err(e) = tx.send(loads) { + error!("Failed to send load update: {}", e); + } + } + } + } + + // Static version of get_worker_load for use in monitoring task + async fn get_worker_load_static(client: &reqwest::Client, worker_url: &str) -> Option { + match client.get(format!("{}/get_load", worker_url)).send().await { + Ok(res) if res.status().is_success() => match res.bytes().await { + Ok(bytes) => match serde_json::from_slice::(&bytes) { + Ok(data) => data + .get("load") + .and_then(|v| v.as_i64()) + .map(|v| v as isize), + Err(e) => { + debug!("Failed to parse load response from {}: {}", worker_url, e); + None + } + }, + Err(e) => { + debug!("Failed to read load response from {}: {}", worker_url, e); + None + } + }, + Ok(res) => { + debug!( + "Worker {} returned non-success status: {}", + worker_url, + res.status() + ); + None + } + Err(e) => { + debug!("Failed to get load from {}: {}", worker_url, e); + None + } + } + } +} + +use async_trait::async_trait; + +#[async_trait] +impl WorkerManagement for Router { + async fn add_worker(&self, worker_url: &str) -> Result { + Router::add_worker(self, worker_url).await + } + + fn remove_worker(&self, worker_url: &str) { + Router::remove_worker(self, worker_url) + } + + fn get_worker_urls(&self) -> Vec { + Router::get_worker_urls(self) + } +} + +#[async_trait] +impl RouterTrait for Router { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + async fn health(&self, _req: Request) -> Response { + let workers = self.worker_registry.get_all(); + let unhealthy_servers: Vec<_> = workers + .iter() + .filter(|w| !w.is_healthy()) + .map(|w| w.url().to_string()) + .collect(); + + if unhealthy_servers.is_empty() { + (StatusCode::OK, "All servers healthy").into_response() + } else { + ( + StatusCode::SERVICE_UNAVAILABLE, + format!("Unhealthy servers: {:?}", unhealthy_servers), + ) + .into_response() + } + } + + async fn health_generate(&self, req: Request) -> Response { + self.proxy_get_request(req, "health_generate").await + } + + async fn get_server_info(&self, req: Request) -> Response { + self.proxy_get_request(req, "get_server_info").await + } + + async fn get_models(&self, req: Request) -> Response { + self.proxy_get_request(req, "v1/models").await + } + + async fn get_model_info(&self, req: Request) -> Response { + self.proxy_get_request(req, "get_model_info").await + } + + async fn route_generate( + &self, + headers: Option<&HeaderMap>, + body: &GenerateRequest, + model_id: Option<&str>, + ) -> Response { + self.route_typed_request(headers, body, "/generate", model_id) + .await + } + + async fn route_chat( + &self, + headers: Option<&HeaderMap>, + body: &ChatCompletionRequest, + model_id: Option<&str>, + ) -> Response { + self.route_typed_request(headers, body, "/v1/chat/completions", model_id) + .await + } + + async fn route_completion( + &self, + headers: Option<&HeaderMap>, + body: &CompletionRequest, + model_id: Option<&str>, + ) -> Response { + self.route_typed_request(headers, body, "/v1/completions", model_id) + .await + } + + async fn route_responses( + &self, + headers: Option<&HeaderMap>, + body: &ResponsesRequest, + model_id: Option<&str>, + ) -> Response { + self.route_typed_request(headers, body, "/v1/responses", model_id) + .await + } + + async fn get_response(&self, headers: Option<&HeaderMap>, response_id: &str) -> Response { + let endpoint = format!("v1/responses/{}", response_id); + self.route_get_request(headers, &endpoint).await + } + + async fn cancel_response(&self, headers: Option<&HeaderMap>, response_id: &str) -> Response { + let endpoint = format!("v1/responses/{}/cancel", response_id); + self.route_post_empty_request(headers, &endpoint).await + } + + async fn route_embeddings( + &self, + headers: Option<&HeaderMap>, + body: &EmbeddingRequest, + model_id: Option<&str>, + ) -> Response { + // Record embeddings-specific metrics in addition to general request metrics + let start = Instant::now(); + let res = self + .route_typed_request(headers, body, "/v1/embeddings", model_id) + .await; + + // Embedding specific metrics + if res.status().is_success() { + RouterMetrics::record_embeddings_request(); + RouterMetrics::record_embeddings_duration(start.elapsed()); + } else { + let error_type = format!("http_{}", res.status().as_u16()); + RouterMetrics::record_embeddings_error(&error_type); + } + + res + } + + async fn flush_cache(&self) -> Response { + // Get all worker URLs + let worker_urls = self.get_worker_urls(); + + // Send requests to all workers concurrently without headers + let mut tasks = Vec::new(); + for worker_url in &worker_urls { + let request_builder = self.client.post(format!("{}/flush_cache", worker_url)); + tasks.push(request_builder.send()); + } + + // Wait for all responses + let results = futures_util::future::join_all(tasks).await; + + // Check if all succeeded + let all_success = results.iter().all(|r| { + r.as_ref() + .map(|res| res.status().is_success()) + .unwrap_or(false) + }); + + if all_success { + (StatusCode::OK, "Cache flushed on all servers").into_response() + } else { + ( + StatusCode::INTERNAL_SERVER_ERROR, + "Cache flush failed on one or more servers", + ) + .into_response() + } + } + + async fn get_worker_loads(&self) -> Response { + let urls = self.get_worker_urls(); + let mut loads = Vec::new(); + + // Get loads from all workers + for url in &urls { + let load = self.get_worker_load(url).await.unwrap_or(-1); + loads.push(serde_json::json!({ + "worker": url, + "load": load + })); + } + + Json(serde_json::json!({ + "workers": loads + })) + .into_response() + } + + fn router_type(&self) -> &'static str { + "regular" + } + + fn readiness(&self) -> Response { + // Regular router is ready if it has at least one healthy worker + let workers = self.worker_registry.get_all(); + let healthy_count = workers.iter().filter(|w| w.is_healthy()).count(); + let total_workers = workers.len(); + + if healthy_count > 0 { + Json(serde_json::json!({ + "status": "ready", + "healthy_workers": healthy_count, + "total_workers": total_workers + })) + .into_response() + } else { + ( + StatusCode::SERVICE_UNAVAILABLE, + Json(serde_json::json!({ + "status": "not_ready", + "reason": "no healthy workers available", + "total_workers": total_workers + })), + ) + .into_response() + } + } + + /// Route a transparent proxy request to a backend worker + /// Forwards the request as-is to a selected worker + async fn route_transparent( + &self, + headers: Option<&HeaderMap>, + path: &str, + method: &Method, + body: serde_json::Value, + ) -> Response { + debug!("Transparent proxy: routing {} {} to backend", method, path); + + // Select a worker (filter by availability like select_worker_for_model) + let all_workers = self.worker_registry.get_all(); + let workers: Vec> = all_workers + .iter() + .filter(|w| w.is_available()) + .cloned() + .collect(); + if workers.is_empty() { + return ( + StatusCode::SERVICE_UNAVAILABLE, + "No available workers".to_string(), + ) + .into_response(); + } + + let policy = self.policy_registry.get_default_policy(); + let request_text = serde_json::to_string(&body).ok(); + let request_headers = Self::headers_to_request_headers(headers); + let worker_idx = match policy.select_worker_with_headers( + &workers, + request_text.as_deref(), + request_headers.as_ref(), + ) { + Some(idx) => idx, + None => { + return ( + StatusCode::SERVICE_UNAVAILABLE, + "Failed to select a worker".to_string(), + ) + .into_response(); + } + }; + + let worker: &dyn Worker = workers[worker_idx].as_ref(); + let url = worker.endpoint_url(path); + + debug!("Transparent proxy: forwarding to {}", url); + + // Build the request + let mut request_builder = match *method { + Method::GET => self.client.get(&url), + Method::POST => self.client.post(&url), + Method::PUT => self.client.put(&url), + Method::DELETE => self.client.delete(&url), + Method::PATCH => self.client.patch(&url), + Method::HEAD => self.client.head(&url), + _ => { + return ( + StatusCode::METHOD_NOT_ALLOWED, + format!("Method {} not supported", method), + ) + .into_response(); + } + }; + + // Add JSON body if not null/empty + if !body.is_null() { + request_builder = request_builder.json(&body); + } + + // Add authorization if configured + if let Some(ref key) = self.api_key { + request_builder = request_builder.header("Authorization", format!("Bearer {}", key)); + } + + // Forward original request headers (e.g. anthropic-version, x-api-key) + // so protocol-specific backends receive the metadata they require. + // Skip content-type/content-length (set by .json(&body)), host and + // hop-by-hop headers, and trace headers (handled by send_client_request). + if let Some(hdrs) = headers { + for (name, value) in hdrs { + let name_lc = name.as_str().to_lowercase(); + if name_lc != "content-type" + && name_lc != "content-length" + && name_lc != "host" + && name_lc != "connection" + && name_lc != "keep-alive" + && name_lc != "proxy-authenticate" + && name_lc != "proxy-authorization" + && name_lc != "te" + && name_lc != "trailers" + && name_lc != "transfer-encoding" + && name_lc != "upgrade" + && !header_utils::TRACE_HEADER_NAMES.contains(&name_lc.as_str()) + { + request_builder = request_builder.header(name, value); + } + } + } + + // Send request + match otel_http::send_client_request( + request_builder, + headers, + ClientRequestOptions { + method: method.as_str(), + url: &url, + route: Some(path), + request_phase: Some("inference"), + }, + ) + .await + { + Ok(response) => { + let status = response.status(); + let headers = response.headers().clone(); + + // Stream the response body + let body = Body::from_stream(response.bytes_stream()); + let mut response_builder = Response::builder().status(status.as_u16()); + + for (name, value) in headers.iter() { + if name != "transfer-encoding" && name != "content-length" { + response_builder = response_builder.header(name, value); + } + } + + match response_builder.body(body) { + Ok(response) => response, + Err(e) => ( + StatusCode::INTERNAL_SERVER_ERROR, + format!("Failed to build response: {}", e), + ) + .into_response(), + } + } + Err(e) => ( + StatusCode::BAD_GATEWAY, + format!("Backend request failed: {}", e), + ) + .into_response(), + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashMap; + + fn create_test_regular_router() -> Router { + // Create registries + let worker_registry = Arc::new(WorkerRegistry::new()); + let policy_registry = Arc::new(PolicyRegistry::new( + crate::config::types::PolicyConfig::RoundRobin, + )); + + // Register test workers + let worker1 = BasicWorker::new("http://worker1:8080".to_string(), WorkerType::Regular); + let worker2 = BasicWorker::new("http://worker2:8080".to_string(), WorkerType::Regular); + worker_registry.register(Arc::new(worker1)); + worker_registry.register(Arc::new(worker2)); + + let (_, rx) = tokio::sync::watch::channel(HashMap::new()); + Router { + worker_registry, + policy_registry, + worker_startup_timeout_secs: 5, + worker_startup_check_interval_secs: 1, + api_key: None, + client: Client::new(), + retry_config: RetryConfig::default(), + circuit_breaker_config: CircuitBreakerConfig::default(), + _worker_loads: Arc::new(rx), + _load_monitor_handle: None, + } + } + + #[test] + fn test_router_get_worker_urls_regular() { + let router = create_test_regular_router(); + let urls = router.get_worker_urls(); + + assert_eq!(urls.len(), 2); + assert!(urls.contains(&"http://worker1:8080".to_string())); + assert!(urls.contains(&"http://worker2:8080".to_string())); + } + + #[test] + fn test_select_first_worker_regular() { + let router = create_test_regular_router(); + let result = router.select_first_worker(); + + assert!(result.is_ok()); + let url = result.unwrap(); + // DashMap doesn't guarantee order, so just check we get one of the workers + assert!(url == "http://worker1:8080" || url == "http://worker2:8080"); + } + + #[tokio::test] + async fn test_wait_for_healthy_workers_empty_list() { + // Empty list will return error immediately + let result = Router::wait_for_healthy_workers(&[], 1, 1).await; + assert!(result.is_err()); + assert!(result.unwrap_err().contains("no workers provided")); + } + + #[tokio::test] + async fn test_wait_for_healthy_workers_invalid_urls() { + // This test will timeout quickly since the URLs are invalid + let result = + Router::wait_for_healthy_workers(&["http://nonexistent:8080".to_string()], 1, 1).await; + assert!(result.is_err()); + assert!(result.unwrap_err().contains("Timeout")); + } + + // ============================= + // Tests for transparent proxy header/availability fixes + // ============================= + + /// Create a test router with ConsistentHash policy instead of RoundRobin + fn create_test_consistent_hash_router() -> Router { + let worker_registry = Arc::new(WorkerRegistry::new()); + let policy_registry = Arc::new(PolicyRegistry::new( + crate::config::types::PolicyConfig::ConsistentHash { virtual_nodes: 100 }, + )); + + let worker1 = BasicWorker::new("http://worker1:8080".to_string(), WorkerType::Regular); + let worker2 = BasicWorker::new("http://worker2:8080".to_string(), WorkerType::Regular); + let worker3 = BasicWorker::new("http://worker3:8080".to_string(), WorkerType::Regular); + worker_registry.register(Arc::new(worker1)); + worker_registry.register(Arc::new(worker2)); + worker_registry.register(Arc::new(worker3)); + + let (_, rx) = tokio::sync::watch::channel(HashMap::new()); + Router { + worker_registry, + policy_registry, + worker_startup_timeout_secs: 5, + worker_startup_check_interval_secs: 1, + api_key: None, + client: Client::new(), + retry_config: RetryConfig::default(), + circuit_breaker_config: CircuitBreakerConfig::default(), + _worker_loads: Arc::new(rx), + _load_monitor_handle: None, + } + } + + #[test] + fn test_headers_to_request_headers_basic() { + // Test that headers_to_request_headers correctly converts HeaderMap to HashMap + let mut header_map = HeaderMap::new(); + header_map.insert("x-session-id", HeaderValue::from_static("session-123")); + header_map.insert("content-type", HeaderValue::from_static("application/json")); + header_map.insert("X-Custom-Header", HeaderValue::from_static("custom-value")); + + let result = Router::headers_to_request_headers(Some(&header_map)); + assert!(result.is_some()); + let headers = result.unwrap(); + + // All keys should be lowercased + assert_eq!(headers.get("x-session-id").unwrap(), "session-123"); + assert_eq!(headers.get("content-type").unwrap(), "application/json"); + assert_eq!(headers.get("x-custom-header").unwrap(), "custom-value"); + } + + #[test] + fn test_headers_to_request_headers_none() { + // Test that None headers produce None output + let result = Router::headers_to_request_headers(None); + assert!(result.is_none()); + } + + #[test] + fn test_headers_to_request_headers_empty() { + // Test that empty HeaderMap produces empty HashMap + let header_map = HeaderMap::new(); + let result = Router::headers_to_request_headers(Some(&header_map)); + assert!(result.is_some()); + assert!(result.unwrap().is_empty()); + } + + #[test] + fn test_select_worker_for_model_with_consistent_hash_uses_headers() { + // Verify that select_worker_for_model passes headers through to the policy, + // producing consistent routing for the same session ID + let router = create_test_consistent_hash_router(); + + let mut header_map = HeaderMap::new(); + header_map.insert("x-session-id", HeaderValue::from_static("sticky-session-1")); + + // Make multiple selections with the same headers - should all pick the same worker + let mut selected_urls: Vec = Vec::new(); + for _ in 0..10 { + let worker = router + .select_worker_for_model(None, Some(r#"{"prompt": "test"}"#), Some(&header_map)) + .expect("Should select a worker"); + selected_urls.push(worker.url().to_string()); + } + + // All selections should go to the same worker (sticky routing) + let first = &selected_urls[0]; + for (i, url) in selected_urls.iter().enumerate() { + assert_eq!( + url, first, + "Request {} routed to {}, expected {} (session stickiness broken)", + i, url, first + ); + } + } + + #[test] + fn test_select_worker_for_model_filters_unavailable_workers() { + // Verify that select_worker_for_model skips unhealthy workers + let router = create_test_consistent_hash_router(); + + // Mark worker1 and worker2 as unhealthy, leaving only worker3 + let all_workers = router.worker_registry.get_all(); + for w in &all_workers { + if w.url() == "http://worker1:8080" || w.url() == "http://worker2:8080" { + w.set_healthy(false); + } + } + + let worker = router + .select_worker_for_model(None, Some(r#"{"prompt": "test"}"#), None) + .expect("Should select the remaining healthy worker"); + + assert_eq!( + worker.url(), + "http://worker3:8080", + "Should only select the healthy worker" + ); + } + + #[test] + fn test_select_worker_for_model_returns_none_when_all_unavailable() { + // Verify that when all workers are unhealthy, None is returned + let router = create_test_consistent_hash_router(); + + // Mark all workers as unhealthy + let all_workers = router.worker_registry.get_all(); + for w in &all_workers { + w.set_healthy(false); + } + + let result = router.select_worker_for_model(None, Some(r#"{"prompt": "test"}"#), None); + assert!( + result.is_none(), + "Should return None when all workers are unavailable" + ); + } + + #[test] + fn test_consistent_hash_different_sessions_can_route_differently() { + // Verify that different session IDs can route to different workers + let router = create_test_consistent_hash_router(); + + let mut worker_urls_seen = std::collections::HashSet::new(); + for i in 0..50 { + let mut header_map = HeaderMap::new(); + let session_id = format!("session-{}", i); + header_map.insert("x-session-id", HeaderValue::from_str(&session_id).unwrap()); + + if let Some(worker) = router.select_worker_for_model( + None, + Some(r#"{"prompt": "test"}"#), + Some(&header_map), + ) { + worker_urls_seen.insert(worker.url().to_string()); + } + } + + // With 50 different sessions and 3 workers, we should see at least 2 workers used + assert!( + worker_urls_seen.len() >= 2, + "Expected distribution across workers, only used: {:?}", + worker_urls_seen + ); + } + + #[test] + fn test_inline_header_conversion_matches_headers_to_request_headers() { + // Verify that the inline header conversion pattern used by transparent + // proxy routes matches Router::headers_to_request_headers. + let mut header_map = HeaderMap::new(); + header_map.insert("X-Session-Id", HeaderValue::from_static("session-abc")); + header_map.insert("Content-Type", HeaderValue::from_static("application/json")); + header_map.insert("x-user-id", HeaderValue::from_static("user-42")); + + // Method 1: Router::headers_to_request_headers (used in router.rs) + let method1 = Router::headers_to_request_headers(Some(&header_map)).unwrap(); + + // Method 2: Inline conversion used by transparent proxy routes. + let method2: HashMap = header_map + .iter() + .filter_map(|(name, value)| { + value + .to_str() + .ok() + .map(|v| (name.as_str().to_lowercase(), v.to_string())) + }) + .collect(); + + assert_eq!( + method1, method2, + "Both header conversion methods should produce identical results" + ); + } + + /// Helper: start a minimal mock server that responds 200 on /health. + async fn start_healthy_mock_server() -> (String, tokio::task::JoinHandle<()>) { + use axum::{routing::get, Router as AxumRouter}; + use tokio::net::TcpListener; + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + let app = AxumRouter::new().route("/health", get(|| async { "ok" })); + let handle = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + (format!("http://{}", addr), handle) + } + + #[tokio::test] + async fn test_wait_for_healthy_workers_all_healthy() { + let (url, _handle) = start_healthy_mock_server().await; + let result = Router::wait_for_healthy_workers(&[url], 5, 1).await; + assert!(result.is_ok()); + } + + #[tokio::test] + async fn test_wait_for_healthy_workers_partial_health() { + // One healthy server + one unreachable URL. + // The new behaviour succeeds when at least one host is healthy. + let (healthy_url, _handle) = start_healthy_mock_server().await; + let unreachable_url = "http://127.0.0.1:1".to_string(); // port 1 is unreachable + + let result = Router::wait_for_healthy_workers(&[healthy_url, unreachable_url], 5, 1).await; + assert!(result.is_ok()); + } + + /// Helper: start a mock server that returns 503 on /health for a given + /// duration, then switches to 200. Simulates a worker with a slow startup. + async fn start_delayed_healthy_mock_server( + delay: std::time::Duration, + ) -> (String, tokio::task::JoinHandle<()>) { + use axum::{extract::State, http::StatusCode, routing::get, Router as AxumRouter}; + use std::sync::Arc; + use tokio::net::TcpListener; + + let start = std::time::Instant::now(); + let ready_after = Arc::new(delay); + + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + let app = AxumRouter::new() + .route( + "/health", + get( + move |State((start, ready_after)): State<( + std::time::Instant, + Arc, + )>| async move { + if start.elapsed() >= *ready_after { + StatusCode::OK + } else { + StatusCode::SERVICE_UNAVAILABLE + } + }, + ), + ) + .with_state((start, ready_after)); + + let handle = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + tokio::time::sleep(std::time::Duration::from_millis(10)).await; + (format!("http://{}", addr), handle) + } + + #[tokio::test] + async fn test_delayed_worker_becomes_routable_via_health_checker() { + use crate::core::{BasicWorker, HealthConfig, WorkerRegistry, WorkerType}; + use std::sync::Arc; + + // Two workers: one immediately healthy, one delayed (503 for 2s, then 200). + let (healthy_url, _h1) = start_healthy_mock_server().await; + let (delayed_url, _h2) = + start_delayed_healthy_mock_server(std::time::Duration::from_secs(2)).await; + + // Verify the delayed worker is genuinely unhealthy right now. + let client = reqwest::Client::new(); + let resp = client + .get(format!("{}/health", delayed_url)) + .send() + .await + .unwrap(); + assert_eq!(resp.status().as_u16(), 503); + + // ── Step 1: wait_for_healthy_workers (mirrors PdRouterBase::new startup) ── + // This succeeds because the healthy worker responds immediately, + // even though the delayed worker is still returning 503. + let result = + Router::wait_for_healthy_workers(&[healthy_url.clone(), delayed_url.clone()], 10, 1) + .await; + assert!( + result.is_ok(), + "Startup should succeed with at least one healthy worker" + ); + + // ── Step 2: register workers in the registry (mirrors PdRouterBase::new) ── + let registry = Arc::new(WorkerRegistry::new()); + + let healthy_worker = Arc::new( + BasicWorker::new(healthy_url, WorkerType::Decode).with_health_config(HealthConfig { + timeout_secs: 2, + check_interval_secs: 1, + endpoint: "/health".to_string(), + failure_threshold: 3, + success_threshold: 1, + }), + ); + registry.register(healthy_worker); + + let delayed_worker = Arc::new( + BasicWorker::new(delayed_url, WorkerType::Decode).with_health_config(HealthConfig { + timeout_secs: 2, + check_interval_secs: 1, + endpoint: "/health".to_string(), + failure_threshold: 3, + success_threshold: 1, + }), + ); + delayed_worker.set_healthy(false); // starts unhealthy + registry.register(delayed_worker.clone()); + + // Only the immediately-healthy worker should be available for routing. + let healthy = registry.get_workers_filtered(None, None, None, true); + assert_eq!( + healthy.len(), + 1, + "Only 1 worker should be healthy initially, got {}", + healthy.len() + ); + + // ── Step 3: start background health checker (mirrors PdRouterBase::new) ── + let health_checker = registry.start_health_checker(1); + + // ── Step 4: wait for delayed worker to recover ── + // The mock server switches to 200 at t≈2s. With a 1s check interval + // and success_threshold=1, the worker should be healthy by t≈3-4s. + tokio::time::sleep(std::time::Duration::from_secs(4)).await; + + // Both workers should now be available for routing. + let healthy = registry.get_workers_filtered(None, None, None, true); + assert_eq!( + healthy.len(), + 2, + "Both workers should be healthy after recovery, got {}", + healthy.len() + ); + assert!( + delayed_worker.is_healthy(), + "Delayed worker should have transitioned to healthy via health checker" + ); + + health_checker.shutdown().await; + } +} diff --git a/lmdeploy/router/src/routers/mod.rs b/lmdeploy/router/src/routers/mod.rs new file mode 100644 index 0000000000..01e32a8bea --- /dev/null +++ b/lmdeploy/router/src/routers/mod.rs @@ -0,0 +1,169 @@ +//! Router implementations + +use async_trait::async_trait; +use axum::{ + body::Body, + extract::Request, + http::{HeaderMap, Method, StatusCode}, + response::{IntoResponse, Response}, +}; +use std::fmt::Debug; + +use crate::protocols::spec::{ + ChatCompletionRequest, CompletionRequest, EmbeddingRequest, GenerateRequest, ResponsesRequest, +}; + +pub mod factory; +pub mod header_utils; +pub mod http; +pub mod router_manager; + +pub use factory::RouterFactory; +// Re-export HTTP routers for convenience (keeps routers::openai_router path working) +pub use http::{openai_router, pd_router, pd_types, router}; + +/// Worker management trait for administrative operations +/// +/// This trait is separate from RouterTrait to allow Send futures +/// for use in service discovery and other background tasks +#[async_trait] +pub trait WorkerManagement: Send + Sync { + /// Add a worker to the router + async fn add_worker(&self, worker_url: &str) -> Result; + + /// Remove a worker from the router + fn remove_worker(&self, worker_url: &str); + + /// Get all worker URLs + fn get_worker_urls(&self) -> Vec; +} + +/// Core trait for all router implementations +/// +/// This trait provides a unified interface for routing requests, +/// regardless of whether it's a regular router or PD router. +#[async_trait] +pub trait RouterTrait: Send + Sync + Debug + WorkerManagement { + /// Get a reference to self as Any for downcasting + fn as_any(&self) -> &dyn std::any::Any; + + /// Route a health check request + async fn health(&self, req: Request) -> Response; + + /// Route a health generate request + async fn health_generate(&self, req: Request) -> Response; + + /// Get server information + async fn get_server_info(&self, req: Request) -> Response; + + /// Get available models + async fn get_models(&self, req: Request) -> Response; + + /// Get model information + async fn get_model_info(&self, req: Request) -> Response; + + /// Route a generate request + async fn route_generate( + &self, + headers: Option<&HeaderMap>, + body: &GenerateRequest, + model_id: Option<&str>, + ) -> Response; + + /// Route a chat completion request + async fn route_chat( + &self, + headers: Option<&HeaderMap>, + body: &ChatCompletionRequest, + model_id: Option<&str>, + ) -> Response; + + /// Route a completion request + async fn route_completion( + &self, + headers: Option<&HeaderMap>, + body: &CompletionRequest, + model_id: Option<&str>, + ) -> Response; + + /// Route a responses request + async fn route_responses( + &self, + headers: Option<&HeaderMap>, + body: &ResponsesRequest, + model_id: Option<&str>, + ) -> Response; + + /// Retrieve a stored/background response by id + async fn get_response(&self, headers: Option<&HeaderMap>, response_id: &str) -> Response; + + /// Cancel a background response by id + async fn cancel_response(&self, headers: Option<&HeaderMap>, response_id: &str) -> Response; + + /// Delete a response by id + async fn delete_response(&self, _headers: Option<&HeaderMap>, _response_id: &str) -> Response { + ( + StatusCode::NOT_IMPLEMENTED, + "Responses delete endpoint not implemented", + ) + .into_response() + } + + /// List input items of a response by id + async fn list_response_input_items( + &self, + _headers: Option<&HeaderMap>, + _response_id: &str, + ) -> Response { + ( + StatusCode::NOT_IMPLEMENTED, + "Responses list input items endpoint not implemented", + ) + .into_response() + } + + /// Route embedding requests (OpenAI-compatible /v1/embeddings) + async fn route_embeddings( + &self, + headers: Option<&HeaderMap>, + body: &EmbeddingRequest, + model_id: Option<&str>, + ) -> Response; + + /// Flush cache on all workers + async fn flush_cache(&self) -> Response; + + /// Get worker loads (for monitoring) + async fn get_worker_loads(&self) -> Response; + + /// Get router type name + fn router_type(&self) -> &'static str; + + /// Check if this is a PD router + fn is_pd_mode(&self) -> bool { + self.router_type() == "pd" + } + + /// Server liveness check - is the server process running + fn liveness(&self) -> Response { + // Simple liveness check - if we can respond, we're alive + (StatusCode::OK, "OK").into_response() + } + + /// Server readiness check - is the server ready to handle requests + fn readiness(&self) -> Response; + + /// Route a transparent proxy request (any path/body) + /// Used for catch-all routing of unmatched paths + /// Returns the response from the backend or an error response + async fn route_transparent( + &self, + _headers: Option<&HeaderMap>, + _path: &str, + _method: &Method, + _body: serde_json::Value, + ) -> Response { + // Default: not supported - return 404 + (StatusCode::NOT_FOUND, "Not Found").into_response() + } +} diff --git a/lmdeploy/router/src/routers/router_manager.rs b/lmdeploy/router/src/routers/router_manager.rs new file mode 100644 index 0000000000..a8a574a948 --- /dev/null +++ b/lmdeploy/router/src/routers/router_manager.rs @@ -0,0 +1,790 @@ +//! Router Manager for coordinating multiple routers and workers +//! +//! Provides centralized management based on enable_igw flag: +//! - Single Router Mode (enable_igw=false): Router owns workers directly +//! - Multi-Router Mode (enable_igw=true): RouterManager coordinates everything + +use crate::config::RouterConfig; +use crate::core::{CircuitBreakerConfig, Worker, WorkerFactory, WorkerRegistry, WorkerType}; +use crate::protocols::spec::{ + ChatCompletionRequest, CompletionRequest, EmbeddingRequest, GenerateRequest, ResponsesRequest, +}; +use crate::protocols::worker_spec::{ + ServerInfo, WorkerApiResponse, WorkerConfigRequest, WorkerErrorResponse, WorkerInfo, + WorkerListResponse, WorkerStats, WorkerTypeStats, +}; +use crate::routers::{RouterTrait, WorkerManagement}; +use async_trait::async_trait; +use axum::{ + body::Body, + extract::Request, + http::{HeaderMap, Method, StatusCode}, + response::{IntoResponse, Response}, +}; +use dashmap::DashMap; +use std::sync::Arc; +use tracing::{info, warn}; + +/// Router identifier +#[derive(Debug, Clone, Hash, Eq, PartialEq)] +pub struct RouterId(String); + +impl RouterId { + pub fn new(id: String) -> Self { + Self(id) + } + + pub fn as_str(&self) -> &str { + &self.0 + } +} + +/// Router Manager - Central coordinator for routers and workers +pub struct RouterManager { + /// Worker registry (single source of truth in multi-router mode) + worker_registry: Arc, + + /// Policy registry for managing model-to-policy mappings + policy_registry: Arc, + + /// All routers managed by this manager + /// RouterId examples: "http-regular", "http-pd" + routers: Arc>>, + + /// Default router for requests without specific routing + default_router: Arc>>, + + /// HTTP client for querying worker info + client: reqwest::Client, + + /// Configuration + #[allow(dead_code)] // May be used in future enhancements + config: RouterConfig, +} + +impl RouterManager { + /// Create a new router manager with shared registries + pub fn new( + config: RouterConfig, + client: reqwest::Client, + worker_registry: Arc, + policy_registry: Arc, + ) -> Self { + Self { + worker_registry, + policy_registry, + routers: Arc::new(DashMap::new()), + default_router: Arc::new(std::sync::RwLock::new(None)), + client, + config, + } + } + + /// Register a router with the manager + pub fn register_router(&self, id: RouterId, router: Arc) { + // Store router + self.routers.insert(id.clone(), router); + + // Set as default if first router + let mut default_router = self.default_router.write().unwrap(); + if default_router.is_none() { + *default_router = Some(id.clone()); + info!("Set default router to {}", id.as_str()); + } + } + + /// Set the default router + pub fn set_default_router(&self, id: RouterId) { + let mut default_router = self.default_router.write().unwrap(); + *default_router = Some(id); + } + + /// Get the number of registered routers + pub fn router_count(&self) -> usize { + self.routers.len() + } + + /// Get router for a specific model based on worker types + pub fn get_router_for_model(&self, model_id: &str) -> Option> { + // Query workers for this model from registry + let workers = self.worker_registry.get_by_model(model_id); + + if !workers.is_empty() { + // Determine router based on worker types + let has_pd_workers = workers + .iter() + .any(|w| matches!(w.worker_type(), WorkerType::Prefill | WorkerType::Decode)); + + let router_id = if has_pd_workers { + RouterId::new("http-pd".to_string()) + } else { + RouterId::new("http-regular".to_string()) + }; + + // Return the router if it exists + if let Some(router) = self.routers.get(&router_id) { + return Some(router.clone()); + } + } + + // Fall back to default router + let default_router = self.default_router.read().unwrap(); + if let Some(ref default_id) = *default_router { + self.routers.get(default_id).map(|r| r.clone()) + } else { + None + } + } + + /// Get workers for routing decision + pub fn get_workers_for_request(&self, model_id: Option<&str>) -> Vec> { + if let Some(model) = model_id { + self.worker_registry.get_by_model(model) + } else { + self.worker_registry.get_all() + } + } + + /// Add a worker to the registry + pub async fn add_worker( + &self, + config: WorkerConfigRequest, + ) -> Result { + // Build labels from configuration + let mut labels = config.labels.clone(); + + // Query server info if model_id not provided + let model_id = if let Some(model_id) = config.model_id { + model_id + } else { + match self.query_server_info(&config.url).await { + Ok(info) => { + // Extract model_id from server info + info.model_id + .or_else(|| { + info.model_path + .as_ref() + .and_then(|path| path.split('/').next_back().map(|s| s.to_string())) + }) + .unwrap_or_else(|| "unknown".to_string()) + } + Err(e) => { + warn!("Failed to query server info from {}: {}", config.url, e); + "unknown".to_string() + } + } + }; + + // Add configuration to labels + labels.insert("model_id".to_string(), model_id.clone()); + + if let Some(priority) = config.priority { + labels.insert("priority".to_string(), priority.to_string()); + } + + if let Some(cost) = config.cost { + labels.insert("cost".to_string(), cost.to_string()); + } + + let worker = match config.worker_type.as_deref() { + Some("prefill") => WorkerFactory::create_prefill_with_labels( + config.url.clone(), + labels.clone(), + CircuitBreakerConfig::default(), + ), + Some("decode") => WorkerFactory::create_decode_with_labels( + config.url.clone(), + labels.clone(), + CircuitBreakerConfig::default(), + ), + _ => WorkerFactory::create_regular_with_labels( + config.url.clone(), + labels.clone(), + CircuitBreakerConfig::default(), + ), + }; + + // Register worker + let worker_arc: Arc = Arc::from(worker); + let worker_id = self.worker_registry.register(worker_arc.clone()); + + // Notify PolicyRegistry about the new worker + // Extract policy hint from labels if provided + let policy_hint = labels.get("policy").map(|s| s.as_str()); + let policy = self.policy_registry.on_worker_added(&model_id, policy_hint); + + // Log which type of router would handle this worker (for debugging) + let expected_router = match config.worker_type.as_deref() { + Some("prefill") | Some("decode") => "http-pd", + _ => "http-regular", + }; + + info!( + "Worker for model '{}' would be handled by '{}' router based on type", + model_id, expected_router + ); + + info!( + "Added worker {} with URL {} for model {} using policy {}", + worker_id.as_str(), + config.url, + model_id, + policy.name() + ); + + // Return worker info + let worker_info = self.worker_to_info(worker_id.as_str(), &worker_arc); + + Ok(WorkerApiResponse { + success: true, + message: format!("Worker {} added successfully", worker_id.as_str()), + worker: Some(worker_info), + }) + } + + /// Remove a worker from the registry + pub fn remove_worker_from_registry( + &self, + url: &str, + ) -> Result { + // Get worker to extract model_id before removing + let model_id = self + .worker_registry + .get_by_url(url) + .map(|worker| worker.model_id().to_string()); + + if let Some(_worker) = self.worker_registry.remove_by_url(url) { + // Notify PolicyRegistry about worker removal + if let Some(ref model_id) = model_id { + self.policy_registry.on_worker_removed(model_id); + + info!("Removed worker with URL {} for model {}", url, model_id); + } else { + info!("Removed worker with URL {}", url); + } + + Ok(WorkerApiResponse { + success: true, + message: format!("Worker {} removed successfully", url), + worker: None, + }) + } else { + Err(WorkerErrorResponse { + error: format!("Worker with URL {} not found", url), + code: "WORKER_NOT_FOUND".to_string(), + }) + } + } + + /// List all workers + pub fn list_workers(&self) -> WorkerListResponse { + let workers = self.worker_registry.get_all_with_ids(); + let worker_infos: Vec = workers + .iter() + .map(|(id, w)| self.worker_to_info(id.as_str(), w)) + .collect(); + + let total = worker_infos.len(); + + // Get stats from the worker registry + let registry_stats = self.worker_registry.stats(); + + // Convert WorkerRegistryStats to WorkerStats + let stats = WorkerStats { + total_workers: registry_stats.total_workers, + healthy_workers: registry_stats.healthy_workers, + total_models: registry_stats.total_models, + total_load: registry_stats.total_load, + by_type: WorkerTypeStats { + regular: registry_stats.regular_workers, + prefill: registry_stats.prefill_workers, + decode: registry_stats.decode_workers, + }, + }; + + WorkerListResponse { + workers: worker_infos, + total, + stats, + } + } + + /// Get worker by URL + pub fn get_worker(&self, url: &str) -> Option { + self.worker_registry + .get_by_url(url) + .map(|w| self.worker_to_info("unknown", &w)) + } + + /// Query server info from a worker URL + async fn query_server_info(&self, url: &str) -> Result { + let info_url = format!("{}/get_server_info", url.trim_end_matches('/')); + + match self.client.get(&info_url).send().await { + Ok(response) => { + if response.status().is_success() { + response + .json::() + .await + .map_err(|e| format!("Failed to parse server info: {}", e)) + } else { + Err(format!("Server returned status: {}", response.status())) + } + } + Err(e) => Err(format!("Failed to connect to server: {}", e)), + } + } + + /// Convert Worker to WorkerInfo + fn worker_to_info(&self, id: &str, worker: &Arc) -> WorkerInfo { + let metadata = worker.metadata(); + + WorkerInfo { + id: id.to_string(), + url: worker.url().to_string(), + model_id: worker.model_id().to_string(), + priority: worker.priority(), + cost: worker.cost(), + worker_type: match worker.worker_type() { + WorkerType::Regular => "regular".to_string(), + WorkerType::Prefill => "prefill".to_string(), + WorkerType::Decode => "decode".to_string(), + }, + is_healthy: worker.is_healthy(), + load: worker.load(), + connection_mode: format!("{:?}", worker.connection_mode()), + metadata: metadata.labels.clone(), + } + } + + /// Get the appropriate router for a request based on headers and request content + pub fn select_router_for_request( + &self, + headers: Option<&HeaderMap>, + model_id: Option<&str>, + ) -> Option> { + // Extract priority and cost preferences from headers if available + let _priority_threshold = headers.and_then(|h| { + h.get("x-worker-priority") + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.parse::().ok()) + }); + + let _max_cost = headers.and_then(|h| { + h.get("x-max-cost") + .and_then(|v| v.to_str().ok()) + .and_then(|s| s.parse::().ok()) + }); + + // Check if PD (prefill-decode) mode is preferred from headers + let prefer_pd = headers + .and_then(|h| { + h.get("x-prefer-pd") + .and_then(|v| v.to_str().ok()) + .map(|s| s == "true" || s == "1") + }) + .unwrap_or(false); + + // If model specified, use get_router_for_model + let candidate_routers = if let Some(model) = model_id { + if let Some(router) = self.get_router_for_model(model) { + vec![router] + } else { + Vec::new() + } + } else { + // No model specified, consider all routers + self.routers + .iter() + .map(|entry| entry.value().clone()) + .collect::>() + }; + + if candidate_routers.is_empty() { + // No routers found for the specified model + return None; + } + + // Score routers based on worker attributes and request preferences + let mut best_router = None; + let mut best_score = 0.0; + + for router in candidate_routers { + let mut score = 1.0; + + // Check if this is a PD router + let is_pd = router.is_pd_mode(); + if prefer_pd && is_pd { + score += 2.0; // Bonus for matching PD preference + } else if !prefer_pd && !is_pd { + score += 1.0; // Bonus for matching regular preference + } + + // Get workers for this router and evaluate based on priority/cost + // Note: This would require routers to expose their workers or stats + // For now, we'll use a simple selection based on router type + + // TODO: Once routers expose worker stats, we can evaluate: + // - Average worker priority vs priority_threshold + // - Average worker cost vs max_cost + // - Current load and health status + + if score > best_score { + best_score = score; + best_router = Some(router); + } + } + + best_router + } +} + +/// RouterManager implements RouterTrait to act as a meta-router +/// that delegates requests to the appropriate underlying router +#[async_trait] +impl WorkerManagement for RouterManager { + /// Add a worker - in multi-router mode, this adds to the registry + async fn add_worker(&self, worker_url: &str) -> Result { + // Create a basic worker config request + let config = WorkerConfigRequest { + url: worker_url.to_string(), + model_id: None, + worker_type: None, + priority: None, + cost: None, + labels: std::collections::HashMap::new(), + }; + + match self.add_worker(config).await { + Ok(response) => Ok(response.message), + Err(e) => Err(e.error), + } + } + + /// Remove a worker from the registry + fn remove_worker(&self, worker_url: &str) { + let _ = self.remove_worker_from_registry(worker_url); + } + + /// Get all worker URLs from the registry + fn get_worker_urls(&self) -> Vec { + self.worker_registry.get_all_urls() + } +} + +#[async_trait] +impl RouterTrait for RouterManager { + fn as_any(&self) -> &dyn std::any::Any { + self + } + + /// Health check - return 503 if no routers available + async fn health(&self, _req: Request) -> Response { + // Health check should succeed if RouterManager exists, even without routers + // Individual router health can be checked via specific endpoints + (StatusCode::OK, "RouterManager is healthy").into_response() + } + + /// Health generate - check if any router can handle generate requests + async fn health_generate(&self, _req: Request) -> Response { + // Return 503 since we have no routers with workers + // TODO: Should check if any router has healthy workers + ( + StatusCode::SERVICE_UNAVAILABLE, + "No routers with healthy workers available", + ) + .into_response() + } + + /// Get server information - aggregate from all routers + async fn get_server_info(&self, _req: Request) -> Response { + // TODO: Aggregate info from all routers with healthy workers + // For now, return basic info about the RouterManager + ( + StatusCode::OK, + serde_json::json!({ + "router_manager": true, + "routers_count": self.routers.len(), + "workers_count": self.worker_registry.get_all().len() + }) + .to_string(), + ) + .into_response() + } + + /// Get available models - query from worker registry + async fn get_models(&self, _req: Request) -> Response { + // Get models from worker registry + let models = self.worker_registry.get_models(); + + if models.is_empty() { + (StatusCode::SERVICE_UNAVAILABLE, "No models available").into_response() + } else { + ( + StatusCode::OK, + serde_json::json!({ + "models": models + }) + .to_string(), + ) + .into_response() + } + } + + /// Get model information + async fn get_model_info(&self, _req: Request) -> Response { + // TODO: Extract model from request and route to appropriate router + // For now, return not implemented + ( + StatusCode::NOT_IMPLEMENTED, + "Model info endpoint not yet implemented in RouterManager", + ) + .into_response() + } + + /// Route a generate request + async fn route_generate( + &self, + headers: Option<&HeaderMap>, + body: &GenerateRequest, + _model_id: Option<&str>, + ) -> Response { + // Select router based on headers + // GenerateRequest doesn't have a model field + let router = self.select_router_for_request(headers, None); + + if let Some(router) = router { + // In multi-model mode, pass None since GenerateRequest doesn't have model field + router.route_generate(headers, body, None).await + } else { + // Return 404 when no router is available for the request + ( + StatusCode::NOT_FOUND, + "No router available for this request", + ) + .into_response() + } + } + + /// Route a chat completion request + async fn route_chat( + &self, + headers: Option<&HeaderMap>, + body: &ChatCompletionRequest, + _model_id: Option<&str>, + ) -> Response { + // Select router based on headers and model. + // When model is None, select_router_for_request considers all registered + // routers (including PD routers), so a single-model deployment will still + // route correctly even without the model field. + let model_id = body.model.as_deref(); + let router = self.select_router_for_request(headers, model_id); + + if let Some(router) = router { + router.route_chat(headers, body, model_id).await + } else { + let msg = match model_id { + Some(m) => format!("Model '{}' not found or no router available", m), + None => "No routers registered to handle this request".to_string(), + }; + (StatusCode::NOT_FOUND, msg).into_response() + } + } + + /// Route a completion request + async fn route_completion( + &self, + headers: Option<&HeaderMap>, + body: &CompletionRequest, + _model_id: Option<&str>, + ) -> Response { + // Select router based on headers and model. + // When model is None, select_router_for_request considers all registered + // routers (including PD routers), so a single-model deployment will still + // route correctly even without the model field. + let model_id = body.model.as_deref(); + let router = self.select_router_for_request(headers, model_id); + + if let Some(router) = router { + router.route_completion(headers, body, model_id).await + } else { + let msg = match model_id { + Some(m) => format!("Model '{}' not found or no router available", m), + None => "No routers registered to handle this request".to_string(), + }; + (StatusCode::NOT_FOUND, msg).into_response() + } + } + + async fn route_responses( + &self, + _headers: Option<&HeaderMap>, + _body: &ResponsesRequest, + _model_id: Option<&str>, + ) -> Response { + ( + StatusCode::NOT_IMPLEMENTED, + "responses api not yet implemented in inference gateway mode", + ) + .into_response() + } + + async fn delete_response(&self, _headers: Option<&HeaderMap>, _response_id: &str) -> Response { + ( + StatusCode::NOT_IMPLEMENTED, + "responses api not yet implemented in inference gateway mode", + ) + .into_response() + } + + async fn list_response_input_items( + &self, + _headers: Option<&HeaderMap>, + _response_id: &str, + ) -> Response { + ( + StatusCode::NOT_IMPLEMENTED, + "responses api not yet implemented in inference gateway mode", + ) + .into_response() + } + + async fn get_response(&self, headers: Option<&HeaderMap>, response_id: &str) -> Response { + let router = self.select_router_for_request(headers, None); + if let Some(router) = router { + router.get_response(headers, response_id).await + } else { + ( + StatusCode::NOT_FOUND, + format!("No router available to get response '{}'", response_id), + ) + .into_response() + } + } + + async fn cancel_response(&self, headers: Option<&HeaderMap>, response_id: &str) -> Response { + let router = self.select_router_for_request(headers, None); + if let Some(router) = router { + router.cancel_response(headers, response_id).await + } else { + ( + StatusCode::NOT_FOUND, + format!("No router available to cancel response '{}'", response_id), + ) + .into_response() + } + } + + /// Route embeddings request + async fn route_embeddings( + &self, + headers: Option<&HeaderMap>, + body: &EmbeddingRequest, + _model_id: Option<&str>, + ) -> Response { + // Select router based on headers and model. + // When model is None, select_router_for_request considers all registered + // routers, so a single-model deployment will still route correctly. + let model_id = body.model.as_deref(); + let router = self.select_router_for_request(headers, model_id); + + if let Some(router) = router { + router.route_embeddings(headers, body, model_id).await + } else { + let msg = match model_id { + Some(m) => format!("Model '{}' not found or no router available", m), + None => "No routers registered to handle this request".to_string(), + }; + (StatusCode::NOT_FOUND, msg).into_response() + } + } + + /// Flush cache on all routers and workers + async fn flush_cache(&self) -> Response { + // TODO: Call flush_cache on all routers that have workers + // For now, return success if we have any routers + if self.routers.is_empty() { + (StatusCode::SERVICE_UNAVAILABLE, "No routers configured").into_response() + } else { + // TODO: Actually flush cache on all routers + (StatusCode::OK, "Cache flush requested").into_response() + } + } + + /// Get worker loads from all routers + async fn get_worker_loads(&self) -> Response { + // Return worker loads from the registry + let workers = self.worker_registry.get_all(); + let loads: Vec = workers + .iter() + .map(|w| { + serde_json::json!({ + "url": w.url(), + "model": w.model_id(), + "load": w.load(), + "is_healthy": w.is_healthy() + }) + }) + .collect(); + + ( + StatusCode::OK, + serde_json::json!({ + "workers": loads + }) + .to_string(), + ) + .into_response() + } + + /// Get router type name + fn router_type(&self) -> &'static str { + "manager" + } + + /// Server readiness check - check if any router is ready + fn readiness(&self) -> Response { + if self.routers.is_empty() { + (StatusCode::SERVICE_UNAVAILABLE, "No routers configured").into_response() + } else { + // TODO: Check readiness of all routers + (StatusCode::OK, "Ready").into_response() + } + } + + /// Route a transparent proxy request through the appropriate router + async fn route_transparent( + &self, + headers: Option<&HeaderMap>, + path: &str, + method: &Method, + body: serde_json::Value, + ) -> Response { + // Select router based on headers (no model info available for transparent proxy) + let router = self.select_router_for_request(headers, None); + + if let Some(router) = router { + router.route_transparent(headers, path, method, body).await + } else { + // Return 404 when no router is available + ( + StatusCode::NOT_FOUND, + "No router available for this request", + ) + .into_response() + } + } +} + +// Note: get_first_available_router removed - we now properly handle +// router selection based on model and worker availability + +impl std::fmt::Debug for RouterManager { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("RouterManager") + .field("routers_count", &self.routers.len()) + .field("workers_count", &self.worker_registry.get_all().len()) + .field("default_router", &*self.default_router.read().unwrap()) + .finish() + } +} diff --git a/lmdeploy/router/src/routes/interface.rs b/lmdeploy/router/src/routes/interface.rs new file mode 100644 index 0000000000..da4386ef76 --- /dev/null +++ b/lmdeploy/router/src/routes/interface.rs @@ -0,0 +1,10 @@ +use std::pin::Pin; + +use crate::types::{AppError, ChatCompletionRequest, ChatCompletionResponse}; + +pub trait RouteHandle { + fn route<'a>( + &'a self, + request: &'a ChatCompletionRequest, + ) -> Pin> + Send + 'a>>; +} diff --git a/lmdeploy/router/src/routes/mod.rs b/lmdeploy/router/src/routes/mod.rs new file mode 100644 index 0000000000..b5ae22f380 --- /dev/null +++ b/lmdeploy/router/src/routes/mod.rs @@ -0,0 +1,7 @@ +pub mod interface; +pub mod round_robin_route; +pub mod routing_tree_builder; +pub mod single_server_route; + +pub use routing_tree_builder::RoutingTreeBuilder; +pub use single_server_route::SingleServerRoute; diff --git a/lmdeploy/router/src/routes/pool_route.rs b/lmdeploy/router/src/routes/pool_route.rs new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lmdeploy/router/src/routes/prefill_decode_route.rs b/lmdeploy/router/src/routes/prefill_decode_route.rs new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lmdeploy/router/src/routes/round_robin_route.rs b/lmdeploy/router/src/routes/round_robin_route.rs new file mode 100644 index 0000000000..d1423e48b8 --- /dev/null +++ b/lmdeploy/router/src/routes/round_robin_route.rs @@ -0,0 +1,36 @@ +use crate::routes::interface::RouteHandle; +use crate::{AppError, ChatCompletionResponse}; +use crate::{ChatCompletionRequest, routes::SingleServerRoute}; +use std::pin::Pin; +use std::sync::atomic::{AtomicUsize, Ordering}; + +pub struct RoundRobinRoute { + children: Vec, + idx: AtomicUsize, +} + +impl RoundRobinRoute { + pub fn new(servers: Vec<(&str, u16)>) -> Self { + Self { + children: servers + .into_iter() + .map(|(host, port)| SingleServerRoute::new(host, port)) + .collect(), + idx: AtomicUsize::new(0), + } + } +} + +impl RouteHandle for RoundRobinRoute { + fn route<'a>( + &'a self, + request: &'a ChatCompletionRequest, + ) -> Pin> + Send + 'a>> { + Box::pin(async move { + let idx = self.idx.fetch_add(1, Ordering::Relaxed) % self.children.len(); + let child = self.children.get(idx).unwrap(); + + Ok(child.route(request).await?) + }) + } +} diff --git a/lmdeploy/router/src/routes/routing_tree_builder.rs b/lmdeploy/router/src/routes/routing_tree_builder.rs new file mode 100644 index 0000000000..b902f377c1 --- /dev/null +++ b/lmdeploy/router/src/routes/routing_tree_builder.rs @@ -0,0 +1,64 @@ +use crate::routes::interface::RouteHandle; +use crate::routes::{SingleServerRoute, round_robin_route::RoundRobinRoute}; +use crate::types::ConfigurationError; +use crate::utils::json::RequireField; +use serde_json::{Map, Value}; + +pub struct RoutingTreeBuilder { + config: String, +} + +impl RoutingTreeBuilder { + pub fn new(config: String) -> Self { + Self { config } + } + + pub fn from_file(config: &std::path::PathBuf) -> Result { + let json = std::fs::read_to_string(config) + .map_err(|e| ConfigurationError(format!( + "Failed to read config file '{}': {}", + config.display(), + e + )))?; + Ok(RoutingTreeBuilder::new(json)) + } + + pub fn build_routing_tree( + self, + ) -> Result, ConfigurationError> { + let json: Value = serde_json::from_str(&self.config) + .map_err(|e| ConfigurationError(format!( + "Invalid JSON config: {}", + e + )))?; + let route: &Map = json.require("route")?; + + let root: Box = match route.require::<&str>("type")? { + "SingleServerRoute" => Box::new(SingleServerRoute::new( + route.require("host")?, + route.require("port")?, + )), + "RoundRobinRoute" => { + let servers = route + .require::<&Vec>("servers")? + .iter() + .map(|value| { + let host = value.require::<&str>("host")?; + let port = value.require::("port")?; + + Ok((host, port)) + }) + .collect::, ConfigurationError>>()?; + Box::new(RoundRobinRoute::new(servers)) + } + unsupported_type => { + return Err(ConfigurationError(format!( + "Unsupported route type: {}", + unsupported_type + ))); + } + }; + + Ok(root) + } +} diff --git a/lmdeploy/router/src/routes/single_server_route.rs b/lmdeploy/router/src/routes/single_server_route.rs new file mode 100644 index 0000000000..b16d4eabea --- /dev/null +++ b/lmdeploy/router/src/routes/single_server_route.rs @@ -0,0 +1,42 @@ +use crate::routes::interface::RouteHandle; +use crate::types::{AppError, ChatCompletionRequest, ChatCompletionResponse}; +use reqwest::header::CONTENT_TYPE; +use std::pin::Pin; + +pub struct SingleServerRoute { + host: String, + port: u16, + client: reqwest::Client, +} + +impl SingleServerRoute { + pub fn new(host: &str, port: u16) -> Self { + Self { + host: host.to_string(), + port, + client: reqwest::Client::new(), + } + } +} + +impl RouteHandle for SingleServerRoute { + fn route<'a>( + &'a self, + request: &'a ChatCompletionRequest, + ) -> Pin> + Send + 'a>> { + Box::pin(async move { + let response = self + .client + .post(format!( + "http://{}:{}/v1/chat/completions", + self.host, self.port + )) + .header(CONTENT_TYPE, "application/json") + .json(&request) + .send() + .await?; + + Ok(response.json::().await?) + }) + } +} diff --git a/lmdeploy/router/src/server.rs b/lmdeploy/router/src/server.rs new file mode 100644 index 0000000000..0f8bca4e27 --- /dev/null +++ b/lmdeploy/router/src/server.rs @@ -0,0 +1,1328 @@ +use crate::{ + config::{HistoryBackend, LMDeployMigrationProtocol, RouterConfig, TraceConfig}, + core::{WorkerRegistry, WorkerType}, + data_connector::{MemoryResponseStorage, NoOpResponseStorage, SharedResponseStorage}, + logging::{self, LoggingConfig}, + metrics::{self, PrometheusConfig}, + middleware::{self, QueuedRequest, TokenBucket}, + policies::PolicyRegistry, + protocols::{ + spec::{ChatCompletionRequest, CompletionRequest, EmbeddingRequest, GenerateRequest}, + worker_spec::{WorkerApiResponse, WorkerConfigRequest, WorkerErrorResponse}, + }, + routers::{ + factory::LMDeployPDRouterParams, + router_manager::{RouterId, RouterManager}, + RouterFactory, RouterTrait, + }, + service_discovery::{start_service_discovery, ServiceDiscoveryConfig}, +}; +use axum::{ + extract::{rejection::JsonRejection, DefaultBodyLimit, Path, Query, Request, State}, + http::StatusCode, + response::{IntoResponse, Response}, + routing::{delete, get, post}, + serve, Json, Router, +}; +use reqwest::Client; +use serde::{Deserialize, Serialize}; +use serde_json::json; +use std::{ + collections::{HashMap, VecDeque}, + sync::atomic::{AtomicBool, Ordering}, + sync::Arc, + time::Duration, +}; +use tokio::{ + net::TcpListener, + signal, spawn, + sync::{Mutex, RwLock}, +}; +use tracing::{error, info, warn, Level}; + +#[derive(Clone)] +pub struct AppContext { + pub client: Client, + pub router_config: RouterConfig, + pub rate_limiter: Arc, + pub worker_registry: Arc, + pub policy_registry: Arc, + pub router_manager: Option>, + pub response_storage: SharedResponseStorage, + pub api_key_cache: Arc>>, + pub api_key_validation_urls: Arc>, + pub(crate) lmdeploy_nodes: Arc>>, +} + +impl AppContext { + pub fn new( + router_config: RouterConfig, + client: Client, + max_concurrent_requests: usize, + rate_limit_tokens_per_second: Option, + api_key_validation_urls: Vec, + ) -> Result { + let rate_limit_tokens = rate_limit_tokens_per_second.unwrap_or(max_concurrent_requests); + let rate_limiter = Arc::new(TokenBucket::new(max_concurrent_requests, rate_limit_tokens)); + + let worker_registry = Arc::new(WorkerRegistry::new()); + let policy_registry = Arc::new(PolicyRegistry::new(router_config.policy.clone())); + + let router_manager = None; + + // Initialize response storage based on configuration + let response_storage: SharedResponseStorage = match router_config.history_backend { + HistoryBackend::Memory => Arc::new(MemoryResponseStorage::new()), + HistoryBackend::None => Arc::new(NoOpResponseStorage::new()), + }; + + Ok(Self { + client, + router_config, + rate_limiter, + worker_registry, + policy_registry, + router_manager, + response_storage, + api_key_cache: Arc::new(RwLock::new(HashMap::new())), + api_key_validation_urls: Arc::new(api_key_validation_urls), + lmdeploy_nodes: Arc::new(Mutex::new(HashMap::new())), + }) + } +} + +const LMDEPLOY_ROLE_HYBRID: u8 = 1; +const LMDEPLOY_ROLE_PREFILL: u8 = 2; +const LMDEPLOY_ROLE_DECODE: u8 = 3; + +fn default_lmdeploy_role() -> u8 { + LMDEPLOY_ROLE_HYBRID +} + +/// Wire-compatible with lmdeploy.serve.proxy.proxy.Status. +#[derive(Debug, Clone, Deserialize, Serialize)] +pub(crate) struct LMDeployNodeStatus { + #[serde(default = "default_lmdeploy_role")] + role: u8, + #[serde(default)] + models: Vec, + #[serde(default)] + unfinished: usize, + #[serde(default)] + latency: VecDeque, + #[serde(default)] + speed: Option, +} + +impl Default for LMDeployNodeStatus { + fn default() -> Self { + Self { + role: LMDEPLOY_ROLE_HYBRID, + models: Vec::new(), + unfinished: 0, + latency: VecDeque::new(), + speed: None, + } + } +} + +/// Wire-compatible with lmdeploy.serve.proxy.proxy.Node. +#[derive(Debug, Deserialize)] +struct LMDeployNode { + url: String, + #[serde(default)] + status: Option, +} + +#[derive(Clone)] +pub struct AppState { + pub router: Arc, + pub context: Arc, + pub concurrency_queue_tx: Option>, + pub router_manager: Option>, +} + +// Fallback handler for unmatched routes +async fn sink_handler() -> Response { + StatusCode::NOT_FOUND.into_response() +} + +/// Transparent proxy handler for unmatched routes +/// Routes requests through the router's route_transparent method +async fn transparent_proxy_handler(State(state): State>, req: Request) -> Response { + let headers = req.headers().clone(); + + // Check authorization + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + // Extract path and method + let path = req.uri().path().to_string(); + let method = req.method().clone(); + + // Read body + let body_bytes = match axum::body::to_bytes(req.into_body(), usize::MAX).await { + Ok(bytes) => bytes, + Err(e) => { + return ( + StatusCode::BAD_REQUEST, + format!("Failed to read request body: {}", e), + ) + .into_response() + } + }; + + // Parse body as JSON + let body_json: serde_json::Value = if body_bytes.is_empty() { + serde_json::Value::Null + } else { + match serde_json::from_slice(&body_bytes) { + Ok(json) => json, + Err(e) => { + return (StatusCode::BAD_REQUEST, format!("Invalid JSON body: {}", e)) + .into_response() + } + } + }; + + // Route through transparent proxy + state + .router + .route_transparent(Some(&headers), &path, &method, body_json) + .await +} + +// Health check endpoints +async fn liveness(State(state): State>, req: Request) -> Response { + let headers = req.headers().clone(); + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + state.router.liveness() +} + +async fn readiness(State(state): State>, req: Request) -> Response { + let headers = req.headers().clone(); + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + state.router.readiness() +} + +async fn health(State(state): State>, req: Request) -> Response { + let headers = req.headers().clone(); + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + state.router.health(req).await +} + +async fn health_generate(State(state): State>, req: Request) -> Response { + let headers = req.headers().clone(); + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + state.router.health_generate(req).await +} + +async fn get_server_info(State(state): State>, req: Request) -> Response { + let headers = req.headers().clone(); + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + state.router.get_server_info(req).await +} + +async fn v1_models(State(state): State>, req: Request) -> Response { + let headers = req.headers().clone(); + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + state.router.get_models(req).await +} + +async fn get_model_info(State(state): State>, req: Request) -> Response { + let headers = req.headers().clone(); + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + state.router.get_model_info(req).await +} + +// Generation endpoints +// The RouterTrait now accepts optional headers and typed body directly +async fn generate( + State(state): State>, + headers: http::HeaderMap, + body: Result, JsonRejection>, +) -> Response { + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + let body = match body { + Ok(Json(body)) => body, + Err(error) => return (StatusCode::BAD_REQUEST, error.to_string()).into_response(), + }; + + state + .router + .route_generate(Some(&headers), &body, None) + .await +} + +async fn v1_chat_completions( + State(state): State>, + headers: http::HeaderMap, + Json(body): Json, +) -> Response { + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + state.router.route_chat(Some(&headers), &body, None).await +} + +async fn v1_completions( + State(state): State>, + headers: http::HeaderMap, + Json(body): Json, +) -> Response { + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + state + .router + .route_completion(Some(&headers), &body, None) + .await +} + +async fn v1_responses( + State(state): State>, + headers: http::HeaderMap, + Json(body): Json, +) -> Response { + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + state + .router + .route_transparent(Some(&headers), "/v1/responses", &http::Method::POST, body) + .await +} + +async fn v1_embeddings( + State(state): State>, + headers: http::HeaderMap, + Json(body): Json, +) -> Response { + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + state + .router + .route_embeddings(Some(&headers), &body, None) + .await +} + +async fn v1_responses_get( + State(state): State>, + Path(response_id): Path, + headers: http::HeaderMap, +) -> Response { + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + state + .router + .get_response(Some(&headers), &response_id) + .await +} + +async fn v1_responses_cancel( + State(state): State>, + Path(response_id): Path, + headers: http::HeaderMap, +) -> Response { + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + state + .router + .cancel_response(Some(&headers), &response_id) + .await +} + +async fn v1_responses_delete( + State(state): State>, + Path(response_id): Path, + headers: http::HeaderMap, +) -> Response { + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + // Python server does not support this yet + state + .router + .delete_response(Some(&headers), &response_id) + .await +} + +async fn v1_responses_list_input_items( + State(state): State>, + Path(response_id): Path, + headers: http::HeaderMap, +) -> Response { + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + // Python server does not support this yet + state + .router + .list_response_input_items(Some(&headers), &response_id) + .await +} + +const AUTH_FAILURE_MESSAGE: &str = + "You must provide a valid API key. Obtain one from http://helmholtz.cloud"; + +// ---------- Worker management endpoints (Legacy) ---------- + +#[derive(Deserialize)] +struct UrlQuery { + url: String, +} + +async fn authorize_request( + state: &Arc, + headers: &http::HeaderMap, +) -> Result<(), Box> { + let validation_urls = state.context.api_key_validation_urls.as_ref(); + if validation_urls.is_empty() { + return Ok(()); + } + + let auth_header = headers + .get(http::header::AUTHORIZATION) + .and_then(|value| value.to_str().ok()) + .unwrap_or_default(); + + let token = auth_header + .strip_prefix("Bearer ") + .map(str::trim) + .filter(|value| !value.is_empty()) + .ok_or_else(|| { + Box::new((StatusCode::UNAUTHORIZED, AUTH_FAILURE_MESSAGE).into_response()) + })?; + + if let Some(valid) = state.context.api_key_cache.read().await.get(token).copied() { + if valid { + return Ok(()); + } + return Err(Box::new( + (StatusCode::UNAUTHORIZED, AUTH_FAILURE_MESSAGE).into_response(), + )); + } + + let mut validated = false; + for url in validation_urls { + match state + .context + .client + .get(url) + .header(http::header::AUTHORIZATION, format!("Bearer {token}")) + .send() + .await + { + Ok(response) if response.status() == StatusCode::OK => { + validated = true; + break; + } + Ok(_) => { + continue; + } + Err(err) => { + warn!("Failed to validate API key against {url}: {err}"); + } + } + } + + state + .context + .api_key_cache + .write() + .await + .insert(token.to_string(), validated); + + if validated { + Ok(()) + } else { + Err(Box::new( + (StatusCode::UNAUTHORIZED, AUTH_FAILURE_MESSAGE).into_response(), + )) + } +} + +fn validate_lmdeploy_role_for_router(state: &AppState, role: u8) -> Result<(), String> { + if state + .router + .as_any() + .is::() + { + return match role { + LMDEPLOY_ROLE_PREFILL | LMDEPLOY_ROLE_DECODE => Ok(()), + LMDEPLOY_ROLE_HYBRID => Err( + "LMDeploy PD mode accepts only Prefill (2) or Decode (3) registrations".to_string(), + ), + _ => Err(format!( + "Invalid LMDeploy engine role {role}; expected Hybrid=1, Prefill=2, or Decode=3" + )), + }; + } + + if state + .router + .as_any() + .is::() + { + return match role { + LMDEPLOY_ROLE_HYBRID => Ok(()), + LMDEPLOY_ROLE_PREFILL | LMDEPLOY_ROLE_DECODE => { + Err("Regular mode accepts only Hybrid (1) LMDeploy registrations".to_string()) + } + _ => Err(format!( + "Invalid LMDeploy engine role {role}; expected Hybrid=1, Prefill=2, or Decode=3" + )), + }; + } + + Err("LMDeploy node registration is supported only in regular or LMDeploy PD mode".to_string()) +} + +fn registered_lmdeploy_role(state: &AppState, node_url: &str) -> Option { + let dp_prefix = format!("{}@", node_url); + state + .context + .worker_registry + .get_all() + .into_iter() + .find(|worker| worker.url() == node_url || worker.url().starts_with(&dp_prefix)) + .map(|worker| match worker.worker_type() { + WorkerType::Regular => LMDEPLOY_ROLE_HYBRID, + WorkerType::Prefill => LMDEPLOY_ROLE_PREFILL, + WorkerType::Decode => LMDEPLOY_ROLE_DECODE, + }) +} + +fn register_lmdeploy_worker_unchecked( + state: &AppState, + node_url: &str, + role: u8, +) -> Result<(), String> { + validate_lmdeploy_role_for_router(state, role)?; + + if let Some(router) = state + .router + .as_any() + .downcast_ref::( + ) { + return match role { + LMDEPLOY_ROLE_PREFILL => router + .register_prefill_server_unchecked(node_url.to_string()) + .map(|_| ()) + .map_err(|error| error.to_string()), + LMDEPLOY_ROLE_DECODE => router + .register_decode_server_unchecked(node_url.to_string()) + .map(|_| ()) + .map_err(|error| error.to_string()), + _ => unreachable!("role was validated above"), + }; + } + + let router = state + .router + .as_any() + .downcast_ref::() + .expect("regular router type was validated above"); + router.register_worker_unchecked(node_url).map(|_| ()) +} + +async fn remove_lmdeploy_worker(state: &AppState, node_url: &str, role: u8) -> Result<(), String> { + validate_lmdeploy_role_for_router(state, role)?; + + if let Some(router) = state + .router + .as_any() + .downcast_ref::( + ) { + return match role { + LMDEPLOY_ROLE_PREFILL => router + .remove_prefill_server(node_url) + .await + .map(|_| ()) + .map_err(|error| error.to_string()), + LMDEPLOY_ROLE_DECODE => router + .remove_decode_server(node_url) + .await + .map(|_| ()) + .map_err(|error| error.to_string()), + _ => unreachable!("role was validated above"), + }; + } + + let router = state + .router + .as_any() + .downcast_ref::() + .expect("regular router type was validated above"); + router.remove_worker(node_url); + Ok(()) +} + +/// LMDeploy-compatible dynamic registration endpoint. This endpoint is +/// intentionally unauthenticated: `lmdeploy serve api_server --proxy-url` +/// does not send an Authorization header. +async fn add_lmdeploy_node( + State(state): State>, + Json(node): Json, +) -> Response { + if node.url.is_empty() { + return ( + StatusCode::BAD_REQUEST, + Json("Failed to add, please check the input url."), + ) + .into_response(); + } + + let status = node.status.unwrap_or_default(); + if let Err(error) = validate_lmdeploy_role_for_router(&state, status.role) { + return (StatusCode::BAD_REQUEST, error).into_response(); + } + + // Serialize registration/removal for a URL. This also makes repeated + // startup registration idempotent instead of surfacing WorkerAlreadyExists. + let mut nodes = state.context.lmdeploy_nodes.lock().await; + let current_role = registered_lmdeploy_role(&state, &node.url); + + if let Some(old_role) = current_role { + if old_role != status.role { + if let Err(error) = remove_lmdeploy_worker(&state, &node.url, old_role).await { + return (StatusCode::BAD_REQUEST, error).into_response(); + } + if let Err(error) = register_lmdeploy_worker_unchecked(&state, &node.url, status.role) { + let rollback = register_lmdeploy_worker_unchecked(&state, &node.url, old_role); + let message = match rollback { + Ok(()) => error, + Err(rollback_error) => { + format!("{error}; failed to restore previous role: {rollback_error}") + } + }; + return (StatusCode::BAD_REQUEST, message).into_response(); + } + } + } else if let Err(error) = register_lmdeploy_worker_unchecked(&state, &node.url, status.role) { + return (StatusCode::BAD_REQUEST, error).into_response(); + } + + info!( + "LMDeploy node registered: url={}, role={}, models={:?}", + node.url, status.role, status.models + ); + nodes.insert(node.url, status); + Json("Added successfully").into_response() +} + +async fn remove_lmdeploy_node( + State(state): State>, + Json(node): Json, +) -> Response { + let mut nodes = state.context.lmdeploy_nodes.lock().await; + let role = registered_lmdeploy_role(&state, &node.url) + .or_else(|| nodes.get(&node.url).map(|status| status.role)) + .or_else(|| node.status.map(|status| status.role)); + + if let Some(role) = role { + // Native LMDeploy treats removal as idempotent. A missing worker can + // race with health-based/manual removal, so it is still a success. + if let Err(error) = remove_lmdeploy_worker(&state, &node.url, role).await { + if registered_lmdeploy_role(&state, &node.url).is_some() { + return (StatusCode::BAD_REQUEST, error).into_response(); + } + } + } + + nodes.remove(&node.url); + info!("LMDeploy node removed: url={}", node.url); + Json("Deleted successfully").into_response() +} + +async fn lmdeploy_node_status(State(state): State>) -> Response { + let dynamic_nodes = state.context.lmdeploy_nodes.lock().await.clone(); + let mut nodes = dynamic_nodes; + + for worker in state.context.worker_registry.get_all() { + let role = match worker.worker_type() { + WorkerType::Regular => LMDEPLOY_ROLE_HYBRID, + WorkerType::Prefill => LMDEPLOY_ROLE_PREFILL, + WorkerType::Decode => LMDEPLOY_ROLE_DECODE, + }; + nodes + .entry(worker.url().to_string()) + .or_insert_with(|| LMDeployNodeStatus { + role, + ..LMDeployNodeStatus::default() + }); + } + + Json(nodes).into_response() +} + +async fn add_worker( + State(state): State>, + Query(UrlQuery { url }): Query, + headers: http::HeaderMap, +) -> Response { + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + match state.router.add_worker(&url).await { + Ok(message) => (StatusCode::OK, message).into_response(), + Err(error) => (StatusCode::BAD_REQUEST, error).into_response(), + } +} + +async fn list_workers(State(state): State>, headers: http::HeaderMap) -> Response { + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + let worker_list = state.router.get_worker_urls(); + Json(serde_json::json!({ "urls": worker_list })).into_response() +} + +async fn remove_worker( + State(state): State>, + Query(UrlQuery { url }): Query, + headers: http::HeaderMap, +) -> Response { + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + state.router.remove_worker(&url); + ( + StatusCode::OK, + format!("Successfully removed worker: {url}"), + ) + .into_response() +} + +async fn flush_cache(State(state): State>, headers: http::HeaderMap) -> Response { + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + state.router.flush_cache().await +} + +async fn get_loads(State(state): State>, headers: http::HeaderMap) -> Response { + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + state.router.get_worker_loads().await +} + +// ---------- Worker management endpoints (RESTful) ---------- + +/// POST /workers - Add a new worker with full configuration +async fn create_worker( + State(state): State>, + headers: http::HeaderMap, + Json(config): Json, +) -> Response { + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + // Check if we have a RouterManager (enable_igw=true) + if let Some(router_manager) = &state.router_manager { + // Call RouterManager's add_worker method directly with the full config + match router_manager.add_worker(config).await { + Ok(response) => (StatusCode::OK, Json(response)).into_response(), + Err(error) => (StatusCode::BAD_REQUEST, Json(error)).into_response(), + } + } else { + // In single router mode, use the router's add_worker with basic config + match state.router.add_worker(&config.url).await { + Ok(message) => { + let response = WorkerApiResponse { + success: true, + message, + worker: None, + }; + (StatusCode::OK, Json(response)).into_response() + } + Err(error) => { + let error_response = WorkerErrorResponse { + error, + code: "ADD_WORKER_FAILED".to_string(), + }; + (StatusCode::BAD_REQUEST, Json(error_response)).into_response() + } + } + } +} + +async fn list_workers_rest( + State(state): State>, + headers: http::HeaderMap, +) -> Response { + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + if let Some(router_manager) = &state.router_manager { + let response = router_manager.list_workers(); + Json(response).into_response() + } else { + // In single router mode, get detailed worker info from registry + let workers = state.context.worker_registry.get_all(); + let response = serde_json::json!({ + "workers": workers.iter().map(|worker| { + let worker_info = serde_json::json!({ + "url": worker.url(), + "model_id": worker.model_id(), + "worker_type": match worker.worker_type() { + WorkerType::Regular => "regular", + WorkerType::Prefill => "prefill", + WorkerType::Decode => "decode", + }, + "is_healthy": worker.is_healthy(), + "load": worker.load(), + "connection_mode": format!("{:?}", worker.connection_mode()), + "priority": worker.priority(), + "cost": worker.cost(), + }); + + worker_info + }).collect::>(), + "total": workers.len(), + "stats": { + "prefill_count": state.context.worker_registry.get_prefill_workers().len(), + "decode_count": state.context.worker_registry.get_decode_workers().len(), + "regular_count": state.context.worker_registry.get_by_type(&WorkerType::Regular).len(), + } + }); + Json(response).into_response() + } +} + +/// GET /workers/{url} - Get specific worker info +async fn get_worker( + State(state): State>, + Path(url): Path, + headers: http::HeaderMap, +) -> Response { + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + if let Some(router_manager) = &state.router_manager { + if let Some(worker) = router_manager.get_worker(&url) { + Json(worker).into_response() + } else { + let error = WorkerErrorResponse { + error: format!("Worker {url} not found"), + code: "WORKER_NOT_FOUND".to_string(), + }; + (StatusCode::NOT_FOUND, Json(error)).into_response() + } + } else { + let workers = state.router.get_worker_urls(); + if workers.contains(&url) { + Json(json!({ + "url": url, + "model_id": "unknown", + "is_healthy": true + })) + .into_response() + } else { + let error = WorkerErrorResponse { + error: format!("Worker {url} not found"), + code: "WORKER_NOT_FOUND".to_string(), + }; + (StatusCode::NOT_FOUND, Json(error)).into_response() + } + } +} + +/// DELETE /workers/{url} - Remove a worker +async fn delete_worker( + State(state): State>, + Path(url): Path, + headers: http::HeaderMap, +) -> Response { + if let Err(response) = authorize_request(&state, &headers).await { + return *response; + } + + if let Some(router_manager) = &state.router_manager { + match router_manager.remove_worker_from_registry(&url) { + Ok(response) => (StatusCode::OK, Json(response)).into_response(), + Err(error) => (StatusCode::BAD_REQUEST, Json(error)).into_response(), + } + } else { + // In single router mode, use router's remove_worker + state.router.remove_worker(&url); + let response = WorkerApiResponse { + success: true, + message: format!("Worker {url} removed successfully"), + worker: None, + }; + (StatusCode::OK, Json(response)).into_response() + } +} + +pub struct ServerConfig { + pub host: String, + pub port: u16, + pub router_config: RouterConfig, + pub max_payload_size: usize, + pub log_dir: Option, + pub log_level: Option, + pub service_discovery_config: Option, + pub prometheus_config: Option, + pub request_timeout_secs: u64, + pub request_id_headers: Option>, + pub trace_config: Option, +} + +/// Build the Axum application with all routes and middleware using the +/// current runtime OpenTelemetry state. +pub fn build_app( + app_state: Arc, + max_payload_size: usize, + request_id_headers: Vec, + cors_allowed_origins: Vec, + enable_transparent_proxy: bool, +) -> Router { + build_app_with_request_tracing( + app_state, + max_payload_size, + request_id_headers, + cors_allowed_origins, + enable_transparent_proxy, + crate::otel_trace::is_otel_enabled(), + ) +} + +/// Build the Axum application with an explicit request-tracing toggle. +pub fn build_app_with_request_tracing( + app_state: Arc, + max_payload_size: usize, + request_id_headers: Vec, + cors_allowed_origins: Vec, + enable_transparent_proxy: bool, + enable_request_tracing: bool, +) -> Router { + // Create routes + let protected_routes = Router::new() + .route("/generate", post(generate)) + .route("/v1/chat/completions", post(v1_chat_completions)) + .route("/v1/completions", post(v1_completions)) + .route("/v1/responses", post(v1_responses)) + .route("/v1/embeddings", post(v1_embeddings)) + .route("/v1/responses/{response_id}", get(v1_responses_get)) + .route( + "/v1/responses/{response_id}/cancel", + post(v1_responses_cancel), + ) + .route("/v1/responses/{response_id}", delete(v1_responses_delete)) + .route( + "/v1/responses/{response_id}/input", + get(v1_responses_list_input_items), + ) + .route_layer(axum::middleware::from_fn_with_state( + app_state.clone(), + middleware::concurrency_limit_middleware, + )); + + let public_routes = Router::new() + .route("/liveness", get(liveness)) + .route("/readiness", get(readiness)) + .route("/health", get(health)) + .route("/health_generate", get(health_generate)) + .route("/v1/models", get(v1_models)) + .route("/get_model_info", get(get_model_info)) + .route("/get_server_info", get(get_server_info)) + .route("/nodes/status", get(lmdeploy_node_status)) + .route("/nodes/add", post(add_lmdeploy_node)) + .route("/nodes/remove", post(remove_lmdeploy_node)); + + let admin_routes = Router::new() + .route("/add_worker", post(add_worker)) + .route("/remove_worker", post(remove_worker)) + .route("/list_workers", get(list_workers)) + .route("/flush_cache", post(flush_cache)) + .route("/get_loads", get(get_loads)); + + // Worker management routes + let worker_routes = Router::new() + .route("/workers", post(create_worker)) + .route("/workers", get(list_workers_rest)) + .route("/workers/{url}", get(get_worker)) + .route("/workers/{url}", delete(delete_worker)); + + // Build base app with all routes and middleware + let base_app = Router::new() + .merge(protected_routes) + .merge(public_routes) + .merge(admin_routes) + .merge(worker_routes) + // Request body size limiting + .layer(DefaultBodyLimit::max(max_payload_size)) + .layer(tower_http::limit::RequestBodyLimitLayer::new( + max_payload_size, + )); + + let base_app = if enable_request_tracing { + base_app.layer(middleware::create_logging_layer()) + } else { + base_app + }; + + let base_app = base_app + .layer(middleware::RequestIdLayer::new(request_id_headers)) + .layer(create_cors_layer(cors_allowed_origins)); + + // Choose fallback based on transparent proxy mode + if enable_transparent_proxy { + base_app + .fallback(transparent_proxy_handler) + .with_state(app_state) + } else { + base_app.fallback(sink_handler).with_state(app_state) + } +} + +pub async fn startup(config: ServerConfig) -> Result<(), Box> { + println!("DEBUG: Server startup function called"); + + // Only initialize logging if not already done (for Python bindings support) + static LOGGING_INITIALIZED: AtomicBool = AtomicBool::new(false); + + println!("DEBUG: Initializing logging"); + let _log_guard = if !LOGGING_INITIALIZED.swap(true, Ordering::SeqCst) { + Some(logging::init_logging( + LoggingConfig { + level: config + .log_level + .as_deref() + .and_then(|s| match s.to_uppercase().parse::() { + Ok(l) => Some(l), + Err(_) => { + warn!("Invalid log level string: '{s}'. Defaulting to INFO."); + None + } + }) + .unwrap_or(Level::INFO), + json_format: false, + log_dir: config.log_dir.clone(), + colorize: true, + log_file_name: "lmdeploy-router".to_string(), + log_targets: None, + }, + config.trace_config.clone(), + )) + } else { + if config.trace_config.is_some() && !crate::otel_trace::is_otel_enabled() { + warn!( + "Tracing was requested after logging was already initialized; \ + the existing subscriber will continue without enabling OpenTelemetry" + ); + } + None + }; + println!("DEBUG: Logging initialized"); + + // Initialize prometheus metrics exporter + println!("DEBUG: Initializing Prometheus metrics"); + if let Some(prometheus_config) = config.prometheus_config { + metrics::start_prometheus(prometheus_config); + } + println!("DEBUG: Prometheus metrics initialized"); + + info!( + "Starting router on {}:{} | mode: {:?} | policy: {:?} | kv_connector: {:?} | max_payload: {}MB", + config.host, + config.port, + config.router_config.mode, + config.router_config.policy, + config.router_config.kv_connector, + config.max_payload_size / (1024 * 1024) + ); + + println!("DEBUG: Creating HTTP client"); + let client = Client::builder() + .pool_idle_timeout(Some(Duration::from_secs(50))) + .pool_max_idle_per_host(500) + .timeout(Duration::from_secs(config.request_timeout_secs)) + .connect_timeout(Duration::from_secs(10)) + .tcp_nodelay(true) + .tcp_keepalive(Some(Duration::from_secs(30))) + .build() + .expect("Failed to create HTTP client"); + println!("DEBUG: HTTP client created"); + + // Create the application context with all dependencies + println!("DEBUG: Creating AppContext"); + let app_context = AppContext::new( + config.router_config.clone(), + client.clone(), + config.router_config.max_concurrent_requests, + config.router_config.rate_limit_tokens_per_second, + config.router_config.api_key_validation_urls.clone(), + )?; + println!("DEBUG: AppContext created"); + + let app_context = Arc::new(app_context); + + // Create the appropriate router based on enable_igw flag + let (router, router_manager): (Arc, Option>) = + if config.router_config.enable_igw { + info!("Multi-router mode enabled (enable_igw=true)"); + + // Create RouterManager with shared registries from AppContext + let router_manager = Arc::new(RouterManager::new( + config.router_config.clone(), + client.clone(), + app_context.worker_registry.clone(), + app_context.policy_registry.clone(), + )); + + // 1. HTTP Regular Router + match RouterFactory::create_regular_router( + &[], // Empty worker list - workers added later + &app_context, + ) + .await + { + Ok(http_regular) => { + info!("Created HTTP Regular router"); + router_manager.register_router( + RouterId::new("http-regular".to_string()), + Arc::from(http_regular), + ); + } + Err(e) => { + warn!("Failed to create HTTP Regular router: {e}"); + } + } + + // 2. HTTP LMDeploy PD Router + match RouterFactory::create_lmdeploy_pd_router(LMDeployPDRouterParams { + prefill_urls: &[], + decode_urls: &[], + migration_protocol: LMDeployMigrationProtocol::default(), + rdma_config: None, + dummy_prefill: false, + prefill_policy_config: None, + decode_policy_config: None, + main_policy_config: &config.router_config.policy, + ctx: &app_context, + }) + .await + { + Ok(http_pd) => { + info!("Created HTTP LMDeploy PD router"); + router_manager + .register_router(RouterId::new("http-pd".to_string()), Arc::from(http_pd)); + } + Err(e) => { + warn!("Failed to create HTTP PD router: {e}"); + } + } + + info!( + "RouterManager initialized with {} routers", + router_manager.router_count() + ); + ( + router_manager.clone() as Arc, + Some(router_manager), + ) + } else { + info!("Single router mode (enable_igw=false)"); + // Create single router with the context + ( + Arc::from(RouterFactory::create_router(&app_context).await?), + None, + ) + }; + + // Start health checker for all workers in the registry + let _health_checker = app_context + .worker_registry + .start_health_checker(config.router_config.health_check.check_interval_secs); + info!( + "Started health checker for workers with {}s interval", + config.router_config.health_check.check_interval_secs + ); + + // Set up concurrency limiter with queue if configured + let (limiter, processor) = middleware::ConcurrencyLimiter::new( + app_context.rate_limiter.clone(), + config.router_config.queue_size, + Duration::from_secs(config.router_config.queue_timeout_secs), + ); + + // Start queue processor if enabled + if let Some(processor) = processor { + tokio::spawn(processor.run()); + info!( + "Started request queue with size: {}, timeout: {}s", + config.router_config.queue_size, config.router_config.queue_timeout_secs + ); + } + + // Create app state with router and context + let app_state = Arc::new(AppState { + router, + context: app_context.clone(), + concurrency_queue_tx: limiter.queue_tx.clone(), + router_manager, + }); + let router_arc = Arc::clone(&app_state.router); + + // Start the service discovery if enabled + if let Some(service_discovery_config) = config.service_discovery_config { + if service_discovery_config.enabled { + match start_service_discovery(service_discovery_config, router_arc).await { + Ok(handle) => { + info!("Service discovery started"); + // Spawn a task to handle the service discovery thread + spawn(async move { + if let Err(e) = handle.await { + error!("Service discovery task failed: {:?}", e); + } + }); + } + Err(e) => { + error!("Failed to start service discovery: {e}"); + warn!("Continuing without service discovery"); + } + } + } + } + + info!( + "Router ready | workers: {:?}", + app_state.router.get_worker_urls() + ); + + let request_id_headers = config.request_id_headers.clone().unwrap_or_else(|| { + vec![ + "x-request-id".to_string(), + "x-correlation-id".to_string(), + "x-trace-id".to_string(), + "request-id".to_string(), + ] + }); + + // Build the application + // Enable transparent proxy for all routing modes + let enable_transparent_proxy = true; + + let app = build_app_with_request_tracing( + app_state, + config.max_payload_size, + request_id_headers, + config.router_config.cors_allowed_origins.clone(), + enable_transparent_proxy, + crate::otel_trace::is_otel_enabled(), + ); + + let addr = format!("{}:{}", config.host, config.port); + let listener = TcpListener::bind(&addr).await?; + info!("Starting server on {}", addr); + serve(listener, app) + .with_graceful_shutdown(shutdown_signal()) + .await + .map_err(|e| Box::new(e) as Box)?; + + Ok(()) +} + +// Graceful shutdown handler +async fn shutdown_signal() { + let ctrl_c = async { + signal::ctrl_c() + .await + .expect("failed to install Ctrl+C handler"); + }; + + #[cfg(unix)] + let terminate = async { + signal::unix::signal(signal::unix::SignalKind::terminate()) + .expect("failed to install signal handler") + .recv() + .await; + }; + + #[cfg(not(unix))] + let terminate = std::future::pending::<()>(); + + tokio::select! { + _ = ctrl_c => { + info!("Received Ctrl+C, starting graceful shutdown"); + }, + _ = terminate => { + info!("Received terminate signal, starting graceful shutdown"); + }, + } +} + +// CORS Layer Creation +fn create_cors_layer(allowed_origins: Vec) -> tower_http::cors::CorsLayer { + use tower_http::cors::Any; + + let cors = if allowed_origins.is_empty() { + // Allow all origins if none specified + tower_http::cors::CorsLayer::new() + .allow_origin(Any) + .allow_methods(Any) + .allow_headers(Any) + .expose_headers(Any) + } else { + // Restrict to specific origins + let origins: Vec = allowed_origins + .into_iter() + .filter_map(|origin| origin.parse().ok()) + .collect(); + + tower_http::cors::CorsLayer::new() + .allow_origin(origins) + .allow_methods([http::Method::GET, http::Method::POST, http::Method::OPTIONS]) + .allow_headers([http::header::CONTENT_TYPE, http::header::AUTHORIZATION]) + .expose_headers([http::header::HeaderName::from_static("x-request-id")]) + }; + + cors.max_age(Duration::from_secs(3600)) +} diff --git a/lmdeploy/router/src/service_discovery.rs b/lmdeploy/router/src/service_discovery.rs new file mode 100644 index 0000000000..ff88c60007 --- /dev/null +++ b/lmdeploy/router/src/service_discovery.rs @@ -0,0 +1,1126 @@ +use crate::routers::RouterTrait; + +use futures::{StreamExt, TryStreamExt}; +use k8s_openapi::api::core::v1::Pod; +use kube::{ + api::Api, + runtime::watcher::{watcher, Config}, + runtime::WatchStreamExt, + Client, +}; +use std::collections::{HashMap, HashSet}; +use std::net::{IpAddr, SocketAddr}; + +use rustls; +use std::sync::{Arc, Mutex}; +use std::time::Duration; +use tokio::task; +use tokio::time; +use tracing::{debug, error, info, warn}; + +/// Represents the service discovery configuration +#[derive(Debug, Clone)] +pub struct ServiceDiscoveryConfig { + pub enabled: bool, + pub selector: HashMap, + pub check_interval: Duration, + pub port: u16, + pub namespace: Option, + // PD mode specific configuration + pub pd_mode: bool, + pub prefill_selector: HashMap, + pub decode_selector: HashMap, +} + +impl Default for ServiceDiscoveryConfig { + fn default() -> Self { + ServiceDiscoveryConfig { + enabled: false, + selector: HashMap::new(), + check_interval: Duration::from_secs(60), + port: 8000, // Standard port for modern services + namespace: None, // None means watch all namespaces + pd_mode: false, + prefill_selector: HashMap::new(), + decode_selector: HashMap::new(), + } + } +} + +/// Pod type for PD mode service discovery +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub enum PodType { + Prefill, + Decode, + Regular, +} + +/// Represents a Kubernetes pod's information used for worker management +#[derive(Debug, Clone, PartialEq, Eq, Hash)] +pub struct PodInfo { + pub name: String, + pub ip: IpAddr, + pub status: String, + pub is_ready: bool, + pub pod_type: Option, +} + +impl PodInfo { + /// Check if a pod matches any of the given selectors + fn matches_selector(pod: &Pod, selector: &HashMap) -> bool { + if selector.is_empty() { + return false; + } + + pod.metadata + .labels + .as_ref() + .is_some_and(|labels| selector.iter().all(|(k, v)| labels.get(k) == Some(v))) + } + + /// Check if a pod should be included in service discovery + pub fn should_include(pod: &Pod, config: &ServiceDiscoveryConfig) -> bool { + if config.pd_mode { + // In PD mode, at least one selector must be non-empty + if config.prefill_selector.is_empty() && config.decode_selector.is_empty() { + warn!("PD mode enabled but both prefill_selector and decode_selector are empty"); + return false; + } + // In PD mode, pod must match either prefill or decode selector + Self::matches_selector(pod, &config.prefill_selector) + || Self::matches_selector(pod, &config.decode_selector) + } else { + // In regular mode, pod must match the general selector + if config.selector.is_empty() { + warn!("Regular mode enabled but selector is empty"); + return false; + } + Self::matches_selector(pod, &config.selector) + } + } + + /// Unified PodInfo creation with optional PD configuration + pub fn from_pod(pod: &Pod, config: Option<&ServiceDiscoveryConfig>) -> Option { + let name = pod.metadata.name.clone()?; + let status = pod.status.clone()?; + let pod_ip: IpAddr = status.pod_ip?.parse().ok()?; + + let is_ready = if let Some(conditions) = &status.conditions { + conditions + .iter() + .any(|condition| condition.type_ == "Ready" && condition.status == "True") + } else { + false + }; + + let pod_status = status.phase.unwrap_or_else(|| "Unknown".to_string()); + + // Determine pod type based on labels if config is provided and in PD mode + let pod_type = if let Some(config) = config { + if config.pd_mode { + // Use simplified helper methods for cleaner logic + if Self::matches_selector(pod, &config.prefill_selector) { + Some(PodType::Prefill) + } else if Self::matches_selector(pod, &config.decode_selector) { + Some(PodType::Decode) + } else { + Some(PodType::Regular) + } + } else { + Some(PodType::Regular) + } + } else { + // No config provided, default to None (for backwards compatibility) + None + }; + + Some(PodInfo { + name, + ip: pod_ip, + status: pod_status, + is_ready, + pod_type, + }) + } + + /// Returns true if the pod is in a state where it can accept traffic + pub fn is_healthy(&self) -> bool { + self.is_ready && self.status == "Running" + } + + /// Generates a worker URL for this pod. + pub fn worker_url(&self, port: u16) -> String { + format!("http://{}", SocketAddr::new(self.ip, port)) + } +} + +pub async fn start_service_discovery( + config: ServiceDiscoveryConfig, + router: Arc, +) -> Result, kube::Error> { + // Don't initialize anything if service discovery is disabled + if !config.enabled { + // Return a generic error when service discovery is disabled + return Err(kube::Error::Api(kube::error::ErrorResponse { + status: "Disabled".to_string(), + message: "Service discovery is disabled".to_string(), + reason: "ConfigurationError".to_string(), + code: 400, + })); + } + + let _ = rustls::crypto::ring::default_provider().install_default(); + + // Initialize Kubernetes client + let client = Client::try_default().await?; + + // Log the appropriate selectors based on mode + if config.pd_mode { + let prefill_selector = config + .prefill_selector + .iter() + .map(|(k, v)| format!("{}={}", k, v)) + .collect::>() + .join(","); + + let decode_selector = config + .decode_selector + .iter() + .map(|(k, v)| format!("{}={}", k, v)) + .collect::>() + .join(","); + + info!( + "Starting K8s service discovery | PD mode | prefill: '{}' | decode: '{}'", + prefill_selector, decode_selector + ); + } else { + let label_selector = config + .selector + .iter() + .map(|(k, v)| format!("{}={}", k, v)) + .collect::>() + .join(","); + + info!( + "Starting K8s service discovery | selector: '{}'", + label_selector + ); + } + + // Create the task that will run in the background + let handle = task::spawn(async move { + // We'll track pods we've already added to avoid duplicates + let tracked_pods = Arc::new(Mutex::new(HashSet::new())); + + // Create a watcher for pods + let pods: Api = if let Some(namespace) = &config.namespace { + Api::namespaced(client, namespace) + } else { + Api::all(client) + }; + + debug!("K8s service discovery initialized"); + + // Create Arcs for configuration data + let config_arc = Arc::new(config.clone()); + let port = config.port; + + let mut retry_delay = Duration::from_secs(1); + const MAX_RETRY_DELAY: Duration = Duration::from_secs(300); // 5 minutes max + + loop { + // Create a watcher with the proper parameters according to the kube-rs API + let watcher_config = Config::default(); + let watcher_stream = watcher(pods.clone(), watcher_config).applied_objects(); + + // Clone Arcs for the closures + let config_clone = Arc::clone(&config_arc); + let tracked_pods_clone = Arc::clone(&tracked_pods); + + // Simplified label selector filter using helper method + let filtered_stream = watcher_stream.filter_map(move |obj_res| { + let config_inner = Arc::clone(&config_clone); + + async move { + match obj_res { + Ok(pod) => { + if PodInfo::should_include(&pod, &config_inner) { + Some(Ok(pod)) + } else { + None + } + } + Err(e) => Some(Err(e)), + } + } + }); + + // Clone again for the next closure + let tracked_pods_clone2 = Arc::clone(&tracked_pods_clone); + let router_clone = Arc::clone(&router); + let config_clone2 = Arc::clone(&config_arc); + + match filtered_stream + .try_for_each(move |pod| { + let tracked_pods_inner = Arc::clone(&tracked_pods_clone2); + let router_inner = Arc::clone(&router_clone); + let config_inner = Arc::clone(&config_clone2); + + async move { + let pod_info = PodInfo::from_pod(&pod, Some(&config_inner)); + + if let Some(pod_info) = pod_info { + if pod.metadata.deletion_timestamp.is_some() { + handle_pod_deletion( + &pod_info, + tracked_pods_inner, + router_inner, + port, + config_inner.pd_mode, + ) + .await; + } else { + handle_pod_event( + &pod_info, + tracked_pods_inner, + router_inner, + port, + config_inner.pd_mode, + ) + .await; + } + } + Ok(()) + } + }) + .await + { + Ok(_) => { + // Reset retry delay on success + retry_delay = Duration::from_secs(1); + } + Err(err) => { + error!("Error in Kubernetes watcher: {}", err); + warn!( + "Retrying in {} seconds with exponential backoff", + retry_delay.as_secs() + ); + time::sleep(retry_delay).await; + + // Exponential backoff with jitter + retry_delay = std::cmp::min(retry_delay * 2, MAX_RETRY_DELAY); + } + } + + // If the watcher exits for some reason, wait a bit before restarting + warn!( + "Kubernetes watcher exited, restarting in {} seconds", + config_arc.check_interval.as_secs() + ); + time::sleep(config_arc.check_interval).await; + } + }); + + Ok(handle) +} + +async fn handle_pod_event( + pod_info: &PodInfo, + tracked_pods: Arc>>, + router: Arc, + port: u16, + pd_mode: bool, +) { + let worker_url = pod_info.worker_url(port); + + // If pod is healthy, try to add it (with atomic check-and-insert) + if pod_info.is_healthy() { + // Atomic check-and-insert to prevent race conditions + let should_add = { + let mut tracker = match tracked_pods.lock() { + Ok(tracker) => tracker, + Err(e) => { + error!("Failed to acquire tracked_pods lock: {}", e); + return; + } + }; + + if tracker.contains(pod_info) { + false // Already tracked + } else { + // Reserve the spot to prevent other threads from adding the same pod + tracker.insert(pod_info.clone()); + true + } + }; + + if should_add { + info!( + "Adding pod: {} | type: {:?} | url: {}", + pod_info.name, pod_info.pod_type, worker_url + ); + + // Handle PD mode with specific pod types + let result = if pd_mode && pod_info.pod_type.is_some() { + use crate::routers::http::lmdeploy_pd_router::LMDeployPDRouter; + + if let Some(lmdeploy_pd_router) = router.as_any().downcast_ref::() + { + match &pod_info.pod_type { + Some(PodType::Prefill) => lmdeploy_pd_router + .add_prefill_server(worker_url.clone()) + .await + .map_err(|e| e.to_string()), + Some(PodType::Decode) => lmdeploy_pd_router + .add_decode_server(worker_url.clone()) + .await + .map_err(|e| e.to_string()), + Some(PodType::Regular) | None => router.add_worker(&worker_url).await, + } + } else { + Err("PD mode enabled but router is not a supported PD router".to_string()) + } + } else { + // Regular mode or no pod type specified + router.add_worker(&worker_url).await + }; + + match result { + Ok(_) => { + debug!("Worker added: {}", worker_url); + } + Err(e) => { + error!("Failed to add worker {} to router: {}", worker_url, e); + // Remove from tracking since addition failed + if let Ok(mut tracker) = tracked_pods.lock() { + tracker.remove(pod_info); + } + } + } + } + } +} + +async fn handle_pod_deletion( + pod_info: &PodInfo, + tracked_pods: Arc>>, + router: Arc, + port: u16, + pd_mode: bool, +) { + let worker_url = pod_info.worker_url(port); + + let was_tracked = { + let mut tracked = match tracked_pods.lock() { + Ok(tracked) => tracked, + Err(e) => { + error!("Failed to acquire tracked_pods lock during deletion: {}", e); + return; + } + }; + tracked.remove(pod_info) + }; + + if was_tracked { + info!( + "Removing pod: {} | type: {:?} | url: {}", + pod_info.name, pod_info.pod_type, worker_url + ); + + // Handle PD mode removal + if pd_mode && pod_info.pod_type.is_some() { + use crate::routers::http::lmdeploy_pd_router::LMDeployPDRouter; + + if let Some(lmdeploy_pd_router) = router.as_any().downcast_ref::() { + match &pod_info.pod_type { + Some(PodType::Prefill) => { + if let Err(e) = lmdeploy_pd_router.remove_prefill_server(&worker_url).await + { + error!( + "Failed to remove lmdeploy prefill server {}: {}", + worker_url, e + ); + } + } + Some(PodType::Decode) => { + if let Err(e) = lmdeploy_pd_router.remove_decode_server(&worker_url).await { + error!( + "Failed to remove lmdeploy decode server {}: {}", + worker_url, e + ); + } + } + Some(PodType::Regular) | None => router.remove_worker(&worker_url), + } + } else { + // Unknown PD implementation: fall back to generic removal. + router.remove_worker(&worker_url); + } + } else { + // Regular mode removal + router.remove_worker(&worker_url); + } + } else { + // This case might occur if a pod is deleted before it was ever marked healthy and added. + // Or if the event is duplicated. No action needed on the router if it wasn't tracked (and thus not added). + debug!( + "Pod deletion event for untracked/already removed pod: {} (type: {:?}). Worker URL: {}", + pod_info.name, pod_info.pod_type, worker_url + ); + } +} + +#[cfg(test)] +mod tests { + use super::*; + use k8s_openapi::api::core::v1::{Pod, PodCondition, PodSpec, PodStatus}; + use k8s_openapi::apimachinery::pkg::apis::meta::v1::ObjectMeta; + use k8s_openapi::apimachinery::pkg::apis::meta::v1::Time; + + // Helper function to create a Pod for testing PodInfo::from_pod + fn create_k8s_pod( + name: Option<&str>, + ip: Option<&str>, + phase: Option<&str>, + ready_status: Option<&str>, + deletion_timestamp: Option")) +} + +/// Factory function to create tokenizer from a model name or path (async version) +pub async fn create_tokenizer_async( + model_name_or_path: &str, +) -> Result> { + // Check if it's a file path + let path = Path::new(model_name_or_path); + if path.exists() { + return create_tokenizer_from_file(model_name_or_path); + } + + // Check if it's a GPT model name that should use Tiktoken + if model_name_or_path.contains("gpt-") + || model_name_or_path.contains("davinci") + || model_name_or_path.contains("curie") + || model_name_or_path.contains("babbage") + || model_name_or_path.contains("ada") + { + let tokenizer = TiktokenTokenizer::from_model_name(model_name_or_path)?; + return Ok(Arc::new(tokenizer)); + } + + // Try to download tokenizer files from HuggingFace + match download_tokenizer_from_hf(model_name_or_path).await { + Ok(cache_dir) => { + // Look for tokenizer.json in the cache directory + let tokenizer_path = cache_dir.join("tokenizer.json"); + if tokenizer_path.exists() { + create_tokenizer_from_file(tokenizer_path.to_str().unwrap()) + } else { + // Try other common tokenizer file names + let possible_files = ["tokenizer_config.json", "vocab.json"]; + for file_name in &possible_files { + let file_path = cache_dir.join(file_name); + if file_path.exists() { + return create_tokenizer_from_file(file_path.to_str().unwrap()); + } + } + Err(Error::msg(format!( + "Downloaded model '{}' but couldn't find a suitable tokenizer file", + model_name_or_path + ))) + } + } + Err(e) => Err(Error::msg(format!( + "Failed to download tokenizer from HuggingFace: {}", + e + ))), + } +} + +/// Factory function to create tokenizer from a model name or path (blocking version) +pub fn create_tokenizer(model_name_or_path: &str) -> Result> { + // Check if it's a file path + let path = Path::new(model_name_or_path); + if path.exists() { + return create_tokenizer_from_file(model_name_or_path); + } + + // Check if it's a GPT model name that should use Tiktoken + if model_name_or_path.contains("gpt-") + || model_name_or_path.contains("davinci") + || model_name_or_path.contains("curie") + || model_name_or_path.contains("babbage") + || model_name_or_path.contains("ada") + { + let tokenizer = TiktokenTokenizer::from_model_name(model_name_or_path)?; + return Ok(Arc::new(tokenizer)); + } + + // Only use tokio for HuggingFace downloads + // Check if we're already in a tokio runtime + if let Ok(handle) = tokio::runtime::Handle::try_current() { + // We're in a runtime, use block_in_place + tokio::task::block_in_place(|| handle.block_on(create_tokenizer_async(model_name_or_path))) + } else { + // No runtime, create a temporary one + let rt = tokio::runtime::Runtime::new()?; + rt.block_on(create_tokenizer_async(model_name_or_path)) + } +} + +/// Get information about a tokenizer file +pub fn get_tokenizer_info(file_path: &str) -> Result { + let path = Path::new(file_path); + + if !path.exists() { + return Err(Error::msg(format!("File not found: {}", file_path))); + } + + let extension = path + .extension() + .and_then(std::ffi::OsStr::to_str) + .map(|s| s.to_lowercase()); + + match extension.as_deref() { + Some("json") => Ok(TokenizerType::HuggingFace(file_path.to_string())), + _ => { + // Try auto-detection + use std::fs::File; + use std::io::Read; + + let mut file = File::open(file_path)?; + let mut buffer = vec![0u8; 512]; + let bytes_read = file.read(&mut buffer)?; + buffer.truncate(bytes_read); + + if is_likely_json(&buffer) { + Ok(TokenizerType::HuggingFace(file_path.to_string())) + } else { + Err(Error::msg("Unknown tokenizer type")) + } + } + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_json_detection() { + assert!(is_likely_json(b"{\"test\": \"value\"}")); + assert!(is_likely_json(b" \n\t{\"test\": \"value\"}")); + assert!(is_likely_json(b"[1, 2, 3]")); + assert!(!is_likely_json(b"not json")); + assert!(!is_likely_json(b"")); + } + + #[test] + fn test_mock_tokenizer_creation() { + let tokenizer = create_tokenizer_from_file("mock").unwrap(); + assert_eq!(tokenizer.vocab_size(), 8); // Mock tokenizer has 8 tokens + } + + #[test] + fn test_file_not_found() { + let result = create_tokenizer_from_file("/nonexistent/file.json"); + assert!(result.is_err()); + if let Err(e) = result { + assert!(e.to_string().contains("File not found")); + } + } + + #[test] + fn test_create_tiktoken_tokenizer() { + // Test creating tokenizer for GPT models + let tokenizer = create_tokenizer("gpt-4").unwrap(); + assert!(tokenizer.vocab_size() > 0); + + // Test encoding and decoding + let text = "Hello, world!"; + let encoding = tokenizer.encode(text).unwrap(); + let decoded = tokenizer.decode(encoding.token_ids(), false).unwrap(); + assert_eq!(decoded, text); + } + + #[tokio::test] + async fn test_download_tokenizer_from_hf() { + // Test with a small model that should have tokenizer files + // Skip this test if HF_TOKEN is not set and we're in CI + if std::env::var("CI").is_ok() && std::env::var("HF_TOKEN").is_err() { + println!("Skipping HF download test in CI without HF_TOKEN"); + return; + } + + // Try to create tokenizer for a known small model + let result = create_tokenizer_async("bert-base-uncased").await; + + // The test might fail due to network issues or rate limiting + // so we just check that the function executes without panic + match result { + Ok(tokenizer) => { + assert!(tokenizer.vocab_size() > 0); + println!("Successfully downloaded and created tokenizer"); + } + Err(e) => { + println!("Download failed (this might be expected): {}", e); + // Don't fail the test - network issues shouldn't break CI + } + } + } +} diff --git a/lmdeploy/router/src/tokenizer/hub.rs b/lmdeploy/router/src/tokenizer/hub.rs new file mode 100644 index 0000000000..c9d2cd1a44 --- /dev/null +++ b/lmdeploy/router/src/tokenizer/hub.rs @@ -0,0 +1,238 @@ +use hf_hub::api::tokio::ApiBuilder; +use std::env; +use std::path::{Path, PathBuf}; + +const IGNORED: [&str; 5] = [ + ".gitattributes", + "LICENSE", + "LICENSE.txt", + "README.md", + "USE_POLICY.md", +]; + +const HF_TOKEN_ENV_VAR: &str = "HF_TOKEN"; + +/// Checks if a file is a model weight file +fn is_weight_file(filename: &str) -> bool { + filename.ends_with(".bin") + || filename.ends_with(".safetensors") + || filename.ends_with(".h5") + || filename.ends_with(".msgpack") + || filename.ends_with(".ckpt.index") +} + +/// Checks if a file is an image file +fn is_image(filename: &str) -> bool { + filename.ends_with(".png") + || filename.ends_with("PNG") + || filename.ends_with(".jpg") + || filename.ends_with("JPG") + || filename.ends_with(".jpeg") + || filename.ends_with("JPEG") +} + +/// Checks if a file is a tokenizer file +fn is_tokenizer_file(filename: &str) -> bool { + filename.ends_with("tokenizer.json") + || filename.ends_with("tokenizer_config.json") + || filename.ends_with("special_tokens_map.json") + || filename.ends_with("vocab.json") + || filename.ends_with("merges.txt") + || filename.ends_with(".model") // SentencePiece models + || filename.ends_with(".tiktoken") +} + +/// Attempt to download tokenizer files from Hugging Face +/// Returns the directory containing the downloaded tokenizer files +pub async fn download_tokenizer_from_hf(model_id: impl AsRef) -> anyhow::Result { + let model_id = model_id.as_ref(); + let token = env::var(HF_TOKEN_ENV_VAR).ok(); + let api = ApiBuilder::new() + .with_progress(true) + .with_token(token) + .build()?; + let model_name = model_id.display().to_string(); + + let repo = api.model(model_name.clone()); + + let info = match repo.info().await { + Ok(info) => info, + Err(e) => { + return Err(anyhow::anyhow!( + "Failed to fetch model '{}' from HuggingFace: {}. Is this a valid HuggingFace ID?", + model_name, + e + )); + } + }; + + if info.siblings.is_empty() { + return Err(anyhow::anyhow!( + "Model '{}' exists but contains no downloadable files.", + model_name + )); + } + + let mut cache_dir = None; + let mut tokenizer_files_found = false; + + // First, identify all tokenizer files to download + let tokenizer_files: Vec<_> = info + .siblings + .iter() + .filter(|sib| { + !IGNORED.contains(&sib.rfilename.as_str()) + && !is_image(&sib.rfilename) + && !is_weight_file(&sib.rfilename) + && is_tokenizer_file(&sib.rfilename) + }) + .collect(); + + if tokenizer_files.is_empty() { + return Err(anyhow::anyhow!( + "No tokenizer files found for model '{}'.", + model_name + )); + } + + // Download all tokenizer files + for sib in tokenizer_files { + match repo.get(&sib.rfilename).await { + Ok(path) => { + if cache_dir.is_none() { + cache_dir = path.parent().map(|p| p.to_path_buf()); + } + tokenizer_files_found = true; + } + Err(e) => { + return Err(anyhow::anyhow!( + "Failed to download tokenizer file '{}' from model '{}': {}", + sib.rfilename, + model_name, + e + )); + } + } + } + + if !tokenizer_files_found { + return Err(anyhow::anyhow!( + "No tokenizer files could be downloaded for model '{}'.", + model_name + )); + } + + match cache_dir { + Some(dir) => Ok(dir), + None => Err(anyhow::anyhow!( + "Invalid HF cache path for model '{}'", + model_name + )), + } +} + +/// Attempt to download a model from Hugging Face (including weights) +/// Returns the directory it is in +/// If ignore_weights is true, model weight files will be skipped +pub async fn from_hf(name: impl AsRef, ignore_weights: bool) -> anyhow::Result { + let name = name.as_ref(); + let token = env::var(HF_TOKEN_ENV_VAR).ok(); + let api = ApiBuilder::new() + .with_progress(true) + .with_token(token) + .build()?; + let model_name = name.display().to_string(); + + let repo = api.model(model_name.clone()); + + let info = match repo.info().await { + Ok(info) => info, + Err(e) => { + return Err(anyhow::anyhow!( + "Failed to fetch model '{}' from HuggingFace: {}. Is this a valid HuggingFace ID?", + model_name, + e + )); + } + }; + + if info.siblings.is_empty() { + return Err(anyhow::anyhow!( + "Model '{}' exists but contains no downloadable files.", + model_name + )); + } + + let mut p = PathBuf::new(); + let mut files_downloaded = false; + + for sib in info.siblings { + if IGNORED.contains(&sib.rfilename.as_str()) || is_image(&sib.rfilename) { + continue; + } + + // If ignore_weights is true, skip weight files + if ignore_weights && is_weight_file(&sib.rfilename) { + continue; + } + + match repo.get(&sib.rfilename).await { + Ok(path) => { + p = path; + files_downloaded = true; + } + Err(e) => { + return Err(anyhow::anyhow!( + "Failed to download file '{}' from model '{}': {}", + sib.rfilename, + model_name, + e + )); + } + } + } + + if !files_downloaded { + let file_type = if ignore_weights { + "non-weight" + } else { + "valid" + }; + return Err(anyhow::anyhow!( + "No {} files found for model '{}'.", + file_type, + model_name + )); + } + + match p.parent() { + Some(p) => Ok(p.to_path_buf()), + None => Err(anyhow::anyhow!("Invalid HF cache path: {}", p.display())), + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_is_tokenizer_file() { + assert!(is_tokenizer_file("tokenizer.json")); + assert!(is_tokenizer_file("tokenizer_config.json")); + assert!(is_tokenizer_file("special_tokens_map.json")); + assert!(is_tokenizer_file("vocab.json")); + assert!(is_tokenizer_file("merges.txt")); + assert!(is_tokenizer_file("spiece.model")); + assert!(!is_tokenizer_file("model.bin")); + assert!(!is_tokenizer_file("README.md")); + } + + #[test] + fn test_is_weight_file() { + assert!(is_weight_file("model.bin")); + assert!(is_weight_file("model.safetensors")); + assert!(is_weight_file("pytorch_model.bin")); + assert!(!is_weight_file("tokenizer.json")); + assert!(!is_weight_file("config.json")); + } +} diff --git a/lmdeploy/router/src/tokenizer/huggingface.rs b/lmdeploy/router/src/tokenizer/huggingface.rs new file mode 100644 index 0000000000..02dce5a0ae --- /dev/null +++ b/lmdeploy/router/src/tokenizer/huggingface.rs @@ -0,0 +1,234 @@ +use super::traits::{ + Decoder, Encoder, Encoding, SpecialTokens, TokenIdType, Tokenizer as TokenizerTrait, +}; +use anyhow::{Error, Result}; +use std::collections::HashMap; +use tokenizers::tokenizer::Tokenizer as HfTokenizer; + +use super::chat_template::{ChatMessage, ChatTemplateProcessor}; + +/// HuggingFace tokenizer wrapper +pub struct HuggingFaceTokenizer { + tokenizer: HfTokenizer, + special_tokens: SpecialTokens, + vocab: HashMap, + reverse_vocab: HashMap, + chat_template: Option, +} + +impl HuggingFaceTokenizer { + /// Create a tokenizer from a HuggingFace tokenizer JSON file + pub fn from_file(file_path: &str) -> Result { + Self::from_file_with_chat_template(file_path, None) + } + + /// Create a tokenizer from a HuggingFace tokenizer JSON file with an optional chat template + pub fn from_file_with_chat_template( + file_path: &str, + chat_template_path: Option<&str>, + ) -> Result { + let tokenizer = HfTokenizer::from_file(file_path) + .map_err(|e| Error::msg(format!("Failed to load tokenizer: {}", e)))?; + + // Extract special tokens + let special_tokens = Self::extract_special_tokens(&tokenizer); + + // Build vocab mappings + let vocab = tokenizer.get_vocab(false); + let reverse_vocab: HashMap = vocab + .iter() + .map(|(token, &id)| (id, token.clone())) + .collect(); + + // Load chat template + let chat_template = if let Some(template_path) = chat_template_path { + // Load from specified .jinja file + Self::load_chat_template_from_file(template_path)? + } else { + // Try to load from tokenizer_config.json + Self::load_chat_template(file_path) + }; + + Ok(HuggingFaceTokenizer { + tokenizer, + special_tokens, + vocab, + reverse_vocab, + chat_template, + }) + } + + /// Create from an existing HuggingFace tokenizer + pub fn from_tokenizer(tokenizer: HfTokenizer) -> Self { + let special_tokens = Self::extract_special_tokens(&tokenizer); + let vocab = tokenizer.get_vocab(false); + let reverse_vocab: HashMap = vocab + .iter() + .map(|(token, &id)| (id, token.clone())) + .collect(); + + HuggingFaceTokenizer { + tokenizer, + special_tokens, + vocab, + reverse_vocab, + chat_template: None, + } + } + + /// Extract special tokens from the tokenizer + fn extract_special_tokens(tokenizer: &HfTokenizer) -> SpecialTokens { + // Try to get special tokens from the tokenizer + // This is a simplified version - actual implementation would need to handle various formats + let vocab = tokenizer.get_vocab(true); + + let find_token = |patterns: &[&str]| -> Option { + for pattern in patterns { + if vocab.contains_key(*pattern) { + return Some(pattern.to_string()); + } + } + None + }; + + SpecialTokens { + bos_token: find_token(&["", "<|startoftext|>", "", "[CLS]"]), + eos_token: find_token(&["", "<|endoftext|>", "", "[SEP]"]), + unk_token: find_token(&["", "", "[UNK]"]), + sep_token: find_token(&["[SEP]", "", ""]), + pad_token: find_token(&["", "", "[PAD]"]), + cls_token: find_token(&["[CLS]", "", ""]), + mask_token: find_token(&["[MASK]", "", ""]), + additional_special_tokens: vec![], + } + } + + /// Try to load chat template from tokenizer_config.json + fn load_chat_template(tokenizer_path: &str) -> Option { + // Try to find tokenizer_config.json in the same directory + let path = std::path::Path::new(tokenizer_path); + let dir = path.parent()?; + let config_path = dir.join("tokenizer_config.json"); + + if config_path.exists() { + if let Ok(template) = + super::chat_template::load_chat_template_from_config(config_path.to_str()?) + { + return template; + } + } + None + } + + /// Load chat template from a .jinja file + fn load_chat_template_from_file(template_path: &str) -> Result> { + use std::fs; + + let content = fs::read_to_string(template_path) + .map_err(|e| Error::msg(format!("Failed to read chat template file: {}", e)))?; + + // Clean up the template (similar to Python implementation) + let template = content.trim().replace("\\n", "\n"); + + Ok(Some(template)) + } + + /// Set or override the chat template + pub fn set_chat_template(&mut self, template: String) { + self.chat_template = Some(template); + } + + /// Apply chat template if available + pub fn apply_chat_template( + &self, + messages: &[ChatMessage], + add_generation_prompt: bool, + ) -> Result { + if let Some(ref template) = self.chat_template { + let processor = ChatTemplateProcessor::new( + template.clone(), + self.special_tokens.bos_token.clone(), + self.special_tokens.eos_token.clone(), + ); + processor.apply_chat_template(messages, add_generation_prompt) + } else { + // Fallback to simple formatting if no template is available + let mut result = String::new(); + for msg in messages { + result.push_str(&format!("{}: {}\n", msg.role, msg.content)); + } + if add_generation_prompt { + result.push_str("assistant: "); + } + Ok(result) + } + } +} + +impl Encoder for HuggingFaceTokenizer { + fn encode(&self, input: &str) -> Result { + self.tokenizer + .encode(input, false) + .map_err(|e| Error::msg(format!("Encoding failed: {}", e))) + .map(|encoding| Encoding::Hf(Box::new(encoding))) + } + + fn encode_batch(&self, inputs: &[&str]) -> Result> { + let encodings = self + .tokenizer + .encode_batch(inputs.to_vec(), false) + .map_err(|e| Error::msg(format!("Batch encoding failed: {}", e)))?; + + Ok(encodings + .into_iter() + .map(|e| Encoding::Hf(Box::new(e))) + .collect()) + } +} + +impl Decoder for HuggingFaceTokenizer { + fn decode(&self, token_ids: &[TokenIdType], skip_special_tokens: bool) -> Result { + self.tokenizer + .decode(token_ids, skip_special_tokens) + .map_err(|e| Error::msg(format!("Decoding failed: {}", e))) + } +} + +impl TokenizerTrait for HuggingFaceTokenizer { + fn vocab_size(&self) -> usize { + self.tokenizer.get_vocab_size(false) + } + + fn get_special_tokens(&self) -> &SpecialTokens { + &self.special_tokens + } + + fn token_to_id(&self, token: &str) -> Option { + self.vocab.get(token).copied() + } + + fn id_to_token(&self, id: TokenIdType) -> Option { + self.reverse_vocab.get(&id).cloned() + } +} + +#[cfg(test)] +mod tests { + use super::ChatMessage; + + #[test] + fn test_chat_message_creation() { + let msg = ChatMessage::system("You are a helpful assistant"); + assert_eq!(msg.role, "system"); + assert_eq!(msg.content, "You are a helpful assistant"); + + let user_msg = ChatMessage::user("Hello!"); + assert_eq!(user_msg.role, "user"); + + let assistant_msg = ChatMessage::assistant("Hi there!"); + assert_eq!(assistant_msg.role, "assistant"); + } + + // Note: Actual tokenizer tests would require a real tokenizer file + // These would be integration tests rather than unit tests +} diff --git a/lmdeploy/router/src/tokenizer/mock.rs b/lmdeploy/router/src/tokenizer/mock.rs new file mode 100644 index 0000000000..afb91543c6 --- /dev/null +++ b/lmdeploy/router/src/tokenizer/mock.rs @@ -0,0 +1,112 @@ +//! Mock tokenizer implementation for testing + +use super::traits::{Decoder, Encoder, Encoding, SpecialTokens, Tokenizer as TokenizerTrait}; +use anyhow::Result; +use std::collections::HashMap; + +/// Mock tokenizer for testing purposes +pub struct MockTokenizer { + vocab: HashMap, + reverse_vocab: HashMap, + special_tokens: SpecialTokens, +} + +impl Default for MockTokenizer { + fn default() -> Self { + Self::new() + } +} + +impl MockTokenizer { + pub fn new() -> Self { + let mut vocab = HashMap::new(); + let mut reverse_vocab = HashMap::new(); + + // Add some basic tokens + let tokens = vec![ + ("Hello", 1), + ("world", 2), + ("test", 3), + ("token", 4), + (" ", 5), + (".", 6), + ("", 999), + ("", 1000), + ]; + + for (token, id) in tokens { + vocab.insert(token.to_string(), id); + reverse_vocab.insert(id, token.to_string()); + } + + let special_tokens = SpecialTokens { + bos_token: Some("".to_string()), + eos_token: Some("".to_string()), + unk_token: Some("".to_string()), + sep_token: None, + pad_token: None, + cls_token: None, + mask_token: None, + additional_special_tokens: vec![], + }; + + Self { + vocab, + reverse_vocab, + special_tokens, + } + } +} + +impl Encoder for MockTokenizer { + fn encode(&self, input: &str) -> Result { + // Simple word-based tokenization for testing + let tokens: Vec = input + .split_whitespace() + .filter_map(|word| self.vocab.get(word).copied()) + .collect(); + + Ok(Encoding::Sp(tokens)) + } + + fn encode_batch(&self, inputs: &[&str]) -> Result> { + inputs.iter().map(|input| self.encode(input)).collect() + } +} + +impl Decoder for MockTokenizer { + fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result { + let tokens: Vec = token_ids + .iter() + .filter_map(|id| { + self.reverse_vocab.get(id).and_then(|token| { + if skip_special_tokens && (token == "" || token == "") { + None + } else { + Some(token.clone()) + } + }) + }) + .collect(); + + Ok(tokens.join(" ")) + } +} + +impl TokenizerTrait for MockTokenizer { + fn vocab_size(&self) -> usize { + self.vocab.len() + } + + fn get_special_tokens(&self) -> &SpecialTokens { + &self.special_tokens + } + + fn token_to_id(&self, token: &str) -> Option { + self.vocab.get(token).copied() + } + + fn id_to_token(&self, id: u32) -> Option { + self.reverse_vocab.get(&id).cloned() + } +} diff --git a/lmdeploy/router/src/tokenizer/mod.rs b/lmdeploy/router/src/tokenizer/mod.rs new file mode 100644 index 0000000000..98a23f7617 --- /dev/null +++ b/lmdeploy/router/src/tokenizer/mod.rs @@ -0,0 +1,123 @@ +use anyhow::Result; +use std::ops::Deref; +use std::sync::Arc; + +pub mod factory; +pub mod hub; +pub mod mock; +pub mod sequence; +pub mod stop; +pub mod stream; +pub mod traits; + +// Feature-gated modules + +pub mod chat_template; + +pub mod huggingface; + +pub mod tiktoken; + +#[cfg(test)] +mod tests; + +// Re-exports +pub use factory::{ + create_tokenizer, create_tokenizer_async, create_tokenizer_from_file, + create_tokenizer_with_chat_template, TokenizerType, +}; +pub use sequence::Sequence; +pub use stop::{SequenceDecoderOutput, StopSequenceConfig, StopSequenceDecoder}; +pub use stream::DecodeStream; +pub use traits::{Decoder, Encoder, Encoding, SpecialTokens, Tokenizer as TokenizerTrait}; + +pub use huggingface::HuggingFaceTokenizer; + +pub use chat_template::ChatMessage; + +pub use tiktoken::{TiktokenModel, TiktokenTokenizer}; + +/// Main tokenizer wrapper that provides a unified interface for different tokenizer implementations +#[derive(Clone)] +pub struct Tokenizer(Arc); + +impl Tokenizer { + /// Create a tokenizer from a file path + pub fn from_file(file_path: &str) -> Result { + Ok(Tokenizer(factory::create_tokenizer_from_file(file_path)?)) + } + + /// Create a tokenizer from a file path with an optional chat template + pub fn from_file_with_chat_template( + file_path: &str, + chat_template_path: Option<&str>, + ) -> Result { + Ok(Tokenizer(factory::create_tokenizer_with_chat_template( + file_path, + chat_template_path, + )?)) + } + + /// Create a tokenizer from an Arc + pub fn from_arc(tokenizer: Arc) -> Self { + Tokenizer(tokenizer) + } + + /// Create a stateful sequence object for decoding token_ids into text + pub fn decode_stream( + &self, + prompt_token_ids: &[u32], + skip_special_tokens: bool, + ) -> DecodeStream { + DecodeStream::new(self.0.clone(), prompt_token_ids, skip_special_tokens) + } + + /// Direct encode method + pub fn encode(&self, input: &str) -> Result { + self.0.encode(input) + } + + /// Direct batch encode method + pub fn encode_batch(&self, inputs: &[&str]) -> Result> { + self.0.encode_batch(inputs) + } + + /// Direct decode method + pub fn decode(&self, token_ids: &[u32], skip_special_tokens: bool) -> Result { + self.0.decode(token_ids, skip_special_tokens) + } + + /// Get vocabulary size + pub fn vocab_size(&self) -> usize { + self.0.vocab_size() + } + + /// Get special tokens + pub fn get_special_tokens(&self) -> &SpecialTokens { + self.0.get_special_tokens() + } + + /// Convert token string to ID + pub fn token_to_id(&self, token: &str) -> Option { + self.0.token_to_id(token) + } + + /// Convert ID to token string + pub fn id_to_token(&self, id: u32) -> Option { + self.0.id_to_token(id) + } +} + +impl Deref for Tokenizer { + type Target = Arc; + + fn deref(&self) -> &Self::Target { + &self.0 + } +} + +impl From> for Tokenizer { + fn from(tokenizer: Arc) -> Self { + Tokenizer(tokenizer) + } +} diff --git a/lmdeploy/router/src/tokenizer/sequence.rs b/lmdeploy/router/src/tokenizer/sequence.rs new file mode 100644 index 0000000000..99801438d0 --- /dev/null +++ b/lmdeploy/router/src/tokenizer/sequence.rs @@ -0,0 +1,238 @@ +use super::traits::{TokenIdType, Tokenizer as TokenizerTrait}; +use anyhow::Result; +use std::sync::Arc; + +/// Maintains state for an ongoing sequence of tokens and their decoded text +/// This provides a cleaner abstraction for managing token sequences +pub struct Sequence { + /// The tokenizer used for encoding/decoding + tokenizer: Arc, + + /// The current sequence of token ids + token_ids: Vec, + + /// The position in the current sequence the last decoded token completed + prefix_offset: usize, + + /// Current position in the sequence + read_offset: usize, +} + +impl std::fmt::Debug for Sequence { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("Sequence") + .field("tokenizer", &"Arc") + .field( + "token_ids", + &format_args!("{}", { + let token_ids = self.token_ids(); + if token_ids.len() <= 20 { + format!("{:?}", token_ids) + } else { + let first_ten = &token_ids[..10]; + let last_ten = &token_ids[token_ids.len() - 10..]; + format!("{:?} ... {:?}", first_ten, last_ten) + } + }), + ) + .field("prefix_offset", &self.prefix_offset) + .field("read_offset", &self.read_offset) + .field("token count", &self.token_ids.len()) + .finish() + } +} + +impl Sequence { + /// Create a new empty sequence + pub fn new(tokenizer: Arc) -> Self { + Self { + tokenizer, + token_ids: Vec::new(), + prefix_offset: 0, + read_offset: 0, + } + } + + /// Create a sequence with initial tokens + pub fn with_tokens(tokenizer: Arc, token_ids: Vec) -> Self { + let len = token_ids.len(); + Self { + tokenizer, + token_ids, + prefix_offset: 0, + read_offset: len, + } + } + + /// Check if the sequence is empty + pub fn is_empty(&self) -> bool { + self.token_ids.is_empty() + } + + /// Get the length of the sequence + pub fn len(&self) -> usize { + self.token_ids.len() + } + + /// Clear the sequence + pub fn clear(&mut self) { + self.token_ids.clear(); + self.prefix_offset = 0; + self.read_offset = 0; + } + + /// Append text to the sequence by encoding it + pub fn append_text(&mut self, input: &str) -> Result<()> { + let encoding = self.tokenizer.encode(input)?; + self.token_ids.extend(encoding.token_ids()); + Ok(()) + } + + /// Append a single token to the sequence and return newly decoded text + /// Based on HuggingFace TGI incremental decoding + pub fn append_token(&mut self, token_id: TokenIdType) -> Result { + // Store the old read offset before adding the new token + let old_read_offset = self.read_offset; + + self.token_ids.push(token_id); + self.read_offset = self.token_ids.len(); + + // If this is the first token or we're at the beginning, decode everything + if self.prefix_offset == 0 && old_read_offset == 0 { + let text = self.tokenizer.decode(&self.token_ids, false)?; + if text.ends_with("�") { + // Incomplete UTF-8 sequence, wait for more tokens + return Ok(String::new()); + } + self.prefix_offset = 0; + return Ok(text); + } + + // Decode the text up to the previous position + let prefix_text = self + .tokenizer + .decode(&self.token_ids[self.prefix_offset..old_read_offset], false)?; + + // Decode the text including the new token + let new_text = self + .tokenizer + .decode(&self.token_ids[self.prefix_offset..], false)?; + + // Handle multi-byte character boundaries + let mut prefix_text_len = prefix_text.len(); + while !new_text.is_char_boundary(prefix_text_len) && prefix_text_len > 0 { + prefix_text_len -= 1; + } + + if new_text.len() > prefix_text.len() { + if new_text.ends_with("�") { + // Incomplete UTF-8 sequence, wait for more tokens + return Ok(String::new()); + } else { + // Return the new text portion + let incremental_text = new_text[prefix_text_len..].to_string().replace("�", ""); + self.prefix_offset = old_read_offset; + return Ok(incremental_text); + } + } + + Ok(String::new()) + } + + /// Get a reference to the tokenizer + pub fn tokenizer(&self) -> &Arc { + &self.tokenizer + } + + /// Get the current token ids + pub fn token_ids(&self) -> &[TokenIdType] { + &self.token_ids + } + + /// Decode the entire sequence to text + pub fn text(&self) -> Result { + self.tokenizer.decode(&self.token_ids, false) + } + + /// Get the prefix offset + pub fn prefix_offset(&self) -> usize { + self.prefix_offset + } + + /// Get the read offset + pub fn read_offset(&self) -> usize { + self.read_offset + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::tokenizer::mock::MockTokenizer; + + #[test] + fn test_sequence_new() { + let tokenizer = Arc::new(MockTokenizer::new()); + let seq = Sequence::new(tokenizer); + assert!(seq.is_empty()); + assert_eq!(seq.len(), 0); + } + + #[test] + fn test_sequence_append_text() { + let tokenizer = Arc::new(MockTokenizer::new()); + let mut seq = Sequence::new(tokenizer); + + seq.append_text("Hello").unwrap(); + assert!(!seq.is_empty()); + assert!(!seq.is_empty()); + + let text = seq.text().unwrap(); + assert_eq!(text, "Hello"); + } + + #[test] + fn test_sequence_append_token() { + let tokenizer = Arc::new(MockTokenizer::new()); + let mut seq = Sequence::new(tokenizer.clone()); + + // Start with an empty sequence and append token 1 ("Hello") + let text1 = seq.append_token(1).unwrap(); + assert_eq!(text1, "Hello"); + + // Now append token 2 ("world") + // The mock tokenizer will decode [1, 2] as "Hello world" (with a space) + let text2 = seq.append_token(2).unwrap(); + // The incremental text should be " world" (with the space that the mock tokenizer adds) + assert_eq!(text2, " world"); + + // Verify the full text + assert_eq!(seq.text().unwrap(), "Hello world"); + } + + #[test] + fn test_sequence_clear() { + let tokenizer = Arc::new(MockTokenizer::new()); + let mut seq = Sequence::new(tokenizer); + + seq.append_text("Hello world").unwrap(); + assert!(!seq.is_empty()); + + seq.clear(); + assert!(seq.is_empty()); + assert_eq!(seq.len(), 0); + assert_eq!(seq.prefix_offset(), 0); + assert_eq!(seq.read_offset(), 0); + } + + #[test] + fn test_sequence_debug() { + let tokenizer = Arc::new(MockTokenizer::new()); + let mut seq = Sequence::new(tokenizer); + + seq.append_text("Test").unwrap(); + let debug_str = format!("{:?}", seq); + assert!(debug_str.contains("Sequence")); + assert!(debug_str.contains("token count")); + } +} diff --git a/lmdeploy/router/src/tokenizer/stop.rs b/lmdeploy/router/src/tokenizer/stop.rs new file mode 100644 index 0000000000..1efda15b64 --- /dev/null +++ b/lmdeploy/router/src/tokenizer/stop.rs @@ -0,0 +1,506 @@ +use super::traits::{self, TokenIdType}; +use anyhow::Result; +use std::collections::HashSet; +use std::sync::Arc; + +/// Output from the sequence decoder +#[derive(Debug, Clone, PartialEq)] +pub enum SequenceDecoderOutput { + /// Normal text output + Text(String), + /// Text is being held due to partial stop sequence match + Held, + /// Stop sequence matched (hidden - not included in output) + Stopped, + /// Stop sequence matched with text (visible - included in output) + StoppedWithText(String), +} + +/// Configuration for stop sequences +#[derive(Debug, Clone, Default)] +pub struct StopSequenceConfig { + /// Token IDs that trigger a stop + pub stop_tokens: HashSet, + /// String sequences that trigger a stop + pub stop_sequences: Vec, + /// Token IDs for visible stops (included in output) + pub visible_stop_tokens: HashSet, + /// String sequences for visible stops (included in output) + pub visible_stop_sequences: Vec, +} + +impl StopSequenceConfig { + /// Builder pattern - add a stop token + pub fn with_stop_token(mut self, token_id: TokenIdType) -> Self { + self.stop_tokens.insert(token_id); + self + } + + /// Builder pattern - add a stop sequence + pub fn with_stop_sequence(mut self, sequence: impl Into) -> Self { + self.stop_sequences.push(sequence.into()); + self + } + + /// Builder pattern - add a visible stop token + pub fn with_visible_stop_token(mut self, token_id: TokenIdType) -> Self { + self.visible_stop_tokens.insert(token_id); + self + } + + /// Builder pattern - add a visible stop sequence + pub fn with_visible_stop_sequence(mut self, sequence: impl Into) -> Self { + self.visible_stop_sequences.push(sequence.into()); + self + } +} + +/// Decoder that handles stop sequences +pub struct StopSequenceDecoder { + tokenizer: Arc, + config: StopSequenceConfig, + /// Buffer for partial matches (the "jail") + jail_buffer: String, + /// Accumulated tokens + token_buffer: Vec, + /// Offset where the prefix text starts (for context) + prefix_offset: usize, + /// Offset marking the end of previously decoded text + read_offset: usize, + /// Whether we've stopped + stopped: bool, + skip_special_tokens: bool, +} + +impl StopSequenceDecoder { + /// Create a new stop sequence decoder + pub fn new( + tokenizer: Arc, + config: StopSequenceConfig, + skip_special_tokens: bool, + ) -> Self { + StopSequenceDecoder { + tokenizer, + config, + jail_buffer: String::new(), + token_buffer: Vec::new(), + prefix_offset: 0, + read_offset: 0, + stopped: false, + skip_special_tokens, + } + } + + /// Process a single token + pub fn process_token(&mut self, token_id: TokenIdType) -> Result { + if self.stopped { + return Ok(SequenceDecoderOutput::Stopped); + } + + // Check for token-level stops first + if self.config.stop_tokens.contains(&token_id) { + self.stopped = true; + + // Flush any jailed text before stopping + if !self.jail_buffer.is_empty() { + let output = self.jail_buffer.clone(); + self.jail_buffer.clear(); + return Ok(SequenceDecoderOutput::StoppedWithText(output)); + } + return Ok(SequenceDecoderOutput::Stopped); + } + + if self.config.visible_stop_tokens.contains(&token_id) { + self.stopped = true; + + // Include jailed text plus the stop token + let stop_text = self + .tokenizer + .decode(&[token_id], self.skip_special_tokens)?; + let output = format!("{}{}", self.jail_buffer, stop_text); + self.jail_buffer.clear(); + return Ok(SequenceDecoderOutput::StoppedWithText(output)); + } + + // Add token to buffer + self.token_buffer.push(token_id); + + // Use incremental decoding like DecodeStream + // First decode the previous context (what we've already output) + let prefix_text = if self.read_offset > self.prefix_offset { + self.tokenizer.decode( + &self.token_buffer[self.prefix_offset..self.read_offset], + self.skip_special_tokens, + )? + } else { + String::new() + }; + + // Now decode from prefix to current position + let new_full_text = self.tokenizer.decode( + &self.token_buffer[self.prefix_offset..], + self.skip_special_tokens, + )?; + + // Check for incomplete UTF-8 sequence + if new_full_text.ends_with("�") { + // Wait for more tokens to complete the sequence + return Ok(SequenceDecoderOutput::Held); + } + + // Calculate only the NEW text since last successful decode + let new_text = if new_full_text.len() > prefix_text.len() { + &new_full_text[prefix_text.len()..] + } else { + // No new text produced (can happen with special tokens) + return Ok(SequenceDecoderOutput::Held); + }; + + // Combine jail buffer with new text for checking + let check_text = format!("{}{}", self.jail_buffer, new_text); + + // Check for complete stop sequences + for stop_seq in &self.config.stop_sequences { + if let Some(pos) = check_text.find(stop_seq) { + self.stopped = true; + + // Output text before the stop sequence + let output = check_text[..pos].to_string(); + self.jail_buffer.clear(); + return Ok(if output.is_empty() { + SequenceDecoderOutput::Stopped + } else { + SequenceDecoderOutput::StoppedWithText(output) + }); + } + } + + // Check for visible stop sequences + for stop_seq in &self.config.visible_stop_sequences { + if let Some(pos) = check_text.find(stop_seq) { + self.stopped = true; + + // Include the stop sequence in output + let end_pos = pos + stop_seq.len(); + let output = check_text[..end_pos].to_string(); + self.jail_buffer.clear(); + return Ok(SequenceDecoderOutput::StoppedWithText(output)); + } + } + + // Check for partial matches at the end of check_text + let mut partial_match_len = 0; + for stop_seq in self + .config + .stop_sequences + .iter() + .chain(&self.config.visible_stop_sequences) + { + // Check all possible suffixes that could be a prefix of stop_seq + for i in 1..=check_text.len().min(stop_seq.len() - 1) { + let suffix = &check_text[check_text.len() - i..]; + if stop_seq.starts_with(suffix) { + partial_match_len = partial_match_len.max(i); + } + } + } + + if partial_match_len > 0 { + // Split: output safe text, jail the potential match + let safe_end = check_text.len() - partial_match_len; + let safe_text = &check_text[..safe_end]; + self.jail_buffer = check_text[safe_end..].to_string(); + + // Update offsets for next iteration + self.prefix_offset = self.read_offset; + self.read_offset = self.token_buffer.len(); + + if safe_text.is_empty() { + Ok(SequenceDecoderOutput::Held) + } else { + Ok(SequenceDecoderOutput::Text(safe_text.to_string())) + } + } else { + // No partial matches - output everything + self.jail_buffer.clear(); + + // Update offsets for next iteration + self.prefix_offset = self.read_offset; + self.read_offset = self.token_buffer.len(); + + Ok(SequenceDecoderOutput::Text(check_text)) + } + } + + /// Process multiple tokens + pub fn process_tokens( + &mut self, + token_ids: &[TokenIdType], + ) -> Result> { + let mut outputs = Vec::new(); + for &token_id in token_ids { + outputs.push(self.process_token(token_id)?); + } + Ok(outputs) + } + + /// Flush any held text + pub fn flush(&mut self) -> SequenceDecoderOutput { + if !self.jail_buffer.is_empty() { + let output = self.jail_buffer.clone(); + self.jail_buffer.clear(); + SequenceDecoderOutput::Text(output) + } else { + SequenceDecoderOutput::Text(String::new()) + } + } + + /// Check if decoding has stopped + pub fn is_stopped(&self) -> bool { + self.stopped + } + + /// Reset the decoder state + pub fn reset(&mut self) { + self.jail_buffer.clear(); + self.token_buffer.clear(); + self.prefix_offset = 0; + self.read_offset = 0; + self.stopped = false; + } +} + +/// Builder for StopSequenceDecoder +pub struct StopSequenceDecoderBuilder { + tokenizer: Arc, + config: StopSequenceConfig, + skip_special_tokens: bool, +} + +impl StopSequenceDecoderBuilder { + pub fn new(tokenizer: Arc) -> Self { + StopSequenceDecoderBuilder { + tokenizer, + config: StopSequenceConfig::default(), + skip_special_tokens: true, + } + } + + pub fn stop_token(mut self, token_id: TokenIdType) -> Self { + self.config.stop_tokens.insert(token_id); + self + } + + pub fn stop_sequence(mut self, sequence: impl Into) -> Self { + self.config.stop_sequences.push(sequence.into()); + self + } + + pub fn visible_stop_token(mut self, token_id: TokenIdType) -> Self { + self.config.visible_stop_tokens.insert(token_id); + self + } + + pub fn visible_stop_sequence(mut self, sequence: impl Into) -> Self { + self.config.visible_stop_sequences.push(sequence.into()); + self + } + + pub fn skip_special_tokens(mut self, skip: bool) -> Self { + self.skip_special_tokens = skip; + self + } + + pub fn build(self) -> StopSequenceDecoder { + StopSequenceDecoder::new(self.tokenizer, self.config, self.skip_special_tokens) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::tokenizer::mock::MockTokenizer; + + #[test] + fn test_stop_token_detection() { + let tokenizer = Arc::new(MockTokenizer::new()); + let config = StopSequenceConfig::default().with_stop_token(999); // token + + let mut decoder = StopSequenceDecoder::new(tokenizer, config, false); + + // Process tokens before stop + let result = decoder.process_token(1).unwrap(); // "Hello" + assert!(matches!(result, SequenceDecoderOutput::Text(_))); + + // Process stop token + let result = decoder.process_token(999).unwrap(); // + assert_eq!(result, SequenceDecoderOutput::Stopped); + + // Further tokens should also return Stopped + let result = decoder.process_token(2).unwrap(); + assert_eq!(result, SequenceDecoderOutput::Stopped); + } + + #[test] + fn test_visible_stop_token() { + let tokenizer = Arc::new(MockTokenizer::new()); + let config = StopSequenceConfig::default().with_visible_stop_token(999); + + let mut decoder = StopSequenceDecoder::new(tokenizer, config, false); + + let result = decoder.process_token(999).unwrap(); + assert!(matches!(result, SequenceDecoderOutput::StoppedWithText(_))); + } + + #[test] + fn test_builder_pattern() { + let tokenizer = Arc::new(MockTokenizer::new()); + + let decoder = StopSequenceDecoderBuilder::new(tokenizer) + .stop_token(999) + .stop_sequence("STOP") + .visible_stop_token(1000) + .skip_special_tokens(true) + .build(); + + assert!(!decoder.is_stopped()); + } + + #[test] + fn test_incremental_decoding_no_repetition() { + // This test verifies the critical fix: no repeated output + let tokenizer = Arc::new(MockTokenizer::new()); + let config = StopSequenceConfig::default(); + let mut decoder = StopSequenceDecoder::new(tokenizer, config, false); + + // Process tokens one by one and collect outputs + let mut outputs = Vec::new(); + + // Token 1: "Hello" + let result = decoder.process_token(1).unwrap(); + if let SequenceDecoderOutput::Text(text) = result { + outputs.push(text.clone()); + } + + // Token 2: "world" + let result = decoder.process_token(2).unwrap(); + if let SequenceDecoderOutput::Text(text) = result { + outputs.push(text.clone()); + } + + // Token 3: "test" + let result = decoder.process_token(3).unwrap(); + if let SequenceDecoderOutput::Text(text) = result { + outputs.push(text.clone()); + } + + // CRITICAL: Each output should be unique (no accumulation) + // The fix ensures we only output NEW text, not accumulated text + assert_eq!(outputs.len(), 3); + + // Verify no text is repeated + for i in 0..outputs.len() { + for j in i + 1..outputs.len() { + // No output should contain another (no accumulation) + assert!(!outputs[j].contains(&outputs[i])); + } + } + } + + #[test] + fn test_stop_sequence_detection() { + let tokenizer = Arc::new(MockTokenizer::new()); + let config = StopSequenceConfig::default().with_stop_sequence("test"); + let mut decoder = StopSequenceDecoder::new(tokenizer, config, false); + + // Process "Hello world" + decoder.process_token(1).unwrap(); // "Hello" + decoder.process_token(2).unwrap(); // "world" + + // Process "test" which should trigger stop + let result = decoder.process_token(3).unwrap(); // "test" + + // Should stop when we hit "test" + assert!(matches!( + result, + SequenceDecoderOutput::Stopped | SequenceDecoderOutput::StoppedWithText(_) + )); + } + + #[test] + fn test_flush_after_partial() { + let tokenizer = Arc::new(MockTokenizer::new()); + let config = StopSequenceConfig::default().with_stop_sequence("NEVER_MATCH"); + let mut decoder = StopSequenceDecoder::new(tokenizer, config, false); + + // Process a token + decoder.process_token(1).unwrap(); // "Hello" + + // Flush should return any remaining text in jail + let result = decoder.flush(); + + // After processing, flush should work + assert!(matches!(result, SequenceDecoderOutput::Text(_))); + } + + #[test] + fn test_reset_functionality() { + let tokenizer = Arc::new(MockTokenizer::new()); + let config = StopSequenceConfig::default().with_stop_token(999); + let mut decoder = StopSequenceDecoder::new(tokenizer, config, false); + + // Process and stop + decoder.process_token(1).unwrap(); + decoder.process_token(999).unwrap(); + assert!(decoder.is_stopped()); + + // Reset should clear everything + decoder.reset(); + assert!(!decoder.is_stopped()); + + // Should be able to process again + let result = decoder.process_token(2).unwrap(); + assert!(matches!(result, SequenceDecoderOutput::Text(_))); + } + + #[test] + fn test_visible_stop_sequence() { + let tokenizer = Arc::new(MockTokenizer::new()); + let config = StopSequenceConfig::default().with_visible_stop_sequence("world"); + let mut decoder = StopSequenceDecoder::new(tokenizer, config, false); + + // Process "Hello" + decoder.process_token(1).unwrap(); + + // Process "world" - should include it in output + let result = decoder.process_token(2).unwrap(); + + if let SequenceDecoderOutput::StoppedWithText(text) = result { + // Should include "world" in the output + assert!(text.contains("world")); + } else { + panic!("Expected StoppedWithText with visible stop sequence"); + } + } + + #[test] + fn test_multiple_tokens_processing() { + let tokenizer = Arc::new(MockTokenizer::new()); + let config = StopSequenceConfig::default(); + let mut decoder = StopSequenceDecoder::new(tokenizer, config, false); + + // Process multiple tokens at once + let results = decoder.process_tokens(&[1, 2, 3]).unwrap(); + + // Should get results for each token + assert_eq!(results.len(), 3); + + // Each result should be Text (no stops configured) + for result in results { + assert!(matches!( + result, + SequenceDecoderOutput::Text(_) | SequenceDecoderOutput::Held + )); + } + } +} diff --git a/lmdeploy/router/src/tokenizer/stream.rs b/lmdeploy/router/src/tokenizer/stream.rs new file mode 100644 index 0000000000..848be8a8c9 --- /dev/null +++ b/lmdeploy/router/src/tokenizer/stream.rs @@ -0,0 +1,105 @@ +// src/tokenizer/stream.rs + +use super::traits::{self, TokenIdType}; +use anyhow::Result; +use std::sync::Arc; + +const INITIAL_INCREMENTAL_DETOKENIZATION_OFFSET: usize = 5; + +/// DecodeStream will keep the state necessary to produce individual chunks of +/// strings given an input stream of token_ids +pub struct DecodeStream { + /// The tokenizer used to decode token_ids + tokenizer: Arc, + + skip_special_tokens: bool, + + /// A temporary buffer of the necessary token_ids needed + /// to produce valid string chunks + all_token_ids: Vec, + + prefix_offset: usize, + read_offset: usize, +} + +impl DecodeStream { + pub fn new( + tokenizer: Arc, + prompt_token_ids: &[TokenIdType], + skip_special_tokens: bool, + ) -> Self { + let num_input_tokens = prompt_token_ids.len(); + let prompt_token_ids = prompt_token_ids.to_vec(); + Self { + tokenizer, + skip_special_tokens, + all_token_ids: prompt_token_ids, + prefix_offset: num_input_tokens + .saturating_sub(INITIAL_INCREMENTAL_DETOKENIZATION_OFFSET), + read_offset: num_input_tokens, + } + } + + /// Step appends a token_id to the internal state and tries to produce a text chunk. + /// Returning `None` means the given id is not enough to produce a chunk. + pub fn step(&mut self, id: TokenIdType) -> Result> { + self.all_token_ids.push(id); + + let prefix_text = self.tokenizer.decode( + &self.all_token_ids[self.prefix_offset..self.read_offset], + self.skip_special_tokens, + )?; + + let new_text = self.tokenizer.decode( + &self.all_token_ids[self.prefix_offset..], + self.skip_special_tokens, + )?; + + if new_text.len() > prefix_text.len() && !new_text.ends_with("�") { + let new_text = new_text[prefix_text.len()..].to_string(); + + self.prefix_offset = self.read_offset; + self.read_offset = self.all_token_ids.len(); + + Ok(Some(new_text)) + } else { + Ok(None) + } + } + + /// Process multiple tokens at once + pub fn step_batch(&mut self, token_ids: &[u32]) -> Result> { + let mut chunks = Vec::new(); + + for &token_id in token_ids { + if let Some(text) = self.step(token_id)? { + chunks.push(text); + } + } + + Ok(chunks) + } + + /// Force flush any remaining text + pub fn flush(&mut self) -> Result> { + if self.read_offset < self.all_token_ids.len() { + let remaining = self.tokenizer.decode( + &self.all_token_ids[self.read_offset..], + self.skip_special_tokens, + )?; + + self.read_offset = self.all_token_ids.len(); + + if !remaining.is_empty() { + return Ok(Some(remaining)); + } + } + + Ok(None) + } + + /// Get all tokens processed so far + pub fn tokens(&self) -> &[u32] { + &self.all_token_ids + } +} diff --git a/lmdeploy/router/src/tokenizer/tests.rs b/lmdeploy/router/src/tokenizer/tests.rs new file mode 100644 index 0000000000..2c4d4b108e --- /dev/null +++ b/lmdeploy/router/src/tokenizer/tests.rs @@ -0,0 +1,143 @@ +#[cfg(test)] +use super::*; +#[cfg(test)] +use std::sync::Arc; + +#[test] +fn test_mock_tokenizer_encode() { + let tokenizer = mock::MockTokenizer::new(); + let encoding = tokenizer.encode("Hello world").unwrap(); + let token_ids = encoding.token_ids(); + assert_eq!(token_ids, &[1, 2]); // "Hello" -> 1, "world" -> 2 +} + +#[test] +fn test_mock_tokenizer_decode() { + let tokenizer = mock::MockTokenizer::new(); + let text = tokenizer.decode(&[1, 2], false).unwrap(); + assert_eq!(text, "Hello world"); +} + +#[test] +fn test_mock_tokenizer_decode_skip_special() { + let tokenizer = mock::MockTokenizer::new(); + + // With special tokens + let text = tokenizer.decode(&[1000, 1, 2, 999], false).unwrap(); + assert_eq!(text, " Hello world "); + + // Without special tokens + let text = tokenizer.decode(&[1000, 1, 2, 999], true).unwrap(); + assert_eq!(text, "Hello world"); +} + +#[test] +fn test_tokenizer_wrapper() { + let mock_tokenizer = Arc::new(mock::MockTokenizer::new()); + let tokenizer = Tokenizer::from_arc(mock_tokenizer); + + // Test encoding + let encoding = tokenizer.encode("Hello world").unwrap(); + assert_eq!(encoding.token_ids(), &[1, 2]); + + // Test decoding + let text = tokenizer.decode(&[1, 2], false).unwrap(); + assert_eq!(text, "Hello world"); + + // Test vocab size + assert_eq!(tokenizer.vocab_size(), 8); + + // Test token to ID + assert_eq!(tokenizer.token_to_id("Hello"), Some(1)); + assert_eq!(tokenizer.token_to_id("unknown"), None); + + // Test ID to token + assert_eq!(tokenizer.id_to_token(1), Some("Hello".to_string())); + assert_eq!(tokenizer.id_to_token(9999), None); +} + +#[test] +fn test_decode_stream_basic() { + let mock_tokenizer = Arc::new(mock::MockTokenizer::new()); + let tokenizer = Tokenizer::from_arc(mock_tokenizer); + + // Create a decode stream with initial tokens + let initial_tokens = vec![1, 2]; // "Hello world" + let mut stream = tokenizer.decode_stream(&initial_tokens, false); + + // Add a new token + let result = stream.step(3).unwrap(); // "test" + // Since we're using a mock, the actual incremental behavior depends on implementation + // For now, we just verify it doesn't crash + assert!(result.is_some() || result.is_none()); +} + +#[test] +fn test_decode_stream_flush() { + let mock_tokenizer = Arc::new(mock::MockTokenizer::new()); + let tokenizer = Tokenizer::from_arc(mock_tokenizer); + + let initial_tokens = vec![1]; + let mut stream = tokenizer.decode_stream(&initial_tokens, false); + + // Add tokens + stream.step(2).unwrap(); + stream.step(3).unwrap(); + + // Flush remaining + let flushed = stream.flush().unwrap(); + // The flush behavior depends on the implementation + assert!(flushed.is_some() || flushed.is_none()); +} + +#[test] +fn test_special_tokens() { + let mock_tokenizer = Arc::new(mock::MockTokenizer::new()); + let tokenizer = Tokenizer::from_arc(mock_tokenizer); + + let special_tokens = tokenizer.get_special_tokens(); + assert_eq!(special_tokens.bos_token, Some("".to_string())); + assert_eq!(special_tokens.eos_token, Some("".to_string())); + assert_eq!(special_tokens.unk_token, Some("".to_string())); + assert!(special_tokens.sep_token.is_none()); + assert!(special_tokens.pad_token.is_none()); +} + +#[test] +fn test_batch_encode() { + let tokenizer = mock::MockTokenizer::new(); + let inputs = vec!["Hello", "world", "test"]; + let encodings = tokenizer.encode_batch(&inputs).unwrap(); + + assert_eq!(encodings.len(), 3); + assert_eq!(encodings[0].token_ids(), &[1]); // "Hello" -> 1 + assert_eq!(encodings[1].token_ids(), &[2]); // "world" -> 2 + assert_eq!(encodings[2].token_ids(), &[3]); // "test" -> 3 +} + +#[test] +fn test_thread_safety() { + use std::thread; + + let mock_tokenizer = Arc::new(mock::MockTokenizer::new()); + let tokenizer = Tokenizer::from_arc(mock_tokenizer); + + // Spawn multiple threads that use the same tokenizer + let handles: Vec<_> = (0..10) + .map(|i| { + let tokenizer_clone = tokenizer.clone(); + thread::spawn(move || { + let text = "Hello test".to_string(); + let encoding = tokenizer_clone.encode(&text).unwrap(); + let decoded = tokenizer_clone.decode(encoding.token_ids(), false).unwrap(); + assert!(decoded.contains("Hello") || decoded.contains("test")); + i + }) + }) + .collect(); + + // Wait for all threads to complete + for handle in handles { + handle.join().unwrap(); + } +} diff --git a/lmdeploy/router/src/tokenizer/tiktoken.rs b/lmdeploy/router/src/tokenizer/tiktoken.rs new file mode 100644 index 0000000000..0af5a97916 --- /dev/null +++ b/lmdeploy/router/src/tokenizer/tiktoken.rs @@ -0,0 +1,276 @@ +use super::traits::{ + Decoder, Encoder, Encoding, SpecialTokens, TokenIdType, Tokenizer as TokenizerTrait, +}; +use anyhow::{Error, Result}; +use tiktoken_rs::{cl100k_base, p50k_base, p50k_edit, r50k_base, CoreBPE}; + +/// Tiktoken tokenizer wrapper for OpenAI GPT models +pub struct TiktokenTokenizer { + tokenizer: CoreBPE, + #[allow(dead_code)] + model: TiktokenModel, + special_tokens: SpecialTokens, + vocab_size: usize, +} + +/// Supported Tiktoken models +#[derive(Debug, Clone, Copy)] +pub enum TiktokenModel { + /// GPT-4, GPT-3.5-turbo, text-embedding-ada-002 + Cl100kBase, + /// Codex models, text-davinci-002, text-davinci-003 + P50kBase, + /// Use for edit models like text-davinci-edit-001, code-davinci-edit-001 + P50kEdit, + /// GPT-3 models like davinci + R50kBase, +} + +impl TiktokenTokenizer { + /// Create a new Tiktoken tokenizer for the specified model + pub fn new(model: TiktokenModel) -> Result { + let tokenizer = + match model { + TiktokenModel::Cl100kBase => cl100k_base() + .map_err(|e| Error::msg(format!("Failed to load cl100k_base: {}", e)))?, + TiktokenModel::P50kBase => p50k_base() + .map_err(|e| Error::msg(format!("Failed to load p50k_base: {}", e)))?, + TiktokenModel::P50kEdit => p50k_edit() + .map_err(|e| Error::msg(format!("Failed to load p50k_edit: {}", e)))?, + TiktokenModel::R50kBase => r50k_base() + .map_err(|e| Error::msg(format!("Failed to load r50k_base: {}", e)))?, + }; + + // Extract special tokens (tiktoken-rs doesn't expose them directly) + // We'll use common ones for GPT models + let special_tokens = Self::get_special_tokens_for_model(model); + + // Get vocabulary size (this is an approximation) + let vocab_size = match model { + TiktokenModel::Cl100kBase => 100256, // cl100k has ~100k tokens + TiktokenModel::P50kBase | TiktokenModel::P50kEdit => 50281, // p50k has ~50k tokens + TiktokenModel::R50kBase => 50257, // r50k has ~50k tokens + }; + + Ok(TiktokenTokenizer { + tokenizer, + model, + special_tokens, + vocab_size, + }) + } + + /// Create a tokenizer from a model string (e.g., "gpt-4", "gpt-3.5-turbo") + pub fn from_model_name(model_name: &str) -> Result { + let model = Self::model_from_name(model_name)?; + Self::new(model) + } + + /// Determine the appropriate model from a model name + fn model_from_name(model_name: &str) -> Result { + // Based on OpenAI's model-to-encoding mapping + if model_name.contains("gpt-4") + || model_name.contains("gpt-3.5") + || model_name.contains("turbo") + { + Ok(TiktokenModel::Cl100kBase) + } else if model_name.contains("davinci-002") + || model_name.contains("davinci-003") + || model_name.contains("codex") + { + Ok(TiktokenModel::P50kBase) + } else if model_name.contains("edit") { + Ok(TiktokenModel::P50kEdit) + } else if model_name.contains("davinci") + || model_name.contains("curie") + || model_name.contains("babbage") + || model_name.contains("ada") + { + Ok(TiktokenModel::R50kBase) + } else { + // Return an error for unrecognized model names to prevent silent failures + Err(anyhow::anyhow!( + "Unrecognized OpenAI model name: '{}'. Expected GPT-3, GPT-3.5, GPT-4, or related model names", + model_name + )) + } + } + + /// Get special tokens for a specific model + fn get_special_tokens_for_model(model: TiktokenModel) -> SpecialTokens { + // These are common special tokens for GPT models + // The actual token IDs might vary by model + match model { + TiktokenModel::Cl100kBase => SpecialTokens { + bos_token: Some("<|endoftext|>".to_string()), + eos_token: Some("<|endoftext|>".to_string()), + unk_token: None, + sep_token: None, + pad_token: Some("<|endoftext|>".to_string()), + cls_token: None, + mask_token: None, + additional_special_tokens: vec![ + "<|fim_prefix|>".to_string(), + "<|fim_middle|>".to_string(), + "<|fim_suffix|>".to_string(), + "<|endofprompt|>".to_string(), + ], + }, + _ => SpecialTokens { + bos_token: Some("<|endoftext|>".to_string()), + eos_token: Some("<|endoftext|>".to_string()), + unk_token: None, + sep_token: None, + pad_token: Some("<|endoftext|>".to_string()), + cls_token: None, + mask_token: None, + additional_special_tokens: vec![], + }, + } + } +} + +impl Encoder for TiktokenTokenizer { + fn encode(&self, input: &str) -> Result { + let tokens = self.tokenizer.encode_ordinary(input); + Ok(Encoding::Tiktoken(tokens)) + } + + fn encode_batch(&self, inputs: &[&str]) -> Result> { + inputs.iter().map(|input| self.encode(input)).collect() + } +} + +impl Decoder for TiktokenTokenizer { + fn decode(&self, token_ids: &[TokenIdType], _skip_special_tokens: bool) -> Result { + // tiktoken-rs 0.7.0 now uses u32 (Rank type) + self.tokenizer + .decode(token_ids.to_vec()) + .map_err(|e| Error::msg(format!("Decoding failed: {}", e))) + } +} + +impl TokenizerTrait for TiktokenTokenizer { + fn vocab_size(&self) -> usize { + self.vocab_size + } + + fn get_special_tokens(&self) -> &SpecialTokens { + &self.special_tokens + } + + fn token_to_id(&self, _token: &str) -> Option { + // Tiktoken doesn't provide direct token-to-id mapping + // We'd need to encode the token and check if it produces a single ID + None + } + + fn id_to_token(&self, _id: TokenIdType) -> Option { + // Tiktoken doesn't provide direct id-to-token mapping + // We can only decode IDs to text + None + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_tiktoken_creation() { + let tokenizer = TiktokenTokenizer::new(TiktokenModel::Cl100kBase).unwrap(); + assert_eq!(tokenizer.vocab_size(), 100256); + } + + #[test] + fn test_model_from_name() { + assert!(matches!( + TiktokenTokenizer::model_from_name("gpt-4").unwrap(), + TiktokenModel::Cl100kBase + )); + assert!(matches!( + TiktokenTokenizer::model_from_name("gpt-3.5-turbo").unwrap(), + TiktokenModel::Cl100kBase + )); + assert!(matches!( + TiktokenTokenizer::model_from_name("text-davinci-003").unwrap(), + TiktokenModel::P50kBase + )); + assert!(matches!( + TiktokenTokenizer::model_from_name("text-davinci-edit-001").unwrap(), + TiktokenModel::P50kEdit + )); + assert!(matches!( + TiktokenTokenizer::model_from_name("davinci").unwrap(), + TiktokenModel::R50kBase + )); + } + + #[test] + fn test_encode_decode() { + let tokenizer = TiktokenTokenizer::new(TiktokenModel::Cl100kBase).unwrap(); + + let text = "Hello, world!"; + let encoding = tokenizer.encode(text).unwrap(); + + let decoded = tokenizer.decode(encoding.token_ids(), false).unwrap(); + assert_eq!(decoded, text); + } + + #[test] + fn test_batch_encode() { + let tokenizer = TiktokenTokenizer::new(TiktokenModel::Cl100kBase).unwrap(); + + let texts = vec!["Hello", "World", "Test"]; + let encodings = tokenizer.encode_batch(&texts).unwrap(); + + assert_eq!(encodings.len(), 3); + for (i, encoding) in encodings.iter().enumerate() { + let decoded = tokenizer.decode(encoding.token_ids(), false).unwrap(); + assert_eq!(decoded, texts[i]); + } + } + + #[test] + fn test_special_tokens() { + let tokenizer = TiktokenTokenizer::new(TiktokenModel::Cl100kBase).unwrap(); + let special_tokens = tokenizer.get_special_tokens(); + + assert!(special_tokens.eos_token.is_some()); + assert_eq!(special_tokens.eos_token.as_ref().unwrap(), "<|endoftext|>"); + } + + #[test] + fn test_unrecognized_model_name_returns_error() { + // Test that unrecognized model names return an error + let result = TiktokenTokenizer::from_model_name("distilgpt-2"); + assert!(result.is_err()); + if let Err(e) = result { + assert!(e.to_string().contains("Unrecognized OpenAI model name")); + } + + let result = TiktokenTokenizer::from_model_name("bert-base-uncased"); + assert!(result.is_err()); + if let Err(e) = result { + assert!(e.to_string().contains("Unrecognized OpenAI model name")); + } + + let result = TiktokenTokenizer::from_model_name("llama-7b"); + assert!(result.is_err()); + if let Err(e) = result { + assert!(e.to_string().contains("Unrecognized OpenAI model name")); + } + } + + #[test] + fn test_recognized_model_names() { + // Test that recognized model names work correctly + assert!(TiktokenTokenizer::from_model_name("gpt-4").is_ok()); + assert!(TiktokenTokenizer::from_model_name("gpt-3.5-turbo").is_ok()); + assert!(TiktokenTokenizer::from_model_name("text-davinci-003").is_ok()); + assert!(TiktokenTokenizer::from_model_name("code-davinci-002").is_ok()); + assert!(TiktokenTokenizer::from_model_name("text-curie-001").is_ok()); + assert!(TiktokenTokenizer::from_model_name("text-babbage-001").is_ok()); + assert!(TiktokenTokenizer::from_model_name("text-ada-001").is_ok()); + } +} diff --git a/lmdeploy/router/src/tokenizer/traits.rs b/lmdeploy/router/src/tokenizer/traits.rs new file mode 100644 index 0000000000..275dd822f4 --- /dev/null +++ b/lmdeploy/router/src/tokenizer/traits.rs @@ -0,0 +1,83 @@ +use anyhow::Result; +use std::collections::hash_map::DefaultHasher; +use std::hash::{Hash, Hasher}; + +/// Type alias for token IDs +pub type TokenIdType = u32; + +/// Core encoding trait - separate from decoding for modularity +pub trait Encoder: Send + Sync { + fn encode(&self, input: &str) -> Result; + fn encode_batch(&self, inputs: &[&str]) -> Result>; +} + +/// Core decoding trait - can be implemented independently +pub trait Decoder: Send + Sync { + fn decode(&self, token_ids: &[TokenIdType], skip_special_tokens: bool) -> Result; +} + +/// Combined tokenizer trait +pub trait Tokenizer: Encoder + Decoder { + fn vocab_size(&self) -> usize; + fn get_special_tokens(&self) -> &SpecialTokens; + fn token_to_id(&self, token: &str) -> Option; + fn id_to_token(&self, id: TokenIdType) -> Option; +} + +/// Contains the results of tokenizing text: token IDs, string tokens, and their spans +#[derive(Debug, Clone)] +pub enum Encoding { + /// Hugging Face + Hf(Box), + /// Sentence Piece + Sp(Vec), + /// Tiktoken (for GPT models) - now uses u32 in tiktoken-rs 0.7.0 + Tiktoken(Vec), +} + +impl Encoding { + /// Returns a reference to token IDs - zero-copy operation + pub fn token_ids(&self) -> &[TokenIdType] { + match self { + Encoding::Hf(inner) => inner.get_ids(), + Encoding::Sp(inner) => inner, + Encoding::Tiktoken(inner) => inner, + } + } + + /// Deprecated: Use token_ids() instead (kept for compatibility) + #[deprecated(since = "0.1.0", note = "Use token_ids() instead")] + pub fn token_ids_ref(&self) -> &[TokenIdType] { + self.token_ids() + } + + /// Get a hash of the token IDs for caching purposes + pub fn get_hash(&self) -> u64 { + let mut hasher = DefaultHasher::new(); + self.hash(&mut hasher); + hasher.finish() + } +} + +/// Hash implementation for Encoding +impl Hash for Encoding { + fn hash(&self, state: &mut H) { + match self { + Encoding::Hf(inner) => inner.get_ids().hash(state), + Encoding::Sp(inner) => inner.hash(state), + Encoding::Tiktoken(inner) => inner.hash(state), + } + } +} + +#[derive(Debug, Clone)] +pub struct SpecialTokens { + pub bos_token: Option, + pub eos_token: Option, + pub unk_token: Option, + pub sep_token: Option, + pub pad_token: Option, + pub cls_token: Option, + pub mask_token: Option, + pub additional_special_tokens: Vec, +} diff --git a/lmdeploy/router/src/tree.rs b/lmdeploy/router/src/tree.rs new file mode 100644 index 0000000000..410224dd72 --- /dev/null +++ b/lmdeploy/router/src/tree.rs @@ -0,0 +1,2320 @@ +use std::{ + cmp::Reverse, + collections::{BinaryHeap, HashMap, VecDeque}, + hash::{BuildHasherDefault, Hasher}, + sync::{ + atomic::{AtomicU64, Ordering}, + Arc, RwLock, + }, +}; + +use dashmap::{mapref::entry::Entry, DashMap}; +use tracing::debug; + +type NodeRef = Arc; + +/// Shard counts for DashMaps to balance concurrency vs allocation overhead. +/// Default DashMap uses num_cpus * 4 shards (e.g., 256 on 64-core machines). +/// +/// Root node uses higher shard count since ALL requests pass through it. +/// Other nodes use lower count as traffic diverges through the tree. +/// +/// This reduces memory by ~90% vs default while maintaining good concurrency. +const ROOT_SHARD_COUNT: usize = 32; +const NODE_SHARD_COUNT: usize = 8; + +/// Create a children DashMap for non-root nodes +#[inline] +fn new_children_map() -> DashMap { + DashMap::with_hasher_and_shard_amount(CharHasherBuilder::default(), NODE_SHARD_COUNT) +} + +/// Create a tenant access time DashMap for non-root nodes +#[inline] +fn new_tenant_map() -> DashMap { + DashMap::with_shard_amount(NODE_SHARD_COUNT) +} + +/// Interned tenant ID to avoid repeated string allocations. +/// Using Arc allows cheap cloning and comparison. +pub type TenantId = Arc; + +/// Result of a prefix match operation, including char counts to avoid recomputation. +#[derive(Debug, Clone)] +pub struct PrefixMatchResult { + /// The tenant that owns the matched prefix (zero-copy) + pub tenant: TenantId, + /// Number of characters matched + pub matched_char_count: usize, + /// Total number of characters in the input text + pub input_char_count: usize, +} + +/// A fast identity hasher for single-character keys (used in children DashMap). +/// Since chars have good distribution already, we use identity hashing with mixing. +#[derive(Default)] +struct CharHasher(u64); + +impl Hasher for CharHasher { + #[inline(always)] + fn finish(&self) -> u64 { + self.0 + } + + #[inline(always)] + fn write(&mut self, bytes: &[u8]) { + // Fast path for 4-byte (char) writes - avoid loop + if bytes.len() == 4 { + let val = u32::from_ne_bytes([bytes[0], bytes[1], bytes[2], bytes[3]]); + // Mix with golden ratio for better distribution + self.0 = (val as u64).wrapping_mul(0x9E3779B97F4A7C15); + return; + } + // Fallback for other sizes (shouldn't happen for char keys) + for &byte in bytes { + self.0 = self.0.wrapping_mul(0x100000001b3).wrapping_add(byte as u64); + } + } + + #[inline(always)] + fn write_u32(&mut self, i: u32) { + // Chars are u32 - use golden ratio multiplication for distribution + self.0 = (i as u64).wrapping_mul(0x9E3779B97F4A7C15); + } +} + +type CharHasherBuilder = BuildHasherDefault; + +/// Advance a string slice by N characters, returning the remaining slice. +/// Returns empty string if n >= char count. +/// Optimized: uses direct byte slicing for ASCII, falls back to char_indices for UTF-8. +#[inline] +fn advance_by_chars(s: &str, n: usize) -> &str { + if n == 0 { + return s; + } + if n >= s.len() { + return ""; + } + // Fast path: if first N bytes are all ASCII, we can slice directly + let bytes = s.as_bytes(); + if bytes[..n].is_ascii() { + // Safe: we verified all bytes in [0..n] are ASCII (valid UTF-8 boundary) + return &s[n..]; + } + // Slow path: UTF-8 requires char-by-char traversal + s.char_indices() + .nth(n) + .map(|(idx, _)| &s[idx..]) + .unwrap_or("") +} + +/// Get the first N characters of a string as a new String. +/// More efficient than chars().take(n).collect() for known bounds. +#[inline] +fn take_chars(s: &str, n: usize) -> String { + if n == 0 { + return String::new(); + } + s.char_indices() + .nth(n) + .map(|(idx, _)| s[..idx].to_string()) + .unwrap_or_else(|| s.to_string()) +} + +/// Node text with cached character count to avoid repeated O(n) chars().count() calls. +#[derive(Debug)] +struct NodeText { + /// The actual text stored in this node + text: String, + /// Cached character count (UTF-8 chars, not bytes) + char_count: usize, +} + +impl NodeText { + #[inline] + fn new(text: String) -> Self { + let char_count = text.chars().count(); + Self { text, char_count } + } + + #[inline] + fn empty() -> Self { + Self { + text: String::new(), + char_count: 0, + } + } + + #[inline] + fn char_count(&self) -> usize { + self.char_count + } + + #[inline] + fn as_str(&self) -> &str { + &self.text + } + + #[inline] + fn first_char(&self) -> Option { + self.text.chars().next() + } + + /// Split the text at a character boundary, returning the prefix and suffix. + /// This is more efficient than slice_by_chars as it computes both at once. + #[inline] + fn split_at_char(&self, char_idx: usize) -> (NodeText, NodeText) { + if char_idx == 0 { + return (NodeText::empty(), self.clone_text()); + } + if char_idx >= self.char_count { + return (self.clone_text(), NodeText::empty()); + } + + // Find byte index for the character boundary + let byte_idx = self + .text + .char_indices() + .nth(char_idx) + .map(|(i, _)| i) + .unwrap_or(self.text.len()); + + let prefix = NodeText { + text: self.text[..byte_idx].to_string(), + char_count: char_idx, + }; + let suffix = NodeText { + text: self.text[byte_idx..].to_string(), + char_count: self.char_count - char_idx, + }; + (prefix, suffix) + } + + #[inline] + fn clone_text(&self) -> NodeText { + NodeText { + text: self.text.clone(), + char_count: self.char_count, + } + } +} + +impl Clone for NodeText { + fn clone(&self) -> Self { + self.clone_text() + } +} + +/// Global epoch counter for LRU ordering. +/// Uses a simple incrementing counter instead of wall clock time. +/// +/// Benefits: +/// - No syscall overhead (vs SystemTime::now()) +/// - Smaller memory footprint (u64 vs u128) +/// - Perfectly monotonic (no clock skew issues) +/// +/// For LRU eviction, relative ordering is all that matters. +static EPOCH_COUNTER: AtomicU64 = AtomicU64::new(0); + +/// Get the next epoch value for LRU timestamp ordering. +/// Uses fetch_add for lock-free, monotonically increasing values. +/// Relaxed ordering is sufficient since we only need eventual consistency +/// for approximate LRU behavior. +#[inline] +fn get_epoch() -> u64 { + EPOCH_COUNTER.fetch_add(1, Ordering::Relaxed) +} + +#[derive(Debug)] +struct Node { + /// Children nodes indexed by first character. + /// Using custom hasher optimized for char keys. + children: DashMap, + /// Node text with cached character count + text: RwLock, + /// Per-tenant last access epoch for LRU ordering. Using TenantId (Arc) for cheap cloning. + tenant_last_access_time: DashMap, + /// Parent pointer for upward traversal during timestamp updates + parent: RwLock>, + /// Cached last-accessed tenant for O(1) lookup during prefix match. + /// Avoids O(shards) DashMap iteration in the common case. + last_tenant: parking_lot::RwLock>, +} + +#[derive(Debug)] +pub struct Tree { + root: NodeRef, + /// Per-tenant character count for size tracking. Using TenantId for consistency. + pub tenant_char_count: DashMap, +} + +// For the heap + +struct EvictionEntry { + timestamp: u64, + tenant: TenantId, + node: NodeRef, +} + +impl Eq for EvictionEntry {} + +#[allow(clippy::non_canonical_partial_ord_impl)] +impl PartialOrd for EvictionEntry { + fn partial_cmp(&self, other: &Self) -> Option { + Some(self.timestamp.cmp(&other.timestamp)) + } +} + +impl Ord for EvictionEntry { + fn cmp(&self, other: &Self) -> std::cmp::Ordering { + self.timestamp.cmp(&other.timestamp) + } +} + +impl PartialEq for EvictionEntry { + fn eq(&self, other: &Self) -> bool { + self.timestamp == other.timestamp + } +} + +// For char operations +// Note that in rust, `.len()` or slice is operated on the "byte" level. It causes issues for UTF-8 characters because one character might use multiple bytes. +// https://en.wikipedia.org/wiki/UTF-8 + +/// Count matching prefix characters between two strings. +/// Returns the number of characters that match from the start. +/// Optimized: uses fast byte comparison for ASCII, falls back to char iteration for UTF-8. +#[inline] +fn shared_prefix_count(a: &str, b: &str) -> usize { + let a_bytes = a.as_bytes(); + let b_bytes = b.as_bytes(); + + // Find common byte prefix length using iterator (potentially SIMD-optimized) + let common_byte_len = a_bytes + .iter() + .zip(b_bytes) + .position(|(&a_byte, &b_byte)| a_byte != b_byte) + .unwrap_or_else(|| a_bytes.len().min(b_bytes.len())); + + // If the common byte prefix is all ASCII, byte count == char count + // Otherwise, fall back to char-by-char comparison for UTF-8 safety + if a_bytes[..common_byte_len].is_ascii() { + common_byte_len + } else { + shared_prefix_count_chars(a, b) + } +} + +/// Fallback char-by-char comparison for strings with non-ASCII characters. +#[inline] +fn shared_prefix_count_chars(a: &str, b: &str) -> usize { + a.chars() + .zip(b.chars()) + .take_while(|(a_char, b_char)| a_char == b_char) + .count() +} + +/// Intern a tenant string into an Arc for efficient storage and comparison. +#[inline] +fn intern_tenant(tenant: &str) -> TenantId { + Arc::from(tenant) +} + +impl Default for Tree { + fn default() -> Self { + Self::new() + } +} + +impl Tree { + /* + Thread-safe multi tenant radix tree + + 1. Storing data for multiple tenants (the overlap of multiple radix tree) + 2. Node-level lock to enable concurrent access on nodes + 3. Leaf LRU eviction based on tenant access time + + Optimizations: + - Cached character counts in NodeText to avoid O(n) chars().count() calls + - Interned tenant IDs (Arc) for cheap cloning and comparison + - Batched timestamp updates to reduce syscalls + - Custom hasher for char keys in children DashMap + */ + + pub fn new() -> Self { + Tree { + // Root uses higher shard count since ALL requests pass through it + root: Arc::new(Node { + children: DashMap::with_hasher_and_shard_amount( + CharHasherBuilder::default(), + ROOT_SHARD_COUNT, + ), + text: RwLock::new(NodeText::empty()), + tenant_last_access_time: DashMap::with_shard_amount(ROOT_SHARD_COUNT), + parent: RwLock::new(None), + last_tenant: parking_lot::RwLock::new(None), + }), + tenant_char_count: DashMap::with_shard_amount(ROOT_SHARD_COUNT), + } + } + + pub fn insert(&self, text: &str, tenant: &str) { + // Insert text into tree with given tenant + // Use slice-based traversal to avoid Vec allocation + + // Intern the tenant ID once for reuse + let tenant_id = intern_tenant(tenant); + + // Ensure tenant exists at root (don't update timestamp - root is never evicted) + self.root + .tenant_last_access_time + .entry(Arc::clone(&tenant_id)) + .or_insert(0); + + self.tenant_char_count + .entry(Arc::clone(&tenant_id)) + .or_insert(0); + + // Track remaining text as a slice - no allocation needed + let mut remaining = text; + let mut prev = Arc::clone(&self.root); + + // Result type to carry state out of the match block + // This allows the entry guard to be dropped before we update prev + enum InsertStep { + Done, + Continue { + next_prev: NodeRef, + advance_chars: usize, + }, + } + + while !remaining.is_empty() { + let first_char = remaining.chars().next().unwrap(); + + // Use entry API for atomic check-and-insert semantics (required for thread safety) + let step = match prev.children.entry(first_char) { + Entry::Vacant(entry) => { + // No match - create new node with remaining text (this is the leaf) + // Compute remaining char count lazily - only here when creating leaf + let remaining_char_count = remaining.chars().count(); + let epoch = get_epoch(); + + let new_node = Arc::new(Node { + children: new_children_map(), + text: RwLock::new(NodeText::new(remaining.to_string())), + tenant_last_access_time: new_tenant_map(), + parent: RwLock::new(Some(Arc::clone(&prev))), + last_tenant: parking_lot::RwLock::new(Some(Arc::clone(&tenant_id))), + }); + + // Attach tenant to the new leaf node with timestamp + self.tenant_char_count + .entry(Arc::clone(&tenant_id)) + .and_modify(|count| *count += remaining_char_count) + .or_insert(remaining_char_count); + new_node + .tenant_last_access_time + .insert(Arc::clone(&tenant_id), epoch); + + entry.insert(new_node); + InsertStep::Done + } + + Entry::Occupied(mut entry) => { + let matched_node = entry.get().clone(); + + let matched_node_text = matched_node.text.read().unwrap(); + let matched_node_text_count = matched_node_text.char_count(); + let matched_node_text_str = matched_node_text.as_str(); + + // Use slice-based comparison - no allocation + let shared_count = shared_prefix_count(remaining, matched_node_text_str); + + if shared_count < matched_node_text_count { + // Split the matched node + let (matched_text, contracted_text) = + matched_node_text.split_at_char(shared_count); + let matched_text_count = shared_count; + + // Drop read lock before creating new node + drop(matched_node_text); + + let new_node = Arc::new(Node { + text: RwLock::new(matched_text), + children: new_children_map(), + parent: RwLock::new(Some(Arc::clone(&prev))), + tenant_last_access_time: matched_node.tenant_last_access_time.clone(), + last_tenant: parking_lot::RwLock::new( + matched_node.last_tenant.read().clone(), + ), + }); + + let first_new_char = contracted_text.first_char().unwrap(); + new_node + .children + .insert(first_new_char, Arc::clone(&matched_node)); + + entry.insert(Arc::clone(&new_node)); + + *matched_node.text.write().unwrap() = contracted_text; + *matched_node.parent.write().unwrap() = Some(Arc::clone(&new_node)); + + // Attach tenant to the new split node (intermediate - no timestamp update) + // The cloned DashMap already has the tenant; just ensure char count is correct + if !new_node + .tenant_last_access_time + .contains_key(tenant_id.as_ref()) + { + self.tenant_char_count + .entry(Arc::clone(&tenant_id)) + .and_modify(|count| *count += matched_text_count) + .or_insert(matched_text_count); + new_node + .tenant_last_access_time + .insert(Arc::clone(&tenant_id), 0); + } + + InsertStep::Continue { + next_prev: new_node, + advance_chars: shared_count, + } + } else { + // Full match - move to next node (intermediate - no timestamp update) + drop(matched_node_text); + + // Ensure tenant exists at this intermediate node + if !matched_node + .tenant_last_access_time + .contains_key(tenant_id.as_ref()) + { + self.tenant_char_count + .entry(Arc::clone(&tenant_id)) + .and_modify(|count| *count += matched_node_text_count) + .or_insert(matched_node_text_count); + matched_node + .tenant_last_access_time + .insert(Arc::clone(&tenant_id), 0); + } + + InsertStep::Continue { + next_prev: matched_node, + advance_chars: shared_count, + } + } + } + }; + + // Entry guard is now dropped - safe to update prev + match step { + InsertStep::Done => return, // New leaf created with timestamp, we're done + InsertStep::Continue { + next_prev, + advance_chars, + } => { + prev = next_prev; + remaining = advance_by_chars(remaining, advance_chars); + } + } + } + + // Loop exited normally (remaining empty) - prev is the leaf node + // Update its timestamp for LRU ordering + let epoch = get_epoch(); + prev.tenant_last_access_time + .insert(Arc::clone(&tenant_id), epoch); + } + + /// Performs prefix matching and returns detailed result with char counts. + /// Optimized: no string allocations, deferred char counting. + pub fn prefix_match_with_counts(&self, text: &str) -> PrefixMatchResult { + let mut remaining = text; + let mut matched_chars = 0; + let mut prev = Arc::clone(&self.root); + + while !remaining.is_empty() { + let first_char = remaining.chars().next().unwrap(); + + let child_node = prev.children.get(&first_char).map(|e| e.value().clone()); + + if let Some(matched_node) = child_node { + let matched_text_guard = matched_node.text.read().unwrap(); + let matched_node_text_count = matched_text_guard.char_count(); + + // Use slice-based comparison - no allocation + let shared_count = shared_prefix_count(remaining, matched_text_guard.as_str()); + drop(matched_text_guard); + + if shared_count == matched_node_text_count { + // Full match with current node's text, continue to next node + matched_chars += shared_count; + remaining = advance_by_chars(remaining, shared_count); + prev = matched_node; + } else { + // Partial match - still use this node for tenant selection + matched_chars += shared_count; + prev = matched_node; + break; + } + } else { + // No match found, stop here + break; + } + } + + let curr = prev; + + // Try cached tenant first (O(1)) before falling back to O(shards) DashMap iteration. + // The cache is valid if the tenant still exists in tenant_last_access_time. + let tenant: TenantId = { + let cached = curr.last_tenant.read(); + if let Some(ref t) = *cached { + if curr.tenant_last_access_time.contains_key(t.as_ref()) { + Arc::clone(t) + } else { + drop(cached); + // Cache stale, fall back to iteration and update cache + let t = curr + .tenant_last_access_time + .iter() + .next() + .map(|kv| Arc::clone(kv.key())) + .unwrap_or_else(|| Arc::from("empty")); + *curr.last_tenant.write() = Some(Arc::clone(&t)); + t + } + } else { + drop(cached); + // No cache, iterate and populate cache + let t = curr + .tenant_last_access_time + .iter() + .next() + .map(|kv| Arc::clone(kv.key())) + .unwrap_or_else(|| Arc::from("empty")); + *curr.last_tenant.write() = Some(Arc::clone(&t)); + t + } + }; + + // Update timestamp probabilistically (1 in 8 matches) to reduce DashMap contention. + // LRU eviction doesn't need perfect accuracy - approximate timestamps suffice. + let epoch = get_epoch(); + if epoch & 0x7 == 0 { + curr.tenant_last_access_time + .insert(Arc::clone(&tenant), epoch); + } + + // Compute input char count directly from input text. + // This is equivalent to matched_chars + remaining.chars().count() but avoids + // needing to track remaining precisely through the traversal. + let input_char_count = text.chars().count(); + + PrefixMatchResult { + tenant, + matched_char_count: matched_chars, + input_char_count, + } + } + + /// Legacy prefix_match API for backward compatibility. + /// Note: This computes matched_text which has allocation overhead. + pub fn prefix_match(&self, text: &str) -> (String, String) { + let result = self.prefix_match_with_counts(text); + let matched_text = take_chars(text, result.matched_char_count); + (matched_text, result.tenant.to_string()) + } + + #[allow(dead_code)] + pub fn prefix_match_tenant(&self, text: &str, tenant: &str) -> String { + // Use slice-based traversal - no Vec allocation + + // Intern tenant ID once for efficient lookups + let tenant_id = intern_tenant(tenant); + + let mut remaining = text; + let mut matched_chars = 0; + let mut prev = Arc::clone(&self.root); + + while !remaining.is_empty() { + let first_char = remaining.chars().next().unwrap(); + + let child_node = prev.children.get(&first_char).map(|e| e.value().clone()); + + if let Some(matched_node) = child_node { + // Only continue matching if this node belongs to the specified tenant + if !matched_node + .tenant_last_access_time + .contains_key(tenant_id.as_ref()) + { + break; + } + + let matched_text_guard = matched_node.text.read().unwrap(); + let matched_node_text_count = matched_text_guard.char_count(); + + // Use slice-based comparison - no allocation + let shared_count = shared_prefix_count(remaining, matched_text_guard.as_str()); + drop(matched_text_guard); + + if shared_count == matched_node_text_count { + // Full match with current node's text, continue to next node + matched_chars += shared_count; + remaining = advance_by_chars(remaining, shared_count); + prev = matched_node; + } else { + // Partial match - still use this node for timestamp update + matched_chars += shared_count; + prev = matched_node; + break; + } + } else { + // No match found, stop here + break; + } + } + + let curr = prev; + + // Only update timestamp if we found a match for the specified tenant. + // Update matched node only - ancestor propagation is unnecessary. + if curr + .tenant_last_access_time + .contains_key(tenant_id.as_ref()) + { + let epoch = get_epoch(); + curr.tenant_last_access_time + .insert(Arc::clone(&tenant_id), epoch); + } + + // Build result from original input using char count + take_chars(text, matched_chars) + } + + /// Return the list of tenants for which this node is a leaf. + /// A tenant is a leaf at this node if no children have that tenant. + fn leaf_of(node: &NodeRef) -> Vec { + let mut candidates: HashMap = node + .tenant_last_access_time + .iter() + .map(|entry| (Arc::clone(entry.key()), true)) + .collect(); + + for child in node.children.iter() { + for tenant in child.value().tenant_last_access_time.iter() { + // Mark as non-leaf if any child has this tenant + candidates.insert(Arc::clone(tenant.key()), false); + } + } + + candidates + .into_iter() + .filter(|(_, is_leaf)| *is_leaf) + .map(|(tenant, _)| tenant) + .collect() + } + + pub fn evict_tenant_by_size(&self, max_size: usize) { + // Calculate used size and collect leaves + let mut stack = vec![Arc::clone(&self.root)]; + let mut pq = BinaryHeap::new(); + + while let Some(curr) = stack.pop() { + for child in curr.children.iter() { + stack.push(Arc::clone(child.value())); + } + + // Add leaves to priority queue + for tenant in Tree::leaf_of(&curr) { + if let Some(timestamp) = curr.tenant_last_access_time.get(tenant.as_ref()) { + pq.push(Reverse(EvictionEntry { + timestamp: *timestamp, + tenant: Arc::clone(&tenant), + node: Arc::clone(&curr), + })); + } + } + } + + debug!("Before eviction - Used size per tenant:"); + for entry in self.tenant_char_count.iter() { + debug!("Tenant: {}, Size: {}", entry.key(), entry.value()); + } + + // Process eviction + while let Some(Reverse(entry)) = pq.pop() { + let EvictionEntry { tenant, node, .. } = entry; + + if let Some(used_size) = self.tenant_char_count.get(tenant.as_ref()) { + if *used_size <= max_size { + continue; + } + } + + // Verify this node is still a leaf for this tenant (may have changed) + // A node is a leaf for a tenant if no children have that tenant + let is_still_leaf = node.tenant_last_access_time.contains_key(tenant.as_ref()) + && !node.children.iter().any(|child| { + child + .value() + .tenant_last_access_time + .contains_key(tenant.as_ref()) + }); + if !is_still_leaf { + continue; + } + + // Decrement when removing tenant from node + let node_len = node.text.read().unwrap().char_count(); + self.tenant_char_count + .entry(Arc::clone(&tenant)) + .and_modify(|count| { + *count = count.saturating_sub(node_len); + }); + + // Remove tenant from node + node.tenant_last_access_time.remove(tenant.as_ref()); + + // Get parent reference outside of the borrow scope + let parent_opt = node.parent.read().unwrap().clone(); + + // Remove empty nodes + if node.children.is_empty() && node.tenant_last_access_time.is_empty() { + if let Some(ref parent) = parent_opt { + if let Some(fc) = node.text.read().unwrap().first_char() { + parent.children.remove(&fc); + } + } + } + + // If parent has this tenant and no other children have it, + // parent becomes a new leaf - add to priority queue + if let Some(ref parent) = parent_opt { + if parent.tenant_last_access_time.contains_key(tenant.as_ref()) { + let has_child_with_tenant = parent.children.iter().any(|child| { + child + .value() + .tenant_last_access_time + .contains_key(tenant.as_ref()) + }); + + if !has_child_with_tenant { + // Add parent to priority queue as new leaf + if let Some(timestamp) = parent.tenant_last_access_time.get(tenant.as_ref()) + { + pq.push(Reverse(EvictionEntry { + timestamp: *timestamp, + tenant: Arc::clone(&tenant), + node: Arc::clone(parent), + })); + } + } + } + } + } + + debug!("After eviction - Used size per tenant:"); + for entry in self.tenant_char_count.iter() { + debug!("Tenant: {}, Size: {}", entry.key(), entry.value()); + } + } + + pub fn remove_tenant(&self, tenant: &str) { + // Intern tenant ID once for efficient lookups + let tenant_id = intern_tenant(tenant); + + // 1. Find all the leaves for the tenant + // A leaf is a node that has this tenant but no children have it + let mut stack = vec![Arc::clone(&self.root)]; + let mut queue = VecDeque::new(); + + while let Some(curr) = stack.pop() { + for child in curr.children.iter() { + stack.push(Arc::clone(child.value())); + } + + // Check if this node is a leaf for the tenant + if curr + .tenant_last_access_time + .contains_key(tenant_id.as_ref()) + { + let has_child_with_tenant = curr.children.iter().any(|child| { + child + .value() + .tenant_last_access_time + .contains_key(tenant_id.as_ref()) + }); + if !has_child_with_tenant { + queue.push_back(Arc::clone(&curr)); + } + } + } + + // 2. Start from the leaves and traverse up to the root, removing the tenant from each node + while let Some(curr) = queue.pop_front() { + // Remove tenant from node + curr.tenant_last_access_time.remove(tenant_id.as_ref()); + + // Get parent reference outside of the borrow scope + let parent_opt = curr.parent.read().unwrap().clone(); + + // Remove empty nodes + if curr.children.is_empty() && curr.tenant_last_access_time.is_empty() { + if let Some(ref parent) = parent_opt { + if let Some(fc) = curr.text.read().unwrap().first_char() { + parent.children.remove(&fc); + } + } + } + + // If parent has this tenant and no other children have it, + // parent becomes a new leaf - add to queue + if let Some(ref parent) = parent_opt { + if parent + .tenant_last_access_time + .contains_key(tenant_id.as_ref()) + { + let has_child_with_tenant = parent.children.iter().any(|child| { + child + .value() + .tenant_last_access_time + .contains_key(tenant_id.as_ref()) + }); + + if !has_child_with_tenant { + queue.push_back(Arc::clone(parent)); + } + } + } + } + + // 3. Remove the tenant from the tenant_char_count map + self.tenant_char_count.remove(tenant_id.as_ref()); + } + + #[allow(dead_code)] + pub fn get_tenant_char_count(&self) -> HashMap { + self.tenant_char_count + .iter() + .map(|entry| (entry.key().to_string(), *entry.value())) + .collect() + } + + #[allow(dead_code)] + pub fn get_used_size_per_tenant(&self) -> HashMap { + // perform a DFS to traverse all nodes and calculate the total size used by each tenant + + let mut used_size_per_tenant: HashMap = HashMap::new(); + let mut stack = vec![Arc::clone(&self.root)]; + + while let Some(curr) = stack.pop() { + // Use cached char count instead of chars().count() + let text_count = curr.text.read().unwrap().char_count(); + + for tenant in curr.tenant_last_access_time.iter() { + let size = used_size_per_tenant + .entry(tenant.key().to_string()) + .or_insert(0); + *size += text_count; + } + + for child in curr.children.iter() { + stack.push(Arc::clone(child.value())); + } + } + + used_size_per_tenant + } + + #[allow(dead_code)] + fn node_to_string(node: &NodeRef, prefix: &str, is_last: bool) -> String { + let mut result = String::new(); + + // Add prefix and branch character + result.push_str(prefix); + result.push_str(if is_last { "└── " } else { "├── " }); + + // Add node text + let node_text = node.text.read().unwrap(); + result.push_str(&format!("'{}' [", node_text.as_str())); + + // Add tenant information with epoch values + let mut tenant_info = Vec::new(); + for entry in node.tenant_last_access_time.iter() { + let tenant_id = entry.key(); + let epoch = entry.value(); + tenant_info.push(format!("{} | epoch:{}", tenant_id, epoch)); + } + + result.push_str(&tenant_info.join(", ")); + result.push_str("]\n"); + + // Process children + let children: Vec<_> = node.children.iter().collect(); + let child_count = children.len(); + + for (i, entry) in children.iter().enumerate() { + let is_last_child = i == child_count - 1; + let new_prefix = format!("{}{}", prefix, if is_last { " " } else { "│ " }); + + result.push_str(&Tree::node_to_string( + entry.value(), + &new_prefix, + is_last_child, + )); + } + + result + } + + #[allow(dead_code)] + pub fn pretty_print(&self) { + if self.root.children.is_empty() { + return; + } + + let mut result = String::new(); + let children: Vec<_> = self.root.children.iter().collect(); + let child_count = children.len(); + + for (i, entry) in children.iter().enumerate() { + let is_last = i == child_count - 1; + result.push_str(&Tree::node_to_string(entry.value(), "", is_last)); + } + + println!("{result}"); + } +} + +// Unit tests +#[cfg(test)] +mod tests { + use std::{ + thread, + time::{Duration, Instant}, + }; + + use rand::{ + distr::{Alphanumeric, SampleString}, + rng as thread_rng, Rng, + }; + + use super::*; + + /// Helper to convert tenant_char_count to HashMap for comparison + fn get_maintained_counts(tree: &Tree) -> HashMap { + tree.tenant_char_count + .iter() + .map(|entry| (entry.key().to_string(), *entry.value())) + .collect() + } + + #[test] + fn test_tenant_char_count() { + let tree = Tree::new(); + + tree.insert("apple", "tenant1"); + tree.insert("apricot", "tenant1"); + tree.insert("banana", "tenant1"); + tree.insert("amplify", "tenant2"); + tree.insert("application", "tenant2"); + + let computed_sizes = tree.get_used_size_per_tenant(); + let maintained_counts = get_maintained_counts(&tree); + + println!("Phase 1 - Maintained vs Computed counts:"); + println!( + "Maintained: {:?}\nComputed: {:?}", + maintained_counts, computed_sizes + ); + assert_eq!( + maintained_counts, computed_sizes, + "Phase 1: Initial insertions" + ); + + tree.insert("apartment", "tenant1"); + tree.insert("appetite", "tenant2"); + tree.insert("ball", "tenant1"); + tree.insert("box", "tenant2"); + + let computed_sizes = tree.get_used_size_per_tenant(); + let maintained_counts = get_maintained_counts(&tree); + + println!("Phase 2 - Maintained vs Computed counts:"); + println!( + "Maintained: {:?}\nComputed: {:?}", + maintained_counts, computed_sizes + ); + assert_eq!( + maintained_counts, computed_sizes, + "Phase 2: Additional insertions" + ); + + tree.insert("zebra", "tenant1"); + tree.insert("zebra", "tenant2"); + tree.insert("zero", "tenant1"); + tree.insert("zero", "tenant2"); + + let computed_sizes = tree.get_used_size_per_tenant(); + let maintained_counts = get_maintained_counts(&tree); + + println!("Phase 3 - Maintained vs Computed counts:"); + println!( + "Maintained: {:?}\nComputed: {:?}", + maintained_counts, computed_sizes + ); + assert_eq!( + maintained_counts, computed_sizes, + "Phase 3: Overlapping insertions" + ); + + tree.evict_tenant_by_size(10); + + let computed_sizes = tree.get_used_size_per_tenant(); + let maintained_counts = get_maintained_counts(&tree); + + println!("Phase 4 - Maintained vs Computed counts:"); + println!( + "Maintained: {:?}\nComputed: {:?}", + maintained_counts, computed_sizes + ); + assert_eq!(maintained_counts, computed_sizes, "Phase 4: After eviction"); + } + + fn random_string(len: usize) -> String { + Alphanumeric.sample_string(&mut thread_rng(), len) + } + + #[test] + fn test_cold_start() { + let tree = Tree::new(); + + let (matched_text, tenant) = tree.prefix_match("hello"); + + assert_eq!(matched_text, ""); + assert_eq!(tenant, "empty"); + } + + #[test] + fn test_exact_match_seq() { + let tree = Tree::new(); + tree.insert("hello", "tenant1"); + tree.pretty_print(); + tree.insert("apple", "tenant2"); + tree.pretty_print(); + tree.insert("banana", "tenant3"); + tree.pretty_print(); + + let (matched_text, tenant) = tree.prefix_match("hello"); + assert_eq!(matched_text, "hello"); + assert_eq!(tenant, "tenant1"); + + let (matched_text, tenant) = tree.prefix_match("apple"); + assert_eq!(matched_text, "apple"); + assert_eq!(tenant, "tenant2"); + + let (matched_text, tenant) = tree.prefix_match("banana"); + assert_eq!(matched_text, "banana"); + assert_eq!(tenant, "tenant3"); + } + + #[test] + fn test_exact_match_concurrent() { + let tree = Arc::new(Tree::new()); + + // spawn 3 threads for insert + let tree_clone = Arc::clone(&tree); + + let texts = ["hello", "apple", "banana"]; + let tenants = ["tenant1", "tenant2", "tenant3"]; + + let mut handles = vec![]; + + for i in 0..3 { + let tree_clone = Arc::clone(&tree_clone); + let text = texts[i]; + let tenant = tenants[i]; + + let handle = thread::spawn(move || { + tree_clone.insert(text, tenant); + }); + + handles.push(handle); + } + + // wait + for handle in handles { + handle.join().unwrap(); + } + + // spawn 3 threads for match + let mut handles = vec![]; + + let tree_clone = Arc::clone(&tree); + + for i in 0..3 { + let tree_clone = Arc::clone(&tree_clone); + let text = texts[i]; + let tenant = tenants[i]; + + let handle = thread::spawn(move || { + let (matched_text, matched_tenant) = tree_clone.prefix_match(text); + assert_eq!(matched_text, text); + assert_eq!(matched_tenant, tenant); + }); + + handles.push(handle); + } + + // wait + for handle in handles { + handle.join().unwrap(); + } + } + + #[test] + fn test_partial_match_concurrent() { + let tree = Arc::new(Tree::new()); + + // spawn 3 threads for insert + let tree_clone = Arc::clone(&tree); + + static TEXTS: [&str; 3] = ["apple", "apabc", "acbdeds"]; + + let mut handles = vec![]; + + for text in TEXTS.iter() { + let tree_clone = Arc::clone(&tree_clone); + let tenant = "tenant0"; + + let handle = thread::spawn(move || { + tree_clone.insert(text, tenant); + }); + + handles.push(handle); + } + + // wait + for handle in handles { + handle.join().unwrap(); + } + + // spawn 3 threads for match + let mut handles = vec![]; + + let tree_clone = Arc::clone(&tree); + + for text in TEXTS.iter() { + let tree_clone = Arc::clone(&tree_clone); + let tenant = "tenant0"; + + let handle = thread::spawn(move || { + let (matched_text, matched_tenant) = tree_clone.prefix_match(text); + assert_eq!(matched_text, *text); + assert_eq!(matched_tenant, tenant); + }); + + handles.push(handle); + } + + // wait + for handle in handles { + handle.join().unwrap(); + } + } + + #[test] + fn test_group_prefix_insert_match_concurrent() { + static PREFIXES: [&str; 4] = [ + "Clock strikes midnight, I'm still wide awake", + "Got dreams bigger than these city lights", + "Time waits for no one, gotta make my move", + "Started from the bottom, that's no metaphor", + ]; + let suffixes = [ + "Got too much to prove, ain't got time to lose", + "History in the making, yeah, you can't erase this", + ]; + let tree = Arc::new(Tree::new()); + + let mut handles = vec![]; + + for (i, prefix) in PREFIXES.iter().enumerate() { + for suffix in suffixes.iter() { + let tree_clone = Arc::clone(&tree); + let text = format!("{} {}", prefix, suffix); + let tenant = format!("tenant{}", i); + + let handle = thread::spawn(move || { + tree_clone.insert(&text, &tenant); + }); + + handles.push(handle); + } + } + + // wait + for handle in handles { + handle.join().unwrap(); + } + + tree.pretty_print(); + + // check matching using multi threads + let mut handles = vec![]; + + for (i, prefix) in PREFIXES.iter().enumerate() { + let tree_clone = Arc::clone(&tree); + + let handle = thread::spawn(move || { + let (matched_text, matched_tenant) = tree_clone.prefix_match(prefix); + let tenant = format!("tenant{}", i); + assert_eq!(matched_text, *prefix); + assert_eq!(matched_tenant, tenant); + }); + + handles.push(handle); + } + + // wait + for handle in handles { + handle.join().unwrap(); + } + } + + #[test] + fn test_mixed_concurrent_insert_match() { + // ensure it does not deadlock instead of doing correctness check + + static PREFIXES: [&str; 4] = [ + "Clock strikes midnight, I'm still wide awake", + "Got dreams bigger than these city lights", + "Time waits for no one, gotta make my move", + "Started from the bottom, that's no metaphor", + ]; + let suffixes = [ + "Got too much to prove, ain't got time to lose", + "History in the making, yeah, you can't erase this", + ]; + let tree = Arc::new(Tree::new()); + + let mut handles = vec![]; + + for (i, prefix) in PREFIXES.iter().enumerate() { + for suffix in suffixes.iter() { + let tree_clone = Arc::clone(&tree); + let text = format!("{} {}", prefix, suffix); + let tenant = format!("tenant{}", i); + + let handle = thread::spawn(move || { + tree_clone.insert(&text, &tenant); + }); + + handles.push(handle); + } + } + + // check matching using multi threads + for prefix in PREFIXES.iter() { + let tree_clone = Arc::clone(&tree); + + let handle = thread::spawn(move || { + let (_matched_text, _matched_tenant) = tree_clone.prefix_match(prefix); + }); + + handles.push(handle); + } + + // wait + for handle in handles { + handle.join().unwrap(); + } + } + + #[test] + fn test_utf8_split_seq() { + // The string should be indexed and split by a utf-8 value basis instead of byte basis + // use .chars() to get the iterator of the utf-8 value + let tree = Arc::new(Tree::new()); + + static TEST_PAIRS: [(&str, &str); 3] = [ + ("你好嗎", "tenant1"), + ("你好喔", "tenant2"), + ("你心情好嗎", "tenant3"), + ]; + + // Insert sequentially + for (text, tenant) in TEST_PAIRS.iter() { + tree.insert(text, tenant); + } + + tree.pretty_print(); + + for (text, tenant) in TEST_PAIRS.iter() { + let (matched_text, matched_tenant) = tree.prefix_match(text); + assert_eq!(matched_text, *text); + assert_eq!(matched_tenant, *tenant); + } + } + + #[test] + fn test_utf8_split_concurrent() { + let tree = Arc::new(Tree::new()); + + static TEST_PAIRS: [(&str, &str); 3] = [ + ("你好嗎", "tenant1"), + ("你好喔", "tenant2"), + ("你心情好嗎", "tenant3"), + ]; + + // Create multiple threads for insertion + let mut handles = vec![]; + + for (text, tenant) in TEST_PAIRS.iter() { + let tree_clone = Arc::clone(&tree); + + let handle = thread::spawn(move || { + tree_clone.insert(text, tenant); + }); + + handles.push(handle); + } + + // Wait for all insertions to complete + for handle in handles { + handle.join().unwrap(); + } + + tree.pretty_print(); + + // Create multiple threads for matching + let mut handles = vec![]; + + for (text, tenant) in TEST_PAIRS.iter() { + let tree_clone = Arc::clone(&tree); + + let handle = thread::spawn(move || { + let (matched_text, matched_tenant) = tree_clone.prefix_match(text); + assert_eq!(matched_text, *text); + assert_eq!(matched_tenant, *tenant); + }); + + handles.push(handle); + } + + // Wait for all matches to complete + for handle in handles { + handle.join().unwrap(); + } + } + + #[test] + fn test_simple_eviction() { + let tree = Tree::new(); + let max_size = 5; + + // Insert strings for both tenants + tree.insert("hello", "tenant1"); // size 5 + + tree.insert("hello", "tenant2"); // size 5 + thread::sleep(Duration::from_millis(10)); + tree.insert("world", "tenant2"); // size 5, total for tenant2 = 10 + + tree.pretty_print(); + + let sizes_before = tree.get_used_size_per_tenant(); + assert_eq!(sizes_before.get("tenant1").unwrap(), &5); // "hello" = 5 + assert_eq!(sizes_before.get("tenant2").unwrap(), &10); // "hello" + "world" = 10 + + // Evict - should remove "hello" from tenant2 as it's the oldest + tree.evict_tenant_by_size(max_size); + + tree.pretty_print(); + + let sizes_after = tree.get_used_size_per_tenant(); + assert_eq!(sizes_after.get("tenant1").unwrap(), &5); // Should be unchanged + assert_eq!(sizes_after.get("tenant2").unwrap(), &5); // Only "world" remains + + let (matched, tenant) = tree.prefix_match("world"); + assert_eq!(matched, "world"); + assert_eq!(tenant, "tenant2"); + } + + #[test] + fn test_advanced_eviction() { + let tree = Tree::new(); + + // Set limits for each tenant + let max_size: usize = 100; + + // Define prefixes + let prefixes = ["aqwefcisdf", "iajsdfkmade", "kjnzxcvewqe", "iejksduqasd"]; + + // Insert strings with shared prefixes + for _i in 0..100 { + for (j, prefix) in prefixes.iter().enumerate() { + let random_suffix = random_string(10); + let text = format!("{}{}", prefix, random_suffix); + let tenant = format!("tenant{}", j + 1); + tree.insert(&text, &tenant); + } + } + + // Perform eviction + tree.evict_tenant_by_size(max_size); + + // Check sizes after eviction + let sizes_after = tree.get_used_size_per_tenant(); + for (tenant, &size) in sizes_after.iter() { + assert!( + size <= max_size, + "Tenant {} exceeds size limit. Current size: {}, Limit: {}", + tenant, + size, + max_size + ); + } + } + + #[test] + fn test_concurrent_operations_with_eviction() { + // Ensure eviction works fine with concurrent insert and match operations for a given period + + let tree = Arc::new(Tree::new()); + let mut handles = vec![]; + let test_duration = Duration::from_secs(10); + let start_time = Instant::now(); + let max_size = 100; // Single max size for all tenants + + // Spawn eviction thread + { + let tree = Arc::clone(&tree); + let handle = thread::spawn(move || { + while start_time.elapsed() < test_duration { + // Run eviction + tree.evict_tenant_by_size(max_size); + + // Sleep for 5 seconds + thread::sleep(Duration::from_secs(5)); + } + }); + handles.push(handle); + } + + // Spawn 4 worker threads + for thread_id in 0..4 { + let tree = Arc::clone(&tree); + let handle = thread::spawn(move || { + let mut rng = rand::rng(); + let tenant = format!("tenant{}", thread_id + 1); + let prefix = format!("prefix{}", thread_id); + + while start_time.elapsed() < test_duration { + // Random decision: match or insert (70% match, 30% insert) + if rng.random_bool(0.7) { + // Perform match operation + let random_len = rng.random_range(3..10); + let search_str = format!("{}{}", prefix, random_string(random_len)); + let (_matched, _) = tree.prefix_match(&search_str); + } else { + // Perform insert operation + let random_len = rng.random_range(5..15); + let insert_str = format!("{}{}", prefix, random_string(random_len)); + tree.insert(&insert_str, &tenant); + // println!("Thread {} inserted: {}", thread_id, insert_str); + } + + // Small random sleep to vary timing + thread::sleep(Duration::from_millis(rng.random_range(10..100))); + } + }); + handles.push(handle); + } + + // Wait for all threads to complete + for handle in handles { + handle.join().unwrap(); + } + + // final eviction + tree.evict_tenant_by_size(max_size); + + // Final size check + let final_sizes = tree.get_used_size_per_tenant(); + println!("Final sizes after test completion: {:?}", final_sizes); + + for &size in final_sizes.values() { + assert!( + size <= max_size, + "Tenant exceeds size limit. Final size: {}, Limit: {}", + size, + max_size + ); + } + } + + #[test] + fn test_leaf_of() { + let tree = Tree::new(); + + // Helper to convert leaves to strings for easier assertion + let leaves_as_strings = + |leaves: &[TenantId]| -> Vec { leaves.iter().map(|t| t.to_string()).collect() }; + + // Single node + tree.insert("hello", "tenant1"); + let leaves = Tree::leaf_of(&tree.root.children.get(&'h').unwrap()); + assert_eq!(leaves_as_strings(&leaves), vec!["tenant1"]); + + // Node with multiple tenants + tree.insert("hello", "tenant2"); + let leaves = Tree::leaf_of(&tree.root.children.get(&'h').unwrap()); + let leaves_str = leaves_as_strings(&leaves); + assert_eq!(leaves_str.len(), 2); + assert!(leaves_str.contains(&"tenant1".to_string())); + assert!(leaves_str.contains(&"tenant2".to_string())); + + // Non-leaf node + tree.insert("hi", "tenant1"); + let leaves = Tree::leaf_of(&tree.root.children.get(&'h').unwrap()); + assert!(leaves.is_empty()); + } + + #[test] + fn test_get_used_size_per_tenant() { + let tree = Tree::new(); + + // Single tenant + tree.insert("hello", "tenant1"); + tree.insert("world", "tenant1"); + let sizes = tree.get_used_size_per_tenant(); + + tree.pretty_print(); + println!("{:?}", sizes); + assert_eq!(sizes.get("tenant1").unwrap(), &10); // "hello" + "world" + + // Multiple tenants sharing nodes + tree.insert("hello", "tenant2"); + tree.insert("help", "tenant2"); + let sizes = tree.get_used_size_per_tenant(); + + tree.pretty_print(); + println!("{:?}", sizes); + assert_eq!(sizes.get("tenant1").unwrap(), &10); + assert_eq!(sizes.get("tenant2").unwrap(), &6); // "hello" + "p" + + // UTF-8 characters + tree.insert("你好", "tenant3"); + let sizes = tree.get_used_size_per_tenant(); + tree.pretty_print(); + println!("{:?}", sizes); + assert_eq!(sizes.get("tenant3").unwrap(), &2); // 2 Chinese characters + + tree.pretty_print(); + } + + #[test] + fn test_prefix_match_tenant() { + let tree = Tree::new(); + + // Insert overlapping prefixes for different tenants + tree.insert("hello", "tenant1"); // tenant1: hello + tree.insert("hello", "tenant2"); // tenant2: hello + tree.insert("hello world", "tenant2"); // tenant2: hello -> world + tree.insert("help", "tenant1"); // tenant1: hel -> p + tree.insert("helicopter", "tenant2"); // tenant2: hel -> icopter + + assert_eq!(tree.prefix_match_tenant("hello", "tenant1"), "hello"); // Full match for tenant1 + assert_eq!(tree.prefix_match_tenant("help", "tenant1"), "help"); // Exclusive to tenant1 + assert_eq!(tree.prefix_match_tenant("hel", "tenant1"), "hel"); // Shared prefix + assert_eq!(tree.prefix_match_tenant("hello world", "tenant1"), "hello"); // Should stop at tenant1's boundary + assert_eq!(tree.prefix_match_tenant("helicopter", "tenant1"), "hel"); // Should stop at tenant1's boundary + + assert_eq!(tree.prefix_match_tenant("hello", "tenant2"), "hello"); // Full match for tenant2 + assert_eq!( + tree.prefix_match_tenant("hello world", "tenant2"), + "hello world" + ); // Exclusive to tenant2 + assert_eq!( + tree.prefix_match_tenant("helicopter", "tenant2"), + "helicopter" + ); // Exclusive to tenant2 + assert_eq!(tree.prefix_match_tenant("hel", "tenant2"), "hel"); // Shared prefix + assert_eq!(tree.prefix_match_tenant("help", "tenant2"), "hel"); // Should stop at tenant2's boundary + + assert_eq!(tree.prefix_match_tenant("hello", "tenant3"), ""); // Non-existent tenant + assert_eq!(tree.prefix_match_tenant("help", "tenant3"), ""); // Non-existent tenant + } + + #[test] + fn test_simple_tenant_eviction() { + let tree = Tree::new(); + + // Insert data for multiple tenants + tree.insert("hello", "tenant1"); + tree.insert("world", "tenant1"); + tree.insert("hello", "tenant2"); + tree.insert("help", "tenant2"); + + let initial_sizes = tree.get_used_size_per_tenant(); + assert_eq!(initial_sizes.get("tenant1").unwrap(), &10); // "hello" + "world" + assert_eq!(initial_sizes.get("tenant2").unwrap(), &6); // "hello" + "p" + + // Evict tenant1 + tree.remove_tenant("tenant1"); + + let final_sizes = tree.get_used_size_per_tenant(); + assert!( + !final_sizes.contains_key("tenant1"), + "tenant1 should be completely removed" + ); + assert_eq!( + final_sizes.get("tenant2").unwrap(), + &6, + "tenant2 should be unaffected" + ); + + assert_eq!(tree.prefix_match_tenant("hello", "tenant1"), ""); + assert_eq!(tree.prefix_match_tenant("world", "tenant1"), ""); + + assert_eq!(tree.prefix_match_tenant("hello", "tenant2"), "hello"); + assert_eq!(tree.prefix_match_tenant("help", "tenant2"), "help"); + } + + #[test] + fn test_complex_tenant_eviction() { + let tree = Tree::new(); + + // Create a more complex tree structure with shared prefixes + tree.insert("apple", "tenant1"); + tree.insert("application", "tenant1"); + tree.insert("apple", "tenant2"); + tree.insert("appetite", "tenant2"); + tree.insert("banana", "tenant1"); + tree.insert("banana", "tenant2"); + tree.insert("ball", "tenant2"); + + let initial_sizes = tree.get_used_size_per_tenant(); + println!("Initial sizes: {:?}", initial_sizes); + tree.pretty_print(); + + // Evict tenant1 + tree.remove_tenant("tenant1"); + + let final_sizes = tree.get_used_size_per_tenant(); + println!("Final sizes: {:?}", final_sizes); + tree.pretty_print(); + + assert!( + !final_sizes.contains_key("tenant1"), + "tenant1 should be completely removed" + ); + + assert_eq!(tree.prefix_match_tenant("apple", "tenant1"), ""); + assert_eq!(tree.prefix_match_tenant("application", "tenant1"), ""); + assert_eq!(tree.prefix_match_tenant("banana", "tenant1"), ""); + + assert_eq!(tree.prefix_match_tenant("apple", "tenant2"), "apple"); + assert_eq!(tree.prefix_match_tenant("appetite", "tenant2"), "appetite"); + assert_eq!(tree.prefix_match_tenant("banana", "tenant2"), "banana"); + assert_eq!(tree.prefix_match_tenant("ball", "tenant2"), "ball"); + + let tenant2_size = final_sizes.get("tenant2").unwrap(); + assert_eq!(tenant2_size, &(5 + 5 + 6 + 2)); // "apple" + "etite" + "banana" + "ll" + } + + // ==================== Edge Case Tests ==================== + + #[test] + fn test_empty_string_input() { + let tree = Tree::new(); + + // Insert empty string + tree.insert("", "tenant1"); + + // Match empty string + let (matched, tenant) = tree.prefix_match(""); + assert_eq!(matched, ""); + assert_eq!(tenant, "tenant1"); + + // Insert non-empty, then match empty + tree.insert("hello", "tenant2"); + let (matched, tenant) = tree.prefix_match(""); + assert_eq!(matched, ""); + assert_eq!(tenant, "tenant1"); + } + + #[test] + fn test_single_character_operations() { + let tree = Tree::new(); + + // Insert single characters + tree.insert("a", "tenant1"); + tree.insert("b", "tenant2"); + tree.insert("c", "tenant1"); + + let (matched, tenant) = tree.prefix_match("a"); + assert_eq!(matched, "a"); + assert_eq!(tenant, "tenant1"); + + let (matched, tenant) = tree.prefix_match("b"); + assert_eq!(matched, "b"); + assert_eq!(tenant, "tenant2"); + + // Match with longer string starting with single char + let (matched, tenant) = tree.prefix_match("abc"); + assert_eq!(matched, "a"); + assert_eq!(tenant, "tenant1"); + } + + #[test] + fn test_prefix_is_subset_of_existing() { + let tree = Tree::new(); + + // Insert longer string first + tree.insert("application", "tenant1"); + + // Now insert prefix of existing + tree.insert("app", "tenant2"); + + // Match the prefix - both tenants own "app" node + let (matched, tenant) = tree.prefix_match("app"); + assert_eq!(matched, "app"); + assert!(tenant == "tenant1" || tenant == "tenant2"); + + // Match longer string + let (matched, tenant) = tree.prefix_match("application"); + assert_eq!(matched, "application"); + assert_eq!(tenant, "tenant1"); + + // Match "apple" - matches "app" + "l" from the child node = "appl" + // Then 'e' doesn't match 'i' in the remaining suffix, so stops at 4 chars + let (matched, _tenant) = tree.prefix_match("apple"); + assert_eq!(matched, "appl"); + } + + #[test] + fn test_existing_is_prefix_of_new() { + let tree = Tree::new(); + + // Insert shorter string first + tree.insert("app", "tenant1"); + + // Now insert longer string with same prefix + tree.insert("application", "tenant2"); + + let (matched, tenant) = tree.prefix_match("app"); + assert_eq!(matched, "app"); + assert!(tenant == "tenant1" || tenant == "tenant2"); + + let (matched, tenant) = tree.prefix_match("application"); + assert_eq!(matched, "application"); + assert_eq!(tenant, "tenant2"); + + // "applesauce" matches "app" + "l" from the child node = "appl" + // Then 'e' in "esauce" doesn't match 'i' in the suffix, so matching stops + let (matched, _tenant) = tree.prefix_match("applesauce"); + assert_eq!(matched, "appl"); + } + + // ==================== prefix_match_with_counts Tests ==================== + + #[test] + fn test_prefix_match_with_counts_accuracy() { + let tree = Tree::new(); + + tree.insert("hello world", "tenant1"); + + // Exact match + let result = tree.prefix_match_with_counts("hello world"); + assert_eq!(result.matched_char_count, 11); + assert_eq!(result.input_char_count, 11); + assert_eq!(&*result.tenant, "tenant1"); + + // Partial match + let result = tree.prefix_match_with_counts("hello"); + assert_eq!(result.matched_char_count, 5); + assert_eq!(result.input_char_count, 5); + + // Extended match + let result = tree.prefix_match_with_counts("hello world and more"); + assert_eq!(result.matched_char_count, 11); + assert_eq!(result.input_char_count, 20); + + // No match + let result = tree.prefix_match_with_counts("goodbye"); + assert_eq!(result.matched_char_count, 0); + assert_eq!(result.input_char_count, 7); + } + + #[test] + fn test_prefix_match_with_counts_utf8() { + let tree = Tree::new(); + + // UTF-8 string: 5 characters, more bytes + tree.insert("你好世界呀", "tenant1"); + + let result = tree.prefix_match_with_counts("你好世界呀"); + assert_eq!(result.matched_char_count, 5); + assert_eq!(result.input_char_count, 5); + + let result = tree.prefix_match_with_counts("你好"); + assert_eq!(result.matched_char_count, 2); + assert_eq!(result.input_char_count, 2); + + // Mixed ASCII and UTF-8 + tree.insert("hello你好", "tenant2"); + let result = tree.prefix_match_with_counts("hello你好世界"); + assert_eq!(result.matched_char_count, 7); // "hello你好" = 7 chars + assert_eq!(result.input_char_count, 9); // "hello你好世界" = 9 chars + } + + // ==================== Node Splitting Edge Cases ==================== + + #[test] + fn test_split_at_first_character() { + let tree = Tree::new(); + + // Insert "abc" + tree.insert("abc", "tenant1"); + + // Insert "aXX" - should split at first char + tree.insert("aXX", "tenant2"); + + let (matched, tenant) = tree.prefix_match("abc"); + assert_eq!(matched, "abc"); + assert_eq!(tenant, "tenant1"); + + let (matched, tenant) = tree.prefix_match("aXX"); + assert_eq!(matched, "aXX"); + assert_eq!(tenant, "tenant2"); + + let (matched, _) = tree.prefix_match("a"); + assert_eq!(matched, "a"); + } + + #[test] + fn test_split_at_last_character() { + let tree = Tree::new(); + + // Insert "abcd" + tree.insert("abcd", "tenant1"); + + // Insert "abcX" - should split at last char of shared prefix + tree.insert("abcX", "tenant2"); + + let (matched, tenant) = tree.prefix_match("abcd"); + assert_eq!(matched, "abcd"); + assert_eq!(tenant, "tenant1"); + + let (matched, tenant) = tree.prefix_match("abcX"); + assert_eq!(matched, "abcX"); + assert_eq!(tenant, "tenant2"); + + let (matched, _) = tree.prefix_match("abc"); + assert_eq!(matched, "abc"); + } + + #[test] + fn test_multiple_splits_same_path() { + let tree = Tree::new(); + + // Create a chain of splits + tree.insert("abcdefgh", "tenant1"); + tree.insert("abcdef", "tenant2"); + tree.insert("abcd", "tenant3"); + tree.insert("ab", "tenant4"); + + // Verify all paths work + assert_eq!(tree.prefix_match("abcdefgh").0, "abcdefgh"); + assert_eq!(tree.prefix_match("abcdef").0, "abcdef"); + assert_eq!(tree.prefix_match("abcd").0, "abcd"); + assert_eq!(tree.prefix_match("ab").0, "ab"); + assert_eq!(tree.prefix_match("a").0, "a"); + } + + // ==================== High Contention Stress Tests ==================== + + #[test] + fn test_high_contention_same_prefix() { + let tree = Arc::new(Tree::new()); + let num_threads = 16; + let ops_per_thread = 100; + let mut handles = vec![]; + + // All threads operate on strings with same prefix + for thread_id in 0..num_threads { + let tree = Arc::clone(&tree); + let handle = thread::spawn(move || { + let tenant = format!("tenant{}", thread_id); + for i in 0..ops_per_thread { + let text = format!("shared_prefix_{}", i); + tree.insert(&text, &tenant); + + // Immediately try to match + let (matched, _) = tree.prefix_match(&text); + assert!( + matched.starts_with("shared_prefix_"), + "Match should start with shared_prefix_" + ); + } + }); + handles.push(handle); + } + + for handle in handles { + handle.join().expect("Thread panicked"); + } + + // Verify tree is still consistent + let sizes = tree.get_used_size_per_tenant(); + assert!(!sizes.is_empty(), "Tree should have entries"); + } + + #[test] + fn test_rapid_insert_remove_cycles() { + let tree = Arc::new(Tree::new()); + let num_cycles = 50; + + for cycle in 0..num_cycles { + let tenant = format!("tenant{}", cycle % 5); + + // Insert several entries + for i in 0..10 { + let text = format!("cycle{}entry{}", cycle, i); + tree.insert(&text, &tenant); + } + + // Remove the tenant + tree.remove_tenant(&tenant); + + // Verify tenant is gone + let sizes = tree.get_used_size_per_tenant(); + assert!( + !sizes.contains_key(&tenant), + "Tenant {} should be removed after cycle {}", + tenant, + cycle + ); + } + } + + // ==================== ASCII/UTF-8 Consistency Tests ==================== + + #[test] + fn test_ascii_utf8_consistency() { + let tree = Tree::new(); + + // Insert ASCII + tree.insert("hello", "tenant1"); + + // Insert UTF-8 with same logical prefix (none) + tree.insert("你好", "tenant2"); + + // Insert mixed + tree.insert("hello你好", "tenant3"); + + // All should be retrievable + assert_eq!(tree.prefix_match("hello").0, "hello"); + assert_eq!(tree.prefix_match("你好").0, "你好"); + assert_eq!(tree.prefix_match("hello你好").0, "hello你好"); + + // Counts should be correct + let result = tree.prefix_match_with_counts("hello"); + assert_eq!(result.matched_char_count, 5); + assert_eq!(result.input_char_count, 5); + + let result = tree.prefix_match_with_counts("你好"); + assert_eq!(result.matched_char_count, 2); + assert_eq!(result.input_char_count, 2); + + let result = tree.prefix_match_with_counts("hello你好"); + assert_eq!(result.matched_char_count, 7); + assert_eq!(result.input_char_count, 7); + } + + #[test] + fn test_emoji_handling() { + let tree = Tree::new(); + + // Emoji are multi-byte UTF-8 + tree.insert("hello 👋", "tenant1"); + tree.insert("hello 👋🌍", "tenant2"); + + let (matched, tenant) = tree.prefix_match("hello 👋"); + assert_eq!(matched, "hello 👋"); + assert_eq!(tenant, "tenant1"); + + let (matched, tenant) = tree.prefix_match("hello 👋🌍"); + assert_eq!(matched, "hello 👋🌍"); + assert_eq!(tenant, "tenant2"); + + // Verify char count (not byte count) + let result = tree.prefix_match_with_counts("hello 👋"); + assert_eq!(result.matched_char_count, 7); + assert_eq!(result.input_char_count, 7); // h-e-l-l-o-space-emoji + } + + // ==================== Eviction Edge Cases ==================== + + #[test] + fn test_eviction_empty_tree() { + let tree = Tree::new(); + + // Should not panic on empty tree + tree.evict_tenant_by_size(100); + + let sizes = tree.get_used_size_per_tenant(); + assert!(sizes.is_empty()); + } + + #[test] + fn test_eviction_zero_max_size() { + let tree = Tree::new(); + + tree.insert("hello", "tenant1"); + tree.insert("world", "tenant1"); + + // Evict with max_size = 0 should remove everything + tree.evict_tenant_by_size(0); + + let sizes = tree.get_used_size_per_tenant(); + assert!( + sizes.is_empty() || sizes.values().all(|&v| v == 0), + "All tenants should be evicted or have zero size" + ); + } + + #[test] + fn test_eviction_single_tenant_all_entries() { + let tree = Tree::new(); + + // Insert many entries for single tenant + for i in 0..100 { + let text = format!("entry{:03}", i); + tree.insert(&text, "tenant1"); + } + + let initial_size = *tree.get_used_size_per_tenant().get("tenant1").unwrap(); + assert!(initial_size > 50, "Should have significant size"); + + // Evict to small size + tree.evict_tenant_by_size(50); + + let final_size = *tree.get_used_size_per_tenant().get("tenant1").unwrap_or(&0); + assert!( + final_size <= 50, + "Size {} should be <= 50 after eviction", + final_size + ); + } + + // ==================== Last Tenant Cache Tests ==================== + + #[test] + fn test_last_tenant_cache_update() { + let tree = Tree::new(); + + // Insert for tenant1 + tree.insert("hello", "tenant1"); + + // First match should return tenant1 + let (_, tenant) = tree.prefix_match("hello"); + assert_eq!(tenant, "tenant1"); + + // Insert for tenant2 on same path + tree.insert("hello", "tenant2"); + + // Match again - should still work (cache or iteration) + let (matched, _) = tree.prefix_match("hello"); + assert_eq!(matched, "hello"); + } + + #[test] + fn test_stale_cache_after_tenant_removal() { + let tree = Tree::new(); + + tree.insert("hello", "tenant1"); + tree.insert("hello", "tenant2"); + + // Access to populate cache + let _ = tree.prefix_match("hello"); + + // Remove tenant1 + tree.remove_tenant("tenant1"); + + // Should still work with tenant2 + let (matched, tenant) = tree.prefix_match("hello"); + assert_eq!(matched, "hello"); + assert_eq!(tenant, "tenant2"); + } + + // ==================== Consistency Verification Tests ==================== + + #[test] + fn test_char_count_consistency_after_operations() { + let tree = Tree::new(); + + // Helper to verify consistency + let verify_consistency = |tree: &Tree| { + let maintained = get_maintained_counts(tree); + let computed = tree.get_used_size_per_tenant(); + assert_eq!( + maintained, computed, + "Maintained counts should match computed counts" + ); + }; + + // Insert phase + for i in 0..50 { + tree.insert(&format!("prefix{}", i), "tenant1"); + tree.insert(&format!("other{}", i), "tenant2"); + } + verify_consistency(&tree); + + // Overlapping inserts + for i in 0..25 { + tree.insert(&format!("prefix{}", i), "tenant2"); + } + verify_consistency(&tree); + + // Eviction + tree.evict_tenant_by_size(100); + verify_consistency(&tree); + + // Tenant removal + tree.remove_tenant("tenant1"); + verify_consistency(&tree); + } + + #[test] + fn test_tree_structure_integrity_after_stress() { + let tree = Arc::new(Tree::new()); + let num_threads = 8; + let mut handles = vec![]; + + for thread_id in 0..num_threads { + let tree = Arc::clone(&tree); + let handle = thread::spawn(move || { + let mut rng = rand::rng(); + let tenant = format!("tenant{}", thread_id); + + for _ in 0..200 { + let op: u8 = rng.random_range(0..10); + let key = format!("key{}", rng.random_range(0..50)); + + match op { + 0..=6 => { + // Insert (70%) + tree.insert(&key, &tenant); + } + 7..=8 => { + // Match (20%) + let _ = tree.prefix_match(&key); + } + _ => { + // Match with counts (10%) + let _ = tree.prefix_match_with_counts(&key); + } + } + } + }); + handles.push(handle); + } + + for handle in handles { + handle.join().expect("Thread panicked during stress test"); + } + + // Verify tree is still functional + let sizes = tree.get_used_size_per_tenant(); + for (tenant, size) in sizes.iter() { + if tenant != "empty" { + assert!(*size > 0, "Tenant {} should have positive size", tenant); + } + } + + // Verify char count consistency (filter out zero-count entries like "empty" + // which appear in DFS traversal but not in the incremental counter) + let maintained: HashMap = get_maintained_counts(&tree) + .into_iter() + .filter(|(_, v)| *v > 0) + .collect(); + let computed: HashMap = tree + .get_used_size_per_tenant() + .into_iter() + .filter(|(_, v)| *v > 0) + .collect(); + assert_eq!( + maintained, computed, + "Counts should be consistent after stress test" + ); + } + + // ==================== Boundary Condition Tests ==================== + + #[test] + fn test_very_long_strings() { + let tree = Tree::new(); + + // Create a very long string (10KB) + let long_string: String = (0..10000) + .map(|i| ((i % 26) as u8 + b'a') as char) + .collect(); + + tree.insert(&long_string, "tenant1"); + + let (matched, tenant) = tree.prefix_match(&long_string); + assert_eq!(matched.len(), long_string.len()); + assert_eq!(tenant, "tenant1"); + + // Partial match of long string + let partial = &long_string[..5000]; + let (matched, _) = tree.prefix_match(partial); + assert_eq!(matched, partial); + } + + #[test] + fn test_many_tenants_same_path() { + let tree = Tree::new(); + + // 100 tenants all insert same string + for i in 0..100 { + tree.insert("shared_path", &format!("tenant{}", i)); + } + + // Match should return one of them + let (matched, _) = tree.prefix_match("shared_path"); + assert_eq!(matched, "shared_path"); + + // Verify all tenants are tracked + let sizes = tree.get_used_size_per_tenant(); + assert_eq!(sizes.len(), 100, "Should have 100 tenants"); + } + + #[test] + fn test_special_characters() { + let tree = Tree::new(); + + // Various special characters + let test_cases = vec![ + ("hello\nworld", "tenant1"), // newline + ("hello\tworld", "tenant2"), // tab + ("hello\0world", "tenant3"), // null byte + ("hello\u{A0}world", "tenant4"), // non-breaking space + ("path/to/file", "tenant5"), // slashes + ("query?param=value", "tenant6"), // URL-like + ]; + + for (text, tenant) in &test_cases { + tree.insert(text, tenant); + } + + for (text, tenant) in &test_cases { + let (matched, matched_tenant) = tree.prefix_match(text); + assert_eq!(matched, *text, "Failed for: {:?}", text); + assert_eq!(matched_tenant, *tenant); + } + } +} diff --git a/lmdeploy/router/src/types.rs b/lmdeploy/router/src/types.rs new file mode 100644 index 0000000000..ad190572ac --- /dev/null +++ b/lmdeploy/router/src/types.rs @@ -0,0 +1,74 @@ +use crate::utils::json; +use bytes::Bytes; +use http_body_util::Full; +use hyper::StatusCode; +use hyper::body::Incoming; +use hyper::header::CONTENT_TYPE; +use serde::{Deserialize, Serialize}; +use std::collections::HashMap; + +pub type Request = hyper::Request; +pub type Response = hyper::Response>; + +// Server failed to start up due to bad config +#[derive(Debug, thiserror::Error)] +#[error("Configuration error: {0}")] +pub struct ConfigurationError(String); + +// Request-specific errors +#[derive(Debug)] +pub enum AppError { + BadRequest(String), + NotFound(String), + MethodNotAllowed(String), + InternalError(String), +} + +impl AppError { + pub fn to_response(self) -> Response { + let (status, message) = match self { + AppError::BadRequest(msg) => (StatusCode::BAD_REQUEST, msg), + AppError::NotFound(msg) => (StatusCode::NOT_FOUND, msg), + AppError::MethodNotAllowed(msg) => (StatusCode::METHOD_NOT_ALLOWED, msg), + AppError::InternalError(msg) => (StatusCode::INTERNAL_SERVER_ERROR, msg), + }; + + let body = serde_json::json!({ + "error": message, + "status": status.as_u16() + }); + + hyper::Response::builder() + .status(status) + .header(CONTENT_TYPE, "application/json") + .body(Full::new(Bytes::from(body.to_string()))) + .unwrap() + } +} + +impl From for AppError { + fn from(error: reqwest::Error) -> Self { + AppError::BadRequest(format!("Bad request: {}", error)) + } +} + +impl From for ConfigurationError { + fn from(error: json::RequireError) -> Self { + ConfigurationError(format!("Failed to configure router: {}", error)) + } +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ChatCompletionRequest { + #[serde(default)] + pub stream: bool, + #[serde(flatten)] + pub rest: HashMap, +} + +#[derive(Debug, Clone, Deserialize, Serialize)] +pub struct ChatCompletionResponse { + pub id: String, + #[serde(flatten)] + pub rest: HashMap, +} diff --git a/lmdeploy/router/src/utils/json.rs b/lmdeploy/router/src/utils/json.rs new file mode 100644 index 0000000000..390f6885c9 --- /dev/null +++ b/lmdeploy/router/src/utils/json.rs @@ -0,0 +1,336 @@ +use serde_json::{Map, Number, Value}; +use std::convert::TryFrom; + +#[derive(Debug, thiserror::Error)] +pub enum RequireError { + #[error("Expected object, found {0}")] + NotAnObject(String), + #[error("Missing required field: {0}")] + MissingField(String), + #[error("Expected field '{field}' to be {expected}, but got {actual}")] + MismatchedFieldType { + field: String, + expected: String, + actual: String, + }, +} + +fn value_type(value: &Value) -> &'static str { + match value { + Value::Null => "null", + Value::Bool(_) => "boolean", + Value::Number(_) => "number", + Value::String(_) => "string", + Value::Array(_) => "array", + Value::Object(_) => "object", + } +} + +pub trait JsonTypeName { + fn type_name() -> &'static str; +} + +macro_rules! json_type_name { + ($type:ty, $name:expr) => { + impl JsonTypeName for $type { + fn type_name() -> &'static str { + $name + } + } + }; +} + +json_type_name!(bool, "boolean"); +json_type_name!(i64, "i64"); +json_type_name!(u16, "u16"); +json_type_name!(u64, "u64"); +json_type_name!(f64, "f64"); +json_type_name!(&str, "string"); +json_type_name!(&Vec, "array"); +json_type_name!(&Map, "object"); + +fn json_type_name() -> &'static str { + T::type_name() +} + +pub trait AsU16 { + fn as_u16(&self) -> Option; +} + +impl AsU16 for Number { + fn as_u16(&self) -> Option { + let ui: u64 = match self.as_u64() { + Some(x) => x, + None => return None, + }; + + u16::try_from(ui).ok() + } +} + +pub struct ValueRef<'a>(&'a Value); + +impl<'a> TryFrom> for bool { + type Error = (); + + fn try_from(value: ValueRef<'a>) -> Result { + match value { + ValueRef(Value::Bool(b)) => Ok(*b), + _ => Err(()), + } + } +} + +impl<'a> TryFrom> for i64 { + type Error = (); + + fn try_from(value: ValueRef<'a>) -> Result { + match value { + ValueRef(Value::Number(n)) => n.as_i64().ok_or(()), + _ => Err(()), + } + } +} + +impl<'a> TryFrom> for u16 { + type Error = (); + + fn try_from(value: ValueRef<'a>) -> Result { + match value { + ValueRef(Value::Number(n)) => n.as_u16().ok_or(()), + _ => Err(()), + } + } +} + +impl<'a> TryFrom> for u64 { + type Error = (); + + fn try_from(value: ValueRef<'a>) -> Result { + match value { + ValueRef(Value::Number(n)) => n.as_u64().ok_or(()), + _ => Err(()), + } + } +} + +impl<'a> TryFrom> for f64 { + type Error = (); + + fn try_from(value: ValueRef<'a>) -> Result { + match value { + ValueRef(Value::Number(n)) => n.as_f64().ok_or(()), + _ => Err(()), + } + } +} + +impl<'a> TryFrom> for &'a str { + type Error = (); + + fn try_from(value: ValueRef<'a>) -> Result { + match value { + ValueRef(Value::String(s)) => Ok(&s), + _ => Err(()), + } + } +} + +impl<'a> TryFrom> for &'a Vec { + type Error = (); + + fn try_from(value: ValueRef<'a>) -> Result { + match value { + ValueRef(Value::Array(a)) => Ok(&a), + _ => Err(()), + } + } +} + +impl<'a> TryFrom> for &'a Map { + type Error = (); + + fn try_from(value: ValueRef<'a>) -> Result { + match value { + ValueRef(Value::Object(o)) => Ok(&o), + _ => Err(()), + } + } +} + +pub trait RequireField { + fn require<'a, T>(&'a self, field: &str) -> Result + where + T: JsonTypeName + TryFrom>; +} + +impl RequireField for Map { + fn require<'a, T>(&'a self, field: &str) -> Result + where + T: JsonTypeName + TryFrom>, + { + match self.get(field) { + Some(value) => { + ValueRef(value) + .try_into() + .map_err(|_| RequireError::MismatchedFieldType { + field: field.to_string(), + expected: json_type_name::().to_string(), + actual: value_type(value).to_string(), + }) + } + None => Err(RequireError::MissingField(field.to_string())), + } + } +} + +impl RequireField for Value { + fn require<'a, T>(&'a self, field: &str) -> Result + where + T: JsonTypeName + TryFrom>, + { + let obj = match self { + Value::Object(obj) => obj, + other => return Err(RequireError::NotAnObject(value_type(other).to_string())), + }; + + obj.require(field) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use serde_json::Number; + + #[test] + fn test_require_bool() { + let json_str = r#" + { + "bool": true + } + "#; + + let json: Value = serde_json::from_str(&json_str).unwrap(); + + let b: bool = json.require("bool").unwrap(); + assert_eq!(b, true); + } + + #[test] + fn test_require_i64() { + let json_str = r#" + { + "i64": -17, + "not_i64": 3.14 + } + "#; + + let json: Value = serde_json::from_str(&json_str).unwrap(); + + let x: i64 = json.require("i64").unwrap(); + assert_eq!(x, -17); + + let y = json.require::("not_i64"); + assert!(y.is_err()); + } + + #[test] + fn test_require_u16() { + let json_str = r#" + { + "u16": 65535, + "not_u16": 100000 + } + "#; + + let json: Value = serde_json::from_str(&json_str).unwrap(); + + let x: u16 = json.require("u16").unwrap(); + assert_eq!(x, u16::MAX); + + let y = json.require::("not_u64"); + assert!(y.is_err()); + } + + #[test] + fn test_require_u64() { + let json_str = r#" + { + "u64": 17, + "not_u64": -42 + } + "#; + + let json: Value = serde_json::from_str(&json_str).unwrap(); + + let x: u64 = json.require("u64").unwrap(); + assert_eq!(x, 17); + + let y = json.require::("not_u64"); + assert!(y.is_err()); + } + + #[test] + fn test_require_f64() { + let json_str = r#" + { + "f64": 3.14159, + "also_f64": -42 + } + "#; + + let json: Value = serde_json::from_str(&json_str).unwrap(); + + let x: f64 = json.require("f64").unwrap(); + assert_eq!(x, 3.14159); + + let x: f64 = json.require("also_f64").unwrap(); + assert_eq!(x, -42.0); + } + + #[test] + fn test_require_string() { + let json_str = r#" + { + "field": "some_string" + } + "#; + + let json: Value = serde_json::from_str(&json_str).unwrap(); + let value: &str = json.require("field").unwrap(); + + assert_eq!(value, "some_string"); + } + + #[test] + fn test_require_array() { + let json_str = r#" + { + "field": [1, "b"] + } + "#; + + let json: Value = serde_json::from_str(&json_str).unwrap(); + let a: &Vec = json.require("field").unwrap(); + + assert_eq!(a[0], Value::Number(Number::from(1))); + assert_eq!(a[1], Value::String("b".to_string())); + } + + #[test] + fn test_require_object() { + let json_str = r#" + { + "field": { + "subfield": "subvalue" + } + } + "#; + + let json: Value = serde_json::from_str(&json_str).unwrap(); + let o: &Map = json.require("field").unwrap(); + + assert_eq!(o["subfield"], "subvalue"); + } +} diff --git a/lmdeploy/router/src/utils/mod.rs b/lmdeploy/router/src/utils/mod.rs new file mode 100644 index 0000000000..22fdbb38c8 --- /dev/null +++ b/lmdeploy/router/src/utils/mod.rs @@ -0,0 +1 @@ +pub mod json; diff --git a/lmdeploy/router/test_prompt_input.json b/lmdeploy/router/test_prompt_input.json new file mode 100644 index 0000000000..aa5dab4b43 --- /dev/null +++ b/lmdeploy/router/test_prompt_input.json @@ -0,0 +1,37 @@ +[ + { + "name": "Single string prompt", + "request": { + "model": "test-model", + "prompt": "Hello, world!", + "max_tokens": 100 + } + }, + { + "name": "Array of strings prompt", + "request": { + "model": "test-model", + "prompt": ["Hello", "world", "test"], + "max_tokens": 100 + } + }, + { + "name": "Single token ID array prompt", + "request": { + "model": "test-model", + "prompt": [128000, 9906, 11, 1917, 0], + "max_tokens": 100 + } + }, + { + "name": "Batch of token ID arrays prompt", + "request": { + "model": "test-model", + "prompt": [ + [128000, 9906, 11, 1917, 0], + [128001, 9906, 11, 1917, 1] + ], + "max_tokens": 100 + } + } +] diff --git a/lmdeploy/router/tests/api_endpoints_test.rs b/lmdeploy/router/tests/api_endpoints_test.rs new file mode 100644 index 0000000000..e95cead305 --- /dev/null +++ b/lmdeploy/router/tests/api_endpoints_test.rs @@ -0,0 +1,1982 @@ +mod common; + +use axum::{ + body::Body, + extract::Request, + http::{header::CONTENT_TYPE, StatusCode}, +}; +use common::mock_worker::{HealthStatus, MockWorker, MockWorkerConfig, WorkerType}; +use lmdeploy_router_rs::config::{ + CircuitBreakerConfig, ConnectionMode, LMDeployMigrationProtocol, PolicyConfig, RetryConfig, + RouterConfig, RoutingMode, +}; +use lmdeploy_router_rs::routers::{RouterFactory, RouterTrait}; +use reqwest::Client; +use serde_json::json; +use std::sync::Arc; +use tower::ServiceExt; + +/// Test context that manages mock workers +struct TestContext { + workers: Vec, + router: Arc, + client: Client, + config: RouterConfig, +} +impl TestContext { + async fn new(worker_configs: Vec) -> Self { + // Create default router config + let config = RouterConfig { + mode: RoutingMode::Regular { + worker_urls: vec![], + }, + policy: PolicyConfig::Random, + host: "127.0.0.1".to_string(), + port: 3002, + max_payload_size: 256 * 1024 * 1024, + request_timeout_secs: 600, + worker_startup_timeout_secs: 1, + worker_startup_check_interval_secs: 1, + discovery: None, + api_key: None, + api_key_validation_urls: vec![], + metrics: None, + log_dir: None, + log_level: None, + request_id_headers: None, + max_concurrent_requests: 64, + queue_size: 0, + queue_timeout_secs: 60, + rate_limit_tokens_per_second: None, + cors_allowed_origins: vec![], + retry: RetryConfig::default(), + circuit_breaker: CircuitBreakerConfig::default(), + disable_retries: false, + disable_circuit_breaker: false, + health_check: lmdeploy_router_rs::config::HealthCheckConfig::default(), + enable_igw: false, + connection_mode: ConnectionMode::Http, + history_backend: lmdeploy_router_rs::config::HistoryBackend::Memory, + enable_profiling: false, + profile_timeout_secs: 30, + kv_connector: lmdeploy_router_rs::config::KvConnector::Nixl, + }; + + Self::new_with_config(config, worker_configs).await + } + + async fn new_with_config( + mut config: RouterConfig, + worker_configs: Vec, + ) -> Self { + let mut workers = Vec::new(); + let mut worker_urls = Vec::new(); + + // Start mock workers if any + for worker_config in worker_configs { + let mut worker = MockWorker::new(worker_config); + let url = worker.start().await.unwrap(); + worker_urls.push(url); + workers.push(worker); + } + + if !workers.is_empty() { + tokio::time::sleep(tokio::time::Duration::from_millis(200)).await; + } + + // Update config with worker URLs if not already set + if let RoutingMode::Regular { + worker_urls: ref mut urls, + } = config.mode + { + if urls.is_empty() { + *urls = worker_urls.clone(); + } + } + + let client = Client::builder() + .timeout(std::time::Duration::from_secs(config.request_timeout_secs)) + .build() + .unwrap(); + + // Create app context + let app_context = common::create_test_context(config.clone()); + + // Create router + let router = RouterFactory::create_router(&app_context).await.unwrap(); + let router = Arc::from(router); + + // Wait for router to discover workers + if !workers.is_empty() { + tokio::time::sleep(tokio::time::Duration::from_millis(500)).await; + } + + Self { + workers, + router, + client, + config, + } + } + + async fn create_app(&self) -> axum::Router { + common::test_app::create_test_app( + Arc::clone(&self.router), + self.client.clone(), + &self.config, + ) + } + + async fn shutdown(mut self) { + for worker in &mut self.workers { + worker.stop().await; + } + } +} + +#[cfg(test)] +mod health_tests { + use super::*; + + #[tokio::test] + async fn test_liveness_endpoint() { + let ctx = TestContext::new(vec![]).await; + let app = ctx.create_app().await; + + let req = Request::builder() + .method("GET") + .uri("/liveness") + .body(Body::empty()) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_readiness_with_healthy_workers() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18001, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let app = ctx.create_app().await; + + let req = Request::builder() + .method("GET") + .uri("/readiness") + .body(Body::empty()) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_readiness_with_unhealthy_workers() { + let ctx = TestContext::new(vec![]).await; + + let app = ctx.create_app().await; + + let req = Request::builder() + .method("GET") + .uri("/readiness") + .body(Body::empty()) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + // With no workers, readiness should return SERVICE_UNAVAILABLE + assert_eq!(resp.status(), StatusCode::SERVICE_UNAVAILABLE); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_health_endpoint_details() { + let ctx = TestContext::new(vec![ + MockWorkerConfig { + port: 18003, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }, + MockWorkerConfig { + port: 18004, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }, + ]) + .await; + + let app = ctx.create_app().await; + + let req = Request::builder() + .method("GET") + .uri("/health") + .body(Body::empty()) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + // The health endpoint returns plain text, not JSON + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body_str = String::from_utf8_lossy(&body); + assert!(body_str.contains("All servers healthy")); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_health_generate_endpoint() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18005, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let app = ctx.create_app().await; + + let req = Request::builder() + .method("GET") + .uri("/health_generate") + .body(Body::empty()) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert!(body_json.is_object()); + + ctx.shutdown().await; + } +} + +#[cfg(test)] +mod generation_tests { + use super::*; + + #[tokio::test] + async fn test_generate_success() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18101, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let app = ctx.create_app().await; + + let payload = json!({ + "text": "Hello, world!", + "stream": false + }); + + let req = Request::builder() + .method("POST") + .uri("/generate") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert!(body_json.get("text").is_some()); + assert!(body_json.get("meta_info").is_some()); + let meta_info = &body_json["meta_info"]; + assert!(meta_info.get("finish_reason").is_some()); + assert_eq!(meta_info["finish_reason"]["type"], "stop"); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_generate_invalid_input_ids_returns_bad_request() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18107, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + let app = ctx.create_app().await; + + let payload = json!({ + "input_ids": "not_a_list", + "max_tokens": 5 + }); + let request = Request::builder() + .method("POST") + .uri("/generate") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let response = app.oneshot(request).await.unwrap(); + + assert_eq!(response.status(), StatusCode::BAD_REQUEST); + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_generate_streaming() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18102, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let app = ctx.create_app().await; + + let payload = json!({ + "text": "Stream test", + "stream": true + }); + + let req = Request::builder() + .method("POST") + .uri("/generate") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + // For streaming responses, the router might use chunked encoding or other streaming mechanisms + // The exact content-type can vary based on the router implementation + // Just verify we got a successful response + // Note: In a real implementation, we'd check for text/event-stream or appropriate streaming headers + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_generate_with_worker_failure() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18103, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 1.0, // Always fail + }]) + .await; + + let app = ctx.create_app().await; + + let payload = json!({ + "text": "This should fail", + "stream": false + }); + + let req = Request::builder() + .method("POST") + .uri("/generate") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_v1_chat_completions_success() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18104, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let app = ctx.create_app().await; + + let payload = json!({ + "model": "test-model", + "messages": [ + {"role": "user", "content": "Hello!"} + ], + "stream": false + }); + + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert!(body_json.get("choices").is_some()); + + ctx.shutdown().await; + } +} + +#[cfg(test)] +mod model_info_tests { + use super::*; + + #[tokio::test] + async fn test_get_server_info() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18201, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let app = ctx.create_app().await; + + let req = Request::builder() + .method("GET") + .uri("/get_server_info") + .body(Body::empty()) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert!(body_json.is_object()); + // Check for actual backend server fields + assert!(body_json.get("version").is_some()); + assert!(body_json.get("model_path").is_some()); + assert!(body_json.get("port").is_some()); + assert!(body_json.get("max_num_batched_tokens").is_some()); + assert!(body_json.get("schedule_policy").is_some()); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_get_model_info() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18202, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let app = ctx.create_app().await; + + let req = Request::builder() + .method("GET") + .uri("/get_model_info") + .body(Body::empty()) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert!(body_json.is_object()); + // Check for actual backend model info fields + assert_eq!( + body_json.get("model_path").and_then(|v| v.as_str()), + Some("mock-model-path") + ); + assert_eq!( + body_json.get("is_generation").and_then(|v| v.as_bool()), + Some(true) + ); + assert!(body_json.get("preferred_sampling_params").is_some()); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_v1_models() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18203, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let app = ctx.create_app().await; + + let req = Request::builder() + .method("GET") + .uri("/v1/models") + .body(Body::empty()) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert!(body_json.get("object").is_some()); + assert_eq!( + body_json.get("object").and_then(|v| v.as_str()), + Some("list") + ); + + let data = body_json.get("data").and_then(|v| v.as_array()); + assert!(data.is_some()); + + let models = data.unwrap(); + assert!(!models.is_empty()); + + let first_model = &models[0]; + assert_eq!( + first_model.get("id").and_then(|v| v.as_str()), + Some("mock-model") + ); + assert_eq!( + first_model.get("object").and_then(|v| v.as_str()), + Some("model") + ); + assert!(first_model.get("created").is_some()); + assert_eq!( + first_model.get("owned_by").and_then(|v| v.as_str()), + Some("organization-owner") + ); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_model_info_with_no_workers() { + let ctx = TestContext::new(vec![]).await; + let app = ctx.create_app().await; + + // Test server info with no workers + let req = Request::builder() + .method("GET") + .uri("/get_server_info") + .body(Body::empty()) + .unwrap(); + let resp = app.clone().oneshot(req).await.unwrap(); + // Router may return various error codes when no workers + assert!( + resp.status() == StatusCode::OK + || resp.status() == StatusCode::SERVICE_UNAVAILABLE + || resp.status() == StatusCode::NOT_FOUND + || resp.status() == StatusCode::INTERNAL_SERVER_ERROR, + "Unexpected status code: {:?}", + resp.status() + ); + + // Test model info with no workers + let req = Request::builder() + .method("GET") + .uri("/get_model_info") + .body(Body::empty()) + .unwrap(); + let resp = app.clone().oneshot(req).await.unwrap(); + // Router may return various error codes when no workers + assert!( + resp.status() == StatusCode::OK + || resp.status() == StatusCode::SERVICE_UNAVAILABLE + || resp.status() == StatusCode::NOT_FOUND + || resp.status() == StatusCode::INTERNAL_SERVER_ERROR, + "Unexpected status code: {:?}", + resp.status() + ); + + // Test v1/models with no workers + let req = Request::builder() + .method("GET") + .uri("/v1/models") + .body(Body::empty()) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + // Router may return various error codes when no workers + assert!( + resp.status() == StatusCode::OK + || resp.status() == StatusCode::SERVICE_UNAVAILABLE + || resp.status() == StatusCode::NOT_FOUND + || resp.status() == StatusCode::INTERNAL_SERVER_ERROR, + "Unexpected status code: {:?}", + resp.status() + ); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_model_info_with_multiple_workers() { + let ctx = TestContext::new(vec![ + MockWorkerConfig { + port: 18204, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }, + MockWorkerConfig { + port: 18205, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }, + ]) + .await; + + let app = ctx.create_app().await; + + // Test that model info is consistent across workers + for _ in 0..5 { + let req = Request::builder() + .method("GET") + .uri("/get_model_info") + .body(Body::empty()) + .unwrap(); + + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!( + body_json.get("model_path").and_then(|v| v.as_str()), + Some("mock-model-path") + ); + } + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_model_info_with_unhealthy_worker() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18206, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 1.0, // Always fail + }]) + .await; + + let app = ctx.create_app().await; + + let req = Request::builder() + .method("GET") + .uri("/get_model_info") + .body(Body::empty()) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + // Worker with fail_rate: 1.0 should always return an error status + assert!( + resp.status() == StatusCode::INTERNAL_SERVER_ERROR + || resp.status() == StatusCode::SERVICE_UNAVAILABLE, + "Expected error status for always-failing worker, got: {:?}", + resp.status() + ); + + ctx.shutdown().await; + } +} + +#[cfg(test)] +mod worker_management_tests { + use super::*; + + #[tokio::test] + async fn test_add_new_worker() { + let ctx = TestContext::new(vec![]).await; + let app = ctx.create_app().await; + + // Start a mock worker + let mut worker = MockWorker::new(MockWorkerConfig { + port: 18301, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }); + let url = worker.start().await.unwrap(); + + // Add the worker + let req = Request::builder() + .method("POST") + .uri(format!("/add_worker?url={}", url)) + .body(Body::empty()) + .unwrap(); + + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + // List workers to verify + let req = Request::builder() + .method("GET") + .uri("/list_workers") + .body(Body::empty()) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + let workers = body_json["urls"].as_array().unwrap(); + assert!(workers.iter().any(|w| w.as_str().unwrap() == url)); + + worker.stop().await; + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_remove_existing_worker() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18302, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let app = ctx.create_app().await; + + // Get the worker URL + let req = Request::builder() + .method("GET") + .uri("/list_workers") + .body(Body::empty()) + .unwrap(); + let resp = app.clone().oneshot(req).await.unwrap(); + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + let workers = body_json["urls"].as_array().unwrap(); + let worker_url = workers[0].as_str().unwrap(); + + // Remove the worker + let req = Request::builder() + .method("POST") + .uri(format!("/remove_worker?url={}", worker_url)) + .body(Body::empty()) + .unwrap(); + + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + // Verify it's removed + let req = Request::builder() + .method("GET") + .uri("/list_workers") + .body(Body::empty()) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + let workers = body_json["urls"].as_array().unwrap(); + assert!(workers.is_empty()); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_add_worker_invalid_url() { + let ctx = TestContext::new(vec![]).await; + let app = ctx.create_app().await; + + // Invalid URL format + let req = Request::builder() + .method("POST") + .uri("/add_worker?url=not-a-valid-url") + .body(Body::empty()) + .unwrap(); + + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::BAD_REQUEST); + + // Missing URL parameter + let req = Request::builder() + .method("POST") + .uri("/add_worker") + .body(Body::empty()) + .unwrap(); + + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::BAD_REQUEST); + + // Empty URL + let req = Request::builder() + .method("POST") + .uri("/add_worker?url=") + .body(Body::empty()) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::BAD_REQUEST); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_add_duplicate_worker() { + // Start a mock worker + let mut worker = MockWorker::new(MockWorkerConfig { + port: 18303, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }); + let url = worker.start().await.unwrap(); + + let ctx = TestContext::new(vec![]).await; + let app = ctx.create_app().await; + + // Add worker first time + let req = Request::builder() + .method("POST") + .uri(format!("/add_worker?url={}", url)) + .body(Body::empty()) + .unwrap(); + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + tokio::time::sleep(tokio::time::Duration::from_millis(500)).await; + + // Try to add same worker again + let req = Request::builder() + .method("POST") + .uri(format!("/add_worker?url={}", url)) + .body(Body::empty()) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + // Should return error for duplicate + assert_eq!(resp.status(), StatusCode::BAD_REQUEST); + + worker.stop().await; + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_add_unhealthy_worker() { + // Start unhealthy worker + let mut worker = MockWorker::new(MockWorkerConfig { + port: 18304, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Unhealthy, + response_delay_ms: 0, + fail_rate: 0.0, + }); + let url = worker.start().await.unwrap(); + + let ctx = TestContext::new(vec![]).await; + let app = ctx.create_app().await; + + // Try to add unhealthy worker + let req = Request::builder() + .method("POST") + .uri(format!("/add_worker?url={}", url)) + .body(Body::empty()) + .unwrap(); + let resp = app.oneshot(req).await.unwrap(); + + // Router should reject unhealthy workers + assert!( + resp.status() == StatusCode::BAD_REQUEST + || resp.status() == StatusCode::SERVICE_UNAVAILABLE + ); + + worker.stop().await; + ctx.shutdown().await; + } +} + +#[cfg(test)] +mod router_policy_tests { + use super::*; + + #[tokio::test] + async fn test_random_policy() { + let ctx = TestContext::new(vec![ + MockWorkerConfig { + port: 18801, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }, + MockWorkerConfig { + port: 18802, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }, + ]) + .await; + + // Send multiple requests and verify they succeed + let app = ctx.create_app().await; + + for i in 0..10 { + let payload = json!({ + "text": format!("Request {}", i), + "stream": false + }); + + let req = Request::builder() + .method("POST") + .uri("/generate") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + } + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_worker_selection() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18207, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let _payload = json!({ + "text": "Test selection", + "stream": false + }); + + // Check that router has the worker + let worker_urls = ctx.router.get_worker_urls(); + assert_eq!(worker_urls.len(), 1); + assert!(worker_urls[0].contains("18207")); + + ctx.shutdown().await; + } +} + +#[cfg(test)] +mod responses_endpoint_tests { + use super::*; + use reqwest::Client as HttpClient; + + #[tokio::test] + async fn test_v1_responses_non_streaming() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18950, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let app = ctx.create_app().await; + + let payload = json!({ + "input": "Hello Responses API", + "model": "mock-model", + "stream": false + }); + + let req = Request::builder() + .method("POST") + .uri("/v1/responses") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(body_json["object"], "response"); + assert_eq!(body_json["status"], "completed"); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_v1_responses_streaming() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18951, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let app = ctx.create_app().await; + + let payload = json!({ + "input": "Hello Responses API", + "model": "mock-model", + "stream": true + }); + + let req = Request::builder() + .method("POST") + .uri("/v1/responses") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + // Check that content-type indicates SSE + let headers = resp.headers().clone(); + let ct = headers + .get("content-type") + .and_then(|v| v.to_str().ok()) + .unwrap_or(""); + assert!(ct.contains("text/event-stream")); + + // We don't fully consume the stream in this test harness. + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_v1_responses_get() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18952, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let app = ctx.create_app().await; + + // First create a response to obtain an id + let resp_id = "test-get-resp-id-123"; + let payload = json!({ + "input": "Hello Responses API", + "model": "mock-model", + "stream": false, + "store": true, + "background": true, + "request_id": resp_id + }); + let req = Request::builder() + .method("POST") + .uri("/v1/responses") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + // Retrieve the response + let req = Request::builder() + .method("GET") + .uri(format!("/v1/responses/{}", resp_id)) + .body(Body::empty()) + .unwrap(); + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let get_json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(get_json["object"], "response"); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_v1_responses_cancel() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18953, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let app = ctx.create_app().await; + + // First create a response to obtain an id + let resp_id = "test-cancel-resp-id-456"; + let payload = json!({ + "input": "Hello Responses API", + "model": "mock-model", + "stream": false, + "store": true, + "background": true, + "request_id": resp_id + }); + let req = Request::builder() + .method("POST") + .uri("/v1/responses") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + // Cancel the response + let req = Request::builder() + .method("POST") + .uri(format!("/v1/responses/{}/cancel", resp_id)) + .body(Body::empty()) + .unwrap(); + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let cancel_json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + assert_eq!(cancel_json["status"], "cancelled"); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_v1_responses_delete_and_list_not_implemented() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18954, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let app = ctx.create_app().await; + + // Use an arbitrary id for delete/list + let resp_id = "resp-test-123"; + + let req = Request::builder() + .method("DELETE") + .uri(format!("/v1/responses/{}", resp_id)) + .body(Body::empty()) + .unwrap(); + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::NOT_IMPLEMENTED); + + let req = Request::builder() + .method("GET") + .uri(format!("/v1/responses/{}/input", resp_id)) + .body(Body::empty()) + .unwrap(); + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::NOT_IMPLEMENTED); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_v1_responses_get_multi_worker_fanout() { + // Start two mock workers + let ctx = TestContext::new(vec![ + MockWorkerConfig { + port: 18960, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }, + MockWorkerConfig { + port: 18961, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }, + ]) + .await; + + let app = ctx.create_app().await; + + // Create a background response with a known id + let rid = format!("resp_{}", 18960); // arbitrary unique id + let payload = json!({ + "input": "Hello Responses API", + "model": "mock-model", + "background": true, + "store": true, + "request_id": rid, + }); + + let req = Request::builder() + .method("POST") + .uri("/v1/responses") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + // Using the router, GET should succeed by fanning out across workers + let req = Request::builder() + .method("GET") + .uri(format!("/v1/responses/{}", rid)) + .body(Body::empty()) + .unwrap(); + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + // Validate only one worker holds the metadata: direct calls + let client = HttpClient::new(); + let mut ok_count = 0usize; + for url in ctx.router.get_worker_urls() { + let get_url = format!("{}/v1/responses/{}", url, rid); + let res = client.get(get_url).send().await.unwrap(); + if res.status() == StatusCode::OK { + ok_count += 1; + } + } + assert_eq!(ok_count, 1, "exactly one worker should store the response"); + + ctx.shutdown().await; + } +} + +#[cfg(test)] +mod error_tests { + use super::*; + + #[tokio::test] + async fn test_404_not_found() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18401, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let app = ctx.create_app().await; + + // Test unknown endpoint + let req = Request::builder() + .method("GET") + .uri("/unknown_endpoint") + .body(Body::empty()) + .unwrap(); + + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::NOT_FOUND); + + // Test POST to unknown endpoint + let req = Request::builder() + .method("POST") + .uri("/api/v2/generate") + .header(CONTENT_TYPE, "application/json") + .body(Body::from( + serde_json::to_string(&json!({"text": "test"})).unwrap(), + )) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::NOT_FOUND); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_method_not_allowed() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18402, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let app = ctx.create_app().await; + + // GET request to POST-only endpoint + let req = Request::builder() + .method("GET") + .uri("/generate") + .body(Body::empty()) + .unwrap(); + + let resp = app.clone().oneshot(req).await.unwrap(); + // Note: Axum returns 405 for wrong methods on matched routes + assert_eq!(resp.status(), StatusCode::METHOD_NOT_ALLOWED); + + // POST request to GET-only endpoint + let req = Request::builder() + .method("POST") + .uri("/health") + .header(CONTENT_TYPE, "application/json") + .body(Body::from("{}")) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::METHOD_NOT_ALLOWED); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_payload_too_large() { + // Create context with small payload limit + let config = RouterConfig { + mode: RoutingMode::Regular { + worker_urls: vec![], + }, + policy: PolicyConfig::Random, + host: "127.0.0.1".to_string(), + port: 3010, + max_payload_size: 1024, // 1KB limit + request_timeout_secs: 600, + worker_startup_timeout_secs: 1, + worker_startup_check_interval_secs: 1, + api_key: None, + api_key_validation_urls: vec![], + discovery: None, + metrics: None, + log_dir: None, + log_level: None, + request_id_headers: None, + max_concurrent_requests: 64, + queue_size: 0, + queue_timeout_secs: 60, + rate_limit_tokens_per_second: None, + cors_allowed_origins: vec![], + retry: RetryConfig::default(), + circuit_breaker: CircuitBreakerConfig::default(), + disable_retries: false, + disable_circuit_breaker: false, + health_check: lmdeploy_router_rs::config::HealthCheckConfig::default(), + enable_igw: false, + connection_mode: ConnectionMode::Http, + history_backend: lmdeploy_router_rs::config::HistoryBackend::Memory, + enable_profiling: false, + profile_timeout_secs: 30, + kv_connector: lmdeploy_router_rs::config::KvConnector::Nixl, + }; + + let ctx = TestContext::new_with_config( + config, + vec![MockWorkerConfig { + port: 18403, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }], + ) + .await; + + // Note: The server would have payload size middleware configured + // but we cannot test it directly through the test app + // This test is kept for documentation purposes + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_invalid_json_payload() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18404, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let app = ctx.create_app().await; + + // Send invalid JSON + let req = Request::builder() + .method("POST") + .uri("/generate") + .header(CONTENT_TYPE, "application/json") + .body(Body::from("{invalid json}")) + .unwrap(); + + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::BAD_REQUEST); + + // Send empty body + let req = Request::builder() + .method("POST") + .uri("/generate") + .header(CONTENT_TYPE, "application/json") + .body(Body::empty()) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::BAD_REQUEST); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_missing_required_fields() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18405, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let app = ctx.create_app().await; + + // Missing messages in chat completion + let payload = json!({ + "model": "test-model" + // missing "messages" + }); + + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + // Axum validates JSON schema - returns 422 for validation errors + assert_eq!(resp.status(), StatusCode::UNPROCESSABLE_ENTITY); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_invalid_model() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18406, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let app = ctx.create_app().await; + + let payload = json!({ + "model": "invalid-model-name-that-does-not-exist", + "messages": [{"role": "user", "content": "Hello"}], + "stream": false + }); + + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + // Mock worker accepts any model, but real implementation might return 400 + assert!(resp.status().is_success() || resp.status() == StatusCode::BAD_REQUEST); + + ctx.shutdown().await; + } +} + +#[cfg(test)] +mod cache_tests { + use super::*; + + #[tokio::test] + async fn test_flush_cache() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18501, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let app = ctx.create_app().await; + + let req = Request::builder() + .method("POST") + .uri("/flush_cache") + .body(Body::empty()) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + // The response might be empty or contain a message + let body_bytes = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + if !body_bytes.is_empty() { + if let Ok(body) = serde_json::from_slice::(&body_bytes) { + // Check that we got a successful response with expected fields + assert!(body.is_object()); + assert!(body.get("message").is_some() || body.get("status").is_some()); + } + } + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_get_loads() { + let ctx = TestContext::new(vec![ + MockWorkerConfig { + port: 18502, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }, + MockWorkerConfig { + port: 18503, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }, + ]) + .await; + + let app = ctx.create_app().await; + + let req = Request::builder() + .method("GET") + .uri("/get_loads") + .body(Body::empty()) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let body = axum::body::to_bytes(resp.into_body(), usize::MAX) + .await + .unwrap(); + let body_json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + + // Verify the response contains load information + assert!(body_json.is_object()); + // The exact structure depends on the implementation + // but should contain worker load information + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_flush_cache_no_workers() { + let ctx = TestContext::new(vec![]).await; + + let app = ctx.create_app().await; + + let req = Request::builder() + .method("POST") + .uri("/flush_cache") + .body(Body::empty()) + .unwrap(); + + let resp = app.oneshot(req).await.unwrap(); + // Should either succeed (no-op) or return service unavailable + assert!( + resp.status() == StatusCode::OK || resp.status() == StatusCode::SERVICE_UNAVAILABLE + ); + + ctx.shutdown().await; + } +} + +#[cfg(test)] +mod load_balancing_tests { + use super::*; + + #[tokio::test] + async fn test_request_distribution() { + // Create multiple workers + let ctx = TestContext::new(vec![ + MockWorkerConfig { + port: 18601, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }, + MockWorkerConfig { + port: 18602, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }, + ]) + .await; + + let app = ctx.create_app().await; + + // Send multiple requests and track distribution + let mut request_count = 0; + for i in 0..10 { + let payload = json!({ + "text": format!("Request {}", i), + "stream": false + }); + + let req = Request::builder() + .method("POST") + .uri("/generate") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.clone().oneshot(req).await.unwrap(); + if resp.status() == StatusCode::OK { + request_count += 1; + } + } + + // With random policy, all requests should succeed + assert_eq!(request_count, 10); + + ctx.shutdown().await; + } +} + +#[cfg(test)] +mod pd_mode_tests { + use super::*; + + #[tokio::test] + async fn test_pd_mode_routing() { + // Create PD mode configuration with prefill and decode workers + let mut prefill_worker = MockWorker::new(MockWorkerConfig { + port: 18701, + worker_type: WorkerType::Prefill, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }); + + let mut decode_worker = MockWorker::new(MockWorkerConfig { + port: 18702, + worker_type: WorkerType::Decode, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }); + + let prefill_url = prefill_worker.start().await.unwrap(); + let decode_url = decode_worker.start().await.unwrap(); + + tokio::time::sleep(tokio::time::Duration::from_millis(200)).await; + + let config = RouterConfig { + mode: RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec![prefill_url], + decode_urls: vec![decode_url], + prefill_policy: None, + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::Rdma, + rdma_config: None, + dummy_prefill: false, + }, + policy: PolicyConfig::Random, + host: "127.0.0.1".to_string(), + port: 3011, + max_payload_size: 256 * 1024 * 1024, + request_timeout_secs: 600, + worker_startup_timeout_secs: 1, + worker_startup_check_interval_secs: 1, + discovery: None, + metrics: None, + log_dir: None, + api_key: None, + api_key_validation_urls: vec![], + log_level: None, + request_id_headers: None, + max_concurrent_requests: 64, + queue_size: 0, + queue_timeout_secs: 60, + rate_limit_tokens_per_second: None, + cors_allowed_origins: vec![], + retry: RetryConfig::default(), + circuit_breaker: CircuitBreakerConfig::default(), + disable_retries: false, + disable_circuit_breaker: false, + health_check: lmdeploy_router_rs::config::HealthCheckConfig::default(), + enable_igw: false, + connection_mode: ConnectionMode::Http, + history_backend: lmdeploy_router_rs::config::HistoryBackend::Memory, + enable_profiling: false, + profile_timeout_secs: 30, + kv_connector: lmdeploy_router_rs::config::KvConnector::Nixl, + }; + + // Create app context + let app_context = common::create_test_context(config); + + // Create router - this might fail due to health check issues + let router_result = RouterFactory::create_router(&app_context).await; + + // Clean up workers + prefill_worker.stop().await; + decode_worker.stop().await; + + // For now, just verify the configuration was attempted + assert!(router_result.is_err() || router_result.is_ok()); + } +} + +#[cfg(test)] +mod request_id_tests { + use super::*; + + #[tokio::test] + async fn test_request_id_generation() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 18901, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let app = ctx.create_app().await; + + // Test 1: Request without any request ID header should generate one + let payload = json!({ + "text": "Test request", + "stream": false + }); + + let req = Request::builder() + .method("POST") + .uri("/generate") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + // Check that response has x-request-id header + let request_id = resp.headers().get("x-request-id"); + assert!( + request_id.is_some(), + "Response should have x-request-id header" + ); + + let id_value = request_id.unwrap().to_str().unwrap(); + assert!( + id_value.starts_with("gnt-"), + "Generate endpoint should have gnt- prefix" + ); + assert!( + id_value.len() > 4, + "Request ID should have content after prefix" + ); + + // Test 2: Request with custom x-request-id should preserve it + let custom_id = "custom-request-id-123"; + let req = Request::builder() + .method("POST") + .uri("/generate") + .header(CONTENT_TYPE, "application/json") + .header("x-request-id", custom_id) + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let response_id = resp.headers().get("x-request-id"); + assert!(response_id.is_some()); + assert_eq!(response_id.unwrap(), custom_id); + + // Test 3: Different endpoints should have different prefixes + let chat_payload = json!({ + "messages": [{"role": "user", "content": "Hello"}], + "model": "test-model" + }); + + let req = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header(CONTENT_TYPE, "application/json") + .body(Body::from(serde_json::to_string(&chat_payload).unwrap())) + .unwrap(); + + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let request_id = resp.headers().get("x-request-id"); + assert!(request_id.is_some()); + assert!(request_id + .unwrap() + .to_str() + .unwrap() + .starts_with("chatcmpl-")); + + // Test 4: Alternative request ID headers should be recognized + let req = Request::builder() + .method("POST") + .uri("/generate") + .header(CONTENT_TYPE, "application/json") + .header("x-correlation-id", "correlation-123") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let response_id = resp.headers().get("x-request-id"); + assert!(response_id.is_some()); + assert_eq!(response_id.unwrap(), "correlation-123"); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_request_id_with_custom_headers() { + // Create config with custom request ID headers + let config = RouterConfig { + mode: RoutingMode::Regular { + worker_urls: vec![], + }, + policy: PolicyConfig::Random, + host: "127.0.0.1".to_string(), + port: 3002, + max_payload_size: 256 * 1024 * 1024, + request_timeout_secs: 600, + worker_startup_timeout_secs: 1, + worker_startup_check_interval_secs: 1, + discovery: None, + metrics: None, + api_key: None, + api_key_validation_urls: vec![], + log_dir: None, + log_level: None, + request_id_headers: Some(vec!["custom-id".to_string(), "trace-id".to_string()]), + max_concurrent_requests: 64, + queue_size: 0, + queue_timeout_secs: 60, + rate_limit_tokens_per_second: None, + cors_allowed_origins: vec![], + retry: RetryConfig::default(), + circuit_breaker: CircuitBreakerConfig::default(), + disable_retries: false, + disable_circuit_breaker: false, + health_check: lmdeploy_router_rs::config::HealthCheckConfig::default(), + enable_igw: false, + connection_mode: ConnectionMode::Http, + history_backend: lmdeploy_router_rs::config::HistoryBackend::Memory, + enable_profiling: false, + profile_timeout_secs: 30, + kv_connector: lmdeploy_router_rs::config::KvConnector::Nixl, + }; + + let ctx = TestContext::new_with_config( + config, + vec![MockWorkerConfig { + port: 18902, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }], + ) + .await; + + let app = ctx.create_app().await; + + let payload = json!({ + "text": "Test request", + "stream": false + }); + + // Test custom header is recognized + let req = Request::builder() + .method("POST") + .uri("/generate") + .header(CONTENT_TYPE, "application/json") + .header("custom-id", "my-custom-id") + .body(Body::from(serde_json::to_string(&payload).unwrap())) + .unwrap(); + + let resp = app.clone().oneshot(req).await.unwrap(); + assert_eq!(resp.status(), StatusCode::OK); + + let response_id = resp.headers().get("x-request-id"); + assert!(response_id.is_some()); + assert_eq!(response_id.unwrap(), "my-custom-id"); + + ctx.shutdown().await; + } +} diff --git a/lmdeploy/router/tests/benchmark_integration.rs b/lmdeploy/router/tests/benchmark_integration.rs new file mode 100644 index 0000000000..077da8dda3 --- /dev/null +++ b/lmdeploy/router/tests/benchmark_integration.rs @@ -0,0 +1,233 @@ +// Integration test to ensure benchmarks compile and basic functionality works +// This prevents benchmarks from breaking in CI +// +// UPDATED: Removed deprecated ToPdRequest usage, now uses direct JSON serialization + +use lmdeploy_router_rs::core::{BasicWorker, WorkerType}; +use lmdeploy_router_rs::protocols::spec::{ + ChatCompletionRequest, ChatMessage, CompletionRequest, GenerateParameters, GenerateRequest, + PromptInput, SamplingParams, UserMessageContent, +}; +use serde_json::{from_str, to_string, to_value}; + +/// Create a default GenerateRequest for benchmarks with minimal fields set +fn default_generate_request() -> GenerateRequest { + GenerateRequest { + text: None, + prompt: None, + input_ids: None, + stream: false, + parameters: None, + sampling_params: None, + return_logprob: false, + rid: None, + other: serde_json::Map::new(), + } +} + +/// Create a default ChatCompletionRequest for benchmarks with minimal fields set +fn default_chat_completion_request() -> ChatCompletionRequest { + ChatCompletionRequest { + model: None, + messages: vec![], + input_ids: None, + max_tokens: None, + max_completion_tokens: None, + temperature: None, + top_p: None, + n: None, + stream: false, + stream_options: None, + stop: None, + presence_penalty: None, + frequency_penalty: None, + logit_bias: None, + logprobs: false, + top_logprobs: None, + user: None, + response_format: None, + seed: None, + tools: None, + tool_choice: None, + parallel_tool_calls: None, + function_call: None, + functions: None, + chat_template_kwargs: None, + // Generation extensions + top_k: None, + min_p: None, + min_tokens: None, + repetition_penalty: None, + regex: None, + ebnf: None, + stop_token_ids: None, + no_stop_trim: false, + ignore_eos: false, + add_generation_prompt: true, + continue_final_message: false, + skip_special_tokens: true, + separate_reasoning: true, + stream_reasoning: true, + echo: None, + reasoning_effort: None, + include_reasoning: true, + structured_outputs: None, + other: serde_json::Map::new(), + } +} + +/// Create a default CompletionRequest for benchmarks with minimal fields set +fn default_completion_request() -> CompletionRequest { + CompletionRequest { + model: None, + prompt: PromptInput::String(String::new()), + suffix: None, + max_tokens: None, + temperature: None, + top_p: None, + n: None, + stream: false, + stream_options: None, + logprobs: None, + echo: false, + stop: None, + presence_penalty: None, + frequency_penalty: None, + best_of: None, + logit_bias: None, + user: None, + seed: None, + // Generation extensions + top_k: None, + min_p: None, + min_tokens: None, + repetition_penalty: None, + regex: None, + ebnf: None, + json_schema: None, + stop_token_ids: None, + no_stop_trim: false, + ignore_eos: false, + skip_special_tokens: true, + other: serde_json::Map::new(), + } +} + +#[allow(dead_code)] +fn create_test_worker() -> BasicWorker { + BasicWorker::new("http://test-server:8000".to_string(), WorkerType::Prefill) +} + +#[test] +fn test_benchmark_request_creation() { + // Ensure all benchmark request types can be created without panicking + + let generate_req = GenerateRequest { + text: Some("Test prompt".to_string()), + parameters: Some(GenerateParameters { + max_new_tokens: Some(100), + temperature: Some(0.8), + top_p: Some(0.9), + top_k: Some(50), + repetition_penalty: Some(1.0), + ..Default::default() + }), + sampling_params: Some(SamplingParams { + temperature: Some(0.8), + top_p: Some(0.9), + top_k: Some(50), + frequency_penalty: Some(0.0), + presence_penalty: Some(0.0), + repetition_penalty: Some(1.0), + ..Default::default() + }), + ..default_generate_request() + }; + + let chat_req = ChatCompletionRequest { + model: Some("test-model".to_string()), + messages: vec![ChatMessage::User { + role: "user".to_string(), + content: UserMessageContent::Text("Test message".to_string()), + name: None, + }], + max_tokens: Some(150), + max_completion_tokens: Some(150), + temperature: Some(0.7), + top_p: Some(1.0), + n: Some(1), + presence_penalty: Some(0.0), + frequency_penalty: Some(0.0), + parallel_tool_calls: Some(true), + structured_outputs: Some( + lmdeploy_router_rs::protocols::spec::StructuredOutputsParams { + json: Some(serde_json::json!({ + "type": "object", + "properties": { + "name": {"type": "string"}, + "age": {"type": "integer"} + }, + "required": ["name", "age"] + })), + ..Default::default() + }, + ), + ..default_chat_completion_request() + }; + + let completion_req = CompletionRequest { + model: Some("test-model".to_string()), + prompt: PromptInput::String("Test prompt".to_string()), + max_tokens: Some(50), + temperature: Some(0.8), + top_p: Some(1.0), + n: Some(1), + presence_penalty: Some(0.0), + frequency_penalty: Some(0.0), + best_of: Some(1), + ..default_completion_request() + }; + + // Test serialization works + assert!(to_string(&generate_req).is_ok()); + assert!(to_string(&chat_req).is_ok()); + assert!(to_string(&completion_req).is_ok()); +} + +#[test] +fn test_benchmark_serialization_roundtrip() { + // Test serialization/deserialization roundtrip for benchmark types + + let generate_req = GenerateRequest { + text: Some("Test prompt".to_string()), + ..default_generate_request() + }; + + // Serialize and deserialize + let json = to_string(&generate_req).expect("Serialization should work"); + let deserialized: GenerateRequest = from_str(&json).expect("Deserialization should work"); + + // Verify basic field equality + assert_eq!(generate_req.text, deserialized.text); + assert_eq!(generate_req.stream, deserialized.stream); + assert_eq!(generate_req.return_logprob, deserialized.return_logprob); +} + +#[test] +fn test_benchmark_direct_json_routing() { + // Test direct JSON routing functionality for benchmark types (replaces regular routing) + + let generate_req = GenerateRequest { + text: Some("Test prompt".to_string()), + ..default_generate_request() + }; + + // Test direct JSON conversion (replaces regular routing methods) + let json = to_value(&generate_req).unwrap(); + let json_string = to_string(&json).unwrap(); + let bytes = json_string.as_bytes(); + + // Verify conversions work + assert!(!json_string.is_empty()); + assert!(!bytes.is_empty()); +} diff --git a/lmdeploy/router/tests/cache_aware_backward_compat_test.rs b/lmdeploy/router/tests/cache_aware_backward_compat_test.rs new file mode 100644 index 0000000000..6876b93850 --- /dev/null +++ b/lmdeploy/router/tests/cache_aware_backward_compat_test.rs @@ -0,0 +1,129 @@ +use lmdeploy_router_rs::core::{BasicWorker, Worker, WorkerType}; +use lmdeploy_router_rs::policies::{CacheAwareConfig, CacheAwarePolicy, LoadBalancingPolicy}; +use std::collections::HashMap; +use std::sync::Arc; + +#[test] +fn test_backward_compatibility_with_empty_model_id() { + let config = CacheAwareConfig { + cache_threshold: 0.5, + balance_abs_threshold: 2, + balance_rel_threshold: 1.5, + eviction_interval_secs: 0, // Disable background eviction for testing + max_tree_size: 100, + }; + + let policy = CacheAwarePolicy::with_config(config); + + // Create workers with empty model_id (simulating existing routers) + let worker1 = BasicWorker::new("http://worker1:8080".to_string(), WorkerType::Regular); + // No model_id label - should default to "unknown" + + let mut labels2 = HashMap::new(); + labels2.insert("model_id".to_string(), "unknown".to_string()); + let worker2 = BasicWorker::new("http://worker2:8080".to_string(), WorkerType::Regular) + .with_labels(labels2); + + // Add workers - should both go to "default" tree + policy.add_worker(&worker1); + policy.add_worker(&worker2); + + // Create worker list + let workers: Vec> = vec![Arc::new(worker1.clone()), Arc::new(worker2.clone())]; + + // Select worker - should work without errors + let selected = policy.select_worker(&workers, Some("test request")); + assert!(selected.is_some(), "Should select a worker"); + + // Remove workers - should work without errors + policy.remove_worker(&worker1); + policy.remove_worker(&worker2); +} + +#[test] +fn test_mixed_model_ids() { + let config = CacheAwareConfig { + cache_threshold: 0.5, + balance_abs_threshold: 2, + balance_rel_threshold: 1.5, + eviction_interval_secs: 0, + max_tree_size: 100, + }; + + let policy = CacheAwarePolicy::with_config(config); + + // Create workers with different model_id scenarios + let worker1 = BasicWorker::new("http://worker1:8080".to_string(), WorkerType::Regular); + // No model_id label - defaults to "unknown" which goes to "default" tree + + let mut labels2 = HashMap::new(); + labels2.insert("model_id".to_string(), "llama-3".to_string()); + let worker2 = BasicWorker::new("http://worker2:8080".to_string(), WorkerType::Regular) + .with_labels(labels2); + + let mut labels3 = HashMap::new(); + labels3.insert("model_id".to_string(), "unknown".to_string()); + let worker3 = BasicWorker::new("http://worker3:8080".to_string(), WorkerType::Regular) + .with_labels(labels3); + + let mut labels4 = HashMap::new(); + labels4.insert("model_id".to_string(), "llama-3".to_string()); + let worker4 = BasicWorker::new("http://worker4:8080".to_string(), WorkerType::Regular) + .with_labels(labels4); + + // Add all workers + policy.add_worker(&worker1); + policy.add_worker(&worker2); + policy.add_worker(&worker3); + policy.add_worker(&worker4); + + // Test selection with default workers only + let default_workers: Vec> = + vec![Arc::new(worker1.clone()), Arc::new(worker3.clone())]; + let selected = policy.select_worker(&default_workers, Some("test request")); + assert!(selected.is_some(), "Should select from default workers"); + + // Test selection with specific model workers only + let llama_workers: Vec> = + vec![Arc::new(worker2.clone()), Arc::new(worker4.clone())]; + let selected = policy.select_worker(&llama_workers, Some("test request")); + assert!(selected.is_some(), "Should select from llama-3 workers"); + + // Test selection with mixed workers + let all_workers: Vec> = vec![ + Arc::new(worker1.clone()), + Arc::new(worker2.clone()), + Arc::new(worker3.clone()), + Arc::new(worker4.clone()), + ]; + let selected = policy.select_worker(&all_workers, Some("test request")); + assert!(selected.is_some(), "Should select from all workers"); +} + +#[test] +fn test_remove_worker_by_url_backward_compat() { + let config = CacheAwareConfig::default(); + let policy = CacheAwarePolicy::with_config(config); + + // Create workers with different model_ids + let mut labels1 = HashMap::new(); + labels1.insert("model_id".to_string(), "llama-3".to_string()); + let worker1 = BasicWorker::new("http://worker1:8080".to_string(), WorkerType::Regular) + .with_labels(labels1); + + let worker2 = BasicWorker::new("http://worker2:8080".to_string(), WorkerType::Regular); + // No model_id label - defaults to "unknown" + + // Add workers + policy.add_worker(&worker1); + policy.add_worker(&worker2); + + // Remove by URL (backward compatibility method) + // Should remove from all trees since we don't know the model + policy.remove_worker_by_url("http://worker1:8080"); + + // Verify removal worked + let workers: Vec> = vec![Arc::new(worker2.clone())]; + let selected = policy.select_worker(&workers, Some("test")); + assert_eq!(selected, Some(0), "Should only have worker2 left"); +} diff --git a/lmdeploy/router/tests/common/mock_openai_server.rs b/lmdeploy/router/tests/common/mock_openai_server.rs new file mode 100644 index 0000000000..19058f6868 --- /dev/null +++ b/lmdeploy/router/tests/common/mock_openai_server.rs @@ -0,0 +1,299 @@ +//! Mock servers for testing + +#![allow(dead_code)] + +use axum::{ + body::Body, + extract::{Request, State}, + http::{HeaderValue, StatusCode}, + response::sse::{Event, KeepAlive}, + response::{IntoResponse, Response, Sse}, + routing::post, + Json, Router, +}; +use futures_util::stream::{self, StreamExt}; +use serde_json::json; +use std::collections::HashMap; +use std::net::SocketAddr; +use std::sync::{Arc, Mutex, OnceLock}; +use tokio::net::TcpListener; + +/// Mock OpenAI API server for testing +pub struct MockOpenAIServer { + addr: SocketAddr, + _handle: tokio::task::JoinHandle<()>, +} + +#[derive(Clone)] +struct MockServerState { + require_auth: bool, + expected_auth: Option, + port: u16, +} + +impl MockOpenAIServer { + /// Create and start a new mock OpenAI server + pub async fn new() -> Self { + Self::new_with_auth(None).await + } + + /// Create and start a new mock OpenAI server with optional auth requirement + pub async fn new_with_auth(expected_auth: Option) -> Self { + let listener = TcpListener::bind("127.0.0.1:0").await.unwrap(); + let addr = listener.local_addr().unwrap(); + + let state = Arc::new(MockServerState { + require_auth: expected_auth.is_some(), + expected_auth, + port: addr.port(), + }); + + let app = Router::new() + .route("/v1/chat/completions", post(mock_chat_completions)) + .route("/v1/completions", post(mock_completions)) + .route("/v1/models", post(mock_models).get(mock_models)) + .with_state(state); + + let handle = tokio::spawn(async move { + axum::serve(listener, app).await.unwrap(); + }); + + // Give the server a moment to start + tokio::time::sleep(tokio::time::Duration::from_millis(10)).await; + + Self { + addr, + _handle: handle, + } + } + + /// Get the base URL for this mock server + pub fn base_url(&self) -> String { + format!("http://{}", self.addr) + } + + pub fn clear_captured_requests(&self) { + clear_captured_requests(self.addr.port()); + } + + pub fn captured_requests(&self) -> Vec { + get_captured_requests(self.addr.port()) + } +} + +#[derive(Debug, Clone)] +pub struct CapturedRequest { + pub path: String, + pub headers: HashMap, +} + +static REQ_CAPTURE_STORE: OnceLock>>> = OnceLock::new(); + +fn get_capture_store() -> &'static Mutex>> { + REQ_CAPTURE_STORE.get_or_init(|| Mutex::new(HashMap::new())) +} + +fn capture_request(port: u16, path: &str, headers: &axum::http::HeaderMap) { + let captured = CapturedRequest { + path: path.to_string(), + headers: headers + .iter() + .filter_map(|(name, value)| { + value + .to_str() + .ok() + .map(|v| (name.as_str().to_string(), v.to_string())) + }) + .collect(), + }; + + let mut store = get_capture_store().lock().unwrap(); + store.entry(port).or_default().push(captured); +} + +fn clear_captured_requests(port: u16) { + let mut store = get_capture_store().lock().unwrap(); + store.remove(&port); +} + +fn get_captured_requests(port: u16) -> Vec { + let store = get_capture_store().lock().unwrap(); + store.get(&port).cloned().unwrap_or_default() +} + +/// Mock chat completions endpoint +async fn mock_chat_completions( + State(state): State>, + req: Request, +) -> Response { + capture_request(state.port, req.uri().path(), req.headers()); + let (_, body) = req.into_parts(); + let body_bytes = match axum::body::to_bytes(body, usize::MAX).await { + Ok(bytes) => bytes, + Err(_) => return StatusCode::BAD_REQUEST.into_response(), + }; + + let request: serde_json::Value = match serde_json::from_slice(&body_bytes) { + Ok(req) => req, + Err(_) => return StatusCode::BAD_REQUEST.into_response(), + }; + + // Extract model from request or use default (owned String to satisfy 'static in stream) + let model: String = request + .get("model") + .and_then(|v| v.as_str()) + .unwrap_or("gpt-3.5-turbo") + .to_string(); + + // If stream requested, return SSE + let is_stream = request + .get("stream") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + + if is_stream { + let created = 1677652288u64; + // Single chunk then [DONE] + let model_chunk = model.clone(); + let event_stream = stream::once(async move { + let chunk = json!({ + "id": "chatcmpl-123456789", + "object": "chat.completion.chunk", + "created": created, + "model": model_chunk, + "choices": [{ + "index": 0, + "delta": { + "content": "Hello!" + }, + "finish_reason": null + }] + }); + Ok::<_, std::convert::Infallible>(Event::default().data(chunk.to_string())) + }) + .chain(stream::once(async { Ok(Event::default().data("[DONE]")) })); + + Sse::new(event_stream) + .keep_alive(KeepAlive::default()) + .into_response() + } else { + // Create a mock non-streaming response + let response = json!({ + "id": "chatcmpl-123456789", + "object": "chat.completion", + "created": 1677652288, + "model": model, + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": "Hello! I'm a mock OpenAI assistant. How can I help you today?" + }, + "finish_reason": "stop" + }], + "usage": { + "prompt_tokens": 9, + "completion_tokens": 12, + "total_tokens": 21 + } + }); + + Json(response).into_response() + } +} + +/// Mock completions endpoint (legacy) +async fn mock_completions( + State(state): State>, + req: Request, +) -> Response { + capture_request(state.port, req.uri().path(), req.headers()); + let (_, body) = req.into_parts(); + let body_bytes = match axum::body::to_bytes(body, usize::MAX).await { + Ok(bytes) => bytes, + Err(_) => return StatusCode::BAD_REQUEST.into_response(), + }; + + let request: serde_json::Value = match serde_json::from_slice(&body_bytes) { + Ok(req) => req, + Err(_) => return StatusCode::BAD_REQUEST.into_response(), + }; + + let model = request["model"].as_str().unwrap_or("text-davinci-003"); + + let response = json!({ + "id": "cmpl-123456789", + "object": "text_completion", + "created": 1677652288, + "model": model, + "choices": [{ + "text": " This is a mock completion response.", + "index": 0, + "logprobs": null, + "finish_reason": "stop" + }], + "usage": { + "prompt_tokens": 5, + "completion_tokens": 7, + "total_tokens": 12 + } + }); + + Json(response).into_response() +} + +/// Mock models endpoint +async fn mock_models(State(state): State>, req: Request) -> Response { + capture_request(state.port, req.uri().path(), req.headers()); + + // Optionally enforce Authorization header + if state.require_auth { + let auth = req + .headers() + .get("authorization") + .or_else(|| req.headers().get("Authorization")) + .and_then(|v| v.to_str().ok()) + .map(|s| s.to_string()); + let auth_ok = match (&state.expected_auth, auth) { + (Some(expected), Some(got)) => &got == expected, + (None, Some(_)) => true, + _ => false, + }; + if !auth_ok { + let mut response = Response::new(Body::from( + json!({ + "error": { + "message": "Unauthorized", + "type": "invalid_request_error" + } + }) + .to_string(), + )); + *response.status_mut() = StatusCode::UNAUTHORIZED; + response + .headers_mut() + .insert("WWW-Authenticate", HeaderValue::from_static("Bearer")); + return response; + } + } + + let response = json!({ + "object": "list", + "data": [ + { + "id": "gpt-4", + "object": "model", + "created": 1677610602, + "owned_by": "openai" + }, + { + "id": "gpt-3.5-turbo", + "object": "model", + "created": 1677610602, + "owned_by": "openai" + } + ] + }); + + Json(response).into_response() +} diff --git a/lmdeploy/router/tests/common/mock_worker.rs b/lmdeploy/router/tests/common/mock_worker.rs new file mode 100644 index 0000000000..3a3e3c8e1b --- /dev/null +++ b/lmdeploy/router/tests/common/mock_worker.rs @@ -0,0 +1,890 @@ +// Mock worker for testing - these functions are used by integration tests +#![allow(dead_code)] + +use axum::{ + extract::{Json, Path, State}, + http::StatusCode, + response::sse::{Event, KeepAlive}, + response::{IntoResponse, Response, Sse}, + routing::{get, post}, + Router, +}; +use futures_util::stream::{self, StreamExt}; +use serde_json::json; +use std::collections::{HashMap, HashSet}; +use std::convert::Infallible; +use std::sync::{Arc, Mutex, OnceLock}; +use std::time::{SystemTime, UNIX_EPOCH}; +use tokio::sync::RwLock; +use uuid::Uuid; + +/// Configuration for mock worker behavior +#[derive(Clone)] +pub struct MockWorkerConfig { + pub port: u16, + pub worker_type: WorkerType, + pub health_status: HealthStatus, + pub response_delay_ms: u64, + pub fail_rate: f32, +} + +#[derive(Clone, Debug)] +pub enum WorkerType { + Regular, + Prefill, + Decode, +} + +#[derive(Clone, Debug)] +pub enum HealthStatus { + Healthy, + Unhealthy, + Degraded, +} + +/// Mock worker server for testing +pub struct MockWorker { + config: Arc>, + shutdown_handle: Option>, + shutdown_tx: Option>, +} + +impl MockWorker { + pub fn new(config: MockWorkerConfig) -> Self { + Self { + config: Arc::new(RwLock::new(config)), + shutdown_handle: None, + shutdown_tx: None, + } + } + + /// Start the mock worker server + pub async fn start(&mut self) -> Result> { + let config = self.config.clone(); + let port = config.read().await.port; + + // If port is 0, find an available port + let port = if port == 0 { + let listener = std::net::TcpListener::bind("127.0.0.1:0")?; + let port = listener.local_addr()?.port(); + drop(listener); + config.write().await.port = port; + port + } else { + port + }; + + let app = Router::new() + .route("/health", get(health_handler)) + .route("/health_generate", get(health_generate_handler)) + .route("/get_server_info", get(server_info_handler)) + .route("/get_model_info", get(model_info_handler)) + .route("/generate", post(generate_handler)) + .route("/v1/chat/completions", post(chat_completions_handler)) + .route("/v1/completions", post(completions_handler)) + .route("/v1/responses", post(responses_handler)) + .route("/v1/responses/{response_id}", get(responses_get_handler)) + .route( + "/v1/responses/{response_id}/cancel", + post(responses_cancel_handler), + ) + .route("/flush_cache", post(flush_cache_handler)) + .route("/v1/models", get(v1_models_handler)) + .with_state(config); + + let (shutdown_tx, shutdown_rx) = tokio::sync::oneshot::channel::<()>(); + self.shutdown_tx = Some(shutdown_tx); + + // Spawn the server in a separate task + let handle = tokio::spawn(async move { + let listener = match tokio::net::TcpListener::bind(("127.0.0.1", port)).await { + Ok(l) => l, + Err(e) => { + eprintln!("Failed to bind to port {}: {}", port, e); + return; + } + }; + + let server = axum::serve(listener, app).with_graceful_shutdown(async move { + let _ = shutdown_rx.await; + }); + + if let Err(e) = server.await { + eprintln!("Server error: {}", e); + } + }); + + self.shutdown_handle = Some(handle); + + // Wait for the server to start + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + + let url = format!("http://127.0.0.1:{}", port); + Ok(url) + } + + /// Stop the mock worker server + pub async fn stop(&mut self) { + if let Some(shutdown_tx) = self.shutdown_tx.take() { + let _ = shutdown_tx.send(()); + } + + if let Some(handle) = self.shutdown_handle.take() { + // Wait for the server to shut down + let _ = tokio::time::timeout(tokio::time::Duration::from_secs(5), handle).await; + } + } +} + +impl Drop for MockWorker { + fn drop(&mut self) { + // Clean shutdown when dropped + if let Some(shutdown_tx) = self.shutdown_tx.take() { + let _ = shutdown_tx.send(()); + } + } +} + +// Handler implementations + +/// Check if request should fail based on configured fail_rate +async fn should_fail(config: &MockWorkerConfig) -> bool { + rand::random::() < config.fail_rate +} + +async fn health_handler(State(config): State>>) -> Response { + let config = config.read().await; + + match config.health_status { + HealthStatus::Healthy => Json(json!({ + "status": "healthy", + "timestamp": SystemTime::now().duration_since(UNIX_EPOCH).unwrap().as_secs(), + "worker_type": format!("{:?}", config.worker_type), + })) + .into_response(), + HealthStatus::Unhealthy => ( + StatusCode::SERVICE_UNAVAILABLE, + Json(json!({ + "status": "unhealthy", + "error": "Worker is not responding" + })), + ) + .into_response(), + HealthStatus::Degraded => Json(json!({ + "status": "degraded", + "warning": "High load detected" + })) + .into_response(), + } +} + +async fn health_generate_handler(State(config): State>>) -> Response { + let config = config.read().await; + + if should_fail(&config).await { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({ + "error": "Random failure for testing" + })), + ) + .into_response(); + } + + if matches!(config.health_status, HealthStatus::Healthy) { + Json(json!({ + "status": "ok", + "queue_length": 0, + "processing_time_ms": config.response_delay_ms + })) + .into_response() + } else { + ( + StatusCode::SERVICE_UNAVAILABLE, + Json(json!({ + "error": "Generation service unavailable" + })), + ) + .into_response() + } +} + +async fn server_info_handler(State(config): State>>) -> Response { + let config = config.read().await; + + if should_fail(&config).await { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({ + "error": "Random failure for testing" + })), + ) + .into_response(); + } + + Json(json!({ + "model_path": "mock-model-path", + "port": config.port, + "host": "127.0.0.1", + "max_num_batched_tokens": 32768, + "max_prefill_tokens": 16384, + "mem_fraction_static": 0.88, + "tp_size": 1, + "dp_size": 1, + "stream_interval": 8, + "dtype": "float16", + "device": "cuda", + "enable_flashinfer": true, + "enable_p2p_check": true, + "context_length": 32768, + "disable_radix_cache": false, + "enable_torch_compile": false, + "trust_remote_code": false, + "show_time_cost": false, + "waiting_queue_size": 0, + "running_queue_size": 0, + "req_to_token_ratio": 1.2, + "min_running_requests": 0, + "max_running_requests": 2048, + "max_req_num": 8192, + "max_batch_tokens": 32768, + "schedule_policy": "lpm", + "schedule_conservativeness": 1.0, + "version": "0.3.0", + "internal_states": [{ + "waiting_queue_size": 0, + "running_queue_size": 0 + }] + })) + .into_response() +} + +async fn model_info_handler(State(config): State>>) -> Response { + let config = config.read().await; + + if should_fail(&config).await { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({ + "error": "Random failure for testing" + })), + ) + .into_response(); + } + + Json(json!({ + "model_path": "mock-model-path", + "is_generation": true, + "preferred_sampling_params": { + "temperature": 0.7, + "top_p": 0.9, + "top_k": 40, + "max_tokens": 2048 + } + })) + .into_response() +} + +async fn generate_handler( + State(config): State>>, + headers: axum::http::HeaderMap, + Json(payload): Json, +) -> Response { + let config = config.read().await; + + // Capture request for test inspection + capture_request(config.port, "/generate", &headers); + + if should_fail(&config).await { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({ + "error": "Random failure for testing" + })), + ) + .into_response(); + } + + if config.response_delay_ms > 0 { + tokio::time::sleep(tokio::time::Duration::from_millis(config.response_delay_ms)).await; + } + + let is_stream = payload + .get("stream") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + let output_ids = payload + .get("input_ids") + .filter(|value| value.is_array()) + .map(|_| vec![9001, 9002]) + .unwrap_or_default(); + + if is_stream { + let stream_delay = config.response_delay_ms; + + // Check if it's a batch request + let is_batch = payload.get("text").and_then(|t| t.as_array()).is_some(); + + let batch_size = if is_batch { + payload + .get("text") + .and_then(|t| t.as_array()) + .map(|arr| arr.len()) + .unwrap_or(1) + } else { + 1 + }; + + let mut events = Vec::new(); + + // Generate events for each item in batch + for i in 0..batch_size { + let timestamp_start = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs_f64(); + + let data = json!({ + "text": format!("Mock response {}", i + 1), + "output_ids": output_ids, + "meta_info": { + "prompt_tokens": 10, + "completion_tokens": 5, + "completion_tokens_wo_jump_forward": 5, + "input_token_logprobs": null, + "output_token_logprobs": null, + "first_token_latency": stream_delay as f64 / 1000.0, + "time_to_first_token": stream_delay as f64 / 1000.0, + "time_per_output_token": 0.01, + "end_time": timestamp_start + (stream_delay as f64 / 1000.0), + "start_time": timestamp_start, + "finish_reason": { + "type": "stop", + "reason": "length" + } + }, + "stage": "mid" + }); + + events.push(Ok::<_, Infallible>(Event::default().data(data.to_string()))); + } + + // Add [DONE] event + events.push(Ok(Event::default().data("[DONE]"))); + + let stream = stream::iter(events); + + Sse::new(stream) + .keep_alive(KeepAlive::default()) + .into_response() + } else { + Json(json!({ + "text": "This is a mock response.", + "output_ids": output_ids, + "meta_info": { + "prompt_tokens": 10, + "completion_tokens": 5, + "completion_tokens_wo_jump_forward": 5, + "input_token_logprobs": null, + "output_token_logprobs": null, + "first_token_latency": config.response_delay_ms as f64 / 1000.0, + "time_to_first_token": config.response_delay_ms as f64 / 1000.0, + "time_per_output_token": 0.01, + "finish_reason": { + "type": "stop", + "reason": "length" + } + } + })) + .into_response() + } +} + +async fn chat_completions_handler( + State(config): State>>, + headers: axum::http::HeaderMap, + Json(payload): Json, +) -> Response { + let config = config.read().await; + + // Capture request for test inspection + capture_request(config.port, "/v1/chat/completions", &headers); + + if should_fail(&config).await { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({ + "error": { + "message": "Random failure for testing", + "type": "internal_error", + "code": "internal_error" + } + })), + ) + .into_response(); + } + + if config.response_delay_ms > 0 { + tokio::time::sleep(tokio::time::Duration::from_millis(config.response_delay_ms)).await; + } + + let is_stream = payload + .get("stream") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + let output_ids = (payload + .get("return_token_ids") + .and_then(|value| value.as_bool()) + .unwrap_or(false) + && payload.get("input_ids").is_some()) + .then_some(vec![9001, 9002]); + + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(); + + if is_stream { + let request_id = format!("chatcmpl-{}", Uuid::new_v4()); + + let stream = stream::once(async move { + let chunk = json!({ + "id": request_id, + "object": "chat.completion.chunk", + "created": timestamp, + "model": "mock-model", + "choices": [{ + "index": 0, + "delta": { + "content": "This is a mock chat response." + }, + "output_ids": output_ids, + "finish_reason": null + }] + }); + + Ok::<_, Infallible>(Event::default().data(chunk.to_string())) + }) + .chain(stream::once(async { Ok(Event::default().data("[DONE]")) })); + + Sse::new(stream) + .keep_alive(KeepAlive::default()) + .into_response() + } else { + Json(json!({ + "id": format!("chatcmpl-{}", Uuid::new_v4()), + "object": "chat.completion", + "created": timestamp, + "model": "mock-model", + "choices": [{ + "index": 0, + "message": { + "role": "assistant", + "content": "This is a mock chat response." + }, + "output_ids": output_ids, + "finish_reason": "stop" + }], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15 + } + })) + .into_response() + } +} + +async fn completions_handler( + State(config): State>>, + headers: axum::http::HeaderMap, + Json(payload): Json, +) -> Response { + let config = config.read().await; + + // Capture request for test inspection + capture_request(config.port, "/v1/completions", &headers); + + if should_fail(&config).await { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({ + "error": { + "message": "Random failure for testing", + "type": "internal_error", + "code": "internal_error" + } + })), + ) + .into_response(); + } + + if config.response_delay_ms > 0 { + tokio::time::sleep(tokio::time::Duration::from_millis(config.response_delay_ms)).await; + } + + let is_stream = payload + .get("stream") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(); + + if is_stream { + let request_id = format!("cmpl-{}", Uuid::new_v4()); + + let stream = stream::once(async move { + let chunk = json!({ + "id": request_id, + "object": "text_completion", + "created": timestamp, + "model": "mock-model", + "choices": [{ + "text": "This is a mock completion.", + "index": 0, + "logprobs": null, + "finish_reason": null + }] + }); + + Ok::<_, Infallible>(Event::default().data(chunk.to_string())) + }) + .chain(stream::once(async { Ok(Event::default().data("[DONE]")) })); + + Sse::new(stream) + .keep_alive(KeepAlive::default()) + .into_response() + } else { + Json(json!({ + "id": format!("cmpl-{}", Uuid::new_v4()), + "object": "text_completion", + "created": timestamp, + "model": "mock-model", + "choices": [{ + "text": "This is a mock completion.", + "index": 0, + "logprobs": null, + "finish_reason": "stop" + }], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 5, + "total_tokens": 15 + } + })) + .into_response() + } +} + +async fn responses_handler( + State(config): State>>, + Json(payload): Json, +) -> Response { + let config = config.read().await; + + if should_fail(&config).await { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({ + "error": { + "message": "Random failure for testing", + "type": "internal_error", + "code": "internal_error" + } + })), + ) + .into_response(); + } + + if config.response_delay_ms > 0 { + tokio::time::sleep(tokio::time::Duration::from_millis(config.response_delay_ms)).await; + } + + let is_stream = payload + .get("stream") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs() as i64; + + // Background storage simulation + let is_background = payload + .get("background") + .and_then(|v| v.as_bool()) + .unwrap_or(false); + let req_id = payload + .get("request_id") + .and_then(|v| v.as_str()) + .map(|s| s.to_string()); + if is_background { + if let Some(id) = &req_id { + store_response_for_port(config.port, id); + } + } + + if is_stream { + let request_id = format!("resp-{}", Uuid::new_v4()); + + let stream = stream::once(async move { + let chunk = json!({ + "id": request_id, + "object": "response", + "created_at": timestamp, + "model": "mock-model", + "status": "in_progress", + "output": [{ + "type": "message", + "role": "assistant", + "content": [{ + "type": "output_text", + "text": "This is a mock responses streamed output." + }] + }] + }); + Ok::<_, Infallible>(Event::default().data(chunk.to_string())) + }) + .chain(stream::once(async { Ok(Event::default().data("[DONE]")) })); + + Sse::new(stream) + .keep_alive(KeepAlive::default()) + .into_response() + } else if is_background { + let rid = req_id.unwrap_or_else(|| format!("resp-{}", Uuid::new_v4())); + Json(json!({ + "id": rid, + "object": "response", + "created_at": timestamp, + "model": "mock-model", + "output": [], + "status": "queued", + "usage": null + })) + .into_response() + } else { + Json(json!({ + "id": format!("resp-{}", Uuid::new_v4()), + "object": "response", + "created_at": timestamp, + "model": "mock-model", + "output": [{ + "type": "message", + "role": "assistant", + "content": [{ + "type": "output_text", + "text": "This is a mock responses output." + }] + }], + "status": "completed", + "usage": { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15 + } + })) + .into_response() + } +} + +async fn flush_cache_handler(State(config): State>>) -> Response { + let config = config.read().await; + + if should_fail(&config).await { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({ + "error": "Random failure for testing" + })), + ) + .into_response(); + } + + Json(json!({ + "message": "Cache flushed successfully" + })) + .into_response() +} + +async fn v1_models_handler(State(config): State>>) -> Response { + let config = config.read().await; + + if should_fail(&config).await { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({ + "error": { + "message": "Random failure for testing", + "type": "internal_error", + "code": "internal_error" + } + })), + ) + .into_response(); + } + + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs(); + + Json(json!({ + "object": "list", + "data": [{ + "id": "mock-model", + "object": "model", + "created": timestamp, + "owned_by": "organization-owner" + }] + })) + .into_response() +} + +async fn responses_get_handler( + State(config): State>>, + Path(response_id): Path, +) -> Response { + let config = config.read().await; + if should_fail(&config).await { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({ "error": "Random failure for testing" })), + ) + .into_response(); + } + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs() as i64; + // Only return 200 if this worker "stores" the response id + if response_exists_for_port(config.port, &response_id) { + Json(json!({ + "id": response_id, + "object": "response", + "created_at": timestamp, + "model": "mock-model", + "output": [], + "status": "completed", + "usage": { + "input_tokens": 0, + "output_tokens": 0, + "total_tokens": 0 + } + })) + .into_response() + } else { + StatusCode::NOT_FOUND.into_response() + } +} + +async fn responses_cancel_handler( + State(config): State>>, + Path(response_id): Path, +) -> Response { + let config = config.read().await; + if should_fail(&config).await { + return ( + StatusCode::INTERNAL_SERVER_ERROR, + Json(json!({ "error": "Random failure for testing" })), + ) + .into_response(); + } + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap() + .as_secs() as i64; + if response_exists_for_port(config.port, &response_id) { + Json(json!({ + "id": response_id, + "object": "response", + "created_at": timestamp, + "model": "mock-model", + "output": [], + "status": "cancelled", + "usage": null + })) + .into_response() + } else { + StatusCode::NOT_FOUND.into_response() + } +} + +// --- Simple in-memory response store per worker port (for tests) --- +static RESP_STORE: OnceLock>>> = OnceLock::new(); + +fn get_store() -> &'static Mutex>> { + RESP_STORE.get_or_init(|| Mutex::new(HashMap::new())) +} + +fn store_response_for_port(port: u16, response_id: &str) { + let mut map = get_store().lock().unwrap(); + map.entry(port).or_default().insert(response_id.to_string()); +} + +fn response_exists_for_port(port: u16, response_id: &str) -> bool { + let map = get_store().lock().unwrap(); + map.get(&port) + .map(|set| set.contains(response_id)) + .unwrap_or(false) +} + +impl Default for MockWorkerConfig { + fn default() -> Self { + Self { + port: 0, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + } + } +} + +// --- Request header capture --- + +/// A captured request with headers and path +#[derive(Debug, Clone)] +pub struct CapturedRequest { + pub path: String, + pub headers: HashMap, +} + +static REQ_CAPTURE_STORE: OnceLock>>> = OnceLock::new(); + +fn get_capture_store() -> &'static Mutex>> { + REQ_CAPTURE_STORE.get_or_init(|| Mutex::new(HashMap::new())) +} + +/// Record a request for a given worker port +pub fn capture_request(port: u16, path: &str, headers: &axum::http::HeaderMap) { + let captured = CapturedRequest { + path: path.to_string(), + headers: headers + .iter() + .filter_map(|(name, value)| { + value + .to_str() + .ok() + .map(|v| (name.as_str().to_string(), v.to_string())) + }) + .collect(), + }; + let mut store = get_capture_store().lock().unwrap(); + store.entry(port).or_default().push(captured); +} + +/// Get all captured requests for a given worker port +pub fn get_captured_requests(port: u16) -> Vec { + let store = get_capture_store().lock().unwrap(); + store.get(&port).cloned().unwrap_or_default() +} + +/// Clear captured requests for a given worker port +pub fn clear_captured_requests(port: u16) { + let mut store = get_capture_store().lock().unwrap(); + store.remove(&port); +} diff --git a/lmdeploy/router/tests/common/mod.rs b/lmdeploy/router/tests/common/mod.rs new file mode 100644 index 0000000000..7dacda334e --- /dev/null +++ b/lmdeploy/router/tests/common/mod.rs @@ -0,0 +1,104 @@ +// These modules are used by tests and benchmarks +#![allow(dead_code)] + +// pub mod mock_mcp_server; +pub mod mock_openai_server; +pub mod mock_worker; +pub mod test_app; + +use lmdeploy_router_rs::config::RouterConfig; +use lmdeploy_router_rs::server::AppContext; +use std::fs; +use std::path::PathBuf; +use std::sync::{Arc, Mutex, OnceLock}; + +/// Helper function to create AppContext for tests +pub fn create_test_context(config: RouterConfig) -> Arc { + Arc::new( + AppContext::new( + config.clone(), + reqwest::Client::new(), + config.max_concurrent_requests, + config.rate_limit_tokens_per_second, + config.api_key_validation_urls.clone(), + ) + .expect("Failed to create AppContext in test"), + ) +} + +// Tokenizer download configuration +const TINYLLAMA_TOKENIZER_URL: &str = + "https://huggingface.co/TinyLlama/TinyLlama-1.1B-Chat-v1.0/resolve/main/tokenizer.json"; +const CACHE_DIR: &str = ".tokenizer_cache"; +const TINYLLAMA_TOKENIZER_FILENAME: &str = "tinyllama_tokenizer.json"; + +// Global mutex to prevent concurrent downloads +static DOWNLOAD_MUTEX: OnceLock> = OnceLock::new(); + +/// Downloads the TinyLlama tokenizer from HuggingFace if not already cached. +/// Returns the path to the cached tokenizer file. +/// +/// This function is thread-safe and will only download the tokenizer once +/// even if called from multiple threads concurrently. +pub fn ensure_tokenizer_cached() -> PathBuf { + // Get or initialize the mutex + let mutex = DOWNLOAD_MUTEX.get_or_init(|| Mutex::new(())); + + // Lock to ensure only one thread downloads at a time + let _guard = mutex.lock().unwrap(); + + let cache_dir = PathBuf::from(CACHE_DIR); + let tokenizer_path = cache_dir.join(TINYLLAMA_TOKENIZER_FILENAME); + + // Create cache directory if it doesn't exist + if !cache_dir.exists() { + fs::create_dir_all(&cache_dir).expect("Failed to create cache directory"); + } + + // Download tokenizer if not already cached + if !tokenizer_path.exists() { + println!("Downloading TinyLlama tokenizer from HuggingFace..."); + + // Use blocking reqwest client since we're in tests/benchmarks + let client = reqwest::blocking::Client::new(); + let response = client + .get(TINYLLAMA_TOKENIZER_URL) + .send() + .expect("Failed to download tokenizer"); + + if !response.status().is_success() { + panic!("Failed to download tokenizer: HTTP {}", response.status()); + } + + let content = response.bytes().expect("Failed to read tokenizer content"); + + // Verify we got actual JSON content + if content.len() < 100 { + panic!("Downloaded content too small: {} bytes", content.len()); + } + + fs::write(&tokenizer_path, content).expect("Failed to write tokenizer to cache"); + println!( + "Tokenizer downloaded and cached successfully ({} bytes)", + tokenizer_path.metadata().unwrap().len() + ); + } + + tokenizer_path +} + +/// Common test prompts for consistency across tests +pub const TEST_PROMPTS: [&str; 4] = [ + "deep learning is", + "Deep learning is", + "has anyone seen nemo lately", + "another prompt", +]; + +/// Pre-computed hashes for verification +pub const EXPECTED_HASHES: [u64; 4] = [ + 1209591529327510910, + 4181375434596349981, + 6245658446118930933, + 5097285695902185237, +]; diff --git a/lmdeploy/router/tests/common/test_app.rs b/lmdeploy/router/tests/common/test_app.rs new file mode 100644 index 0000000000..77a73e86b1 --- /dev/null +++ b/lmdeploy/router/tests/common/test_app.rs @@ -0,0 +1,68 @@ +use axum::Router; +use lmdeploy_router_rs::{ + config::RouterConfig, + otel_trace, + routers::RouterTrait, + server::{build_app_with_request_tracing, AppContext, AppState}, +}; +use reqwest::Client; +use std::sync::Arc; + +/// Create a test Axum application using the actual server's build_app function +#[allow(dead_code)] +pub fn create_test_app( + router: Arc, + client: Client, + router_config: &RouterConfig, +) -> Router { + create_test_app_with_tracing(router, client, router_config, otel_trace::is_otel_enabled()) +} + +/// Create a test Axum application with an explicit request-tracing toggle. +#[allow(dead_code)] +pub fn create_test_app_with_tracing( + router: Arc, + client: Client, + router_config: &RouterConfig, + enable_request_tracing: bool, +) -> Router { + // Create AppContext + let app_context = Arc::new( + AppContext::new( + router_config.clone(), + client, + router_config.max_concurrent_requests, + router_config.rate_limit_tokens_per_second, + router_config.api_key_validation_urls.clone(), + ) + .expect("Failed to create AppContext in test"), + ); + + // Create AppState with the test router and context + let app_state = Arc::new(AppState { + router, + context: app_context, + concurrency_queue_tx: None, + router_manager: None, + }); + + // Configure request ID headers (use defaults if not specified) + let request_id_headers = router_config.request_id_headers.clone().unwrap_or_else(|| { + vec![ + "x-request-id".to_string(), + "x-correlation-id".to_string(), + "x-trace-id".to_string(), + "request-id".to_string(), + ] + }); + + // Use the actual server's build_app function + build_app_with_request_tracing( + app_state, + router_config.max_payload_size, + request_id_headers, + router_config.cors_allowed_origins.clone(), + true, // enable_transparent_proxy + enable_request_tracing, + ) +} diff --git a/lmdeploy/router/tests/lmdeploy_registration_test.rs b/lmdeploy/router/tests/lmdeploy_registration_test.rs new file mode 100644 index 0000000000..64ecc0bd8b --- /dev/null +++ b/lmdeploy/router/tests/lmdeploy_registration_test.rs @@ -0,0 +1,208 @@ +use axum::{ + body::{to_bytes, Body}, + http::{Request, StatusCode}, + Router, +}; +use lmdeploy_router_rs::{ + config::{LMDeployMigrationProtocol, RouterConfig, RoutingMode}, + core::{BasicWorker, WorkerType}, + routers::{http::lmdeploy_pd_router::LMDeployPDRouter, http::router::Router as HttpRouter}, + server::{build_app_with_request_tracing, AppContext, AppState}, +}; +use reqwest::Client; +use serde_json::{json, Value}; +use std::sync::Arc; +use tower::ServiceExt; + +async fn regular_app() -> (Router, Arc) { + let config = RouterConfig { + mode: RoutingMode::Regular { + worker_urls: Vec::new(), + }, + ..RouterConfig::default() + }; + let context = Arc::new( + AppContext::new( + config.clone(), + Client::new(), + config.max_concurrent_requests, + config.rate_limit_tokens_per_second, + Vec::new(), + ) + .unwrap(), + ); + let router = Arc::new(HttpRouter::new(Vec::new(), &context).await.unwrap()); + let state = Arc::new(AppState { + router, + context: context.clone(), + concurrency_queue_tx: None, + router_manager: None, + }); + ( + build_app_with_request_tracing(state, 1024 * 1024, Vec::new(), Vec::new(), true, false), + context, + ) +} + +async fn lmdeploy_pd_app() -> (Router, Arc) { + let config = RouterConfig { + mode: RoutingMode::LMDeployPrefillDecode { + prefill_urls: Vec::new(), + decode_urls: Vec::new(), + prefill_policy: None, + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::Rdma, + rdma_config: None, + dummy_prefill: false, + }, + ..RouterConfig::default() + }; + let context = Arc::new( + AppContext::new( + config.clone(), + Client::new(), + config.max_concurrent_requests, + config.rate_limit_tokens_per_second, + Vec::new(), + ) + .unwrap(), + ); + let router = Arc::new( + LMDeployPDRouter::new( + Vec::new(), + Vec::new(), + LMDeployMigrationProtocol::Rdma, + None, + false, + &context, + ) + .await + .unwrap(), + ); + let state = Arc::new(AppState { + router, + context: context.clone(), + concurrency_queue_tx: None, + router_manager: None, + }); + ( + build_app_with_request_tracing(state, 1024 * 1024, Vec::new(), Vec::new(), true, false), + context, + ) +} + +async fn json_request(app: Router, method: &str, uri: &str, body: Value) -> (StatusCode, Value) { + let response = app + .oneshot( + Request::builder() + .method(method) + .uri(uri) + .header("content-type", "application/json") + .body(Body::from(body.to_string())) + .unwrap(), + ) + .await + .unwrap(); + let status = response.status(); + let body = to_bytes(response.into_body(), usize::MAX).await.unwrap(); + let body = serde_json::from_slice(&body) + .unwrap_or_else(|_| Value::String(String::from_utf8_lossy(&body).into_owned())); + (status, body) +} + +#[tokio::test] +async fn lmdeploy_hybrid_registration_is_immediate_and_idempotent() { + let (app, context) = regular_app().await; + let url = "http://127.0.0.1:49151"; + let node = json!({ + "url": url, + "status": {"models": ["test-model"], "role": 1} + }); + + let (status, body) = json_request(app.clone(), "POST", "/nodes/add", node.clone()).await; + assert_eq!(status, StatusCode::OK); + assert_eq!(body, json!("Added successfully")); + assert_eq!(context.worker_registry.get_all().len(), 1); + + let (status, _) = json_request(app.clone(), "POST", "/nodes/add", node).await; + assert_eq!(status, StatusCode::OK); + assert_eq!(context.worker_registry.get_all().len(), 1); + + let (status, body) = json_request(app.clone(), "GET", "/nodes/status", Value::Null).await; + assert_eq!(status, StatusCode::OK); + assert_eq!(body[url]["role"], 1); + assert_eq!(body[url]["models"], json!(["test-model"])); + assert_eq!(body[url]["unfinished"], 0); + assert_eq!(body[url]["latency"], json!([])); + assert!(body[url]["speed"].is_null()); + + let (status, body) = json_request(app, "POST", "/nodes/remove", json!({"url": url})).await; + assert_eq!(status, StatusCode::OK); + assert_eq!(body, json!("Deleted successfully")); + assert!(context.worker_registry.get_all().is_empty()); +} + +#[tokio::test] +async fn nodes_status_includes_registry_workers() { + let (app, context) = regular_app().await; + let workers = [ + ("http://127.0.0.1:49152", WorkerType::Regular, 1), + ("http://127.0.0.1:49153", WorkerType::Prefill, 2), + ("http://127.0.0.1:49154", WorkerType::Decode, 3), + ]; + + for (url, worker_type, _) in &workers { + context.worker_registry.register(Arc::new(BasicWorker::new( + (*url).to_string(), + worker_type.clone(), + ))); + } + + let (status, body) = json_request(app, "GET", "/nodes/status", Value::Null).await; + + assert_eq!(status, StatusCode::OK); + for (url, _, role) in workers { + assert_eq!(body[url]["role"], role); + assert_eq!(body[url]["models"], json!([])); + assert_eq!(body[url]["unfinished"], 0); + assert_eq!(body[url]["latency"], json!([])); + assert!(body[url]["speed"].is_null()); + } +} + +#[tokio::test] +async fn lmdeploy_pd_registration_dispatches_prefill_and_decode_roles() { + let (app, context) = lmdeploy_pd_app().await; + let prefill = "http://127.0.0.1:49152"; + let decode = "http://127.0.0.1:49153"; + + for (url, role) in [(prefill, 2), (decode, 3)] { + let (status, body) = json_request( + app.clone(), + "POST", + "/nodes/add", + json!({ + "url": url, + "status": {"models": ["test-model"], "role": role} + }), + ) + .await; + assert_eq!(status, StatusCode::OK); + assert_eq!(body, json!("Added successfully")); + } + + assert_eq!(context.worker_registry.get_prefill_workers().len(), 1); + assert_eq!(context.worker_registry.get_decode_workers().len(), 1); + + let (status, _) = json_request( + app, + "POST", + "/nodes/add", + json!({ + "url": "http://127.0.0.1:49154", + "status": {"models": ["test-model"], "role": 1} + }), + ) + .await; + assert_eq!(status, StatusCode::BAD_REQUEST); +} diff --git a/lmdeploy/router/tests/mock_backend_server.py b/lmdeploy/router/tests/mock_backend_server.py new file mode 100755 index 0000000000..2e867880f9 --- /dev/null +++ b/lmdeploy/router/tests/mock_backend_server.py @@ -0,0 +1,219 @@ +# Copyright (c) OpenMMLab. All rights reserved. +#!/usr/bin/env python3 +"""Simple mock backend server for testing the router's transparent proxy +feature. + +Usage: + python mock_backend_server.py [--port PORT] [--host HOST] + +Example: + # Start server on port 8081 + python mock_backend_server.py --port 8081 + + # Test with curl + curl -X POST http://localhost:8081/generate \ + -H "Content-Type: application/json" \ + -d '{"prompt": "Hello world"}' +""" + +import argparse +import json +import sys +from datetime import datetime +from http.server import BaseHTTPRequestHandler, HTTPServer + + +class MockBackendHandler(BaseHTTPRequestHandler): + """Handler for mock backend server requests.""" + + def log_request_info(self, method: str): + """Log request information to stdout.""" + timestamp = datetime.now().isoformat() + print(f"\n{'='*60}") + print(f'[{timestamp}] {method} {self.path}') + print(f"{'='*60}") + print(f'Client: {self.client_address[0]}:{self.client_address[1]}') + print(f'Path: {self.path}') + print('Headers:') + for header, value in self.headers.items(): + print(f' {header}: {value}') + + def send_json_response(self, status_code: int, data: dict): + """Send a JSON response.""" + response_body = json.dumps(data, indent=2).encode('utf-8') + self.send_response(status_code) + self.send_header('Content-Type', 'application/json') + self.send_header('Content-Length', str(len(response_body))) + self.end_headers() + self.wfile.write(response_body) + + def do_GET(self): + """Handle GET requests.""" + self.log_request_info('GET') + + if self.path == '/health' or self.path == '/healthz': + print('Body: (none)') + print('\n[SUCCESS] Health check passed') + self.send_json_response(200, {'status': 'healthy'}) + + elif self.path == '/v1/models': + print('Body: (none)') + print('\n[SUCCESS] Models list requested') + self.send_json_response(200, { + 'object': 'list', + 'data': [ + { + 'id': 'mock-model', + 'object': 'model', + 'owned_by': 'mock-server' + } + ] + }) + + else: + print('Body: (none)') + print(f'\n[SUCCESS] GET request received at {self.path}') + self.send_json_response(200, { + 'message': f'GET request received at {self.path}', + 'path': self.path + }) + + def do_POST(self): + """Handle POST requests.""" + self.log_request_info('POST') + + # Read request body + content_length = int(self.headers.get('Content-Length', 0)) + body = self.rfile.read(content_length).decode('utf-8') if content_length > 0 else '' + + # Try to parse as JSON for pretty printing + try: + body_json = json.loads(body) if body else {} + print('Body (JSON):') + print(json.dumps(body_json, indent=2)) + except json.JSONDecodeError: + print(f"Body (raw): {body[:500]}{'...' if len(body) > 500 else ''}") + body_json = {'raw': body} + + # Handle different endpoints + if self.path == '/generate': + print('\n[SUCCESS] /generate endpoint called') + response = { + 'text': 'This is a mock response from the generate endpoint.', + 'prompt': body_json.get('prompt', ''), + 'model': body_json.get('model', 'mock-model'), + 'usage': { + 'prompt_tokens': 10, + 'completion_tokens': 15, + 'total_tokens': 25 + } + } + self.send_json_response(200, response) + + elif self.path == '/v1/chat/completions': + print('\n[SUCCESS] /v1/chat/completions endpoint called') + response = { + 'id': 'mock-chat-completion', + 'object': 'chat.completion', + 'model': body_json.get('model', 'mock-model'), + 'choices': [ + { + 'index': 0, + 'message': { + 'role': 'assistant', + 'content': 'This is a mock response from chat completions.' + }, + 'finish_reason': 'stop' + } + ], + 'usage': { + 'prompt_tokens': 10, + 'completion_tokens': 12, + 'total_tokens': 22 + } + } + self.send_json_response(200, response) + + elif self.path == '/v1/completions': + print('\n[SUCCESS] /v1/completions endpoint called') + response = { + 'id': 'mock-completion', + 'object': 'text_completion', + 'model': body_json.get('model', 'mock-model'), + 'choices': [ + { + 'index': 0, + 'text': 'This is a mock completion response.', + 'finish_reason': 'stop' + } + ], + 'usage': { + 'prompt_tokens': 5, + 'completion_tokens': 8, + 'total_tokens': 13 + } + } + self.send_json_response(200, response) + + else: + # Handle any other POST endpoint (transparent proxy test) + print(f'\n[SUCCESS] Custom endpoint {self.path} called (transparent proxy)') + response = { + 'message': f'POST request received at {self.path}', + 'path': self.path, + 'body_received': body_json, + 'server': 'mock-backend-server' + } + self.send_json_response(200, response) + + def log_message(self, format, *args): + """Suppress default logging (we do our own).""" + pass + + +def main(): + parser = argparse.ArgumentParser( + description='Mock backend server for testing router transparent proxy' + ) + parser.add_argument( + '--port', '-p', + type=int, + default=8081, + help='Port to listen on (default: 8081)' + ) + parser.add_argument( + '--host', '-H', + type=str, + default='0.0.0.0', + help='Host to bind to (default: 0.0.0.0)' + ) + args = parser.parse_args() + + server_address = (args.host, args.port) + httpd = HTTPServer(server_address, MockBackendHandler) + + print('Mock Backend Server') + print('================') + print(f'Listening on {args.host}:{args.port}') + print('') + print('Available endpoints:') + print(' GET /health - Health check') + print(' GET /v1/models - List models') + print(' POST /generate - Generate endpoint') + print(' POST /v1/completions - Completions endpoint') + print(' POST /v1/chat/completions - Chat completions endpoint') + print(' POST /* - Any other path (transparent proxy test)') + print('') + print('Press Ctrl+C to stop') + print('') + + try: + httpd.serve_forever() + except KeyboardInterrupt: + print('\nShutting down...') + httpd.shutdown() + sys.exit(0) + + +if __name__ == '__main__': + main() diff --git a/lmdeploy/router/tests/otel_config_test.rs b/lmdeploy/router/tests/otel_config_test.rs new file mode 100644 index 0000000000..e5e5f59d7f --- /dev/null +++ b/lmdeploy/router/tests/otel_config_test.rs @@ -0,0 +1,60 @@ +use axum::http::Request; +use lmdeploy_router_rs::{ + config::TraceConfig, + middleware::RequestSpan, + otel_trace::{self, PreparedOtel}, +}; +use opentelemetry::trace::TracerProvider as _; +use opentelemetry_sdk::{testing::trace::InMemorySpanExporter, trace::TracerProvider}; +use tower_http::trace::MakeSpan; +use tracing_subscriber::layer::SubscriberExt; + +#[test] +fn test_trace_config_controls() { + let invalid_config = TraceConfig { + sampling_ratio: 1.5, + ..Default::default() + }; + assert!( + otel_trace::prepare_otel(&invalid_config).is_err(), + "sampling ratios outside [0.0, 1.0] should be rejected" + ); + + let exporter = InMemorySpanExporter::default(); + let provider = TracerProvider::builder() + .with_simple_exporter(exporter) + .build(); + let tracer = provider.tracer("test-tracer"); + let prepared = + PreparedOtel::from_parts_with_excluded_paths(provider, tracer, vec!["/health".to_string()]); + + let subscriber = tracing_subscriber::registry().with(prepared.layer()); + let _guard = tracing::subscriber::set_default(subscriber); + otel_trace::activate_otel(prepared); + + let mut request_span = RequestSpan; + + let excluded_request = Request::builder() + .method("GET") + .uri("/health") + .body(()) + .unwrap(); + let excluded_span = request_span.make_span(&excluded_request); + assert!( + excluded_span.is_none(), + "Excluded paths should not create OTel spans" + ); + + let traced_request = Request::builder() + .method("GET") + .uri("/v1/models") + .body(()) + .unwrap(); + let traced_span = request_span.make_span(&traced_request); + assert!( + !traced_span.is_none(), + "Non-excluded paths should still create OTel spans" + ); + + otel_trace::shutdown_otel(); +} diff --git a/lmdeploy/router/tests/otel_disabled_path_test.rs b/lmdeploy/router/tests/otel_disabled_path_test.rs new file mode 100644 index 0000000000..1fd1843a01 --- /dev/null +++ b/lmdeploy/router/tests/otel_disabled_path_test.rs @@ -0,0 +1,145 @@ +mod common; + +use std::sync::Arc; + +use axum::{ + body::Body, + http::{HeaderMap, HeaderValue, Request, StatusCode}, +}; +use common::{ + create_test_context, + mock_worker::{MockWorker, MockWorkerConfig}, + test_app::create_test_app_with_tracing, +}; +use lmdeploy_router_rs::{ + config::{PolicyConfig, RouterConfig, RoutingMode}, + middleware::RequestSpan, + otel_trace, + routers::{header_utils, RouterFactory}, +}; +use reqwest::Client; +use tower::ServiceExt; +use tower_http::trace::MakeSpan; + +fn test_router_config(worker_url: &str) -> RouterConfig { + RouterConfig { + mode: RoutingMode::Regular { + worker_urls: vec![worker_url.to_string()], + }, + policy: PolicyConfig::RoundRobin, + host: "127.0.0.1".to_string(), + port: 0, + ..Default::default() + } +} + +#[test] +fn request_span_is_none_when_otel_disabled() { + otel_trace::shutdown_otel(); + + let request = Request::builder() + .method("GET") + .uri("/health") + .header( + "traceparent", + "00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01", + ) + .body(()) + .unwrap(); + + let mut request_span = RequestSpan; + let span = request_span.make_span(&request); + + assert!( + span.is_none(), + "RequestSpan should return Span::none() when OTel is disabled" + ); +} + +#[test] +fn inject_trace_context_http_is_noop_when_otel_disabled() { + otel_trace::shutdown_otel(); + + let mut headers = HeaderMap::new(); + otel_trace::inject_trace_context_http(&mut headers); + + assert!( + headers.is_empty(), + "No trace context should be injected when OTel is disabled" + ); +} + +#[test] +fn propagate_trace_headers_passively_forwards_only_w3c_headers_when_otel_disabled() { + otel_trace::shutdown_otel(); + + let mut incoming = HeaderMap::new(); + incoming.insert( + "traceparent", + HeaderValue::from_static("00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"), + ); + incoming.insert("tracestate", HeaderValue::from_static("vendor=value")); + incoming.insert("baggage", HeaderValue::from_static("userId=alice")); + incoming.insert("x-extra-header", HeaderValue::from_static("ignore-me")); + + let request = header_utils::propagate_trace_headers( + reqwest::Client::new().post("http://example.com"), + Some(&incoming), + ) + .build() + .unwrap(); + + assert_eq!( + request.headers().get("traceparent"), + incoming.get("traceparent"), + "traceparent should be forwarded unchanged" + ); + assert_eq!( + request.headers().get("tracestate"), + incoming.get("tracestate"), + "tracestate should be forwarded unchanged" + ); + assert_eq!( + request.headers().get("baggage"), + incoming.get("baggage"), + "baggage should be forwarded unchanged" + ); + assert!( + request.headers().get("x-extra-header").is_none(), + "Non-trace headers should not be forwarded by propagate_trace_headers" + ); +} + +#[tokio::test(flavor = "current_thread")] +async fn app_routes_requests_without_trace_layer_when_otel_disabled() { + otel_trace::shutdown_otel(); + + let mut worker = MockWorker::new(MockWorkerConfig::default()); + let worker_url = worker.start().await.expect("Failed to start mock worker"); + + let config = test_router_config(&worker_url); + let ctx = create_test_context(config.clone()); + let router = RouterFactory::create_regular_router(std::slice::from_ref(&worker_url), &ctx) + .await + .expect("Failed to create router"); + let app = create_test_app_with_tracing(Arc::from(router), Client::new(), &config, false); + + let response = app + .oneshot( + Request::builder() + .method("GET") + .uri("/liveness") + .body(Body::empty()) + .unwrap(), + ) + .await + .expect("Request should succeed with tracing disabled"); + + assert_eq!(response.status(), StatusCode::OK); + assert!( + response.headers().contains_key("x-request-id"), + "RequestIdLayer should still run when request tracing is disabled" + ); + + worker.stop().await; +} diff --git a/lmdeploy/router/tests/otel_integration_test.rs b/lmdeploy/router/tests/otel_integration_test.rs new file mode 100644 index 0000000000..c87d8f9f21 --- /dev/null +++ b/lmdeploy/router/tests/otel_integration_test.rs @@ -0,0 +1,237 @@ +//! Integration tests for OpenTelemetry distributed tracing through the real router. +//! +//! These tests exercise the full request path: middleware -> router -> backend -> +//! header propagation. They use MockWorker as a real TCP backend server and +//! send requests through the actual axum app via tower::ServiceExt::oneshot. + +mod common; + +use axum::http::{Request, StatusCode}; +use common::mock_worker::{self, MockWorker, MockWorkerConfig}; +use common::test_app::create_test_app; +use lmdeploy_router_rs::{ + config::{PolicyConfig, RouterConfig, RoutingMode}, + otel_trace, + routers::RouterFactory, +}; +use opentelemetry::{global, propagation::TextMapCompositePropagator, trace::TracerProvider as _}; +use opentelemetry_sdk::{ + propagation::{BaggagePropagator, TraceContextPropagator}, + testing::trace::InMemorySpanExporter, + trace::TracerProvider, +}; +use tower::ServiceExt; +use tracing_subscriber::layer::SubscriberExt; + +/// Set up OTel with an in-memory exporter for one test section. +fn setup_otel_harness() -> ( + InMemorySpanExporter, + TracerProvider, + tracing::subscriber::DefaultGuard, +) { + let exporter = InMemorySpanExporter::default(); + let provider = TracerProvider::builder() + .with_simple_exporter(exporter.clone()) + .build(); + let tracer = provider.tracer("test-tracer"); + + let otel_layer = tracing_opentelemetry::layer().with_tracer(tracer); + let subscriber = tracing_subscriber::registry().with(otel_layer); + let guard = tracing::subscriber::set_default(subscriber); + + global::set_text_map_propagator(TextMapCompositePropagator::new(vec![ + Box::new(TraceContextPropagator::new()), + Box::new(BaggagePropagator::new()), + ])); + otel_trace::mark_otel_enabled(); + + (exporter, provider, guard) +} + +fn test_router_config(worker_url: &str) -> RouterConfig { + RouterConfig { + mode: RoutingMode::Regular { + worker_urls: vec![worker_url.to_string()], + }, + policy: PolicyConfig::RoundRobin, + host: "127.0.0.1".to_string(), + port: 0, + ..Default::default() + } +} + +/// All OTel integration tests run sequentially in one function to avoid +/// global state conflicts. +#[tokio::test(flavor = "current_thread")] +async fn test_otel_integration() { + // Start a real backend server + let mut worker = MockWorker::new(MockWorkerConfig::default()); + let worker_url = worker.start().await.expect("Failed to start mock worker"); + let port: u16 = worker_url.rsplit(':').next().unwrap().parse().unwrap(); + mock_worker::clear_captured_requests(port); + + let config = test_router_config(&worker_url); + let ctx = common::create_test_context(config.clone()); + + // ================================================================== + // Section 1: Span attributes on a route through the real app + // ================================================================== + { + let (exporter, provider, _guard) = setup_otel_harness(); + + let router = RouterFactory::create_regular_router(std::slice::from_ref(&worker_url), &ctx) + .await + .unwrap(); + let app = create_test_app( + std::sync::Arc::from(router), + reqwest::Client::new(), + &config, + ); + + let upstream_trace_id = "baf7651916cd43dd8448eb211c80319c"; + let upstream_span_id = "e7ad6b7169203331"; + let traceparent = format!("00-{upstream_trace_id}-{upstream_span_id}-01"); + + let request = Request::builder() + .method("GET") + .uri("/health") + .header("traceparent", &traceparent) + .body(axum::body::Body::empty()) + .unwrap(); + + let response = app.oneshot(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let _ = axum::body::to_bytes(response.into_body(), usize::MAX).await; + tokio::task::yield_now().await; + + provider.force_flush(); + let spans = exporter.get_finished_spans().unwrap(); + assert!(!spans.is_empty(), "Should have at least one exported span"); + + // Find the HTTP server span + let http_span = spans + .iter() + .find(|s| s.span_kind == opentelemetry::trace::SpanKind::Server) + .expect("Should have a SpanKind::Server span"); + + // Verify trace_id matches upstream + let expected_trace_id = opentelemetry::trace::TraceId::from_hex(upstream_trace_id).unwrap(); + assert_eq!( + http_span.span_context.trace_id(), + expected_trace_id, + "Span trace_id should match upstream" + ); + + // Verify parent_span_id matches upstream + let expected_parent_id = opentelemetry::trace::SpanId::from_hex(upstream_span_id).unwrap(); + assert_eq!( + http_span.parent_span_id, expected_parent_id, + "Span parent_span_id should match upstream" + ); + + // Verify span attributes + let attrs: std::collections::HashMap = http_span + .attributes + .iter() + .map(|kv| (kv.key.to_string(), kv.value.to_string())) + .collect(); + + assert_eq!( + attrs.get("http.request.method").map(|s| s.as_str()), + Some("GET"), + "Should record http.request.method" + ); + assert_eq!( + attrs.get("url.path").map(|s| s.as_str()), + Some("/health"), + "Should record url.path" + ); + assert!( + attrs.contains_key("http.response.status_code"), + "Should record http.response.status_code" + ); + eprintln!("PASS: Section 1 - Span attributes on health route"); + } + + // ================================================================== + // Section 2: Trace propagation to backend via chat completions + // ================================================================== + { + let (_exporter, _provider, _guard) = setup_otel_harness(); + + mock_worker::clear_captured_requests(port); + + let router = RouterFactory::create_regular_router(std::slice::from_ref(&worker_url), &ctx) + .await + .unwrap(); + let app = create_test_app( + std::sync::Arc::from(router), + reqwest::Client::new(), + &config, + ); + + let upstream_trace_id = "aaf7651916cd43dd8448eb211c80319c"; + let upstream_span_id = "c7ad6b7169203331"; + let traceparent = format!("00-{upstream_trace_id}-{upstream_span_id}-01"); + + let request = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .header("traceparent", &traceparent) + .header("baggage", "userId=alice,env=test") + .body(axum::body::Body::from( + serde_json::to_string(&serde_json::json!({ + "model": "mock-model", + "messages": [{"role": "user", "content": "Hello"}] + })) + .unwrap(), + )) + .unwrap(); + + let response = app.oneshot(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + let _ = axum::body::to_bytes(response.into_body(), usize::MAX).await; + + // Verify backend received fresh trace context + let captured = mock_worker::get_captured_requests(port); + assert!( + !captured.is_empty(), + "Backend should have received at least one request" + ); + + let backend_req = &captured[0]; + + // The backend should have traceparent with the SAME trace_id but a + // DIFFERENT span_id (the router's span, not the upstream's) + let backend_tp = backend_req + .headers + .get("traceparent") + .expect("Backend should receive traceparent header"); + let parts: Vec<&str> = backend_tp.split('-').collect(); + assert_eq!(parts.len(), 4, "traceparent should have 4 parts"); + assert_eq!( + parts[1], upstream_trace_id, + "trace_id should be preserved from upstream" + ); + assert_ne!( + parts[2], upstream_span_id, + "span_id should be the router's own span, not the upstream's" + ); + + // Baggage should be propagated + let backend_baggage = backend_req + .headers + .get("baggage") + .expect("Backend should receive baggage header"); + assert!( + backend_baggage.contains("userId=alice"), + "baggage should contain userId=alice, got: {backend_baggage}" + ); + eprintln!("PASS: Section 2 - Trace propagation to backend"); + } + + // Clean up + worker.stop().await; + otel_trace::shutdown_otel(); +} diff --git a/lmdeploy/router/tests/otel_openai_test.rs b/lmdeploy/router/tests/otel_openai_test.rs new file mode 100644 index 0000000000..4c6304eb6e --- /dev/null +++ b/lmdeploy/router/tests/otel_openai_test.rs @@ -0,0 +1,134 @@ +mod common; + +use axum::{ + body::Body, + extract::Request, + http::{HeaderMap, HeaderValue, StatusCode}, +}; +use common::mock_openai_server::MockOpenAIServer; +use lmdeploy_router_rs::{ + middleware::RequestSpan, + otel_trace, + protocols::spec::ChatCompletionRequest, + routers::{openai_router::OpenAIRouter, RouterTrait}, +}; +use opentelemetry::{global, propagation::TextMapCompositePropagator, trace::TracerProvider as _}; +use opentelemetry_sdk::{ + propagation::{BaggagePropagator, TraceContextPropagator}, + testing::trace::InMemorySpanExporter, + trace::TracerProvider, +}; +use tower_http::trace::MakeSpan; +use tracing_subscriber::layer::SubscriberExt; + +fn setup_otel_harness() -> ( + InMemorySpanExporter, + TracerProvider, + tracing::subscriber::DefaultGuard, +) { + let exporter = InMemorySpanExporter::default(); + let provider = TracerProvider::builder() + .with_simple_exporter(exporter.clone()) + .build(); + let tracer = provider.tracer("test-tracer"); + + let otel_layer = tracing_opentelemetry::layer().with_tracer(tracer); + let subscriber = tracing_subscriber::registry().with(otel_layer); + let guard = tracing::subscriber::set_default(subscriber); + + global::set_text_map_propagator(TextMapCompositePropagator::new(vec![ + Box::new(TraceContextPropagator::new()), + Box::new(BaggagePropagator::new()), + ])); + otel_trace::mark_otel_enabled(); + + (exporter, provider, guard) +} + +fn create_chat_request() -> ChatCompletionRequest { + serde_json::from_value(serde_json::json!({ + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 32 + })) + .unwrap() +} + +#[tokio::test(flavor = "current_thread")] +async fn test_openai_router_propagates_trace_context() { + let (_exporter, _provider, _guard) = setup_otel_harness(); + + let mock_server = MockOpenAIServer::new().await; + mock_server.clear_captured_requests(); + + let router = OpenAIRouter::new(mock_server.base_url(), None) + .await + .expect("OpenAI router should be created"); + + let upstream_trace_id = "daf7651916cd43dd8448eb211c80319c"; + let upstream_span_id = "a7ad6b7169203331"; + let traceparent = format!("00-{upstream_trace_id}-{upstream_span_id}-01"); + + let mut headers = HeaderMap::new(); + headers.insert("traceparent", HeaderValue::from_str(&traceparent).unwrap()); + headers.insert("baggage", HeaderValue::from_static("userId=alice,env=test")); + + let request = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("traceparent", &traceparent) + .header("baggage", "userId=alice,env=test") + .body(Body::empty()) + .unwrap(); + let mut request_span = RequestSpan; + let span = request_span.make_span(&request); + assert!( + !span.is_none(), + "OpenAI chat requests should create a real server span when OTel is enabled" + ); + + let response = { + let _enter = span.enter(); + router + .route_chat(Some(&headers), &create_chat_request(), None) + .await + }; + assert_eq!(response.status(), StatusCode::OK); + let _ = axum::body::to_bytes(response.into_body(), usize::MAX).await; + + let captured = mock_server.captured_requests(); + assert!( + !captured.is_empty(), + "Mock OpenAI server should capture at least one request" + ); + let backend_request = &captured[0]; + assert_eq!( + backend_request.path, "/v1/chat/completions", + "Trace propagation test should hit the chat completions endpoint" + ); + + let backend_traceparent = backend_request + .headers + .get("traceparent") + .expect("OpenAI upstream should receive traceparent"); + let parts: Vec<&str> = backend_traceparent.split('-').collect(); + assert_eq!(parts.len(), 4, "traceparent should have 4 parts"); + assert_eq!( + parts[1], upstream_trace_id, + "OpenAI upstream trace_id should stay on the upstream trace" + ); + assert_ne!( + parts[2], upstream_span_id, + "OpenAI upstream span_id should be a child span, not the original caller span" + ); + let baggage = backend_request + .headers + .get("baggage") + .expect("OpenAI upstream should receive baggage"); + assert!( + baggage.contains("userId=alice") && baggage.contains("env=test"), + "OpenAI upstream should receive both baggage items, got: {baggage}" + ); + + otel_trace::shutdown_otel(); +} diff --git a/lmdeploy/router/tests/otel_trace_test.rs b/lmdeploy/router/tests/otel_trace_test.rs new file mode 100644 index 0000000000..f968c44f14 --- /dev/null +++ b/lmdeploy/router/tests/otel_trace_test.rs @@ -0,0 +1,294 @@ +//! Integration tests for OpenTelemetry distributed tracing. +//! +//! These tests verify that: +//! 1. Inbound `traceparent` is extracted into the router span and the router +//! span becomes a child of the upstream caller's span. +//! 2. `inject_trace_context_http` (the real crate helper) correctly injects +//! the router's span context into outgoing backend requests. +//! 3. Inbound `baggage` survives the enabled path via the composite propagator. +//! +//! NOTE: These tests share global OTel state (propagator, ENABLED flag) and +//! must run sequentially. They are combined into a single test function to +//! avoid parallel execution issues with `set_text_map_propagator` and the +//! global ENABLED flag. + +use axum::http::{HeaderMap, HeaderValue}; +use lmdeploy_router_rs::otel_trace; +use opentelemetry::{ + global, + propagation::TextMapCompositePropagator, + trace::{SpanId, TraceId, TracerProvider as _}, +}; +use opentelemetry_sdk::{ + propagation::{BaggagePropagator, TraceContextPropagator}, + testing::trace::InMemorySpanExporter, + trace::TracerProvider, +}; +use tracing::info_span; +use tracing_opentelemetry::OpenTelemetrySpanExt; +use tracing_subscriber::layer::SubscriberExt; + +/// Helper: set up a test OTel environment with in-memory exporter. +fn setup_test_otel() -> ( + InMemorySpanExporter, + TracerProvider, + tracing::subscriber::DefaultGuard, +) { + let exporter = InMemorySpanExporter::default(); + let provider = TracerProvider::builder() + .with_simple_exporter(exporter.clone()) + .build(); + let tracer = provider.tracer("test-tracer"); + + // Use the same composite propagator as production code + global::set_text_map_propagator(TextMapCompositePropagator::new(vec![ + Box::new(TraceContextPropagator::new()), + Box::new(BaggagePropagator::new()), + ])); + + let otel_layer = tracing_opentelemetry::layer().with_tracer(tracer); + let subscriber = tracing_subscriber::registry().with(otel_layer); + let guard = tracing::subscriber::set_default(subscriber); + + // Enable the flag so real code paths (make_span, inject_trace_context_http) activate + otel_trace::mark_otel_enabled(); + + (exporter, provider, guard) +} + +/// All OTel tests run sequentially in one function to avoid global state +/// conflicts (propagator, ENABLED flag, tracer provider). +#[test] +fn test_otel_distributed_tracing() { + // ===================================================================== + // Section 1: End-to-end trace context propagation using real crate helpers + // ===================================================================== + { + let (exporter, provider, _guard) = setup_test_otel(); + + let upstream_trace_id = "0af7651916cd43dd8448eb211c80319c"; + let upstream_span_id = "b7ad6b7169203331"; + let traceparent = format!("00-{}-{}-01", upstream_trace_id, upstream_span_id); + + let mut incoming_headers = HeaderMap::new(); + incoming_headers.insert("traceparent", HeaderValue::from_str(&traceparent).unwrap()); + + // Extract parent context from headers (same as make_span does) + let parent_cx = global::get_text_map_propagator(|propagator| { + propagator.extract(&opentelemetry_http::HeaderExtractor(&incoming_headers)) + }); + + let span = info_span!( + target: "otel_trace", + "http_request", + method = "POST", + uri = "/v1/chat/completions" + ); + span.set_parent(parent_cx); + + // Outbound injection using the REAL crate helper + let outgoing_headers = { + let _enter = span.enter(); + let mut headers = HeaderMap::new(); + otel_trace::inject_trace_context_http(&mut headers); + headers + }; + + drop(span); + provider.force_flush(); + + // Verify inbound context was correctly linked + let spans = exporter.get_finished_spans().unwrap(); + assert!( + !spans.is_empty(), + "Expected at least one span to be exported" + ); + + let exported_span = &spans[0]; + let expected_trace_id = TraceId::from_hex(upstream_trace_id).unwrap(); + assert_eq!( + exported_span.span_context.trace_id(), + expected_trace_id, + "Span's trace_id should match the incoming traceparent's trace_id" + ); + + let expected_parent_id = SpanId::from_hex(upstream_span_id).unwrap(); + assert_eq!( + exported_span.parent_span_id, expected_parent_id, + "Span's parent_span_id should match the incoming traceparent's span_id" + ); + + // Verify outbound context was correctly injected + let outgoing_tp = outgoing_headers + .get("traceparent") + .expect("Expected outgoing traceparent header") + .to_str() + .unwrap(); + + let parts: Vec<&str> = outgoing_tp.split('-').collect(); + assert_eq!(parts.len(), 4, "traceparent should have 4 parts"); + assert_eq!(parts[0], "00", "version should be 00"); + assert_eq!( + parts[1], upstream_trace_id, + "trace_id in outgoing request should match the original upstream trace_id" + ); + assert_ne!( + parts[2], upstream_span_id, + "outgoing span_id should be the router's own span, not the upstream's" + ); + assert_eq!(parts[3], "01", "trace flags should indicate sampled"); + + eprintln!("PASS: Section 1 - End-to-end trace context propagation"); + } + + // ===================================================================== + // Section 2: RequestSpan::make_span exercises parent extraction path + // ===================================================================== + { + use axum::http::Request; + use lmdeploy_router_rs::middleware::RequestSpan; + use tower_http::trace::MakeSpan; + + let (exporter, provider, _guard) = setup_test_otel(); + + let upstream_trace_id = "2af7651916cd43dd8448eb211c80319c"; + let upstream_span_id = "d7ad6b7169203331"; + let traceparent = format!("00-{}-{}-01", upstream_trace_id, upstream_span_id); + + let request = Request::builder() + .method("POST") + .uri("/v1/chat/completions") + .header("traceparent", &traceparent) + .body(()) + .unwrap(); + + let mut request_span = RequestSpan; + let span = request_span.make_span(&request); + + { + let _enter = span.enter(); + } + drop(span); + provider.force_flush(); + + let spans = exporter.get_finished_spans().unwrap(); + assert!(!spans.is_empty(), "Expected at least one span"); + + // Verify parent context was extracted (this was previously skipped!) + let exported_span = &spans[0]; + let expected_trace_id = TraceId::from_hex(upstream_trace_id).unwrap(); + assert_eq!( + exported_span.span_context.trace_id(), + expected_trace_id, + "make_span should extract the upstream trace_id from traceparent" + ); + + let expected_parent_id = SpanId::from_hex(upstream_span_id).unwrap(); + assert_eq!( + exported_span.parent_span_id, expected_parent_id, + "make_span should set the upstream span as parent" + ); + + eprintln!("PASS: Section 2 - RequestSpan::make_span parent extraction"); + } + + // ===================================================================== + // Section 3: Baggage survives the enabled path + // ===================================================================== + { + let (_exporter, _provider, _guard) = setup_test_otel(); + + let mut incoming_headers = HeaderMap::new(); + incoming_headers.insert( + "traceparent", + HeaderValue::from_static("00-0af7651916cd43dd8448eb211c80319c-b7ad6b7169203331-01"), + ); + incoming_headers.insert( + "baggage", + HeaderValue::from_static("userId=alice,serverNode=DF%2028"), + ); + + let parent_cx = global::get_text_map_propagator(|propagator| { + propagator.extract(&opentelemetry_http::HeaderExtractor(&incoming_headers)) + }); + + let span = info_span!(target: "otel_trace", "http_request"); + span.set_parent(parent_cx); + + let outgoing_headers = { + let _enter = span.enter(); + let mut headers = HeaderMap::new(); + otel_trace::inject_trace_context_http(&mut headers); + headers + }; + + let baggage = outgoing_headers + .get("baggage") + .expect("baggage header should be propagated when tracing is enabled") + .to_str() + .unwrap(); + + assert!( + baggage.contains("userId=alice"), + "baggage should contain userId=alice, got: {baggage}" + ); + + eprintln!("PASS: Section 3 - Baggage survives enabled path"); + } + + // ===================================================================== + // Section 4: Client spans can be exported with proper kind and attrs + // ===================================================================== + { + let (exporter, provider, _guard) = setup_test_otel(); + + let span = info_span!( + target: "otel_trace", + "http_client_request", + otel.kind = "client", + otel.name = "POST /v1/chat/completions", + http.request.method = "POST", + http.route = "/v1/chat/completions", + url.full = "http://127.0.0.1:3000/v1/chat/completions", + ); + + { + let _enter = span.enter(); + } + drop(span); + provider.force_flush(); + + let spans = exporter.get_finished_spans().unwrap(); + let client_span = spans + .iter() + .find(|span| span.span_kind == opentelemetry::trace::SpanKind::Client) + .expect("Expected a client span to be exported"); + + let attrs: std::collections::HashMap = client_span + .attributes + .iter() + .map(|kv| (kv.key.to_string(), kv.value.to_string())) + .collect(); + + assert_eq!( + attrs.get("http.request.method").map(|value| value.as_str()), + Some("POST"), + "Client span should record the HTTP method" + ); + assert_eq!( + attrs.get("http.route").map(|value| value.as_str()), + Some("/v1/chat/completions"), + "Client span should record the logical route" + ); + assert_eq!( + attrs.get("url.full").map(|value| value.as_str()), + Some("http://127.0.0.1:3000/v1/chat/completions"), + "Client span should record the backend URL" + ); + + eprintln!("PASS: Section 4 - Client span kind and attrs"); + } + + // Clean up global state + otel_trace::shutdown_otel(); +} diff --git a/lmdeploy/router/tests/policy_registry_integration.rs b/lmdeploy/router/tests/policy_registry_integration.rs new file mode 100644 index 0000000000..93bc6ab7b9 --- /dev/null +++ b/lmdeploy/router/tests/policy_registry_integration.rs @@ -0,0 +1,153 @@ +//! Integration tests for PolicyRegistry with RouterManager + +use lmdeploy_router_rs::config::{PolicyConfig, RouterConfig}; +use lmdeploy_router_rs::core::WorkerRegistry; +use lmdeploy_router_rs::policies::PolicyRegistry; +use lmdeploy_router_rs::protocols::worker_spec::WorkerConfigRequest; +use lmdeploy_router_rs::routers::router_manager::RouterManager; +use std::collections::HashMap; +use std::sync::Arc; + +#[tokio::test] +async fn test_policy_registry_with_router_manager() { + // Create RouterConfig + let config = RouterConfig { + enable_igw: true, + policy: PolicyConfig::RoundRobin, + ..Default::default() + }; + + // Create HTTP client + let client = reqwest::Client::new(); + + // Create shared registries + let worker_registry = Arc::new(WorkerRegistry::new()); + let policy_registry = Arc::new(PolicyRegistry::new(PolicyConfig::RoundRobin)); + + // Create RouterManager with shared registries + let _router_manager = RouterManager::new( + config, + client, + worker_registry.clone(), + policy_registry.clone(), + ); + + // Test adding workers with different models and policies + + // Add first worker for llama-3 with cache_aware policy hint + let mut labels1 = HashMap::new(); + labels1.insert("policy".to_string(), "cache_aware".to_string()); + + let _worker1_config = WorkerConfigRequest { + url: "http://worker1:8000".to_string(), + model_id: Some("llama-3".to_string()), + worker_type: None, + priority: None, + cost: None, + labels: labels1, + }; + + // This would normally connect to a real worker, but for testing we'll just verify the structure + // In a real test, we'd need to mock the worker or use a test server + + // Verify PolicyRegistry has the correct policy for llama-3 + let _llama_policy = policy_registry.get_policy("llama-3"); + // After first worker is added, llama-3 should have a policy + + // Add second worker for llama-3 with different policy hint (should be ignored) + let mut labels2 = HashMap::new(); + labels2.insert("policy".to_string(), "random".to_string()); + + let _worker2_config = WorkerConfigRequest { + url: "http://worker2:8000".to_string(), + model_id: Some("llama-3".to_string()), + worker_type: None, + priority: None, + cost: None, + labels: labels2, + }; + + // The second worker should use the same policy as the first (cache_aware) + + // Add worker for different model (gpt-4) with random policy + let mut labels3 = HashMap::new(); + labels3.insert("policy".to_string(), "random".to_string()); + + let _worker3_config = WorkerConfigRequest { + url: "http://worker3:8000".to_string(), + model_id: Some("gpt-4".to_string()), + worker_type: None, + priority: None, + cost: None, + labels: labels3, + }; + + // Verify gpt-4 has random policy + let _gpt_policy = policy_registry.get_policy("gpt-4"); + + // Test removing workers + // When we remove both llama-3 workers, the policy should be cleaned up + + println!("PolicyRegistry integration test structure created"); + println!("Note: This test requires mocking or test servers to fully execute"); +} + +#[test] +fn test_policy_registry_cleanup() { + use lmdeploy_router_rs::config::PolicyConfig; + use lmdeploy_router_rs::policies::PolicyRegistry; + + let registry = PolicyRegistry::new(PolicyConfig::RoundRobin); + + // Add workers for a model + let policy1 = registry.on_worker_added("model-1", Some("cache_aware")); + assert_eq!(policy1.name(), "cache_aware"); + + // Second worker uses existing policy + let policy2 = registry.on_worker_added("model-1", Some("random")); + assert_eq!(policy2.name(), "cache_aware"); // Should still be cache_aware + + // Verify policy exists + assert!(registry.get_policy("model-1").is_some()); + + // Remove first worker - policy should remain + registry.on_worker_removed("model-1"); + assert!(registry.get_policy("model-1").is_some()); + + // Remove second worker - policy should be cleaned up + registry.on_worker_removed("model-1"); + assert!(registry.get_policy("model-1").is_none()); + + println!("✓ PolicyRegistry cleanup test passed"); +} + +#[test] +fn test_policy_registry_multiple_models() { + use lmdeploy_router_rs::config::PolicyConfig; + use lmdeploy_router_rs::policies::PolicyRegistry; + + let registry = PolicyRegistry::new(PolicyConfig::RoundRobin); + + // Add workers for different models with different policies + let llama_policy = registry.on_worker_added("llama-3", Some("cache_aware")); + let gpt_policy = registry.on_worker_added("gpt-4", Some("random")); + let mistral_policy = registry.on_worker_added("mistral", None); // Uses default + + assert_eq!(llama_policy.name(), "cache_aware"); + assert_eq!(gpt_policy.name(), "random"); + assert_eq!(mistral_policy.name(), "round_robin"); // Default + + // Verify all policies are stored + assert!(registry.get_policy("llama-3").is_some()); + assert!(registry.get_policy("gpt-4").is_some()); + assert!(registry.get_policy("mistral").is_some()); + + // Get all mappings + let mappings = registry.get_all_mappings(); + assert_eq!(mappings.len(), 3); + assert_eq!(mappings.get("llama-3").unwrap(), "cache_aware"); + assert_eq!(mappings.get("gpt-4").unwrap(), "random"); + assert_eq!(mappings.get("mistral").unwrap(), "round_robin"); + + println!("✓ PolicyRegistry multiple models test passed"); +} diff --git a/lmdeploy/router/tests/request_formats_test.rs b/lmdeploy/router/tests/request_formats_test.rs new file mode 100644 index 0000000000..1000694075 --- /dev/null +++ b/lmdeploy/router/tests/request_formats_test.rs @@ -0,0 +1,418 @@ +mod common; + +use common::mock_worker::{HealthStatus, MockWorker, MockWorkerConfig, WorkerType}; +use lmdeploy_router_rs::config::{RouterConfig, RoutingMode}; +use lmdeploy_router_rs::routers::{RouterFactory, RouterTrait}; +use reqwest::Client; +use serde_json::json; +use std::sync::Arc; + +/// Test context that manages mock workers +struct TestContext { + workers: Vec, + router: Arc, +} + +impl TestContext { + async fn new(worker_configs: Vec) -> Self { + let mut config = RouterConfig { + mode: RoutingMode::Regular { + worker_urls: vec![], + }, + port: 3003, + worker_startup_timeout_secs: 1, + worker_startup_check_interval_secs: 1, + ..Default::default() + }; + + let mut workers = Vec::new(); + let mut worker_urls = Vec::new(); + + for worker_config in worker_configs { + let mut worker = MockWorker::new(worker_config); + let url = worker.start().await.unwrap(); + worker_urls.push(url); + workers.push(worker); + } + + if !workers.is_empty() { + tokio::time::sleep(tokio::time::Duration::from_millis(200)).await; + } + + config.mode = RoutingMode::Regular { worker_urls }; + + let app_context = common::create_test_context(config); + let router = RouterFactory::create_router(&app_context).await.unwrap(); + let router = Arc::from(router); + + if !workers.is_empty() { + tokio::time::sleep(tokio::time::Duration::from_millis(500)).await; + } + + Self { workers, router } + } + + async fn shutdown(mut self) { + // Small delay to ensure any pending operations complete + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + + for worker in &mut self.workers { + worker.stop().await; + } + + // Another small delay to ensure cleanup completes + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + } + + async fn make_request( + &self, + endpoint: &str, + body: serde_json::Value, + ) -> Result { + let client = Client::new(); + + // Get any worker URL for testing + let worker_urls = self.router.get_worker_urls(); + if worker_urls.is_empty() { + return Err("No available workers".to_string()); + } + + let worker_url = &worker_urls[0]; + + let response = client + .post(format!("{}{}", worker_url, endpoint)) + .json(&body) + .send() + .await + .map_err(|e| format!("Request failed: {}", e))?; + + if !response.status().is_success() { + return Err(format!("Request failed with status: {}", response.status())); + } + + response + .json::() + .await + .map_err(|e| format!("Failed to parse response: {}", e)) + } +} + +#[cfg(test)] +mod request_format_tests { + use super::*; + use axum::body::to_bytes; + use lmdeploy_router_rs::protocols::spec::{ + ChatCompletionRequest, GenerateRequest, GenerateResponse, + }; + + #[tokio::test] + async fn test_generate_request_formats() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 19001, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + // Test 1: Basic text request + let payload = json!({ + "text": "Hello, world!", + "stream": false + }); + + let result = ctx.make_request("/generate", payload).await; + assert!(result.is_ok()); + + // Test 2: Request with sampling parameters + let payload = json!({ + "text": "Tell me a story", + "sampling_params": { + "temperature": 0.7, + "max_new_tokens": 100, + "top_p": 0.9 + }, + "stream": false + }); + + let result = ctx.make_request("/generate", payload).await; + assert!(result.is_ok()); + + // Test 3: Request with input_ids + let payload = json!({ + "input_ids": [1, 2, 3, 4, 5], + "sampling_params": { + "temperature": 0.0, + "max_new_tokens": 50 + }, + "stream": false + }); + + let result = ctx.make_request("/generate", payload).await; + assert!(result.is_ok()); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_lmdeploy_token_io_survives_router_http_path() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 0, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + let chat_request: ChatCompletionRequest = serde_json::from_value(json!({ + "model": "test-model", + "messages": [], + "input_ids": [151644, 8948, 198], + "do_preprocess": false, + "return_token_ids": true, + "stream": false + })) + .unwrap(); + let chat_response = ctx.router.route_chat(None, &chat_request, None).await; + assert!(chat_response.status().is_success()); + let chat_body = to_bytes(chat_response.into_body(), usize::MAX) + .await + .unwrap(); + let chat_json: serde_json::Value = serde_json::from_slice(&chat_body).unwrap(); + assert_eq!(chat_json["choices"][0]["output_ids"], json!([9001, 9002])); + + let generate_request: GenerateRequest = serde_json::from_value(json!({ + "input_ids": [151644, 8948, 198], + "max_tokens": 8, + "stream": false + })) + .unwrap(); + let generate_response = ctx + .router + .route_generate(None, &generate_request, None) + .await; + assert!(generate_response.status().is_success()); + let generate_body = to_bytes(generate_response.into_body(), usize::MAX) + .await + .unwrap(); + let generate: GenerateResponse = serde_json::from_slice(&generate_body).unwrap(); + assert_eq!(generate.output_ids, vec![9001, 9002]); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_v1_chat_completions_formats() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 19002, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + // Test 1: Basic chat completion + let payload = json!({ + "model": "test-model", + "messages": [ + {"role": "system", "content": "You are a helpful assistant."}, + {"role": "user", "content": "Hello!"} + ], + "stream": false + }); + + let result = ctx.make_request("/v1/chat/completions", payload).await; + assert!(result.is_ok()); + + let response = result.unwrap(); + assert!(response.get("choices").is_some()); + assert!(response.get("id").is_some()); + assert_eq!( + response.get("object").and_then(|v| v.as_str()), + Some("chat.completion") + ); + + // Test 2: Chat completion with parameters + let payload = json!({ + "model": "test-model", + "messages": [ + {"role": "user", "content": "Tell me a joke"} + ], + "temperature": 0.8, + "max_tokens": 150, + "top_p": 0.95, + "stream": false + }); + + let result = ctx.make_request("/v1/chat/completions", payload).await; + assert!(result.is_ok()); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_v1_completions_formats() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 19003, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + // Test 1: Basic completion + let payload = json!({ + "model": "test-model", + "prompt": "Once upon a time", + "max_tokens": 50, + "stream": false + }); + + let result = ctx.make_request("/v1/completions", payload).await; + assert!(result.is_ok()); + + let response = result.unwrap(); + assert!(response.get("choices").is_some()); + assert_eq!( + response.get("object").and_then(|v| v.as_str()), + Some("text_completion") + ); + + // Test 2: Completion with array prompt + let payload = json!({ + "model": "test-model", + "prompt": ["First prompt", "Second prompt"], + "temperature": 0.5, + "stream": false + }); + + let result = ctx.make_request("/v1/completions", payload).await; + assert!(result.is_ok()); + + // Test 3: Completion with logprobs + let payload = json!({ + "model": "test-model", + "prompt": "The capital of France is", + "max_tokens": 10, + "logprobs": 5, + "stream": false + }); + + let result = ctx.make_request("/v1/completions", payload).await; + assert!(result.is_ok()); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_batch_requests() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 19004, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + // Test batch text generation + let payload = json!({ + "text": ["First text", "Second text", "Third text"], + "sampling_params": { + "temperature": 0.7, + "max_new_tokens": 50 + }, + "stream": false + }); + + let result = ctx.make_request("/generate", payload).await; + assert!(result.is_ok()); + + // Test batch with input_ids + let payload = json!({ + "input_ids": [[1, 2, 3], [4, 5, 6], [7, 8, 9]], + "stream": false + }); + + let result = ctx.make_request("/generate", payload).await; + assert!(result.is_ok()); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_special_parameters() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 19005, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + // Test with return_logprob + let payload = json!({ + "text": "Test", + "return_logprob": true, + "stream": false + }); + + let result = ctx.make_request("/generate", payload).await; + assert!(result.is_ok()); + + // Test with json_schema + let payload = json!({ + "text": "Generate JSON", + "sampling_params": { + "temperature": 0.0, + "json_schema": "$$ANY$$" + }, + "stream": false + }); + + let result = ctx.make_request("/generate", payload).await; + assert!(result.is_ok()); + + // Test with ignore_eos + let payload = json!({ + "text": "Continue forever", + "sampling_params": { + "temperature": 0.7, + "max_new_tokens": 100, + "ignore_eos": true + }, + "stream": false + }); + + let result = ctx.make_request("/generate", payload).await; + assert!(result.is_ok()); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_error_handling() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 19006, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 0.0, + }]) + .await; + + // Test with empty body - should still work with mock worker + let payload = json!({}); + + let result = ctx.make_request("/generate", payload).await; + // Mock worker accepts empty body + assert!(result.is_ok()); + + ctx.shutdown().await; + } +} diff --git a/lmdeploy/router/tests/responses_api_test.rs b/lmdeploy/router/tests/responses_api_test.rs new file mode 100644 index 0000000000..8ab54185f7 --- /dev/null +++ b/lmdeploy/router/tests/responses_api_test.rs @@ -0,0 +1,210 @@ +// Integration test for Responses API + +use lmdeploy_router_rs::protocols::spec::{ + GenerationRequest, ReasoningEffort, ResponseInput, ResponseReasoningParam, ResponseStatus, + ResponseTool, ResponseToolType, ResponsesRequest, ResponsesResponse, ServiceTier, ToolChoice, + ToolChoiceValue, Truncation, UsageInfo, +}; + +#[test] +fn test_responses_request_creation() { + let request = ResponsesRequest { + background: false, + include: None, + input: ResponseInput::Text("Hello, world!".to_string()), + instructions: Some("Be helpful".to_string()), + max_output_tokens: Some(100), + max_tool_calls: None, + metadata: None, + model: Some("test-model".to_string()), + parallel_tool_calls: true, + previous_response_id: None, + reasoning: Some(ResponseReasoningParam { + effort: Some(ReasoningEffort::Medium), + }), + service_tier: ServiceTier::Auto, + store: true, + stream: false, + temperature: Some(0.7), + tool_choice: ToolChoice::Value(ToolChoiceValue::Auto), + tools: vec![ResponseTool { + r#type: ResponseToolType::WebSearchPreview, + }], + top_logprobs: 5, + top_p: Some(0.9), + truncation: Truncation::Disabled, + user: Some("test-user".to_string()), + request_id: "resp_test123".to_string(), + priority: 0, + frequency_penalty: 0.0, + presence_penalty: 0.0, + stop: None, + top_k: -1, + min_p: 0.0, + repetition_penalty: 1.0, + }; + + // Test GenerationRequest trait implementation + assert!(!request.is_stream()); + assert_eq!(request.get_model(), Some("test-model")); + let routing_text = request.extract_text_for_routing(); + assert_eq!(routing_text, "Hello, world!"); +} + +#[test] +fn test_sampling_params_conversion() { + let request = ResponsesRequest { + background: false, + include: None, + input: ResponseInput::Text("Test".to_string()), + instructions: None, + max_output_tokens: Some(50), + max_tool_calls: None, + metadata: None, + model: Some("test-model".to_string()), + parallel_tool_calls: true, // Use default true + previous_response_id: None, + reasoning: None, + service_tier: ServiceTier::Auto, + store: true, // Use default true + stream: false, + temperature: Some(0.8), + tool_choice: ToolChoice::Value(ToolChoiceValue::Auto), + tools: vec![], + top_logprobs: 0, // Use default 0 + top_p: Some(0.95), + truncation: Truncation::Auto, + user: None, + request_id: "resp_test456".to_string(), + priority: 0, + frequency_penalty: 0.1, + presence_penalty: 0.2, + stop: None, + top_k: 10, + min_p: 0.05, + repetition_penalty: 1.1, + }; + + let params = request.to_sampling_params(1000, None); + + // Check that parameters are converted correctly + assert!(params.contains_key("temperature")); + assert!(params.contains_key("top_p")); + assert!(params.contains_key("frequency_penalty")); + assert!(params.contains_key("max_new_tokens")); +} + +#[test] +fn test_responses_response_creation() { + let response = ResponsesResponse::new( + "resp_test789".to_string(), + "test-model".to_string(), + ResponseStatus::Completed, + ); + + assert_eq!(response.id, "resp_test789"); + assert_eq!(response.model, "test-model"); + assert!(response.is_complete()); + assert!(!response.is_in_progress()); + assert!(!response.is_failed()); +} + +#[test] +fn test_usage_conversion() { + let usage_info = UsageInfo::new_with_cached(15, 25, Some(8), 3); + let response_usage = usage_info.to_response_usage(); + + assert_eq!(response_usage.input_tokens, 15); + assert_eq!(response_usage.output_tokens, 25); + assert_eq!(response_usage.total_tokens, 40); + + // Check details are converted correctly + assert!(response_usage.input_tokens_details.is_some()); + assert_eq!( + response_usage + .input_tokens_details + .as_ref() + .unwrap() + .cached_tokens, + 3 + ); + + assert!(response_usage.output_tokens_details.is_some()); + assert_eq!( + response_usage + .output_tokens_details + .as_ref() + .unwrap() + .reasoning_tokens, + 8 + ); + + // Test reverse conversion + let back_to_usage = response_usage.to_usage_info(); + assert_eq!(back_to_usage.prompt_tokens, 15); + assert_eq!(back_to_usage.completion_tokens, 25); + assert_eq!(back_to_usage.reasoning_tokens, Some(8)); +} + +#[test] +fn test_reasoning_param_default() { + let param = ResponseReasoningParam { + effort: Some(ReasoningEffort::Medium), + }; + + // Test JSON serialization/deserialization preserves default + let json = serde_json::to_string(¶m).unwrap(); + let parsed: ResponseReasoningParam = serde_json::from_str(&json).unwrap(); + + assert!(matches!(parsed.effort, Some(ReasoningEffort::Medium))); +} + +#[test] +fn test_json_serialization() { + let request = ResponsesRequest { + background: true, + include: None, + input: ResponseInput::Text("Test input".to_string()), + instructions: Some("Test instructions".to_string()), + max_output_tokens: Some(200), + max_tool_calls: Some(5), + metadata: None, + model: Some("gpt-4".to_string()), + parallel_tool_calls: false, + previous_response_id: None, + reasoning: Some(ResponseReasoningParam { + effort: Some(ReasoningEffort::High), + }), + service_tier: ServiceTier::Priority, + store: false, + stream: true, + temperature: Some(0.9), + tool_choice: ToolChoice::Value(ToolChoiceValue::Required), + tools: vec![ResponseTool { + r#type: ResponseToolType::CodeInterpreter, + }], + top_logprobs: 10, + top_p: Some(0.8), + truncation: Truncation::Auto, + user: Some("test_user".to_string()), + request_id: "resp_comprehensive_test".to_string(), + priority: 1, + frequency_penalty: 0.3, + presence_penalty: 0.4, + stop: None, + top_k: 50, + min_p: 0.1, + repetition_penalty: 1.2, + }; + + // Test that everything can be serialized to JSON and back + let json = serde_json::to_string(&request).expect("Serialization should work"); + let parsed: ResponsesRequest = + serde_json::from_str(&json).expect("Deserialization should work"); + + assert_eq!(parsed.request_id, "resp_comprehensive_test"); + assert_eq!(parsed.model, Some("gpt-4".to_string())); + assert!(parsed.background); + assert!(parsed.stream); + assert_eq!(parsed.tools.len(), 1); +} diff --git a/lmdeploy/router/tests/streaming_tests.rs b/lmdeploy/router/tests/streaming_tests.rs new file mode 100644 index 0000000000..7bef49717c --- /dev/null +++ b/lmdeploy/router/tests/streaming_tests.rs @@ -0,0 +1,344 @@ +mod common; + +use common::mock_worker::{HealthStatus, MockWorker, MockWorkerConfig, WorkerType}; +use futures_util::StreamExt; +use lmdeploy_router_rs::config::{RouterConfig, RoutingMode}; +use lmdeploy_router_rs::routers::{RouterFactory, RouterTrait}; +use reqwest::Client; +use serde_json::json; +use std::sync::Arc; + +/// Test context that manages mock workers +struct TestContext { + workers: Vec, + router: Arc, +} + +impl TestContext { + async fn new(worker_configs: Vec) -> Self { + let mut config = RouterConfig { + mode: RoutingMode::Regular { + worker_urls: vec![], + }, + port: 3004, + worker_startup_timeout_secs: 1, + worker_startup_check_interval_secs: 1, + ..Default::default() + }; + + let mut workers = Vec::new(); + let mut worker_urls = Vec::new(); + + for worker_config in worker_configs { + let mut worker = MockWorker::new(worker_config); + let url = worker.start().await.unwrap(); + worker_urls.push(url); + workers.push(worker); + } + + if !workers.is_empty() { + tokio::time::sleep(tokio::time::Duration::from_millis(200)).await; + } + + config.mode = RoutingMode::Regular { worker_urls }; + + let app_context = common::create_test_context(config); + let router = RouterFactory::create_router(&app_context).await.unwrap(); + let router = Arc::from(router); + + if !workers.is_empty() { + tokio::time::sleep(tokio::time::Duration::from_millis(500)).await; + } + + Self { workers, router } + } + + async fn shutdown(mut self) { + // Small delay to ensure any pending operations complete + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + + for worker in &mut self.workers { + worker.stop().await; + } + + // Another small delay to ensure cleanup completes + tokio::time::sleep(tokio::time::Duration::from_millis(100)).await; + } + + async fn make_streaming_request( + &self, + endpoint: &str, + body: serde_json::Value, + ) -> Result, String> { + let client = Client::new(); + + // Get any worker URL for testing + let worker_urls = self.router.get_worker_urls(); + if worker_urls.is_empty() { + return Err("No available workers".to_string()); + } + + let worker_url = &worker_urls[0]; + + let response = client + .post(format!("{}{}", worker_url, endpoint)) + .json(&body) + .send() + .await + .map_err(|e| format!("Request failed: {}", e))?; + + if !response.status().is_success() { + return Err(format!("Request failed with status: {}", response.status())); + } + + // Check if it's a streaming response + let content_type = response + .headers() + .get("content-type") + .and_then(|v| v.to_str().ok()) + .unwrap_or(""); + + if !content_type.contains("text/event-stream") { + return Err("Response is not a stream".to_string()); + } + + let mut stream = response.bytes_stream(); + let mut events = Vec::new(); + + while let Some(chunk) = stream.next().await { + if let Ok(bytes) = chunk { + let text = String::from_utf8_lossy(&bytes); + for line in text.lines() { + if let Some(stripped) = line.strip_prefix("data: ") { + events.push(stripped.to_string()); + } + } + } + } + + Ok(events) + } +} + +#[cfg(test)] +mod streaming_tests { + use super::*; + + #[tokio::test] + async fn test_generate_streaming() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 20001, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 10, + fail_rate: 0.0, + }]) + .await; + + let payload = json!({ + "text": "Stream test", + "stream": true, + "sampling_params": { + "temperature": 0.7, + "max_new_tokens": 10 + } + }); + + let result = ctx.make_streaming_request("/generate", payload).await; + assert!(result.is_ok()); + + let events = result.unwrap(); + // Should have at least one data chunk and [DONE] + assert!(events.len() >= 2); + assert_eq!(events.last().unwrap(), "[DONE]"); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_v1_chat_completions_streaming() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 20002, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 10, + fail_rate: 0.0, + }]) + .await; + + let payload = json!({ + "model": "test-model", + "messages": [ + {"role": "user", "content": "Count to 3"} + ], + "stream": true, + "max_tokens": 20 + }); + + let result = ctx + .make_streaming_request("/v1/chat/completions", payload) + .await; + assert!(result.is_ok()); + + let events = result.unwrap(); + assert!(events.len() >= 2); // At least one chunk + [DONE] + + // Verify events are valid JSON (except [DONE]) + for event in &events { + if event != "[DONE]" { + let parsed: Result = serde_json::from_str(event); + assert!(parsed.is_ok(), "Invalid JSON in SSE event: {}", event); + + let json = parsed.unwrap(); + assert_eq!( + json.get("object").and_then(|v| v.as_str()), + Some("chat.completion.chunk") + ); + } + } + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_v1_completions_streaming() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 20003, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 10, + fail_rate: 0.0, + }]) + .await; + + let payload = json!({ + "model": "test-model", + "prompt": "Once upon a time", + "stream": true, + "max_tokens": 15 + }); + + let result = ctx.make_streaming_request("/v1/completions", payload).await; + assert!(result.is_ok()); + + let events = result.unwrap(); + assert!(events.len() >= 2); // At least one chunk + [DONE] + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_streaming_with_error() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 20004, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 0, + fail_rate: 1.0, // Always fail + }]) + .await; + + let payload = json!({ + "text": "This should fail", + "stream": true + }); + + let result = ctx.make_streaming_request("/generate", payload).await; + // With fail_rate: 1.0, the request should fail + assert!(result.is_err()); + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_streaming_timeouts() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 20005, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 100, // Slow response + fail_rate: 0.0, + }]) + .await; + + let payload = json!({ + "text": "Slow stream", + "stream": true, + "sampling_params": { + "max_new_tokens": 5 + } + }); + + let start = std::time::Instant::now(); + let result = ctx.make_streaming_request("/generate", payload).await; + let elapsed = start.elapsed(); + + assert!(result.is_ok()); + let events = result.unwrap(); + + // Should have received multiple chunks over time + assert!(!events.is_empty()); + assert!(elapsed.as_millis() >= 100); // At least one delay + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_batch_streaming() { + let ctx = TestContext::new(vec![MockWorkerConfig { + port: 20006, + worker_type: WorkerType::Regular, + health_status: HealthStatus::Healthy, + response_delay_ms: 10, + fail_rate: 0.0, + }]) + .await; + + // Batch request with streaming + let payload = json!({ + "text": ["First", "Second", "Third"], + "stream": true, + "sampling_params": { + "max_new_tokens": 5 + } + }); + + let result = ctx.make_streaming_request("/generate", payload).await; + assert!(result.is_ok()); + + let events = result.unwrap(); + // Should have multiple events for batch + assert!(events.len() >= 4); // At least 3 responses + [DONE] + + ctx.shutdown().await; + } + + #[tokio::test] + async fn test_sse_format_parsing() { + // Test SSE format parsing + let parse_sse_chunk = |chunk: &[u8]| -> Vec { + let text = String::from_utf8_lossy(chunk); + text.lines() + .filter(|line| line.starts_with("data: ")) + .map(|line| line[6..].to_string()) + .collect() + }; + + let sse_data = + b"data: {\"text\":\"Hello\"}\n\ndata: {\"text\":\" world\"}\n\ndata: [DONE]\n\n"; + let events = parse_sse_chunk(sse_data); + + assert_eq!(events.len(), 3); + assert_eq!(events[0], "{\"text\":\"Hello\"}"); + assert_eq!(events[1], "{\"text\":\" world\"}"); + assert_eq!(events[2], "[DONE]"); + + // Test with mixed content + let mixed = b"event: message\ndata: {\"test\":true}\n\n: comment\ndata: [DONE]\n\n"; + let events = parse_sse_chunk(mixed); + + assert_eq!(events.len(), 2); + assert_eq!(events[0], "{\"test\":true}"); + assert_eq!(events[1], "[DONE]"); + } +} diff --git a/lmdeploy/router/tests/test_chat_template.rs b/lmdeploy/router/tests/test_chat_template.rs new file mode 100644 index 0000000000..04035c691a --- /dev/null +++ b/lmdeploy/router/tests/test_chat_template.rs @@ -0,0 +1,150 @@ +#[cfg(test)] +mod tests { + use lmdeploy_router_rs::tokenizer::chat_template::{ChatMessage, ChatTemplateProcessor}; + + #[test] + fn test_chat_message_helpers() { + let system_msg = ChatMessage::system("You are a helpful assistant"); + assert_eq!(system_msg.role, "system"); + assert_eq!(system_msg.content, "You are a helpful assistant"); + + let user_msg = ChatMessage::user("Hello!"); + assert_eq!(user_msg.role, "user"); + assert_eq!(user_msg.content, "Hello!"); + + let assistant_msg = ChatMessage::assistant("Hi there!"); + assert_eq!(assistant_msg.role, "assistant"); + assert_eq!(assistant_msg.content, "Hi there!"); + } + + #[test] + fn test_llama_style_template() { + // Test a Llama-style chat template + let template = r#" +{%- if messages[0]['role'] == 'system' -%} + {%- set system_message = messages[0]['content'] -%} + {%- set messages = messages[1:] -%} +{%- else -%} + {%- set system_message = '' -%} +{%- endif -%} + +{{- bos_token }} +{%- if system_message %} +{{- '<|start_header_id|>system<|end_header_id|>\n\n' + system_message + '<|eot_id|>' }} +{%- endif %} + +{%- for message in messages %} + {{- '<|start_header_id|>' + message['role'] + '<|end_header_id|>\n\n' + message['content'] + '<|eot_id|>' }} +{%- endfor %} + +{%- if add_generation_prompt %} + {{- '<|start_header_id|>assistant<|end_header_id|>\n\n' }} +{%- endif %} +"#; + + let processor = ChatTemplateProcessor::new( + template.to_string(), + Some("<|begin_of_text|>".to_string()), + Some("<|end_of_text|>".to_string()), + ); + + let messages = vec![ + ChatMessage::system("You are a helpful assistant"), + ChatMessage::user("What is 2+2?"), + ]; + + let result = processor.apply_chat_template(&messages, true).unwrap(); + + // Check that the result contains expected markers + assert!(result.contains("<|begin_of_text|>")); + assert!(result.contains("<|start_header_id|>system<|end_header_id|>")); + assert!(result.contains("You are a helpful assistant")); + assert!(result.contains("<|start_header_id|>user<|end_header_id|>")); + assert!(result.contains("What is 2+2?")); + assert!(result.contains("<|start_header_id|>assistant<|end_header_id|>")); + } + + #[test] + fn test_chatml_template() { + // Test a ChatML-style template + let template = r#" +{%- for message in messages %} + {{- '<|im_start|>' + message['role'] + '\n' + message['content'] + '<|im_end|>\n' }} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|im_start|>assistant\n' }} +{%- endif %} +"#; + + let processor = ChatTemplateProcessor::new(template.to_string(), None, None); + + let messages = vec![ + ChatMessage::user("Hello"), + ChatMessage::assistant("Hi there!"), + ChatMessage::user("How are you?"), + ]; + + let result = processor.apply_chat_template(&messages, true).unwrap(); + + // Check ChatML format + assert!(result.contains("<|im_start|>user\nHello<|im_end|>")); + assert!(result.contains("<|im_start|>assistant\nHi there!<|im_end|>")); + assert!(result.contains("<|im_start|>user\nHow are you?<|im_end|>")); + assert!(result.ends_with("<|im_start|>assistant\n")); + } + + #[test] + fn test_template_without_generation_prompt() { + let template = r#" +{%- for message in messages -%} +{{ message.role }}: {{ message.content }} +{% endfor -%} +{%- if add_generation_prompt -%} +assistant: +{%- endif -%} +"#; + + let processor = ChatTemplateProcessor::new(template.to_string(), None, None); + + let messages = vec![ChatMessage::user("Test")]; + + // Test without generation prompt + let result = processor.apply_chat_template(&messages, false).unwrap(); + assert_eq!(result.trim(), "user: Test"); + + // Test with generation prompt + let result_with_prompt = processor.apply_chat_template(&messages, true).unwrap(); + assert!(result_with_prompt.contains("assistant:")); + } + + #[test] + fn test_template_with_special_tokens() { + let template = r#"{{ bos_token }}{% for msg in messages %}{{ msg.content }}{{ eos_token }}{% endfor %}"#; + + let processor = ChatTemplateProcessor::new( + template.to_string(), + Some("".to_string()), + Some("".to_string()), + ); + + let messages = vec![ChatMessage::user("Hello")]; + + let result = processor.apply_chat_template(&messages, false).unwrap(); + assert_eq!(result, "Hello"); + } + + #[test] + fn test_empty_messages() { + let template = + r#"{% for msg in messages %}{{ msg.role }}: {{ msg.content }}\n{% endfor %}"#; + + let processor = ChatTemplateProcessor::new(template.to_string(), None, None); + + let messages = vec![]; + let result = processor.apply_chat_template(&messages, false).unwrap(); + assert_eq!(result, ""); + } + + // Integration test with actual tokenizer file loading would go here + // but requires a real tokenizer_config.json file +} diff --git a/lmdeploy/router/tests/test_chat_template_loading.rs b/lmdeploy/router/tests/test_chat_template_loading.rs new file mode 100644 index 0000000000..5f719ab8fe --- /dev/null +++ b/lmdeploy/router/tests/test_chat_template_loading.rs @@ -0,0 +1,183 @@ +#[cfg(test)] +mod tests { + use std::fs; + use tempfile::TempDir; + + #[test] + fn test_load_chat_template_from_file() { + use lmdeploy_router_rs::tokenizer::chat_template::ChatMessage; + use lmdeploy_router_rs::tokenizer::huggingface::HuggingFaceTokenizer; + + // Create temporary directory + let temp_dir = TempDir::new().unwrap(); + let template_path = temp_dir.path().join("template.jinja"); + + // Write a test template + let template_content = r#" +{%- for message in messages %} + {{- '<|' + message['role'] + '|>' + message['content'] }} +{%- endfor %} +{%- if add_generation_prompt %} + {{- '<|assistant|>' }} +{%- endif %} +"#; + fs::write(&template_path, template_content).unwrap(); + + // Create a mock tokenizer config + let tokenizer_config = r#"{ + "version": "1.0", + "truncation": null, + "padding": null, + "added_tokens": [], + "normalizer": null, + "pre_tokenizer": { + "type": "Whitespace" + }, + "post_processor": null, + "decoder": null, + "model": { + "type": "BPE", + "vocab": { + "hello": 0, + "world": 1, + "": 2, + "": 3 + }, + "merges": [] + } + }"#; + + let tokenizer_path = temp_dir.path().join("tokenizer.json"); + fs::write(&tokenizer_path, tokenizer_config).unwrap(); + + // Load tokenizer with custom chat template + let tokenizer = HuggingFaceTokenizer::from_file_with_chat_template( + tokenizer_path.to_str().unwrap(), + Some(template_path.to_str().unwrap()), + ) + .unwrap(); + + // Test that the custom template is used + let messages = vec![ + ChatMessage::user("Hello"), + ChatMessage::assistant("Hi there"), + ]; + + let result = tokenizer.apply_chat_template(&messages, true).unwrap(); + + // Verify the custom template format + assert!(result.contains("<|user|>Hello")); + assert!(result.contains("<|assistant|>Hi there")); + assert!(result.ends_with("<|assistant|>")); + } + + #[test] + fn test_override_existing_template() { + use lmdeploy_router_rs::tokenizer::chat_template::ChatMessage; + use lmdeploy_router_rs::tokenizer::huggingface::HuggingFaceTokenizer; + + // Create temporary directory + let temp_dir = TempDir::new().unwrap(); + + // Create tokenizer config with a built-in template + let tokenizer_config_path = temp_dir.path().join("tokenizer_config.json"); + let config_with_template = r#"{ + "chat_template": "built-in: {% for msg in messages %}{{ msg.content }}{% endfor %}" + }"#; + fs::write(&tokenizer_config_path, config_with_template).unwrap(); + + // Create the actual tokenizer file + let tokenizer_json = r#"{ + "version": "1.0", + "truncation": null, + "padding": null, + "added_tokens": [], + "normalizer": null, + "pre_tokenizer": { + "type": "Whitespace" + }, + "post_processor": null, + "decoder": null, + "model": { + "type": "BPE", + "vocab": { + "test": 0, + "": 1, + "": 2 + }, + "merges": [] + } + }"#; + let tokenizer_path = temp_dir.path().join("tokenizer.json"); + fs::write(&tokenizer_path, tokenizer_json).unwrap(); + + // Create custom template that should override + let custom_template_path = temp_dir.path().join("custom.jinja"); + let custom_template = + r#"CUSTOM: {% for msg in messages %}[{{ msg.role }}]: {{ msg.content }}{% endfor %}"#; + fs::write(&custom_template_path, custom_template).unwrap(); + + // Load with custom template - should override the built-in one + let tokenizer = HuggingFaceTokenizer::from_file_with_chat_template( + tokenizer_path.to_str().unwrap(), + Some(custom_template_path.to_str().unwrap()), + ) + .unwrap(); + + let messages = vec![ChatMessage::user("Test")]; + let result = tokenizer.apply_chat_template(&messages, false).unwrap(); + + // Should use CUSTOM template, not built-in + assert!(result.starts_with("CUSTOM:")); + assert!(result.contains("[user]: Test")); + assert!(!result.contains("built-in:")); + } + + #[test] + fn test_set_chat_template_after_creation() { + use lmdeploy_router_rs::tokenizer::chat_template::ChatMessage; + use lmdeploy_router_rs::tokenizer::huggingface::HuggingFaceTokenizer; + + // Create temporary directory and tokenizer file + let temp_dir = TempDir::new().unwrap(); + let tokenizer_json = r#"{ + "version": "1.0", + "truncation": null, + "padding": null, + "added_tokens": [], + "normalizer": null, + "pre_tokenizer": { + "type": "Whitespace" + }, + "post_processor": null, + "decoder": null, + "model": { + "type": "BPE", + "vocab": { + "test": 0, + "": 1, + "": 2 + }, + "merges": [] + } + }"#; + let tokenizer_path = temp_dir.path().join("tokenizer.json"); + fs::write(&tokenizer_path, tokenizer_json).unwrap(); + + // Load tokenizer without custom template + let mut tokenizer = + HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap()).unwrap(); + + // Set a template after creation (mimics Python's behavior) + let new_template = + "NEW: {% for msg in messages %}{{ msg.role }}: {{ msg.content }}; {% endfor %}"; + tokenizer.set_chat_template(new_template.to_string()); + + let messages = vec![ChatMessage::user("Hello"), ChatMessage::assistant("World")]; + let result = tokenizer.apply_chat_template(&messages, false).unwrap(); + + assert!(result.starts_with("NEW:")); + assert!(result.contains("user: Hello;")); + assert!(result.contains("assistant: World;")); + } +} diff --git a/lmdeploy/router/tests/test_consistent_hash_policy.rs b/lmdeploy/router/tests/test_consistent_hash_policy.rs new file mode 100644 index 0000000000..a29bbf6f75 --- /dev/null +++ b/lmdeploy/router/tests/test_consistent_hash_policy.rs @@ -0,0 +1,756 @@ +#[cfg(test)] +mod consistent_hash_policy_tests { + use std::collections::HashMap; + use std::sync::Arc; + + use lmdeploy_router_rs::core::BasicWorker; + use lmdeploy_router_rs::core::Worker; + use lmdeploy_router_rs::core::WorkerType; + use lmdeploy_router_rs::policies::ConsistentHashPolicy; + use lmdeploy_router_rs::policies::LoadBalancingPolicy; + use lmdeploy_router_rs::policies::RequestHeaders; + + /// Helper function to create test workers + fn create_test_workers() -> Vec> { + vec![ + Arc::new(BasicWorker::new( + "http://worker1:8000".to_string(), + WorkerType::Regular, + )), + Arc::new(BasicWorker::new( + "http://worker2:8000".to_string(), + WorkerType::Regular, + )), + Arc::new(BasicWorker::new( + "http://worker3:8000".to_string(), + WorkerType::Regular, + )), + ] + } + + #[test] + fn test_consistent_hash_policy_creation() { + let policy = ConsistentHashPolicy::new(); + assert_eq!(policy.name(), "consistent_hash"); + assert!(policy.needs_request_text()); + } + + #[test] + fn test_consistent_routing_same_session() { + let policy = ConsistentHashPolicy::new(); + let workers = create_test_workers(); + + let session_id = "test_session_123"; + + // Make multiple requests with the same top-level session_id + let mut selected_workers = Vec::new(); + for i in 0..10 { + let prompt = format!( + r#"{{"session_id": "{}", "prompt": "request {}}}"#, + session_id, i + ); + if let Some(worker_idx) = policy.select_worker(&workers, Some(&prompt)) { + selected_workers.push(worker_idx); + } + } + + // All requests should go to the same worker + assert!( + !selected_workers.is_empty(), + "Should have selected at least one worker" + ); + + let first_worker = selected_workers[0]; + for (i, &worker_idx) in selected_workers.iter().enumerate() { + assert_eq!( + worker_idx, first_worker, + "Request {} went to worker {}, expected worker {} (same as first request)", + i, worker_idx, first_worker + ); + } + } + + #[test] + fn test_distribution_across_workers() { + let policy = ConsistentHashPolicy::new(); + let workers = create_test_workers(); + + let mut worker_counts = HashMap::new(); + let num_sessions = 100; + + // Create many different sessions and see how they distribute + for i in 0..num_sessions { + let session_id = format!("session_{}", i); + let request_json = format!(r#"{{"session_id": "{}", "prompt": "test"}}"#, session_id); + + if let Some(worker_idx) = policy.select_worker(&workers, Some(&request_json)) { + *worker_counts.entry(worker_idx).or_insert(0) += 1; + } + } + + // Check that at least 2 workers were used (distribution test) + assert!( + worker_counts.len() >= 2, + "Expected distribution across multiple workers, only used: {:?}", + worker_counts + ); + + // Check that no single worker got all requests (basic distribution check) + let max_count = worker_counts.values().max().unwrap(); + assert!( + *max_count < num_sessions, + "One worker got all requests, no distribution occurred" + ); + + // Check that distribution is reasonably balanced (allow for some variance) + let min_count = worker_counts.values().min().unwrap(); + let expected_per_worker = num_sessions / workers.len(); + + // Allow 50% variance from expected + assert!( + *min_count >= expected_per_worker / 3, + "Distribution too uneven, min count: {}, expected around: {}", + min_count, + expected_per_worker + ); + } + + #[test] + fn test_select_worker_pair_pd_mode() { + let policy = ConsistentHashPolicy::new(); + let prefill_workers = create_test_workers(); + let decode_workers = create_test_workers(); + + let session_id = "pd_session_test"; + let request_json = format!(r#"{{"session_id": "{}", "prompt": "pd test"}}"#, session_id); + + // Test PD mode worker pair selection + if let Some((prefill_idx, decode_idx)) = + policy.select_worker_pair(&prefill_workers, &decode_workers, Some(&request_json)) + { + assert!( + prefill_idx < prefill_workers.len(), + "Prefill worker index should be valid" + ); + assert!( + decode_idx < decode_workers.len(), + "Decode worker index should be valid" + ); + + // For the same session, should get consistent results + if let Some((prefill_idx2, decode_idx2)) = + policy.select_worker_pair(&prefill_workers, &decode_workers, Some(&request_json)) + { + assert_eq!( + prefill_idx, prefill_idx2, + "Prefill worker should be consistent" + ); + assert_eq!( + decode_idx, decode_idx2, + "Decode worker should be consistent" + ); + } + } + } + + #[test] + fn test_no_workers_available() { + let policy = ConsistentHashPolicy::new(); + let empty_workers: Vec> = vec![]; + + // Should return None when no workers are available + let result = policy.select_worker(&empty_workers, Some(r#"{"session_id": "test"}"#)); + assert!( + result.is_none(), + "Should return None when no workers available" + ); + } + + #[test] + fn test_unhealthy_workers_fallback() { + let policy = ConsistentHashPolicy::new(); + + // Create workers where some are unhealthy + let worker1 = BasicWorker::new("http://worker1:8000".to_string(), WorkerType::Regular); + let worker2 = BasicWorker::new("http://worker2:8000".to_string(), WorkerType::Regular); + + // Mark first worker as unhealthy + worker1.set_healthy(false); + + let workers: Vec> = vec![Arc::new(worker1), Arc::new(worker2)]; + + // Should still work and select healthy workers + let result = policy.select_worker(&workers, Some(r#"{"session_id": "test"}"#)); + assert!(result.is_some(), "Should fallback to healthy workers"); + + // The selected worker should be healthy + if let Some(idx) = result { + assert!( + workers[idx].is_healthy(), + "Selected worker should be healthy" + ); + } + } + + #[test] + fn test_comprehensive_routing_consistency() { + let policy = ConsistentHashPolicy::new(); + let workers = create_test_workers(); + + // Test multiple sessions with multiple requests each + let sessions = ["session_A", "session_B", "session_C"]; + let mut session_mappings = HashMap::new(); + + for session in &sessions { + let mut worker_indices = Vec::new(); + + // Make 5 requests per session + for i in 0..5 { + let request = format!( + r#"{{"session_id": "{}", "prompt": "request {}}}"#, + session, i + ); + if let Some(idx) = policy.select_worker(&workers, Some(&request)) { + worker_indices.push(idx); + } + } + + // All requests from the same session should go to the same worker + assert!( + !worker_indices.is_empty(), + "Should have selected workers for session: {}", + session + ); + + let first_idx = worker_indices[0]; + for (i, &idx) in worker_indices.iter().enumerate() { + assert_eq!( + idx, first_idx, + "Session {} request {} went to worker {}, expected worker {} (consistency violation)", + session, i, idx, first_idx + ); + } + + session_mappings.insert(*session, first_idx); + } + + // Verify that different sessions can go to different workers (distribution) + let unique_workers: std::collections::HashSet<_> = session_mappings.values().collect(); + println!("Session mappings: {:?}", session_mappings); + println!("Unique workers used: {}", unique_workers.len()); + + // With 3 sessions and 3 workers, we should ideally see some distribution + // but we'll be lenient and just ensure not everything goes to one worker + assert!( + !unique_workers.is_empty(), + "At least one worker should be used" + ); + } + + #[test] + fn test_fallback_without_session_or_user() { + let policy = ConsistentHashPolicy::new(); + let workers = create_test_workers(); + + // Test requests without session_id or user_id (should use request content) + let test_cases = [ + r#"{"prompt": "hello world"}"#, + r#"{"text": "test request", "sampling_params": {}}"#, + "", + "plain text request", + ]; + + for (i, request) in test_cases.iter().enumerate() { + let result = policy.select_worker(&workers, Some(request)); + assert!( + result.is_some(), + "Fallback should work for request {}: {}", + i, + request + ); + } + } + + // Session IDs use the top-level field; user identity uses the `user` field. + + #[test] + fn test_openai_user_field_consistency() { + let policy = ConsistentHashPolicy::new(); + let workers = create_test_workers(); + + let user = "openai_user_123"; + + // Make multiple requests with the same user field (OpenAI format) + let mut selected_workers = Vec::new(); + for i in 0..5 { + let prompt = format!(r#"{{"user": "{}", "prompt": "request {}}}"#, user, i); + if let Some(worker_idx) = policy.select_worker(&workers, Some(&prompt)) { + selected_workers.push(worker_idx); + } + } + + // All requests should go to the same worker + assert!( + !selected_workers.is_empty(), + "Should have selected at least one worker" + ); + + let first_worker = selected_workers[0]; + for (i, &worker_idx) in selected_workers.iter().enumerate() { + assert_eq!( + worker_idx, first_worker, + "OpenAI user request {} went to worker {}, expected worker {} (consistency violation)", + i, worker_idx, first_worker + ); + } + } + + #[test] + fn test_session_id_priority_over_user_id() { + let policy = ConsistentHashPolicy::new(); + let workers = create_test_workers(); + + let session_id = "priority_session"; + let user_id = "priority_user"; + + // Make requests with both session_id and user_id + let mut selected_workers = Vec::new(); + for i in 0..3 { + let prompt = format!( + r#"{{"session_id": "{}", "user_id": "{}", "prompt": "request {}}}"#, + session_id, user_id, i + ); + if let Some(worker_idx) = policy.select_worker(&workers, Some(&prompt)) { + selected_workers.push(worker_idx); + } + } + + // All requests should consistently go to the same worker (based on session_id) + assert!( + !selected_workers.is_empty(), + "Should have selected at least one worker" + ); + + let first_worker = selected_workers[0]; + for (i, &worker_idx) in selected_workers.iter().enumerate() { + assert_eq!( + worker_idx, first_worker, + "Priority test request {} went to worker {}, expected worker {} (inconsistent routing)", + i, worker_idx, first_worker + ); + } + } + + #[test] + fn test_policy_reset() { + let policy = ConsistentHashPolicy::new(); + + // Reset should work without issues + policy.reset(); + + // After reset, should still be able to route requests + let workers = create_test_workers(); + let result = policy.select_worker(&workers, Some(r#"{"session_id": "post_reset_test"}"#)); + assert!(result.is_some(), "Should work after reset"); + } + + #[test] + fn test_different_request_formats() { + let policy = ConsistentHashPolicy::new(); + let workers = create_test_workers(); + + let session_id = "format_test_session"; + + // Test different JSON formats that should all extract the same session_id + let formats = [ + format!(r#"{{"session_id": "{}", "prompt": "test"}}"#, session_id), + format!( + r#"{{ "session_id" : "{}" , "prompt" : "test" }}"#, + session_id + ), + format!(r#"{{"prompt": "test", "session_id": "{}"}}"#, session_id), + ]; + + let mut selected_workers = Vec::new(); + for (i, request) in formats.iter().enumerate() { + if let Some(worker_idx) = policy.select_worker(&workers, Some(request)) { + selected_workers.push(worker_idx); + } else { + panic!("Format {} failed to route: {}", i, request); + } + } + + // All different formats should route to the same worker + let first_worker = selected_workers[0]; + for (i, &worker_idx) in selected_workers.iter().enumerate() { + assert_eq!( + worker_idx, first_worker, + "Format {} routed to worker {}, expected worker {} (format inconsistency)", + i, worker_idx, first_worker + ); + } + } + + // ============================= + // HTTP Header-based routing tests + // ============================= + + /// Helper function to create RequestHeaders from key-value pairs + fn create_headers(pairs: &[(&str, &str)]) -> RequestHeaders { + pairs + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect() + } + + #[test] + fn test_http_header_x_session_id_routing() { + let policy = ConsistentHashPolicy::new(); + let workers = create_test_workers(); + + let session_id = "header-session-123"; + let headers = create_headers(&[("x-session-id", session_id)]); + + // Make multiple requests with the same x-session-id header + let mut selected_workers = Vec::new(); + for i in 0..10 { + let prompt = format!(r#"{{"prompt": "request {}"}}"#, i); + if let Some(worker_idx) = + policy.select_worker_with_headers(&workers, Some(&prompt), Some(&headers)) + { + selected_workers.push(worker_idx); + } + } + + // All requests should go to the same worker + assert!( + !selected_workers.is_empty(), + "Should have selected at least one worker" + ); + + let first_worker = selected_workers[0]; + for (i, &worker_idx) in selected_workers.iter().enumerate() { + assert_eq!( + worker_idx, first_worker, + "Request {} with x-session-id header went to worker {}, expected worker {}", + i, worker_idx, first_worker + ); + } + } + + #[test] + fn test_http_header_priority_over_body() { + let policy = ConsistentHashPolicy::new(); + let workers = create_test_workers(); + + // Header session ID should take priority over body session ID + let header_session = "header-session-priority"; + let body_session = "body-session-ignored"; + + let headers = create_headers(&[("x-session-id", header_session)]); + let body = format!(r#"{{"session_id": "{}", "prompt": "test"}}"#, body_session); + + // Get worker for header-only request + let header_only_headers = create_headers(&[("x-session-id", header_session)]); + let header_only_worker = policy.select_worker_with_headers( + &workers, + Some(r#"{"prompt": "test"}"#), + Some(&header_only_headers), + ); + + // Get worker for request with both header and body + let both_worker = policy.select_worker_with_headers(&workers, Some(&body), Some(&headers)); + + // Both should route to the same worker (header takes priority) + assert_eq!( + header_only_worker, both_worker, + "HTTP header should take priority over body session_id" + ); + } + + #[test] + fn test_http_header_x_user_id_routing() { + let policy = ConsistentHashPolicy::new(); + let workers = create_test_workers(); + + let user_id = "header-user-456"; + let headers = create_headers(&[("x-user-id", user_id)]); + + // Make multiple requests with the same x-user-id header + let mut selected_workers = Vec::new(); + for i in 0..5 { + let prompt = format!(r#"{{"prompt": "request {}"}}"#, i); + if let Some(worker_idx) = + policy.select_worker_with_headers(&workers, Some(&prompt), Some(&headers)) + { + selected_workers.push(worker_idx); + } + } + + // All requests should go to the same worker + assert!( + !selected_workers.is_empty(), + "Should have selected at least one worker" + ); + + let first_worker = selected_workers[0]; + for (i, &worker_idx) in selected_workers.iter().enumerate() { + assert_eq!( + worker_idx, first_worker, + "Request {} with x-user-id header went to worker {}, expected worker {}", + i, worker_idx, first_worker + ); + } + } + + #[test] + fn test_http_header_priority_order() { + let policy = ConsistentHashPolicy::new(); + let workers = create_test_workers(); + + // x-session-id should take priority over x-user-id + let session_id = "priority-session"; + let user_id = "ignored-user"; + + // Test with x-session-id only + let session_headers = create_headers(&[("x-session-id", session_id)]); + let session_worker = policy + .select_worker_with_headers( + &workers, + Some(r#"{"prompt": "test"}"#), + Some(&session_headers), + ) + .expect("Should select worker with x-session-id"); + + // Test with both headers - x-session-id should win + let both_headers = create_headers(&[("x-session-id", session_id), ("x-user-id", user_id)]); + let both_worker = policy + .select_worker_with_headers( + &workers, + Some(r#"{"prompt": "test"}"#), + Some(&both_headers), + ) + .expect("Should select worker with both headers"); + + assert_eq!( + session_worker, both_worker, + "x-session-id should take priority over x-user-id" + ); + } + + #[test] + fn test_http_header_fallback_to_body_when_no_header() { + let policy = ConsistentHashPolicy::new(); + let workers = create_test_workers(); + + let body_session = "body-only-session"; + let body = format!(r#"{{"session_id": "{}", "prompt": "test"}}"#, body_session); + + // Empty headers should fall back to body + let empty_headers: RequestHeaders = HashMap::new(); + let with_empty_headers = policy + .select_worker_with_headers(&workers, Some(&body), Some(&empty_headers)) + .expect("Should route with empty headers"); + + // No headers should also fall back to body + let without_headers = policy + .select_worker_with_headers(&workers, Some(&body), None) + .expect("Should route without headers"); + + // Body-only routing (using old interface) for comparison + let body_only = policy + .select_worker(&workers, Some(&body)) + .expect("Should route with body only"); + + assert_eq!( + with_empty_headers, without_headers, + "Empty headers and None headers should behave the same" + ); + assert_eq!( + with_empty_headers, body_only, + "Should fall back to body when no headers present" + ); + } + + #[test] + fn test_http_header_x_tenant_id_routing() { + let policy = ConsistentHashPolicy::new(); + let workers = create_test_workers(); + + let tenant_id = "tenant-abc"; + let headers = create_headers(&[("x-tenant-id", tenant_id)]); + + let mut selected_workers = Vec::new(); + for i in 0..5 { + let prompt = format!(r#"{{"prompt": "request {}"}}"#, i); + if let Some(worker_idx) = + policy.select_worker_with_headers(&workers, Some(&prompt), Some(&headers)) + { + selected_workers.push(worker_idx); + } + } + + // All requests should go to the same worker + assert!(!selected_workers.is_empty(), "Should have selected workers"); + + let first_worker = selected_workers[0]; + for (i, &worker_idx) in selected_workers.iter().enumerate() { + assert_eq!( + worker_idx, first_worker, + "Request {} with x-tenant-id header went to worker {}, expected worker {}", + i, worker_idx, first_worker + ); + } + } + + #[test] + fn test_http_header_distribution_across_workers() { + let policy = ConsistentHashPolicy::new(); + let workers = create_test_workers(); + + let mut worker_counts = HashMap::new(); + let num_sessions = 100; + + // Create many different session IDs via headers + for i in 0..num_sessions { + let session_id = format!("header-session-{}", i); + let headers = create_headers(&[("x-session-id", &session_id)]); + + if let Some(worker_idx) = policy.select_worker_with_headers( + &workers, + Some(r#"{"prompt": "test"}"#), + Some(&headers), + ) { + *worker_counts.entry(worker_idx).or_insert(0) += 1; + } + } + + // Check distribution - should use at least 2 workers + assert!( + worker_counts.len() >= 2, + "Expected distribution across multiple workers, only used: {:?}", + worker_counts + ); + + // No single worker should get all requests + let max_count = worker_counts.values().max().unwrap(); + assert!( + *max_count < num_sessions, + "One worker got all requests, no distribution occurred" + ); + } + + #[test] + fn test_needs_headers_flag() { + let policy = ConsistentHashPolicy::new(); + assert!( + policy.needs_headers(), + "ConsistentHashPolicy should report that it needs headers" + ); + } + + #[test] + fn test_correlation_id_priority_over_request_id() { + let policy = ConsistentHashPolicy::new(); + let workers = create_test_workers(); + + let correlation_id = "session-correlation-abc"; + let request_id = "unique-request-xyz"; + + // Test with x-correlation-id only + let correlation_headers = create_headers(&[("x-correlation-id", correlation_id)]); + let correlation_worker = policy + .select_worker_with_headers( + &workers, + Some(r#"{"prompt": "test"}"#), + Some(&correlation_headers), + ) + .expect("Should select worker with x-correlation-id"); + + // Test with both headers - x-correlation-id should win + let both_headers = create_headers(&[ + ("x-correlation-id", correlation_id), + ("x-request-id", request_id), + ]); + let both_worker = policy + .select_worker_with_headers( + &workers, + Some(r#"{"prompt": "test"}"#), + Some(&both_headers), + ) + .expect("Should select worker with both headers"); + + assert_eq!( + correlation_worker, both_worker, + "x-correlation-id should take priority over x-request-id" + ); + } + + #[test] + fn test_correlation_id_consistent_across_turns() { + let policy = ConsistentHashPolicy::new(); + let workers = create_test_workers(); + + let correlation_id = "multi-turn-session-123"; + + // Simulate multiple turns with same correlation ID but different request IDs + let mut selected_workers = Vec::new(); + for i in 0..10 { + let headers = create_headers(&[ + ("x-correlation-id", correlation_id), + ("x-request-id", &format!("request-turn-{}", i)), + ]); + let worker = policy + .select_worker_with_headers(&workers, Some(r#"{"prompt": "test"}"#), Some(&headers)) + .expect("Should select worker"); + selected_workers.push(worker); + } + + // All turns should route to the same worker + let first_worker = &selected_workers[0]; + for (i, worker) in selected_workers.iter().enumerate() { + assert_eq!( + first_worker, worker, + "Turn {} should route to same worker as turn 0 when x-correlation-id is shared", + i + ); + } + } + + #[test] + fn test_tenant_id_priority_over_correlation_id() { + let policy = ConsistentHashPolicy::new(); + let workers = create_test_workers(); + + let tenant_id = "tenant-abc"; + let correlation_id = "correlation-xyz"; + + // Test with x-tenant-id only + let tenant_headers = create_headers(&[("x-tenant-id", tenant_id)]); + let tenant_worker = policy + .select_worker_with_headers( + &workers, + Some(r#"{"prompt": "test"}"#), + Some(&tenant_headers), + ) + .expect("Should select worker with x-tenant-id"); + + // Test with both headers - x-tenant-id should win + let both_headers = create_headers(&[ + ("x-tenant-id", tenant_id), + ("x-correlation-id", correlation_id), + ]); + let both_worker = policy + .select_worker_with_headers( + &workers, + Some(r#"{"prompt": "test"}"#), + Some(&both_headers), + ) + .expect("Should select worker with both headers"); + + assert_eq!( + tenant_worker, both_worker, + "x-tenant-id should take priority over x-correlation-id" + ); + } +} diff --git a/lmdeploy/router/tests/test_extra_args_chat.rs b/lmdeploy/router/tests/test_extra_args_chat.rs new file mode 100644 index 0000000000..7633155480 --- /dev/null +++ b/lmdeploy/router/tests/test_extra_args_chat.rs @@ -0,0 +1,88 @@ +//! Tests that unknown/extra fields in ChatCompletionRequest are preserved through serde roundtrip. +use lmdeploy_router_rs::protocols::spec::ChatCompletionRequest; + +#[test] +fn test_extra_fields_preserved_on_deserialize() { + let json = r#"{ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 100, + "return_token_ids": true, + "some_custom_param": "value" + }"#; + + let req: ChatCompletionRequest = serde_json::from_str(json).unwrap(); + + assert!( + req.other.contains_key("return_token_ids"), + "return_token_ids should be captured in other" + ); + assert!( + req.other.contains_key("some_custom_param"), + "some_custom_param should be captured in other" + ); + assert_eq!( + req.other["some_custom_param"], + serde_json::Value::String("value".to_string()) + ); +} + +#[test] +fn test_extra_fields_survive_roundtrip() { + let json = r#"{ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}], + "max_tokens": 100, + "return_token_ids": true + }"#; + + let req: ChatCompletionRequest = serde_json::from_str(json).unwrap(); + let serialized = serde_json::to_string(&req).unwrap(); + let roundtrip: ChatCompletionRequest = serde_json::from_str(&serialized).unwrap(); + + assert_eq!(roundtrip.other["return_token_ids"], serde_json::json!(true)); +} + +#[test] +fn test_extra_fields_appear_at_top_level_in_serialized_json() { + let json = r#"{ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}], + "return_token_ids": true + }"#; + + let req: ChatCompletionRequest = serde_json::from_str(json).unwrap(); + let value: serde_json::Value = serde_json::to_value(&req).unwrap(); + + assert!( + value.get("return_token_ids").is_some(), + "return_token_ids should be at the top level of the serialized JSON" + ); + assert_eq!(value["return_token_ids"], serde_json::json!(true)); +} + +#[test] +fn test_no_other_fields_gives_empty_map() { + let json = r#"{ + "model": "test-model", + "messages": [{"role": "user", "content": "Hello"}] + }"#; + + let req: ChatCompletionRequest = serde_json::from_str(json).unwrap(); + assert!(req.other.is_empty()); +} + +#[test] +fn test_lmdeploy_input_ids_is_typed_not_flattened() { + let json = r#"{ + "model": "test-model", + "messages": [], + "input_ids": [1, 2, 3], + "return_token_ids": true + }"#; + + let req: ChatCompletionRequest = serde_json::from_str(json).unwrap(); + assert_eq!(req.input_ids, Some(vec![1, 2, 3])); + assert!(!req.other.contains_key("input_ids")); + assert_eq!(req.other["return_token_ids"], serde_json::json!(true)); +} diff --git a/lmdeploy/router/tests/test_openai_routing.rs b/lmdeploy/router/tests/test_openai_routing.rs new file mode 100644 index 0000000000..93de447538 --- /dev/null +++ b/lmdeploy/router/tests/test_openai_routing.rs @@ -0,0 +1,465 @@ +//! Comprehensive integration tests for OpenAI backend functionality + +use axum::{ + body::Body, + extract::Request, + http::{Method, StatusCode}, + routing::post, + Router, +}; +use lmdeploy_router_rs::{ + config::{RouterConfig, RoutingMode}, + protocols::spec::{ + ChatCompletionRequest, ChatMessage, CompletionRequest, GenerateRequest, PromptInput, + UserMessageContent, + }, + routers::{openai_router::OpenAIRouter, RouterTrait}, +}; +use serde_json::json; +use std::sync::Arc; +use tower::ServiceExt; + +mod common; +use common::mock_openai_server::MockOpenAIServer; + +/// Helper function to create a minimal chat completion request for testing +fn create_minimal_chat_request() -> ChatCompletionRequest { + let val = json!({ + "model": "gpt-3.5-turbo", + "messages": [ + {"role": "user", "content": "Hello"} + ], + "max_tokens": 100 + }); + serde_json::from_value(val).unwrap() +} + +/// Helper function to create a minimal completion request for testing +fn create_minimal_completion_request() -> CompletionRequest { + CompletionRequest { + model: Some("gpt-3.5-turbo".to_string()), + prompt: PromptInput::String("Hello".to_string()), + suffix: None, + max_tokens: Some(100), + temperature: None, + top_p: None, + n: None, + stream: false, + stream_options: None, + logprobs: None, + echo: false, + stop: None, + presence_penalty: None, + frequency_penalty: None, + best_of: None, + logit_bias: None, + user: None, + seed: None, + top_k: None, + min_p: None, + min_tokens: None, + repetition_penalty: None, + regex: None, + ebnf: None, + json_schema: None, + stop_token_ids: None, + no_stop_trim: false, + ignore_eos: false, + skip_special_tokens: true, + other: serde_json::Map::new(), + } +} + +// ============= Basic Unit Tests ============= + +/// Test basic OpenAI router creation and configuration +#[tokio::test] +async fn test_openai_router_creation() { + let router = OpenAIRouter::new("https://api.openai.com".to_string(), None).await; + + assert!(router.is_ok(), "Router creation should succeed"); + + let router = router.unwrap(); + assert_eq!(router.router_type(), "openai"); + assert!(!router.is_pd_mode()); +} + +/// Test health endpoints +#[tokio::test] +async fn test_openai_router_health() { + let router = OpenAIRouter::new("https://api.openai.com".to_string(), None) + .await + .unwrap(); + + let req = Request::builder() + .method(Method::GET) + .uri("/health") + .body(Body::empty()) + .unwrap(); + + let response = router.health(req).await; + assert_eq!(response.status(), StatusCode::OK); +} + +/// Test server info endpoint +#[tokio::test] +async fn test_openai_router_server_info() { + let router = OpenAIRouter::new("https://api.openai.com".to_string(), None) + .await + .unwrap(); + + let req = Request::builder() + .method(Method::GET) + .uri("/info") + .body(Body::empty()) + .unwrap(); + + let response = router.get_server_info(req).await; + assert_eq!(response.status(), StatusCode::OK); + + let (_, body) = response.into_parts(); + let body_bytes = axum::body::to_bytes(body, usize::MAX).await.unwrap(); + let body_str = String::from_utf8(body_bytes.to_vec()).unwrap(); + + assert!(body_str.contains("openai")); +} + +/// Test models endpoint +#[tokio::test] +async fn test_openai_router_models() { + // Use mock server for deterministic models response + let mock_server = MockOpenAIServer::new().await; + let router = OpenAIRouter::new(mock_server.base_url(), None) + .await + .unwrap(); + + let req = Request::builder() + .method(Method::GET) + .uri("/models") + .body(Body::empty()) + .unwrap(); + + let response = router.get_models(req).await; + assert_eq!(response.status(), StatusCode::OK); + + let (_, body) = response.into_parts(); + let body_bytes = axum::body::to_bytes(body, usize::MAX).await.unwrap(); + let body_str = String::from_utf8(body_bytes.to_vec()).unwrap(); + let models: serde_json::Value = serde_json::from_str(&body_str).unwrap(); + + assert_eq!(models["object"], "list"); + assert!(models["data"].is_array()); +} + +/// Test router factory with OpenAI routing mode +#[tokio::test] +async fn test_router_factory_openai_mode() { + let routing_mode = RoutingMode::OpenAI { + worker_urls: vec!["https://api.openai.com".to_string()], + }; + + let router_config = RouterConfig::new( + routing_mode, + lmdeploy_router_rs::config::PolicyConfig::Random, + ); + + let app_context = common::create_test_context(router_config); + + let router = lmdeploy_router_rs::routers::RouterFactory::create_router(&app_context).await; + assert!( + router.is_ok(), + "Router factory should create OpenAI router successfully" + ); + + let router = router.unwrap(); + assert_eq!(router.router_type(), "openai"); +} + +/// Test that unsupported endpoints return proper error codes +#[tokio::test] +async fn test_unsupported_endpoints() { + let router = OpenAIRouter::new("https://api.openai.com".to_string(), None) + .await + .unwrap(); + + // The OpenAI backend does not support the generate endpoint. + let generate_request = GenerateRequest { + prompt: None, + text: Some("Hello world".to_string()), + input_ids: None, + parameters: None, + sampling_params: None, + stream: false, + return_logprob: false, + rid: None, + other: serde_json::Map::new(), + }; + + let response = router.route_generate(None, &generate_request, None).await; + assert_eq!(response.status(), StatusCode::NOT_IMPLEMENTED); + + // Test completion endpoint (should also not be supported) + let completion_request = create_minimal_completion_request(); + let response = router + .route_completion(None, &completion_request, None) + .await; + assert_eq!(response.status(), StatusCode::NOT_IMPLEMENTED); +} + +// ============= Mock Server E2E Tests ============= + +/// Test chat completion with mock OpenAI server +#[tokio::test] +async fn test_openai_router_chat_completion_with_mock() { + // Start a mock OpenAI server + let mock_server = MockOpenAIServer::new().await; + let base_url = mock_server.base_url(); + + // Create router pointing to mock server + let router = OpenAIRouter::new(base_url, None).await.unwrap(); + + // Create a minimal chat completion request + let mut chat_request = create_minimal_chat_request(); + chat_request.messages = vec![ChatMessage::User { + role: "user".to_string(), + content: UserMessageContent::Text("Hello, how are you?".to_string()), + name: None, + }]; + chat_request.temperature = Some(0.7); + + // Route the request + let response = router.route_chat(None, &chat_request, None).await; + + // Should get a successful response from mock server + assert_eq!(response.status(), StatusCode::OK); + + let (_, body) = response.into_parts(); + let body_bytes = axum::body::to_bytes(body, usize::MAX).await.unwrap(); + let body_str = String::from_utf8(body_bytes.to_vec()).unwrap(); + let chat_response: serde_json::Value = serde_json::from_str(&body_str).unwrap(); + + // Verify it's a valid chat completion response + assert_eq!(chat_response["object"], "chat.completion"); + assert_eq!(chat_response["model"], "gpt-3.5-turbo"); + assert!(!chat_response["choices"].as_array().unwrap().is_empty()); +} + +/// Test full E2E flow with Axum server +#[tokio::test] +async fn test_openai_e2e_with_server() { + // Start mock OpenAI server + let mock_server = MockOpenAIServer::new().await; + let base_url = mock_server.base_url(); + + // Create router + let router = OpenAIRouter::new(base_url, None).await.unwrap(); + + // Create Axum app with chat completions endpoint + let app = Router::new().route( + "/v1/chat/completions", + post({ + let router = Arc::new(router); + move |req: Request| { + let router = router.clone(); + async move { + let (parts, body) = req.into_parts(); + let body_bytes = axum::body::to_bytes(body, usize::MAX).await.unwrap(); + let body_str = String::from_utf8(body_bytes.to_vec()).unwrap(); + + let chat_request: ChatCompletionRequest = + serde_json::from_str(&body_str).unwrap(); + + router + .route_chat(Some(&parts.headers), &chat_request, None) + .await + } + } + }), + ); + + // Make a request to the server + let request = Request::builder() + .method(Method::POST) + .uri("/v1/chat/completions") + .header("content-type", "application/json") + .body(Body::from( + json!({ + "model": "gpt-3.5-turbo", + "messages": [ + { + "role": "user", + "content": "Hello, world!" + } + ], + "max_tokens": 100 + }) + .to_string(), + )) + .unwrap(); + + let response = app.oneshot(request).await.unwrap(); + assert_eq!(response.status(), StatusCode::OK); + + let body = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + let response_json: serde_json::Value = serde_json::from_slice(&body).unwrap(); + + // Verify the response structure + assert_eq!(response_json["object"], "chat.completion"); + assert_eq!(response_json["model"], "gpt-3.5-turbo"); + assert!(!response_json["choices"].as_array().unwrap().is_empty()); +} + +/// Test streaming chat completions pass-through with mock server +#[tokio::test] +async fn test_openai_router_chat_streaming_with_mock() { + let mock_server = MockOpenAIServer::new().await; + let base_url = mock_server.base_url(); + let router = OpenAIRouter::new(base_url, None).await.unwrap(); + + // Build a streaming chat request + let val = json!({ + "model": "gpt-3.5-turbo", + "messages": [ + {"role": "user", "content": "Hello"} + ], + "max_tokens": 10, + "stream": true + }); + let chat_request: ChatCompletionRequest = serde_json::from_value(val).unwrap(); + + let response = router.route_chat(None, &chat_request, None).await; + assert_eq!(response.status(), StatusCode::OK); + + // Should be SSE + let headers = response.headers(); + let ct = headers + .get("content-type") + .unwrap() + .to_str() + .unwrap() + .to_ascii_lowercase(); + assert!(ct.contains("text/event-stream")); + + // Read entire stream body and assert chunks + DONE + let body = axum::body::to_bytes(response.into_body(), usize::MAX) + .await + .unwrap(); + let text = String::from_utf8(body.to_vec()).unwrap(); + assert!(text.contains("chat.completion.chunk")); + assert!(text.contains("[DONE]")); +} + +/// Test circuit breaker functionality +#[tokio::test] +async fn test_openai_router_circuit_breaker() { + // Create router with circuit breaker config + let cb_config = lmdeploy_router_rs::config::CircuitBreakerConfig { + failure_threshold: 2, + success_threshold: 1, + timeout_duration_secs: 1, + window_duration_secs: 10, + }; + + let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap(); + let base_url = format!("http://{}", listener.local_addr().unwrap()); + drop(listener); + + let router = OpenAIRouter::new(base_url, Some(cb_config)).await.unwrap(); + + let chat_request = create_minimal_chat_request(); + + // First few requests should fail and record failures + for _ in 0..3 { + let response = router.route_chat(None, &chat_request, None).await; + // Should get either an error or circuit breaker response + assert!( + response.status() == StatusCode::INTERNAL_SERVER_ERROR + || response.status() == StatusCode::SERVICE_UNAVAILABLE + ); + } +} + +/// Test that Authorization header is forwarded in /v1/models +#[tokio::test] +async fn test_openai_router_models_auth_forwarding() { + // Start a mock server that requires Authorization + let expected_auth = "Bearer test-token".to_string(); + let mock_server = MockOpenAIServer::new_with_auth(Some(expected_auth.clone())).await; + let router = OpenAIRouter::new(mock_server.base_url(), None) + .await + .unwrap(); + + // 1) Without auth header -> expect 401 + let req = Request::builder() + .method(Method::GET) + .uri("/models") + .body(Body::empty()) + .unwrap(); + + let response = router.get_models(req).await; + assert_eq!(response.status(), StatusCode::UNAUTHORIZED); + + // 2) With auth header -> expect 200 + let req = Request::builder() + .method(Method::GET) + .uri("/models") + .header("Authorization", expected_auth) + .body(Body::empty()) + .unwrap(); + + let response = router.get_models(req).await; + assert_eq!(response.status(), StatusCode::OK); + + let (_, body) = response.into_parts(); + let body_bytes = axum::body::to_bytes(body, usize::MAX).await.unwrap(); + let body_str = String::from_utf8(body_bytes.to_vec()).unwrap(); + let models: serde_json::Value = serde_json::from_str(&body_str).unwrap(); + assert_eq!(models["object"], "list"); +} + +// ============= Reasoning Effort Integration Tests ============= + +/// Test chat completion with the new reasoning fields (echo, reasoning_effort, include_reasoning) +#[tokio::test] +async fn test_openai_router_chat_with_reasoning_fields() { + let mock_server = MockOpenAIServer::new().await; + let base_url = mock_server.base_url(); + let router = OpenAIRouter::new(base_url, None).await.unwrap(); + + // Create a chat request with all three new reasoning fields + let val = json!({ + "model": "gpt-3.5-turbo", + "messages": [ + {"role": "user", "content": "What is 2+2?"} + ], + "max_tokens": 100, + "echo": true, + "reasoning_effort": "low", + "include_reasoning": false + }); + let chat_request: ChatCompletionRequest = serde_json::from_value(val).unwrap(); + + // Verify the fields are correctly deserialized + assert_eq!(chat_request.echo, Some(true)); + assert!(!chat_request.include_reasoning); + assert!(matches!( + chat_request.reasoning_effort, + Some(lmdeploy_router_rs::protocols::spec::ReasoningEffort::Low) + )); + + // Route the request to mock server + let response = router.route_chat(None, &chat_request, None).await; + + // Should get a successful response + assert_eq!(response.status(), StatusCode::OK); + + let (_, body) = response.into_parts(); + let body_bytes = axum::body::to_bytes(body, usize::MAX).await.unwrap(); + let body_str = String::from_utf8(body_bytes.to_vec()).unwrap(); + let chat_response: serde_json::Value = serde_json::from_str(&body_str).unwrap(); + + // Verify it's a valid chat completion response + assert_eq!(chat_response["object"], "chat.completion"); +} diff --git a/lmdeploy/router/tests/test_pd_routing.rs b/lmdeploy/router/tests/test_pd_routing.rs new file mode 100644 index 0000000000..2d447645bf --- /dev/null +++ b/lmdeploy/router/tests/test_pd_routing.rs @@ -0,0 +1,156 @@ +#[cfg(test)] +mod test_pd_routing { + use lmdeploy_router_rs::config::{ + CircuitBreakerConfig, ConnectionMode, LMDeployMigrationProtocol, PolicyConfig, RetryConfig, + RouterConfig, RoutingMode, + }; + use lmdeploy_router_rs::routers::RouterFactory; + + // ======================================================================== + // Phase 1: Basic PD Components and Router Creation + // ======================================================================== + + #[test] + fn test_worker_types() { + use lmdeploy_router_rs::core::{WorkerFactory, WorkerType}; + + // Test worker creation for prefill servers + let prefill_worker = WorkerFactory::create_prefill("http://prefill:8080".to_string()); + assert_eq!(prefill_worker.url(), "http://prefill:8080"); + assert_eq!(prefill_worker.worker_type(), WorkerType::Prefill); + + // Test worker creation for decode servers + let decode_worker = WorkerFactory::create_decode("http://decode:8080".to_string()); + assert_eq!(decode_worker.url(), "http://decode:8080"); + match decode_worker.worker_type() { + WorkerType::Decode => (), + _ => panic!("Expected Decode worker type"), + } + + // Test regular worker creation + let regular_worker = WorkerFactory::create_regular("http://regular:8080".to_string()); + assert_eq!(regular_worker.url(), "http://regular:8080"); + match regular_worker.worker_type() { + WorkerType::Regular => (), + _ => panic!("Expected Regular worker type"), + } + } + + #[tokio::test] + async fn test_pd_router_configuration() { + // Test PD router configuration with various policies + // In the new structure, RoutingMode and PolicyConfig are separate + let test_cases = vec![ + ( + RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec![ + "http://prefill1:8080".to_string(), + "http://prefill2:8080".to_string(), + ], + decode_urls: vec![ + "http://decode1:8080".to_string(), + "http://decode2:8080".to_string(), + ], + prefill_policy: None, + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::Rdma, + rdma_config: None, + dummy_prefill: false, + }, + PolicyConfig::Random, + ), + ( + RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec!["http://prefill:8080".to_string()], + decode_urls: vec!["http://decode:8080".to_string()], + prefill_policy: None, + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::Rdma, + rdma_config: None, + dummy_prefill: false, + }, + PolicyConfig::PowerOfTwo { + load_check_interval_secs: 5, + }, + ), + ( + RoutingMode::LMDeployPrefillDecode { + prefill_urls: vec![ + "http://p1:8080".to_string(), + "http://p2:8080".to_string(), + "http://p3:8080".to_string(), + ], + decode_urls: vec!["http://d1:8080".to_string(), "http://d2:8080".to_string()], + prefill_policy: None, + decode_policy: None, + migration_protocol: LMDeployMigrationProtocol::Rdma, + rdma_config: None, + dummy_prefill: false, + }, + PolicyConfig::CacheAware { + cache_threshold: 0.7, + balance_abs_threshold: 20, + balance_rel_threshold: 1.2, + eviction_interval_secs: 60, + max_tree_size: 1000000, + }, + ), + ]; + + for (mode, policy) in test_cases { + let config = RouterConfig { + mode, + policy, + host: "127.0.0.1".to_string(), + port: 3001, + max_payload_size: 1024 * 1024, + request_timeout_secs: 60, + worker_startup_timeout_secs: 10, + worker_startup_check_interval_secs: 1, + api_key: None, + api_key_validation_urls: vec![], + discovery: None, + metrics: None, + log_dir: None, + log_level: None, + request_id_headers: None, + max_concurrent_requests: 64, + queue_size: 0, + queue_timeout_secs: 60, + cors_allowed_origins: vec![], + retry: RetryConfig::default(), + circuit_breaker: CircuitBreakerConfig::default(), + disable_retries: false, + disable_circuit_breaker: false, + health_check: lmdeploy_router_rs::config::HealthCheckConfig::default(), + enable_igw: false, + rate_limit_tokens_per_second: None, + connection_mode: ConnectionMode::Http, + history_backend: lmdeploy_router_rs::config::HistoryBackend::Memory, + enable_profiling: false, + profile_timeout_secs: 30, + kv_connector: lmdeploy_router_rs::config::KvConnector::Nixl, + }; + + // Router creation will fail due to health checks, but config should be valid + let app_context = lmdeploy_router_rs::server::AppContext::new( + config.clone(), + reqwest::Client::new(), + 64, + None, + config.api_key_validation_urls.clone(), + ) + .expect("Failed to create AppContext"); + let app_context = std::sync::Arc::new(app_context); + let result = RouterFactory::create_router(&app_context).await; + assert!(result.is_err()); + let error_msg = result.unwrap_err(); + // Error should be about health/timeout, not configuration + assert!( + error_msg.contains("healthy") || error_msg.contains("timeout"), + "Unexpected error: {}", + error_msg + ); + } + } +} diff --git a/lmdeploy/router/tests/test_prompt_input.rs b/lmdeploy/router/tests/test_prompt_input.rs new file mode 100644 index 0000000000..4ea69da9b1 --- /dev/null +++ b/lmdeploy/router/tests/test_prompt_input.rs @@ -0,0 +1,254 @@ +/// Tests for PromptInput enum supporting str, list[str], list[int], list[list[int]] +use lmdeploy_router_rs::protocols::spec::{CompletionRequest, PromptInput}; + +#[test] +fn test_prompt_input_single_string() { + let json = r#"{ + "model": "test-model", + "prompt": "Hello, world!", + "max_tokens": 100 + }"#; + + let req: CompletionRequest = serde_json::from_str(json).unwrap(); + assert!(matches!(req.prompt, PromptInput::String(_))); + + if let PromptInput::String(s) = &req.prompt { + assert_eq!(s, "Hello, world!"); + } + + assert_eq!(req.prompt.len(), 1); + assert!(!req.prompt.is_empty()); + assert!(!req.prompt.is_token_based()); + assert_eq!(req.prompt.extract_text_for_routing(), "Hello, world!"); +} + +#[test] +fn test_prompt_input_string_array() { + let json = r#"{ + "model": "test-model", + "prompt": ["Hello", "world", "test"], + "max_tokens": 100 + }"#; + + let req: CompletionRequest = serde_json::from_str(json).unwrap(); + assert!(matches!(req.prompt, PromptInput::StringArray(_))); + + if let PromptInput::StringArray(arr) = &req.prompt { + assert_eq!(arr.len(), 3); + assert_eq!(arr[0], "Hello"); + assert_eq!(arr[1], "world"); + assert_eq!(arr[2], "test"); + } + + assert_eq!(req.prompt.len(), 3); + assert!(!req.prompt.is_empty()); + assert!(!req.prompt.is_token_based()); + assert_eq!(req.prompt.extract_text_for_routing(), "Hello world test"); +} + +#[test] +fn test_prompt_input_single_int_array() { + let json = r#"{ + "model": "test-model", + "prompt": [128000, 9906, 11, 1917, 0], + "max_tokens": 100 + }"#; + + let req: CompletionRequest = serde_json::from_str(json).unwrap(); + assert!(matches!(req.prompt, PromptInput::IntArray(_))); + + if let PromptInput::IntArray(ids) = &req.prompt { + assert_eq!(ids.len(), 5); + assert_eq!(ids[0], 128000); + assert_eq!(ids[1], 9906); + assert_eq!(ids[4], 0); + } + + assert_eq!(req.prompt.len(), 1); + assert!(!req.prompt.is_empty()); + assert!(req.prompt.is_token_based()); + assert_eq!(req.prompt.extract_text_for_routing(), "token_ids:5"); + assert_eq!(req.prompt.estimated_token_count(), 5); +} + +#[test] +fn test_prompt_input_int_batch() { + let json = r#"{ + "model": "test-model", + "prompt": [ + [128000, 9906, 11, 1917, 0], + [128001, 9906, 11, 1917, 1] + ], + "max_tokens": 100 + }"#; + + let req: CompletionRequest = serde_json::from_str(json).unwrap(); + assert!(matches!(req.prompt, PromptInput::IntBatch(_))); + + if let PromptInput::IntBatch(batches) = &req.prompt { + assert_eq!(batches.len(), 2); + assert_eq!(batches[0].len(), 5); + assert_eq!(batches[1].len(), 5); + assert_eq!(batches[0][0], 128000); + assert_eq!(batches[1][0], 128001); + } + + assert_eq!(req.prompt.len(), 2); + assert!(!req.prompt.is_empty()); + assert!(req.prompt.is_token_based()); + assert_eq!( + req.prompt.extract_text_for_routing(), + "token_ids_batch:2:10" + ); + assert_eq!(req.prompt.estimated_token_count(), 10); +} + +#[test] +fn test_prompt_input_serialization_roundtrip_string() { + let req = CompletionRequest { + model: Some("test-model".to_string()), + prompt: PromptInput::String("Hello, world!".to_string()), + suffix: None, + max_tokens: Some(100), + temperature: None, + top_p: None, + n: None, + stream: false, + stream_options: None, + logprobs: None, + echo: false, + stop: None, + presence_penalty: None, + frequency_penalty: None, + best_of: None, + logit_bias: None, + user: None, + seed: None, + top_k: None, + min_p: None, + min_tokens: None, + repetition_penalty: None, + regex: None, + ebnf: None, + json_schema: None, + stop_token_ids: None, + no_stop_trim: false, + ignore_eos: false, + skip_special_tokens: true, + other: serde_json::Map::new(), + }; + + let json = serde_json::to_string(&req).unwrap(); + let parsed: CompletionRequest = serde_json::from_str(&json).unwrap(); + + assert_eq!(req.model, parsed.model); + assert_eq!(req.prompt, parsed.prompt); +} + +#[test] +fn test_prompt_input_serialization_roundtrip_int_array() { + let req = CompletionRequest { + model: Some("test-model".to_string()), + prompt: PromptInput::IntArray(vec![128000, 9906, 11, 1917, 0]), + suffix: None, + max_tokens: Some(100), + temperature: None, + top_p: None, + n: None, + stream: false, + stream_options: None, + logprobs: None, + echo: false, + stop: None, + presence_penalty: None, + frequency_penalty: None, + best_of: None, + logit_bias: None, + user: None, + seed: None, + top_k: None, + min_p: None, + min_tokens: None, + repetition_penalty: None, + regex: None, + ebnf: None, + json_schema: None, + stop_token_ids: None, + no_stop_trim: false, + ignore_eos: false, + skip_special_tokens: true, + other: serde_json::Map::new(), + }; + + let json = serde_json::to_string(&req).unwrap(); + let parsed: CompletionRequest = serde_json::from_str(&json).unwrap(); + + assert_eq!(req.model, parsed.model); + assert_eq!(req.prompt, parsed.prompt); +} + +#[test] +fn test_prompt_input_estimated_token_count() { + // String: estimate ~4 chars per token + let prompt1 = PromptInput::String("Hello, world! This is a test.".to_string()); // 30 chars + assert_eq!(prompt1.estimated_token_count(), 7); // 30 / 4 = 7 + + // String array + let prompt2 = PromptInput::StringArray(vec![ + "Hello".to_string(), // 5 chars + "world".to_string(), // 5 chars + ]); + assert_eq!(prompt2.estimated_token_count(), 2); // (5 + 5) / 4 = 2 + + // Int array: exact count + let prompt3 = PromptInput::IntArray(vec![1, 2, 3, 4, 5]); + assert_eq!(prompt3.estimated_token_count(), 5); + + // Int batch: sum of all arrays + let prompt4 = PromptInput::IntBatch(vec![vec![1, 2, 3], vec![4, 5, 6, 7]]); + assert_eq!(prompt4.estimated_token_count(), 7); +} + +#[test] +fn test_prompt_input_empty() { + let prompt1 = PromptInput::String("".to_string()); + assert!(prompt1.is_empty()); + + let prompt2 = PromptInput::StringArray(vec![]); + assert!(prompt2.is_empty()); + + let prompt3 = PromptInput::IntArray(vec![]); + assert!(prompt3.is_empty()); + + let prompt4 = PromptInput::IntBatch(vec![]); + assert!(prompt4.is_empty()); +} + +#[test] +fn test_prompt_input_deserialization_disambiguation() { + // Test that serde correctly disambiguates between different prompt types + + // Empty array should be IntArray (not IntBatch or StringArray) + // Actually, empty array could be any of them due to untagged enum + // Let's test non-empty cases + + // Array with strings + let json1 = r#"["hello", "world"]"#; + let prompt1: PromptInput = serde_json::from_str(json1).unwrap(); + assert!(matches!(prompt1, PromptInput::StringArray(_))); + + // Array with numbers (single-level) + let json2 = r#"[1, 2, 3]"#; + let prompt2: PromptInput = serde_json::from_str(json2).unwrap(); + assert!(matches!(prompt2, PromptInput::IntArray(_))); + + // Array with nested arrays (batch) + let json3 = r#"[[1, 2], [3, 4]]"#; + let prompt3: PromptInput = serde_json::from_str(json3).unwrap(); + assert!(matches!(prompt3, PromptInput::IntBatch(_))); + + // Single string + let json4 = r#""hello world""#; + let prompt4: PromptInput = serde_json::from_str(json4).unwrap(); + assert!(matches!(prompt4, PromptInput::String(_))); +} diff --git a/lmdeploy/router/tests/test_transparent_proxy_routing.rs b/lmdeploy/router/tests/test_transparent_proxy_routing.rs new file mode 100644 index 0000000000..a10daabeca --- /dev/null +++ b/lmdeploy/router/tests/test_transparent_proxy_routing.rs @@ -0,0 +1,508 @@ +//! Tests for transparent proxy routing with headers and availability filtering. +//! +//! These tests verify the fixes to route_transparent() across router +//! implementations (Router, LMDeployPDRouter): +//! 1. Headers are passed to select_worker_with_headers() for consistent hash routing +//! 2. Workers are filtered by is_available() before selection +//! 3. The inline header conversion pattern used by transparent proxy routes matches +//! the Router::headers_to_request_headers() output + +#[cfg(test)] +mod transparent_proxy_routing_tests { + use std::collections::HashMap; + use std::collections::HashSet; + use std::sync::Arc; + + use lmdeploy_router_rs::core::BasicWorker; + use lmdeploy_router_rs::core::Worker; + use lmdeploy_router_rs::core::WorkerType; + use lmdeploy_router_rs::policies::ConsistentHashPolicy; + use lmdeploy_router_rs::policies::LoadBalancingPolicy; + use lmdeploy_router_rs::policies::RequestHeaders; + + /// Helper to create test workers + fn create_workers(n: usize) -> Vec> { + (0..n) + .map(|i| { + Arc::new(BasicWorker::new( + format!("http://worker{}:8080", i + 1), + WorkerType::Regular, + )) as Arc + }) + .collect() + } + + /// Helper to build RequestHeaders from key-value pairs + fn make_headers(pairs: &[(&str, &str)]) -> RequestHeaders { + pairs + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect() + } + + // ===================================================================== + // Test 1: Availability filtering before policy selection + // ===================================================================== + // The route_transparent fix filters workers by is_available() BEFORE + // passing them to the policy. These tests verify that behavior. + + #[test] + fn test_availability_filter_excludes_unhealthy_workers() { + let policy = ConsistentHashPolicy::new(); + let workers = create_workers(3); + + // Mark worker1 as unhealthy + workers[0].set_healthy(false); + + // Filter by availability (same pattern as route_transparent) + let available: Vec> = workers + .iter() + .filter(|w| w.is_available()) + .cloned() + .collect(); + + assert_eq!( + available.len(), + 2, + "Only 2 of 3 workers should be available" + ); + + // Policy should still work with filtered workers + let headers = make_headers(&[("x-session-id", "test-session")]); + let result = policy.select_worker_with_headers( + &available, + Some(r#"{"prompt": "test"}"#), + Some(&headers), + ); + assert!(result.is_some(), "Should select from available workers"); + + let idx = result.unwrap(); + assert!( + available[idx].is_healthy(), + "Selected worker must be healthy" + ); + } + + #[test] + fn test_availability_filter_returns_empty_when_all_unhealthy() { + let workers = create_workers(3); + + // Mark all workers as unhealthy + for w in &workers { + w.set_healthy(false); + } + + let available: Vec> = workers + .iter() + .filter(|w| w.is_available()) + .cloned() + .collect(); + + assert!( + available.is_empty(), + "No workers should be available when all are unhealthy" + ); + + // Policy should return None for empty list + let policy = ConsistentHashPolicy::new(); + let result = policy.select_worker_with_headers(&available, None, None); + assert!(result.is_none(), "Should return None for empty worker list"); + } + + // ===================================================================== + // Test 2: Headers are forwarded to the consistent hash policy + // ===================================================================== + // Before the fix, route_transparent called select_worker() which ignores + // headers entirely, making consistent hash fall back to request body hash. + // After the fix, it calls select_worker_with_headers() so x-session-id + // in headers produces sticky routing. + + #[test] + fn test_session_id_header_produces_sticky_routing() { + let policy = ConsistentHashPolicy::new(); + let workers = create_workers(5); + + let headers = make_headers(&[("x-session-id", "my-sticky-session")]); + + // Simulate what route_transparent now does: pass headers to policy + let mut selected: Vec = Vec::new(); + for i in 0..20 { + let body = format!( + r#"{{"prompt": "different prompt {}", "model": "default"}}"#, + i + ); + if let Some(idx) = + policy.select_worker_with_headers(&workers, Some(&body), Some(&headers)) + { + selected.push(idx); + } + } + + // All requests with the same session ID should route to the same worker + assert!(!selected.is_empty()); + let first = selected[0]; + for (i, &idx) in selected.iter().enumerate() { + assert_eq!( + idx, first, + "Request {} routed to worker {}, expected {} (session stickiness broken)", + i, idx, first + ); + } + } + + #[test] + fn test_without_headers_uses_body_fallback() { + let policy = ConsistentHashPolicy::new(); + let workers = create_workers(3); + + // Without headers, consistent hash should fall back to body content + let body = r#"{"session_id": "body-session", "prompt": "test"}"#; + + let mut selected: Vec = Vec::new(); + for _ in 0..10 { + if let Some(idx) = policy.select_worker_with_headers(&workers, Some(body), None) { + selected.push(idx); + } + } + + // Same body → same worker (body fallback is deterministic) + assert!(!selected.is_empty()); + let first = selected[0]; + for &idx in &selected { + assert_eq!(idx, first, "Body fallback should be deterministic"); + } + } + + #[test] + fn test_header_session_id_takes_priority_over_body_session_id() { + let policy = ConsistentHashPolicy::new(); + let workers = create_workers(5); + + let header_session = "header-session-wins"; + let body_session = "body-session-ignored"; + + let headers = make_headers(&[("x-session-id", header_session)]); + let body = format!(r#"{{"session_id": "{}", "prompt": "test"}}"#, body_session); + + // Route with both header and body session ID + let with_both = policy + .select_worker_with_headers(&workers, Some(&body), Some(&headers)) + .expect("Should select a worker"); + + // Route with header only (different body) + let header_only = policy + .select_worker_with_headers( + &workers, + Some(r#"{"prompt": "completely different"}"#), + Some(&headers), + ) + .expect("Should select a worker"); + + assert_eq!( + with_both, header_only, + "Header x-session-id should take priority - different body should not change routing" + ); + } + + // ===================================================================== + // Test 3: Availability + headers work together correctly + // ===================================================================== + // These tests verify the combined effect: availability filtering happens + // BEFORE policy selection, and headers are passed to the policy. + + #[test] + fn test_sticky_routing_with_availability_filtering() { + let policy = ConsistentHashPolicy::new(); + let workers = create_workers(4); + + // Initial routing with all workers available + let headers = make_headers(&[("x-session-id", "stable-session")]); + let body = r#"{"prompt": "test"}"#; + + let initial_idx = policy + .select_worker_with_headers(&workers, Some(body), Some(&headers)) + .expect("Should select a worker"); + + // Now filter by availability (as route_transparent does) + let available: Vec> = workers + .iter() + .filter(|w| w.is_available()) + .cloned() + .collect(); + + // Should get the same worker since all are still available + let after_filter_idx = policy + .select_worker_with_headers(&available, Some(body), Some(&headers)) + .expect("Should select a worker"); + + assert_eq!( + workers[initial_idx].url(), + available[after_filter_idx].url(), + "Same session ID should route to same worker URL when all workers are available" + ); + } + + #[test] + fn test_distribution_with_headers_across_sessions() { + let policy = ConsistentHashPolicy::new(); + let workers = create_workers(3); + + let mut worker_urls: HashSet = HashSet::new(); + + for i in 0..100 { + let session = format!("unique-session-{}", i); + let headers = make_headers(&[("x-session-id", &session)]); + + if let Some(idx) = policy.select_worker_with_headers( + &workers, + Some(r#"{"prompt": "test"}"#), + Some(&headers), + ) { + worker_urls.insert(workers[idx].url().to_string()); + } + } + + // With 100 different sessions and 3 workers, all workers should be used + assert!( + worker_urls.len() >= 2, + "Expected distribution, only used {} workers: {:?}", + worker_urls.len(), + worker_urls + ); + } + + // ===================================================================== + // Test 4: Inline header conversion correctness + // ===================================================================== + // Transparent proxy routes use an inline pattern to convert + // HeaderMap → HashMap. Verify it produces correct output. + + #[test] + fn test_inline_header_conversion_lowercases_keys() { + use axum::http::HeaderMap; + use axum::http::HeaderValue; + + let mut header_map = HeaderMap::new(); + header_map.insert("X-Session-Id", HeaderValue::from_static("abc-123")); + header_map.insert("X-USER-ID", HeaderValue::from_static("user-456")); + header_map.insert("content-type", HeaderValue::from_static("application/json")); + + // Simulate the transparent proxy inline conversion. + let request_headers: Option> = Some(&header_map).map(|h| { + h.iter() + .filter_map(|(name, value)| { + value + .to_str() + .ok() + .map(|v| (name.as_str().to_lowercase(), v.to_string())) + }) + .collect() + }); + + let headers = request_headers.unwrap(); + + // axum normalizes header names to lowercase already, but verify the + // to_lowercase() call in our conversion is idempotent and correct + assert_eq!(headers.get("x-session-id").unwrap(), "abc-123"); + assert_eq!(headers.get("x-user-id").unwrap(), "user-456"); + assert_eq!(headers.get("content-type").unwrap(), "application/json"); + } + + #[test] + fn test_inline_header_conversion_used_by_policy() { + use axum::http::HeaderMap; + use axum::http::HeaderValue; + + let policy = ConsistentHashPolicy::new(); + let workers = create_workers(3); + + // Convert via the transparent proxy inline pattern. + let mut header_map = HeaderMap::new(); + header_map.insert( + "x-session-id", + HeaderValue::from_static("policy-test-session"), + ); + + let request_headers: Option> = Some(&header_map).map(|h| { + h.iter() + .filter_map(|(name, value)| { + value + .to_str() + .ok() + .map(|v| (name.as_str().to_lowercase(), v.to_string())) + }) + .collect() + }); + + // Use the converted headers with the policy (same as route_transparent now does) + let mut selected: Vec = Vec::new(); + for i in 0..10 { + let body = format!(r#"{{"prompt": "request {}"}}"#, i); + if let Some(idx) = + policy.select_worker_with_headers(&workers, Some(&body), request_headers.as_ref()) + { + selected.push(idx); + } + } + + // All should go to the same worker (session stickiness via header) + assert!(!selected.is_empty()); + let first = selected[0]; + for &idx in &selected { + assert_eq!( + idx, first, + "Inline header conversion should produce sticky routing" + ); + } + + // Also verify with make_headers helper (which matches RequestHeaders directly) + let direct_headers = make_headers(&[("x-session-id", "policy-test-session")]); + let direct_result = policy + .select_worker_with_headers( + &workers, + Some(r#"{"prompt": "request 0"}"#), + Some(&direct_headers), + ) + .expect("Direct headers should work"); + + assert_eq!( + workers[first].url(), + workers[direct_result].url(), + "Inline conversion and direct headers should route to the same worker" + ); + } + + // ===================================================================== + // Test 5: PD mode worker pair selection with headers + // ===================================================================== + // For PD routers, both prefill and decode workers need headers. + + #[test] + fn test_pd_mode_worker_pair_with_headers() { + let policy = ConsistentHashPolicy::new(); + + let prefill_workers: Vec> = (0..3) + .map(|i| { + Arc::new(BasicWorker::new( + format!("http://prefill{}:8080", i + 1), + WorkerType::Prefill, + )) as Arc + }) + .collect(); + + let decode_workers: Vec> = (0..3) + .map(|i| { + Arc::new(BasicWorker::new( + format!("http://decode{}:8080", i + 1), + WorkerType::Decode, + )) as Arc + }) + .collect(); + + let headers = make_headers(&[("x-session-id", "pd-session")]); + let body = r#"{"prompt": "test"}"#; + + // select_worker_with_headers on each pool + let prefill_idx = policy + .select_worker_with_headers(&prefill_workers, Some(body), Some(&headers)) + .expect("Should select prefill worker"); + + let decode_idx = policy + .select_worker_with_headers(&decode_workers, Some(body), Some(&headers)) + .expect("Should select decode worker"); + + // Verify consistency: same session → same workers + let prefill_idx2 = policy + .select_worker_with_headers(&prefill_workers, Some(body), Some(&headers)) + .expect("Should select prefill worker again"); + + let decode_idx2 = policy + .select_worker_with_headers(&decode_workers, Some(body), Some(&headers)) + .expect("Should select decode worker again"); + + assert_eq!( + prefill_idx, prefill_idx2, + "Prefill selection should be consistent" + ); + assert_eq!( + decode_idx, decode_idx2, + "Decode selection should be consistent" + ); + + // Verify the workers are valid + assert!(prefill_idx < prefill_workers.len()); + assert!(decode_idx < decode_workers.len()); + } + + #[test] + fn test_pd_mode_availability_filtering() { + let prefill_workers: Vec> = (0..3) + .map(|i| { + Arc::new(BasicWorker::new( + format!("http://prefill{}:8080", i + 1), + WorkerType::Prefill, + )) as Arc + }) + .collect(); + + let decode_workers: Vec> = (0..3) + .map(|i| { + Arc::new(BasicWorker::new( + format!("http://decode{}:8080", i + 1), + WorkerType::Decode, + )) as Arc + }) + .collect(); + + // Mark one prefill and one decode worker as unhealthy + prefill_workers[0].set_healthy(false); + decode_workers[1].set_healthy(false); + + // Filter by availability (as route_transparent now does) + let available_prefill: Vec> = prefill_workers + .iter() + .filter(|w| w.is_available()) + .cloned() + .collect(); + + let available_decode: Vec> = decode_workers + .iter() + .filter(|w| w.is_available()) + .cloned() + .collect(); + + assert_eq!( + available_prefill.len(), + 2, + "2 of 3 prefill workers should be available" + ); + assert_eq!( + available_decode.len(), + 2, + "2 of 3 decode workers should be available" + ); + + // Policy should select from available workers only + let policy = ConsistentHashPolicy::new(); + let headers = make_headers(&[("x-session-id", "pd-avail-test")]); + + let prefill_idx = policy + .select_worker_with_headers( + &available_prefill, + Some(r#"{"prompt": "test"}"#), + Some(&headers), + ) + .expect("Should select available prefill worker"); + + let decode_idx = policy + .select_worker_with_headers( + &available_decode, + Some(r#"{"prompt": "test"}"#), + Some(&headers), + ) + .expect("Should select available decode worker"); + + assert!(available_prefill[prefill_idx].is_healthy()); + assert!(available_decode[decode_idx].is_healthy()); + } +} diff --git a/lmdeploy/router/tests/tokenizer_integration.rs b/lmdeploy/router/tests/tokenizer_integration.rs new file mode 100644 index 0000000000..2f9f9907af --- /dev/null +++ b/lmdeploy/router/tests/tokenizer_integration.rs @@ -0,0 +1,320 @@ +//! Integration tests for tokenizers using real tokenizer data +//! +//! These tests download the TinyLlama tokenizer from HuggingFace to verify our tokenizer +//! implementation works correctly with real-world tokenizer files. + +mod common; +use common::{ensure_tokenizer_cached, EXPECTED_HASHES, TEST_PROMPTS}; + +use lmdeploy_router_rs::tokenizer::{ + factory, huggingface::HuggingFaceTokenizer, sequence::Sequence, stop::*, stream::DecodeStream, + traits::*, +}; +use std::sync::Arc; + +const LONG_TEST_PROMPTS: [(&str, &str); 6] = [ + ("Tell me about the following text.", "Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat."), + ("Tell me about the following text.", "Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum."), + ("Tell me about the following text.", "Sed ut perspiciatis unde omnis iste natus error sit voluptatem accusantium doloremque laudantium, totam rem aperiam, eaque ipsa quae ab illo inventore veritatis et quasi architecto beatae vitae dicta sunt explicabo. Nemo enim ipsam voluptatem quia voluptas sit aspernatur aut odit aut fugit, sed quia consequuntur magni dolores eos qui ratione voluptatem sequi nesciunt."), + ("Tell me about the following text.", "Neque porro quisquam est, qui dolorem ipsum quia dolor sit amet, consectetur, adipisci velit, sed quia non numquam eius modi tempora incidunt ut labore et dolore magnam aliquam quaerat voluptatem."), + // Tennis-themed prompt for variety + ("Tell me about the following text.", "In the ancient realm of Tennisia, the very magic of the land is drawn from the sport itself. Forehands light the skies, backhands carve the earth, and serves rumble like thunder across kingdoms. At the center of this balance lie four sacred Grand Slam relics: the Sapphire Trophy of Melbourne, the Emerald Chalice of Paris, the Ruby Crown of London, and the Diamond Orb of New York. Together, they keep the game's spirit alive. + But the relics are scattered, guarded by champions of legendary skill. The first is the Fire King of Clay, ruler of the crimson courts, whose topspin arcs blaze high and heavy, scorching all who dare stand across from him. The second is the Tempest Trickster, master of the baseline fortress, whose footwork and precision can turn back any storm, and whose returns arrive as if pulled by invisible strings. The third is the Shadow-Dancer of the Highlands, a tactician who thrives in the long rallies of twilight, changing pace and spin until opponents lose their rhythm. The fourth and final guardian is a towering Diamond Titan, a net-charging colossus whose volleys shatter the air itself. + Into this arena of gods steps the Silver-Wristed Knight — a player of impossible grace, whose game is an art form. His quest: to claim each relic not for glory, but to restore harmony to the rankings of the realm. + He travels across the Kingdom of Clay, where the points stretch like marathons and the air tastes of iron; through the Grasslands of London, where the ball skids low and the margins are razor-thin; over the Hard Courts of the East, where rallies turn into duels of endurance; and finally to the Cathedral of Lights in New York, where night matches burn with fevered energy. + Each battle is played under enchanted floodlights, the lines patrolled by spectral line judges whose calls are final. The crowd's roar swells with every break point, and the Silver-Wristed Knight's racket glows brightest when the match teeters at deuce. There are moments when doubt grips him — when his serve falters or his touch deserts him — but each challenge teaches a new stroke, culminating in the legendary Forehand of Dawn. + When the last relic is claimed, he stands not as a conqueror but as a custodian of the game, knowing that rivalries forge the very magic he protects. The balance is restored — until the next season begins."), + // Emoji stress test + ("Tell me about the following text.", "😀😃😄😁😆🥹😅😂🤣🥲☺️😊😇🙂🙃😉🤩😎 🤪🥳🤓🙄🤪😵👻") +]; + +fn compute_hashes_for_tokenizer(tokenizer: &E, prompts: &[&str]) -> Vec { + prompts + .iter() + .map(|&prompt| { + tokenizer + .encode(prompt) + .expect("Failed to encode prompt") + .get_hash() + }) + .collect() +} + +#[test] +fn test_huggingface_tokenizer_hashes() { + let tokenizer_path = ensure_tokenizer_cached(); + let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap()) + .expect("Failed to load HuggingFace tokenizer"); + + let prompt_hashes = compute_hashes_for_tokenizer(&tokenizer, &TEST_PROMPTS); + + println!( + "HF Tokenizer: {:?}\nComputed Hashes: {:?}\nExpected Hashes: {:?}", + tokenizer_path, prompt_hashes, EXPECTED_HASHES + ); + + assert_eq!(prompt_hashes, EXPECTED_HASHES); +} + +#[test] +fn test_tokenizer_encode_decode_lifecycle() { + let tokenizer_path = ensure_tokenizer_cached(); + let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap()) + .expect("Failed to load HuggingFace tokenizer"); + + for prompt in TEST_PROMPTS.iter() { + let encoding = tokenizer.encode(prompt).expect("Failed to encode prompt"); + + let decoded = tokenizer + .decode(encoding.token_ids(), false) + .expect("Failed to decode token_ids"); + + assert_eq!(decoded, *prompt, "Encode-decode mismatch for: {}", prompt); + } +} + +#[test] +fn test_sequence_operations() { + let tokenizer_path = ensure_tokenizer_cached(); + let tokenizer = Arc::new( + HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap()) + .expect("Failed to load tokenizer"), + ); + + for prompt in TEST_PROMPTS.iter() { + let encoding = tokenizer.encode(prompt).expect("Failed to encode prompt"); + + // Test Sequence with append_text + let mut sequence = Sequence::new(tokenizer.clone()); + sequence.append_text(prompt).expect("Failed to append text"); + + assert_eq!( + sequence.len(), + encoding.token_ids().len(), + "Sequence length mismatch" + ); + assert_eq!(sequence.text().unwrap(), *prompt, "Sequence text mismatch"); + + // Test incremental decoding with append_token + let mut decoder = Sequence::new(tokenizer.clone()); + let mut output = String::new(); + + for token_id in encoding.token_ids() { + let text = decoder + .append_token(*token_id) + .expect("Failed to append token"); + output.push_str(&text); + } + + assert_eq!(decoder.len(), sequence.len(), "Decoder length mismatch"); + assert_eq!( + decoder.token_ids(), + sequence.token_ids(), + "Token IDs mismatch" + ); + assert_eq!(output, *prompt, "Incremental decode mismatch"); + } +} + +#[test] +fn test_decode_stream() { + let tokenizer_path = ensure_tokenizer_cached(); + let tokenizer = Arc::new( + HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap()) + .expect("Failed to load tokenizer"), + ); + + for prompt in TEST_PROMPTS.iter() { + let encoding = tokenizer.encode(prompt).expect("Failed to encode prompt"); + + let mut decoder = DecodeStream::new(tokenizer.clone(), &[], false); + let mut output = String::new(); + + for token_id in encoding.token_ids() { + if let Some(text) = decoder.step(*token_id).expect("Failed to decode token") { + output.push_str(&text); + } + } + + assert_eq!(output, *prompt, "DecodeStream output mismatch"); + } +} + +#[test] +fn test_long_sequence_incremental_decode_with_prefill() { + let tokenizer_path = ensure_tokenizer_cached(); + let tokenizer = Arc::new( + HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap()) + .expect("Failed to load tokenizer"), + ); + + for (input_text, output_text) in LONG_TEST_PROMPTS.iter() { + let input_encoding = tokenizer + .encode(input_text) + .expect("Failed to encode input"); + + let output_encoding = tokenizer + .encode(output_text) + .expect("Failed to encode output"); + + let mut decoder = DecodeStream::new(tokenizer.clone(), input_encoding.token_ids(), false); + + let mut output = String::new(); + for token_id in output_encoding.token_ids() { + if let Some(text) = decoder.step(*token_id).expect("Failed to decode token") { + output.push_str(&text); + } + } + + assert_eq!(output.trim(), *output_text, "Long sequence decode mismatch"); + } +} + +#[test] +fn test_stop_sequence_decoder() { + let tokenizer_path = ensure_tokenizer_cached(); + let tokenizer = Arc::new( + HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap()) + .expect("Failed to load tokenizer"), + ); + + // Test with various stop sequences + let test_cases = vec![ + ( + "Hello world! Stop here. Continue after.", + "Stop", + "Hello world! ", + ), + ("Testing stop sequences.", ".", "Testing stop sequences"), + ("No stop sequence here", "xyz", "No stop sequence here"), + ]; + + for (input, stop_seq, expected) in test_cases { + let config = StopSequenceConfig::default().with_stop_sequence(stop_seq); + + let mut decoder = StopSequenceDecoder::new(tokenizer.clone(), config, false); + + let encoding = tokenizer.encode(input).expect("Failed to encode"); + let mut output = String::new(); + let mut stopped = false; + + for token_id in encoding.token_ids() { + match decoder.process_token(*token_id).unwrap() { + SequenceDecoderOutput::Text(text) => output.push_str(&text), + SequenceDecoderOutput::StoppedWithText(text) => { + output.push_str(&text); + stopped = true; + break; + } + SequenceDecoderOutput::Stopped => { + stopped = true; + break; + } + SequenceDecoderOutput::Held => {} + } + } + + if !stopped { + // Flush any remaining text + if let SequenceDecoderOutput::Text(text) = decoder.flush() { + output.push_str(&text); + } + } + + println!( + "Input: '{}', Stop: '{}', Output: '{}', Expected: '{}'", + input, stop_seq, output, expected + ); + + // The test should check if output starts with expected + // since stop sequences might not be perfectly aligned with token boundaries + assert!( + output.starts_with(expected) || output == input, + "Stop sequence test failed" + ); + } +} + +#[test] +fn test_factory_creation() { + // Test factory creation method + let tokenizer_path = ensure_tokenizer_cached(); + let tokenizer = factory::create_tokenizer(tokenizer_path.to_str().unwrap()) + .expect("Failed to create tokenizer via factory"); + + let encoding = tokenizer.encode(TEST_PROMPTS[0]).expect("Failed to encode"); + + let decoded = tokenizer + .decode(encoding.token_ids(), false) + .expect("Failed to decode"); + + assert_eq!(decoded, TEST_PROMPTS[0]); +} + +#[test] +fn test_batch_encoding() { + let tokenizer_path = ensure_tokenizer_cached(); + let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap()) + .expect("Failed to load tokenizer"); + + let encodings = tokenizer + .encode_batch(&TEST_PROMPTS) + .expect("Failed to batch encode"); + + assert_eq!(encodings.len(), TEST_PROMPTS.len()); + + for (i, encoding) in encodings.iter().enumerate() { + let decoded = tokenizer + .decode(encoding.token_ids(), false) + .expect("Failed to decode"); + assert_eq!(decoded, TEST_PROMPTS[i]); + } +} + +#[test] +fn test_special_tokens() { + use lmdeploy_router_rs::tokenizer::traits::Tokenizer as TokenizerTrait; + + let tokenizer_path = ensure_tokenizer_cached(); + let tokenizer = HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap()) + .expect("Failed to load tokenizer"); + + let special_tokens = tokenizer.get_special_tokens(); + + // TinyLlama should have at least BOS and EOS tokens + assert!(special_tokens.bos_token.is_some()); + assert!(special_tokens.eos_token.is_some()); + + println!("Special tokens: {:?}", special_tokens); +} + +#[test] +fn test_thread_safety() { + use std::thread; + + let tokenizer_path = ensure_tokenizer_cached(); + let tokenizer = Arc::new( + HuggingFaceTokenizer::from_file(tokenizer_path.to_str().unwrap()) + .expect("Failed to load tokenizer"), + ); + + let handles: Vec<_> = TEST_PROMPTS + .iter() + .map(|&prompt| { + let tokenizer_clone = tokenizer.clone(); + thread::spawn(move || { + let encoding = tokenizer_clone + .encode(prompt) + .expect("Failed to encode in thread"); + let decoded = tokenizer_clone + .decode(encoding.token_ids(), false) + .expect("Failed to decode in thread"); + assert_eq!(decoded, prompt); + }) + }) + .collect(); + + for handle in handles { + handle.join().expect("Thread panicked"); + } +} diff --git a/lmdeploy/serve/core/chat_runner.py b/lmdeploy/serve/core/chat_runner.py index 2246fac2e0..8b9d5674e6 100644 --- a/lmdeploy/serve/core/chat_runner.py +++ b/lmdeploy/serve/core/chat_runner.py @@ -41,6 +41,7 @@ class ChatStreamChunk: cached_tokens: int routed_experts: Any = None cache_block_ids: list[int] | None = None + remote_session_id: int | None = None reasoning_tokens: int | None = None is_last_delta: bool = True @@ -61,6 +62,7 @@ class ChatResult: routed_experts: Any = None cache_block_ids: list[Any] = field(default_factory=list) remote_token_ids: list[Any] = field(default_factory=list) + remote_session_id: int | None = None reasoning_tokens: int | None = None @@ -205,6 +207,7 @@ async def stream(self) -> AsyncGenerator[ChatStreamChunk, None]: cached_tokens=res.cached_tokens, routed_experts=res.routed_experts if finish_reason is not None else None, cache_block_ids=res.cache_block_ids, + remote_session_id=self.session.session_id, reasoning_tokens=self.response_parser.reasoning_tokens, is_last_delta=is_last_delta, ) @@ -258,6 +261,7 @@ async def collect(self, raw_request=None) -> ChatResult: routed_experts=final_res.routed_experts, cache_block_ids=cache_block_ids, remote_token_ids=remote_token_ids, + remote_session_id=self.session.session_id, reasoning_tokens=self.response_parser.reasoning_tokens, ) diff --git a/lmdeploy/serve/openai/chat_completions/serving.py b/lmdeploy/serve/openai/chat_completions/serving.py index ea57c43221..4db8d3d67c 100644 --- a/lmdeploy/serve/openai/chat_completions/serving.py +++ b/lmdeploy/serve/openai/chat_completions/serving.py @@ -305,6 +305,7 @@ async def _completion_stream_generator() -> AsyncGenerator[str, None]: if chunk.cache_block_ids is not None and chunk.is_last_delta: response_json['cache_block_ids'] = chunk.cache_block_ids response_json['remote_token_ids'] = chunk.token_ids + response_json['remote_session_id'] = chunk.remote_session_id yield f'data: {json.dumps(response_json)}\n\n' if final_usage is not None: yield f'data: {create_stream_usage_response_json(final_usage)}\n\n' @@ -385,5 +386,6 @@ async def completion_stream_generator() -> AsyncGenerator[str, None]: if with_cache: response['cache_block_ids'] = cache_block_ids response['remote_token_ids'] = remote_token_ids + response['remote_session_id'] = res.remote_session_id return response diff --git a/lmdeploy/serve/proxy/proxy.py b/lmdeploy/serve/proxy/proxy.py deleted file mode 100644 index 7b39790ae0..0000000000 --- a/lmdeploy/serve/proxy/proxy.py +++ /dev/null @@ -1,952 +0,0 @@ -# Copyright (c) OpenMMLab. All rights reserved. - -import asyncio -import copy -import json -import os -import os.path as osp -import random -import threading -import time -from collections import deque -from http import HTTPStatus -from typing import Literal - -import aiohttp -import numpy as np -import requests -import uvicorn -from fastapi import BackgroundTasks, Depends, FastAPI, Request -from fastapi.middleware.cors import CORSMiddleware -from fastapi.responses import JSONResponse, StreamingResponse -from pydantic import BaseModel, Field - -from lmdeploy.pytorch.disagg.config import DistServeRDMAConfig, EngineRole, RDMALinkType, ServingStrategy -from lmdeploy.pytorch.disagg.conn.protocol import MigrationProtocol, MigrationRequest -from lmdeploy.pytorch.disagg.conn.proxy_conn import PDConnectionPool -from lmdeploy.pytorch.disagg.messages import PDConnectionMessage -from lmdeploy.serve.openai.errors import create_error_response -from lmdeploy.serve.openai.protocol import ( - ChatCompletionRequest, - CompletionRequest, - ModelCard, # noqa: E501 - ModelList, - ModelPermission, -) -from lmdeploy.serve.proxy.utils import AIOHTTP_TIMEOUT, LATENCY_DEQUE_LEN, ErrorCodes, RoutingStrategy, err_msg -from lmdeploy.serve.utils.server_utils import validate_json_request -from lmdeploy.utils import get_logger - -from .streaming_response import ProxyStreamingResponse -from .utils import APIServerException - -logger = get_logger('lmdeploy') - - -class Status(BaseModel): - """Status protocol consists of models' information.""" - role: EngineRole = EngineRole.Hybrid - models: list[str] = Field(default=[], examples=[[]]) - unfinished: int = 0 - latency: deque = Field(default=deque(maxlen=LATENCY_DEQUE_LEN), examples=[[]]) - speed: int | None = Field(default=None, examples=[None]) - - -class Node(BaseModel): - """Node protocol consists of url and status.""" - url: str - status: Status | None = None - - -CONTROLLER_HEART_BEAT_EXPIRATION = int(os.getenv('LMDEPLOY_CONTROLLER_HEART_BEAT_EXPIRATION', 90)) - - -def heart_beat_controller(proxy_controller): - while True: - time.sleep(CONTROLLER_HEART_BEAT_EXPIRATION) - logger.info('Start heart beat check') - proxy_controller.remove_stale_nodes_by_expiration() - - -class NodeManager: - """Manage all the sub nodes. - - Args: - config_path (str): the path of the config file. - strategy (str): the strategy to dispatch node to handle the requests. - - **random**: not fully radom, but decided by the speed of nodes. - - **min_expected_latency**: will compute the expected latency to - process the requests. The sooner of the node, the more requests - will be dispatched to it. - - **min_observed_latency**: Based on previous finished requests. The - sooner they get processed, the more requests will be dispatched - to. - """ - - def __init__(self, - config_path: str | None = None, - serving_strategy: str = 'Hybrid', - routing_strategy: str = 'min_expected_latency', - migration_protocol: str = 'RDMA', - link_type: str = 'RoCE', - with_gdr: bool = True, - cache_status: bool = True) -> None: - self.nodes = dict() - self.serving_strategy = ServingStrategy[serving_strategy] - self.routing_strategy = RoutingStrategy.from_str(routing_strategy) - - self.cache_status = cache_status - self.latencies = dict() - self.config_path = osp.join(osp.dirname(osp.realpath(__file__)), 'proxy_config.json') - if config_path is not None: - self.config_path = config_path - if osp.exists(self.config_path) and self.cache_status: - with open(self.config_path) as config_file: - if os.path.getsize(self.config_path) > 0: - logger.info(f'loading node configuration: {self.config_path}') - config = json.load(config_file) - self.nodes = { - node_url: Status.model_validate_json(node_status) - for node_url, node_status in config.items() - } - self.heart_beat_thread = threading.Thread(target=heart_beat_controller, args=(self, ), daemon=True) - self.heart_beat_thread.start() - self.aiotimeout = aiohttp.ClientTimeout(total=AIOHTTP_TIMEOUT) - - # For PD Disaggregation - self.migration_protocol = MigrationProtocol[migration_protocol] - self.rdma_config = DistServeRDMAConfig(with_gdr=with_gdr, link_type=RDMALinkType[link_type]) - self.pd_connection_pool = PDConnectionPool() - self.dummy_prefill = False - - def get_nodes(self, role: EngineRole) -> dict[str, Status]: - items = list(self.nodes.items()) - return {node_url: node_status for (node_url, node_status) in items if node_status.role == role} - - @property - def hybrid_nodes(self): - return self.get_nodes(EngineRole.Hybrid) - - @property - def prefill_nodes(self): - return self.get_nodes(EngineRole.Prefill) - - @property - def decode_nodes(self): - return self.get_nodes(EngineRole.Decode) - - def update_config_file(self): - """Update the config file.""" - nodes = copy.deepcopy(self.nodes) - for _, status in nodes.items(): - status.latency = deque(list(status.latency)[-LATENCY_DEQUE_LEN:]) - if self.cache_status: - with open(self.config_path, 'w') as config_file: # update cfg yml - json.dump({ - node_url: node_status.model_dump_json() - for node_url, node_status in nodes.items() - }, - config_file, - indent=2) - - def add(self, node_url: str, status: Status | None = None): - """Add a node to the manager. - - Args: - node_url (str): A http url. Can be the url generated by - `lmdeploy serve api_server`. - status (dict): The status of the node. An example: - {'models': ['intern-s2-preview'], 'speed': -1}. - The speed here can be RPM or other metric. All the - values of nodes should be the same metric. - """ - if status is None: - status = self.nodes.get(node_url, Status()) - if status.models != []: # force register directly - self.remove(node_url) - self.nodes[node_url] = status - self.update_config_file() - return - try: - from lmdeploy.serve.openai.api_client import APIClient - client = APIClient(api_server_url=node_url) - status.models = client.available_models - self.nodes[node_url] = status - except requests.exceptions.RequestException as e: # noqa - logger.error(f'exception happened when adding node {node_url}, {e}') - return self.handle_api_timeout(node_url) - self.update_config_file() - - def remove(self, node_url: str): - """Remove a node.""" - if node_url in self.nodes.keys(): - self.nodes.pop(node_url) - self.update_config_file() - self.pd_connection_pool.dereg_instance(node_url) - - def terminate_node(self, node_url: str): - """Terminate a node.""" - success = True - if node_url in self.nodes: - self.nodes.pop(node_url) - headers = {'accept': 'application/json'} - try: - response = requests.get(f'{node_url}/terminate', headers=headers) - if response.status_code != 200: - success = False - logger.error(f'Failed to terminate node {node_url}, ' - f'error_code={response.status_code}, ' - f'error_msg={response.text}') - except Exception as e: # noqa - logger.error(f'exception happened when terminating node {node_url}, {e}') - success = False - else: - logger.error(f'terminating node {node_url} failed since it does not exist. ' - 'May try /nodes/status to check the node list') - success = False - self.update_config_file() - return success - - def terminate_all_nodes(self): - """Terminate all nodes.""" - node_url_li = list(self.nodes.keys()) - all_success = True - for node_url in node_url_li: - if not self.terminate_node(node_url): - all_success = False - return all_success - - def remove_stale_nodes_by_expiration(self): - """Remove stale nodes.""" - to_be_deleted = [] - node_urls = list(self.nodes.keys()) - for node_url in node_urls: - url = f'{node_url}/health' - headers = {'accept': 'application/json'} - try: - response = requests.get(url, headers=headers) - if response.status_code != 200: - to_be_deleted.append(node_url) - except: # noqa - to_be_deleted.append(node_url) - for node_url in to_be_deleted: - self.remove(node_url) - logger.info(f'Removed node_url: {node_url} ' - 'due to heart beat expiration') - - @property - def model_list(self): - """Supported model list.""" - model_names = [] - items = list(self.nodes.items()) - for _, status in items: - model_names.extend(status.models) - return model_names - - @property - def status(self): - """Return the status.""" - return self.nodes - - def get_node_url(self, model_name: str, role: EngineRole = EngineRole.Hybrid): - """Add a node to the manager. - - Args: - model_name (str): A http url. Can be the url generated by - `lmdeploy serve api_server`. - Return: - A node url or None. - """ - - def get_matched_urls(): - urls_with_speeds, speeds, urls_without_speeds = [], [], [] - for node_url, status in self.get_nodes(role).items(): - if model_name in status.models: - if status.speed is not None: - urls_with_speeds.append(node_url) - speeds.append(status.speed) - else: - urls_without_speeds.append(node_url) - all_matched_urls = urls_with_speeds + urls_without_speeds - if len(all_matched_urls) == 0: - return None - # some nodes does not contain speed - # we can set them the average speed value - average_speed = sum(speeds) / len(speeds) if len(speeds) else 1 - all_the_speeds = speeds + [average_speed] * len(urls_without_speeds) - return all_matched_urls, all_the_speeds - - if self.routing_strategy == RoutingStrategy.RANDOM: - all_matched_urls, all_the_speeds = get_matched_urls() - if len(all_matched_urls) == 0: - return None - speed_sum = sum(all_the_speeds) - weights = [speed / speed_sum for speed in all_the_speeds] - index = random.choices(range(len(all_matched_urls)), weights=weights)[0] - url = all_matched_urls[index] - return url - elif self.routing_strategy == RoutingStrategy.MIN_EXPECTED_LATENCY: - all_matched_urls, all_the_speeds = get_matched_urls() - if len(all_matched_urls) == 0: - return None - min_latency = float('inf') - min_index = 0 - # random traverse nodes for low concurrency situation - all_indexes = [i for i in range(len(all_the_speeds))] - random.shuffle(all_indexes) - for index in all_indexes: - latency = self.get_nodes(role)[all_matched_urls[index]].unfinished / all_the_speeds[index] - if min_latency > latency: - min_latency = latency - min_index = index - url = all_matched_urls[min_index] - return url - elif self.routing_strategy == RoutingStrategy.MIN_OBSERVED_LATENCY: - all_matched_urls, latencies = [], [] - for node_url, node_status in self.get_nodes(role).items(): - if model_name in node_status.models: - if len(node_status.latency): - latencies.append(np.mean(np.array(node_status.latency))) - else: - latencies.append(float('inf')) - all_matched_urls.append(node_url) - if len(all_matched_urls) == 0: - return None - index = np.argmin(np.array(latencies)) - return all_matched_urls[index] - else: - raise ValueError(f'Invalid strategy: {self.routing_strategy}') - - async def check_request_model(self, model_name) -> JSONResponse | None: - """Check if a request is valid.""" - if model_name in self.model_list: - return - ret = create_error_response(HTTPStatus.NOT_FOUND, f'The model {model_name!r} does not exist.') - return ret - - def handle_unavailable_model(self, model_name): - """Handle unavailable model. - - Args: - model_name (str): the model in the request. - """ - logger.warning(f'no model name: {model_name}') - ret = { - 'error_code': ErrorCodes.MODEL_NOT_FOUND, - 'text': err_msg[ErrorCodes.MODEL_NOT_FOUND], - } - return json.dumps(ret).encode() + b'\n' - - def handle_api_timeout(self, node_url): - """Handle the api time out.""" - logger.warning(f'api timeout: {node_url}') - ret = { - 'error_code': ErrorCodes.API_TIMEOUT.value, - 'text': err_msg[ErrorCodes.API_TIMEOUT], - } - return json.dumps(ret).encode() + b'\n' - - async def stream_generate(self, request: dict, node_url: str, endpoint: str): - """Return a generator to handle the input request. - - Args: - request (dict): the input request. - node_url (str): the node url. - endpoint (str): the endpoint. Such as `/v1/chat/completions`. - """ - try: - async with aiohttp.ClientSession() as session: - async with session.post(node_url + endpoint, json=request, timeout=self.aiotimeout) as response: - async for line in response.content: - if line.strip(): - yield line + b'\n\n' - except (Exception, GeneratorExit, aiohttp.ClientError) as e: # noqa - logger.error(f'caught an exception: {e}') - # exception happened, reduce unfinished num - yield self.handle_api_timeout(node_url) - - async def generate(self, request: dict, node_url: str, endpoint: str): - """Return a the response of the input request. - - Args: - request (dict): the input request. - node_url (str): the node url. - endpoint (str): the endpoint. Such as `/v1/chat/completions`. - """ - try: - async with aiohttp.ClientSession() as session: - async with session.post(node_url + endpoint, json=request, timeout=self.aiotimeout) as response: - return await response.text() - except (Exception, GeneratorExit, aiohttp.ClientError, asyncio.CancelledError) as e: # noqa # yapf: disable - logger.error(f'caught an exception: {e}') - return self.handle_api_timeout(node_url) - - async def forward_raw_request_stream_generate(self, raw_request: Request, node_url: str, endpoint: str): - try: - target_url = node_url.rstrip('/') + endpoint - headers = self._prepare_headers(raw_request) - body_bytes = await raw_request.body() - async with aiohttp.ClientSession() as session: - async with session.post(target_url, headers=headers, data=body_bytes, - timeout=self.aiotimeout) as response: - if response.status != 200: - error_body = await response.read() - raise APIServerException(status_code=response.status, body=error_body) - async for line in response.content: - if line.strip(): - yield line + b'\n\n' - except APIServerException: - # raise APIServerException again to be caught by the outer layer - raise - except (Exception, GeneratorExit, aiohttp.ClientError) as e: # noqa - logger.error(f'caught an exception: {e}') - # exception happened, reduce unfinished num - yield self.handle_api_timeout(node_url) - - async def forward_raw_request_generate(self, raw_request: Request, node_url: str, endpoint: str): - try: - target_url = node_url.rstrip('/') + endpoint - headers = self._prepare_headers(raw_request) - body_bytes = await raw_request.body() - async with aiohttp.ClientSession() as session: - async with session.post(target_url, headers=headers, data=body_bytes, - timeout=self.aiotimeout) as response: - return await response.text() - except (Exception, GeneratorExit, aiohttp.ClientError, asyncio.CancelledError) as e: # noqa # yapf: disable - logger.error(f'caught an exception: {e}') - return self.handle_api_timeout(node_url) - - def pre_call(self, node_url): - """Preprocess before the request get processed. - - Args: - node_url (str): the node url. - """ - self.nodes[node_url].unfinished += 1 - return time.time() - - def post_call(self, node_url: str, start: int): - """Post process after the response finished. - - Args: - node_url (str): the node url. - start (int): the start time point. time.time() - """ - if node_url in self.nodes: - self.nodes[node_url].unfinished -= 1 - self.nodes[node_url].latency.append(time.time() - start) - - def create_background_tasks(self, url: str, start: int): - """To create a background task. - - Args: - url (str): the node url. - start (int): the start time point. time.time() - """ - background_tasks = BackgroundTasks() - background_tasks.add_task(self.post_call, url, start) - return background_tasks - - def _prepare_headers(self, raw_request: Request) -> dict[str, str]: - headers = dict((name, value) for name, value in raw_request.headers.items() if name.lower() != 'host') - - client_ip = raw_request.client.host if raw_request.client else 'unknown' - headers.update({ - 'X-Forwarded-For': client_ip, - 'X-Forwarded-Host': raw_request.headers.get('host', ''), - 'X-Forwarded-Proto': raw_request.url.scheme, - }) - - return headers - - -app = FastAPI(docs_url='/') -app.add_middleware( - CORSMiddleware, - allow_origins=['*'], - allow_credentials=True, - allow_methods=['*'], - allow_headers=['*'], -) -node_manager = NodeManager() - - -@app.get('/v1/models') -def available_models(): - """Show available models.""" - model_cards = [] - for model_name in node_manager.model_list: - model_cards.append(ModelCard(id=model_name, root=model_name, permission=[ModelPermission()])) - return ModelList(data=model_cards) - - -@app.get('/nodes/status') -def node_status(): - """Show nodes status.""" - try: - return node_manager.status - except: # noqa - return False - - -@app.post('/nodes/add', dependencies=[Depends(validate_json_request)]) -def add_node(node: Node, raw_request: Request = None): - """Add a node to the manager. - - - **url** (str): A http url. Can be the url generated by - `lmdeploy serve api_server`. - - **status** (dict): The description of the node. An example: - ``{"models": ["intern-s2-preview"], "speed": 1}``. The speed here can be - RPM or other metric. All the values of nodes should be the same metric. - """ - try: - res = node_manager.add(node.url, node.status) - if res is not None: - logger.error(f'add node {node.url} failed, {res}') - return res - logger.info(f'add node {node.url} successfully') - return 'Added successfully' - except: # noqa - return 'Failed to add, please check the input url.' - - -@app.post('/nodes/remove', dependencies=[Depends(validate_json_request)]) -def remove_node(node: Node): - """Show available models.""" - try: - node_url = node.url - node_manager.remove(node_url) - logger.info(f'delete node {node_url} successfully') - return 'Deleted successfully' - except: # noqa - logger.error(f'delete node {node.url} failed.') - return 'Failed to delete, please check the input url.' - - -@app.post('/nodes/terminate', dependencies=[Depends(validate_json_request)]) -def terminate_node(node: Node): - """Terminate nodes.""" - try: - node_url = node.url - success = node_manager.terminate_node(node_url) - if not success: - return f'Failed to terminate node {node_url}' - return 'Terminated successfully' - except: # noqa - logger.error(f'Terminate node {node_url} failed.') - return f'Failed to terminate node {node_url}, please check the input url.' - - -@app.get('/nodes/terminate_all', dependencies=[Depends(validate_json_request)]) -def terminate_node_all(): - """Terminate nodes.""" - try: - success = node_manager.terminate_all_nodes() - if not success: - return 'Failed to terminate all nodes' - return 'All nodes terminated successfully' - except: # noqa - logger.error('Failed to terminate all nodes') - return 'Failed to terminate all nodes.' - - -@app.post('/distserve/connection_warmup', dependencies=[Depends(validate_json_request)]) -async def connection_warmup(): - await asyncio.gather(*[ - node_manager.pd_connection_pool.connect( - PDConnectionMessage( - p_url=p_url, - d_url=d_url, - protocol=node_manager.migration_protocol, - rdma_config=node_manager.rdma_config, - )) for p_url in node_manager.prefill_nodes for d_url in node_manager.decode_nodes - ]) - return JSONResponse({'SUCCESS': True}) - - -@app.post('/distserve/gc', dependencies=[Depends(validate_json_request)]) -async def cache_block_gc_to_be_migrated(): - # TODO (JimyMa): add garbage collection of to be migrated request - raise NotImplementedError - - -@app.post('/v1/chat/completions', dependencies=[Depends(validate_json_request)]) -async def chat_completions_v1(request: ChatCompletionRequest, raw_request: Request = None): - """Completion API similar to OpenAI's API. - - Refer to https://platform.openai.com/docs/api-reference/chat/create - for the API specification. - - The request should be a JSON object with the following fields: - - - **model**: model name. Available from /v1/models. - - **messages**: chat history in OpenAI format. Example: - `[{"role": "user", "content": "hi"}]`. - - **temperature** (float): to modulate the next token probability - - **top_p** (float): If set to float < 1, only the smallest set of most - probable tokens with probabilities that add up to top_p or higher - are kept for generation. - - **n** (int): How many chat completion choices to generate for each input - message. Accepts values from 1 to 128, except in DistServe mode. - - **stream**: whether to stream the results or not. Default to false. - - **max_completion_tokens** (int | None): output token nums. Default to None. - - **max_tokens** (int | None): output token nums. Default to None. - Deprecated: Use max_completion_tokens instead. - - **repetition_penalty** (float): The parameter for repetition penalty. - 1.0 means no penalty - - **stop** (str | list[str] | None): To stop generating further - tokens. Only accept stop words that's encoded to one token idex. - - **response_format** (dict | None): To generate response according to given - schema. Examples: - - .. code-block:: json - - { - "type": "json_schema", - "json_schema":{ - "name": "test", - "schema":{ - "properties":{ - "name":{"type":"string"} - }, - "required":["name"], - "type":"object" - } - } - } - - or - ``{"type": "regex_schema", "regex_schema": "call me [A-Za-z]{1,10}"}`` - - **logit_bias** (dict): Bias to logits. Only supported in pytorch engine. - - **tools** (list): A list of tools the model may call. Currently, only - internlm2 functions are supported as a tool. Use this to specify a - list of functions for which the model can generate JSON inputs. - - **tool_choice** (str | object): Controls which (if any) tool is called by - the model. `none` means the model will not call any tool and instead - generates a message. Specifying a particular tool via - ``{"type": "function", "function": {"name": "my_function"}}`` - forces the model to call that tool. `auto` or `required` will put all - the tools information to the model. - - Additional arguments supported by LMDeploy: - - - **top_k** (int): The number of the highest probability vocabulary - tokens to keep for top-k-filtering - - **ignore_eos** (bool): indicator for ignoring eos - - **skip_special_tokens** (bool): Whether or not to remove special tokens - in the decoding. Default to be True. - - **min_new_tokens** (int): To generate at least numbers of tokens. - - **min_p** (float): Minimum token probability, which will be scaled by the - probability of the most likely token. It must be a value between - 0 and 1. Typical values are in the 0.01-0.2 range, comparably - selective as setting `top_p` in the 0.99-0.8 range (use the - opposite of normal `top_p` values) - - Currently we do not support the following features: - - - **presence_penalty** (replaced with repetition_penalty) - - **frequency_penalty** (replaced with repetition_penalty) - """ - check_response = await node_manager.check_request_model(request.model) - if check_response is not None: - return check_response - if (node_manager.serving_strategy == ServingStrategy.DistServe - and request.n is not None and request.n > 1): - return create_error_response( - HTTPStatus.BAD_REQUEST, - 'n > 1 is not supported with the DistServe serving strategy.') - - if node_manager.serving_strategy == ServingStrategy.Hybrid: - node_url = node_manager.get_node_url(request.model) - if not node_url: - return node_manager.handle_unavailable_model(request.model) - - logger.info(f'A request is dispatched to {node_url}') - start = node_manager.pre_call(node_url) - if request.stream is True: - response = node_manager.forward_raw_request_stream_generate(raw_request, node_url, '/v1/chat/completions') - background_task = node_manager.create_background_tasks(node_url, start) - return ProxyStreamingResponse(response, background=background_task, media_type='text/event-stream') - else: - response = await node_manager.forward_raw_request_generate(raw_request, node_url, '/v1/chat/completions') - node_manager.post_call(node_url, start) - return JSONResponse(json.loads(response)) - elif node_manager.serving_strategy == ServingStrategy.DistServe: - request_dict = request.model_dump() - - # Prefill - prefill_request_dict = copy.deepcopy(request_dict) - prefill_request_dict['max_tokens'] = 1 - prefill_request_dict['max_completion_tokens'] = 1 - prefill_request_dict['stream'] = False - prefill_request_dict['with_cache'] = True - prefill_request_dict['preserve_cache'] = True - - prefill_info = {} - p_url = 'dummy:dummy' - if not node_manager.dummy_prefill: - p_url = node_manager.get_node_url(request.model, EngineRole.Prefill) - if not p_url: - return node_manager.handle_unavailable_model(request.model) - logger.info(f'A Prefill request is dispatched to {p_url}') - - start = node_manager.pre_call(p_url) - prefill_info = json.loads(await node_manager.generate(prefill_request_dict, p_url, '/v1/chat/completions')) - node_manager.post_call(p_url, start) - - # # Decode - d_url = node_manager.get_node_url(request.model, EngineRole.Decode) - if not d_url: - return node_manager.handle_unavailable_model(request.model) - logger.info(f'A Decode request is dispatched to {d_url}') - - if not node_manager.dummy_prefill: - if not node_manager.pd_connection_pool.is_connected(p_url, d_url): - await node_manager.pd_connection_pool.connect( - PDConnectionMessage( - p_url=p_url, - d_url=d_url, - protocol=node_manager.migration_protocol, - rdma_config=node_manager.rdma_config, - )) - - remote_session_id = int(prefill_info.get('id')) if prefill_info.get('id') else 0 - remote_block_ids = prefill_info.get('cache_block_ids') or [] - remote_token_id = prefill_info.get('remote_token_ids')[-1] if prefill_info.get('remote_token_ids') else 0 - - request_dict['migration_request'] = MigrationRequest( - protocol=node_manager.migration_protocol, - remote_engine_id=p_url, - remote_session_id=remote_session_id, - remote_block_ids=remote_block_ids, - remote_token_id=remote_token_id, - is_dummy_prefill=node_manager.dummy_prefill).model_dump(mode='json') - - start = node_manager.pre_call(d_url) - if not node_manager.dummy_prefill: - node_manager.pd_connection_pool.shelf_prefill_session((p_url, d_url), prefill_info['id']) - if request.stream is True: - response = node_manager.stream_generate(request_dict, d_url, '/v1/chat/completions') - background_task = node_manager.create_background_tasks(d_url, start) - resp = StreamingResponse(response, background=background_task, media_type='text/event-stream') - else: - response = await node_manager.generate(request_dict, d_url, '/v1/chat/completions') - node_manager.post_call(d_url, start) - resp = JSONResponse(json.loads(response)) - - if not node_manager.dummy_prefill: - node_manager.pd_connection_pool.unshelf_prefill_session((p_url, d_url), prefill_info['id']) - - return resp - - else: - raise ValueError(f'No serving strategy named {node_manager.serving_strategy}') - - -@app.post('/v1/completions', dependencies=[Depends(validate_json_request)]) -async def completions_v1(request: CompletionRequest, raw_request: Request = None): - """Completion API similar to OpenAI's API. - - Go to https://platform.openai.com/docs/api-reference/completions/create - for the API specification. - - The request should be a JSON object with the following fields: - - - **model** (str): model name. Available from /v1/models. - - **prompt** (str): the input prompt. - - **suffix** (str): The suffix that comes after a completion of inserted text. - - **max_completion_tokens** (int | None): output token nums. Default to None. - - **max_tokens** (int): output token nums. Default to 16. - Deprecated: Use max_completion_tokens instead. - - **temperature** (float): to modulate the next token probability - - **top_p** (float): If set to float < 1, only the smallest set of most - probable tokens with probabilities that add up to top_p or higher - are kept for generation. - - **n** (int): How many chat completion choices to generate for each input - message. **Only support one here**. - - **stream**: whether to stream the results or not. Default to false. - - **repetition_penalty** (float): The parameter for repetition penalty. - 1.0 means no penalty - - **user** (str): A unique identifier representing your end-user. - - **stop** (str | list[str] | None): To stop generating further - tokens. Only accept stop words that's encoded to one token idex. - - Additional arguments supported by LMDeploy: - - - **ignore_eos** (bool): indicator for ignoring eos - - **skip_special_tokens** (bool): Whether or not to remove special tokens - in the decoding. Default to be True. - - **top_k** (int): The number of the highest probability vocabulary - tokens to keep for top-k-filtering - - Currently we do not support the following features: - - - **logprobs** (not supported yet) - - **presence_penalty** (replaced with repetition_penalty) - - **frequency_penalty** (replaced with repetition_penalty) - """ - check_response = await node_manager.check_request_model(request.model) - if check_response is not None: - return check_response - if node_manager.serving_strategy == ServingStrategy.Hybrid: - node_url = node_manager.get_node_url(request.model) - if not node_url: - return node_manager.handle_unavailable_model(request.model) - - logger.info(f'A request is dispatched to {node_url}') - start = node_manager.pre_call(node_url) - if request.stream is True: - response = node_manager.forward_raw_request_stream_generate(raw_request, node_url, '/v1/completions') - background_task = node_manager.create_background_tasks(node_url, start) - return ProxyStreamingResponse(response, background=background_task, media_type='text/event-stream') - else: - response = await node_manager.forward_raw_request_generate(raw_request, node_url, '/v1/completions') - node_manager.post_call(node_url, start) - return JSONResponse(json.loads(response)) - elif node_manager.serving_strategy == ServingStrategy.DistServe: - request_dict = request.model_dump() - - # Prefill - prefill_request_dict = copy.deepcopy(request_dict) - prefill_request_dict['max_tokens'] = 1 - prefill_request_dict['stream'] = False - prefill_request_dict['with_cache'] = True - prefill_request_dict['preserve_cache'] = True - - if not node_manager.dummy_prefill: - try: - p_url = node_manager.get_node_url(request.model, EngineRole.Prefill) - except Exception as e: - logger.error(f'error Msg: {str(e)}') - return {'status': 'Instance sch error, cannot find available p_url'} - - if not p_url: - return node_manager.handle_unavailable_model(request.model) - logger.info(f'A Prefill request is dispatched to {p_url}') - - start = node_manager.pre_call(p_url) - prefill_info = json.loads(await node_manager.generate(prefill_request_dict, p_url, '/v1/completions')) - node_manager.post_call(p_url, start) - else: - p_url = 'dummy:dummy' - prefill_info = {} - - # Decode - try: - d_url = node_manager.get_node_url(request.model, EngineRole.Decode) - except Exception as e: - logger.error(f'error Msg: {str(e)}') - return {'status': 'Instance sch error, cannot find available p_url'} - - if not d_url: - return node_manager.handle_unavailable_model(request.model) - logger.info(f'A Decode request is dispatched to {d_url}') - - if not node_manager.dummy_prefill: - if not node_manager.pd_connection_pool.is_connected(p_url, d_url): - try: - await node_manager.pd_connection_pool.connect( - PDConnectionMessage( - p_url=p_url, - d_url=d_url, - protocol=node_manager.migration_protocol, - rdma_config=node_manager.rdma_config, - )) - except Exception as e: - logger.error(f'error Msg: {str(e)}') - return {'status': f'Connection error, cannot establish connection {(p_url, d_url)}'} - node_manager.pd_connection_pool.shelf_prefill_session((p_url, d_url), prefill_info['id']) - - remote_session_id = int(prefill_info.get('id')) if prefill_info.get('id') else 0 - remote_block_ids = prefill_info.get('cache_block_ids') or [] - remote_token_id = prefill_info.get('remote_token_ids')[-1] if prefill_info.get('remote_token_ids') else 0 - request_dict['migration_request'] = MigrationRequest( - protocol=node_manager.migration_protocol, - remote_engine_id=p_url, - remote_session_id=remote_session_id, - remote_block_ids=remote_block_ids, - remote_token_id=remote_token_id, - is_dummy_prefill=node_manager.dummy_prefill).model_dump(mode='json') - - start = node_manager.pre_call(d_url) - if not node_manager.dummy_prefill: - node_manager.pd_connection_pool.shelf_prefill_session((p_url, d_url), prefill_info['id']) - if request.stream is True: - response = node_manager.stream_generate(request_dict, d_url, '/v1/completions') - background_task = node_manager.create_background_tasks(d_url, start) - resp = StreamingResponse(response, background=background_task, media_type='text/event-stream') - else: - response = await node_manager.generate(request_dict, d_url, '/v1/completions') - node_manager.post_call(d_url, start) - resp = JSONResponse(json.loads(response)) - if not node_manager.dummy_prefill: - node_manager.pd_connection_pool.unshelf_prefill_session((p_url, d_url), prefill_info.get('id')) - return resp - else: - raise ValueError(f'No serving strategy named {node_manager.serving_strategy}') - - -def proxy(server_name: str = '0.0.0.0', - server_port: int = 8000, - serving_strategy: Literal['Hybrid', 'DistServe'] = 'Hybrid', - routing_strategy: Literal['random', 'min_expected_latency', 'min_observed_latency'] = 'min_expected_latency', - api_keys: list[str] | str | None = None, - ssl: bool = False, - log_level: str = 'INFO', - disable_cache_status: bool = False, - link_type: Literal['RoCE', 'IB'] = 'RoCE', - migration_protocol: Literal['RDMA'] = 'RDMA', - dummy_prefill: bool = False, - **kwargs): - """To launch the proxy server. - - Args: - server_name (str): the server name of the proxy. Default to '0.0.0.0'. - server_port (str): the server port. Default to 8000. - serving_strategy ('Hybrid' | 'DistServe'): the strategy to serving. Hybrid default. - DistServe for PD Disaggregation. - route_strategy ('random' | 'min_expected_latency' | 'min_observed_latency'): - the strategy to dispatch requests to nodes. Default to - 'min_expected_latency' - api_keys (list[str] | str | None): Optional list of API keys. Accepts string type as - a single api_key. Default to None, which means no api key applied. - ssl (bool): Enable SSL. Requires OS Environment variables 'SSL_KEYFILE' and 'SSL_CERTFILE'. - log_level (str): Set the log level. Default to INFO. - disable_cache_status (str): Whether to cache the proxy status to - proxy_config.yml. - migration_protocol: migration protocol when PD disaggregation. RDMA default. - """ # noqa - node_manager.serving_strategy = ServingStrategy[serving_strategy] - node_manager.routing_strategy = RoutingStrategy.from_str(routing_strategy) - node_manager.migration_protocol = MigrationProtocol[migration_protocol] - node_manager.dummy_prefill = dummy_prefill - - node_manager.rdma_config = DistServeRDMAConfig( - link_type=RDMALinkType[link_type], - with_gdr=True, - ) - node_manager.cache_status = not disable_cache_status - if api_keys is not None and (tokens := [key for key in api_keys if key]): - from lmdeploy.serve.utils.server_utils import AuthenticationMiddleware - - app.add_middleware(AuthenticationMiddleware, tokens=tokens) - - ssl_keyfile, ssl_certfile = None, None - if ssl: - ssl_keyfile = os.environ['SSL_KEYFILE'] - ssl_certfile = os.environ['SSL_CERTFILE'] - - logger.setLevel(log_level) - uvicorn_log_level = os.getenv('UVICORN_LOG_LEVEL', 'info').lower() - uvicorn.run(app=app, - host=server_name, - port=server_port, - log_level=uvicorn_log_level, - ssl_keyfile=ssl_keyfile, - ssl_certfile=ssl_certfile) - - -if __name__ == '__main__': - import fire - - fire.Fire(proxy) diff --git a/lmdeploy/serve/proxy/streaming_response.py b/lmdeploy/serve/proxy/streaming_response.py deleted file mode 100644 index 5f0965eef6..0000000000 --- a/lmdeploy/serve/proxy/streaming_response.py +++ /dev/null @@ -1,71 +0,0 @@ -# Copyright (c) OpenMMLab. All rights reserved. - -import json - -from fastapi.responses import StreamingResponse - -from .utils import APIServerException - - -class ProxyStreamingResponse(StreamingResponse): - """StreamingResponse that can handle exceptions thrown by the generator.""" - - def __init__(self, content, **kwargs): - super().__init__(content, **kwargs) - - async def stream_response(self, send) -> None: - iterator = self.body_iterator.__aiter__() - try: - # get the first chunk(stream_generate's first yield) - first_chunk = await iterator.__anext__() - - except APIServerException as e: - headers = self._convert_headers_to_asgi(e.headers) if e.headers else self.raw_headers - await send({'type': 'http.response.start', 'status': e.status_code, 'headers': headers}) - await send({ - 'type': 'http.response.body', - 'body': e.body, - 'more_body': False, - }) - return - - # normal case, send the header first - await send({ - 'type': 'http.response.start', - 'status': self.status_code, - 'headers': self.raw_headers, - }) - - # send body with the first chunk - await send({ - 'type': 'http.response.body', - 'body': first_chunk, - 'more_body': True, - }) - - # continue streaming output - try: - async for chunk in iterator: - await send({ - 'type': 'http.response.body', - 'body': chunk, - 'more_body': True, - }) - except Exception: - error_data = {'error': True, 'status': 500, 'message': 'Internal streaming error'} - await send({ - 'type': 'http.response.body', - 'body': json.dumps(error_data).encode('utf-8'), - 'more_body': False, - }) - return - - await send({ - 'type': 'http.response.body', - 'body': b'', - 'more_body': False, - }) - - def _convert_headers_to_asgi(self, headers: dict) -> list[tuple[bytes, bytes]]: - """Convert dict headers to ASGI raw header tuples.""" - return [(name.lower().encode('latin-1'), str(value).encode('latin-1')) for name, value in headers.items()] diff --git a/lmdeploy/serve/proxy/utils.py b/lmdeploy/serve/proxy/utils.py deleted file mode 100644 index 09344cd2f0..0000000000 --- a/lmdeploy/serve/proxy/utils.py +++ /dev/null @@ -1,59 +0,0 @@ -# Copyright (c) OpenMMLab. All rights reserved. - -import enum -import os - -from lmdeploy.utils import get_logger - -logger = get_logger('lmdeploy') - -LATENCY_DEQUE_LEN = 15 -AIOHTTP_TIMEOUT = os.getenv('AIOHTTP_TIMEOUT', None) -if AIOHTTP_TIMEOUT is not None: - AIOHTTP_TIMEOUT = int(AIOHTTP_TIMEOUT) -logger.info(f'AIOHTTP_TIMEOUT set to {AIOHTTP_TIMEOUT}. It can be modified before launching the proxy server ' - 'through env variable AIOHTTP_TIMEOUT') - - -class RoutingStrategy(enum.Enum): - """Strategy to dispatch requests to nodes.""" - RANDOM = enum.auto() - MIN_EXPECTED_LATENCY = enum.auto() - MIN_OBSERVED_LATENCY = enum.auto() - - @classmethod - def from_str(cls, name): - """Get strategy from string.""" - if name == 'random': - return cls.RANDOM - elif name == 'min_expected_latency': - return cls.MIN_EXPECTED_LATENCY - elif name == 'min_observed_latency': - return cls.MIN_OBSERVED_LATENCY - else: - raise ValueError(f'Invalid strategy: {name}. Supported: random, ' - f'min_expected_latency, min_observed_latency.') - - -class ErrorCodes(enum.Enum): - """Error codes.""" - MODEL_NOT_FOUND = 10400 - SERVICE_UNAVAILABLE = 10401 - API_TIMEOUT = 10402 - - -err_msg = { - ErrorCodes.MODEL_NOT_FOUND: 'The request model name does not exist in the model list.', - ErrorCodes.SERVICE_UNAVAILABLE: 'The service is unavailable now. May retry later.', - ErrorCodes.API_TIMEOUT: 'Failed to get response after a period of time' -} - - -class APIServerException(Exception): - - def __init__(self, status_code: int, body: bytes, headers: dict | None = None): - self.status_code = status_code - self.body = body - self.headers = headers or {} - if 'content-type' not in self.headers: - self.headers['content-type'] = 'application/json' diff --git a/tests/test_lmdeploy/serve/core/test_chat_runner.py b/tests/test_lmdeploy/serve/core/test_chat_runner.py index 57b15a9c01..8b9f4ad724 100644 --- a/tests/test_lmdeploy/serve/core/test_chat_runner.py +++ b/tests/test_lmdeploy/serve/core/test_chat_runner.py @@ -227,11 +227,27 @@ async def _run(): assert chunks[0].token_ids == [1] assert chunks[0].logprobs == [{1: -0.1}] assert chunks[0].routed_experts == [[1]] + assert chunks[0].cache_block_ids == ['cache'] + assert chunks[0].remote_session_id == -1 assert chunks[0].reasoning_tokens == 2 assert context.async_engine.generator_closed is True assert context.session_manager.removed == [context.session_manager.session] +def test_runner_collect_preserves_remote_session_id(): + context = _FakeServerContext(_Parser) + + async def _run(): + chat_runner = await ChatRunner.prepare(context, _request()) + return await chat_runner.collect() + + result = asyncio.run(_run()) + + assert result.cache_block_ids == [['cache']] + assert result.remote_token_ids == [[1]] + assert result.remote_session_id == -1 + + def test_runner_close_cleans_prepared_unconsumed_request(): context = _FakeServerContext(_Parser) diff --git a/tests/test_lmdeploy/serve/openai/chat_completions/test_n_completions.py b/tests/test_lmdeploy/serve/openai/chat_completions/test_n_completions.py index ae22efaa1a..49c02c8fa4 100644 --- a/tests/test_lmdeploy/serve/openai/chat_completions/test_n_completions.py +++ b/tests/test_lmdeploy/serve/openai/chat_completions/test_n_completions.py @@ -9,6 +9,7 @@ import pytest from fastapi.responses import JSONResponse +from lmdeploy.pytorch.disagg.config import ServingStrategy from lmdeploy.serve.openai.chat_completions.fanout import _batch_stream_payloads from lmdeploy.serve.openai.protocol import ChatCompletionRequest @@ -194,21 +195,16 @@ def test_handler_rejects_cache_migration_for_multiple_choices( assert context.async_engine.call_count == 0 -def test_distserve_proxy_rejects_multiple_choices(monkeypatch): - from lmdeploy.pytorch.disagg.config import ServingStrategy - from lmdeploy.serve.proxy import proxy - - async def model_exists(model): - return None +def test_distserve_validation_rejects_multiple_choices(chat_endpoint, fake_raw_request): + endpoint, context = chat_endpoint + request = _request(n=2) + fake_raw_request._payload = {'with_cache': True} - monkeypatch.setattr(proxy.node_manager, 'check_request_model', - model_exists) - monkeypatch.setattr(proxy.node_manager, 'serving_strategy', - ServingStrategy.DistServe) - response = asyncio.run(proxy.chat_completions_v1(_request(n=2))) + context.serving_strategy = ServingStrategy.DistServe + response = asyncio.run(endpoint(request, fake_raw_request)) assert isinstance(response, JSONResponse) assert response.status_code == 400 - assert 'DistServe' in response.body.decode() + assert 'cache migration' in response.body.decode() def test_prompt_cache_usage_is_counted_once(chat_endpoint, fake_raw_request): diff --git a/tests/test_lmdeploy/serve/test_cli_serve_proxy.py b/tests/test_lmdeploy/serve/test_cli_serve_proxy.py new file mode 100644 index 0000000000..febaca7b23 --- /dev/null +++ b/tests/test_lmdeploy/serve/test_cli_serve_proxy.py @@ -0,0 +1,71 @@ +import argparse + +import pytest + +from lmdeploy.cli.serve import SubCliServe + + +def proxy_args(**overrides): + args = { + 'server_name': 'router.local', + 'server_port': 30000, + 'routing_strategy': 'cache_aware', + 'serving_strategy': 'Hybrid', + 'dummy_prefill': False, + 'disable_cache_status': False, + 'migration_protocol': 'RDMA', + 'link_type': 'RoCE', + 'disable_gdr': False, + 'api_keys': None, + 'ssl': False, + 'log_level': 'INFO', + } + args.update(overrides) + return argparse.Namespace(**args) + + +@pytest.mark.parametrize( + 'serving_strategy,expected_pd_disaggregation', + [('Hybrid', False), ('DistServe', True)]) +def test_maps_serving_strategy_to_router(serving_strategy, expected_pd_disaggregation): + router_args = SubCliServe.build_router_args(proxy_args(serving_strategy=serving_strategy)) + + assert router_args.lmdeploy_pd_disaggregation is expected_pd_disaggregation + + +@pytest.mark.parametrize('legacy_strategy', ['min_expected_latency', 'min_observed_latency']) +def test_maps_legacy_routing_strategy_to_cache_aware(legacy_strategy): + router_args = SubCliServe.build_router_args(proxy_args(routing_strategy=legacy_strategy)) + + assert router_args.policy == 'cache_aware' + + +def test_maps_router_arguments(): + router_args = SubCliServe.build_router_args( + proxy_args( + dummy_prefill=True, + migration_protocol='NVLINK', + link_type='IB', + disable_gdr=True, + api_keys=['secret'], + log_level='WARNING', + )) + + assert router_args.host == 'router.local' + assert router_args.port == 30000 + assert router_args.lmdeploy_migration_protocol == 'nvlink' + assert router_args.lmdeploy_rdma_link_type == 'ib' + assert router_args.lmdeploy_dummy_prefill is True + assert router_args.lmdeploy_disable_gdr is True + assert router_args.api_key == 'secret' + assert router_args.log_level == 'warning' + + +def test_rejects_multiple_api_keys(): + with pytest.raises(ValueError, match='one API key'): + SubCliServe.build_router_args(proxy_args(api_keys=['secret-one', 'secret-two'])) + + +def test_rejects_ssl(): + with pytest.raises(ValueError, match='does not currently support TLS'): + SubCliServe.build_router_args(proxy_args(ssl=True))