You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
199 lines
8.6 KiB
Python
199 lines
8.6 KiB
Python
import multiprocessing as mp
|
|
import multiprocessing.pool
|
|
import threading
|
|
from collections import defaultdict
|
|
from itertools import chain
|
|
from queue import SimpleQueue
|
|
from selectors import EVENT_READ, DefaultSelector
|
|
from statistics import mean
|
|
from time import time
|
|
from typing import Dict, NamedTuple, Optional
|
|
|
|
import torch
|
|
from hivemind.moe.server.module_backend import ModuleBackend
|
|
from hivemind.utils import get_logger
|
|
from prefetch_generator import BackgroundGenerator
|
|
|
|
logger = get_logger(__name__)
|
|
|
|
|
|
class Runtime(threading.Thread):
|
|
"""
|
|
A group of processes that processes incoming requests for multiple module backends on a shared device.
|
|
Runtime is usually created and managed by Server, humans need not apply.
|
|
|
|
For debugging, you can start runtime manually with .start() or .run()
|
|
|
|
>>> module_backends = {'block_uid': ModuleBackend(**kwargs)}
|
|
>>> runtime = Runtime(module_backends)
|
|
>>> runtime.start() # start runtime in background thread. To start in current thread, use runtime.run()
|
|
>>> runtime.ready.wait() # await for runtime to load all blocks on device and create request pools
|
|
>>> future = runtime.module_backends['block_uid'].forward_pool.submit_task(*module_inputs)
|
|
>>> print("Returned:", future.result())
|
|
>>> runtime.shutdown()
|
|
|
|
:param module_backends: a dict [block uid -> ModuleBackend]
|
|
:param prefetch_batches: form up to this many batches in advance
|
|
:param sender_threads: dispatches outputs from finished batches using this many asynchronous threads
|
|
:param device: if specified, moves all blocks and data to this device via .to(device=device).
|
|
If you want to manually specify devices for each block (in their forward pass), leave device=None (default)
|
|
|
|
:param stats_report_interval: interval to collect and log statistics about runtime performance
|
|
"""
|
|
|
|
SHUTDOWN_TRIGGER = "RUNTIME SHUTDOWN TRIGGERED"
|
|
|
|
def __init__(
|
|
self,
|
|
module_backends: Dict[str, ModuleBackend],
|
|
prefetch_batches: int = 1,
|
|
sender_threads: int = 1,
|
|
device: torch.device = None,
|
|
stats_report_interval: Optional[int] = None,
|
|
):
|
|
super().__init__()
|
|
self.module_backends = module_backends
|
|
self.pools = tuple(chain(*(backend.get_pools() for backend in module_backends.values())))
|
|
self.device, self.prefetch_batches, self.sender_threads = device, prefetch_batches, sender_threads
|
|
self.shutdown_recv, self.shutdown_send = mp.Pipe(duplex=False)
|
|
self.shutdown_trigger = mp.Event()
|
|
self.ready = mp.Event() # event is set iff server is currently running and ready to accept batches
|
|
|
|
self.stats_report_interval = stats_report_interval
|
|
if self.stats_report_interval is not None:
|
|
self.stats_reporter = StatsReporter(self.stats_report_interval)
|
|
|
|
def run(self):
|
|
for pool in self.pools:
|
|
if not pool.is_alive():
|
|
pool.start()
|
|
if self.device is not None:
|
|
for backend in self.module_backends.values():
|
|
backend.module.to(self.device)
|
|
|
|
with mp.pool.ThreadPool(self.sender_threads) as output_sender_pool:
|
|
try:
|
|
self.ready.set()
|
|
if self.stats_report_interval is not None:
|
|
self.stats_reporter.start()
|
|
logger.info("Started")
|
|
|
|
batch_iterator = self.iterate_minibatches_from_pools()
|
|
if self.prefetch_batches > 0:
|
|
batch_iterator = BackgroundGenerator(batch_iterator, self.prefetch_batches)
|
|
|
|
for pool, batch_index, batch in batch_iterator:
|
|
logger.debug(f"Processing batch {batch_index} from pool {pool.name}")
|
|
|
|
start = time()
|
|
try:
|
|
outputs = pool.process_func(*batch)
|
|
output_sender_pool.apply_async(pool.send_outputs_from_runtime, args=[batch_index, outputs])
|
|
|
|
batch_processing_time = time() - start
|
|
|
|
batch_size = outputs[0].size(0)
|
|
logger.debug(f"Pool {pool.name}: batch {batch_index} processed, size {batch_size}")
|
|
|
|
if self.stats_report_interval is not None:
|
|
self.stats_reporter.report_stats(pool.name, batch_size, batch_processing_time)
|
|
|
|
except KeyboardInterrupt:
|
|
raise
|
|
except BaseException as exception:
|
|
logger.exception(f"Caught {exception}, attempting to recover")
|
|
output_sender_pool.apply_async(pool.send_exception_from_runtime, args=[batch_index, exception])
|
|
|
|
finally:
|
|
if not self.shutdown_trigger.is_set():
|
|
self.shutdown()
|
|
|
|
def shutdown(self):
|
|
"""Gracefully terminate a running runtime."""
|
|
logger.info("Shutting down")
|
|
self.ready.clear()
|
|
|
|
if self.stats_report_interval is not None:
|
|
self.stats_reporter.stop.set()
|
|
self.stats_reporter.join()
|
|
|
|
logger.debug("Terminating pools")
|
|
for pool in self.pools:
|
|
if pool.is_alive():
|
|
pool.shutdown()
|
|
logger.debug("Pools terminated")
|
|
|
|
# trigger background thread to shutdown
|
|
self.shutdown_send.send(self.SHUTDOWN_TRIGGER)
|
|
self.shutdown_trigger.set()
|
|
|
|
def iterate_minibatches_from_pools(self, timeout=None):
|
|
"""
|
|
Chooses pool according to priority, then copies exposed batch and frees the buffer
|
|
"""
|
|
with DefaultSelector() as selector:
|
|
for pool in self.pools:
|
|
selector.register(pool.batch_receiver, EVENT_READ, pool)
|
|
selector.register(self.shutdown_recv, EVENT_READ, self.SHUTDOWN_TRIGGER)
|
|
|
|
while True:
|
|
# wait until at least one batch_receiver becomes available
|
|
logger.debug("Waiting for inputs from task pools")
|
|
ready_fds = selector.select()
|
|
ready_objects = {key.data for (key, events) in ready_fds}
|
|
if self.SHUTDOWN_TRIGGER in ready_objects:
|
|
break # someone asked us to shutdown, break from the loop
|
|
|
|
logger.debug("Choosing the pool with first priority")
|
|
|
|
pool = min(ready_objects, key=lambda pool: pool.priority)
|
|
|
|
logger.debug(f"Loading batch from {pool.name}")
|
|
batch_index, batch_tensors = pool.load_batch_to_runtime(timeout, self.device)
|
|
logger.debug(f"Loaded batch from {pool.name}")
|
|
yield pool, batch_index, batch_tensors
|
|
|
|
|
|
BatchStats = NamedTuple("BatchStats", (("batch_size", int), ("processing_time", float)))
|
|
|
|
|
|
class StatsReporter(threading.Thread):
|
|
def __init__(self, report_interval: int):
|
|
super().__init__()
|
|
self.report_interval = report_interval
|
|
self.stop = threading.Event()
|
|
self.stats_queue = SimpleQueue()
|
|
|
|
def run(self):
|
|
while not self.stop.wait(self.report_interval):
|
|
pool_batch_stats = defaultdict(list)
|
|
while not self.stats_queue.empty():
|
|
pool_uid, batch_stats = self.stats_queue.get()
|
|
pool_batch_stats[pool_uid].append(batch_stats)
|
|
|
|
total_processed_batches = sum(len(pool_stats) for pool_stats in pool_batch_stats.values())
|
|
logger.info(f"Processed {total_processed_batches} batches in last {self.report_interval} seconds:")
|
|
for pool_uid, pool_stats in pool_batch_stats.items():
|
|
total_batches = len(pool_stats)
|
|
total_examples = sum(batch_stats.batch_size for batch_stats in pool_stats)
|
|
avg_batch_size = mean(batch_stats.batch_size for batch_stats in pool_stats)
|
|
total_time = sum(batch_stats.processing_time for batch_stats in pool_stats)
|
|
batches_to_time = total_batches / total_time
|
|
batch_performance = f"{batches_to_time:.2f} " + ("batches/s" if batches_to_time > 1 else "s/batch")
|
|
|
|
examples_to_time = total_examples / total_time
|
|
example_performance = f"{examples_to_time:.2f} " + (
|
|
"examples/s" if examples_to_time > 1 else "s/example"
|
|
)
|
|
|
|
logger.info(
|
|
f"{pool_uid}: "
|
|
f"{total_batches} batches ({batch_performance}), "
|
|
f"{total_examples} examples ({example_performance}), "
|
|
f"avg batch size {avg_batch_size:.2f}"
|
|
)
|
|
|
|
def report_stats(self, pool_uid, batch_size, processing_time):
|
|
batch_stats = BatchStats(batch_size, processing_time)
|
|
self.stats_queue.put_nowait((pool_uid, batch_stats))
|