wehub-resource-sync
301 lines
12 KiB
Python
301 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import math
|
|
from collections.abc import Awaitable, Callable
|
|
from multiprocessing.context import BaseContext
|
|
from typing import Any, Literal
|
|
|
|
from .. import utils
|
|
from ..job import JobContext, JobExecutorType, JobProcess, RunningJobInfo
|
|
from ..log import logger
|
|
from ..utils import aio
|
|
from ..utils.hw.cpu import get_cpu_monitor
|
|
from . import inference_executor, job_proc_executor, job_thread_executor
|
|
from .job_executor import JobExecutor
|
|
|
|
EventTypes = Literal[
|
|
"process_created",
|
|
"process_started",
|
|
"process_ready",
|
|
"process_closed",
|
|
"process_job_launched",
|
|
]
|
|
|
|
MAX_CONCURRENT_INITIALIZATIONS = min(math.ceil(get_cpu_monitor().cpu_count()), 4)
|
|
|
|
|
|
class ProcPool(utils.EventEmitter[EventTypes]):
|
|
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,
|
|
num_idle_processes: int,
|
|
initialize_timeout: float,
|
|
close_timeout: float,
|
|
session_end_timeout: float,
|
|
inference_executor: inference_executor.InferenceExecutor | None,
|
|
job_executor_type: JobExecutorType,
|
|
mp_ctx: BaseContext,
|
|
memory_warn_mb: float,
|
|
memory_limit_mb: float,
|
|
http_proxy: str | None,
|
|
loop: asyncio.AbstractEventLoop,
|
|
) -> None:
|
|
super().__init__()
|
|
self._job_executor_type = job_executor_type
|
|
self._mp_ctx = mp_ctx
|
|
self._initialize_process_fnc = initialize_process_fnc
|
|
self._job_entrypoint_fnc = job_entrypoint_fnc
|
|
self._session_end_fnc = session_end_fnc
|
|
self._simulation_end_fnc = simulation_end_fnc
|
|
self._close_timeout = close_timeout
|
|
self._session_end_timeout = session_end_timeout
|
|
self._inf_executor = inference_executor
|
|
self._initialize_timeout = initialize_timeout
|
|
self._loop = loop
|
|
self._memory_limit_mb = memory_limit_mb
|
|
self._memory_warn_mb = memory_warn_mb
|
|
self._default_num_idle_processes = num_idle_processes
|
|
self._http_proxy = http_proxy
|
|
self._target_idle_processes = num_idle_processes
|
|
|
|
self._init_sem = asyncio.Semaphore(MAX_CONCURRENT_INITIALIZATIONS)
|
|
self._warmed_proc_queue = asyncio.Queue[JobExecutor]()
|
|
self._executors: list[JobExecutor] = []
|
|
self._spawn_tasks: set[asyncio.Task[None]] = set()
|
|
self._close_tasks: set[asyncio.Task[None]] = set()
|
|
self._monitor_tasks: set[asyncio.Task[None]] = set()
|
|
self._started = False
|
|
self._closed = False
|
|
|
|
self._idle_ready = asyncio.Event()
|
|
self._jobs_waiting_for_process = 0
|
|
|
|
@property
|
|
def processes(self) -> list[JobExecutor]:
|
|
return self._executors
|
|
|
|
def get_by_job_id(self, job_id: str) -> JobExecutor | None:
|
|
return next(
|
|
(x for x in self._executors if x.running_job and x.running_job.job.id == job_id),
|
|
None,
|
|
)
|
|
|
|
async def start(self) -> None:
|
|
if self._started:
|
|
return
|
|
|
|
self._started = True
|
|
self._main_atask = asyncio.create_task(self._main_task())
|
|
|
|
if self._default_num_idle_processes > 0:
|
|
# wait for the idle processes to be warmed up (by the main task)
|
|
# use a timeout so start() doesn't block forever if initialization fails
|
|
try:
|
|
await asyncio.wait_for(
|
|
self._idle_ready.wait(),
|
|
timeout=self._initialize_timeout + 2,
|
|
)
|
|
except asyncio.TimeoutError:
|
|
logger.warning("timed out waiting for idle processes to initialize")
|
|
|
|
async def aclose(self) -> None:
|
|
if not self._started:
|
|
return
|
|
|
|
self._closed = True
|
|
await aio.cancel_and_wait(self._main_atask)
|
|
|
|
async def _acquire_proc(self, job_id: str) -> JobExecutor:
|
|
MAX_ACQUIRE_ATTEMPTS = 3
|
|
|
|
for attempt in range(MAX_ACQUIRE_ATTEMPTS):
|
|
if (
|
|
self._warmed_proc_queue.empty()
|
|
and len(self._spawn_tasks) < self._jobs_waiting_for_process
|
|
):
|
|
# spawn a new process if there are no idle processes
|
|
task = asyncio.create_task(self._proc_spawn_task())
|
|
self._spawn_tasks.add(task)
|
|
task.add_done_callback(self._spawn_tasks.discard)
|
|
|
|
if self._warmed_proc_queue.empty():
|
|
logger.warning(
|
|
"no warmed process available for job, waiting for one to be created",
|
|
extra={"job_id": job_id},
|
|
)
|
|
|
|
# race the queue against every in-flight spawn task
|
|
while True:
|
|
if not self._warmed_proc_queue.empty():
|
|
return self._warmed_proc_queue.get_nowait()
|
|
|
|
spawns = [t for t in self._spawn_tasks if not t.done()]
|
|
# retry if all in-flight spawns have completed without producing a proc
|
|
if not spawns:
|
|
break
|
|
|
|
get_task = asyncio.ensure_future(self._warmed_proc_queue.get())
|
|
try:
|
|
await asyncio.wait([get_task, *spawns], return_when=asyncio.FIRST_COMPLETED)
|
|
finally:
|
|
if not get_task.done():
|
|
get_task.cancel()
|
|
|
|
if get_task.done() and not get_task.cancelled():
|
|
return get_task.result()
|
|
|
|
logger.warning(
|
|
"all in-flight spawns failed to initialize, retrying",
|
|
extra={"job_id": job_id, "attempt": attempt + 1},
|
|
)
|
|
|
|
logger.error(
|
|
"failed to acquire process for job after %d attempts",
|
|
MAX_ACQUIRE_ATTEMPTS,
|
|
extra={"job_id": job_id},
|
|
)
|
|
raise RuntimeError(f"no process became available after {MAX_ACQUIRE_ATTEMPTS} attempts")
|
|
|
|
async def launch_job(self, info: RunningJobInfo) -> None:
|
|
MAX_LAUNCH_ATTEMPTS = 3
|
|
|
|
for attempt in range(MAX_LAUNCH_ATTEMPTS):
|
|
self._jobs_waiting_for_process += 1
|
|
try:
|
|
proc = await self._acquire_proc(info.job.id)
|
|
finally:
|
|
self._jobs_waiting_for_process -= 1
|
|
|
|
try:
|
|
await proc.launch_job(info)
|
|
self.emit("process_job_launched", proc)
|
|
return
|
|
except Exception:
|
|
close_task = asyncio.create_task(proc.aclose())
|
|
self._close_tasks.add(close_task)
|
|
close_task.add_done_callback(self._close_tasks.discard)
|
|
if attempt == MAX_LAUNCH_ATTEMPTS - 1:
|
|
logger.error(
|
|
"failed to launch job on process after %d attempts",
|
|
MAX_LAUNCH_ATTEMPTS,
|
|
extra={"job_id": info.job.id},
|
|
)
|
|
raise
|
|
logger.warning(
|
|
"failed to launch job on process, retrying with a new process",
|
|
extra={"job_id": info.job.id, "attempt": attempt + 1},
|
|
)
|
|
|
|
def set_target_idle_processes(self, num_idle_processes: int) -> None:
|
|
self._target_idle_processes = num_idle_processes
|
|
|
|
@property
|
|
def target_idle_processes(self) -> int:
|
|
return self._target_idle_processes
|
|
|
|
@utils.log_exceptions(logger=logger)
|
|
async def _proc_spawn_task(self) -> None:
|
|
proc: JobExecutor
|
|
if self._job_executor_type == JobExecutorType.THREAD:
|
|
proc = job_thread_executor.ThreadJobExecutor(
|
|
initialize_process_fnc=self._initialize_process_fnc,
|
|
job_entrypoint_fnc=self._job_entrypoint_fnc,
|
|
session_end_fnc=self._session_end_fnc,
|
|
simulation_end_fnc=self._simulation_end_fnc,
|
|
initialize_timeout=self._initialize_timeout,
|
|
close_timeout=self._close_timeout,
|
|
session_end_timeout=self._session_end_timeout,
|
|
inference_executor=self._inf_executor,
|
|
ping_interval=2.5,
|
|
high_ping_threshold=0.5,
|
|
http_proxy=self._http_proxy,
|
|
loop=self._loop,
|
|
)
|
|
elif self._job_executor_type == JobExecutorType.PROCESS:
|
|
proc = job_proc_executor.ProcJobExecutor(
|
|
initialize_process_fnc=self._initialize_process_fnc,
|
|
job_entrypoint_fnc=self._job_entrypoint_fnc,
|
|
session_end_fnc=self._session_end_fnc,
|
|
simulation_end_fnc=self._simulation_end_fnc,
|
|
initialize_timeout=self._initialize_timeout,
|
|
close_timeout=self._close_timeout,
|
|
session_end_timeout=self._session_end_timeout,
|
|
inference_executor=self._inf_executor,
|
|
mp_ctx=self._mp_ctx,
|
|
loop=self._loop,
|
|
ping_interval=2.5,
|
|
ping_timeout=60,
|
|
high_ping_threshold=0.5,
|
|
memory_warn_mb=self._memory_warn_mb,
|
|
memory_limit_mb=self._memory_limit_mb,
|
|
http_proxy=self._http_proxy,
|
|
)
|
|
else:
|
|
raise ValueError(f"unsupported job executor: {self._job_executor_type}")
|
|
|
|
self._executors.append(proc)
|
|
initialized = False
|
|
try:
|
|
async with self._init_sem:
|
|
if not self._closed:
|
|
self.emit("process_created", proc)
|
|
await proc.start()
|
|
self.emit("process_started", proc)
|
|
await proc.initialize()
|
|
self.emit("process_ready", proc)
|
|
self._warmed_proc_queue.put_nowait(proc)
|
|
if self._warmed_proc_queue.qsize() >= self._default_num_idle_processes:
|
|
self._idle_ready.set()
|
|
|
|
initialized = True
|
|
except Exception:
|
|
logger.exception("error initializing process", extra=proc.logging_extra())
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
if not initialized:
|
|
self._executors.remove(proc)
|
|
await proc.aclose()
|
|
self.emit("process_closed", proc)
|
|
return
|
|
|
|
monitor_task = asyncio.create_task(self._monitor_process_task(proc))
|
|
self._monitor_tasks.add(monitor_task)
|
|
monitor_task.add_done_callback(self._monitor_tasks.discard)
|
|
|
|
@utils.log_exceptions(logger=logger)
|
|
async def _monitor_process_task(self, proc: JobExecutor) -> None:
|
|
try:
|
|
await proc.join()
|
|
self.emit("process_closed", proc)
|
|
finally:
|
|
self._executors.remove(proc)
|
|
|
|
@utils.log_exceptions(logger=logger)
|
|
async def _main_task(self) -> None:
|
|
try:
|
|
while not self._closed:
|
|
current_pending = self._warmed_proc_queue.qsize() + len(self._spawn_tasks)
|
|
target = max(
|
|
min(self._target_idle_processes, self._default_num_idle_processes),
|
|
self._jobs_waiting_for_process,
|
|
)
|
|
to_spawn = target - current_pending
|
|
|
|
for _ in range(to_spawn):
|
|
task = asyncio.create_task(self._proc_spawn_task())
|
|
self._spawn_tasks.add(task)
|
|
task.add_done_callback(self._spawn_tasks.discard)
|
|
|
|
await asyncio.sleep(0.1)
|
|
except asyncio.CancelledError:
|
|
await aio.cancel_and_wait(*self._spawn_tasks)
|
|
await asyncio.gather(*[proc.aclose() for proc in self._executors])
|
|
await asyncio.gather(*self._close_tasks)
|
|
await asyncio.gather(*self._monitor_tasks)
|