443 lines
15 KiB
Python
443 lines
15 KiB
Python
"""Run isolated, billable cold/warm Modal LLM benchmarks."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
from dataclasses import asdict, dataclass
|
|
from datetime import datetime, timezone
|
|
import hashlib
|
|
import importlib
|
|
import json
|
|
import logging
|
|
import os
|
|
from pathlib import Path
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import time
|
|
from types import ModuleType
|
|
from typing import Any, Sequence
|
|
import uuid
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
DEFAULT_MODEL = "Blackfrost-AI/Qwen3.8-27B-ABLITERATED-NVFP4"
|
|
DEFAULT_PROMPT = (
|
|
"Explain how persistent weight and compilation caches improve repeated LLM "
|
|
"inference. Use approximately 120 words and finish with one practical "
|
|
"recommendation."
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BenchmarkConfig:
|
|
"""Describe one reproducible benchmark workload and deployment profile."""
|
|
|
|
app_name: str
|
|
cold_cycles: int
|
|
comfyui_root: str
|
|
enable_reasoning: bool
|
|
gpu: str
|
|
keep_app: bool
|
|
max_new_tokens: int
|
|
mode: str
|
|
model: str
|
|
prompt: str
|
|
seed: int
|
|
synthetic_image: bool
|
|
warm_runs: int
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class BenchmarkInvocation:
|
|
"""Record one cold-container or resident-engine inference result."""
|
|
|
|
cycle: int
|
|
kind: str
|
|
run: int
|
|
wall_seconds: float
|
|
response_characters: int
|
|
response_sha256: str
|
|
reasoning_characters: int
|
|
metadata: dict[str, Any]
|
|
|
|
|
|
def _repository_root() -> Path:
|
|
"""Return the source checkout containing this benchmark script."""
|
|
return Path(__file__).resolve().parents[1]
|
|
|
|
|
|
def _git_value(*arguments: str) -> str:
|
|
"""Return one best-effort Git value for benchmark provenance."""
|
|
completed = subprocess.run(
|
|
["git", *arguments],
|
|
cwd=_repository_root(),
|
|
check=False,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
return completed.stdout.strip() if completed.returncode == 0 else ""
|
|
|
|
|
|
def _load_remote_module() -> Any:
|
|
"""Load the source package without importing the ComfyUI extension entrypoint."""
|
|
package_name = f"comfyui_modal_benchmark_{uuid.uuid4().hex}"
|
|
package = ModuleType(package_name)
|
|
package.__path__ = [str(_repository_root())]
|
|
package.__package__ = package_name
|
|
sys.modules[package_name] = package
|
|
return importlib.import_module(f"{package_name}.remote.modal_app")
|
|
|
|
|
|
def _configure_environment(config: BenchmarkConfig) -> None:
|
|
"""Set isolated deployment controls before importing Modal-Sync modules."""
|
|
environment = {
|
|
"COMFY_MODAL_APP_NAME": config.app_name,
|
|
"COMFY_MODAL_ALLOW_EPHEMERAL_FALLBACK": "false",
|
|
"COMFY_MODAL_AUTO_DEPLOY": "true",
|
|
"COMFY_MODAL_COMFYUI_ROOT": config.comfyui_root,
|
|
"COMFY_MODAL_EXECUTION_MODE": "remote",
|
|
"COMFY_MODAL_GPU": config.gpu,
|
|
"COMFY_MODAL_LLM_MAX_RESIDENT_MODELS": "1",
|
|
"COMFY_MODAL_LLM_VLLM_EXECUTION_MODE": config.mode,
|
|
"COMFY_MODAL_MAX_CONTAINERS": "1",
|
|
"COMFY_MODAL_SCALEDOWN_WINDOW": "600",
|
|
"COMFY_MODAL_SYNC_CUSTOM_NODES": "false",
|
|
}
|
|
os.environ.update(environment)
|
|
|
|
|
|
def _synthetic_image(enabled: bool) -> Any | None:
|
|
"""Return a deterministic image tensor when multimodal testing is enabled."""
|
|
if not enabled:
|
|
return None
|
|
import torch
|
|
|
|
image = torch.zeros((1, 64, 64, 3), dtype=torch.float32)
|
|
image[:, 8:56, 8:56, 0] = 1.0
|
|
image[:, 20:44, 20:44, 1] = 0.5
|
|
return image
|
|
|
|
|
|
def _payload(
|
|
config: BenchmarkConfig,
|
|
cycle: int,
|
|
kind: str,
|
|
run: int,
|
|
) -> dict[str, Any]:
|
|
"""Build one direct Modal LLM payload with deterministic worker affinity."""
|
|
prompt_id = f"llm-benchmark-{config.mode}-{cycle}-{kind}-{run}-{uuid.uuid4().hex}"
|
|
return {
|
|
"class_type": "ModalLLM",
|
|
"component_id": "modal-llm-benchmark",
|
|
"modal_gpu": config.gpu,
|
|
"prompt_id": prompt_id,
|
|
"requires_volume_reload": False,
|
|
"subgraph_prompt": {
|
|
"modal-llm-benchmark": {
|
|
"class_type": "ModalLLM",
|
|
"inputs": {"model_profile": config.model},
|
|
}
|
|
},
|
|
"terminate_container_on_error": True,
|
|
}
|
|
|
|
|
|
def _node_inputs(config: BenchmarkConfig) -> dict[str, Any]:
|
|
"""Return deterministic node inputs shared by every benchmark invocation."""
|
|
inputs: dict[str, Any] = {
|
|
"enable_reasoning": config.enable_reasoning,
|
|
"keep_model_loaded": True,
|
|
"max_new_tokens": config.max_new_tokens,
|
|
"model_profile": config.model,
|
|
"prompt": config.prompt,
|
|
"reserve_free_vram_gb": 24.0,
|
|
"seed": config.seed,
|
|
"system_prompt": "Answer directly and follow the requested length.",
|
|
"temperature": 0.0,
|
|
"top_p": 1.0,
|
|
"video_frames": 1,
|
|
}
|
|
image = _synthetic_image(config.synthetic_image)
|
|
if image is not None:
|
|
inputs["images"] = image
|
|
return inputs
|
|
|
|
|
|
def _invoke(
|
|
remote_module: Any,
|
|
config: BenchmarkConfig,
|
|
*,
|
|
cycle: int,
|
|
kind: str,
|
|
run: int,
|
|
) -> BenchmarkInvocation:
|
|
"""Invoke one direct node request and retain timing plus backend telemetry."""
|
|
started_at = time.perf_counter()
|
|
response = remote_module.invoke_remote_engine(
|
|
_payload(config, cycle, kind, run),
|
|
remote_module.serialize_node_inputs(_node_inputs(config)),
|
|
allow_implicit_mapping=False,
|
|
)
|
|
wall_seconds = time.perf_counter() - started_at
|
|
outputs = remote_module.deserialize_node_outputs(response)
|
|
if len(outputs) < 3:
|
|
raise RuntimeError(f"Modal LLM benchmark returned {len(outputs)} outputs.")
|
|
response_text = str(outputs[0])
|
|
metadata = json.loads(str(outputs[1]))
|
|
reasoning = str(outputs[2])
|
|
_validate_invocation_metadata(config, kind, metadata)
|
|
return BenchmarkInvocation(
|
|
cycle=cycle,
|
|
kind=kind,
|
|
run=run,
|
|
wall_seconds=wall_seconds,
|
|
response_characters=len(response_text),
|
|
response_sha256=hashlib.sha256(response_text.encode("utf-8")).hexdigest(),
|
|
reasoning_characters=len(reasoning),
|
|
metadata=metadata,
|
|
)
|
|
|
|
|
|
def _validate_invocation_metadata(
|
|
config: BenchmarkConfig,
|
|
kind: str,
|
|
metadata: dict[str, Any],
|
|
) -> None:
|
|
"""Reject benchmark samples that do not have the requested engine state."""
|
|
observed_mode = metadata.get("vllm_execution_mode")
|
|
if observed_mode != config.mode:
|
|
raise RuntimeError(
|
|
f"Expected vLLM mode {config.mode!r}, observed {observed_mode!r}."
|
|
)
|
|
expected_cache_hit = kind == "warm"
|
|
observed_cache_hit = metadata.get("cache_hit")
|
|
if observed_cache_hit is not expected_cache_hit:
|
|
raise RuntimeError(
|
|
f"Expected {kind!r} cache_hit={expected_cache_hit}, "
|
|
f"observed {observed_cache_hit!r}."
|
|
)
|
|
|
|
|
|
def _reset_deployment_state(remote_module: Any, app_name: str) -> None:
|
|
"""Stop only the dedicated benchmark app and clear process-local lookup state."""
|
|
if not remote_module._stop_modal_app_via_sdk(app_name):
|
|
if not remote_module._stop_modal_app_via_cli(app_name):
|
|
raise RuntimeError(f"Unable to stop dedicated benchmark app {app_name!r}.")
|
|
_wait_for_deployment_stopped(app_name)
|
|
with remote_module._MODAL_AUTO_DEPLOY_LOCK:
|
|
remote_module._MODAL_AUTO_DEPLOY_STATES.clear()
|
|
remote_module._MODAL_REMOTE_APP_VERSION_OK.clear()
|
|
|
|
|
|
def _has_active_deployment(rows: Sequence[dict[str, Any]], app_name: str) -> bool:
|
|
"""Return whether Modal still reports live tasks for this exact deployment."""
|
|
return any(
|
|
str(row.get("Description")) == app_name
|
|
and (
|
|
str(row.get("State", "")).lower() != "stopped"
|
|
or str(row.get("Tasks", "0")) != "0"
|
|
)
|
|
for row in rows
|
|
)
|
|
|
|
|
|
def _wait_for_deployment_stopped(app_name: str, timeout_seconds: float = 120.0) -> None:
|
|
"""Wait until Modal confirms the old deployment has no live containers."""
|
|
modal_cli = shutil.which("modal")
|
|
if modal_cli is None:
|
|
raise RuntimeError("The Modal CLI is required to verify cold benchmark cycles.")
|
|
deadline = time.monotonic() + timeout_seconds
|
|
while True:
|
|
completed = subprocess.run(
|
|
[modal_cli, "app", "list", "--json"],
|
|
check=False,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
if completed.returncode != 0:
|
|
raise RuntimeError(
|
|
f"Unable to list Modal apps while stopping {app_name!r}: "
|
|
f"{completed.stderr.strip()}"
|
|
)
|
|
rows = json.loads(completed.stdout)
|
|
if not _has_active_deployment(rows, app_name):
|
|
return
|
|
if time.monotonic() >= deadline:
|
|
raise TimeoutError(
|
|
f"Modal deployment {app_name!r} still has live tasks after "
|
|
f"{timeout_seconds:.0f} seconds."
|
|
)
|
|
time.sleep(1.0)
|
|
|
|
|
|
def _deployment_app_name(remote_module: Any) -> str:
|
|
"""Resolve the GPU-specific deployed app that owns benchmark containers."""
|
|
return str(
|
|
remote_module.modal_deployment_app_name(remote_module.get_settings())
|
|
)
|
|
|
|
|
|
def _report(
|
|
config: BenchmarkConfig,
|
|
runs: Sequence[BenchmarkInvocation],
|
|
) -> dict[str, Any]:
|
|
"""Build the stable JSON benchmark artifact."""
|
|
return {
|
|
"benchmark_schema_version": 1,
|
|
"completed_at": datetime.now(timezone.utc).isoformat(),
|
|
"config": asdict(config),
|
|
"runs": [asdict(run) for run in runs],
|
|
"source": {
|
|
"commit": _git_value("rev-parse", "HEAD"),
|
|
"dirty": bool(_git_value("status", "--short")),
|
|
},
|
|
}
|
|
|
|
|
|
def _summarize_report(report: dict[str, Any]) -> str:
|
|
"""Render compact comparable cold/warm rows from one report."""
|
|
rows = ["cycle kind wall_s load_s gen_s ttft_s tok_s output_tokens cache_hit"]
|
|
for run in report["runs"]:
|
|
metadata = run["metadata"]
|
|
rows.append(
|
|
" ".join(
|
|
[
|
|
str(run["cycle"]),
|
|
str(run["kind"]),
|
|
f"{run['wall_seconds']:.3f}",
|
|
f"{float(metadata['load_seconds']):.3f}",
|
|
f"{float(metadata['generation_seconds']):.3f}",
|
|
f"{float(metadata['time_to_first_token_seconds']):.3f}",
|
|
f"{float(metadata['tokens_per_second']):.3f}",
|
|
str(metadata["output_tokens"]),
|
|
str(metadata["cache_hit"]).lower(),
|
|
]
|
|
)
|
|
)
|
|
return "\n".join(rows)
|
|
|
|
|
|
def run_benchmark(config: BenchmarkConfig) -> dict[str, Any]:
|
|
"""Run cold-container cycles followed by immediate resident-engine repeats."""
|
|
_configure_environment(config)
|
|
remote_module = _load_remote_module()
|
|
deployment_app_name = _deployment_app_name(remote_module)
|
|
runs: list[BenchmarkInvocation] = []
|
|
try:
|
|
for cycle in range(1, config.cold_cycles + 1):
|
|
_reset_deployment_state(remote_module, deployment_app_name)
|
|
runs.append(
|
|
_invoke(
|
|
remote_module,
|
|
config,
|
|
cycle=cycle,
|
|
kind="cold",
|
|
run=1,
|
|
)
|
|
)
|
|
for warm_run in range(1, config.warm_runs + 1):
|
|
runs.append(
|
|
_invoke(
|
|
remote_module,
|
|
config,
|
|
cycle=cycle,
|
|
kind="warm",
|
|
run=warm_run,
|
|
)
|
|
)
|
|
finally:
|
|
if not config.keep_app:
|
|
_reset_deployment_state(remote_module, deployment_app_name)
|
|
return _report(config, runs)
|
|
|
|
|
|
def _default_comfyui_root() -> Path:
|
|
"""Return the live sibling checkout when present, otherwise the canonical source."""
|
|
sibling = _repository_root().parent / "Latest_ComfyUI"
|
|
if sibling.exists():
|
|
return sibling
|
|
return Path.home() / "git" / "ComfyUI"
|
|
|
|
|
|
def _default_app_name(gpu: str, mode: str) -> str:
|
|
"""Return a short deterministic app name safe for all derived Modal objects."""
|
|
gpu_digest = hashlib.sha256(gpu.encode("utf-8")).hexdigest()[:8]
|
|
mode_slug = "tp" if mode == "throughput" else "eg"
|
|
return f"cm-llm-bench-{gpu_digest}-{mode_slug}"
|
|
|
|
|
|
def _validate_app_name(app_name: str) -> str:
|
|
"""Reject names that cannot accommodate Modal's derived Dict suffixes."""
|
|
if len(f"{app_name}-session-bridges") >= 64:
|
|
raise ValueError(
|
|
"--app-name is too long; it must leave room for Modal's "
|
|
"'-session-bridges' suffix."
|
|
)
|
|
return app_name
|
|
|
|
|
|
def _arguments() -> argparse.Namespace:
|
|
"""Parse benchmark controls from the command line."""
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
|
parser.add_argument("--mode", choices=("eager", "throughput"), required=True)
|
|
parser.add_argument("--gpu", default="RTX-PRO-6000")
|
|
parser.add_argument("--model", default=DEFAULT_MODEL)
|
|
parser.add_argument("--prompt", default=DEFAULT_PROMPT)
|
|
parser.add_argument("--max-new-tokens", type=int, default=256)
|
|
parser.add_argument("--seed", type=int, default=0)
|
|
parser.add_argument("--cold-cycles", type=int, default=2)
|
|
parser.add_argument("--warm-runs", type=int, default=1)
|
|
parser.add_argument("--enable-reasoning", action="store_true")
|
|
parser.add_argument("--text-only", action="store_true")
|
|
parser.add_argument("--keep-app", action="store_true")
|
|
parser.add_argument("--comfyui-root", type=Path, default=_default_comfyui_root())
|
|
parser.add_argument("--app-name")
|
|
parser.add_argument("--output", type=Path)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
"""Execute the benchmark and write its machine-readable artifact."""
|
|
arguments = _arguments()
|
|
if arguments.cold_cycles <= 0 or arguments.warm_runs <= 0:
|
|
raise ValueError("--cold-cycles and --warm-runs must be positive.")
|
|
if arguments.max_new_tokens <= 0:
|
|
raise ValueError("--max-new-tokens must be positive.")
|
|
gpu = str(arguments.gpu).strip()
|
|
gpu_slug = gpu.lower().replace("_", "-")
|
|
app_name = _validate_app_name(
|
|
arguments.app_name or _default_app_name(gpu, arguments.mode)
|
|
)
|
|
output_path = arguments.output or Path(
|
|
f"modal-llm-benchmark-{gpu_slug}-{arguments.mode}.json"
|
|
)
|
|
config = BenchmarkConfig(
|
|
app_name=app_name,
|
|
cold_cycles=arguments.cold_cycles,
|
|
comfyui_root=str(arguments.comfyui_root.expanduser().resolve()),
|
|
enable_reasoning=bool(arguments.enable_reasoning),
|
|
gpu=gpu,
|
|
keep_app=bool(arguments.keep_app),
|
|
max_new_tokens=arguments.max_new_tokens,
|
|
mode=arguments.mode,
|
|
model=arguments.model,
|
|
prompt=arguments.prompt,
|
|
seed=arguments.seed,
|
|
synthetic_image=not arguments.text_only,
|
|
warm_runs=arguments.warm_runs,
|
|
)
|
|
report = run_benchmark(config)
|
|
output_path.parent.mkdir(parents=True, exist_ok=True)
|
|
output_path.write_text(
|
|
json.dumps(report, indent=2, sort_keys=True) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
print(_summarize_report(report))
|
|
print(f"artifact {output_path.resolve()}")
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
raise SystemExit(main())
|