lmcache--lmcache
wehub-resource-sync
85 行
2.5 KiB
Python
85 行
2.5 KiB
Python
#!/usr/bin/env python3
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
# This script implements a health check probe for the LMCache server.
|
|
# It performs the following functions:
|
|
# 1. Sends a health check request to the LMCache server
|
|
# 2. Verifies the server's response indicates successful operation
|
|
# 3. Serves as a Kubernetes liveness probe to monitor server health
|
|
# 4. Helps ensure the server is running and responsive
|
|
|
|
# Standard
|
|
import socket
|
|
import sys
|
|
|
|
# Third Party
|
|
import torch
|
|
|
|
# First Party
|
|
from lmcache.v1.protocol import (
|
|
CacheEngineKey,
|
|
ClientCommand,
|
|
ClientMetaMessage,
|
|
MemoryFormat,
|
|
ServerMetaMessage,
|
|
ServerReturnCode,
|
|
)
|
|
|
|
|
|
def main():
|
|
if len(sys.argv) != 3:
|
|
print("Usage: health_probe.py <host> <port>", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
host = sys.argv[1]
|
|
port = int(sys.argv[2])
|
|
try:
|
|
# Create connection with timeout and ensure proper cleanup with context manager
|
|
with socket.create_connection((host, port), timeout=5) as s:
|
|
# Create and send health check message
|
|
msg = ClientMetaMessage(
|
|
ClientCommand.HEALTH,
|
|
key=CacheEngineKey(
|
|
model_name="",
|
|
world_size=0,
|
|
worker_id=0,
|
|
chunk_hash=0,
|
|
dtype=torch.float16,
|
|
),
|
|
length=0,
|
|
fmt=MemoryFormat(1),
|
|
dtype=torch.float16,
|
|
shape=torch.Size((0, 0, 0, 0)),
|
|
)
|
|
s.sendall(msg.serialize())
|
|
|
|
# Receive and parse response
|
|
resp = s.recv(ServerMetaMessage.packlength())
|
|
if not resp:
|
|
print("No response received from server", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
# Parse the response message
|
|
meta = ServerMetaMessage.deserialize(resp)
|
|
|
|
# Check if server responded with success
|
|
if meta.code == ServerReturnCode.SUCCESS:
|
|
sys.exit(0)
|
|
else:
|
|
print(f"Server returned error code: {meta.code}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
except socket.timeout:
|
|
print("Connection timed out", file=sys.stderr)
|
|
sys.exit(1)
|
|
except ConnectionRefusedError:
|
|
print("Connection refused - server may not be running", file=sys.stderr)
|
|
sys.exit(1)
|
|
except Exception as e:
|
|
print(f"Error during health check: {str(e)}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|