mooncake-track
/
dynamo
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

chore: restructure mocker cli args handling, to include prefill/decode (#3847) 

Signed-off-by: PeaBrane <yanrpei@gmail.com>
This commit is contained in:
Yan Ru Pei 2025-10-23 10:41:23 -07:00 committed by GitHub
parent a1b38af25e
commit 41ff394f51
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
9 changed files with 265 additions and 180 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
benchmarks/router/run_engines.sh
View File

@ -198,6 +198,11 @@ for i in $(seq 1 $NUM_WORKERS); do
if [ "$DATA_PARALLEL_SIZE" -gt 1 ]; then
MOCKER_ARGS+=("--data-parallel-size" "$DATA_PARALLEL_SIZE")
fi
if [ "$MODE" = "prefill" ]; then
MOCKER_ARGS+=("--is-prefill-worker")
elif [ "$MODE" = "decode" ]; then
MOCKER_ARGS+=("--is-decode-worker")
fi
MOCKER_ARGS+=("${EXTRA_ARGS[@]}")
exec python -m dynamo.mocker "${MOCKER_ARGS[@]}"

13
components/src/dynamo/mocker/README.md
View File

@ -40,14 +40,5 @@ python -m dynamo.mocker \
python -m dynamo.frontend --http-port 8000
```
### Legacy JSON file support:
For backward compatibility, you can still provide configuration via a JSON file:
```bash
echo '{"speedup_ratio": 10.0, "num_gpu_blocks": 8192}' > mocker_args.json
python -m dynamo.mocker \
--model-path TinyLlama/TinyLlama-1.1B-Chat-v1.0 \
--extra-engine-args mocker_args.json
```
Note: If `--extra-engine-args` is provided, it overrides all individual CLI arguments.
> [!Note]
> Each mocker instance runs as a single process, and each DP worker (specified by `--data-parallel-size`) is spawned as a lightweight async task within that process. For benchmarking (e.g., router testing), you would much prefer launching one mocker instance with a large `--data-parallel-size` rather than multiple separate mocker instances to reduce overhead.

202
components/src/dynamo/mocker/args.py Normal file
View File

@ -0,0 +1,202 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
import argparse
import json
import logging
import os
import tempfile
from pathlib import Path
from . import __version__
DYN_NAMESPACE = os.environ.get("DYN_NAMESPACE", "dynamo")
DEFAULT_ENDPOINT = f"dyn://{DYN_NAMESPACE}.backend.generate"
logger = logging.getLogger(__name__)
def create_temp_engine_args_file(args) -> Path:
"""
Create a temporary JSON file with MockEngineArgs from CLI arguments.
Returns the path to the temporary file.
"""
engine_args = {}
# Only include non-None values that differ from defaults
# Note: argparse converts hyphens to underscores in attribute names
# Extract all potential engine arguments, using None as default for missing attributes
engine_args = {
"num_gpu_blocks": getattr(args, "num_gpu_blocks", None),
"block_size": getattr(args, "block_size", None),
"max_num_seqs": getattr(args, "max_num_seqs", None),
"max_num_batched_tokens": getattr(args, "max_num_batched_tokens", None),
"enable_prefix_caching": getattr(args, "enable_prefix_caching", None),
"enable_chunked_prefill": getattr(args, "enable_chunked_prefill", None),
"watermark": getattr(args, "watermark", None),
"speedup_ratio": getattr(args, "speedup_ratio", None),
"dp_size": getattr(args, "dp_size", None),
"startup_time": getattr(args, "startup_time", None),
"is_prefill": getattr(args, "is_prefill_worker", None),
"is_decode": getattr(args, "is_decode_worker", None),
}
# Remove None values to only include explicitly set arguments
engine_args = {k: v for k, v in engine_args.items() if v is not None}
# Create temporary file
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
json.dump(engine_args, f, indent=2)
temp_path = Path(f.name)
logger.debug(f"Created temporary MockEngineArgs file at {temp_path}")
logger.debug(f"MockEngineArgs: {engine_args}")
return temp_path
def validate_worker_type_args(args):
"""
Validate that is_prefill_worker and is_decode_worker are not both True.
Raises ValueError if validation fails.
"""
if args.is_prefill_worker and args.is_decode_worker:
raise ValueError(
"Cannot specify both --is-prefill-worker and --is-decode-worker. "
"A worker must be either prefill, decode, or aggregated (neither flag set)."
)
def parse_args():
parser = argparse.ArgumentParser(
description="Mocker engine for testing Dynamo LLM infrastructure with vLLM-style CLI.",
formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument(
"--version", action="version", version=f"Dynamo Mocker {__version__}"
)
# Basic configuration
parser.add_argument(
"--model-path",
type=str,
help="Path to model directory or HuggingFace model ID for tokenizer",
)
parser.add_argument(
"--endpoint",
type=str,
default=DEFAULT_ENDPOINT,
help=f"Dynamo endpoint string (default: {DEFAULT_ENDPOINT})",
)
parser.add_argument(
"--model-name",
type=str,
default=None,
help="Model name for API responses (default: derived from model-path)",
)
# MockEngineArgs parameters (similar to vLLM style)
parser.add_argument(
"--num-gpu-blocks-override",
type=int,
dest="num_gpu_blocks", # Maps to num_gpu_blocks in MockEngineArgs
default=None,
help="Number of GPU blocks for KV cache (default: 16384)",
)
parser.add_argument(
"--block-size",
type=int,
default=None,
help="Token block size for KV cache blocks (default: 64)",
)
parser.add_argument(
"--max-num-seqs",
type=int,
default=None,
help="Maximum number of sequences per iteration (default: 256)",
)
parser.add_argument(
"--max-num-batched-tokens",
type=int,
default=None,
help="Maximum number of batched tokens per iteration (default: 8192)",
)
parser.add_argument(
"--enable-prefix-caching",
action="store_true",
dest="enable_prefix_caching",
default=None,
help="Enable automatic prefix caching (default: True)",
)
parser.add_argument(
"--no-enable-prefix-caching",
action="store_false",
dest="enable_prefix_caching",
default=None,
help="Disable automatic prefix caching",
)
parser.add_argument(
"--enable-chunked-prefill",
action="store_true",
dest="enable_chunked_prefill",
default=None,
help="Enable chunked prefill (default: True)",
)
parser.add_argument(
"--no-enable-chunked-prefill",
action="store_false",
dest="enable_chunked_prefill",
default=None,
help="Disable chunked prefill",
)
parser.add_argument(
"--watermark",
type=float,
default=None,
help="Watermark value for the mocker engine (default: 0.01)",
)
parser.add_argument(
"--speedup-ratio",
type=float,
default=None,
help="Speedup ratio for mock execution (default: 1.0)",
)
parser.add_argument(
"--data-parallel-size",
type=int,
dest="dp_size",
default=None,
help="Number of data parallel replicas (default: 1)",
)
parser.add_argument(
"--startup-time",
type=float,
default=None,
help="Simulated engine startup time in seconds (default: None)",
)
# Legacy support - allow direct JSON file specification
parser.add_argument(
"--extra-engine-args",
type=Path,
help="Path to JSON file with mocker configuration. "
"If provided, overrides individual CLI arguments.",
)
# Worker type configuration
parser.add_argument(
"--is-prefill-worker",
action="store_true",
default=False,
help="Register as Prefill model type instead of Chat+Completions (default: False)",
)
parser.add_argument(
"--is-decode-worker",
action="store_true",
default=False,
help="Mark this as a decode worker which does not publish KV events and skips prefill cost estimation (default: False)",
)
args = parser.parse_args()
validate_worker_type_args(args)
return args

169
components/src/dynamo/mocker/main.py
View File

@ -4,12 +4,7 @@
# Usage: `python -m dynamo.mocker --model-path /data/models/Qwen3-0.6B`
# Now supports vLLM-style individual arguments for MockEngineArgs
import argparse
import json
import logging
import os
import tempfile
from pathlib import Path
import uvloop
@ -17,55 +12,15 @@ from dynamo.llm import EngineType, EntrypointArgs, make_engine, run_input
from dynamo.runtime import DistributedRuntime, dynamo_worker
from dynamo.runtime.logging import configure_dynamo_logging
from . import __version__
DYN_NAMESPACE = os.environ.get("DYN_NAMESPACE", "dynamo")
DEFAULT_ENDPOINT = f"dyn://{DYN_NAMESPACE}.backend.generate"
from .args import create_temp_engine_args_file, parse_args
configure_dynamo_logging()
logger = logging.getLogger(__name__)
def create_temp_engine_args_file(args) -> Path:
"""
Create a temporary JSON file with MockEngineArgs from CLI arguments.
Returns the path to the temporary file.
"""
engine_args = {}
# Only include non-None values that differ from defaults
# Note: argparse converts hyphens to underscores in attribute names
# Extract all potential engine arguments, using None as default for missing attributes
engine_args = {
"num_gpu_blocks": getattr(args, "num_gpu_blocks", None),
"block_size": getattr(args, "block_size", None),
"max_num_seqs": getattr(args, "max_num_seqs", None),
"max_num_batched_tokens": getattr(args, "max_num_batched_tokens", None),
"enable_prefix_caching": getattr(args, "enable_prefix_caching", None),
"enable_chunked_prefill": getattr(args, "enable_chunked_prefill", None),
"watermark": getattr(args, "watermark", None),
"speedup_ratio": getattr(args, "speedup_ratio", None),
"dp_size": getattr(args, "dp_size", None),
"startup_time": getattr(args, "startup_time", None),
}
# Remove None values to only include explicitly set arguments
engine_args = {k: v for k, v in engine_args.items() if v is not None}
# Create temporary file
with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
json.dump(engine_args, f, indent=2)
temp_path = Path(f.name)
logger.debug(f"Created temporary MockEngineArgs file at {temp_path}")
logger.debug(f"MockEngineArgs: {engine_args}")
return temp_path
@dynamo_worker(static=False)
async def worker(runtime: DistributedRuntime):
args = cmd_line_args()
args = parse_args()
# Handle extra_engine_args: either use provided file or create from CLI args
if args.extra_engine_args:
@ -85,6 +40,7 @@ async def worker(runtime: DistributedRuntime):
model_name=args.model_name,
endpoint_id=args.endpoint,
extra_engine_args=extra_engine_args_path,
is_prefill=args.is_prefill_worker,
)
# Create and run the engine
@ -101,125 +57,6 @@ async def worker(runtime: DistributedRuntime):
logger.warning(f"Failed to clean up temporary file: {e}")
def cmd_line_args():
parser = argparse.ArgumentParser(
description="Mocker engine for testing Dynamo LLM infrastructure with vLLM-style CLI.",
formatter_class=argparse.RawDescriptionHelpFormatter,
)
parser.add_argument(
"--version", action="version", version=f"Dynamo Mocker {__version__}"
)
# Basic configuration
parser.add_argument(
"--model-path",
type=str,
help="Path to model directory or HuggingFace model ID for tokenizer",
)
parser.add_argument(
"--endpoint",
type=str,
default=DEFAULT_ENDPOINT,
help=f"Dynamo endpoint string (default: {DEFAULT_ENDPOINT})",
)
parser.add_argument(
"--model-name",
type=str,
default=None,
help="Model name for API responses (default: derived from model-path)",
)
# MockEngineArgs parameters (similar to vLLM style)
parser.add_argument(
"--num-gpu-blocks-override",
type=int,
dest="num_gpu_blocks", # Maps to num_gpu_blocks in MockEngineArgs
default=None,
help="Number of GPU blocks for KV cache (default: 16384)",
)
parser.add_argument(
"--block-size",
type=int,
default=None,
help="Token block size for KV cache blocks (default: 64)",
)
parser.add_argument(
"--max-num-seqs",
type=int,
default=None,
help="Maximum number of sequences per iteration (default: 256)",
)
parser.add_argument(
"--max-num-batched-tokens",
type=int,
default=None,
help="Maximum number of batched tokens per iteration (default: 8192)",
)
parser.add_argument(
"--enable-prefix-caching",
action="store_true",
dest="enable_prefix_caching",
default=None,
help="Enable automatic prefix caching (default: True)",
)
parser.add_argument(
"--no-enable-prefix-caching",
action="store_false",
dest="enable_prefix_caching",
default=None,
help="Disable automatic prefix caching",
)
parser.add_argument(
"--enable-chunked-prefill",
action="store_true",
dest="enable_chunked_prefill",
default=None,
help="Enable chunked prefill (default: True)",
)
parser.add_argument(
"--no-enable-chunked-prefill",
action="store_false",
dest="enable_chunked_prefill",
default=None,
help="Disable chunked prefill",
)
parser.add_argument(
"--watermark",
type=float,
default=None,
help="Watermark value for the mocker engine (default: 0.01)",
)
parser.add_argument(
"--speedup-ratio",
type=float,
default=None,
help="Speedup ratio for mock execution (default: 1.0)",
)
parser.add_argument(
"--data-parallel-size",
type=int,
dest="dp_size",
default=None,
help="Number of data parallel replicas (default: 1)",
)
parser.add_argument(
"--startup-time",
type=float,
default=None,
help="Simulated engine startup time in seconds (default: None)",
)
# Legacy support - allow direct JSON file specification
parser.add_argument(
"--extra-engine-args",
type=Path,
help="Path to JSON file with mocker configuration. "
"If provided, overrides individual CLI arguments.",
)
return parser.parse_args()
def main():
uvloop.run(worker())

1
launch/dynamo-run/src/lib.rs
View File

@ -145,6 +145,7 @@ async fn engine_for(
engine,
model: Box::new(local_model),
is_static: flags.static_worker,
is_prefill: false,
})
}
}

6
lib/bindings/python/rust/llm/entrypoint.rs
View File

@ -120,13 +120,14 @@ pub(crate) struct EntrypointArgs {
namespace: Option<String>,
custom_backend_metrics_endpoint: Option<String>,
custom_backend_metrics_polling_interval: Option<f64>,
is_prefill: bool,
}
#[pymethods]
impl EntrypointArgs {
#[allow(clippy::too_many_arguments)]
#[new]
#[pyo3(signature = (engine_type, model_path=None, model_name=None, endpoint_id=None, context_length=None, template_file=None, router_config=None, kv_cache_block_size=None, http_host=None, http_port=None, tls_cert_path=None, tls_key_path=None, extra_engine_args=None, namespace=None, custom_backend_metrics_endpoint=None, custom_backend_metrics_polling_interval=None))]
#[pyo3(signature = (engine_type, model_path=None, model_name=None, endpoint_id=None, context_length=None, template_file=None, router_config=None, kv_cache_block_size=None, http_host=None, http_port=None, tls_cert_path=None, tls_key_path=None, extra_engine_args=None, namespace=None, custom_backend_metrics_endpoint=None, custom_backend_metrics_polling_interval=None, is_prefill=false))]
pub fn new(
engine_type: EngineType,
model_path: Option<PathBuf>,
@ -144,6 +145,7 @@ impl EntrypointArgs {
namespace: Option<String>,
custom_backend_metrics_endpoint: Option<String>,
custom_backend_metrics_polling_interval: Option<f64>,
is_prefill: bool,
) -> PyResult<Self> {
let endpoint_id_obj: Option<EndpointId> = endpoint_id.as_deref().map(EndpointId::from);
if (tls_cert_path.is_some() && tls_key_path.is_none())
@ -170,6 +172,7 @@ impl EntrypointArgs {
namespace,
custom_backend_metrics_endpoint,
custom_backend_metrics_polling_interval,
is_prefill,
})
}
}
@ -277,6 +280,7 @@ async fn select_engine(
engine,
model: Box::new(local_model),
is_static: false,
is_prefill: args.is_prefill,
}
}
};

1
lib/llm/src/entrypoint.rs
View File

@ -59,6 +59,7 @@ pub enum EngineConfig {
engine: ExecutionContext,
model: Box<LocalModel>,
is_static: bool,
is_prefill: bool,
},
}

8
lib/llm/src/entrypoint/input/endpoint.rs
View File

@ -67,6 +67,7 @@ pub async fn run(
engine: inner_engine,
mut model,
is_static,
is_prefill,
} => {
// Pre-processing is done ingress-side, so it should be already done.
let frontend = SegmentSource::<
@ -83,8 +84,11 @@ pub async fn run(
let ingress = Ingress::for_pipeline(pipeline)?;
if !is_static {
// Default to supporting both Chat and Completions endpoints
let model_type = ModelType::Chat | ModelType::Completions;
let model_type = if is_prefill {
ModelType::Prefill
} else {
ModelType::Chat | ModelType::Completions
};
model
.attach(&endpoint, model_type, ModelInput::Tokens)
.await?;

40
lib/llm/src/mocker/protocols.rs
View File

@ -61,6 +61,18 @@ pub struct OutputSignal {
pub completed: bool,
}
/// Worker type for disaggregated serving configurations
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize, Default)]
pub enum WorkerType {
/// Standard aggregated worker handling both prefill and decode
#[default]
Aggregated,
/// Dedicated prefill worker in disaggregated mode
Prefill,
/// Dedicated decode worker in disaggregated mode
Decode,
}
/// Configuration arguments for MockVllmEngine
#[derive(Debug, Clone, Serialize, Deserialize, Builder)]
#[builder(pattern = "owned", build_fn(public))]
@ -97,6 +109,10 @@ pub struct MockEngineArgs {
/// Optional startup time in seconds to simulate engine initialization delay
#[builder(default = "None")]
pub startup_time: Option<f64>,
/// Worker type for disaggregated serving (Aggregated, Prefill, or Decode)
#[builder(default = "WorkerType::Aggregated")]
pub worker_type: WorkerType,
}
impl Default for MockEngineArgs {
@ -132,6 +148,8 @@ impl MockEngineArgs {
"speedup_ratio",
"dp_size",
"startup_time",
"is_prefill",
"is_decode",
]
.iter()
.cloned()
@ -213,6 +231,28 @@ impl MockEngineArgs {
builder = builder.startup_time(Some(num));
}
// Parse worker type from is_prefill and is_decode flags
let is_prefill = extra_args
.get("is_prefill")
.and_then(|v| v.as_bool())
.unwrap_or(false);
let is_decode = extra_args
.get("is_decode")
.and_then(|v| v.as_bool())
.unwrap_or(false);
// Determine worker type based on flags
let worker_type = match (is_prefill, is_decode) {
(false, false) => WorkerType::Aggregated,
(true, false) => WorkerType::Prefill,
(false, true) => WorkerType::Decode,
(true, true) => panic!(
"Invalid worker configuration: is_prefill and is_decode cannot both be true. \
Worker must be either Aggregated (both false), Prefill (is_prefill=true), or Decode (is_decode=true)."
),
};
builder = builder.worker_type(worker_type);
// Build the MockEngineArgs with either defaults or overridden values
builder
.build()