wehub-resource-sync
110 lines
3.8 KiB
Python
110 lines
3.8 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import struct
|
|
|
|
from livekit.protocol.agent_pb import agent_dev
|
|
|
|
from .. import utils
|
|
from ..log import logger
|
|
from ..worker import AgentServer
|
|
from . import proto
|
|
|
|
|
|
async def _send_proto(writer: asyncio.StreamWriter, msg: bytes) -> None:
|
|
writer.write(struct.pack("!I", len(msg)))
|
|
writer.write(msg)
|
|
await writer.drain()
|
|
|
|
|
|
async def _recv_proto(reader: asyncio.StreamReader) -> bytes:
|
|
len_bytes = await reader.readexactly(4)
|
|
length = struct.unpack("!I", len_bytes)[0]
|
|
return await reader.readexactly(length)
|
|
|
|
|
|
class WatchClient:
|
|
"""Connects to Go CLI's reload server over TCP using DevMessage protobuf."""
|
|
|
|
def __init__(
|
|
self,
|
|
worker: AgentServer,
|
|
cli_addr: str,
|
|
loop: asyncio.AbstractEventLoop | None = None,
|
|
) -> None:
|
|
self._loop = loop or asyncio.get_event_loop()
|
|
self._worker = worker
|
|
self._cli_addr = cli_addr
|
|
self._main_task: asyncio.Task | None = None
|
|
|
|
def start(self) -> None:
|
|
self._main_task = self._loop.create_task(self._run())
|
|
|
|
@utils.log_exceptions(logger=logger)
|
|
async def _run(self) -> None:
|
|
host, port_str = self._cli_addr.rsplit(":", 1)
|
|
reader, writer = await asyncio.open_connection(host, int(port_str))
|
|
|
|
try:
|
|
# On startup: tell the CLI about this agent server (agent name + URL) so
|
|
# it can surface things like a Cloud console link, then sync running jobs.
|
|
info = agent_dev.AgentDevMessage(
|
|
server_info=agent_dev.ServerInfo(
|
|
agent_name=self._worker._agent_name,
|
|
url=self._worker._ws_url,
|
|
)
|
|
)
|
|
await _send_proto(writer, info.SerializeToString())
|
|
|
|
# send GetRunningJobsRequest to Go, recv response, reload jobs
|
|
req = agent_dev.AgentDevMessage(
|
|
get_running_jobs_request=agent_dev.GetRunningAgentJobsRequest()
|
|
)
|
|
await _send_proto(writer, req.SerializeToString())
|
|
|
|
data = await _recv_proto(reader)
|
|
resp = agent_dev.AgentDevMessage()
|
|
resp.ParseFromString(data)
|
|
|
|
if resp.HasField("get_running_jobs_response"):
|
|
jobs_resp = resp.get_running_jobs_response
|
|
if jobs_resp.jobs:
|
|
running_jobs = [proto.running_job_from_proto(j) for j in jobs_resp.jobs]
|
|
logger.info(f"reloading {len(running_jobs)} job(s)")
|
|
await self._worker._reload_jobs(running_jobs)
|
|
|
|
# Listen for GetRunningJobsRequest from Go (capture before restart)
|
|
while True:
|
|
try:
|
|
data = await _recv_proto(reader)
|
|
except (asyncio.IncompleteReadError, ConnectionError, OSError):
|
|
break
|
|
|
|
msg = agent_dev.AgentDevMessage()
|
|
msg.ParseFromString(data)
|
|
|
|
if msg.HasField("get_running_jobs_request"):
|
|
jobs = self._worker.active_jobs
|
|
job_protos = [proto.running_job_to_proto(j) for j in jobs]
|
|
resp = agent_dev.AgentDevMessage(
|
|
get_running_jobs_response=agent_dev.GetRunningAgentJobsResponse(
|
|
jobs=job_protos
|
|
)
|
|
)
|
|
await _send_proto(writer, resp.SerializeToString())
|
|
except (asyncio.IncompleteReadError, ConnectionError, OSError):
|
|
pass
|
|
finally:
|
|
writer.close()
|
|
with contextlib.suppress(Exception):
|
|
await writer.wait_closed()
|
|
|
|
async def aclose(self) -> None:
|
|
if not self._main_task:
|
|
return
|
|
|
|
self._main_task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await self._main_task
|