test_multi_lora.py

"""GovOn Multi-LoRA + LMCache Integration Test.

Phases:
  1. vLLM Multi-LoRA 기본 서빙 (베이스 vs LoRA 추론 비교)
  2. LMCache LoRA-aware caching 검증 (캐시 격리 + TTFT 감소)
  3. 결과 보고
"""

import time
import json
import sys
import os

import torch
from vllm import LLM, SamplingParams
from vllm.lora.request import LoRARequest

# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------
BASE_MODEL = "LGAI-EXAONE/EXAONE-4.0-32B-AWQ"
LORA_ADAPTER = "umyunsang/govon-civil-adapter"
LORA_NAME = "civil"

SYSTEM_PROMPT = "당신은 대한민국 공무원 민원 답변 전문가입니다."
TEST_PROMPT = "국민연금 수령 나이가 어떻게 되나요?"

# EXAONE chat template
PROMPT_TEMPLATE = (
    "[|system|]{system}[|endofturn|]\n"
    "[|user|]{user}[|endofturn|]\n"
    "[|assistant|]"
)


def build_prompt(system: str = SYSTEM_PROMPT, user: str = TEST_PROMPT) -> str:
    return PROMPT_TEMPLATE.format(system=system, user=user)


# ---------------------------------------------------------------------------
# Phase 1: Multi-LoRA Basic Serving
# ---------------------------------------------------------------------------
def phase1_multi_lora_basic() -> tuple:
    """vLLM Multi-LoRA 기본 동작 테스트."""
    print("=" * 60)
    print("Phase 1: Multi-LoRA Basic Serving Test")
    print("=" * 60)

    llm = LLM(
        model=BASE_MODEL,
        trust_remote_code=True,
        dtype="half",
        enforce_eager=True,
        gpu_memory_utilization=0.95,
        max_model_len=512,
        enable_lora=True,
        max_loras=2,
        max_lora_rank=64,
    )

    sampling_params = SamplingParams(temperature=0.7, max_tokens=128)
    prompt = build_prompt()

    # --- Base model inference ---
    print("\n[Base Model] Generating...")
    t0 = time.time()
    base_output = llm.generate([prompt], sampling_params)
    base_time = time.time() - t0
    base_text = base_output[0].outputs[0].text
    print(f"  Time: {base_time:.2f}s")
    print(f"  Output: {base_text[:200]}...")

    # --- LoRA adapter inference ---
    print(f"\n[LoRA: {LORA_NAME}] Generating...")
    lora_request = LoRARequest(LORA_NAME, 1, LORA_ADAPTER)
    t0 = time.time()
    lora_output = llm.generate(
        [prompt], sampling_params, lora_request=lora_request
    )
    lora_time = time.time() - t0
    lora_text = lora_output[0].outputs[0].text
    print(f"  Time: {lora_time:.2f}s")
    print(f"  Output: {lora_text[:200]}...")

    # --- Verify outputs differ ---
    outputs_differ = base_text != lora_text
    status = "PASS" if outputs_differ else "WARN - outputs identical"
    print(f"\n  Outputs differ (LoRA effect): {status}")

    results = {
        "status": "PASS" if outputs_differ else "WARN",
        "base_time_s": round(base_time, 3),
        "lora_time_s": round(lora_time, 3),
        "outputs_differ": outputs_differ,
    }
    return llm, results


# ---------------------------------------------------------------------------
# Phase 2: LMCache LoRA-aware Caching
# ---------------------------------------------------------------------------
def phase2_lmcache_test(llm: LLM) -> dict:
    """LMCache LoRA-aware caching 테스트."""
    print("\n" + "=" * 60)
    print("Phase 2: LMCache LoRA-aware Caching Test")
    print("=" * 60)

    try:
        import lmcache  # noqa: F401
        version = getattr(lmcache, "__version__", "unknown")
        print(f"  LMCache version: {version}")
    except ImportError:
        print("  SKIP: LMCache not installed")
        return {"status": "SKIP", "reason": "lmcache not installed"}

    sampling_params = SamplingParams(temperature=0.0, max_tokens=128)
    prompt = build_prompt()

    results: dict = {}

    # --- 2-a. Base model cache test (same prompt x2) ---
    print("\n[2-a] Base model: same prompt x2 (expect TTFT drop on 2nd)")
    ttfts_base = []
    for i in range(2):
        t0 = time.time()
        out = llm.generate([prompt], sampling_params)
        elapsed = time.time() - t0
        ttfts_base.append(round(elapsed, 4))
        print(f"  Run {i+1}: {elapsed:.4f}s")

    base_speedup = (
        (ttfts_base[0] - ttfts_base[1]) / ttfts_base[0] * 100
        if ttfts_base[0] > 0
        else 0
    )
    results["base_ttft_1"] = ttfts_base[0]
    results["base_ttft_2"] = ttfts_base[1]
    results["base_speedup_pct"] = round(base_speedup, 2)
    print(f"  Speedup: {base_speedup:.1f}%")

    # --- 2-b. LoRA model cache test (same prompt x2) ---
    print(f"\n[2-b] LoRA ({LORA_NAME}): same prompt x2 (expect TTFT drop)")
    lora_request = LoRARequest(LORA_NAME, 1, LORA_ADAPTER)
    ttfts_lora = []
    for i in range(2):
        t0 = time.time()
        out = llm.generate(
            [prompt], sampling_params, lora_request=lora_request
        )
        elapsed = time.time() - t0
        ttfts_lora.append(round(elapsed, 4))
        print(f"  Run {i+1}: {elapsed:.4f}s")

    lora_speedup = (
        (ttfts_lora[0] - ttfts_lora[1]) / ttfts_lora[0] * 100
        if ttfts_lora[0] > 0
        else 0
    )
    results["lora_ttft_1"] = ttfts_lora[0]
    results["lora_ttft_2"] = ttfts_lora[1]
    results["lora_speedup_pct"] = round(lora_speedup, 2)
    print(f"  Speedup: {lora_speedup:.1f}%")

    # --- 2-c. Cache isolation test ---
    # 베이스 모델 캐시가 LoRA 요청에 재사용되면 안 됨
    # 이미 base 캐시가 warm 상태 -> LoRA로 fresh prompt 실행
    print("\n[2-c] Cache isolation: base cache must NOT leak to LoRA")
    fresh_prompt = build_prompt(
        user="건강보험 피부양자 등록 조건이 무엇인가요?"
    )

    # Warm base cache
    _ = llm.generate([fresh_prompt], sampling_params)

    # LoRA with same prompt — should NOT benefit from base cache
    t0 = time.time()
    lora_fresh = llm.generate(
        [fresh_prompt], sampling_params, lora_request=lora_request
    )
    lora_fresh_time = time.time() - t0

    # LoRA 2nd call with same prompt — should benefit from LoRA cache
    t0 = time.time()
    lora_cached = llm.generate(
        [fresh_prompt], sampling_params, lora_request=lora_request
    )
    lora_cached_time = time.time() - t0

    isolation_ok = lora_fresh_time >= lora_cached_time * 0.8
    results["cache_isolation"] = "PASS" if isolation_ok else "FAIL"
    results["lora_fresh_time"] = round(lora_fresh_time, 4)
    results["lora_cached_time"] = round(lora_cached_time, 4)
    print(f"  LoRA fresh (after base warm): {lora_fresh_time:.4f}s")
    print(f"  LoRA cached (2nd call):       {lora_cached_time:.4f}s")
    print(f"  Cache isolation: {results['cache_isolation']}")

    results["status"] = "PASS" if isolation_ok else "FAIL"
    return results


# ---------------------------------------------------------------------------
# Phase 3: Summary
# ---------------------------------------------------------------------------
def main():
    print(f"PyTorch: {torch.__version__}")
    print(f"CUDA available: {torch.cuda.is_available()}")
    if torch.cuda.is_available():
        print(f"GPU: {torch.cuda.get_device_name(0)}")
        vram = torch.cuda.get_device_properties(0).total_memory / (1024**3)
        print(f"VRAM: {vram:.1f} GB")
    print()

    results = {}

    # Phase 1
    try:
        llm, phase1_results = phase1_multi_lora_basic()
        results["phase1"] = phase1_results
    except Exception as e:
        print(f"\nPhase 1 FAILED: {e}")
        results["phase1"] = {"status": "FAIL", "error": str(e)}
        llm = None

    # Phase 2
    if llm is not None:
        try:
            phase2_results = phase2_lmcache_test(llm)
            results["phase2"] = phase2_results
        except Exception as e:
            print(f"\nPhase 2 FAILED: {e}")
            results["phase2"] = {"status": "FAIL", "error": str(e)}
    else:
        results["phase2"] = {"status": "SKIP", "reason": "phase1 failed"}

    # Summary
    print("\n" + "=" * 60)
    print("Test Summary")
    print("=" * 60)
    summary = json.dumps(results, indent=2, default=str, ensure_ascii=False)
    print(summary)

    sys.stdout.flush()

    # 결과를 파일로 저장
    result_path = "/tmp/test_results.json"
    with open(result_path, "w") as f:
        f.write(summary)

    # 결과를 HF repo에 업로드
    try:
        from huggingface_hub import HfApi
        api = HfApi()
        api.upload_file(
            path_or_fileobj=result_path,
            path_in_repo="test_results.json",
            repo_id="umyunsang/govon-multi-lora-test",
            repo_type="space",
        )
        print("\nResults uploaded to repo as test_results.json")
    except Exception as e:
        print(f"\nFailed to upload results: {e}")

    sys.stdout.flush()

    # 간단한 HTTP 서버로 Space RUNNING 상태 유지 + 결과 제공
    print("\nTest complete. Starting HTTP server on port 7860...")
    from http.server import HTTPServer, BaseHTTPRequestHandler

    class ResultHandler(BaseHTTPRequestHandler):
        def do_GET(self):
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.end_headers()
            with open(result_path, "rb") as f:
                self.wfile.write(f.read())
        def log_message(self, format, *args):
            pass  # suppress access logs

    server = HTTPServer(("0.0.0.0", 7860), ResultHandler)
    server.serve_forever()


if __name__ == "__main__":
    main()