__init__.py

# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Benchmark helpers for kv-prune / compactor kernels.

Upstream snapshot (``compactor-vllm/src/compactor_vllm/benchmark``) contained **only**
an empty ``__init__.py`` — no additional ``.py`` scripts. Those files are merged here
as-is; there is nothing else to list under that directory in upstream.

Use :data:`BENCHMARK_REGISTRY` to register microbenchmarks or CLI entrypoints you
add under ``vllm.kvprune.benchmark``.
"""

from __future__ import annotations

from typing import Any, Callable

# Files copied from upstream ``compactor_vllm/benchmark/`` (relative to that dir).
UPSTREAM_BENCHMARK_FILES: tuple[str, ...] = ("__init__.py",)

# Optional: name -> benchmark callable or import path string (e.g. "mymod:main").
# Populated when you add real benchmarks beside this package.
BENCHMARK_REGISTRY: dict[str, Callable[..., Any] | str] = {}


def list_upstream_benchmark_files() -> tuple[str, ...]:
    """Return the list of filenames that existed in upstream ``benchmark/``."""
    return UPSTREAM_BENCHMARK_FILES


def register_benchmark(name: str, target: Callable[..., Any] | str) -> None:
    """Register a benchmark by name (callable or ``"module:attr"`` import path)."""
    BENCHMARK_REGISTRY[name] = target


def iter_registered_benchmarks() -> list[tuple[str, Callable[..., Any] | str]]:
    """Return ``(name, target)`` pairs from :data:`BENCHMARK_REGISTRY`."""
    return list(BENCHMARK_REGISTRY.items())


__all__ = [
    "BENCHMARK_REGISTRY",
    "UPSTREAM_BENCHMARK_FILES",
    "iter_registered_benchmarks",
    "list_upstream_benchmark_files",
    "register_benchmark",
]