diff --git a/qa/L0_torch_aoti_hstu/test.sh b/qa/L0_torch_aoti_hstu/test.sh new file mode 100755 index 0000000000..38b64dea03 --- /dev/null +++ b/qa/L0_torch_aoti_hstu/test.sh @@ -0,0 +1,217 @@ +#!/bin/bash +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: BSD-3-Clause + +# End-to-end HSTU (Generative Recommenders) Torch AOTI test. +# +# 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. +# +# 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 + +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 + +export CUDA_VISIBLE_DEVICES=0 + +TESTDIR=`pwd` + +TRITON_DIR=${TRITON_DIR:="/opt/tritonserver"} +SERVER=${TRITON_DIR}/bin/tritonserver +BACKEND_DIR=${BACKEND_DIR:=${TRITON_DIR}/backends} +SERVER_TIMEOUT=${SERVER_TIMEOUT:=300} + +# 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 + +# 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}" +} + +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 +} + +rm -rf ${MODELDIR} ${EXPORTED_MODEL} ${DUMP_DIR} + +# 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}" +# 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 + 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 + +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} -ne 0 ]]; then + echo -e "${COLOR_ERROR}\n***\n*** Test Suite FAILED\n***${COLOR_RESET}" 1>&2 +else + echo -e "${COLOR_SUCCESS}\n***\n*** Test Suite PASSED\n***${COLOR_RESET}" +fi + +exit ${RET}