Support patching source code (#19561)

python/sglang/srt/debug_utils/source_patcher/__init__.py

@@ -0,0 +1,12 @@
+from sglang.srt.debug_utils.source_patcher.code_patcher import (
+    CodePatcher,
+    apply_patches_from_config,
+    patch_function,
+)
+from sglang.srt.debug_utils.source_patcher.types import (
+    EditSpec,
+    PatchApplicationError,
+    PatchConfig,
+    PatchSpec,
+    PatchState,
+)

python/sglang/srt/debug_utils/source_patcher/code_patcher.py

@@ -0,0 +1,188 @@
+import importlib
+import inspect
+import textwrap
+import types
+from collections.abc import Callable
+from typing import Any, Optional
+
+import yaml
+
+from sglang.srt.debug_utils.source_patcher.source_editor import apply_edits
+from sglang.srt.debug_utils.source_patcher.types import (
+    EditSpec,
+    PatchConfig,
+    PatchSpec,
+    PatchState,
+)
+
+
+def apply_patches_from_config(
+    yaml_content: str,
+    *,
+    extra_imports: Optional[list[str]] = None,
+) -> list[PatchState]:
+    """Parse a YAML config string and apply all patches.
+
+    Args:
+        yaml_content: YAML string with patch specifications.
+        extra_imports: Import lines inserted once at the top of each patched
+            function body (e.g. ["from pkg import foo"]).  The caller (dumper)
+            uses this so users don't have to write boilerplate in YAML.
+    """
+    raw: dict[str, Any] = yaml.safe_load(yaml_content)
+    config: PatchConfig = PatchConfig(**raw)
+
+    if extra_imports:
+        config = _inject_preamble(config=config, extra_imports=extra_imports)
+
+    return _apply_specs(config.patches)
+
+
+class CodePatcher:
+    """Context manager that patches functions on enter and restores on exit."""
+
+    def __init__(self, *, patches: list[PatchSpec]) -> None:
+        self._patches = patches
+        self._states: list[PatchState] = []
+
+    def __enter__(self) -> "CodePatcher":
+        self._states = _apply_specs(self._patches)
+        return self
+
+    def __exit__(
+        self,
+        exc_type: Optional[type],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[Any],
+    ) -> None:
+        for state in reversed(self._states):
+            state.restore()
+        self._states.clear()
+
+
+def patch_function(
+    *,
+    target: Callable[..., Any],
+    edits: list[EditSpec],
+    preamble: str = "",
+) -> PatchState:
+    """Patch a function by modifying its source and replacing __code__.
+
+    1. inspect.getsource -> get original source
+    2. apply_edits -> modify source text
+    3. optionally prepend preamble (e.g. import lines) inside the function body
+    4. compile + exec -> get new code object
+    5. replace target.__code__
+
+    Returns PatchState that can restore the original code.
+    """
+    original_code: types.CodeType = target.__code__
+
+    source: str = inspect.getsource(target)
+    modified_source: str = apply_edits(source=source, edits=edits)
+    modified_source = textwrap.dedent(modified_source)
+
+    if preamble.strip():
+        modified_source = _insert_preamble(source=modified_source, preamble=preamble)
+
+    code: types.CodeType = compile(modified_source, inspect.getfile(target), "exec")
+    temp_namespace: dict[str, Any] = {}
+    exec(code, target.__globals__, temp_namespace)
+
+    new_fn: Any = temp_namespace[target.__name__]
+    target.__code__ = new_fn.__code__
+
+    return PatchState(target_fn=target, original_code=original_code)
+
+
+# --------------------------------- private ---------------------------------
+
+
+def _apply_specs(specs: list[PatchSpec]) -> list[PatchState]:
+    states: list[PatchState] = []
+    for spec in specs:
+        target_fn: Callable[..., Any] = _resolve_target(spec.target)
+        print(f"[source_patcher] patching {spec.target}")
+        state: PatchState = patch_function(
+            target=target_fn, edits=spec.edits, preamble=spec.preamble
+        )
+        states.append(state)
+    return states
+
+
+def _inject_preamble(*, config: PatchConfig, extra_imports: list[str]) -> PatchConfig:
+    """Set preamble on every PatchSpec so imports are inserted once at function top."""
+    import_block: str = "\n".join(extra_imports)
+    new_patches: list[PatchSpec] = []
+
+    for spec in config.patches:
+        existing: str = spec.preamble
+        combined: str = (
+            existing + "\n" + import_block if existing.strip() else import_block
+        )
+        new_patches.append(
+            PatchSpec(target=spec.target, edits=spec.edits, preamble=combined)
+        )
+
+    return PatchConfig(patches=new_patches)
+
+
+def _insert_preamble(*, source: str, preamble: str) -> str:
+    """Insert preamble lines right after the function signature (and optional docstring)."""
+    lines: list[str] = source.splitlines()
+
+    signature_end: int = _find_signature_end(lines)
+
+    body_start: int = signature_end + 1
+    body_indent: str = ""
+    for i in range(body_start, len(lines)):
+        if lines[i].strip():
+            body_indent = " " * (len(lines[i]) - len(lines[i].lstrip()))
+            body_start = i
+            break
+
+    preamble_lines: list[str] = [
+        body_indent + pl for pl in preamble.strip().splitlines()
+    ]
+    return "\n".join(lines[:body_start] + preamble_lines + lines[body_start:])
+
+
+def _find_signature_end(lines: list[str]) -> int:
+    """Find the line index where the function signature ends (the line with trailing colon)."""
+    for i, line in enumerate(lines):
+        if line.rstrip().endswith(":"):
+            return i
+    return 0
+
+
+def _resolve_target(qualified_name: str) -> Callable[..., Any]:
+    """Resolve 'pkg.mod.Class.method' to the actual function object.
+
+    Tries progressively shorter module paths from right to left,
+    then uses getattr for the remaining attribute chain.
+    """
+    parts: list[str] = qualified_name.split(".")
+
+    target: Any = None
+    for split_idx in range(len(parts), 0, -1):
+        module_path: str = ".".join(parts[:split_idx])
+        try:
+            target = importlib.import_module(module_path)
+            attr_parts: list[str] = parts[split_idx:]
+            break
+        except ImportError:
+            continue
+    else:
+        raise ImportError(f"could not import any module prefix of '{qualified_name}'")
+
+    for attr_name in attr_parts:
+        target = getattr(target, attr_name)
+
+    if isinstance(target, classmethod):
+        target = target.__func__
+    if not callable(target):
+        raise TypeError(
+            f"resolved target '{qualified_name}' is not callable: {type(target)}"
+        )
+
+    return target

python/sglang/srt/debug_utils/source_patcher/source_editor.py

@@ -0,0 +1,108 @@
+from sglang.srt.debug_utils.source_patcher.types import EditSpec, PatchApplicationError
+
+
+def apply_edits(*, source: str, edits: list[EditSpec]) -> str:
+    """Apply a sequence of match/replacement edits to source text.
+
+    Each edit is applied sequentially so later edits see the result of earlier ones.
+    """
+    result: str = source
+    for edit in edits:
+        result = _apply_single_edit(source=result, edit=edit)
+    return result
+
+
+def _apply_single_edit(*, source: str, edit: EditSpec) -> str:
+    """Apply a single match/replacement edit to the source text."""
+    match_text: str = edit.match.strip()
+    if not match_text:
+        raise PatchApplicationError("empty match text")
+
+    source_lines: list[str] = source.splitlines()
+    match_lines: list[str] = match_text.splitlines()
+
+    start_idx: int = _find_match(source_lines=source_lines, match_lines=match_lines)
+    match_len: int = len(match_lines)
+
+    original_indent: int = _leading_spaces(source_lines[start_idx])
+
+    effective_replacement: str = _resolve_replacement(edit=edit, match_text=match_text)
+    replacement_lines: list[str] = (
+        effective_replacement.splitlines() if effective_replacement else []
+    )
+    aligned: list[str] = _realign_replacement(
+        replacement_lines=replacement_lines, original_indent=original_indent
+    )
+    new_lines: list[str] = (
+        source_lines[:start_idx] + aligned + source_lines[start_idx + match_len :]
+    )
+
+    trailing_newline: str = "\n" if source.endswith("\n") else ""
+    return "\n".join(new_lines) + trailing_newline
+
+
+def _resolve_replacement(*, edit: EditSpec, match_text: str) -> str:
+    """Return the effective replacement text, handling replacement, prepend, and append modes."""
+    if edit.prepend.strip():
+        return edit.prepend.strip() + "\n" + match_text
+    if edit.append.strip():
+        return match_text + "\n" + edit.append.strip()
+    return edit.replacement.strip()
+
+
+def _find_match(*, source_lines: list[str], match_lines: list[str]) -> int:
+    """Find the start index of match_lines in source_lines (strip-compared).
+
+    Returns the index of the first matching line.
+    Raises PatchApplicationError if not found or found multiple times.
+    """
+    stripped_source: list[str] = [line.strip() for line in source_lines]
+    stripped_match: list[str] = [line.strip() for line in match_lines]
+    match_len: int = len(stripped_match)
+
+    found_indices: list[int] = [
+        i
+        for i in range(len(stripped_source) - match_len + 1)
+        if stripped_source[i : i + match_len] == stripped_match
+    ]
+
+    if len(found_indices) == 0:
+        preview: str = "\n".join(match_lines)
+        raise PatchApplicationError(f"match text not found in source:\n{preview}")
+    if len(found_indices) > 1:
+        preview = "\n".join(match_lines)
+        raise PatchApplicationError(
+            f"match text found multiple times ({len(found_indices)} occurrences) in source:\n{preview}"
+        )
+
+    return found_indices[0]
+
+
+def _realign_replacement(
+    *, replacement_lines: list[str], original_indent: int
+) -> list[str]:
+    """Realign replacement lines to the original indentation level.
+
+    Strategy:
+    - Take the leading spaces of the first non-empty replacement line as base_indent
+    - For each replacement line: remove base_indent, add original_indent
+    """
+    non_empty: list[str] = [line for line in replacement_lines if line.strip()]
+    if not non_empty:
+        return []
+
+    base_indent: int = _leading_spaces(non_empty[0])
+    result: list[str] = []
+
+    for line in replacement_lines:
+        if not line.strip():
+            result.append("")
+        else:
+            stripped = line[min(base_indent, len(line) - len(line.lstrip())) :]
+            result.append(" " * original_indent + stripped)
+
+    return result
+
+
+def _leading_spaces(line: str) -> int:
+    return len(line) - len(line.lstrip(" "))

python/sglang/srt/debug_utils/source_patcher/types.py

@@ -0,0 +1,63 @@
+import types
+from collections.abc import Callable
+from typing import Any
+
+from pydantic import BaseModel, ConfigDict, model_validator
+
+
+class PatchApplicationError(Exception):
+    """match text not found or not unique in source."""
+
+
+class _StrictBase(BaseModel):
+    model_config = ConfigDict(extra="forbid")
+
+
+class EditSpec(_StrictBase):
+    """Specify one edit: replace, prepend before, or append after the matched text.
+
+    Use ``replacement`` to substitute the matched text (empty string = delete).
+    Use ``prepend`` to keep the matched text and add lines before it.
+    Use ``append`` to keep the matched text and add lines after it.
+    Only one of ``replacement``, ``prepend``, and ``append`` may be set.
+    """
+
+    match: str
+    replacement: str = ""
+    prepend: str = ""
+    append: str = ""
+
+    @model_validator(mode="after")
+    def _check_modes_mutually_exclusive(self) -> "EditSpec":
+        active: list[str] = [
+            name
+            for name in ("replacement", "prepend", "append")
+            if getattr(self, name).strip()
+        ]
+        if len(active) > 1:
+            raise ValueError(
+                f"only one of 'replacement', 'prepend', 'append' may be set, "
+                f"got: {', '.join(active)}"
+            )
+        return self
+
+
+class PatchSpec(_StrictBase):
+    target: str
+    edits: list[EditSpec]
+    preamble: str = ""
+
+
+class PatchConfig(_StrictBase):
+    patches: list[PatchSpec]
+
+
+class PatchState:
+    def __init__(
+        self, *, target_fn: Callable[..., Any], original_code: types.CodeType
+    ) -> None:
+        self.target_fn = target_fn
+        self.original_code = original_code
+
+    def restore(self) -> None:
+        self.target_fn.__code__ = self.original_code