381 lines
16 KiB
Python
381 lines
16 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import argparse
|
|
import ast
|
|
import os
|
|
from typing import Any, Dict
|
|
|
|
import yaml
|
|
|
|
from benchmarks.profiler.utils.planner_utils import add_planner_arguments_to_parser
|
|
from benchmarks.profiler.utils.search_space_autogen import auto_generate_search_space
|
|
|
|
|
|
def _get(cfg: Dict[str, Any], camel: str, snake: str, default: Any = None) -> Any:
|
|
"""Get config value with camelCase preferred, snake_case fallback."""
|
|
if camel in cfg:
|
|
return cfg[camel]
|
|
return cfg.get(snake, default)
|
|
|
|
|
|
def _camel_to_snake(name: str) -> str:
|
|
"""Convert camelCase to snake_case."""
|
|
import re
|
|
|
|
# Insert underscore before uppercase letters and lowercase
|
|
s1 = re.sub("(.)([A-Z][a-z]+)", r"\1_\2", name)
|
|
return re.sub("([a-z0-9])([A-Z])", r"\1_\2", s1).lower()
|
|
|
|
|
|
def parse_config_string(config_str: str) -> Dict[str, Any]:
|
|
"""Parse configuration string as Python dict literal, YAML, or JSON.
|
|
|
|
Supports multiple input formats:
|
|
1. Python dict literal: "{'engine': {'backend': 'vllm'}, 'sla': {'isl': 3000}}"
|
|
2. YAML string: "engine:\n backend: vllm\nsla:\n isl: 3000"
|
|
3. JSON string: '{"engine": {"backend": "vllm"}, "sla": {"isl": 3000}}'
|
|
|
|
Args:
|
|
config_str: Configuration string in one of the supported formats
|
|
|
|
Returns:
|
|
Dictionary containing the configuration
|
|
|
|
Raises:
|
|
ValueError: If config cannot be parsed or is not a dictionary
|
|
"""
|
|
config = None
|
|
|
|
# Try 1: Parse as Python dict literal (most direct for CLI)
|
|
try:
|
|
config = ast.literal_eval(config_str)
|
|
if isinstance(config, dict):
|
|
return config
|
|
except (ValueError, SyntaxError):
|
|
pass
|
|
|
|
# Try 2: Parse as YAML/JSON (for K8s ConfigMaps and files)
|
|
try:
|
|
config = yaml.safe_load(config_str)
|
|
if config is not None and isinstance(config, dict):
|
|
return config
|
|
except yaml.YAMLError:
|
|
pass
|
|
|
|
# If we got here, parsing failed
|
|
raise ValueError(
|
|
"Failed to parse config string. Expected Python dict literal, YAML, or JSON format. "
|
|
f"Examples:\n"
|
|
f" Python dict: \"{'engine': {'backend': 'vllm'}}\"\n"
|
|
f' YAML: "engine:\\n backend: vllm"\n'
|
|
f' JSON: \'{{"engine": {{"backend": "vllm"}}}}\''
|
|
)
|
|
|
|
|
|
def create_profiler_parser() -> argparse.Namespace:
|
|
"""
|
|
Create argument parser with support for YAML config string.
|
|
|
|
Config structure (camelCase preferred, snake_case supported for backwards compat):
|
|
outputDir: String (path to the output results directory, default: profiling_results)
|
|
deployment:
|
|
namespace: String (kubernetes namespace, default: dynamo-sla-profiler)
|
|
serviceName: String (service name, default: "")
|
|
model: String (served model name)
|
|
dgdImage: String (container image to use for DGD components (frontend, planner, workers), overrides images in config file)
|
|
deploymentTimeout: Int (maximum time to wait for deployment to become ready in seconds, default: 1800)
|
|
modelCache:
|
|
pvcName: String (name of the PVC to mount the model cache,
|
|
if not provided, model must be HF name and will download from HF, default: "")
|
|
pvcPath: String (path to the model cache in the PVC, default: "")
|
|
mountPath: String (path to the model cache in the container,
|
|
note that the PVC must be mounted to the same path for the profiling job,
|
|
default: "/opt/model-cache")
|
|
engine:
|
|
backend: String (backend type, currently support [vllm, sglang, trtllm], default: vllm)
|
|
config: String (path to the DynamoGraphDeployment config file, default: "")
|
|
maxContextLength: Int (maximum context length supported by the served model, default: 0)
|
|
isMoeModel: Boolean (enable MoE (Mixture of Experts) model support, use TEP for prefill and DEP for decode, default: False)
|
|
hardware:
|
|
minNumGpusPerEngine: Int (minimum number of GPUs per engine, default: 0)
|
|
maxNumGpusPerEngine: Int (maximum number of GPUs per engine, default: 0)
|
|
numGpusPerNode: Int (number of GPUs per node for MoE models - this will be the granularity when searching for the best TEP/DEP size, default: 0)
|
|
enableGpuDiscovery: Boolean (enable automatic GPU discovery from Kubernetes cluster nodes, when enabled overrides any manually specified hardware configuration, requires cluster-wide node access permissions, default: False)
|
|
sweep:
|
|
prefillInterpolationGranularity: Int (how many samples to benchmark to interpolate TTFT under different ISL, default: 16)
|
|
decodeInterpolationGranularity: Int (how many samples to benchmark to interpolate ITL under different active kv cache size and decode context length, default: 6)
|
|
useAiConfigurator: Boolean (use ai-configurator to estimate benchmarking results instead of running actual deployment, default: False)
|
|
aicSystem: String (target system for use with aiconfigurator, default: None)
|
|
aicHfId: String (aiconfigurator huggingface id of the target model, default: None)
|
|
aicBackend: String (aiconfigurator backend of the target model, if not provided, will use args.backend, default: "")
|
|
aicBackendVersion: String (specify backend version when using aiconfigurator to estimate perf, default: None)
|
|
dryRun: Boolean (dry run the profile job, default: False)
|
|
pickWithWebui: Boolean (pick the best parallelization mapping using webUI, default: False)
|
|
webuiPort: Int (webUI port, default: $PROFILER_WEBUI_PORT or 8000)
|
|
sla:
|
|
isl: Int (target input sequence length, default: 3000)
|
|
osl: Int (target output sequence length, default: 500)
|
|
ttft: Float (target Time To First Token in milliseconds, default: 50)
|
|
itl: Float (target Inter Token Latency in milliseconds, default: 10)
|
|
planner: (planner arguments)
|
|
e.g., plannerMinEndpoint: 2
|
|
"""
|
|
# Step 1: Pre-parse to check if --profile-config is provided
|
|
pre_parser = argparse.ArgumentParser(add_help=False)
|
|
pre_parser.add_argument("--profile-config", type=str)
|
|
pre_args, _ = pre_parser.parse_known_args()
|
|
|
|
# Step 2: Parse config if provided
|
|
config = {}
|
|
if pre_args.profile_config:
|
|
config = parse_config_string(pre_args.profile_config)
|
|
|
|
# Step 3: Create main parser with config-aware defaults
|
|
parser = argparse.ArgumentParser(
|
|
description="Profile the TTFT and ITL of the Prefill and Decode engine with different parallelization mapping. When profiling prefill we mock/fix decode,when profiling decode we mock/fix prefill."
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--profile-config",
|
|
type=str,
|
|
help="Configuration as Python dict literal, YAML, or JSON string. CLI args override config values. "
|
|
"Example: \"{'engine': {'backend': 'vllm', 'config': '/path'}, 'sla': {'isl': 3000}}\"",
|
|
)
|
|
|
|
# CLI arguments with config-aware defaults (using nested .get() for cleaner code)
|
|
parser.add_argument(
|
|
"--model",
|
|
type=str,
|
|
default=config.get("deployment", {}).get("model", ""),
|
|
help="Served model name",
|
|
)
|
|
model_cache_config = config.get("deployment", {}).get("modelCache", {})
|
|
parser.add_argument(
|
|
"--model-cache-pvc-name",
|
|
type=str,
|
|
default=model_cache_config.get("pvcName", ""),
|
|
help="Name of the PVC that contains the model weights. If not provided, args.model must be a HF model name and will download from HF",
|
|
)
|
|
parser.add_argument(
|
|
"--model-cache-pvc-path",
|
|
type=str,
|
|
default=model_cache_config.get("pvcPath", ""),
|
|
help="Path to the model cache in the PVC",
|
|
)
|
|
parser.add_argument(
|
|
"--model-cache-pvc-mount-path",
|
|
type=str,
|
|
default=model_cache_config.get("mountPath", "/opt/model-cache"),
|
|
help="Path to the model cache in the container, note that the PVC must be mounted to the same path for the profiling job",
|
|
)
|
|
deployment_cfg = config.get("deployment", {})
|
|
parser.add_argument(
|
|
"--dgd-image",
|
|
type=str,
|
|
default=_get(deployment_cfg, "dgdImage", "dgd_image", ""),
|
|
help="Container image to use for DGD components (frontend, planner, workers). Overrides images in config file.",
|
|
)
|
|
parser.add_argument(
|
|
"--deployment-timeout",
|
|
type=int,
|
|
default=_get(deployment_cfg, "deploymentTimeout", "deployment_timeout", 1800),
|
|
help="Maximum time to wait for deployment to become ready in seconds (default: 1800)",
|
|
)
|
|
|
|
parser.add_argument(
|
|
"--namespace",
|
|
type=str,
|
|
default=deployment_cfg.get("namespace", "dynamo-sla-profiler"),
|
|
help="Kubernetes namespace to deploy the DynamoGraphDeployment",
|
|
)
|
|
parser.add_argument(
|
|
"--backend",
|
|
type=str,
|
|
default=config.get("engine", {}).get("backend", "vllm"),
|
|
choices=["vllm", "sglang", "trtllm"],
|
|
help="backend type, currently support [vllm, sglang, trtllm]",
|
|
)
|
|
parser.add_argument(
|
|
"--config",
|
|
type=str,
|
|
default=config.get("engine", {}).get("config", ""),
|
|
required=False,
|
|
help="Path to the DynamoGraphDeployment config file (required, can be provided via CLI or config)",
|
|
)
|
|
parser.add_argument(
|
|
"--output-dir",
|
|
type=str,
|
|
default=_get(config, "outputDir", "output_dir", "profiling_results"),
|
|
help="Path to the output results directory",
|
|
)
|
|
hardware_cfg = config.get("hardware", {})
|
|
parser.add_argument(
|
|
"--min-num-gpus-per-engine",
|
|
type=int,
|
|
default=_get(hardware_cfg, "minNumGpusPerEngine", "min_num_gpus_per_engine", 0),
|
|
help="minimum number of GPUs per engine",
|
|
)
|
|
parser.add_argument(
|
|
"--max-num-gpus-per-engine",
|
|
type=int,
|
|
default=_get(hardware_cfg, "maxNumGpusPerEngine", "max_num_gpus_per_engine", 0),
|
|
help="maximum number of GPUs per engine",
|
|
)
|
|
parser.add_argument(
|
|
"--num-gpus-per-node",
|
|
type=int,
|
|
default=_get(hardware_cfg, "numGpusPerNode", "num_gpus_per_node", 0),
|
|
help="Number of GPUs per node for MoE models - this will be the granularity when searching for the best TEP/DEP size",
|
|
)
|
|
parser.add_argument(
|
|
"--isl",
|
|
type=int,
|
|
default=config.get("sla", {}).get("isl", 3000),
|
|
help="target input sequence length",
|
|
)
|
|
parser.add_argument(
|
|
"--osl",
|
|
type=int,
|
|
default=config.get("sla", {}).get("osl", 500),
|
|
help="target output sequence length",
|
|
)
|
|
parser.add_argument(
|
|
"--ttft",
|
|
type=float,
|
|
default=config.get("sla", {}).get("ttft", 50.0),
|
|
help="target Time To First Token (float, in milliseconds)",
|
|
)
|
|
parser.add_argument(
|
|
"--itl",
|
|
type=float,
|
|
default=config.get("sla", {}).get("itl", 10.0),
|
|
help="target Inter Token Latency (float, in milliseconds)",
|
|
)
|
|
|
|
# arguments used for interpolating TTFT and ITL under different ISL/OSL
|
|
engine_cfg = config.get("engine", {})
|
|
parser.add_argument(
|
|
"--max-context-length",
|
|
type=int,
|
|
default=_get(engine_cfg, "maxContextLength", "max_context_length", 0),
|
|
help="maximum context length supported by the served model",
|
|
)
|
|
sweep_cfg = config.get("sweep", {})
|
|
parser.add_argument(
|
|
"--prefill-interpolation-granularity",
|
|
type=int,
|
|
default=_get(
|
|
sweep_cfg,
|
|
"prefillInterpolationGranularity",
|
|
"prefill_interpolation_granularity",
|
|
16,
|
|
),
|
|
help="how many samples to benchmark to interpolate TTFT under different ISL",
|
|
)
|
|
parser.add_argument(
|
|
"--decode-interpolation-granularity",
|
|
type=int,
|
|
default=_get(
|
|
sweep_cfg,
|
|
"decodeInterpolationGranularity",
|
|
"decode_interpolation_granularity",
|
|
6,
|
|
),
|
|
help="how many samples to benchmark to interpolate ITL under different active kv cache size and decode context length",
|
|
)
|
|
parser.add_argument(
|
|
"--service-name",
|
|
type=str,
|
|
default=_get(deployment_cfg, "serviceName", "service_name", ""),
|
|
help="Service name for port forwarding (default: {deployment_name}-frontend)",
|
|
)
|
|
parser.add_argument(
|
|
"--dry-run",
|
|
action="store_true",
|
|
default=_get(sweep_cfg, "dryRun", "dry_run", False),
|
|
help="Dry run the profile job",
|
|
)
|
|
parser.add_argument(
|
|
"--enable-gpu-discovery",
|
|
action="store_true",
|
|
default=_get(hardware_cfg, "enableGpuDiscovery", "enable_gpu_discovery", False),
|
|
help="Enable automatic GPU discovery from Kubernetes cluster nodes. When enabled, overrides any manually specified hardware configuration. Requires cluster-wide node access permissions.",
|
|
)
|
|
parser.add_argument(
|
|
"--pick-with-webui",
|
|
action="store_true",
|
|
default=_get(sweep_cfg, "pickWithWebui", "pick_with_webui", False),
|
|
help="Pick the best parallelization mapping using webUI",
|
|
)
|
|
|
|
default_webui_port = 8000
|
|
webui_port_env = os.environ.get("PROFILER_WEBUI_PORT")
|
|
if webui_port_env:
|
|
default_webui_port = int(webui_port_env)
|
|
parser.add_argument(
|
|
"--webui-port",
|
|
type=int,
|
|
default=_get(sweep_cfg, "webuiPort", "webui_port", default_webui_port),
|
|
help="WebUI port",
|
|
)
|
|
|
|
# Dynamically add all planner arguments from planner_argparse.py
|
|
add_planner_arguments_to_parser(parser, prefix="planner-")
|
|
# Set defaults for any planner arguments found in config.planner
|
|
# Normalize keys: camelCase -> snake_case, hyphens -> underscores
|
|
planner_config = config.get("planner", {})
|
|
if planner_config:
|
|
normalized_planner_config = {
|
|
_camel_to_snake(key).replace("-", "_"): value
|
|
for key, value in planner_config.items()
|
|
}
|
|
parser.set_defaults(**normalized_planner_config)
|
|
|
|
# arguments if using aiconfigurator
|
|
parser.add_argument(
|
|
"--use-ai-configurator",
|
|
action="store_true",
|
|
default=_get(sweep_cfg, "useAiConfigurator", "use_ai_configurator", False),
|
|
help="Use ai-configurator to estimate benchmarking results instead of running actual deployment.",
|
|
)
|
|
parser.add_argument(
|
|
"--aic-system",
|
|
type=str,
|
|
default=_get(sweep_cfg, "aicSystem", "aic_system", None),
|
|
help="Target system for use with aiconfigurator (e.g. h100_sxm, h200_sxm)",
|
|
)
|
|
parser.add_argument(
|
|
"--aic-hf-id",
|
|
type=str,
|
|
default=_get(sweep_cfg, "aicHfId", "aic_hf_id", None),
|
|
help="aiconfigurator name of the target model (e.g. Qwen/Qwen3-32B, meta-llama/Llama-3.1-405B)",
|
|
)
|
|
parser.add_argument(
|
|
"--aic-backend",
|
|
type=str,
|
|
default=_get(sweep_cfg, "aicBackend", "aic_backend", ""),
|
|
help="aiconfigurator backend of the target model, if not provided, will use args.backend",
|
|
)
|
|
parser.add_argument(
|
|
"--aic-backend-version",
|
|
type=str,
|
|
default=_get(sweep_cfg, "aicBackendVersion", "aic_backend_version", None),
|
|
help="Specify backend version when using aiconfigurator to estimate perf.",
|
|
)
|
|
|
|
# Parse arguments
|
|
args = parser.parse_args()
|
|
|
|
# remove --profile-config from args
|
|
if hasattr(args, "profile_config"):
|
|
delattr(args, "profile_config")
|
|
|
|
# Validate required arguments
|
|
# Either --model or --config (or both) must be provided
|
|
if not args.model and not args.config:
|
|
parser.error("--model or --config is required (provide at least one)")
|
|
|
|
auto_generate_search_space(args)
|
|
return args
|