From e1e61914c7bc41ca8eb54623b40898713053b9fa Mon Sep 17 00:00:00 2001 From: Yingge He Date: Tue, 28 Jul 2026 17:32:37 -0700 Subject: [PATCH 1/4] test: Add HSTU (Generative Recommenders) Torch AOTI CI test Add server/qa/L0_torch_aoti_hstu with export.sh (AOTI export + C++ KV-cache reference dump) and test.sh (assemble torch_aoti model repo, start FlexKV, launch tritonserver, run the HSTU client). The in-container command sequences mirror the DevTech reference test (Devtech-Compute/distributed-recommender: ci/tritonserver_test.sh). Docker orchestration of the two recsys-examples images lives in the Triton CI job, so these scripts stay Docker-free. Refs: TRI-1615 Signed-off-by: Yingge He --- qa/L0_torch_aoti_hstu/export.sh | 42 +++++++++++++++++++++++ qa/L0_torch_aoti_hstu/test.sh | 59 +++++++++++++++++++++++++++++++++ 2 files changed, 101 insertions(+) create mode 100755 qa/L0_torch_aoti_hstu/export.sh create mode 100755 qa/L0_torch_aoti_hstu/test.sh diff --git a/qa/L0_torch_aoti_hstu/export.sh b/qa/L0_torch_aoti_hstu/export.sh new file mode 100755 index 0000000000..ae8e8e9201 --- /dev/null +++ b/qa/L0_torch_aoti_hstu/export.sh @@ -0,0 +1,42 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# HSTU Torch AOTI test, step 1 (runs inside the recsys-examples inference +# image). Exports the AOTI model from the KuaiRand-1K checkpoint and runs the +# C++ KV-cache inference to produce a reference dump, then copies both to +# /exported_hstu_model for step 2 (test.sh). The CI job launches this via +# `docker run`; this script itself does not use Docker. + +set -e + +RECSYS_DIR=${RECSYS_DIR:="/workspace/recsys-examples/examples/hstu"} +HSTU_CKPT_NAME=${HSTU_CKPT_NAME:="fused_kuairand_1k_ckpt"} +EXPORTED_DIR=${EXPORTED_DIR:="/exported_hstu_model"} + +export FLEXKV_LOG_LEVEL=WARNING +export DYNAMICEMB_OPS_LIB_DIR=/workspace/recsys-examples/corelib/dynamicemb/torch_binding_build/ +export PYTHONPATH=${PYTHONPATH}:/workspace/recsys-examples/examples/ + +cd "${RECSYS_DIR}" +export KVCACHE_MANAGER_CONFIG_FILE=./inference_aoti/kvcache_cpp_runtime.yaml + +python3 ./inference_aoti/export_inference_gr_ranking_kvcache.py \ + --gin_config_file ./inference/configs/kuairand_1k_inference_ranking.gin \ + --checkpoint_dir "${RECSYS_DIR}/ckpt/${HSTU_CKPT_NAME}" \ + --max_bs 2 --kvcache_config_file "${KVCACHE_MANAGER_CONFIG_FILE}" + +python3 ./inference_aoti/start_flexkv_server_for_kvcache_cpp.py \ + --config_file "${KVCACHE_MANAGER_CONFIG_FILE}" > cpp_kvcache_server.log 2>&1 & +kvserver_pid=$! +sleep 10 +kill -0 ${kvserver_pid} + +./inference_aoti/cpp_inference/build/inference_hstu_gr_ranking_kvcache_exported_model \ + ./inference_aoti/hstu_gr_ranking_kvcache_model \ + ./inference_aoti/export_test_dump +kill ${kvserver_pid} || true + +mkdir -p "${EXPORTED_DIR}" +cp -apr "${RECSYS_DIR}/inference_aoti/hstu_gr_ranking_kvcache_model" "${EXPORTED_DIR}/" +cp -apr "${RECSYS_DIR}/inference_aoti/export_test_dump" "${EXPORTED_DIR}/" diff --git a/qa/L0_torch_aoti_hstu/test.sh b/qa/L0_torch_aoti_hstu/test.sh new file mode 100755 index 0000000000..780b9ef805 --- /dev/null +++ b/qa/L0_torch_aoti_hstu/test.sh @@ -0,0 +1,59 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# HSTU (Generative Recommenders) Torch AOTI test, step 2 (runs inside the +# recsys-examples tritonserver image). Assembles the `platform: "torch_aoti"` +# model repository from the model exported in step 1 (export.sh), starts the +# FlexKV KV-cache server, launches tritonserver, and runs the HSTU client. The +# CI job launches this via `docker run`; this script itself does not use Docker. +# +# The overall flow mirrors the DevTech reference test +# (Devtech-Compute/distributed-recommender: ci/tritonserver_test.sh). + +set -e + +RECSYS_DIR=${RECSYS_DIR:="/workspace/recsys-examples/examples/hstu"} +EXPORTED_DIR=${EXPORTED_DIR:="/exported_hstu_model"} +MODEL_REPO=${MODEL_REPO:="/triton_model_repo"} + +COLOR_ERROR="\033[31m" +COLOR_RESET="\033[0m" +COLOR_SUCCESS="\033[32m" +RET=0 + +# Assemble the torch_aoti model repository: config from the recsys tree, the +# exported AOTI package as model version 1. +mkdir -p "${MODEL_REPO}" +cp -apr "${RECSYS_DIR}/inference_aoti/triton_aoti/hstu_gr_ranking_kvcache" "${MODEL_REPO}/" +cp -apr "${EXPORTED_DIR}/hstu_gr_ranking_kvcache_model" "${MODEL_REPO}/hstu_gr_ranking_kvcache/1" +cp -apr "${EXPORTED_DIR}/export_test_dump" "${RECSYS_DIR}/inference_aoti" + +cd "${RECSYS_DIR}/inference_aoti" +export FLEXKV_LOG_LEVEL=WARNING +export KVCACHE_MANAGER_CONFIG_FILE=${PWD}/kvcache_cpp_runtime.yaml + +python3 start_flexkv_server_for_kvcache_cpp.py \ + --config_file "${KVCACHE_MANAGER_CONFIG_FILE}" 2>&1 & +kvserver_pid=$! +sleep 10 +kill -0 ${kvserver_pid} + +tritonserver --model-repository="${MODEL_REPO}/" & +triton_pid=$! +sleep 30 + +python3 test_tritonserver_aoti_hstu_model.py > test_benchmark.log || RET=1 +cat test_benchmark.log + +kill ${triton_pid} || true +kill -9 ${triton_pid} || true +kill ${kvserver_pid} || true + +if [[ ${RET} -eq 0 ]]; then + echo -e "${COLOR_SUCCESS}\n***\n*** HSTU Torch AOTI Test Passed\n***${COLOR_RESET}" +else + echo -e "${COLOR_ERROR}\n***\n*** HSTU Torch AOTI Test FAILED\n***${COLOR_RESET}" +fi + +exit ${RET} From 06f2e2dacd357ab02e7d71113fb2346aa62d2ff0 Mon Sep 17 00:00:00 2001 From: Yingge He Date: Fri, 31 Jul 2026 01:05:02 -0700 Subject: [PATCH 2/4] test: Run HSTU Torch AOTI test in a single container Fold the export phase into test.sh so the whole flow (AOTI export, reference dump, serving, client) runs in one container like the other L0 tests, instead of orchestrating two images from the runner host. Serve with the tritonserver tree at TRITON_DIR so CI can point the test at the build under test rather than the one bundled in the test image, and clear errexit around the util.sh helpers so failures still reach teardown. Signed-off-by: Yingge He --- qa/L0_torch_aoti_hstu/export.sh | 42 ------ qa/L0_torch_aoti_hstu/test.sh | 234 +++++++++++++++++++++++++++----- 2 files changed, 198 insertions(+), 78 deletions(-) delete mode 100755 qa/L0_torch_aoti_hstu/export.sh diff --git a/qa/L0_torch_aoti_hstu/export.sh b/qa/L0_torch_aoti_hstu/export.sh deleted file mode 100755 index ae8e8e9201..0000000000 --- a/qa/L0_torch_aoti_hstu/export.sh +++ /dev/null @@ -1,42 +0,0 @@ -#!/bin/bash -# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: BSD-3-Clause - -# HSTU Torch AOTI test, step 1 (runs inside the recsys-examples inference -# image). Exports the AOTI model from the KuaiRand-1K checkpoint and runs the -# C++ KV-cache inference to produce a reference dump, then copies both to -# /exported_hstu_model for step 2 (test.sh). The CI job launches this via -# `docker run`; this script itself does not use Docker. - -set -e - -RECSYS_DIR=${RECSYS_DIR:="/workspace/recsys-examples/examples/hstu"} -HSTU_CKPT_NAME=${HSTU_CKPT_NAME:="fused_kuairand_1k_ckpt"} -EXPORTED_DIR=${EXPORTED_DIR:="/exported_hstu_model"} - -export FLEXKV_LOG_LEVEL=WARNING -export DYNAMICEMB_OPS_LIB_DIR=/workspace/recsys-examples/corelib/dynamicemb/torch_binding_build/ -export PYTHONPATH=${PYTHONPATH}:/workspace/recsys-examples/examples/ - -cd "${RECSYS_DIR}" -export KVCACHE_MANAGER_CONFIG_FILE=./inference_aoti/kvcache_cpp_runtime.yaml - -python3 ./inference_aoti/export_inference_gr_ranking_kvcache.py \ - --gin_config_file ./inference/configs/kuairand_1k_inference_ranking.gin \ - --checkpoint_dir "${RECSYS_DIR}/ckpt/${HSTU_CKPT_NAME}" \ - --max_bs 2 --kvcache_config_file "${KVCACHE_MANAGER_CONFIG_FILE}" - -python3 ./inference_aoti/start_flexkv_server_for_kvcache_cpp.py \ - --config_file "${KVCACHE_MANAGER_CONFIG_FILE}" > cpp_kvcache_server.log 2>&1 & -kvserver_pid=$! -sleep 10 -kill -0 ${kvserver_pid} - -./inference_aoti/cpp_inference/build/inference_hstu_gr_ranking_kvcache_exported_model \ - ./inference_aoti/hstu_gr_ranking_kvcache_model \ - ./inference_aoti/export_test_dump -kill ${kvserver_pid} || true - -mkdir -p "${EXPORTED_DIR}" -cp -apr "${RECSYS_DIR}/inference_aoti/hstu_gr_ranking_kvcache_model" "${EXPORTED_DIR}/" -cp -apr "${RECSYS_DIR}/inference_aoti/export_test_dump" "${EXPORTED_DIR}/" diff --git a/qa/L0_torch_aoti_hstu/test.sh b/qa/L0_torch_aoti_hstu/test.sh index 780b9ef805..3e5147d162 100755 --- a/qa/L0_torch_aoti_hstu/test.sh +++ b/qa/L0_torch_aoti_hstu/test.sh @@ -2,58 +2,220 @@ # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: BSD-3-Clause -# HSTU (Generative Recommenders) Torch AOTI test, step 2 (runs inside the -# recsys-examples tritonserver image). Assembles the `platform: "torch_aoti"` -# model repository from the model exported in step 1 (export.sh), starts the -# FlexKV KV-cache server, launches tritonserver, and runs the HSTU client. The -# CI job launches this via `docker run`; this script itself does not use Docker. +# End-to-end HSTU (Generative Recommenders) Torch AOTI test. # -# The overall flow mirrors the DevTech reference test -# (Devtech-Compute/distributed-recommender: ci/tritonserver_test.sh). +# Exports the HSTU ranking model to an AOTI package from the KuaiRand-1K +# checkpoint, produces a reference output dump with the C++ KV-cache runtime, +# then serves the package with `platform: "torch_aoti"` and runs the HSTU +# client against it. The command sequence mirrors the DevTech reference test +# (Devtech-Compute/distributed-recommender: ci/tritonserver_test.sh), which +# splits these phases across two container images. +# +# The test image must provide the recsys-examples stack (export tooling, +# dynamicemb, the C++ KV-cache runtime, FlexKV, the torch_aoti model +# configuration, and the HSTU client), with the HSTU dataset and checkpoint +# mounted into RECSYS_DIR. tritonserver itself comes from TRITON_DIR, which CI +# points at a build from the pipeline under test. -set -e +source ../common/util.sh -RECSYS_DIR=${RECSYS_DIR:="/workspace/recsys-examples/examples/hstu"} -EXPORTED_DIR=${EXPORTED_DIR:="/exported_hstu_model"} -MODEL_REPO=${MODEL_REPO:="/triton_model_repo"} +if [[ "${DEBUG}" == "true" ]]; then + set -x +else + set +x +fi +# The CI harness runs this with `bash -ex`. Errexit is disabled because the test +# checks exit codes itself and has to reach the FlexKV and tritonserver teardown +# on failure. util.sh helpers re-enable it, so it is cleared again after each one. +set +e + +COLOR_DARK="\033[90m" COLOR_ERROR="\033[31m" +COLOR_INFO="\033[94m" COLOR_RESET="\033[0m" COLOR_SUCCESS="\033[32m" RET=0 -# Assemble the torch_aoti model repository: config from the recsys tree, the -# exported AOTI package as model version 1. -mkdir -p "${MODEL_REPO}" -cp -apr "${RECSYS_DIR}/inference_aoti/triton_aoti/hstu_gr_ranking_kvcache" "${MODEL_REPO}/" -cp -apr "${EXPORTED_DIR}/hstu_gr_ranking_kvcache_model" "${MODEL_REPO}/hstu_gr_ranking_kvcache/1" -cp -apr "${EXPORTED_DIR}/export_test_dump" "${RECSYS_DIR}/inference_aoti" +export CUDA_VISIBLE_DEVICES=0 + +TESTDIR=`pwd` + +# TRITON_DIR selects which Triton is exercised. The recsys image ships a released +# tritonserver with a pinned PyTorch backend, so CI points TRITON_DIR at a mounted +# tree built by this pipeline in order to catch regressions in Triton itself. +TRITON_DIR=${TRITON_DIR:="/opt/tritonserver"} +SERVER=${TRITON_DIR}/bin/tritonserver +BACKEND_DIR=${BACKEND_DIR:=${TRITON_DIR}/backends} +SERVER_TIMEOUT=${SERVER_TIMEOUT:=300} +export LD_LIBRARY_PATH=${TRITON_DIR}/lib:${LD_LIBRARY_PATH} + +# recsys-examples HSTU tree shipped in the test image. inference_aoti holds the +# export script, the C++ KV-cache runtime, the FlexKV launcher, the torch_aoti +# model configuration, and the client. +RECSYS_DIR=${RECSYS_DIR:="/workspace/recsys-examples/examples/hstu"} +AOTI_DIR=${AOTI_DIR:=${RECSYS_DIR}/inference_aoti} + +MODEL_NAME=${MODEL_NAME:="hstu_gr_ranking_kvcache"} +HSTU_CKPT_NAME=${HSTU_CKPT_NAME:="fused_kuairand_1k_ckpt"} +HSTU_CKPT_DIR=${HSTU_CKPT_DIR:=${RECSYS_DIR}/ckpt/${HSTU_CKPT_NAME}} +GIN_CONFIG=${GIN_CONFIG:=${RECSYS_DIR}/inference/configs/kuairand_1k_inference_ranking.gin} +MAX_BS=${MAX_BS:=2} + +# Everything the test produces is kept under the test directory so CI collects it +# and reruns start from a clean slate. +MODELDIR=${MODELDIR:=${TESTDIR}/models} +EXPORTED_MODEL=${EXPORTED_MODEL:=${TESTDIR}/${MODEL_NAME}_model} +DUMP_DIR=${DUMP_DIR:=${TESTDIR}/export_test_dump} + +CLIENT_LOG="${TESTDIR}/${MODEL_NAME}-client.log" +SERVER_LOG="${TESTDIR}/${MODEL_NAME}-server.log" +EXPORT_LOG="${TESTDIR}/${MODEL_NAME}-export.log" +KVCACHE_LOG="${TESTDIR}/${MODEL_NAME}-kvcache.log" + +BACKENDS=${BACKENDS:="pytorch"} +export BACKENDS + +# dynamicemb ops lib and the recsys examples package are needed by the export +# tooling and the client. +export FLEXKV_LOG_LEVEL=${FLEXKV_LOG_LEVEL:="WARNING"} +export DYNAMICEMB_OPS_LIB_DIR=${DYNAMICEMB_OPS_LIB_DIR:="/workspace/recsys-examples/corelib/dynamicemb/torch_binding_build/"} +export PYTHONPATH=${PYTHONPATH}:/workspace/recsys-examples/examples/ +export KVCACHE_MANAGER_CONFIG_FILE=${KVCACHE_MANAGER_CONFIG_FILE:=${AOTI_DIR}/kvcache_cpp_runtime.yaml} + +KVCACHE_PID=0 -cd "${RECSYS_DIR}/inference_aoti" -export FLEXKV_LOG_LEVEL=WARNING -export KVCACHE_MANAGER_CONFIG_FILE=${PWD}/kvcache_cpp_runtime.yaml +# The KV-cache runtime and the torch_aoti model both talk to a FlexKV server. +# It is restarted between the reference run and serving so Triton sees a clean +# cache, matching the reference test's two-phase flow. +function start_kvcache_server () { + KVCACHE_PID=0 + python3 ${AOTI_DIR}/start_flexkv_server_for_kvcache_cpp.py \ + --config_file ${KVCACHE_MANAGER_CONFIG_FILE} >> ${KVCACHE_LOG} 2>&1 & + local pid=$! + sleep 10 + if ! kill -0 ${pid} > /dev/null 2>&1; then + echo -e "${COLOR_ERROR}\n***\n*** Failed to start FlexKV KV-cache server\n***${COLOR_RESET}" 1>&2 + cat ${KVCACHE_LOG} 1>&2 + return 1 + fi + KVCACHE_PID=${pid} + echo -e "${COLOR_DARK}FlexKV KV-cache server running (pid: ${KVCACHE_PID})${COLOR_RESET}" +} -python3 start_flexkv_server_for_kvcache_cpp.py \ - --config_file "${KVCACHE_MANAGER_CONFIG_FILE}" 2>&1 & -kvserver_pid=$! -sleep 10 -kill -0 ${kvserver_pid} +function stop_kvcache_server () { + if [[ "${KVCACHE_PID}" -ne 0 ]]; then + echo -e "${COLOR_DARK}Killing FlexKV KV-cache server (pid: ${KVCACHE_PID})${COLOR_RESET}" + kill ${KVCACHE_PID} > /dev/null 2>&1 || true + wait ${KVCACHE_PID} > /dev/null 2>&1 || true + KVCACHE_PID=0 + fi +} -tritonserver --model-repository="${MODEL_REPO}/" & -triton_pid=$! -sleep 30 +rm -rf ${MODELDIR} ${EXPORTED_MODEL} ${DUMP_DIR} -python3 test_tritonserver_aoti_hstu_model.py > test_benchmark.log || RET=1 -cat test_benchmark.log +# The export tooling and the client resolve dataset and config paths relative to +# the recsys tree. +cd ${RECSYS_DIR} + +# Export the AOTI package from the ranking checkpoint. The export script writes +# it under AOTI_DIR; it is relocated to the test directory below. +echo -e "${COLOR_DARK}Exporting ${MODEL_NAME} from ${HSTU_CKPT_DIR}${COLOR_RESET}" +python3 ${AOTI_DIR}/export_inference_gr_ranking_kvcache.py \ + --gin_config_file ${GIN_CONFIG} \ + --checkpoint_dir ${HSTU_CKPT_DIR} \ + --max_bs ${MAX_BS} \ + --kvcache_config_file ${KVCACHE_MANAGER_CONFIG_FILE} > ${EXPORT_LOG} 2>&1 +EXIT_CODE=$? +if [[ ${EXIT_CODE} -ne 0 ]]; then + echo -e "${COLOR_ERROR}\n***\n*** AOTI export failed with exit code ${EXIT_CODE}\n***${COLOR_RESET}" 1>&2 + cat ${EXPORT_LOG} 1>&2 + echo -e "${COLOR_ERROR}\n***\n*** Test Suite FAILED\n***${COLOR_RESET}" 1>&2 + exit 1 +fi +if [[ ! -d ${AOTI_DIR}/${MODEL_NAME}_model ]]; then + echo -e "${COLOR_ERROR}\n***\n*** Export did not produce ${AOTI_DIR}/${MODEL_NAME}_model\n***${COLOR_RESET}" 1>&2 + cat ${EXPORT_LOG} 1>&2 + echo -e "${COLOR_ERROR}\n***\n*** Test Suite FAILED\n***${COLOR_RESET}" 1>&2 + exit 1 +fi +mv ${AOTI_DIR}/${MODEL_NAME}_model ${EXPORTED_MODEL} + +# Generate the reference output dump the client compares Triton against. The C++ +# KV-cache runtime needs a FlexKV server to talk to. +echo -e "${COLOR_DARK}Generating reference dump with the C++ KV-cache runtime${COLOR_RESET}" +start_kvcache_server || exit 1 +${AOTI_DIR}/cpp_inference/build/inference_hstu_gr_ranking_kvcache_exported_model \ + ${EXPORTED_MODEL} \ + ${DUMP_DIR} >> ${EXPORT_LOG} 2>&1 +EXIT_CODE=$? +stop_kvcache_server +if [[ ${EXIT_CODE} -ne 0 ]]; then + echo -e "${COLOR_ERROR}\n***\n*** Reference inference failed with exit code ${EXIT_CODE}\n***${COLOR_RESET}" 1>&2 + cat ${EXPORT_LOG} 1>&2 + echo -e "${COLOR_ERROR}\n***\n*** Test Suite FAILED\n***${COLOR_RESET}" 1>&2 + exit 1 +fi + +# The client reads the dump from inference_aoti, so link it to the copy held with +# the test artifacts. The link name is removed first: `ln -sfn` would otherwise +# create the link inside a pre-existing directory of that name. +rm -rf ${AOTI_DIR}/export_test_dump +ln -sfn ${DUMP_DIR} ${AOTI_DIR}/export_test_dump + +# Assemble the model repository: torch_aoti configuration from the recsys tree, +# exported package as version 1. +echo -e "${COLOR_DARK}Setting up model repository in ${MODELDIR}${COLOR_RESET}" +mkdir -p ${MODELDIR} +cp -r ${AOTI_DIR}/triton_aoti/${MODEL_NAME} ${MODELDIR}/${MODEL_NAME} +cp -r ${EXPORTED_MODEL} ${MODELDIR}/${MODEL_NAME}/1 +echo -e "${COLOR_DARK}ls ${MODELDIR}/${MODEL_NAME}${COLOR_RESET}" +ls -lha ${MODELDIR}/${MODEL_NAME} + +start_kvcache_server || exit 1 + +SERVER_ARGS="--model-repository=${MODELDIR} --backend-directory=${BACKEND_DIR} --log-verbose=1" +echo -e "${COLOR_DARK}Running ${SERVER} (backends: ${BACKEND_DIR})${COLOR_RESET}" +# Reports a dynamic-linker mismatch between the overlaid tritonserver and the +# image's libraries directly, instead of as an opaque startup timeout. +if ! ${SERVER} --version; then + echo -e "${COLOR_ERROR}\n***\n*** ${SERVER} failed to run\n***${COLOR_RESET}" 1>&2 + stop_kvcache_server + echo -e "${COLOR_ERROR}\n***\n*** Test Suite FAILED\n***${COLOR_RESET}" 1>&2 + exit 1 +fi +run_server +set +e +if [[ "${SERVER_PID}" -eq 0 ]]; then + echo -e "${COLOR_ERROR}\n***\n*** Failed to start ${SERVER}\n***${COLOR_RESET}" 1>&2 + cat ${SERVER_LOG} 1>&2 + stop_kvcache_server + echo -e "${COLOR_ERROR}\n***\n*** Test Suite FAILED\n***${COLOR_RESET}" 1>&2 + exit 1 +fi + +# The client resolves the reference dump relative to inference_aoti. +cd ${AOTI_DIR} +TEST_NAME="test_tritonserver_aoti_hstu_model" +python3 ./${TEST_NAME}.py > ${CLIENT_LOG} 2>&1 +EXIT_CODE=$? +cat ${CLIENT_LOG} +if [[ ${EXIT_CODE} -ne 0 ]]; then + echo -e "${COLOR_ERROR}\n***\n*** Test '${TEST_NAME}' Failed with exit code ${EXIT_CODE}\n***${COLOR_RESET}" 1>&2 + RET=1 +else + echo -e "${COLOR_INFO}\n***\n*** Test '${TEST_NAME}' Passed\n***${COLOR_RESET}" +fi -kill ${triton_pid} || true -kill -9 ${triton_pid} || true -kill ${kvserver_pid} || true +echo -e "${COLOR_DARK}Killing server (pid: ${SERVER_PID})${COLOR_RESET}" +kill -s SIGINT ${SERVER_PID} +wait ${SERVER_PID} || true +stop_kvcache_server -if [[ ${RET} -eq 0 ]]; then - echo -e "${COLOR_SUCCESS}\n***\n*** HSTU Torch AOTI Test Passed\n***${COLOR_RESET}" +if [[ ${RET} -ne 0 ]]; then + echo -e "${COLOR_ERROR}\n***\n*** Test Suite FAILED\n***${COLOR_RESET}" 1>&2 else - echo -e "${COLOR_ERROR}\n***\n*** HSTU Torch AOTI Test FAILED\n***${COLOR_RESET}" + echo -e "${COLOR_SUCCESS}\n***\n*** Test Suite PASSED\n***${COLOR_RESET}" fi exit ${RET} From 7f5de74ac6d620b5db9a80db8a5b6769dc5f22d7 Mon Sep 17 00:00:00 2001 From: Yingge He Date: Fri, 31 Jul 2026 01:22:19 -0700 Subject: [PATCH 3/4] test: Serve HSTU AOTI test from the image's own tritonserver The test image is built on top of the tritonserver under test, so TRITON_DIR resolves to it directly and the LD_LIBRARY_PATH override for a mounted install tree is no longer needed. Signed-off-by: Yingge He --- qa/L0_torch_aoti_hstu/test.sh | 20 +++++++++----------- 1 file changed, 9 insertions(+), 11 deletions(-) diff --git a/qa/L0_torch_aoti_hstu/test.sh b/qa/L0_torch_aoti_hstu/test.sh index 3e5147d162..1332765db7 100755 --- a/qa/L0_torch_aoti_hstu/test.sh +++ b/qa/L0_torch_aoti_hstu/test.sh @@ -11,11 +11,12 @@ # (Devtech-Compute/distributed-recommender: ci/tritonserver_test.sh), which # splits these phases across two container images. # -# The test image must provide the recsys-examples stack (export tooling, -# dynamicemb, the C++ KV-cache runtime, FlexKV, the torch_aoti model -# configuration, and the HSTU client), with the HSTU dataset and checkpoint -# mounted into RECSYS_DIR. tritonserver itself comes from TRITON_DIR, which CI -# points at a build from the pipeline under test. +# The test image must layer the recsys-examples stack (export tooling, dynamicemb, +# the C++ KV-cache runtime, FlexKV, the torch_aoti model configuration, and the +# HSTU client) on top of the tritonserver under test, with the HSTU dataset and +# checkpoint mounted into RECSYS_DIR. CI builds that image from the +# recsys-examples Dockerfile.tritonserver with BASE_IMAGE set to the server image +# built by the pipeline. source ../common/util.sh @@ -41,14 +42,10 @@ export CUDA_VISIBLE_DEVICES=0 TESTDIR=`pwd` -# TRITON_DIR selects which Triton is exercised. The recsys image ships a released -# tritonserver with a pinned PyTorch backend, so CI points TRITON_DIR at a mounted -# tree built by this pipeline in order to catch regressions in Triton itself. TRITON_DIR=${TRITON_DIR:="/opt/tritonserver"} SERVER=${TRITON_DIR}/bin/tritonserver BACKEND_DIR=${BACKEND_DIR:=${TRITON_DIR}/backends} SERVER_TIMEOUT=${SERVER_TIMEOUT:=300} -export LD_LIBRARY_PATH=${TRITON_DIR}/lib:${LD_LIBRARY_PATH} # recsys-examples HSTU tree shipped in the test image. inference_aoti holds the # export script, the C++ KV-cache runtime, the FlexKV launcher, the torch_aoti @@ -176,8 +173,9 @@ start_kvcache_server || exit 1 SERVER_ARGS="--model-repository=${MODELDIR} --backend-directory=${BACKEND_DIR} --log-verbose=1" echo -e "${COLOR_DARK}Running ${SERVER} (backends: ${BACKEND_DIR})${COLOR_RESET}" -# Reports a dynamic-linker mismatch between the overlaid tritonserver and the -# image's libraries directly, instead of as an opaque startup timeout. +# The image LD_PRELOADs the HSTU ops libraries into every process, so run the +# server once up front: a link error surfaces here rather than as an opaque +# startup timeout. if ! ${SERVER} --version; then echo -e "${COLOR_ERROR}\n***\n*** ${SERVER} failed to run\n***${COLOR_RESET}" 1>&2 stop_kvcache_server From 649d7d1764ae49c2162648f7b35a77fd92602db3 Mon Sep 17 00:00:00 2001 From: Yingge He Date: Thu, 6 Aug 2026 18:34:31 -0700 Subject: [PATCH 4/4] test: Align HSTU AOTI test comments with two-image CI split --- qa/L0_torch_aoti_hstu/test.sh | 10 ++++------ 1 file changed, 4 insertions(+), 6 deletions(-) diff --git a/qa/L0_torch_aoti_hstu/test.sh b/qa/L0_torch_aoti_hstu/test.sh index 1332765db7..38b64dea03 100755 --- a/qa/L0_torch_aoti_hstu/test.sh +++ b/qa/L0_torch_aoti_hstu/test.sh @@ -11,12 +11,10 @@ # (Devtech-Compute/distributed-recommender: ci/tritonserver_test.sh), which # splits these phases across two container images. # -# The test image must layer the recsys-examples stack (export tooling, dynamicemb, -# the C++ KV-cache runtime, FlexKV, the torch_aoti model configuration, and the -# HSTU client) on top of the tritonserver under test, with the HSTU dataset and -# checkpoint mounted into RECSYS_DIR. CI builds that image from the -# recsys-examples Dockerfile.tritonserver with BASE_IMAGE set to the server image -# built by the pipeline. +# It runs in the recsys-examples tritonserver image, which carries the export +# tooling, dynamicemb, the C++ KV-cache runtime, FlexKV, the torch_aoti model +# configuration, the HSTU client, and the tritonserver to serve with. The HSTU +# dataset and checkpoint are mounted into RECSYS_DIR. source ../common/util.sh