wehub-resource-sync
355 lines
13 KiB
Python
355 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextlib
|
|
import socket
|
|
import threading
|
|
import time
|
|
from collections.abc import Awaitable, Callable
|
|
from dataclasses import dataclass
|
|
from typing import Any
|
|
|
|
from .. import utils
|
|
from ..job import JobContext, JobProcess, RunningJobInfo
|
|
from ..log import logger
|
|
from ..utils.aio import duplex_unix
|
|
from . import channel, job_proc_lazy_main, proto
|
|
from .inference_executor import InferenceExecutor
|
|
from .job_executor import JobStatus
|
|
|
|
|
|
@dataclass
|
|
class _ProcOpts:
|
|
initialize_process_fnc: Callable[[JobProcess], Any]
|
|
job_entrypoint_fnc: Callable[[JobContext], Awaitable[None]]
|
|
session_end_fnc: Callable[[JobContext], Awaitable[None]] | None
|
|
simulation_end_fnc: Callable[[Any], Any] | None
|
|
initialize_timeout: float
|
|
close_timeout: float
|
|
session_end_timeout: float
|
|
ping_interval: float
|
|
high_ping_threshold: float
|
|
http_proxy: str | None
|
|
|
|
|
|
class ThreadJobExecutor:
|
|
def __init__(
|
|
self,
|
|
*,
|
|
initialize_process_fnc: Callable[[JobProcess], Any],
|
|
job_entrypoint_fnc: Callable[[JobContext], Awaitable[None]],
|
|
session_end_fnc: Callable[[JobContext], Awaitable[None]] | None,
|
|
simulation_end_fnc: Callable[[Any], Any] | None,
|
|
inference_executor: InferenceExecutor | None,
|
|
initialize_timeout: float,
|
|
close_timeout: float,
|
|
session_end_timeout: float,
|
|
ping_interval: float,
|
|
high_ping_threshold: float,
|
|
http_proxy: str | None,
|
|
loop: asyncio.AbstractEventLoop,
|
|
) -> None:
|
|
self._loop = loop
|
|
self._opts = _ProcOpts(
|
|
initialize_process_fnc=initialize_process_fnc,
|
|
job_entrypoint_fnc=job_entrypoint_fnc,
|
|
session_end_fnc=session_end_fnc,
|
|
simulation_end_fnc=simulation_end_fnc,
|
|
initialize_timeout=initialize_timeout,
|
|
close_timeout=close_timeout,
|
|
session_end_timeout=session_end_timeout,
|
|
ping_interval=ping_interval,
|
|
high_ping_threshold=high_ping_threshold,
|
|
http_proxy=http_proxy,
|
|
)
|
|
|
|
self._user_args: Any | None = None
|
|
self._job_status: JobStatus | None = None
|
|
self._running_job: RunningJobInfo | None = None
|
|
|
|
self._main_atask: asyncio.Task[None] | None = None
|
|
self._initialize_fut = asyncio.Future[None]()
|
|
self._closing = False
|
|
self._lock = asyncio.Lock()
|
|
self._shutdown_ack_fut = asyncio.Future[None]()
|
|
self._shutting_down_fut = asyncio.Future[None]()
|
|
|
|
self._thread: threading.Thread | None = None
|
|
self._inference_executor = inference_executor
|
|
self._inference_tasks: set[asyncio.Task[None]] = set()
|
|
self._id = utils.shortuuid("THEXEC_")
|
|
|
|
@property
|
|
def id(self) -> str:
|
|
return self._id
|
|
|
|
@property
|
|
def status(self) -> JobStatus:
|
|
if self._job_status is None:
|
|
raise RuntimeError("job status not available")
|
|
|
|
return self._job_status
|
|
|
|
@property
|
|
def started(self) -> bool:
|
|
return self._main_atask is not None
|
|
|
|
@property
|
|
def user_arguments(self) -> Any | None:
|
|
return self._user_args
|
|
|
|
@user_arguments.setter
|
|
def user_arguments(self, value: Any | None) -> None:
|
|
self._user_args = value
|
|
|
|
@property
|
|
def running_job(self) -> RunningJobInfo | None:
|
|
return self._running_job
|
|
|
|
async def start(self) -> None:
|
|
if self.started:
|
|
raise RuntimeError("runner already started")
|
|
|
|
if self._closing:
|
|
raise RuntimeError("runner is closed")
|
|
|
|
await asyncio.shield(self._start())
|
|
|
|
async def _start(self) -> None:
|
|
async with self._lock:
|
|
# to simplify the runners implementation, we also use a duplex in the threaded executor
|
|
# (ThreadedRunners), so we can use the same protocol
|
|
mp_pch, mp_cch = socket.socketpair()
|
|
pch: duplex_unix._AsyncDuplex | None = None
|
|
try:
|
|
pch = await duplex_unix._AsyncDuplex.open(mp_pch)
|
|
self._pch = pch
|
|
|
|
self._join_fut = asyncio.Future[None]()
|
|
|
|
def _on_join() -> None:
|
|
with contextlib.suppress(RuntimeError):
|
|
self._loop.call_soon_threadsafe(self._join_fut.set_result, None)
|
|
|
|
targs = job_proc_lazy_main.ThreadStartArgs(
|
|
mp_cch=mp_cch,
|
|
initialize_process_fnc=self._opts.initialize_process_fnc,
|
|
job_entrypoint_fnc=self._opts.job_entrypoint_fnc,
|
|
session_end_fnc=self._opts.session_end_fnc,
|
|
simulation_end_fnc=self._opts.simulation_end_fnc,
|
|
session_end_timeout=self._opts.session_end_timeout,
|
|
user_arguments=self._user_args,
|
|
join_fnc=_on_join,
|
|
)
|
|
|
|
self._thread = t = threading.Thread(
|
|
target=job_proc_lazy_main.thread_main,
|
|
args=(targs,),
|
|
name="job_thread_runner",
|
|
)
|
|
t.start()
|
|
|
|
self._main_atask = asyncio.create_task(self._main_task())
|
|
except Exception:
|
|
with contextlib.suppress(OSError):
|
|
mp_cch.close()
|
|
|
|
if pch is not None:
|
|
with contextlib.suppress(duplex_unix.DuplexClosed):
|
|
await pch.aclose()
|
|
else:
|
|
with contextlib.suppress(OSError):
|
|
mp_pch.close()
|
|
|
|
raise
|
|
|
|
async def join(self) -> None:
|
|
"""wait for the thread to finish"""
|
|
if not self.started:
|
|
raise RuntimeError("runner not started")
|
|
|
|
async with self._lock:
|
|
if self._main_atask:
|
|
await asyncio.shield(self._main_atask)
|
|
|
|
async def initialize(self) -> None:
|
|
await channel.asend_message(
|
|
self._pch, proto.InitializeRequest(http_proxy=self._opts.http_proxy or "")
|
|
)
|
|
|
|
try:
|
|
logger.info("initializing job runner", extra=self.logging_extra())
|
|
start_time = time.perf_counter()
|
|
init_res = await asyncio.wait_for(
|
|
channel.arecv_message(self._pch, proto.IPC_MESSAGES),
|
|
timeout=self._opts.initialize_timeout,
|
|
)
|
|
assert isinstance(init_res, proto.InitializeResponse), (
|
|
"first message must be InitializeResponse"
|
|
)
|
|
logger.info(
|
|
"job runner initialized",
|
|
extra={
|
|
**self.logging_extra(),
|
|
"elapsed_time": round(time.perf_counter() - start_time, 2),
|
|
},
|
|
)
|
|
except asyncio.TimeoutError:
|
|
self._initialize_fut.set_exception(
|
|
asyncio.TimeoutError("runner initialization timed out")
|
|
)
|
|
raise
|
|
except Exception as e: # should be channel.ChannelClosed most of the time
|
|
self._initialize_fut.set_exception(e)
|
|
raise
|
|
else:
|
|
self._initialize_fut.set_result(None)
|
|
|
|
async def aclose(self) -> None:
|
|
"""
|
|
attempt to gracefully close the job. warn if it takes too long to close
|
|
(in the threaded executor, the job can't be "killed")
|
|
"""
|
|
if not self.started:
|
|
return
|
|
|
|
self._closing = True
|
|
with contextlib.suppress(utils.aio.duplex_unix.DuplexClosed):
|
|
await channel.asend_message(self._pch, proto.ShutdownRequest())
|
|
|
|
try:
|
|
await asyncio.wait_for(self._shutdown_ack_fut, timeout=self._opts.close_timeout)
|
|
except asyncio.TimeoutError:
|
|
logger.error("job did not ack shutdown in time", extra=self.logging_extra())
|
|
|
|
if not self._shutting_down_fut.done():
|
|
await self._shutting_down_fut
|
|
|
|
try:
|
|
if self._main_atask:
|
|
await asyncio.wait_for(
|
|
asyncio.shield(self._main_atask), timeout=self._opts.close_timeout
|
|
)
|
|
except asyncio.TimeoutError:
|
|
logger.error("job shutdown is taking too much time..", extra=self.logging_extra())
|
|
|
|
async with self._lock:
|
|
if self._main_atask:
|
|
await asyncio.shield(self._main_atask)
|
|
|
|
async def _do_inference_task(self, inf_req: proto.InferenceRequest) -> None:
|
|
if self._inference_executor is None:
|
|
logger.warning("inference request received but no inference executor")
|
|
await channel.asend_message(
|
|
self._pch,
|
|
proto.InferenceResponse(
|
|
request_id=inf_req.request_id, error="no inference executor"
|
|
),
|
|
)
|
|
return
|
|
|
|
try:
|
|
inf_res = await self._inference_executor.do_inference(inf_req.method, inf_req.data)
|
|
await channel.asend_message(
|
|
self._pch,
|
|
proto.InferenceResponse(request_id=inf_req.request_id, data=inf_res),
|
|
)
|
|
except Exception as e:
|
|
await channel.asend_message(
|
|
self._pch,
|
|
proto.InferenceResponse(request_id=inf_req.request_id, error=str(e)),
|
|
)
|
|
|
|
async def launch_job(self, info: RunningJobInfo) -> None:
|
|
"""start/assign a job to the executor"""
|
|
if self._running_job is not None:
|
|
raise RuntimeError("executor already has a running job")
|
|
|
|
if not self._initialize_fut.done():
|
|
raise RuntimeError("executor not initialized")
|
|
|
|
self._running_job = info
|
|
self._job_status = JobStatus.RUNNING
|
|
|
|
start_req = proto.StartJobRequest()
|
|
start_req.running_job = info
|
|
await channel.asend_message(self._pch, start_req)
|
|
|
|
@utils.log_exceptions(logger=logger)
|
|
async def _main_task(self) -> None:
|
|
try:
|
|
await self._initialize_fut
|
|
except asyncio.TimeoutError:
|
|
pass # this happens when the initialization takes longer than self._initialize_timeout
|
|
except Exception:
|
|
pass # initialization failed
|
|
|
|
ping_task = asyncio.create_task(self._ping_task())
|
|
monitor_task = asyncio.create_task(self._monitor_task())
|
|
|
|
await self._join_fut
|
|
|
|
await utils.aio.cancel_and_wait(ping_task, monitor_task)
|
|
await utils.aio.cancel_and_wait(*self._inference_tasks)
|
|
|
|
with contextlib.suppress(duplex_unix.DuplexClosed):
|
|
await self._pch.aclose()
|
|
|
|
self._job_status = JobStatus.SUCCESS
|
|
|
|
@utils.log_exceptions(logger=logger)
|
|
async def _monitor_task(self) -> None:
|
|
while True:
|
|
try:
|
|
msg = await channel.arecv_message(self._pch, proto.IPC_MESSAGES)
|
|
except utils.aio.duplex_unix.DuplexClosed:
|
|
break
|
|
|
|
if isinstance(msg, proto.PongResponse):
|
|
delay = utils.time_ms() - msg.timestamp
|
|
if delay > self._opts.high_ping_threshold * 1000:
|
|
logger.warning(
|
|
"job executor is unresponsive",
|
|
extra={"delay": delay, **self.logging_extra()},
|
|
)
|
|
|
|
if isinstance(msg, proto.ShutdownRequestAck):
|
|
if not self._shutdown_ack_fut.done():
|
|
self._shutdown_ack_fut.set_result(None)
|
|
|
|
if isinstance(msg, proto.ShuttingDown):
|
|
if not self._shutting_down_fut.done():
|
|
self._shutting_down_fut.set_result(None)
|
|
|
|
if isinstance(msg, proto.Exiting):
|
|
logger.debug("job exiting", extra={"reason": msg.reason, **self.logging_extra()})
|
|
|
|
if isinstance(msg, proto.InferenceRequest):
|
|
task = asyncio.create_task(self._do_inference_task(msg))
|
|
self._inference_tasks.add(task)
|
|
task.add_done_callback(self._inference_tasks.discard)
|
|
|
|
# resolve pending futures when the channel closes
|
|
if not self._shutdown_ack_fut.done():
|
|
self._shutdown_ack_fut.set_result(None)
|
|
|
|
@utils.log_exceptions(logger=logger)
|
|
async def _ping_task(self) -> None:
|
|
ping_interval = utils.aio.interval(self._opts.ping_interval)
|
|
while True:
|
|
await ping_interval.tick()
|
|
try:
|
|
await channel.asend_message(self._pch, proto.PingRequest(timestamp=utils.time_ms()))
|
|
except utils.aio.duplex_unix.DuplexClosed:
|
|
break
|
|
|
|
def logging_extra(self) -> dict[str, Any]:
|
|
extra: dict[str, Any] = {
|
|
"tid": self._thread.native_id if self._thread else None,
|
|
}
|
|
if self._running_job:
|
|
extra["job_id"] = self._running_job.job.id
|
|
extra["room"] = self._running_job.job.room.name
|
|
|
|
return extra
|