wehub-resource-sync
666 lines
26 KiB
Python
666 lines
26 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# Standard
|
||
from typing import TYPE_CHECKING, Optional
|
||
import asyncio
|
||
import threading
|
||
|
||
# Third Party
|
||
import msgspec
|
||
import zmq
|
||
import zmq.asyncio
|
||
|
||
# First Party
|
||
from lmcache.logging import init_logger
|
||
from lmcache.v1.cache_controller.full_sync_sender import FullSyncSender
|
||
from lmcache.v1.cache_controller.message import (
|
||
BatchedP2PLookupMsg,
|
||
BatchedP2PLookupRetMsg,
|
||
ClearWorkerMsg,
|
||
ClearWorkerRetMsg,
|
||
CompressWorkerMsg,
|
||
CompressWorkerRetMsg,
|
||
DecompressWorkerMsg,
|
||
DecompressWorkerRetMsg,
|
||
DeRegisterMsg,
|
||
ErrorMsg,
|
||
FullSyncStartMsg,
|
||
FullSyncStartRetMsg,
|
||
FullSyncStatusMsg,
|
||
FullSyncStatusRetMsg,
|
||
HealthWorkerMsg,
|
||
HealthWorkerRetMsg,
|
||
HeartbeatMsg,
|
||
HeartbeatRetMsg,
|
||
MoveWorkerMsg,
|
||
MoveWorkerRetMsg,
|
||
Msg,
|
||
PinWorkerMsg,
|
||
PinWorkerRetMsg,
|
||
RegisterMsg,
|
||
RegisterRetMsg,
|
||
WorkerMsg,
|
||
WorkerReqMsg,
|
||
WorkerReqRetMsg,
|
||
)
|
||
from lmcache.v1.config import LMCacheEngineConfig
|
||
from lmcache.v1.metadata import LMCacheMetadata
|
||
from lmcache.v1.rpc_utils import (
|
||
DEFAULT_SOCKET_RECV_TIMEOUT_MS,
|
||
DEFAULT_SOCKET_SEND_TIMEOUT_MS,
|
||
close_zmq_socket,
|
||
get_ip,
|
||
get_zmq_context,
|
||
get_zmq_socket,
|
||
get_zmq_socket_with_timeout,
|
||
)
|
||
|
||
if TYPE_CHECKING:
|
||
# First Party
|
||
from lmcache.v1.cache_engine import LMCacheEngine
|
||
|
||
logger = init_logger(__name__)
|
||
|
||
|
||
class LMCacheWorker:
|
||
"""
|
||
LMCache Worker class to handle the execution of cache operations.
|
||
This class is responsible for receiving requests from the executor and
|
||
executing the corresponding operations on the LMCache engine.
|
||
Each worker is associated with a specific LMCache instance and a worker id.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
config: LMCacheEngineConfig,
|
||
metadata: LMCacheMetadata,
|
||
lmcache_engine: "LMCacheEngine",
|
||
):
|
||
# TODO (Jiayi): "instance_id" might not be needed anymore.
|
||
# Please consider removing it.
|
||
self.config = config
|
||
self.lmcache_instance_id = config.lmcache_instance_id
|
||
if self.lmcache_instance_id is None:
|
||
raise ValueError(
|
||
"lmcache_instance_id is required when enable_controller=True"
|
||
)
|
||
self.lmcache_engine = lmcache_engine
|
||
self.worker_id = metadata.worker_id
|
||
|
||
self.context = get_zmq_context()
|
||
|
||
# Load timeout configurations from extra_config (in milliseconds)
|
||
self.socket_recv_timeout_ms = config.get_extra_config_value(
|
||
"worker_socket_recv_timeout_ms", DEFAULT_SOCKET_RECV_TIMEOUT_MS
|
||
)
|
||
self.socket_send_timeout_ms = config.get_extra_config_value(
|
||
"worker_socket_send_timeout_ms", DEFAULT_SOCKET_SEND_TIMEOUT_MS
|
||
)
|
||
|
||
if config.controller_pull_url is None:
|
||
raise ValueError(
|
||
"controller_pull_url is required when enable_controller=True"
|
||
)
|
||
|
||
controller_pull_url = config.controller_pull_url
|
||
self.push_socket = get_zmq_socket(
|
||
self.context,
|
||
controller_pull_url,
|
||
protocol="tcp",
|
||
role=zmq.PUSH, # type: ignore[attr-defined]
|
||
bind_or_connect="connect",
|
||
)
|
||
|
||
if config.controller_reply_url is not None:
|
||
self.controller_rep_url = config.controller_reply_url
|
||
self._create_req_socket()
|
||
|
||
# Heartbeat socket will be created dynamically after register
|
||
# based on heartbeat_url returned from controller
|
||
self.heartbeat_socket: Optional[zmq.asyncio.Socket] = None
|
||
self.controller_heartbeat_url: Optional[str] = None
|
||
|
||
# metadata.world_size comes from vLLM's parallel_config.world_size.
|
||
# For MLA models, vLLM divides this by tp_size (e.g. TP=8 PP=1 on
|
||
# 8 GPUs → world_size=1), so it may be much smaller than the total
|
||
# GPU count. For non-MLA models it equals TP × PP.
|
||
#
|
||
# get_lmcache_worker_ids() decides which workers run an LMCache
|
||
# instance: [0] for MLA (only one worker needed since KV caches
|
||
# are not TP-sharded), or range(world_size) for non-MLA.
|
||
#
|
||
# We use >= because extra ports are harmless — port selection
|
||
# indexes by worker_id or lmcache_worker_ids position, so
|
||
# trailing entries are never bound to sockets.
|
||
lmcache_worker_ids = config.get_lmcache_worker_ids(
|
||
metadata.use_mla, metadata.world_size
|
||
)
|
||
if not lmcache_worker_ids:
|
||
# start lmcache worker on all ranks;
|
||
# need at least one port per rank (world_size)
|
||
if len(config.lmcache_worker_ports) < metadata.world_size:
|
||
raise ValueError(
|
||
f"lmcache_worker_ports must have at least {metadata.world_size} "
|
||
f"port(s) (world_size), got {len(config.lmcache_worker_ports)}"
|
||
)
|
||
lmcache_worker_port = config.lmcache_worker_ports[self.worker_id]
|
||
else:
|
||
# start lmcache worker on given worker ids;
|
||
# need at least one port per explicitly listed worker
|
||
if len(config.lmcache_worker_ports) < len(lmcache_worker_ids):
|
||
raise ValueError(
|
||
f"lmcache_worker_ports must have at least "
|
||
f"{len(lmcache_worker_ids)} "
|
||
f"port(s) (one per lmcache worker), "
|
||
f"got {len(config.lmcache_worker_ports)}"
|
||
)
|
||
index = lmcache_worker_ids.index(self.worker_id)
|
||
lmcache_worker_port = config.lmcache_worker_ports[index]
|
||
|
||
self.lmcache_worker_internal_url = f"*:{lmcache_worker_port}"
|
||
self.lmcache_worker_ip = get_ip()
|
||
self.lmcache_worker_port = lmcache_worker_port
|
||
|
||
self.p2p_init_url = None
|
||
if config.enable_p2p:
|
||
self.p2p_host = config.p2p_host
|
||
self.p2p_init_port = config.p2p_init_ports[self.worker_id]
|
||
self.p2p_init_url = f"{self.p2p_host}:{self.p2p_init_port}"
|
||
|
||
self.reply_socket = get_zmq_socket(
|
||
self.context,
|
||
self.lmcache_worker_internal_url,
|
||
protocol="tcp",
|
||
role=zmq.REP, # type: ignore[attr-defined]
|
||
bind_or_connect="bind",
|
||
)
|
||
|
||
logger.info(f"Reply socket established at {self.lmcache_worker_internal_url}")
|
||
|
||
self.loop = asyncio.new_event_loop()
|
||
self.thread = threading.Thread(
|
||
target=self.loop.run_forever, daemon=True, name="lmcache-worker-thread"
|
||
)
|
||
self.thread.start()
|
||
asyncio.run_coroutine_threadsafe(self.start_all(), self.loop)
|
||
|
||
self.msg_queue: asyncio.Queue[WorkerMsg] = asyncio.Queue()
|
||
|
||
# Full sync sender (initialized lazily when needed)
|
||
self._full_sync_sender: Optional["FullSyncSender"] = None
|
||
|
||
async def register(self):
|
||
"""
|
||
Register the lmcache worker with the controller via DEALER-ROUTER.
|
||
|
||
This method sends a RegisterMsg and waits for RegisterRetMsg
|
||
which contains extra_config (e.g., heartbeat_url).
|
||
"""
|
||
assert self.lmcache_instance_id is not None
|
||
logger.info(
|
||
"Registering lmcache instance-worker: %s",
|
||
(self.lmcache_instance_id, self.worker_id),
|
||
)
|
||
|
||
register_msg = RegisterMsg(
|
||
instance_id=self.lmcache_instance_id,
|
||
worker_id=self.worker_id,
|
||
ip=self.lmcache_worker_ip,
|
||
port=self.lmcache_worker_port,
|
||
peer_init_url=self.p2p_init_url,
|
||
)
|
||
|
||
# Send via DEALER socket (empty frame + payload) and wait for response
|
||
try:
|
||
await self.req_socket.send_multipart(
|
||
[b"", msgspec.msgpack.encode(register_msg)]
|
||
)
|
||
# DEALER receives: [empty_frame, payload]
|
||
frames = await self.req_socket.recv_multipart()
|
||
serialized_ret_msg = frames[-1]
|
||
ret_msg = msgspec.msgpack.decode(serialized_ret_msg, type=Msg)
|
||
|
||
if isinstance(ret_msg, RegisterRetMsg):
|
||
self._process_register_response(ret_msg)
|
||
else:
|
||
logger.warning("Unexpected register response type: %s", type(ret_msg))
|
||
except zmq.ZMQError as e:
|
||
logger.error("Failed to register with controller: %s", e)
|
||
raise
|
||
|
||
def _process_register_response(self, ret_msg: RegisterRetMsg):
|
||
"""Process RegisterRetMsg and initialize components based on extra_config."""
|
||
extra_config = ret_msg.extra_config
|
||
|
||
# Initialize heartbeat socket if heartbeat_url is provided
|
||
heartbeat_url = extra_config.get("heartbeat_url")
|
||
if heartbeat_url:
|
||
logger.info("Received heartbeat_url from controller: %s", heartbeat_url)
|
||
self.controller_heartbeat_url = heartbeat_url
|
||
self._create_heartbeat_socket()
|
||
else:
|
||
logger.info("No dedicated heartbeat_url provided by controller")
|
||
|
||
def deregister(self):
|
||
"""
|
||
De-register the lmcache worker from the controller.
|
||
"""
|
||
assert self.lmcache_instance_id is not None
|
||
self.put_msg(
|
||
DeRegisterMsg(
|
||
instance_id=self.lmcache_instance_id,
|
||
worker_id=self.worker_id,
|
||
ip=self.lmcache_worker_ip,
|
||
port=self.lmcache_worker_port,
|
||
)
|
||
)
|
||
|
||
async def async_put_and_wait_msg(
|
||
self,
|
||
msg: WorkerReqMsg,
|
||
) -> WorkerReqRetMsg:
|
||
"""
|
||
Send a message to the controller and wait for the response.
|
||
|
||
This method handles different types of WorkerReqMsg using appropriate sockets:
|
||
- HeartbeatMsg: Uses dedicated heartbeat socket
|
||
- Other messages (RegisterMsg, BatchedP2PLookupMsg, FullSyncStartMsg,
|
||
FullSyncStatusMsg): Uses DEALER socket (req_socket)
|
||
|
||
Note: With DEALER-ROUTER mode, we no longer need to separate FullSync
|
||
messages to heartbeat socket since DEALER supports async concurrent requests.
|
||
"""
|
||
# Send heartbeat via dedicated heartbeat socket
|
||
if isinstance(msg, HeartbeatMsg):
|
||
return await self._send_heartbeat_msg(msg)
|
||
|
||
# Send other messages via DEALER socket
|
||
try:
|
||
# DEALER socket: send [empty_frame, payload]
|
||
await self.req_socket.send_multipart([b"", msgspec.msgpack.encode(msg)])
|
||
frames = await self.req_socket.recv_multipart()
|
||
# DEALER receives: [empty_frame, payload]
|
||
serialized_ret_msg = frames[-1]
|
||
ret_msg = msgspec.msgpack.decode(serialized_ret_msg, type=Msg)
|
||
return ret_msg
|
||
except zmq.Again as e:
|
||
logger.error("Timeout occurred, recreating socket. Error: %s", e)
|
||
self._recreate_req_socket()
|
||
return self._on_request_failure(msg)
|
||
except zmq.ZMQError as e:
|
||
logger.error("ZMQ error occurred, recreating socket. Error: %s", e)
|
||
self._recreate_req_socket()
|
||
return self._on_request_failure(msg)
|
||
except Exception as e:
|
||
logger.error("Error happens in lmcache worker req_socket. Error: %s", e)
|
||
return self._on_request_failure(msg)
|
||
|
||
def _create_req_socket(self):
|
||
self.req_socket = get_zmq_socket_with_timeout(
|
||
self.context,
|
||
self.controller_rep_url,
|
||
"tcp",
|
||
zmq.DEALER, # type: ignore[attr-defined]
|
||
"connect",
|
||
self.socket_recv_timeout_ms,
|
||
self.socket_send_timeout_ms,
|
||
)
|
||
|
||
def _create_heartbeat_socket(self):
|
||
logger.info(
|
||
"Creating heartbeat socket to connect to: %s, "
|
||
"recv_timeout: %dms, send_timeout: %dms",
|
||
self.controller_heartbeat_url,
|
||
self.socket_recv_timeout_ms,
|
||
self.socket_send_timeout_ms,
|
||
)
|
||
self.heartbeat_socket = get_zmq_socket_with_timeout(
|
||
self.context,
|
||
self.controller_heartbeat_url,
|
||
"tcp",
|
||
zmq.DEALER, # type: ignore[attr-defined]
|
||
"connect",
|
||
self.socket_recv_timeout_ms,
|
||
self.socket_send_timeout_ms,
|
||
)
|
||
|
||
def _recreate_heartbeat_socket(self):
|
||
try:
|
||
self.heartbeat_socket.close(linger=0)
|
||
except Exception as e:
|
||
logger.error("Error closing heartbeat socket: %s", e)
|
||
self._create_heartbeat_socket()
|
||
|
||
def _recreate_req_socket(self):
|
||
try:
|
||
self.req_socket.close(linger=0)
|
||
except Exception as e:
|
||
logger.error("Error closing req socket: %s", e)
|
||
self._create_req_socket()
|
||
|
||
def _get_full_sync_sender(self):
|
||
"""Lazy initialization of FullSyncSender"""
|
||
if self._full_sync_sender is None:
|
||
# Get the local_cpu_backend from lmcache_engine
|
||
local_cpu_backend = self.lmcache_engine.storage_manager.local_cpu_backend
|
||
self._full_sync_sender = FullSyncSender(
|
||
config=self.config,
|
||
worker=self,
|
||
lmcache_engine=self.lmcache_engine,
|
||
local_cpu_backend=local_cpu_backend,
|
||
)
|
||
return self._full_sync_sender
|
||
|
||
def _on_request_failure(self, msg: WorkerReqMsg) -> WorkerReqRetMsg:
|
||
"""
|
||
Create a default return message when worker -> controller
|
||
request encounters an error (e.g., timeout, ZMQ error).
|
||
"""
|
||
if isinstance(msg, BatchedP2PLookupMsg):
|
||
return BatchedP2PLookupRetMsg(layout_info=[("", "", 0, "")])
|
||
elif isinstance(msg, HeartbeatMsg):
|
||
return HeartbeatRetMsg() # No command by default
|
||
elif isinstance(msg, FullSyncStartMsg):
|
||
return FullSyncStartRetMsg(
|
||
sync_id=msg.sync_id,
|
||
accepted=False,
|
||
error_msg="Communication error",
|
||
)
|
||
elif isinstance(msg, FullSyncStatusMsg):
|
||
return FullSyncStatusRetMsg(
|
||
sync_id=msg.sync_id,
|
||
is_complete=False,
|
||
global_progress=0.0,
|
||
can_exit_freeze=False,
|
||
)
|
||
else:
|
||
raise ValueError(f"Unknown message type: {type(msg)}")
|
||
|
||
def put_msg(self, msg: WorkerMsg):
|
||
"""
|
||
Put a message into the message queue.
|
||
"""
|
||
# TODO(Jiayi): This might introduce ~0.05ms latency than
|
||
# a normal function call.
|
||
# Not sure how much overhead is blocking though.
|
||
self.loop.call_soon_threadsafe(self.msg_queue.put_nowait, msg)
|
||
|
||
async def batched_get_msg(self, max_bsz: int = 50) -> list[WorkerMsg]:
|
||
"""
|
||
Get a batch of messages from the message queue.
|
||
"""
|
||
batch = []
|
||
|
||
# use blocking get for the first msg
|
||
try:
|
||
item = await self.msg_queue.get()
|
||
batch.append(item)
|
||
except asyncio.CancelledError:
|
||
return batch # shutdown path
|
||
|
||
for _ in range(max_bsz - 1):
|
||
try:
|
||
item = self.msg_queue.get_nowait()
|
||
batch.append(item)
|
||
except asyncio.QueueEmpty:
|
||
break
|
||
return batch
|
||
|
||
async def _send_heartbeat_msg(self, msg: HeartbeatMsg) -> HeartbeatRetMsg:
|
||
"""
|
||
Send heartbeat message via dedicated heartbeat DEALER socket.
|
||
This is separate from async_put_and_wait_msg to keep heartbeat independent.
|
||
"""
|
||
if self.heartbeat_socket is None:
|
||
logger.warning("Heartbeat socket is not initialized")
|
||
return HeartbeatRetMsg()
|
||
try:
|
||
# DEALER socket: send [empty_frame, payload]
|
||
logger.info("Sending heartbeat message to controller...")
|
||
await self.heartbeat_socket.send_multipart(
|
||
[b"", msgspec.msgpack.encode(msg)]
|
||
)
|
||
logger.info("Heartbeat message sent, waiting for response...")
|
||
frames = await self.heartbeat_socket.recv_multipart()
|
||
logger.info("Received heartbeat response with %d frames", len(frames))
|
||
# DEALER receives: [empty_frame, payload]
|
||
serialized_ret_msg = frames[-1]
|
||
ret_msg = msgspec.msgpack.decode(serialized_ret_msg, type=Msg)
|
||
return ret_msg
|
||
except zmq.Again as e:
|
||
logger.error("Heartbeat timeout occurred, recreating socket. Error: %s", e)
|
||
self._recreate_heartbeat_socket()
|
||
return HeartbeatRetMsg()
|
||
except zmq.ZMQError as e:
|
||
logger.error(
|
||
"Heartbeat ZMQ error occurred, recreating socket. Error: %s", e
|
||
)
|
||
self._recreate_heartbeat_socket()
|
||
return HeartbeatRetMsg()
|
||
except Exception as e:
|
||
logger.error("Error happens in heartbeat socket. Error: %s", e)
|
||
return HeartbeatRetMsg()
|
||
|
||
async def heartbeat(self):
|
||
"""
|
||
Send periodic heartbeats to the controller (DEALER-ROUTER mode).
|
||
|
||
Process any commands received in the heartbeat response.
|
||
Uses dedicated heartbeat socket to avoid blocking from other requests.
|
||
"""
|
||
enable_heartbeat = (
|
||
self.config.lmcache_worker_heartbeat_time is not None
|
||
and self.config.lmcache_worker_heartbeat_time > 0
|
||
and self.heartbeat_socket is not None
|
||
)
|
||
if enable_heartbeat:
|
||
await asyncio.sleep(self.config.lmcache_worker_heartbeat_delay_time)
|
||
logger.info(
|
||
"Start heartbeat in %s : %s, delay time: %ss, heartbeat time: %ss",
|
||
self.lmcache_instance_id,
|
||
self.worker_id,
|
||
self.config.lmcache_worker_heartbeat_delay_time,
|
||
self.config.lmcache_worker_heartbeat_time,
|
||
)
|
||
while True:
|
||
# Send heartbeat via dedicated heartbeat socket
|
||
heartbeat_msg = HeartbeatMsg(
|
||
instance_id=self.lmcache_instance_id,
|
||
worker_id=self.worker_id,
|
||
ip=self.lmcache_worker_ip,
|
||
port=self.lmcache_worker_port,
|
||
peer_init_url=self.p2p_init_url,
|
||
)
|
||
|
||
try:
|
||
ret_msg = await self._send_heartbeat_msg(heartbeat_msg)
|
||
|
||
if isinstance(ret_msg, HeartbeatRetMsg):
|
||
self._handle_heartbeat_commands(ret_msg)
|
||
else:
|
||
logger.warning(
|
||
"Unexpected heartbeat response type: %s", type(ret_msg)
|
||
)
|
||
except Exception as e:
|
||
logger.error("Error during heartbeat: %s", e)
|
||
|
||
await asyncio.sleep(self.config.lmcache_worker_heartbeat_time)
|
||
|
||
def _handle_heartbeat_commands(self, ret_msg: HeartbeatRetMsg):
|
||
"""
|
||
Handle commands received in heartbeat response.
|
||
|
||
Uses polymorphic dispatch - each command class implements its own
|
||
execute() method. Commands are executed sequentially.
|
||
"""
|
||
if not ret_msg.has_commands():
|
||
return
|
||
|
||
for command in ret_msg.commands:
|
||
logger.info(
|
||
"Executing heartbeat command: %s",
|
||
command.describe(),
|
||
)
|
||
try:
|
||
command.execute(self)
|
||
except NotImplementedError:
|
||
logger.warning(
|
||
"Command %s.execute() not implemented yet",
|
||
type(command).__name__,
|
||
)
|
||
except Exception as e:
|
||
logger.error(
|
||
"Error executing command %s: %s",
|
||
type(command).__name__,
|
||
e,
|
||
)
|
||
|
||
async def push(self):
|
||
while True:
|
||
try:
|
||
msgs = await self.batched_get_msg()
|
||
logger.debug(f"Sending {len(msgs)} messages")
|
||
self.push_socket.send_multipart(
|
||
[msgspec.msgpack.encode(msg) for msg in msgs]
|
||
)
|
||
|
||
except Exception as e:
|
||
logger.error(f"Push error: {e}")
|
||
|
||
async def handle_request(self):
|
||
"""
|
||
Handle incoming requests (control msgs) from the controller.
|
||
"""
|
||
while True:
|
||
try:
|
||
serialized_request = await self.reply_socket.recv()
|
||
request = msgspec.msgpack.decode(serialized_request, type=Msg)
|
||
logger.debug(f"Received message: {request}")
|
||
if isinstance(request, MoveWorkerMsg):
|
||
tokens = request.tokens
|
||
old_position = request.old_position
|
||
new_position = request.new_position
|
||
do_copy = request.copy
|
||
worker_event_id = request.worker_event_id
|
||
|
||
# Intra node move
|
||
if new_position[0] == self.lmcache_worker_internal_url:
|
||
# TODO(Jiayi): currently we only support moving from
|
||
# local disk to local cpu.
|
||
assert old_position[1] == "LocalDiskBackend"
|
||
assert new_position[1] == "LocalCPUBackend"
|
||
assert do_copy
|
||
|
||
# TODO(Jiayi): We need to align prefetch and move.
|
||
logger.debug("Executing prefetch operation.")
|
||
raise NotImplementedError(
|
||
"Prefetch from controller is not implemented yet."
|
||
)
|
||
else:
|
||
assert new_position[1] == "LocalCPUBackend", (
|
||
"Only support moving to cpu for now."
|
||
)
|
||
logger.debug("Executing cross-node move operation.")
|
||
num_tokens = self.lmcache_engine.move(
|
||
tokens=tokens,
|
||
old_position=old_position,
|
||
new_position=new_position,
|
||
event_id=worker_event_id,
|
||
do_copy=do_copy,
|
||
)
|
||
|
||
# TODO(Jiayi): LMCache needs to have an event tracking
|
||
# pool to enable more advanced control-plane optims.
|
||
# For now, we use a dummy `event_id`.
|
||
serialized_ret_msg = msgspec.msgpack.encode(
|
||
MoveWorkerRetMsg(num_tokens=num_tokens)
|
||
)
|
||
elif isinstance(request, CompressWorkerMsg):
|
||
num_compressed_tokens = self.lmcache_engine.compress(
|
||
tokens=request.tokens,
|
||
method=request.method,
|
||
location=request.location,
|
||
event_id=request.worker_event_id,
|
||
)
|
||
serialized_ret_msg = msgspec.msgpack.encode(
|
||
CompressWorkerRetMsg(num_tokens=num_compressed_tokens)
|
||
)
|
||
elif isinstance(request, DecompressWorkerMsg):
|
||
num_decompressed_tokens = self.lmcache_engine.decompress(
|
||
tokens=request.tokens,
|
||
method=request.method,
|
||
location=request.location,
|
||
event_id=request.worker_event_id,
|
||
)
|
||
serialized_ret_msg = msgspec.msgpack.encode(
|
||
DecompressWorkerRetMsg(num_tokens=num_decompressed_tokens)
|
||
)
|
||
elif isinstance(request, PinWorkerMsg):
|
||
num_pinned_tokens = self.lmcache_engine.lookup(
|
||
tokens=request.tokens,
|
||
search_range=[request.location],
|
||
lookup_id=request.worker_event_id,
|
||
pin=True,
|
||
)
|
||
serialized_ret_msg = msgspec.msgpack.encode(
|
||
PinWorkerRetMsg(num_tokens=num_pinned_tokens)
|
||
)
|
||
elif isinstance(request, ClearWorkerMsg):
|
||
num_cleared_tokens = self.lmcache_engine.clear(
|
||
locations=[request.location],
|
||
)
|
||
serialized_ret_msg = msgspec.msgpack.encode(
|
||
ClearWorkerRetMsg(num_tokens=num_cleared_tokens)
|
||
)
|
||
elif isinstance(request, HealthWorkerMsg):
|
||
error_code = self.lmcache_engine.health()
|
||
serialized_ret_msg = msgspec.msgpack.encode(
|
||
HealthWorkerRetMsg(error_code=error_code)
|
||
)
|
||
else:
|
||
logger.error(f"Unknown message: {request}")
|
||
serialized_ret_msg = msgspec.msgpack.encode(
|
||
ErrorMsg(error=f"Unknown message: {request}")
|
||
)
|
||
|
||
await self.reply_socket.send(serialized_ret_msg)
|
||
except Exception as e:
|
||
logger.error(f"Worker error: {e}")
|
||
serialized_ret_msg = msgspec.msgpack.encode(
|
||
ErrorMsg(error=f"Worker error: {e}")
|
||
)
|
||
await self.reply_socket.send(serialized_ret_msg)
|
||
|
||
async def start_all(self):
|
||
try:
|
||
# Register first to get heartbeat_url before starting heartbeat task
|
||
await self.register()
|
||
|
||
logger.info(
|
||
f"Starting lmcache worker {self.worker_id}"
|
||
f"for instance {self.lmcache_instance_id}"
|
||
)
|
||
await asyncio.gather(
|
||
self.push(),
|
||
self.handle_request(),
|
||
self.heartbeat(),
|
||
)
|
||
except Exception as e:
|
||
logger.error(
|
||
f"Instance {self.lmcache_instance_id}, "
|
||
f"worker {self.worker_id} error: {e}"
|
||
)
|
||
|
||
def close(self):
|
||
self.deregister()
|
||
if self.loop.is_running():
|
||
self.loop.call_soon_threadsafe(self.loop.stop)
|
||
if self.thread.is_alive():
|
||
self.thread.join()
|
||
self.loop.close()
|
||
close_zmq_socket(self.push_socket)
|
||
close_zmq_socket(self.reply_socket)
|
||
if self.heartbeat_socket is not None:
|
||
close_zmq_socket(self.heartbeat_socket)
|
||
if hasattr(self, "req_socket"):
|
||
close_zmq_socket(self.req_socket)
|