Hanrui / SpecForge /tests /utils.py
Lekr0's picture
Add files using upload-large-folder tool
2d67aa6 verified
import os
import socket
import subprocess
import time
import requests
from sglang.utils import print_highlight
def is_port_in_use(port: int) -> bool:
"""Check if a port is in use"""
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
try:
s.bind(("localhost", port))
return False
except OSError:
return True
def get_available_port():
# get a random available port
# and try to find a port that is not in use
for port in range(10000, 65535):
if not is_port_in_use(port):
return port
raise RuntimeError("No available port found")
def execute_shell_command(
command: str, disable_proxy: bool = False, enable_hf_mirror: bool = False
):
"""
Execute a shell command and return its process handle.
"""
command = command.replace("\\\n", " ").replace("\\", " ")
parts = command.split()
env = os.environ.copy()
if disable_proxy:
env.pop("http_proxy", None)
env.pop("https_proxy", None)
env.pop("no_proxy", None)
env.pop("HTTP_PROXY", None)
env.pop("HTTPS_PROXY", None)
env.pop("NO_PROXY", None)
if enable_hf_mirror:
env["HF_ENDPOINT"] = "https://hf-mirror.com"
return subprocess.Popen(parts, text=True, stderr=subprocess.STDOUT, env=env)
def wait_for_server(
base_url: str, timeout: int = None, disable_proxy: bool = False
) -> None:
"""Wait for the server to be ready by polling the /v1/models endpoint.
Args:
base_url: The base URL of the server
timeout: Maximum time to wait in seconds. None means wait forever.
"""
start_time = time.perf_counter()
if disable_proxy:
http_proxy = os.environ.pop("http_proxy", None)
https_proxy = os.environ.pop("https_proxy", None)
no_proxy = os.environ.pop("no_proxy", None)
http_proxy_capitalized = os.environ.pop("HTTP_PROXY", None)
https_proxy_capitalized = os.environ.pop("HTTPS_PROXY", None)
no_proxy_capitalized = os.environ.pop("NO_PROXY", None)
while True:
try:
response = requests.get(
f"{base_url}/v1/models",
headers={"Authorization": "Bearer None"},
)
if response.status_code == 200:
time.sleep(5)
print_highlight(
"""\n
NOTE: Typically, the server runs in a separate terminal.
In this notebook, we run the server and notebook code together, so their outputs are combined.
To improve clarity, the server logs are displayed in the original black color, while the notebook outputs are highlighted in blue.
To reduce the log length, we set the log level to warning for the server, the default log level is info.
We are running those notebooks in a CI environment, so the throughput is not representative of the actual performance.
"""
)
break
if timeout and time.perf_counter() - start_time > timeout:
raise TimeoutError("Server did not become ready within timeout period")
except requests.exceptions.RequestException:
time.sleep(1)
if disable_proxy:
if http_proxy:
os.environ["http_proxy"] = http_proxy
if https_proxy:
os.environ["https_proxy"] = https_proxy
if no_proxy:
os.environ["no_proxy"] = no_proxy
if http_proxy_capitalized:
os.environ["HTTP_PROXY"] = http_proxy_capitalized
if https_proxy_capitalized:
os.environ["HTTPS_PROXY"] = https_proxy_capitalized
if no_proxy_capitalized:
os.environ["NO_PROXY"] = no_proxy_capitalized