[proto] Add script to compile python protos (#11171)

python/sglang/srt/grpc/compile_proto.py

+#!/usr/bin/env python3
+"""
+Compile protobuf files for SGLang gRPC server.
+
+This script compiles .proto files to Python code using grpc_tools.protoc.
+It generates:
+- *_pb2.py (protobuf message classes)
+- *_pb2_grpc.py (gRPC service classes)
+- *_pb2.pyi (type hints for mypy/IDEs)
+
+Usage:
+    python compile_proto.py [--check] [--proto-file PROTO_FILE]
+
+Options:
+    --check         Check if regeneration is needed (exit 1 if needed)
+    --proto-file    Specify proto file (default: sglang_scheduler.proto)
+
+### Install Dependencies
+pip install grpcio grpcio-tools
+
+### Run Script
+cd python/sglang/srt/grpc
+python compile_proto.py
+"""
+
+import argparse
+import os
+import sys
+import subprocess
+import time
+from pathlib import Path
+from typing import List, Optional
+
+
+def get_file_mtime(path: Path) -> float:
+    """Get file modification time, return 0 if file doesn't exist."""
+    try:
+        return path.stat().st_mtime
+    except FileNotFoundError:
+        return 0.0
+
+
+def check_regeneration_needed(proto_file: Path, output_dir: Path) -> bool:
+    """Check if proto files are newer than generated files."""
+    proto_mtime = get_file_mtime(proto_file)
+
+    generated_files = [
+        output_dir / f"{proto_file.stem}_pb2.py",
+        output_dir / f"{proto_file.stem}_pb2_grpc.py",
+        output_dir / f"{proto_file.stem}_pb2.pyi",
+    ]
+
+    for gen_file in generated_files:
+        if get_file_mtime(gen_file) < proto_mtime:
+            return True
+
+    return False
+
+
+def compile_proto(proto_file: Path, output_dir: Path, verbose: bool = True) -> bool:
+    """Compile the protobuf file to Python."""
+
+    if not proto_file.exists():
+        print(f"Error: Proto file not found: {proto_file}")
+        return False
+
+    if verbose:
+        print(f"Found proto file: {proto_file}")
+
+    # Check if grpc_tools is available
+    try:
+        import grpc_tools.protoc
+    except ImportError:
+        print("Error: grpcio-tools not installed")
+        print("Install with: pip install grpcio-tools")
+        return False
+
+    # Compile command
+    cmd = [
+        sys.executable, "-m", "grpc_tools.protoc",
+        f"-I{proto_file.parent}",
+        f"--python_out={output_dir}",
+        f"--grpc_python_out={output_dir}",
+        f"--pyi_out={output_dir}",  # Generate type stubs
+        str(proto_file.name)
+    ]
+
+    if verbose:
+        print(f"Running: {' '.join(cmd)}")
+
+    # Run protoc
+    result = subprocess.run(cmd, capture_output=True, text=True, cwd=proto_file.parent)
+
+    if result.returncode != 0:
+        print(f"Error compiling proto:")
+        print(result.stderr)
+        if result.stdout:
+            print(result.stdout)
+        return False
+
+    # Verify generated files exist
+    generated_files = [
+        f"{proto_file.stem}_pb2.py",
+        f"{proto_file.stem}_pb2_grpc.py",
+        f"{proto_file.stem}_pb2.pyi"
+    ]
+
+    missing_files = []
+    for gen_file in generated_files:
+        if not (output_dir / gen_file).exists():
+            missing_files.append(gen_file)
+
+    if missing_files:
+        print(f"Error: Expected generated files not found: {missing_files}")
+        return False
+
+    if verbose:
+        print("Successfully compiled protobuf files:")
+        for gen_file in generated_files:
+            print(f"  - {output_dir}/{gen_file}")
+
+    # Fix imports in generated files
+    fix_imports(output_dir, proto_file.stem, verbose)
+
+    return True
+
+
+def fix_imports(output_dir: Path, proto_stem: str, verbose: bool = True) -> None:
+    """Fix imports in generated files to use relative imports."""
+    grpc_file = output_dir / f"{proto_stem}_pb2_grpc.py"
+
+    if grpc_file.exists():
+        content = grpc_file.read_text()
+        # Change absolute import to relative import
+        old_import = f"import {proto_stem}_pb2"
+        new_import = f"from . import {proto_stem}_pb2"
+
+        if old_import in content:
+            content = content.replace(old_import, new_import)
+            grpc_file.write_text(content)
+            if verbose:
+                print("Fixed imports in generated files")
+
+
+def add_generation_header(output_dir: Path, proto_stem: str) -> None:
+    """Add header to generated files indicating they are auto-generated."""
+    header = """# This file is auto-generated. Do not edit manually.
+# Regenerate with: python compile_proto.py
+
+"""
+
+    files_to_update = [
+        f"{proto_stem}_pb2.py",
+        f"{proto_stem}_pb2_grpc.py"
+    ]
+
+    for filename in files_to_update:
+        file_path = output_dir / filename
+        if file_path.exists():
+            content = file_path.read_text()
+            if not content.startswith("# This file is auto-generated"):
+                file_path.write_text(header + content)
+
+
+def main():
+    """Main entry point."""
+    parser = argparse.ArgumentParser(
+        description="Compile protobuf files for SGLang gRPC server",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog=__doc__
+    )
+
+    parser.add_argument(
+        "--check",
+        action="store_true",
+        help="Check if regeneration is needed (exit 1 if needed)"
+    )
+
+    parser.add_argument(
+        "--proto-file",
+        type=str,
+        default="sglang_scheduler.proto",
+        help="Proto file to compile (default: sglang_scheduler.proto)"
+    )
+
+    parser.add_argument(
+        "-v", "--verbose",
+        action="store_true",
+        default=True,
+        help="Verbose output (default: True)"
+    )
+
+    parser.add_argument(
+        "-q", "--quiet",
+        action="store_true",
+        help="Quiet mode (overrides verbose)"
+    )
+
+    args = parser.parse_args()
+
+    # Handle verbosity
+    verbose = args.verbose and not args.quiet
+
+    # Get paths
+    script_dir = Path(__file__).parent
+    proto_file = script_dir / args.proto_file
+    output_dir = script_dir
+
+    # Check mode
+    if args.check:
+        if check_regeneration_needed(proto_file, output_dir):
+            if verbose:
+                print("Proto files need regeneration")
+            sys.exit(1)
+        else:
+            if verbose:
+                print("Generated files are up to date")
+            sys.exit(0)
+
+    # Compile mode
+    success = compile_proto(proto_file, output_dir, verbose)
+
+    if success:
+        # Add generation headers
+        add_generation_header(output_dir, proto_file.stem)
+
+        if verbose:
+            print("\n✅ Protobuf compilation successful!")
+            print("Generated files are ready for use")
+    else:
+        if verbose:
+            print("\n❌ Protobuf compilation failed!")
+        sys.exit(1)
+
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file