Unverified Commit 7a05d2dc authored by TJian's avatar TJian Committed by GitHub
Browse files

[CI] [ROCm] Fix `tests/entrypoints/test_grpc_server.py` on ROCm (#31970)


Signed-off-by: default avatartjtanaa <tunjian.tan@embeddedllm.com>
parent a1648c40
...@@ -16,3 +16,4 @@ runai-model-streamer[s3,gcs]==0.15.3 ...@@ -16,3 +16,4 @@ runai-model-streamer[s3,gcs]==0.15.3
conch-triton-kernels==1.2.1 conch-triton-kernels==1.2.1
timm>=1.0.17 timm>=1.0.17
fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@d6f998a03432b2452f8de2bb5cefb5af9795d459 fastsafetensors @ git+https://github.com/foundation-model-stack/fastsafetensors.git@d6f998a03432b2452f8de2bb5cefb5af9795d459
grpcio-tools>=1.76.0
\ No newline at end of file
...@@ -19,6 +19,7 @@ from packaging.version import Version, parse ...@@ -19,6 +19,7 @@ from packaging.version import Version, parse
from setuptools import Extension, setup from setuptools import Extension, setup
from setuptools.command.build_ext import build_ext from setuptools.command.build_ext import build_ext
from setuptools.command.build_py import build_py from setuptools.command.build_py import build_py
from setuptools.command.develop import develop
from setuptools_scm import get_version from setuptools_scm import get_version
from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME from torch.utils.cpp_extension import CUDA_HOME, ROCM_HOME
...@@ -147,6 +148,14 @@ class BuildPyAndGenerateGrpc(build_py): ...@@ -147,6 +148,14 @@ class BuildPyAndGenerateGrpc(build_py):
super().run() super().run()
class DevelopAndGenerateGrpc(develop):
"""Develop mode that also generates gRPC stubs from proto files."""
def run(self):
compile_grpc_protos()
super().run()
class CMakeExtension(Extension): class CMakeExtension(Extension):
def __init__(self, name: str, cmake_lists_dir: str = ".", **kwa) -> None: def __init__(self, name: str, cmake_lists_dir: str = ".", **kwa) -> None:
super().__init__(name, sources=[], py_limited_api=not is_freethreaded(), **kwa) super().__init__(name, sources=[], py_limited_api=not is_freethreaded(), **kwa)
...@@ -950,13 +959,17 @@ if _no_device(): ...@@ -950,13 +959,17 @@ if _no_device():
ext_modules = [] ext_modules = []
if not ext_modules: if not ext_modules:
cmdclass = {"build_py": BuildPyAndGenerateGrpc} cmdclass = {
"build_py": BuildPyAndGenerateGrpc,
"develop": DevelopAndGenerateGrpc,
}
else: else:
cmdclass = { cmdclass = {
"build_ext": precompiled_build_ext "build_ext": precompiled_build_ext
if envs.VLLM_USE_PRECOMPILED if envs.VLLM_USE_PRECOMPILED
else cmake_build_ext, else cmake_build_ext,
"build_py": BuildPyAndGenerateGrpc, "build_py": BuildPyAndGenerateGrpc,
"develop": DevelopAndGenerateGrpc,
} }
setup( setup(
......
...@@ -29,7 +29,7 @@ def find_free_port() -> int: ...@@ -29,7 +29,7 @@ def find_free_port() -> int:
return port return port
async def wait_for_server(port: int, timeout: float = 30.0) -> bool: async def wait_for_server(port: int, timeout: float = 60.0) -> bool:
"""Wait for the gRPC server to be ready by trying health checks.""" """Wait for the gRPC server to be ready by trying health checks."""
start_time = time.time() start_time = time.time()
print("waiting for server to start...") print("waiting for server to start...")
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment