246 lines
6.8 KiB
Python
Executable File
246 lines
6.8 KiB
Python
Executable File
#!/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==1.74.0" "grpcio-tools==1.74.0"
|
|
|
|
### Run Script
|
|
cd python/sglang/srt/grpc
|
|
python compile_proto.py
|
|
"""
|
|
|
|
|
|
import argparse
|
|
import subprocess
|
|
import sys
|
|
from importlib.metadata import version
|
|
from pathlib import Path
|
|
|
|
GRPC_VERSION = "1.74.0"
|
|
|
|
|
|
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(
|
|
f'Install with: pip install "grpcio-tools=={GRPC_VERSION}" "grpcio=={GRPC_VERSION}"'
|
|
)
|
|
return False
|
|
|
|
grpc_tools_version = version("grpcio-tools")
|
|
grpc_version = version("grpcio")
|
|
if grpc_tools_version != GRPC_VERSION or grpc_version != GRPC_VERSION:
|
|
raise RuntimeError(
|
|
f"Error: grpcio-tools version {grpc_tools_version} and grpcio version {grpc_version} detected, but {GRPC_VERSION} is required."
|
|
)
|
|
|
|
# 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()
|