feat: add latency benchmarking and HTML reporting for trainer MCP tools by haroon0x · Pull Request #26 · kubeflow/mcp-server

haroon0x · 2026-05-18T06:52:53Z

First thing: this is not a vibe coded AI slop PR.

I have added the initial benchmark suite for trainer MCP tools as mentioned in #10 . This PR focuses only on latency benchmarks.

Only tests/benchmarks/report.py is AI generated. The benchmark cases, pytest setup, JSON format, and benchmark flow are written in a simple way so it is easy to review and extend later.

What is implemented

This PR adds latency benchmarks under:

tests/benchmarks/test_latency.py

The benchmarks use pytest-benchmark for timing. After pytest finishes, tests/benchmarks/conftest.py collects the benchmark samples and writes the local result files.

Run:

uv run pytest tests/benchmarks/

Generated files:

benchmark-results/latency.json
benchmark-results/index.html

benchmark-results/ is ignored because these files are generated locally.

Latency benchmarks added

Current latency suite covers:

server init in full mode
server init in progressive mode
server init in semantic mode
dynamic tool registry init
dynamic tool discovery:
- list_tools
- describe_tools
- find_tools
trainer schema and metadata scan
preview tool paths:
- fine_tune
- run_custom_training
- run_container_training
security validation:
- validate_k8s_name
- validate_resource_limits
- is_safe_python_code

Preview benchmarks use confirmed=False, so they do not submit jobs to Kubernetes.

For fine_tune, the GPU pre-check is patched inside the benchmark so the benchmark can run without needing a real cluster or GPU.

Metrics

The JSON output contains:

P50
P95
P99
min
max

P95 and P99 are included because average or P50 alone will not show tail latency clearly.

Example output shape:

{
  "suite": "latency",
  "iterations": 100,
  "warmup": 5,
  "results": [
    {
      "name": "server_init_full",
      "unit": "ms",
      "p50": 18.1,
      "p95": 24.5,
      "p99": 31.2,
      "min": 15.9,
      "max": 40.8
    }
  ]
}

HTML report

The HTML report is generated at:

benchmark-results/index.html

It shows:

summary cards
latency table
P50, P95, P99
min and max
P99/P50 tail ratio
visual spread bars

The report code is kept separate so future benchmark suites can also be added to the same report.

Future plan

Next benchmark suites can be added as separate files:

tests/benchmarks/test_token_usage.py
tests/benchmarks/test_cpu_profile.py
tests/benchmarks/test_memory.py

Rough plan:

token usage: custom estimator because it is not a timing benchmark
CPU profile: Python cProfile
memory profile: Python tracemalloc

All of them can still run through pytest and write JSON files into benchmark-results/.

Validation done

I ran:

uv run pytest tests/benchmarks/
uv run ruff check tests/benchmarks/test_latency.py tests/benchmarks/conftest.py tests/benchmarks/report.py
uv run ruff format --check tests/benchmarks/test_latency.py tests/benchmarks/conftest.py tests/benchmarks/report.py
uv run python -m tests.benchmarks.report
uv lock --check

All passed locally.

There are 3 warnings during the benchmark run. These warnings come from the installed Kubeflow dependency using old Pydantic class-based config. They are not from this benchmark code.

Type of Change

feat: New feature
fix: Bug fix
revert: Revert a change
chore: Maintenance / tooling

Checklist

Benchmark tests pass locally
Linting passes for benchmark files
Documentation updated

google-oss-prow · 2026-05-18T06:53:01Z

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by:
Once this PR has been reviewed and has the lgtm label, please assign astefanutti for approval. For more information see the Kubernetes Code Review Process.

The full list of commands accepted by this bot can be found here.

Details

Needs approval from an approver in each of these files:

OWNERS

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

abhijeet-dhumal · 2026-05-18T13:23:04Z

Thanks for this @haroon0x 🚀
scope is right for a first PR. Just few nit picks

Copilot

Pull request overview

Adds an initial latency benchmark suite for trainer MCP tools, built on pytest-benchmark, plus a small JSON aggregation layer and a self-contained HTML report generator. Output is written under benchmark-results/ and is git-ignored.

Changes:

New tests/benchmarks/ package with test_latency.py (server init, dynamic tool discovery, schema scan, preview tool paths, security validation), conftest.py (latency fixture + percentile aggregation + terminal-summary hook), and report.py (HTML report rendering).
Add pytest-benchmark>=5.2.3 dev dependency in pyproject.toml / uv.lock.
Document benchmark usage in docs/benchmarks.md and ignore benchmark-results/.

Reviewed changes

Copilot reviewed 6 out of 8 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
pyproject.toml	Adds `pytest-benchmark` to dev extras.
uv.lock	Locks `pytest-benchmark` and transitive `py-cpuinfo`.
.gitignore	Ignores generated `benchmark-results/` directory.
docs/benchmarks.md	New documentation describing how to run benchmarks, output format, and planned suites.
tests/benchmarks/test_latency.py	Defines latency benchmark cases for server init, dynamic tools, security validation, and preview tool paths.
tests/benchmarks/conftest.py	Provides `record_latency_benchmark` fixture and writes `latency.json` + triggers report generation in `pytest_terminal_summary`.
tests/benchmarks/report.py	Renders a static HTML report (summary cards, latency table, spread bars) from suite JSON files.
tests/benchmarks/benchmarks_utils.py	Empty placeholder file.

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

abhijeet-dhumal · 2026-05-25T08:00:28Z

@@ -0,0 +1 @@
+


@haroon0x Why this empty file exists ?

@haroon0x Why this empty file exists ?

Did you get time to resolve this and cleanup empty file?

haroon0x · 2026-05-19T18:14:41Z

@abhijeet-dhumal I have made the changes , can you review ?

haroon0x · 2026-05-25T18:48:38Z

@abhijeet-dhumal I have resolved the issues as you have mentioned. Could you take a look and let me know if this is ready to be merged?

After this i was thinking on working issue #5 . Do you have any specific thoughts or req before i start working on that?

Copilot

Pull request overview

Copilot reviewed 6 out of 8 changed files in this pull request and generated 6 comments.

haroon0x · 2026-06-02T03:13:08Z

+        benchmark.pedantic(
+            func,
+            rounds=DEFAULT_ITERATIONS,
+            warmup_rounds=DEFAULT_WARMUP,
+            iterations=1,
+        )
+        samples_ms = [sample * 1_000 for sample in benchmark.stats["data"]]
+        LATENCY_RESULTS.append(_latency_result(name, samples_ms))


@abhijeet-dhumal Copilot is suggesting to use benchmark.stats.data

Hmm I see.. IIUC benchmark.stats["data"].. that’s valid on pytest-benchmark 5.2

+def test_dynamic_tool_registry_init_latency(
+    record_latency_benchmark: Callable[[str, Callable[[], object]], None],
+) -> None:
+    def initialize_registry() -> None:
+        init_dynamic_tools(TOOLS, CLIENT_TOOL_DESCRIPTIONS)
+
+    record_latency_benchmark("dynamic_tool_registry_init", initialize_registry)


+def test_preview_tools_latency(
+    record_latency_benchmark: Callable[[str, Callable[[], object]], None],
+    name: str,
+    benchmark_func: Callable[[], object],
+) -> None:
+    assert hasattr(training, "_check_gpu_available"), (
+        "_check_gpu_available renamed — update benchmark patch"
+    )
+    original_gpu_check = training._check_gpu_available
+    training._check_gpu_available = lambda: None
+    try:
+        record_latency_benchmark(name, benchmark_func)
+    finally:
+        training._check_gpu_available = original_gpu_check


+def _latency_summary(payload: dict[str, Any]) -> str:
+    results = [result for result in payload.get("results", []) if _number(result.get("p50"))]
+    if not results:
+        return ""


+def generate_report(results_dir: Path = RESULTS_DIR) -> Path:
+    payloads = load_benchmark_payloads(results_dir)
+    results_dir.mkdir(exist_ok=True)
+    output_path = results_dir / "index.html"
+    output_path.write_text(_render_html(payloads))
+    return output_path


@@ -0,0 +1 @@
+


abhijeet-dhumal · 2026-07-10T12:58:03Z

Benchmark suite is valuable but this is XL scope. Please confirm: (1) benchmark-results/ stays gitignored, (2) GPU pre-check patching is acceptable for CI-less runs, (3) split latency report from token/CPU profiling into follow-up PRs if possible.

/hold

haroon0x · 2026-07-11T18:44:17Z

@abhijeet-dhumal yes all these are confirmed

abhijeet-dhumal · 2026-07-12T10:32:52Z

/ok-to-test

abhijeet-dhumal · 2026-07-12T10:33:14Z

@haroon0x can you rebase this PR?

abhijeet-dhumal · 2026-07-12T10:35:16Z

Why this empty file is needed?

When we want to add more util func , we can use it. Currently there is no use. i will remove it for now.

abhijeet-dhumal · 2026-07-12T11:47:45Z

@haroon0x can you check why pr checks are failing ? Have you rebased your PR on main branch ?

abhijeet-dhumal · 2026-07-18T11:58:23Z

@haroon0x can you fix PR title for PR checks pass

haroon0x · 2026-07-18T13:55:19Z

@abhijeet-dhumal renamed the pr title , now ready for merging , lmk what you want me to work on next.

abhijeet-dhumal · 2026-07-19T08:34:16Z

Thanks @haroon0x for this, just few nits after this Im happy to unhold once these land.

/hold

abhijeet-dhumal · 2026-07-19T08:36:42Z

Default pytest / make test-python will collect tests/benchmarks/ and run pedantic around 100 round benches in every CI job. It seems bit overkill wdyt?
If possible can we mark/exclude them and add 'make benchmark' (or equivalent) so unit CI stays fast.

yeah that is overkill , i missed it. I willl exclude benchmark run from the tests.

abhijeet-dhumal · 2026-07-19T08:38:35Z

+DEFAULT_ITERATIONS = 100
+DEFAULT_WARMUP = 5
+OUTPUT_DIR = Path(__file__).parent.parent.parent / "benchmark-results"
+LATENCY_RESULTS: list[dict[str, object]] = []


Remaining unresolved: Still needs a reset in pytest_configure so results don’t accumulate across in-process re-runs: LATENCY_RESULTS.clear() (or reassign) at session start..

abhijeet-dhumal · 2026-07-23T14:11:28Z

@haroon0x can we resolve above nits? We are very close!
I would like to add this PR in v0.1 milestone 🚀

… testing and HTML reporting Signed-off-by: haroon0x <haroonbmc0@gmail.com>

Signed-off-by: haroon0x <haroonbmc0@gmail.com>

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Haroon <106879583+haroon0x@users.noreply.github.com> Signed-off-by: haroon0x <haroonbmc0@gmail.com>

Co-authored-by: Abhijeet Dhumal <84722973+abhijeet-dhumal@users.noreply.github.com> Signed-off-by: Haroon <106879583+haroon0x@users.noreply.github.com> Signed-off-by: haroon0x <haroonbmc0@gmail.com>

Signed-off-by: haroon0x <haroonbmc0@gmail.com>

- Exclude benchmarks from default pytest runs to keep unit CI fast. - Add `make benchmark` as the explicit suite entry point. - Mark tests by directory so new benchmarks inherit the exclusion. - Reset latency samples per session to keep in-process reruns clean. - Use pytest's pathlib item API instead of the legacy fspath bridge. - Keep zero-valued P50 samples in report summaries. - Document default and explicit benchmark behavior. Signed-off-by: haroon0x <haroonbmc0@gmail.com>

haroon0x · 2026-07-23T18:00:48Z

@abhijeet-dhumal Lmk if there are any other nits as well

haroon0x · 2026-07-30T10:25:18Z

@abhijeet-dhumal are there any other changes ?

Resolve the Makefile conflict while retaining the benchmark target and upstream release tooling. Signed-off-by: haroon0x <haroonbmc0@gmail.com>

Signed-off-by: Haroon <106879583+haroon0x@users.noreply.github.com>

google-oss-prow Bot requested review from Electronic-Waste, andreyvelich and szaher May 18, 2026 06:53

google-oss-prow Bot added the size/XL label May 18, 2026

haroon0x changed the title ~~feat: initial implementation of benchmarking (latency)~~ feat: Add Custom Latency Benchmarking and HTML Reporting for Trainer MCP Tools May 18, 2026

abhijeet-dhumal reviewed May 18, 2026

View reviewed changes

Comment thread tests/benchmarks/test_latency.py Outdated

abhijeet-dhumal reviewed May 18, 2026

View reviewed changes

Comment thread tests/benchmarks/benchmarks_runner.py Outdated

Copilot AI review requested due to automatic review settings May 18, 2026 17:27

Copilot started reviewing on behalf of haroon0x May 18, 2026 17:27 View session

haroon0x requested a review from abhijeet-dhumal May 18, 2026 17:28

Copilot AI reviewed May 18, 2026

View reviewed changes

abhijeet-dhumal reviewed May 25, 2026

View reviewed changes

Comment thread tests/benchmarks/conftest.py Outdated

abhijeet-dhumal reviewed May 25, 2026

View reviewed changes

Comment thread tests/benchmarks/conftest.py

abhijeet-dhumal reviewed May 25, 2026

View reviewed changes

Comment thread tests/benchmarks/report.py Outdated

abhijeet-dhumal reviewed May 25, 2026

View reviewed changes

Comment thread tests/benchmarks/test_latency.py

haroon0x requested a review from abhijeet-dhumal May 25, 2026 18:40

haroon0x requested a review from Copilot May 29, 2026 17:36

Copilot AI reviewed May 29, 2026

View reviewed changes

google-oss-prow Bot added the do-not-merge/hold Blocked — do not merge label Jul 10, 2026

google-oss-prow Bot added the ok-to-test Approve CI for external contributors label Jul 12, 2026

abhijeet-dhumal reviewed Jul 12, 2026

View reviewed changes

haroon0x force-pushed the trainer/benchmark-suite branch from 3f4657b to 186495a Compare July 12, 2026 10:43

google-oss-prow Bot added size/XL and removed size/XXL labels Jul 12, 2026

haroon0x force-pushed the trainer/benchmark-suite branch from aa16296 to b49b42b Compare July 12, 2026 10:53

haroon0x requested a review from abhijeet-dhumal July 12, 2026 10:53

This was referenced Jul 12, 2026

test(core): expand unit coverage for core MCP modules #67

Open

test(trainer): expand unit coverage for trainer MCP tools #68

Open

Unit test suite for trainer MCP tools" #9

Open

haroon0x force-pushed the trainer/benchmark-suite branch from 3ad24cd to 9aff785 Compare July 16, 2026 05:19

haroon0x changed the title ~~feat: Add Custom Latency Benchmarking and HTML Reporting for Trainer MCP Tools~~ feat: add latency benchmarking and HTML reporting for trainer MCP tools Jul 18, 2026

abhijeet-dhumal reviewed Jul 19, 2026

View reviewed changes

haroon0x and others added 8 commits July 23, 2026 17:54

feat: initial implementation of benchmarking framework with automated…

ed81f74

… testing and HTML reporting Signed-off-by: haroon0x <haroonbmc0@gmail.com>

add more latency benchmarking and supporting documentation

61c947b

Signed-off-by: haroon0x <haroonbmc0@gmail.com>

refactor latency benchmarking suite using pytest-benchmark

ab9bce8

Signed-off-by: haroon0x <haroonbmc0@gmail.com>

remove unused code

fc93857

Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com> Signed-off-by: Haroon <106879583+haroon0x@users.noreply.github.com> Signed-off-by: haroon0x <haroonbmc0@gmail.com>

Update tests/benchmarks/report.py

671ac60

Co-authored-by: Abhijeet Dhumal <84722973+abhijeet-dhumal@users.noreply.github.com> Signed-off-by: Haroon <106879583+haroon0x@users.noreply.github.com> Signed-off-by: haroon0x <haroonbmc0@gmail.com>

resolve review issues

4a52439

Signed-off-by: haroon0x <haroonbmc0@gmail.com>

chore: remove unused benchmark utility

62fedab

Signed-off-by: haroon0x <haroonbmc0@gmail.com>

haroon0x force-pushed the trainer/benchmark-suite branch from 7df0de5 to 281f542 Compare July 23, 2026 17:56

haroon0x requested a review from abhijeet-dhumal July 23, 2026 17:56

haroon0x and others added 2 commits August 3, 2026 13:47

chore: sync benchmark branch with upstream main

cdc58b8

Resolve the Makefile conflict while retaining the benchmark target and upstream release tooling. Signed-off-by: haroon0x <haroonbmc0@gmail.com>

Merge branch 'main' into trainer/benchmark-suite

824cc76

Signed-off-by: Haroon <106879583+haroon0x@users.noreply.github.com>

Uh oh!

Conversation

haroon0x commented May 18, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What is implemented

Latency benchmarks added

Metrics

HTML report

Future plan

Validation done

Type of Change

Checklist

Uh oh!

google-oss-prow Bot commented May 18, 2026

Uh oh!

abhijeet-dhumal commented May 18, 2026

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

haroon0x commented May 19, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

haroon0x commented May 25, 2026

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

abhijeet-dhumal commented Jul 10, 2026

Uh oh!

haroon0x commented Jul 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

abhijeet-dhumal commented Jul 12, 2026

Uh oh!

abhijeet-dhumal commented Jul 12, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

abhijeet-dhumal commented Jul 12, 2026

Uh oh!

abhijeet-dhumal commented Jul 18, 2026

Uh oh!

haroon0x commented Jul 18, 2026

Uh oh!

abhijeet-dhumal commented Jul 19, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

abhijeet-dhumal commented Jul 23, 2026

Uh oh!

haroon0x commented Jul 23, 2026

Uh oh!

haroon0x commented Jul 30, 2026

Uh oh!

Reviewers

haroon0x commented May 18, 2026 •

edited

Loading

haroon0x commented Jul 11, 2026 •

edited

Loading