Enhance unit tests with timeout handling and stability improvements

.azure-pipelines/scripts/cuda_unit_test/run_cuda_ut.sh

@@ -68,7 +68,7 @@ function run_unit_test() {
     uv pip install -r test/unit/test_cuda/requirements.txt
     uv pip install -r test/unit/test_cuda/requirements_diffusion.txt
     uv pip install -U transformers chardet
-    uv pip install -U pytest-cov
+    uv pip install -U pytest-cov pytest-timeout
     uv pip install kernels==0.15.2 # For sm120: https://github.com/huggingface/transformers/blob/v5.13.1/setup.py#L93
     uv pip uninstall torch torchvision
     uv pip install torch==2.13.0 torchvision torchao --index-url https://download.pytorch.org/whl/cu130
@@ -100,7 +100,10 @@ function run_unit_test() {
         local test_basename=$(basename ${test_file} .py)
         local ut_log_name=${LOG_DIR}/unittest_cuda_${test_basename}.log
 
-        pytest -m "not skip_ci" --cov=auto_round --cov-report= --cov-append -vs --disable-warnings --junitxml="${ut_log_name%.log}.xml" ${test_file} 2>&1 | tee ${ut_log_name}
+        COVERAGE_CORE=sysmon pytest -m "not skip_ci" \
+            --cov=auto_round --cov-report= --cov-append --timeout=60 --session-timeout=720 \
+            -vs --disable-warnings --durations=0 --durations-min=1 --junitxml="${ut_log_name%.log}.xml" \
+            ${test_file} 2>&1 | tee ${ut_log_name}
         echo "##[endgroup]"
     done
     [ -f .coverage ] && cp .coverage "${LOG_DIR}/.coverage.part${test_part}"
@@ -113,7 +116,7 @@ function run_unit_test_llmc() {
     cd "${BUILD_SOURCESDIRECTORY}" || exit 1
     rm -rf /root/.venv
     uv venv --python=3.12 /root/.venv
-    uv pip install -U pytest-cov
+    uv pip install -U pytest-cov pytest-timeout
     BUILD_TYPE="nightly" uv pip install \
         -r test/integration/test_cuda/requirements_llmc.txt \
         --extra-index-url https://download.pytorch.org/whl/cu130 \
@@ -131,7 +134,10 @@ function run_unit_test_llmc() {
         echo "##[group]Running ${test_file}..."
         local test_basename=$(basename ${test_file} .py)
         local ut_log_name=${LOG_DIR}/unittest_cuda_llmc_${test_basename}.log
-        pytest -m "not skip_ci" --cov=auto_round --cov-report= --cov-append -vs --disable-warnings --junitxml="${ut_log_name%.log}.xml" ${test_file} 2>&1 | tee ${ut_log_name}
+        COVERAGE_CORE=sysmon pytest -m "not skip_ci" \
+            --cov=auto_round --cov-report= --cov-append -vs --disable-warnings \
+            --durations=0 --durations-min=1 --junitxml="${ut_log_name%.log}.xml" \
+            ${test_file} 2>&1 | tee ${ut_log_name}
         echo "##[endgroup]"
     done
     [ -f .coverage ] && cp .coverage "${LOG_DIR}/.coverage.llmc"
@@ -144,7 +150,7 @@ function run_unit_test_sglang() {
     cd "${BUILD_SOURCESDIRECTORY}" || exit 1
     rm -rf /root/.venv
     uv venv --python=3.12 /root/.venv
-    uv pip install -U pytest-cov
+    uv pip install -U pytest-cov pytest-timeout
     uv pip install -r test/integration/test_cuda/requirements_sglang.txt \
         --prerelease=allow \
         --extra-index-url https://download.pytorch.org/whl/cu130 \
@@ -162,7 +168,10 @@ function run_unit_test_sglang() {
         echo "##[group]Running ${test_file}..."
         local test_basename=$(basename ${test_file} .py)
         local ut_log_name=${LOG_DIR}/unittest_cuda_sglang_${test_basename}.log
-        pytest -m "not skip_ci" --cov=auto_round --cov-report= --cov-append -vs --disable-warnings --junitxml="${ut_log_name%.log}.xml" ${test_file} 2>&1 | tee ${ut_log_name}
+        COVERAGE_CORE=sysmon pytest -m "not skip_ci" \
+            --cov=auto_round --cov-report= --cov-append -vs --disable-warnings \
+            --durations=0 --durations-min=1 --junitxml="${ut_log_name%.log}.xml" \
+             ${test_file} 2>&1 | tee ${ut_log_name}
         echo "##[endgroup]"
     done
     [ -f .coverage ] && cp .coverage "${LOG_DIR}/.coverage.sglang"
@@ -175,7 +184,7 @@ function run_unit_test_vllm() {
     cd "${BUILD_SOURCESDIRECTORY}" || exit 1
     rm -rf /root/.venv
     uv venv --python=3.12 /root/.venv
-    uv pip install -U pytest-cov
+    uv pip install -U pytest-cov pytest-timeout
     uv pip install -r test/integration/test_cuda/requirements_vllm.txt \
         --extra-index-url https://download.pytorch.org/whl/cu130 \
         --index-strategy unsafe-best-match
@@ -193,7 +202,10 @@ function run_unit_test_vllm() {
         echo "##[group]Running ${test_file}..."
         local test_basename=$(basename ${test_file} .py)
         local ut_log_name=${LOG_DIR}/unittest_cuda_vllm_${test_basename}.log
-        pytest -m "not skip_ci" --cov=auto_round --cov-report= --cov-append -vs --disable-warnings --junitxml="${ut_log_name%.log}.xml" ${test_file} 2>&1 | tee ${ut_log_name}
+        COVERAGE_CORE=sysmon pytest -m "not skip_ci" \
+            --cov=auto_round --cov-report= --cov-append -vs --disable-warnings \
+            --durations=0 --durations-min=1 --junitxml="${ut_log_name%.log}.xml" \
+            ${test_file} 2>&1 | tee ${ut_log_name}
         echo "##[endgroup]"
     done
     [ -f .coverage ] && cp .coverage "${LOG_DIR}/.coverage.vllm"

.azure-pipelines/scripts/ut/run_ark_ut.sh

@@ -25,7 +25,7 @@ function run_unit_test() {
 
         echo "##[group]Running ark ${test_file}..."
         local ut_log_name="${LOG_DIR}/unittest_ark_${test_basename}.log"
-        pytest -v ${test_file} 2>&1 | tee ${ut_log_name}
+        COVERAGE_CORE=sysmon pytest -v ${test_file} 2>&1 | tee ${ut_log_name}
         echo "##[endgroup]"
     done
 }

.azure-pipelines/scripts/ut/run_ut.sh

@@ -22,7 +22,7 @@ function setup_environment() {
     git clone -b master --quiet --single-branch https://github.com/ggml-org/llama.cpp.git && cd llama.cpp/gguf-py && uv pip install .
 
     # install unit report dependencies
-    uv pip install pytest-cov
+    uv pip install pytest-cov pytest-timeout
     uv pip install -U chardet
     uv pip list
 
@@ -90,8 +90,9 @@ function run_unit_test() {
         local test_basename=$(basename ${test_file} .py)
         local ut_log_name=${LOG_DIR}/unittest_${test_basename}.log
 
-        numactl --physcpubind="${NUMA_CPUSET:-0-15}" --membind="${NUMA_NODE:-0}" \
-            pytest --cov=auto_round --cov-report= --cov-append \
+        COVERAGE_CORE=sysmon numactl --physcpubind="${NUMA_CPUSET:-0-15}" --membind="${NUMA_NODE:-0}" \
+            pytest -m "not skip_ci" --timeout=30 --session-timeout=600 --durations=0 --durations-min=1 \
+                --cov=auto_round --cov-report= --cov-append \
                 -vs --disable-warnings --junitxml="${ut_log_name%.log}.xml" ${test_file} 2>&1 | tee ${ut_log_name}
         echo "##[endgroup]"
     done
@@ -102,15 +103,16 @@ function run_inc_unit_test() {
     INC_PT_ONLY=1 uv pip install -r /auto-round/test/integration/test_cpu/requirements_inc.txt --extra-index-url https://download.pytorch.org/whl/cpu
     echo "##[endgroup]"
 
-    cd /auto-round/test || exit 1
+    cd /auto-round/test/integration || exit 1
 
     for test_file in $(find ./test_cpu -name "test_inc*.py" | sort); do
         echo "##[group]Running ${test_file}..."
         local test_basename=$(basename ${test_file} .py)
         local ut_log_name=${LOG_DIR}/unittest_${test_basename}.log
 
-        numactl --physcpubind="${NUMA_CPUSET:-0-15}" --membind="${NUMA_NODE:-0}" \
-            pytest --cov=auto_round --cov-report= --cov-append \
+        COVERAGE_CORE=sysmon numactl --physcpubind="${NUMA_CPUSET:-0-15}" --membind="${NUMA_NODE:-0}" \
+            pytest --timeout=30 --session-timeout=600 --durations=0 --durations-min=1 \
+                --cov=auto_round --cov-report= --cov-append \
                 -vs --disable-warnings --junitxml="${ut_log_name%.log}.xml" ${test_file} 2>&1 | tee ${ut_log_name}
         echo "##[endgroup]"
     done
@@ -123,15 +125,16 @@ function run_llmc_unit_test() {
     cd /auto-round && uv pip install .
     echo "##[endgroup]"
 
-    cd /auto-round/test || exit 1
+    cd /auto-round/test/integration || exit 1
 
     for test_file in $(find ./test_cpu -name "test_llmc*.py" | sort); do
         echo "##[group]Running ${test_file}..."
         local test_basename=$(basename ${test_file} .py)
         local ut_log_name=${LOG_DIR}/unittest_${test_basename}.log
 
-        numactl --physcpubind="${NUMA_CPUSET:-0-15}" --membind="${NUMA_NODE:-0}" \
-            pytest --cov=auto_round --cov-report= --cov-append \
+        COVERAGE_CORE=sysmon numactl --physcpubind="${NUMA_CPUSET:-0-15}" --membind="${NUMA_NODE:-0}" \
+            pytest --timeout=30 --session-timeout=600 --durations=0 --durations-min=1 \
+                --cov=auto_round --cov-report= --cov-append \
                 -vs --disable-warnings --junitxml="${ut_log_name%.log}.xml" ${test_file} 2>&1 | tee ${ut_log_name}
         echo "##[endgroup]"
     done
@@ -147,6 +150,10 @@ function collect_log() {
 function main() {
     setup_environment
     run_unit_test
+    if [ "$test_part" -eq 5 ] && [ "$NIGHTLY_TEST" = 1 ]; then
+        run_inc_unit_test
+        run_llmc_unit_test
+    fi
     collect_log
     check_storage_usage
     print_summary

.azure-pipelines/scripts/ut/run_ut_cuda.sh

@@ -31,7 +31,7 @@ function create_conda_env() {
         export PATH=/usr/local/cuda/bin${PATH:+:${PATH}}
         export LD_LIBRARY_PATH=$(python -c "import site; print(site.getsitepackages()[0])")/nvidia/nvjitlink/lib:$LD_LIBRARY_PATH
     fi
-    uv pip install pytest-cov cmake requests ninja psutil
+    uv pip install pytest-cov pytest-timeout cmake requests ninja psutil
     uv cache prune
 }
 
@@ -111,7 +111,8 @@ function run_unit_test() {
         local ut_log_name=${LOG_DIR}/unittest_cuda_${test_basename}.log
         echo "Running ${test_file}..."
 
-        pytest --cov=auto_round --cov-report= --cov-append -vs --disable-warnings ${test_file} 2>&1 | tee ${ut_log_name}
+        COVERAGE_CORE=sysmon pytest --cov=auto_round --cov-report= --cov-append \
+            --durations=0 --durations-min=1 -vs --disable-warnings ${test_file} 2>&1 | tee ${ut_log_name}
     done
     [ -f .coverage ] && cp .coverage ${LOG_DIR}/.coverage.unit
 
@@ -145,7 +146,8 @@ function run_unit_test_vlm() {
         local ut_log_name=${LOG_DIR}/unittest_cuda_vlm_${test_basename}.log
         echo "Running ${test_file}..."
 
-        pytest --cov=auto_round --cov-report= --cov-append -vs --disable-warnings ${test_file} 2>&1 | tee ${ut_log_name}
+        COVERAGE_CORE=sysmon pytest --cov=auto_round --cov-report= --cov-append \
+            --durations=0 --durations-min=1 -vs --disable-warnings ${test_file} 2>&1 | tee ${ut_log_name}
     done
     [ -f .coverage ] && cp .coverage ${LOG_DIR}/.coverage.vlm
 
@@ -174,7 +176,8 @@ function run_unit_test_llmc() {
         local ut_log_name=${LOG_DIR}/unittest_cuda_llmc_${test_basename}.log
         echo "Running ${test_file}..."
 
-        pytest --cov=auto_round --cov-report= --cov-append -vs --disable-warnings ${test_file} 2>&1 | tee ${ut_log_name}
+        COVERAGE_CORE=sysmon pytest --cov=auto_round --cov-report= --cov-append \
+            --durations=0 --durations-min=1 -vs --disable-warnings ${test_file} 2>&1 | tee ${ut_log_name}
     done
     [ -f .coverage ] && cp .coverage ${LOG_DIR}/.coverage.llmc
 
@@ -207,7 +210,8 @@ function run_unit_test_sglang() {
         local ut_log_name=${LOG_DIR}/unittest_cuda_sglang_${test_basename}.log
         echo "Running ${test_file}..."
 
-        pytest --cov=auto_round --cov-report= --cov-append -vs --disable-warnings ${test_file} 2>&1 | tee ${ut_log_name}
+        COVERAGE_CORE=sysmon pytest --cov=auto_round --cov-report= --cov-append \
+            --durations=0 --durations-min=1 -vs --disable-warnings ${test_file} 2>&1 | tee ${ut_log_name}
     done
     [ -f .coverage ] && cp .coverage ${LOG_DIR}/.coverage.sglang
 
@@ -241,7 +245,8 @@ function run_unit_test_vllm() {
         local ut_log_name=${LOG_DIR}/unittest_cuda_vllm_${test_basename}.log
         echo "Running ${test_file}..."
 
-        pytest --cov=auto_round --cov-report= --cov-append -vs --disable-warnings ${test_file} 2>&1 | tee ${ut_log_name}
+        COVERAGE_CORE=sysmon pytest --cov=auto_round --cov-report= --cov-append \
+            --durations=0 --durations-min=1 -vs --disable-warnings ${test_file} 2>&1 | tee ${ut_log_name}
     done
     [ -f .coverage ] && cp .coverage ${LOG_DIR}/.coverage.vllm
 

.azure-pipelines/scripts/ut/run_ut_hpu.sh

@@ -8,7 +8,7 @@ function setup_environment() {
     export TZ='Asia/Shanghai'
     export TQDM_MININTERVAL=60
     export HF_HUB_DISABLE_PROGRESS_BARS=1
-    pip install pytest-cov
+    pip install pytest-cov pytest-timeout
     pip list
     echo "##[endgroup]"
 
@@ -33,14 +33,16 @@ function run_unit_test() {
 
         echo "##[group]Running ${test_file} in HPU lazy mode..."
         local ut_log_name="${LOG_DIR}/unittest_lazy_${test_basename}.log"
-        PT_HPU_LAZY_MODE=1 pytest --cov="${auto_round_path}" \
+        COVERAGE_CORE=sysmon PT_HPU_LAZY_MODE=1 pytest --cov="${auto_round_path}" \
+            --timeout=30 --session-timeout=600 --durations=0 --durations-min=1 \
             --cov-report= --cov-append -vs --disable-warnings \
             --junitxml="${ut_log_name%.log}.xml" ${test_file} 2>&1 | tee ${ut_log_name}
         echo "##[endgroup]"
 
         echo "##[group]Running ${test_file} in HPU compile mode..."
         local ut_log_name="${LOG_DIR}/unittest_compile_${test_basename}.log"
-        PT_HPU_LAZY_MODE=0 pytest --mode compile --cov="${auto_round_path}" \
+        COVERAGE_CORE=sysmon PT_HPU_LAZY_MODE=0 pytest --mode compile --cov="${auto_round_path}" \
+            --timeout=30 --session-timeout=600 --durations=0 --durations-min=1 \
             --cov-report= --cov-append -vs --disable-warnings \
             --junitxml="${ut_log_name%.log}.xml" ${test_file} 2>&1 | tee ${ut_log_name}
         echo "##[endgroup]"

.azure-pipelines/scripts/ut/run_ut_xpu.sh

@@ -5,7 +5,7 @@ source /auto-round/.azure-pipelines/scripts/change_color.sh
 
 function setup_environment() {
     echo "##[group]set up UT env..."
-    uv pip install pytest-cov
+    uv pip install pytest-cov pytest-timeout
     uv pip list
     echo "##[endgroup]"
 
@@ -38,9 +38,10 @@ function run_unit_test() {
 
         echo "##[group]Running ark ${test_file}..."
         local ut_log_name="${LOG_DIR}/unittest_ark_${test_basename}.log"
-        numactl --physcpubind="${NUMA_CPUSET:-0-27}" --membind="${NUMA_NODE:-0}" \
-            pytest --cov="${auto_round_path}" --cov-report= --cov-append -vs --disable-warnings \
-            --junitxml="${ut_log_name%.log}.xml" ${test_file} 2>&1 | tee ${ut_log_name}
+        COVERAGE_CORE=sysmon numactl --physcpubind="${NUMA_CPUSET:-0-27}" --membind="${NUMA_NODE:-0}" \
+            pytest --timeout=30 --session-timeout=600 --durations=0 --durations-min=1 \
+                --cov="${auto_round_path}" --cov-report= --cov-append -vs --disable-warnings \
+                --junitxml="${ut_log_name%.log}.xml" ${test_file} 2>&1 | tee ${ut_log_name}
         echo "##[endgroup]"
     done
 
@@ -49,9 +50,10 @@ function run_unit_test() {
 
         echo "##[group]Running xpu ${test_file}..."
         local ut_log_name="${LOG_DIR}/unittest_xpu_${test_basename}.log"
-        numactl --physcpubind="${NUMA_CPUSET:-0-27}" --membind="${NUMA_NODE:-0}" \
-            pytest --cov="${auto_round_path}" --cov-report= --cov-append -vs --disable-warnings \
-            --junitxml="${ut_log_name%.log}.xml" ${test_file} 2>&1 | tee ${ut_log_name}
+        COVERAGE_CORE=sysmon numactl --physcpubind="${NUMA_CPUSET:-0-27}" --membind="${NUMA_NODE:-0}" \
+            pytest --timeout=30 --session-timeout=600 --durations=0 --durations-min=1 \
+                --cov="${auto_round_path}" --cov-report= --cov-append -vs --disable-warnings \
+                --junitxml="${ut_log_name%.log}.xml" ${test_file} 2>&1 | tee ${ut_log_name}
         echo "##[endgroup]"
     done
 }
@@ -69,9 +71,10 @@ function run_unit_test_llmc() {
 
         echo "##[group]Running xpu llmc ${test_file}..."
         local ut_log_name="${LOG_DIR}/unittest_xpu_${test_basename}.log"
-        numactl --physcpubind="${NUMA_CPUSET:-0-27}" --membind="${NUMA_NODE:-0}" \
-            pytest --cov="${auto_round_path}" --cov-report= --cov-append -vs --disable-warnings \
-            --junitxml="${ut_log_name%.log}.xml" ${test_file} 2>&1 | tee ${ut_log_name}
+        COVERAGE_CORE=sysmon numactl --physcpubind="${NUMA_CPUSET:-0-27}" --membind="${NUMA_NODE:-0}" \
+            pytest --timeout=30 --session-timeout=600 --durations=0 --durations-min=1 \
+                --cov="${auto_round_path}" --cov-report= --cov-append -vs --disable-warnings \
+                --junitxml="${ut_log_name%.log}.xml" ${test_file} 2>&1 | tee ${ut_log_name}
         echo "##[endgroup]"
     done
 }

test/README.md

@@ -140,6 +140,25 @@ def test_model_inference(tiny_opt_model_path):
 - **CPU-specific** → `*/test_cpu/`, **CUDA-specific** → `*/test_cuda/`
 - Import from parent: `from test.helpers import ...`
 
+### CI Markers and Timeouts
+
+- For long-running or non-critical tests under `test_cuda/`, use
+  `@pytest.mark.skip_ci(reason="...")` and provide a clear reason why the test should not run in CI.
+- In CI, each test function has a default timeout of **30 seconds**, while each test file has a timeout of
+  **10 minutes**.
+- Prefer simplifying a test so that it completes within the default timeout. If the test cannot be reduced further,
+  extend its timeout with `@pytest.mark.timeout(seconds)`. If a test file exceeds the 10-minute limit, split its tests
+  into smaller files whenever possible.
+
+```python
+@pytest.mark.timeout(120)
+def test_long_running_case(): ...
+
+
+@pytest.mark.skip_ci(reason="Time-consuming accuracy evaluation; covered by nightly tests")
+def test_optional_accuracy_evaluation(): ...
+```
+
 ## 5. Running Tests
 
 ```sh

test/integration/test_xpu/test_llmc_integration.py

@@ -75,6 +75,7 @@
 
 
 @pytest.mark.skipif(torch.xpu.device_count() < 1, reason="test requires at least 1 XPU")
+@pytest.mark.timeout(60)
 @pytest.mark.parametrize(
     "recipe",
     [

test/unit/test_ark/test_model.py

@@ -58,20 +58,23 @@ def main_op(self, format, bits, group_size, sym, dtype, device, fast_cfg=True, t
         evaluate_accuracy(model, tokenizer, threshold=tar_acc, batch_size=32, limit=limit)
         torch.xpu.empty_cache()
 
+    @pytest.mark.timeout(60)
     @pytest.mark.parametrize("format", ["auto_round", "auto_round:gptqmodel"])
     @pytest.mark.parametrize("bits, group_size, sym", [(4, 128, True), (8, 128, True)])
     @pytest.mark.parametrize("dtype", [torch.bfloat16])
     @pytest.mark.parametrize("device", ["cpu", "xpu"])
     def test_formats(self, format, bits, group_size, sym, dtype, device):
         self.main_op(format, bits, group_size, sym, dtype, device)
 
+    @pytest.mark.timeout(60)
     @pytest.mark.parametrize("format", ["auto_round:auto_awq"])
     @pytest.mark.parametrize("bits, group_size, sym", [(4, 32, True)])
     @pytest.mark.parametrize("dtype", [torch.float16])
     @pytest.mark.parametrize("device", ["cpu", "xpu"])
     def test_awq_fp16(self, format, bits, group_size, sym, dtype, device):
         self.main_op(format, bits, group_size, sym, dtype, device)
 
+    @pytest.mark.timeout(300)
     @pytest.mark.parametrize("format", ["auto_round"])
     @pytest.mark.parametrize("bits, group_size, sym", [(2, 32, False)])
     @pytest.mark.parametrize("dtype", [torch.bfloat16])

test/unit/test_cpu/advanced/test_low_precision_input_model.py

@@ -86,6 +86,7 @@ def test_w4a16(self):
             model.model.layers[0].mlp.up_proj.weight.dtype == torch.bfloat16
         ), "CompressedLinear layer was not converted to Linear"
 
+    @pytest.mark.timeout(60)
     def test_w4a16_to_mxfp4(self, tmp_path):
         model = get_tiny_model(get_model_path(self.w4a16_model_path))
         model.config.name_or_path = None  # Clear the name_or_path to avoid MTP copying issues