diff --git a/python/sglang/srt/grpc/compile_proto.py b/python/sglang/srt/grpc/compile_proto.py new file mode 100644 index 000000000..51721446b --- /dev/null +++ b/python/sglang/srt/grpc/compile_proto.py @@ -0,0 +1,237 @@ +#!/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