|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import logging |
|
|
import multiprocessing as mp |
|
|
import os |
|
|
import signal |
|
|
import time |
|
|
from queue import Empty |
|
|
from typing import Any, Dict, List, Set, Tuple |
|
|
|
|
|
from .api import RequestQueue, TimerClient, TimerRequest, TimerServer |
|
|
|
|
|
__all__ = ['LocalTimerClient', 'MultiprocessingRequestQueue', 'LocalTimerServer'] |
|
|
|
|
|
class LocalTimerClient(TimerClient): |
|
|
""" |
|
|
Client side of ``LocalTimerServer``. This client is meant to be used |
|
|
on the same host that the ``LocalTimerServer`` is running on and uses |
|
|
pid to uniquely identify a worker. This is particularly useful in situations |
|
|
where one spawns a subprocess (trainer) per GPU on a host with multiple |
|
|
GPU devices. |
|
|
""" |
|
|
|
|
|
def __init__(self, mp_queue): |
|
|
super().__init__() |
|
|
self._mp_queue = mp_queue |
|
|
|
|
|
def acquire(self, scope_id, expiration_time): |
|
|
pid = os.getpid() |
|
|
acquire_request = TimerRequest(pid, scope_id, expiration_time) |
|
|
self._mp_queue.put(acquire_request) |
|
|
|
|
|
def release(self, scope_id): |
|
|
pid = os.getpid() |
|
|
release_request = TimerRequest(pid, scope_id, -1) |
|
|
self._mp_queue.put(release_request) |
|
|
|
|
|
|
|
|
class MultiprocessingRequestQueue(RequestQueue): |
|
|
""" |
|
|
A ``RequestQueue`` backed by python ``multiprocessing.Queue`` |
|
|
""" |
|
|
|
|
|
def __init__(self, mp_queue: mp.Queue): |
|
|
super().__init__() |
|
|
self._mp_queue = mp_queue |
|
|
|
|
|
def size(self) -> int: |
|
|
return self._mp_queue.qsize() |
|
|
|
|
|
def get(self, size, timeout: float) -> List[TimerRequest]: |
|
|
requests = [] |
|
|
wait = timeout |
|
|
for _ in range(0, size): |
|
|
start = time.time() |
|
|
|
|
|
try: |
|
|
r = self._mp_queue.get(block=True, timeout=wait) |
|
|
except Empty: |
|
|
break |
|
|
|
|
|
requests.append(r) |
|
|
wait = wait - (time.time() - start) |
|
|
if wait <= 0: |
|
|
break |
|
|
|
|
|
return requests |
|
|
|
|
|
|
|
|
class LocalTimerServer(TimerServer): |
|
|
""" |
|
|
Server that works with ``LocalTimerClient``. Clients are expected to be |
|
|
subprocesses to the parent process that is running this server. Each host |
|
|
in the job is expected to start its own timer server locally and each |
|
|
server instance manages timers for local workers (running on processes |
|
|
on the same host). |
|
|
""" |
|
|
|
|
|
def __init__( |
|
|
self, mp_queue: mp.Queue, max_interval: float = 60, daemon: bool = True |
|
|
): |
|
|
super().__init__(MultiprocessingRequestQueue(mp_queue), max_interval, daemon) |
|
|
self._timers: Dict[Tuple[Any, str], TimerRequest] = {} |
|
|
|
|
|
def register_timers(self, timer_requests: List[TimerRequest]) -> None: |
|
|
for request in timer_requests: |
|
|
pid = request.worker_id |
|
|
scope_id = request.scope_id |
|
|
expiration_time = request.expiration_time |
|
|
|
|
|
|
|
|
if expiration_time < 0: |
|
|
self._timers.pop((pid, scope_id), None) |
|
|
else: |
|
|
self._timers[(pid, scope_id)] = request |
|
|
|
|
|
def clear_timers(self, worker_ids: Set[int]) -> None: |
|
|
for (pid, scope_id) in list(self._timers.keys()): |
|
|
if pid in worker_ids: |
|
|
self._timers.pop((pid, scope_id)) |
|
|
|
|
|
def get_expired_timers(self, deadline: float) -> Dict[Any, List[TimerRequest]]: |
|
|
|
|
|
expired_timers: Dict[Any, List[TimerRequest]] = {} |
|
|
for request in self._timers.values(): |
|
|
if request.expiration_time <= deadline: |
|
|
expired_scopes = expired_timers.setdefault(request.worker_id, []) |
|
|
expired_scopes.append(request) |
|
|
return expired_timers |
|
|
|
|
|
def _reap_worker(self, worker_id: int) -> bool: |
|
|
try: |
|
|
os.kill(worker_id, signal.SIGKILL) |
|
|
return True |
|
|
except ProcessLookupError: |
|
|
logging.info(f"Process with pid={worker_id} does not exist. Skipping") |
|
|
return True |
|
|
except Exception as e: |
|
|
logging.error(f"Error terminating pid={worker_id}", exc_info=e) |
|
|
return False |
|
|
|
