Merge branch 'main' into aflowers/fix-cross-arch-uv

Author: ishandhanani

pyproject.toml

@@ -40,6 +40,8 @@ dev = [
     "ty",  # Astral's fast type checker (replaces mypy)
     "fastapi>=0.109.0",
     "httpx>=0.27.0",  # Required by FastAPI TestClient
+    "uvicorn>=0.27.0",  # Required for integration test mock server
+    "aiperf",  # Benchmark tool for trace replay integration tests
 ]
 
 # =============================================================================
@@ -77,7 +79,10 @@ testpaths = ["tests"]
 pythonpath = ["tests"]
 python_files = ["test_*.py"]
 python_functions = ["test_*"]
-addopts = "-v --tb=short"
+addopts = "-v --tb=short -m 'not integration'"
+markers = [
+    "integration: slow tests that run real aiperf against a mock server",
+]
 
 # =============================================================================
 # ty - Astral's fast type checker (10-100x faster than mypy)

src/srtctl/benchmarks/__init__.py

@@ -4,7 +4,17 @@
 """Benchmark runners for srtctl."""
 
 # Import runners to trigger registration
-from srtctl.benchmarks import gpqa, gsm8k, longbenchv2, mmlu, mooncake_router, router, sa_bench, sglang_bench
+from srtctl.benchmarks import (
+    gpqa,
+    gsm8k,
+    longbenchv2,
+    mmlu,
+    mooncake_router,
+    router,
+    sa_bench,
+    sglang_bench,
+    trace_replay,
+)
 from srtctl.benchmarks.base import (
     BenchmarkRunner,
     get_runner,
@@ -26,4 +36,5 @@
     "longbenchv2",
     "router",
     "mooncake_router",
+    "trace_replay",
 ]

src/srtctl/benchmarks/sa_bench.py

@@ -99,5 +99,7 @@ def build_command(
             str(b.random_range_ratio) if b.random_range_ratio is not None else "0.8",
             str(b.num_prompts_mult) if b.num_prompts_mult is not None else "10",
             str(b.num_warmup_mult) if b.num_warmup_mult is not None else "2",
+            b.custom_tokenizer or "",
+            str(b.use_chat_template).lower(),
         ]
         return cmd

src/srtctl/benchmarks/scripts/mooncake-router/bench.sh

@@ -26,10 +26,10 @@ if [ -n "${AIPERF_SERVER_METRICS_URLS:-}" ]; then
     fi
 fi
 
-# Setup directories
-BASE_DIR="/logs"
-TRACE_DIR="${BASE_DIR}/traces"
-ARTIFACT_DIR="${BASE_DIR}/artifacts"
+# Setup directories (BASE_DIR defaults to /logs inside container, overridable for testing)
+BASE_DIR="${BASE_DIR:-/logs}"
+TRACE_DIR="${TRACE_DIR:-${BASE_DIR}/traces}"
+ARTIFACT_DIR="${ARTIFACT_DIR:-${BASE_DIR}/artifacts}"
 mkdir -p "${TRACE_DIR}"
 mkdir -p "${ARTIFACT_DIR}"
 

src/srtctl/benchmarks/scripts/sa-bench/backend_request_func.py

@@ -565,10 +565,52 @@ def _fix_v5_tokenizer_components(tokenizer, model_name_or_path):
         backend.decoder = raw.decoder
 
 
+def _load_glm_moe_dsa_tokenizer(pretrained_model_name_or_path: str) -> "PreTrainedTokenizerFast":
+    """Load GLM-Moe-Dsa / GLM-5 tokenizer directly from tokenizer.json.
+
+    Works around incompatibilities when the checkpoint was saved with
+    transformers 5.x (TokenizersBackend / list-style extra_special_tokens).
+    """
+    import json
+    from pathlib import Path
+
+    from tokenizers import Tokenizer as RustTokenizer
+    from transformers import PreTrainedTokenizerFast
+
+    _SAFE_CONFIG_KEYS = (
+        "pad_token", "pad_token_id", "eos_token", "eos_token_id",
+        "bos_token", "bos_token_id", "unk_token", "unk_token_id",
+        "model_max_length", "padding_side", "truncation_side",
+    )
+
+    path = Path(pretrained_model_name_or_path)
+    tokenizer_json = path / "tokenizer.json"
+    if not tokenizer_json.exists():
+        raise FileNotFoundError(
+            f"Expected tokenizer.json at {tokenizer_json}. "
+            "GlmMoeDsaTokenizer loads from tokenizer.json only."
+        )
+
+    rust_tok = RustTokenizer.from_file(str(tokenizer_json))
+    init_kwargs = {}
+    config_path = path / "tokenizer_config.json"
+    if config_path.exists():
+        with open(config_path, encoding="utf-8") as f:
+            config = json.load(f)
+        for key in _SAFE_CONFIG_KEYS:
+            if key in config:
+                init_kwargs[key] = config[key]
+        if "extra_special_tokens" in config:
+            init_kwargs["additional_special_tokens"] = config["extra_special_tokens"]
+
+    return PreTrainedTokenizerFast(tokenizer_object=rust_tok, **init_kwargs)
+
+
 def get_tokenizer(
     pretrained_model_name_or_path: str,
     tokenizer_mode: str = "auto",
     trust_remote_code: bool = False,
+    custom_tokenizer: str | None = None,
     **kwargs,
 ) -> PreTrainedTokenizer | PreTrainedTokenizerFast:
     if pretrained_model_name_or_path is not None and not os.path.exists(pretrained_model_name_or_path):
@@ -587,14 +629,31 @@ def get_tokenizer(
                 "to use mistral tokenizer mode."
             ) from e
         return MistralTokenizer.from_pretrained(str(pretrained_model_name_or_path))
-    else:
-        tokenizer = AutoTokenizer.from_pretrained(
-            pretrained_model_name_or_path,
-            trust_remote_code=trust_remote_code,
-            **kwargs,
-        )
-        _fix_v5_tokenizer_components(tokenizer, pretrained_model_name_or_path)
-        return tokenizer
+    if custom_tokenizer:
+        if custom_tokenizer == "glm_moe_dsa":
+            return _load_glm_moe_dsa_tokenizer(pretrained_model_name_or_path)
+        from importlib import import_module
+        try:
+            module_path, class_name = custom_tokenizer.rsplit('.', 1)
+            module = import_module(module_path)
+            tokenizer_class = getattr(module, class_name)
+            return tokenizer_class.from_pretrained(
+                pretrained_model_name_or_path,
+                trust_remote_code=trust_remote_code,
+                **kwargs,
+            )
+        except (ValueError, ImportError, AttributeError) as e:
+            raise ValueError(
+                f"Failed to load custom_tokenizer '{custom_tokenizer}'. "
+                "Expected 'glm_moe_dsa' or 'module.path.ClassName'.") from e
+
+    tokenizer = AutoTokenizer.from_pretrained(
+        pretrained_model_name_or_path,
+        trust_remote_code=trust_remote_code,
+        **kwargs,
+    )
+    _fix_v5_tokenizer_components(tokenizer, pretrained_model_name_or_path)
+    return tokenizer
 
 
 ASYNC_REQUEST_FUNCS = {

src/srtctl/benchmarks/scripts/sa-bench/bench.sh

@@ -62,6 +62,20 @@ DECODE_GPUS=${11:-0}
 RANDOM_RANGE_RATIO=${12:-0.8}
 NUM_PROMPTS_MULT=${13:-10}
 NUM_WARMUP_MULT=${14:-2}
+CUSTOM_TOKENIZER=${15:-}
+USE_CHAT_TEMPLATE=${16:-true}
+
+# Build optional custom tokenizer args
+CUSTOM_TOKENIZER_ARGS=()
+if [ -n "$CUSTOM_TOKENIZER" ]; then
+    CUSTOM_TOKENIZER_ARGS=(--custom-tokenizer "$CUSTOM_TOKENIZER")
+fi
+
+# Build optional chat template args
+CHAT_TEMPLATE_ARGS=()
+if [ "$USE_CHAT_TEMPLATE" = "true" ]; then
+    CHAT_TEMPLATE_ARGS=(--use-chat-template)
+fi
 
 # Parse endpoint into host:port
 HOST=$(echo "$ENDPOINT" | sed 's|http://||' | cut -d: -f1)
@@ -121,7 +135,8 @@ for concurrency in "${CONCURRENCY_LIST[@]}"; do
         --request-rate 250 \
         --percentile-metrics ttft,tpot,itl,e2el \
         --max-concurrency "$concurrency" \
-        --trust-remote-code
+        --trust-remote-code \
+        "${CUSTOM_TOKENIZER_ARGS[@]}"
 
     num_prompts=$((concurrency * NUM_PROMPTS_MULT))
 
@@ -151,7 +166,8 @@ for concurrency in "${CONCURRENCY_LIST[@]}"; do
         --percentile-metrics ttft,tpot,itl,e2el \
         --max-concurrency "$concurrency" \
         --trust-remote-code \
-        --use-chat-template \
+        "${CHAT_TEMPLATE_ARGS[@]}" \
+        "${CUSTOM_TOKENIZER_ARGS[@]}" \
         --save-result --result-dir "$result_dir" --result-filename "$result_filename"
     set +x
 

src/srtctl/benchmarks/scripts/sa-bench/benchmark_serving.py

@@ -837,6 +837,7 @@ def main(args: argparse.Namespace):
         tokenizer_id,
         tokenizer_mode=tokenizer_mode,
         trust_remote_code=args.trust_remote_code,
+        custom_tokenizer=args.custom_tokenizer,
     )
 
     if args.dataset is not None:
@@ -1279,6 +1280,14 @@ def main(args: argparse.Namespace):
         '"custom" will use --tokenizer to select the preregistered tokenizer.',
     )
 
+    parser.add_argument(
+        "--custom-tokenizer",
+        type=str,
+        default=None,
+        help="Custom tokenizer to use (e.g., 'glm_moe_dsa' or 'module.path.ClassName'). "
+        "When set, overrides the default tokenizer loading.",
+    )
+
     parser.add_argument(
         "--served-model-name",
         type=str,

src/srtctl/benchmarks/scripts/trace-replay/bench.sh

@@ -0,0 +1,117 @@
+#!/bin/bash
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+# Trace Replay Benchmark using aiperf
+# Replays a user-provided JSONL trace dataset at configurable concurrency levels.
+# Uses aiperf with --custom-dataset-type mooncake_trace.
+#
+# Usage: bench.sh ENDPOINT MODEL_NAME TRACE_FILE CONCURRENCIES [TTFT_THRESHOLD] [ITL_THRESHOLD] [TOKENIZER_PATH]
+
+set -e
+
+ENDPOINT=$1
+MODEL_NAME=${2:-"test-model"}
+TRACE_FILE=$3
+CONCURRENCIES=${4:-"1"}
+TTFT_THRESHOLD=${5:-2000}
+ITL_THRESHOLD=${6:-25}
+TOKENIZER_PATH=${7:-"/model"}
+
+# Optional: extra Prometheus endpoints for AIPerf server metrics
+SERVER_METRICS_ARGS=()
+if [ -n "${AIPERF_SERVER_METRICS_URLS:-}" ]; then
+    IFS=',' read -r -a server_metrics_urls <<< "${AIPERF_SERVER_METRICS_URLS}"
+    if [ ${#server_metrics_urls[@]} -gt 0 ]; then
+        SERVER_METRICS_ARGS+=(--server-metrics "${server_metrics_urls[@]}")
+    fi
+fi
+
+# Setup directories (BASE_DIR defaults to /logs inside container, overridable for testing)
+BASE_DIR="${BASE_DIR:-/logs}"
+ARTIFACT_DIR="${ARTIFACT_DIR:-${BASE_DIR}/artifacts}"
+mkdir -p "${ARTIFACT_DIR}"
+
+# Increase aiperf HTTP timeout
+export AIPERF_HTTP_SO_RCVTIMEO=120
+
+echo "=============================================="
+echo "Trace Replay Benchmark (aiperf)"
+echo "=============================================="
+echo "Endpoint: ${ENDPOINT}"
+echo "Model: ${MODEL_NAME}"
+echo "Trace File: ${TRACE_FILE}"
+echo "Concurrencies: ${CONCURRENCIES}"
+echo "TTFT Threshold: ${TTFT_THRESHOLD}ms"
+echo "ITL Threshold: ${ITL_THRESHOLD}ms"
+echo "Tokenizer Path: ${TOKENIZER_PATH}"
+echo "=============================================="
+
+# Validate trace file exists
+if [ ! -f "${TRACE_FILE}" ]; then
+    echo "ERROR: Trace file not found: ${TRACE_FILE}"
+    exit 1
+fi
+
+# Install aiperf if not present
+if ! command -v aiperf &> /dev/null; then
+    echo "Installing aiperf..."
+    pip install aiperf
+fi
+
+# Run small benchmark for warmup
+echo "Running warmup..."
+aiperf profile \
+    -m "${MODEL_NAME}" \
+    --tokenizer "${TOKENIZER_PATH}" \
+    --url "${ENDPOINT}" \
+    --streaming \
+    --ui simple \
+    --extra-inputs ignore_eos:true \
+    --concurrency 1 \
+    --request-count 5
+echo "Warmup complete"
+
+# Setup artifact directory
+MODEL_BASE_NAME="${MODEL_NAME##*/}"
+TIMESTAMP=$(date '+%Y%m%d_%H%M%S')
+
+# Parse concurrencies (comma-separated)
+IFS=',' read -r -a CONCURRENCY_LIST <<< "${CONCURRENCIES}"
+
+for C in "${CONCURRENCY_LIST[@]}"; do
+    echo ""
+    echo "=============================================="
+    echo "Running concurrency=${C}"
+    echo "=============================================="
+    echo "$(date '+%Y-%m-%d %H:%M:%S') - Starting benchmark at concurrency ${C}"
+
+    RUN_ARTIFACT_DIR="${ARTIFACT_DIR}/${MODEL_BASE_NAME}_trace_c${C}_${TIMESTAMP}"
+    mkdir -p "${RUN_ARTIFACT_DIR}"
+
+    aiperf profile \
+        -m "${MODEL_NAME}" \
+        --tokenizer "${TOKENIZER_PATH}" \
+        --input-file "${TRACE_FILE}" \
+        --custom-dataset-type mooncake_trace \
+        --url "${ENDPOINT}" \
+        --streaming \
+        --extra-inputs ignore_eos:true \
+        --concurrency "${C}" \
+        --random-seed 42 \
+        --ui simple \
+        --artifact-dir "${RUN_ARTIFACT_DIR}" \
+        "${SERVER_METRICS_ARGS[@]}" \
+        --goodput "time_to_first_token:${TTFT_THRESHOLD} inter_token_latency:${ITL_THRESHOLD}"
+
+    echo "$(date '+%Y-%m-%d %H:%M:%S') - Concurrency ${C} complete"
+
+    # List artifacts
+    ls -la "${RUN_ARTIFACT_DIR}" 2>/dev/null || true
+done
+
+echo ""
+echo "=============================================="
+echo "Trace Replay Benchmark Complete"
+echo "Results saved to: ${ARTIFACT_DIR}"
+echo "=============================================="