Skip to content

Commit 2ca23e3

Browse files
authored
feat: add client-side profiling trigger (#326)
* feat: add client-side vLLM profiling trigger Adds an optional client-side trigger that fires POST /start_profile at the performance phase start and /stop_profile at run end, so a profiled run can be driven from a YAML/CLI flag without coupling endpoints to any vendor harness. Schema: ProfilerEngine enum (currently {vllm}) and ProfilingConfig hung off Settings. URLs are auto-derived per entry in endpoint_config.endpoints (strip /v1, append engine-specific path). Default-off; warn-don't-fail throughout. Report.txt gets a Profiling section and a sibling profiling.json is written next to result_summary.json when the trigger is enabled. * fix: close unclosed Field() for metrics_tokenizer_workers The profiler-trigger commit (a4fe30b) left the Field( call for metrics_tokenizer_workers unterminated, so config/schema.py raised SyntaxError and the inference-endpoint CLI could not import. Add the missing ) so the module compiles. * feat: allow separate profiling endpoint override Add an optional profiling.endpoints (CLI --profile-endpoints) field so the profiler start/stop triggers can target a different host than the inference endpoint. When unset, URLs are still derived from endpoint_config.endpoints; when set, derivation runs over the override list using the same engine-specific protocol. Adds a scheme validator mirroring EndpointConfig and a matching --profile-endpoints override on the from-config subcommand. * refactor: drop --profile/--profile-urls overrides from from-config Keep the from-config CLI surface minimal per review feedback: profiling is configured via the YAML settings.profiling block for from-config runs. This also removes the model_copy(update=...) path that bypassed ProfilingConfig URL-scheme validation. offline/online keep the schema-generated --profile/--profile-urls flags, which validate normally. * test: cover profiling trigger config and helpers Adds unit tests for the client-side profiling trigger (review finding #3): - TestProfilingConfig (test_schema.py): defaults, engine enum coercion, and URL-scheme validation on both the direct-construction (offline/online) and model_validate (from-config YAML) paths. - TestProfilingHelpers (test_benchmark.py): _derive_profile_urls /v1 stripping and empty-endpoints ValueError, _post_profile 200/404/connection-failure via mocked urlopen, _render_profile_status, and _write_profiling_section output plus profiling.json serializability.
1 parent 8480f39 commit 2ca23e3

7 files changed

Lines changed: 407 additions & 5 deletions

File tree

src/inference_endpoint/commands/benchmark/execute.py

Lines changed: 186 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -30,13 +30,16 @@
3030
import shutil
3131
import signal
3232
import tempfile
33+
import time
3334
import uuid
3435
from collections.abc import Callable
3536
from dataclasses import dataclass, field
3637
from dataclasses import replace as dataclass_replace
3738
from datetime import datetime
3839
from pathlib import Path
39-
from typing import Any
40+
from typing import Any, TextIO
41+
from urllib import error as urllib_error
42+
from urllib import request as urllib_request
4043
from urllib.parse import urljoin
4144

4245
import msgspec
@@ -65,6 +68,7 @@
6568
DatasetType,
6669
LoadPattern,
6770
LoadPatternType,
71+
ProfilerEngine,
6872
StreamingMode,
6973
TestMode,
7074
TestType,
@@ -140,6 +144,10 @@ class BenchmarkResult:
140144
collector: ResponseCollector
141145
report: Report | None
142146
tmpfs_dir: Path
147+
# Profile trigger payload {engine: str, starts: [...], stops: [...]} when
148+
# settings.profiling.engine is set; None otherwise. Rendered into
149+
# report.txt and a sibling profiling.json by finalize_benchmark.
150+
profiling: dict[str, Any] | None = None
143151

144152

145153
@dataclass
@@ -538,6 +546,110 @@ def _load_final_snapshot_from_disk(path: Path) -> dict[str, Any] | None:
538546
return None
539547

540548

549+
# (start_path, stop_path) for each supported inference engine's profiling
550+
# protocol. Add a row when introducing a new ProfilerEngine variant.
551+
_PROFILE_PATHS: dict[ProfilerEngine, tuple[str, str]] = {
552+
ProfilerEngine.VLLM: ("/start_profile", "/stop_profile"),
553+
}
554+
555+
556+
def _derive_profile_urls(
557+
endpoints: list[str], engine: ProfilerEngine, action: str
558+
) -> list[str]:
559+
"""One profile URL per endpoint, derived from the engine's HTTP protocol.
560+
561+
For vLLM: strip a trailing ``/v1`` from each endpoint and append
562+
``/{start,stop}_profile``. ``action`` is ``"start"`` or ``"stop"``.
563+
"""
564+
if not endpoints:
565+
raise ValueError(
566+
f"profiling.engine={engine.value} but endpoint_config.endpoints "
567+
f"is empty; cannot derive {action} URLs"
568+
)
569+
start_path, stop_path = _PROFILE_PATHS[engine]
570+
path = start_path if action == "start" else stop_path
571+
urls: list[str] = []
572+
for ep in endpoints:
573+
base = ep.rstrip("/")
574+
if base.endswith("/v1"):
575+
base = base[:-3]
576+
urls.append(f"{base.rstrip('/')}{path}")
577+
return urls
578+
579+
580+
def _post_profile(url: str) -> dict[str, Any]:
581+
"""POST {url} with empty body; never raises. Returns a record dict suitable
582+
for report.txt rendering and profiling.json serialization."""
583+
record: dict[str, Any] = {
584+
"url": url,
585+
"sent_at_ns": time.monotonic_ns(),
586+
"sent_at_iso": datetime.now().isoformat(timespec="milliseconds"),
587+
"status": None,
588+
"error": None,
589+
}
590+
req = urllib_request.Request(url, method="POST", data=b"")
591+
try:
592+
with urllib_request.urlopen(req, timeout=2) as resp:
593+
record["status"] = resp.status
594+
except urllib_error.HTTPError as e:
595+
record["status"] = e.code
596+
record["error"] = f"{e.code} {e.reason}"
597+
except Exception as e: # noqa: BLE001 — profile failures must never abort a run
598+
record["error"] = f"{type(e).__name__}: {e}"
599+
return record
600+
601+
602+
def _render_profile_status(rec: dict[str, Any]) -> str:
603+
status = rec.get("status")
604+
error = rec.get("error")
605+
if status == 200:
606+
return "200 OK"
607+
if status == 404:
608+
return (
609+
"404 (profiling not enabled on server — pass "
610+
"--profiler-config.profiler=... to server)"
611+
)
612+
if error:
613+
return error
614+
if status is not None:
615+
return str(status)
616+
return "ERROR"
617+
618+
619+
def _write_profiling_section(f: TextIO, profiling: dict[str, Any]) -> None:
620+
"""Append the Profiling section to report.txt (called after report.display)."""
621+
starts = profiling.get("starts", [])
622+
stops = profiling.get("stops", [])
623+
f.write("\n------------------- Profiling -------------------\n")
624+
f.write(f"Engine: {profiling.get('engine', 'unknown')}\n")
625+
f.write("Start:\n")
626+
for rec in starts:
627+
f.write(
628+
f" POST {rec['url']} @ {rec['sent_at_iso']} → "
629+
f"{_render_profile_status(rec)}\n"
630+
)
631+
if stops:
632+
f.write("Stop:\n")
633+
for rec in stops:
634+
suffix = (
635+
" (from abort handler)" if rec.get("stop_reason") == "abort" else ""
636+
)
637+
f.write(
638+
f" POST {rec['url']} @ {rec['sent_at_iso']} → "
639+
f"{_render_profile_status(rec)}{suffix}\n"
640+
)
641+
if starts and stops:
642+
first_start = min(r["sent_at_ns"] for r in starts)
643+
last_stop = max(r["sent_at_ns"] for r in stops)
644+
f.write(f"Trigger span: {(last_stop - first_start) / 1e9:.2f} s\n")
645+
f.write(
646+
"\nNote: actual trace window is bounded by server-side "
647+
"--profiler-config.delay_iterations and "
648+
"--profiler-config.max_iterations.\n"
649+
"Trace artifact path is in server stdout.\n"
650+
)
651+
652+
541653
async def _run_benchmark_async(
542654
ctx: BenchmarkContext,
543655
loop: asyncio.AbstractEventLoop,
@@ -735,6 +847,23 @@ def _on_sample_complete(result: QueryResult) -> None:
735847
_timeout_done = False
736848
max_duration_ms = ctx.rt_settings.max_duration_ms
737849

850+
# Profile trigger state. Pre-derive URLs once so a bad config
851+
# (engine set but no endpoints) fails before the run.
852+
profiling_cfg = config.settings.profiling
853+
profile_start_urls: list[str] = []
854+
profile_stop_urls: list[str] = []
855+
profile_starts: list[dict[str, Any]] = []
856+
profile_stops: list[dict[str, Any]] = []
857+
if profiling_cfg.engine is not None:
858+
profile_endpoints = profiling_cfg.urls or config.endpoint_config.endpoints
859+
profile_start_urls = _derive_profile_urls(
860+
profile_endpoints, profiling_cfg.engine, "start"
861+
)
862+
profile_stop_urls = _derive_profile_urls(
863+
profile_endpoints, profiling_cfg.engine, "stop"
864+
)
865+
session_completed_normally = False
866+
738867
def _on_global_timeout() -> None:
739868
if not _timeout_done:
740869
logger.warning(
@@ -745,24 +874,58 @@ def _on_global_timeout() -> None:
745874

746875
def _on_phase_start(phase: PhaseConfig) -> None:
747876
nonlocal global_timeout_handle
748-
if (
749-
phase.phase_type == PhaseType.PERFORMANCE
750-
and max_duration_ms is not None
751-
):
877+
if phase.phase_type != PhaseType.PERFORMANCE:
878+
return
879+
if max_duration_ms is not None:
752880
global_timeout_handle = loop.call_later(
753881
max_duration_ms / 1000.0, _on_global_timeout
754882
)
883+
# Fire /start_profile sequentially before any perf request is
884+
# issued, so the server is armed when traffic begins. Blocks
885+
# the loop briefly (sub-100ms per URL); strategy task hasn't
886+
# been created yet so nothing is starved.
887+
for url in profile_start_urls:
888+
rec = _post_profile(url)
889+
if rec["status"] == 200:
890+
logger.info("Profile start: %s -> 200 OK", url)
891+
else:
892+
logger.warning(
893+
"Profile start: %s -> %s",
894+
url,
895+
rec["error"] or rec["status"],
896+
)
897+
profile_starts.append(rec)
755898

756899
loop.add_signal_handler(signal.SIGINT, session.stop)
757900
try:
758901
result = await session.run(phases, on_phase_start=_on_phase_start)
902+
session_completed_normally = True
759903
except Exception as e:
760904
raise ExecutionError(f"Benchmark execution failed: {e}") from e
761905
finally:
762906
_timeout_done = True
763907
if global_timeout_handle is not None:
764908
global_timeout_handle.cancel()
765909
loop.remove_signal_handler(signal.SIGINT)
910+
# Fire /stop_profile for URLs whose /start_profile succeeded.
911+
# Unifies the clean phase-end path and the abort path —
912+
# both reach this block, both fire stops.
913+
if profile_starts:
914+
stop_reason = "phase_end" if session_completed_normally else "abort"
915+
for i, start_rec in enumerate(profile_starts):
916+
if start_rec["status"] != 200 or i >= len(profile_stop_urls):
917+
continue
918+
rec = _post_profile(profile_stop_urls[i])
919+
rec["stop_reason"] = stop_reason
920+
if rec["status"] == 200:
921+
logger.info("Profile stop: %s -> 200 OK", profile_stop_urls[i])
922+
else:
923+
logger.warning(
924+
"Profile stop: %s -> %s",
925+
profile_stop_urls[i],
926+
rec["error"] or rec["status"],
927+
)
928+
profile_stops.append(rec)
766929
logger.info("Cleaning up...")
767930
try:
768931
if http_client:
@@ -824,11 +987,20 @@ def _on_phase_start(phase: PhaseConfig) -> None:
824987
metrics_subscriber.close()
825988
pbar.close()
826989

990+
profiling_payload: dict[str, Any] | None = None
991+
if profiling_cfg.engine is not None:
992+
profiling_payload = {
993+
"engine": profiling_cfg.engine.value,
994+
"starts": profile_starts,
995+
"stops": profile_stops,
996+
}
997+
827998
return BenchmarkResult(
828999
session=result,
8291000
collector=collector,
8301001
report=report,
8311002
tmpfs_dir=tmpfs_dir,
1003+
profiling=profiling_payload,
8321004
)
8331005

8341006

@@ -899,8 +1071,17 @@ def finalize_benchmark(ctx: BenchmarkContext, bench: BenchmarkResult) -> None:
8991071
report_txt = ctx.report_dir / "report.txt"
9001072
with report_txt.open("w") as f:
9011073
report.display(fn=lambda s: print(s, file=f))
1074+
if bench.profiling is not None:
1075+
_write_profiling_section(f, bench.profiling)
9021076
logger.info("Report written to %s", report_txt)
9031077

1078+
# Sibling profiling.json — kept separate so Report stays a pure
1079+
# snapshot-derived struct.
1080+
if bench.profiling is not None:
1081+
(ctx.report_dir / "profiling.json").write_text(
1082+
json.dumps(bench.profiling, indent=2)
1083+
)
1084+
9041085
# Write scoring artifacts + copy event log from tmpfs to disk
9051086
_write_scoring_artifacts(ctx, result, bench.tmpfs_dir)
9061087

src/inference_endpoint/config/schema.py

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -609,6 +609,72 @@ class DrainConfig(BaseModel):
609609
)
610610

611611

612+
class ProfilerEngine(str, Enum):
613+
"""Inference engine whose profiling protocol the client should drive.
614+
615+
Selects the HTTP path layout used to derive start/stop URLs from
616+
``endpoint_config.endpoints``. Each value corresponds to one server-side
617+
profiling protocol; add a new variant + ``_PROFILE_PATHS`` row to support
618+
another engine.
619+
"""
620+
621+
VLLM = "vllm"
622+
623+
624+
@cyclopts.Parameter(name="*")
625+
class ProfilingConfig(BaseModel):
626+
"""Client-side trigger for the server's profiler.
627+
628+
When ``engine`` is set, fires POST ``<start_path>`` at performance-phase
629+
begin and POST ``<stop_path>`` at performance-phase end. URLs are derived
630+
using the engine-specific protocol from ``urls`` when set, otherwise
631+
from ``endpoint_config.endpoints``.
632+
Server must be launched with profiling enabled (e.g. vLLM's
633+
``--profiler-config.profiler=cuda|torch``); the schedule
634+
(``delay_iterations``, ``max_iterations``) is set there, not here.
635+
"""
636+
637+
model_config = ConfigDict(extra="forbid", frozen=True)
638+
639+
engine: Annotated[
640+
ProfilerEngine | None,
641+
cyclopts.Parameter(
642+
alias="--profile",
643+
help="Profile the named inference engine around the performance phase",
644+
),
645+
] = Field(
646+
None,
647+
description="Profile the named inference engine around the performance phase",
648+
)
649+
urls: Annotated[
650+
list[str] | None,
651+
cyclopts.Parameter(
652+
alias="--profile-urls",
653+
help="Override URL(s) for profiler triggers; "
654+
"defaults to endpoint_config.endpoints",
655+
negative="",
656+
),
657+
] = Field(
658+
None,
659+
description="URL(s) the profiler start/stop triggers are derived from. "
660+
"When None, derived from endpoint_config.endpoints instead. Use when "
661+
"the profiler admin endpoint differs from the inference endpoint.",
662+
)
663+
664+
@field_validator("urls", mode="after")
665+
@classmethod
666+
def _validate_url_scheme(cls, v: list[str] | None) -> list[str] | None:
667+
if v is None:
668+
return v
669+
for url in v:
670+
if not url.startswith(("http://", "https://")):
671+
raise ValueError(
672+
f"Profiling endpoint URL must include scheme "
673+
f"(http:// or https://), got: {url!r}"
674+
)
675+
return v
676+
677+
612678
@cyclopts.Parameter(name="*")
613679
class Settings(BaseModel):
614680
"""Test settings."""
@@ -623,6 +689,7 @@ class Settings(BaseModel):
623689
description="Per-phase in-flight response drain timeout configuration",
624690
)
625691
warmup: WarmupConfig = Field(default_factory=WarmupConfig)
692+
profiling: ProfilingConfig = Field(default_factory=ProfilingConfig)
626693

627694

628695
class OfflineSettings(Settings):

src/inference_endpoint/config/templates/concurrency_template_full.yaml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,9 @@ settings:
8787
salt: false # Prepend a unique random hex salt to each warmup prompt
8888
drain: false
8989
warmup_random_seed: 42 # RNG seed for warmup scheduling and sample ordering
90+
profiling:
91+
engine: null # Profile the named inference engine around the performance phase | options: vllm
92+
urls: null # URL(s) the profiler start/stop triggers are derived from. When None, derived from endpoint_config.endpoints instead. Use when the profiler admin endpoint differs from the inference endpoint.
9093
endpoint_config:
9194
endpoints: # Endpoint URL(s). Must include scheme, e.g. 'http://host:port'.
9295
- http://localhost:8000

src/inference_endpoint/config/templates/offline_template_full.yaml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,9 @@ settings:
8787
salt: false # Prepend a unique random hex salt to each warmup prompt
8888
drain: false
8989
warmup_random_seed: 42 # RNG seed for warmup scheduling and sample ordering
90+
profiling:
91+
engine: null # Profile the named inference engine around the performance phase | options: vllm
92+
urls: null # URL(s) the profiler start/stop triggers are derived from. When None, derived from endpoint_config.endpoints instead. Use when the profiler admin endpoint differs from the inference endpoint.
9093
endpoint_config:
9194
endpoints: # Endpoint URL(s). Must include scheme, e.g. 'http://host:port'.
9295
- http://localhost:8000

src/inference_endpoint/config/templates/online_template_full.yaml

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,9 @@ settings:
8787
salt: false # Prepend a unique random hex salt to each warmup prompt
8888
drain: false
8989
warmup_random_seed: 42 # RNG seed for warmup scheduling and sample ordering
90+
profiling:
91+
engine: null # Profile the named inference engine around the performance phase | options: vllm
92+
urls: null # URL(s) the profiler start/stop triggers are derived from. When None, derived from endpoint_config.endpoints instead. Use when the profiler admin endpoint differs from the inference endpoint.
9093
endpoint_config:
9194
endpoints: # Endpoint URL(s). Must include scheme, e.g. 'http://host:port'.
9295
- http://localhost:8000

0 commit comments

Comments
 (0)