system_test.py

#!/usr/bin/env python3
"""
System test runner for quantize_compare.

Usage (env vars):
  PIPELINE_KIND=flux FP16_PATH=/path/to/fp16 FP8_PATH=/path/to/fp8 \
    python system_test.py

Usage (args):
  python system_test.py --pipeline flux --fp16_path /path --fp8_path /path
"""

from __future__ import annotations

import argparse
import os
import sys
import traceback
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Tuple


@dataclass
class CaseResult:
    name: str
    status: str  # PASS | FAIL | SKIP
    detail: str = ""


class TestRunner:
    def __init__(self) -> None:
        self.results: List[CaseResult] = []

    def run(
        self,
        name: str,
        fn: Callable[[], None],
        expect_fail: bool = False,
        skip: bool = False,
        reason: str = "",
    ) -> None:
        if skip:
            self.results.append(CaseResult(name=name, status="SKIP", detail=reason))
            print(f"[SKIP] {name}: {reason}")
            return

        try:
            fn()
            if expect_fail:
                self.results.append(
                    CaseResult(name=name, status="FAIL", detail="expected failure")
                )
                print(f"[FAIL] {name}: expected failure, but passed")
            else:
                self.results.append(CaseResult(name=name, status="PASS"))
                print(f"[PASS] {name}")
        except Exception as exc:  # noqa: BLE001
            if expect_fail:
                self.results.append(
                    CaseResult(name=name, status="PASS", detail=str(exc))
                )
                print(f"[PASS] {name} (expected failure): {exc}")
            else:
                self.results.append(
                    CaseResult(name=name, status="FAIL", detail=str(exc))
                )
                print(f"[FAIL] {name}: {exc}")
                traceback.print_exc()

    def summarize(self) -> int:
        print("\n" + "=" * 80)
        print("System Test Summary")
        print("=" * 80)
        for r in self.results:
            detail = f" - {r.detail}" if r.detail else ""
            print(f"{r.status:>4}  {r.name}{detail}")

        failed = [r for r in self.results if r.status == "FAIL"]
        skipped = [r for r in self.results if r.status == "SKIP"]
        print("-" * 80)
        print(f"Total: {len(self.results)}, Failed: {len(failed)}, Skipped: {len(skipped)}")
        return 1 if failed else 0


def _print_header(title: str) -> None:
    print("\n" + "=" * 80)
    print(title)
    print("=" * 80)


def _require(condition: bool, message: str) -> None:
    if not condition:
        raise AssertionError(message)


def _resolve_fp16_dtype(torch_module) -> Any:
    if torch_module.cuda.is_available() and torch_module.cuda.is_bf16_supported():
        return torch_module.bfloat16
    return torch_module.float16


def _resolve_dtype(torch_module, dtype_name: Optional[str]) -> Optional[Any]:
    if not dtype_name:
        return None
    if not hasattr(torch_module, dtype_name):
        raise ValueError(f"Unknown torch dtype: {dtype_name}")
    return getattr(torch_module, dtype_name)


def _load_pipelines(
    pipeline_kind: str,
    fp16_path: str,
    fp8_path: str,
    device: str,
    fp8_dtype_name: Optional[str],
) -> Tuple[Any, Any, Dict[str, Any]]:
    import torch

    fp8_dtype = _resolve_dtype(torch, fp8_dtype_name)

    if pipeline_kind == "flux":
        from diffusers import FluxPipeline

        pipe_fp16 = FluxPipeline.from_pretrained(
            fp16_path,
            torch_dtype=torch.float16,
        )
        pipe_fp8 = FluxPipeline.from_pretrained(fp8_path, torch_dtype=fp8_dtype)
        inputs = {
            "prompt": "A beautiful sunset over the ocean",
            "num_inference_steps": 4,
        }
        return pipe_fp16, pipe_fp8, inputs

    if pipeline_kind == "qwen_image":
        from diffusers import QwenImagePipeline

        fp16_dtype = _resolve_fp16_dtype(torch)
        pipe_fp16 = QwenImagePipeline.from_pretrained(
            fp16_path,
            torch_dtype=fp16_dtype,
        )
        pipe_fp8 = QwenImagePipeline.from_pretrained(fp8_path, torch_dtype=fp8_dtype)
        inputs = {
            "prompt": "A futuristic city skyline at sunset",
        }
        return pipe_fp16, pipe_fp8, inputs

    raise ValueError(f"Unknown pipeline kind: {pipeline_kind}")


def _ensure_outputs(output_dir: Path) -> None:
    _require(output_dir.exists(), f"missing output dir: {output_dir}")
    _require((output_dir / "summary.json").exists(), "summary.json not found")
    _require((output_dir / "block_metrics.json").exists(), "block_metrics.json not found")
    _require((output_dir / "report.html").exists(), "report.html not found")


def _maybe_empty_cache() -> None:
    try:
        import torch

        if torch.cuda.is_available():
            torch.cuda.empty_cache()
    except Exception:
        pass


def _format_bytes(num_bytes: int) -> str:
    units = ["B", "KiB", "MiB", "GiB", "TiB"]
    value = float(num_bytes)
    for unit in units:
        if value < 1024 or unit == units[-1]:
            return f"{value:.2f} {unit}"
        value /= 1024
    return f"{value:.2f} TiB"


def _estimate_pipeline_bytes(pipe: Any) -> int:
    import torch.nn as nn

    total = 0
    if not hasattr(pipe, "components"):
        return total
    for component in pipe.components.values():
        if not isinstance(component, nn.Module):
            continue
        for param in component.parameters(recurse=True):
            total += param.numel() * param.element_size()
        for buf in component.buffers(recurse=True):
            total += buf.numel() * buf.element_size()
    return total


def _build_dummy_pipeline(
    container_name: str,
    block_count: int,
) -> Any:
    import torch.nn as nn

    class DummyComponent(nn.Module):
        def __init__(self) -> None:
            super().__init__()
            self.register_module(
                container_name,
                nn.ModuleList([nn.Linear(2, 2) for _ in range(block_count)]),
            )

    class DummyPipeline:
        def __init__(self) -> None:
            self.components = {"transformer": DummyComponent()}

    return DummyPipeline()


def main() -> int:
    parser = argparse.ArgumentParser(description="Quantize Compare system tests")
    parser.add_argument(
        "--pipeline",
        default=os.environ.get("PIPELINE_KIND", "qwen_image"),
        choices=["flux", "qwen_image"],
        help="Pipeline kind (flux or qwen_image)",
    )
    parser.add_argument(
        "--fp16_path",
        default=os.environ.get("FP16_PATH"),
        help="FP16 model path",
    )
    parser.add_argument(
        "--fp8_path",
        default=os.environ.get("FP8_PATH"),
        help="FP8 model path",
    )
    parser.add_argument(
        "--device",
        default=os.environ.get("QC_DEVICE", "cuda"),
        help="Device for fp16_resident mode",
    )
    parser.add_argument(
        "--fp8_dtype",
        default=os.environ.get("FP8_DTYPE"),
        help="Optional torch dtype name for FP8 model",
    )
    parser.add_argument(
        "--output_root",
        default=os.environ.get("QC_OUTPUT_ROOT", "system_test_outputs"),
        help="Root directory for test outputs",
    )
    parser.add_argument(
        "--skip_pytest",
        action="store_true",
        default=os.environ.get("SKIP_PYTEST", "0") == "1",
        help="Skip pytest tests",
    )
    parser.add_argument(
        "--level",
        choices=["quick", "full"],
        default=os.environ.get("QC_TEST_LEVEL", "full"),
        help="Test level",
    )
    args = parser.parse_args()

    root = Path(__file__).resolve().parent
    output_root = (root / args.output_root).resolve()
    output_root.mkdir(parents=True, exist_ok=True)

    runner = TestRunner()

    _print_header("Environment")

    def case_env() -> None:
        import torch
        import diffusers
        import qc

        print(f"Python: {sys.version.split()[0]}")
        print(f"torch: {torch.__version__}")
        print(f"diffusers: {diffusers.__version__}")
        print(f"qc: {getattr(qc, '__version__', 'unknown')}")
        print(f"cuda available: {torch.cuda.is_available()}")
        if torch.cuda.is_available():
            print(f"cuda device: {torch.cuda.get_device_name(0)}")

    runner.run("env_info", case_env)

    if not args.skip_pytest:
        _print_header("Unit Tests (pytest)")

        def case_pytest() -> None:
            import subprocess

            subprocess.run(
                [sys.executable, "-m", "pytest", "tests"],
                cwd=str(root),
                check=True,
            )

        runner.run("pytest", case_pytest)
    else:
        runner.run("pytest", lambda: None, skip=True, reason="SKIP_PYTEST=1")

    _print_header("Load Pipelines")

    if not args.fp16_path or not args.fp8_path:
        print("ERROR: FP16_PATH/FP8_PATH not set. Cannot run pipeline tests.")
        return 1

    try:
        pipe_fp16, pipe_fp8, base_inputs = _load_pipelines(
            args.pipeline,
            args.fp16_path,
            args.fp8_path,
            args.device,
            args.fp8_dtype,
        )
    except Exception as exc:
        print(f"ERROR: failed to load pipelines: {exc}")
        traceback.print_exc()
        return 1

    from qc import compare, list_targets, DeviceConfig

    targets_info = list_targets(pipe_fp16)
    block_targets = [t for t in targets_info if t["has_blocks"]]
    _require(block_targets, "No components with blocks found")

    fast_target = min(block_targets, key=lambda t: t["param_count"])["name"]
    all_targets = [t["name"] for t in block_targets]

    print(f"Fast target: {fast_target}")
    print(f"All targets: {all_targets}")

    use_defaults = args.pipeline == "qwen_image" and args.level == "full"
    default_targets: Optional[Any] = None if use_defaults else [fast_target]
    default_sample_steps: Optional[List[int]] = None if use_defaults else [0]
    print(f"Using pipeline defaults: {use_defaults}")

    fast_inputs = dict(base_inputs)
    if args.pipeline == "qwen_image":
        # Speed up non-baseline cases while keeping baseline defaults.
        fast_inputs.update({"num_inference_steps": 4, "height": 512, "width": 512})
    else:
        fast_inputs.update({"num_inference_steps": 4})

    def _run_compare(
        case_name: str,
        kwargs: Dict[str, Any],
        expect_success: bool = True,
        check_outputs: bool = True,
    ) -> None:
        output_dir = output_root / case_name
        kwargs = dict(kwargs)
        for key in ("targets", "sample_steps", "kurtosis_mode"):
            if key in kwargs and kwargs[key] is None:
                kwargs.pop(key)
        kwargs.setdefault("output_dir", str(output_dir))
        report = compare(**kwargs)
        if expect_success:
            _require(report.is_success(), f"compare failed: {report.error}")
            if check_outputs:
                _ensure_outputs(output_dir)
        else:
            _require(not report.is_success(), "expected failure but got success")

    _print_header("Pipeline Tests")

    def _fp16_resident_skip_reason() -> Optional[str]:
        import gc
        import torch

        if not torch.cuda.is_available():
            return "CUDA not available"

        for pipe in (pipe_fp16, pipe_fp8):
            if hasattr(pipe, "remove_all_hooks"):
                try:
                    pipe.remove_all_hooks()
                except Exception:
                    pass
            try:
                pipe.to("cpu", silence_dtype_warnings=True)
            except TypeError:
                try:
                    pipe.to("cpu")
                except Exception:
                    pass

        gc.collect()
        torch.cuda.empty_cache()

        est_bytes = _estimate_pipeline_bytes(pipe_fp16)
        free_bytes, total_bytes = torch.cuda.mem_get_info()
        headroom = 1.1

        if est_bytes == 0:
            return "could not estimate pipeline size"
        if est_bytes * headroom > total_bytes:
            return (
                f"estimated FP16 size {_format_bytes(est_bytes)} exceeds total GPU memory "
                f"{_format_bytes(total_bytes)}"
            )
        if est_bytes * headroom > free_bytes:
            return (
                f"insufficient free GPU memory {_format_bytes(free_bytes)} "
                f"(need ~{_format_bytes(int(est_bytes * headroom))})"
            )
        return None

    runner.run(
        "compare_per_tensor",
        lambda: _run_compare(
            "compare_per_tensor",
            {
                "pipe_fp16": pipe_fp16,
                "pipe_fp8": pipe_fp8,
                "inputs": dict(base_inputs),
                "targets": default_targets,
                "device_config": DeviceConfig.both_offload(device=args.device),
                "analyze_kurtosis": True,
                "sample_steps": default_sample_steps,
                "seed": 42,
                "verbose": True,
            },
        ),
    )
    _maybe_empty_cache()

    runner.run(
        "compare_per_channel",
        lambda: _run_compare(
            "compare_per_channel",
            {
                "pipe_fp16": pipe_fp16,
                "pipe_fp8": pipe_fp8,
                "inputs": dict(base_inputs),
                "targets": default_targets,
                "device_config": DeviceConfig.both_offload(device=args.device),
                "analyze_kurtosis": True,
                "kurtosis_mode": "per_channel",
                "sample_steps": default_sample_steps,
                "seed": 42,
                "verbose": True,
            },
        ),
    )
    _maybe_empty_cache()

    import torch

    fp16_resident_skip = _fp16_resident_skip_reason()
    runner.run(
        "compare_fp16_resident",
        lambda: _run_compare(
            "compare_fp16_resident",
            {
                "pipe_fp16": pipe_fp16,
                "pipe_fp8": pipe_fp8,
                "inputs": dict(base_inputs),
                "targets": default_targets,
                "device_config": DeviceConfig.fp16_resident(device=args.device),
                "analyze_kurtosis": False,
                "sample_steps": default_sample_steps,
                "seed": 42,
                "verbose": True,
            },
        ),
        skip=fp16_resident_skip is not None,
        reason=fp16_resident_skip or "",
    )
    _maybe_empty_cache()

    runner.run(
        "compare_targets_all",
        lambda: _run_compare(
            "compare_targets_all",
            {
                "pipe_fp16": pipe_fp16,
                "pipe_fp8": pipe_fp8,
                "inputs": dict(fast_inputs),
                "targets": "all",
                "device_config": DeviceConfig.both_offload(device=args.device),
                "analyze_kurtosis": False,
                "sample_steps": [0],
                "seed": 42,
                "verbose": True,
            },
        ),
        skip=(args.level == "quick"),
        reason="quick mode",
    )
    _maybe_empty_cache()

    def case_metrics_custom() -> None:
        output_dir = output_root / "compare_metrics_custom"
        report = compare(
            pipe_fp16=pipe_fp16,
            pipe_fp8=pipe_fp8,
            inputs=dict(fast_inputs),
            targets=[fast_target],
            device_config=DeviceConfig.both_offload(device=args.device),
            analyze_kurtosis=False,
            metrics=["rel_l2", "cosine", "sqnr", "mse", "mae", "max_abs"],
            sample_steps=[0],
            seed=42,
            output_dir=str(output_dir),
            verbose=True,
        )
        _require(report.is_success(), f"compare failed: {report.error}")
        _ensure_outputs(output_dir)
        first_summary = next(iter(report.block_summary.values()))
        for key in ["rel_l2", "cosine", "sqnr", "mse", "mae", "max_abs"]:
            _require(key in first_summary, f"missing metric {key}")

    runner.run("compare_metrics_custom", case_metrics_custom)
    _maybe_empty_cache()

    def case_sample_steps_filter() -> None:
        output_dir = output_root / "compare_sample_steps"
        report = compare(
            pipe_fp16=pipe_fp16,
            pipe_fp8=pipe_fp8,
            inputs=dict(fast_inputs),
            targets=[fast_target],
            device_config=DeviceConfig.both_offload(device=args.device),
            analyze_kurtosis=False,
            sample_steps=[0, 2],
            seed=42,
            output_dir=str(output_dir),
            verbose=True,
        )
        _require(report.is_success(), f"compare failed: {report.error}")
        _ensure_outputs(output_dir)
        for metrics_list in report.block_results.values():
            for entry in metrics_list:
                _require(entry["step_idx"] in (0, 2), "unexpected step_idx")

    runner.run("compare_sample_steps_filter", case_sample_steps_filter)
    _maybe_empty_cache()

    def case_invalid_target() -> None:
        report = compare(
            pipe_fp16=pipe_fp16,
            pipe_fp8=pipe_fp8,
            inputs=dict(base_inputs),
            targets=["__not_exists__"],
            device_config=DeviceConfig.both_offload(device=args.device),
            analyze_kurtosis=False,
            seed=42,
            verbose=True,
        )
        _require(not report.is_success(), "expected failure for invalid target")

    runner.run("invalid_target", case_invalid_target)

    def case_invalid_kurtosis_mode() -> None:
        report = compare(
            pipe_fp16=pipe_fp16,
            pipe_fp8=pipe_fp8,
            inputs=dict(base_inputs),
            targets=[fast_target],
            device_config=DeviceConfig.both_offload(device=args.device),
            analyze_kurtosis=True,
            kurtosis_mode="bad_mode",
            seed=42,
            verbose=True,
        )
        _require(not report.is_success(), "expected failure for invalid kurtosis_mode")

    runner.run("invalid_kurtosis_mode", case_invalid_kurtosis_mode)

    def case_alignment_mismatch() -> None:
        pipe_a = _build_dummy_pipeline("blocks", 2)
        pipe_b = _build_dummy_pipeline("blocks", 3)
        report = compare(
            pipe_fp16=pipe_a,
            pipe_fp8=pipe_b,
            inputs={},
            targets=["transformer"],
            analyze_kurtosis=False,
            verbose=False,
        )
        _require(not report.is_success(), "expected alignment failure")

    runner.run("alignment_mismatch", case_alignment_mismatch)

    def case_no_blocks() -> None:
        pipe_a = _build_dummy_pipeline("blocks", 1)
        pipe_b = _build_dummy_pipeline("blocks", 1)
        report = compare(
            pipe_fp16=pipe_a,
            pipe_fp8=pipe_b,
            inputs={},
            targets=["transformer"],
            analyze_kurtosis=False,
            verbose=False,
        )
        _require(not report.is_success(), "expected no-blocks failure")

    runner.run("no_blocks_component", case_no_blocks)

    return runner.summarize()


if __name__ == "__main__":
    raise SystemExit(main())