xpyd-sim 0.1.0__tar.gz

xpyd_sim-0.1.0/PKG-INFO

@@ -0,0 +1,55 @@
+Metadata-Version: 2.4
+Name: xpyd-sim
+Version: 0.1.0
+Summary: OpenAI-compatible LLM inference simulator for xPyD
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: fastapi>=0.111.0
+Requires-Dist: uvicorn>=0.30.0
+Requires-Dist: pydantic>=2.0.0
+Requires-Dist: httpx>=0.27.0
+Requires-Dist: pyyaml>=6.0
+Requires-Dist: numpy>=1.26.0
+Requires-Dist: scipy>=1.12.0
+Requires-Dist: matplotlib>=3.8.0
+Provides-Extra: dev
+Requires-Dist: pytest>=8.0.0; extra == "dev"
+Requires-Dist: pytest-asyncio>=0.23.0; extra == "dev"
+Requires-Dist: ruff>=0.3.0; extra == "dev"
+Requires-Dist: isort>=5.13.0; extra == "dev"
+
+# xPyD-sim
+
+OpenAI-compatible LLM inference simulator for [xPyD](https://github.com/xPyD-hub).
+
+Simulates prefill and decode nodes with realistic behavior for testing xPyD-bench and xPyD-proxy without real GPU hardware.
+
+## Features
+
+- Separate prefill and decode node simulators
+- Full OpenAI API compatibility (/v1/completions, /v1/chat/completions, /v1/models)
+- Configurable latency (prefill delay, decode delay per token)
+- Streaming support with realistic token-by-token delivery
+- EOS simulation with configurable output length distribution
+- All OpenAI API parameters accepted
+- Spec-compliant response formats
+
+## Install
+
+```bash
+pip install xpyd-sim
+```
+
+## Quick Start
+
+```bash
+# Start prefill node
+xpyd-sim prefill --port 8001
+
+# Start decode node
+xpyd-sim decode --port 8002
+```
+
+## License
+
+TBD

xpyd_sim-0.1.0/README.md

@@ -0,0 +1,35 @@
+# xPyD-sim
+
+OpenAI-compatible LLM inference simulator for [xPyD](https://github.com/xPyD-hub).
+
+Simulates prefill and decode nodes with realistic behavior for testing xPyD-bench and xPyD-proxy without real GPU hardware.
+
+## Features
+
+- Separate prefill and decode node simulators
+- Full OpenAI API compatibility (/v1/completions, /v1/chat/completions, /v1/models)
+- Configurable latency (prefill delay, decode delay per token)
+- Streaming support with realistic token-by-token delivery
+- EOS simulation with configurable output length distribution
+- All OpenAI API parameters accepted
+- Spec-compliant response formats
+
+## Install
+
+```bash
+pip install xpyd-sim
+```
+
+## Quick Start
+
+```bash
+# Start prefill node
+xpyd-sim prefill --port 8001
+
+# Start decode node
+xpyd-sim decode --port 8002
+```
+
+## License
+
+TBD

xpyd_sim-0.1.0/pyproject.toml

@@ -0,0 +1,49 @@
+[build-system]
+requires = ["setuptools>=68.0", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "xpyd-sim"
+version = "0.1.0"
+description = "OpenAI-compatible LLM inference simulator for xPyD"
+readme = "README.md"
+requires-python = ">=3.10"
+dependencies = [
+    "fastapi>=0.111.0",
+    "uvicorn>=0.30.0",
+    "pydantic>=2.0.0",
+    "httpx>=0.27.0",
+    "pyyaml>=6.0",
+    "numpy>=1.26.0",
+    "scipy>=1.12.0",
+    "matplotlib>=3.8.0",
+]
+
+[project.optional-dependencies]
+dev = [
+    "pytest>=8.0.0",
+    "pytest-asyncio>=0.23.0",
+    "ruff>=0.3.0",
+    "isort>=5.13.0",
+]
+
+[project.scripts]
+xpyd-sim = "xpyd_sim.cli:main"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.ruff]
+line-length = 100
+target-version = "py310"
+
+[tool.ruff.lint]
+select = ["E", "F", "W", "I"]
+
+[tool.isort]
+profile = "black"
+line_length = 100
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]
+asyncio_mode = "auto"

xpyd_sim-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

xpyd_sim-0.1.0/src/xpyd_sim/__init__.py

@@ -0,0 +1,3 @@
+"""xPyD-sim: OpenAI-compatible LLM inference simulator."""
+
+__version__ = "0.1.0"

xpyd_sim-0.1.0/src/xpyd_sim/calibrate.py

@@ -0,0 +1,180 @@
+"""Calibrate tool: fit latency curves from sample points."""
+
+from __future__ import annotations
+
+import sys
+from typing import Any
+
+import numpy as np
+import yaml
+
+
+def _fit_1d(points: list[dict[str, float]], x_key: str, y_key: str) -> dict[str, Any]:
+    """Fit a 1D polynomial (degree 2) to sample points.
+
+    Returns coefficients for: y = a*x^2 + b*x + c
+    """
+    xs = np.array([p[x_key] for p in points], dtype=np.float64)
+    ys = np.array([p[y_key] for p in points], dtype=np.float64)
+    degree = min(2, len(xs) - 1)
+    coeffs = np.polyfit(xs, ys, degree)
+    # Pad to always have 3 coefficients (a, b, c)
+    padded = [0.0] * (3 - len(coeffs)) + list(coeffs)
+    return {
+        "type": "poly1d",
+        "coefficients": [float(c) for c in padded],
+        "x_key": x_key,
+        "y_key": y_key,
+        "x_range": [float(xs.min()), float(xs.max())],
+    }
+
+
+def _fit_2d(points: list[dict[str, float]]) -> dict[str, Any]:
+    """Fit a 2D surface: delay = a + b*bs + c*ctx + d*bs*ctx + e*bs^2 + f*ctx^2.
+
+    Uses least-squares on polynomial features.
+    """
+    bs = np.array([p["batch_size"] for p in points], dtype=np.float64)
+    ctx = np.array([p["context_length"] for p in points], dtype=np.float64)
+    y = np.array([p["delay_per_token_ms"] for p in points], dtype=np.float64)
+
+    # Build design matrix: [1, bs, ctx, bs*ctx, bs^2, ctx^2]
+    X = np.column_stack([  # noqa: N806
+        np.ones_like(bs),
+        bs,
+        ctx,
+        bs * ctx,
+        bs**2,
+        ctx**2,
+    ])
+
+    # Least squares fit
+    coeffs, _, _, _ = np.linalg.lstsq(X, y, rcond=None)
+
+    return {
+        "type": "poly2d",
+        "coefficients": [float(c) for c in coeffs],
+        "bs_range": [float(bs.min()), float(bs.max())],
+        "ctx_range": [float(ctx.min()), float(ctx.max())],
+    }
+
+
+def _validate_points(points: list[dict], min_count: int, label: str) -> None:
+    """Validate minimum number of sample points."""
+    if len(points) < min_count:
+        print(
+            f"Error: {label} requires at least {min_count} sample points, got {len(points)}",
+            file=sys.stderr,
+        )
+        sys.exit(1)
+
+
+def calibrate(input_path: str, output_path: str, plot_path: str | None = None) -> None:
+    """Run calibration: load samples, fit curves, write profile + optional plot."""
+    with open(input_path) as f:
+        data = yaml.safe_load(f)
+
+    profile: dict[str, Any] = {}
+
+    # Prefill: 1D fit (batch_size → delay_ms)
+    if "prefill" in data:
+        points = data["prefill"]
+        _validate_points(points, 3, "prefill")
+        profile["prefill"] = _fit_1d(points, "batch_size", "delay_ms")
+
+    # KV transfer: 1D fit (batch_size → delay_ms)
+    if "kv_transfer" in data:
+        points = data["kv_transfer"]
+        _validate_points(points, 3, "kv_transfer")
+        profile["kv_transfer"] = _fit_1d(points, "batch_size", "delay_ms")
+
+    # Decode: 2D fit (batch_size, context_length → delay_per_token_ms)
+    if "decode" in data:
+        points = data["decode"]
+        _validate_points(points, 9, "decode (2D)")
+        profile["decode"] = _fit_2d(points)
+
+    # Write profile
+    with open(output_path, "w") as f:
+        yaml.dump(profile, f, default_flow_style=False, sort_keys=False)
+
+    print(f"Profile written to {output_path}")
+
+    # Optional visualization
+    if plot_path:
+        _plot(data, profile, plot_path)
+        print(f"Plot written to {plot_path}")
+
+
+def _plot(
+    data: dict[str, Any], profile: dict[str, Any], plot_path: str
+) -> None:
+    """Generate visualization PNG with sample points and fitted curves."""
+    import matplotlib
+
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    num_plots = sum(1 for k in ("prefill", "kv_transfer", "decode") if k in data)
+    fig, axes = plt.subplots(1, num_plots, figsize=(6 * num_plots, 5))
+    if num_plots == 1:
+        axes = [axes]
+
+    idx = 0
+
+    for key, label in [("prefill", "Prefill"), ("kv_transfer", "KV Transfer")]:
+        if key not in data:
+            continue
+        ax = axes[idx]
+        idx += 1
+        points = data[key]
+        xs = [p["batch_size"] for p in points]
+        ys = [p["delay_ms"] for p in points]
+        ax.scatter(xs, ys, color="red", zorder=5, label="Sample points")
+
+        # Plot fitted curve
+        coeffs = profile[key]["coefficients"]
+        x_fit = np.linspace(min(xs) * 0.8, max(xs) * 1.2, 100)
+        y_fit = np.polyval(coeffs, x_fit)
+        ax.plot(x_fit, y_fit, color="blue", label="Fitted curve")
+        ax.set_xlabel("Batch size (tokens)")
+        ax.set_ylabel("Delay (ms)")
+        ax.set_title(f"{label} Latency")
+        ax.legend()
+
+    if "decode" in data:
+        ax = axes[idx]
+        points = data["decode"]
+        bs_vals = sorted(set(p["batch_size"] for p in points))
+        for b in bs_vals:
+            pts = [p for p in points if p["batch_size"] == b]
+            ctxs = [p["context_length"] for p in pts]
+            delays = [p["delay_per_token_ms"] for p in pts]
+            ax.scatter(ctxs, delays, zorder=5, label=f"bs={b} (data)")
+
+        # Fitted curves per batch_size
+        coeffs = profile["decode"]["coefficients"]
+        ctx_fit = np.linspace(
+            min(p["context_length"] for p in points) * 0.8,
+            max(p["context_length"] for p in points) * 1.2,
+            100,
+        )
+        for b in bs_vals:
+            y_fit = (
+                coeffs[0]
+                + coeffs[1] * b
+                + coeffs[2] * ctx_fit
+                + coeffs[3] * b * ctx_fit
+                + coeffs[4] * b**2
+                + coeffs[5] * ctx_fit**2
+            )
+            ax.plot(ctx_fit, y_fit, linestyle="--", label=f"bs={b} (fit)")
+
+        ax.set_xlabel("Context length")
+        ax.set_ylabel("Delay per token (ms)")
+        ax.set_title("Decode Latency (2D)")
+        ax.legend(fontsize=7)
+
+    plt.tight_layout()
+    plt.savefig(plot_path, dpi=150)
+    plt.close()

xpyd_sim-0.1.0/src/xpyd_sim/cli.py

@@ -0,0 +1,248 @@
+"""CLI entry point with YAML config and environment variable support."""
+
+from __future__ import annotations
+
+import argparse
+import os
+from pathlib import Path
+from typing import Any
+
+import yaml
+
+
+def _load_yaml_config(path: str | Path) -> dict[str, Any]:
+    """Load and flatten a YAML config file into CLI-compatible keys."""
+    with open(path) as f:
+        raw = yaml.safe_load(f) or {}
+
+    flat: dict[str, Any] = {}
+    # Top-level keys
+    for key in ("mode", "port", "host", "model", "max_model_len"):
+        if key in raw:
+            flat[key] = raw[key]
+
+    # Nested latency section
+    latency = raw.get("latency", {})
+    for key in ("prefill_delay_ms", "kv_transfer_delay_ms", "decode_delay_per_token_ms"):
+        if key in latency:
+            flat[key] = latency[key]
+
+    # Nested eos section
+    eos = raw.get("eos", {})
+    if "min_ratio" in eos:
+        flat["eos_min_ratio"] = eos["min_ratio"]
+
+    # Nested warmup section
+    warmup = raw.get("warmup", {})
+    if "requests" in warmup:
+        flat["warmup_requests"] = warmup["requests"]
+    if "penalty_ms" in warmup:
+        flat["warmup_penalty_ms"] = warmup["penalty_ms"]
+
+    # Nested logging section
+    logging_cfg = raw.get("logging", {})
+    if "request_log" in logging_cfg:
+        flat["log_requests"] = logging_cfg["request_log"]
+
+    # Profile
+    if "profile" in raw:
+        flat["profile"] = raw["profile"]
+
+    # Scheduling section
+    scheduling = raw.get("scheduling", {})
+    for key in ("max_num_batched_tokens", "max_num_seqs"):
+        if key in scheduling:
+            flat[key] = scheduling[key]
+    if "max_model_len" in scheduling:
+        flat["max_model_len"] = scheduling["max_model_len"]
+    if "enabled" in scheduling:
+        flat["scheduling_enabled"] = scheduling["enabled"]
+
+    return flat
+
+
+# Mapping: config key -> (env var name, type converter)
+_ENV_MAP: dict[str, tuple[str, type]] = {
+    "mode": ("XPYD_SIM_MODE", str),
+    "port": ("XPYD_SIM_PORT", int),
+    "host": ("XPYD_SIM_HOST", str),
+    "model": ("XPYD_SIM_MODEL", str),
+    "prefill_delay_ms": ("XPYD_SIM_PREFILL_DELAY_MS", float),
+    "kv_transfer_delay_ms": ("XPYD_SIM_KV_TRANSFER_DELAY_MS", float),
+    "decode_delay_per_token_ms": ("XPYD_SIM_DECODE_DELAY_PER_TOKEN_MS", float),
+    "eos_min_ratio": ("XPYD_SIM_EOS_MIN_RATIO", float),
+    "max_model_len": ("XPYD_SIM_MAX_MODEL_LEN", int),
+    "warmup_requests": ("XPYD_SIM_WARMUP_REQUESTS", int),
+    "warmup_penalty_ms": ("XPYD_SIM_WARMUP_PENALTY_MS", float),
+    "log_requests": ("XPYD_SIM_LOG_REQUESTS", str),
+    "profile": ("XPYD_SIM_PROFILE", str),
+    "max_num_batched_tokens": ("XPYD_SIM_MAX_NUM_BATCHED_TOKENS", int),
+    "max_num_seqs": ("XPYD_SIM_MAX_NUM_SEQS", int),
+    "scheduling_enabled": ("XPYD_SIM_SCHEDULING_ENABLED", lambda v: v.lower() in ("1", "true")),
+}
+
+# Default values for all config keys
+_DEFAULTS: dict[str, Any] = {
+    "mode": "dual",
+    "port": 8000,
+    "host": "0.0.0.0",
+    "model": "dummy",
+    "prefill_delay_ms": 50.0,
+    "kv_transfer_delay_ms": 5.0,
+    "decode_delay_per_token_ms": 10.0,
+    "eos_min_ratio": 0.5,
+    "max_model_len": 131072,
+    "warmup_requests": 0,
+    "warmup_penalty_ms": 0.0,
+    "log_requests": None,
+    "profile": None,
+    "max_num_batched_tokens": 8192,
+    "max_num_seqs": 256,
+    "scheduling_enabled": False,
+}
+
+
+def _resolve_config(
+    cli_args: argparse.Namespace,
+    yaml_config: dict[str, Any] | None = None,
+) -> dict[str, Any]:
+    """Resolve config with priority: CLI > env vars > YAML > defaults."""
+    result: dict[str, Any] = {}
+
+    # Map CLI arg names to config keys
+    cli_to_key = {
+        "mode": "mode",
+        "port": "port",
+        "host": "host",
+        "model": "model",
+        "prefill_delay_ms": "prefill_delay_ms",
+        "kv_transfer_delay_ms": "kv_transfer_delay_ms",
+        "decode_delay_per_token_ms": "decode_delay_per_token_ms",
+        "eos_min_ratio": "eos_min_ratio",
+        "max_model_len": "max_model_len",
+        "warmup_requests": "warmup_requests",
+        "warmup_penalty_ms": "warmup_penalty_ms",
+        "log_requests": "log_requests",
+        "profile": "profile",
+    }
+
+    for cli_attr, key in cli_to_key.items():
+        # 1. Check if CLI arg was explicitly provided
+        cli_val = getattr(cli_args, cli_attr, None)
+        default_val = _DEFAULTS[key]
+        # argparse sets unset args to their default — we use None defaults
+        # and check _explicitly_set to detect user intent
+        explicitly_set = key in getattr(cli_args, "_explicitly_set", set())
+
+        if explicitly_set:
+            result[key] = cli_val
+            continue
+
+        # 2. Check environment variable
+        env_name, converter = _ENV_MAP[key]
+        env_val = os.environ.get(env_name)
+        if env_val is not None:
+            try:
+                result[key] = converter(env_val)
+            except (ValueError, TypeError):
+                result[key] = default_val
+            continue
+
+        # 3. Check YAML config
+        if yaml_config and key in yaml_config:
+            result[key] = yaml_config[key]
+            continue
+
+        # 4. Use default
+        result[key] = default_val
+
+    return result
+
+
+class _TrackingNamespace(argparse.Namespace):
+    """Namespace that tracks which arguments were explicitly set on CLI."""
+
+    def __init__(self, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._explicitly_set: set[str] = set()
+
+
+class _TrackAction(argparse.Action):
+    """Custom action that records when an arg is explicitly provided."""
+
+    def __call__(  # type: ignore[override]
+        self, parser: argparse.ArgumentParser, namespace: argparse.Namespace,
+        values: Any, option_string: str | None = None,
+    ) -> None:
+        setattr(namespace, self.dest, values)
+        if hasattr(namespace, "_explicitly_set"):
+            namespace._explicitly_set.add(self.dest)
+
+
+def main(argv: list[str] | None = None) -> None:
+    parser = argparse.ArgumentParser(prog="xpyd-sim", description="xPyD inference simulator")
+    sub = parser.add_subparsers(dest="command")
+
+    serve = sub.add_parser("serve", help="Start the unified simulator server")
+    serve.add_argument("--mode", choices=["dual", "prefill", "decode"], default="dual",
+                       action=_TrackAction)
+    serve.add_argument("--port", type=int, default=8000, action=_TrackAction)
+    serve.add_argument("--host", default="0.0.0.0", action=_TrackAction)
+    serve.add_argument("--model", default="dummy", action=_TrackAction)
+    serve.add_argument("--prefill-delay-ms", type=float, default=50.0, action=_TrackAction)
+    serve.add_argument("--kv-transfer-delay-ms", type=float, default=5.0, action=_TrackAction)
+    serve.add_argument("--decode-delay-per-token-ms", type=float, default=10.0,
+                       action=_TrackAction)
+    serve.add_argument("--eos-min-ratio", type=float, default=0.5, action=_TrackAction)
+    serve.add_argument("--max-model-len", type=int, default=131072, action=_TrackAction)
+    serve.add_argument("--warmup-requests", type=int, default=0, action=_TrackAction)
+    serve.add_argument("--warmup-penalty-ms", type=float, default=0.0, action=_TrackAction)
+    serve.add_argument("--log-requests", type=str, default=None, action=_TrackAction)
+    serve.add_argument("--profile", type=str, default=None, action=_TrackAction)
+    serve.add_argument("--config", type=str, default=None, help="YAML config file path",
+                       action=_TrackAction)
+
+    # Calibrate subcommand
+    cal = sub.add_parser("calibrate", help="Fit latency curves from sample data")
+    cal.add_argument("--input", required=True, help="Path to sample points YAML")
+    cal.add_argument("--output", required=True, help="Path to write profile YAML")
+    cal.add_argument("--plot", default=None, help="Path to write visualization PNG")
+
+    args = parser.parse_args(argv, namespace=_TrackingNamespace())
+    if not args.command:
+        parser.print_help()
+        return
+
+    if args.command == "calibrate":
+        from xpyd_sim.calibrate import calibrate
+
+        calibrate(args.input, args.output, args.plot)
+        return
+
+    if args.command == "serve":
+        import uvicorn
+
+        from xpyd_sim.server import ServerConfig, create_app
+
+        # Load YAML config if provided
+        yaml_config = None
+        if args.config:
+            yaml_config = _load_yaml_config(args.config)
+
+        cfg = _resolve_config(args, yaml_config)
+
+        config = ServerConfig(
+            mode=cfg["mode"],
+            model_name=cfg["model"],
+            prefill_delay_ms=cfg["prefill_delay_ms"],
+            kv_transfer_delay_ms=cfg["kv_transfer_delay_ms"],
+            decode_delay_per_token_ms=cfg["decode_delay_per_token_ms"],
+            eos_min_ratio=cfg["eos_min_ratio"],
+            max_model_len=cfg["max_model_len"],
+            warmup_requests=cfg["warmup_requests"],
+            warmup_penalty_ms=cfg["warmup_penalty_ms"],
+            log_requests=cfg["log_requests"],
+            profile=cfg["profile"],
+        )
+        app = create_app(config)
+        uvicorn.run(app, host=cfg["host"], port=cfg["port"])

xpyd_sim-0.1.0/src/xpyd_sim/common/__init__.py

@@ -0,0 +1 @@
+"""Common models and helpers."""

xpyd_sim-0.1.0/src/xpyd_sim/common/helpers.py

@@ -0,0 +1,47 @@
+"""Shared helper functions."""
+
+from __future__ import annotations
+
+import time
+import uuid
+from typing import Any, Optional
+
+DUMMY_TOKENS = list("The quick brown fox jumps over the lazy dog. " * 20)
+DEFAULT_MAX_TOKENS = 16
+
+
+def generate_id(prefix: str = "chatcmpl") -> str:
+    return f"{prefix}-{uuid.uuid4().hex[:12]}"
+
+
+def now_ts() -> int:
+    return int(time.time())
+
+
+def get_effective_max_tokens(*values: Optional[int]) -> int:
+    for v in values:
+        if v is not None:
+            return v
+    return DEFAULT_MAX_TOKENS
+
+
+def count_prompt_tokens(prompt: Any = None, messages: list | None = None) -> int:
+    if messages is not None:
+        total = sum(
+            len(str(getattr(m, "content", ""))) + len(getattr(m, "role", ""))
+            for m in messages
+        )
+        return max(1, total // 4)
+    if prompt is None:
+        return 1
+    if isinstance(prompt, str):
+        return max(1, len(prompt) // 4)
+    if isinstance(prompt, list):
+        if all(isinstance(i, int) for i in prompt):
+            return len(prompt)
+        return max(1, sum(len(str(i)) for i in prompt) // 4)
+    return max(1, len(str(prompt)) // 4)
+
+
+def render_dummy_text(n_tokens: int) -> str:
+    return "".join(DUMMY_TOKENS[: min(n_tokens, len(DUMMY_TOKENS))])