diff --git a/src/llmflux/cli.py b/src/llmflux/cli.py index e5175ab..4f80baf 100644 --- a/src/llmflux/cli.py +++ b/src/llmflux/cli.py @@ -11,12 +11,21 @@ import subprocess import logging from typing import Optional, List, Dict +from importlib.metadata import PackageNotFoundError, version as pkg_version from .slurm.runner import SlurmRunner from .processors import BatchProcessor from .core.config import Config, EngineConfig from .benchmark_utils import create_test_prompts_file + +def _get_llmflux_version() -> str: + """Return installed llmflux version (best effort).""" + try: + return pkg_version("llmflux") + except PackageNotFoundError: + return "unknown" + def _parse_sbatch_args(sbatch_arg_list: Optional[List[str]]) -> Optional[Dict[str, str]]: """Parse --sbatch-arg arguments into a dictionary. @@ -296,6 +305,14 @@ def _run_command(args: argparse.Namespace) -> int: def build_parser() -> argparse.ArgumentParser: parser = argparse.ArgumentParser(prog="llmflux", description="LLMFlux CLI") + parser.add_argument( + "--version", + "-V", + "-v", + action="version", + version=f"%(prog)s {_get_llmflux_version()}", + help="Show llmflux version and exit", + ) subparsers = parser.add_subparsers(dest="command", required=True) # run subcommand diff --git a/tests/test_cli_version.py b/tests/test_cli_version.py new file mode 100644 index 0000000..88c2337 --- /dev/null +++ b/tests/test_cli_version.py @@ -0,0 +1,24 @@ +import pytest + + +def test_llmflux_version_flag_prints_and_exits(capsys): + from llmflux.cli import main + + with pytest.raises(SystemExit) as excinfo: + main(["--version"]) + + assert excinfo.value.code == 0 + out = capsys.readouterr().out.strip() + assert out.startswith("llmflux ") + + +def test_llmflux_short_version_flag_prints_and_exits(capsys): + from llmflux.cli import main + + with pytest.raises(SystemExit) as excinfo: + main(["-V"]) + + assert excinfo.value.code == 0 + out = capsys.readouterr().out.strip() + assert out.startswith("llmflux ") +