# Copyright (c) Meta Platforms, Inc. and affiliates. import json import logging import multiprocessing as mp from multiprocessing.synchronize import Event as EventClass from queue import Empty, Full import numpy as np from pydantic import BaseModel, ConfigDict from bytelatent.data.data_types import Batch from bytelatent.data.iterators.abstract_iterator import IteratorState, StatefulIterator from bytelatent.data.iterators.packing_iterator import PackingIteratorState logger = logging.getLogger() class MultiprocessIteratorState(BaseModel, IteratorState): model_config = ConfigDict(extra="forbid") base_iterator_state: PackingIteratorState n_batches_to_prefetch: int serialized_prefetch_buffer: str def build(self): base_iterator = self.base_iterator_state.build() data = json.loads(self.serialized_prefetch_buffer) prefetch_buffer = [Batch.from_python_dict(item) for item in data] return MultiprocessIterator( base_iterator, n_batches_to_prefetch=self.n_batches_to_prefetch, prefetch_buffer=prefetch_buffer, ) def start_work_from_state( batch_queue: mp.Queue, state_queue: mp.Queue, stop_event: EventClass, state_dumped_event: EventClass, state: IteratorState, ): logging.info("Worker thread: Starting base_iterator work") stateful_iterator = state.build() iterator = stateful_iterator.create_iter() for item in iterator: while not stop_event.is_set(): try: # Attempt to put on queue or timeout to try again (maybe main thread is busy) batch_queue.put(item, timeout=0.1) # On success, stop trying break except Full: pass if stop_event.is_set(): # Signal the end of output, this ensures that even if the queue takes a while to # buffer, that the main thread receives everything (and tosses this fake batch) logging.info( "Worker thread: Stop event detected, outputting is_final=True batch" ) batch_queue.put( Batch( x=np.zeros((1, 1)), y=np.zeros((1, 1)), is_final=True, mask=None, patch_lengths=None, ngram_ids=None, ) ) break try: logging.info("Worker thread: outputting state") state_queue.put(iterator.get_state(), timeout=1) logging.info("Worker thread: state dump complete") state_dumped_event.set() logging.info("Worker thread: set state_dump_event") except Full: raise ValueError( "Attempted to dump state into the state queue, but it was full" ) class MultiprocessIterator(StatefulIterator): """ Design sketch of the multiprocess iterator: Given the base_iterator, the only thing we do with this is call get_state() so that we can pass that through to the background worker process. The background process will receive this, rebuild the iterator, then start yielding from it. However, in order to implement MultiprocessIterator.get_state(), we need to be able to accurately get (1) the state of the iterator in the worker process (2) the currently buffered items in the Queue To do this, we use: - batch_queue: This is the prefetch buffer the worker yields to and the main loop yields from - state_queue: This size 1 queue will be how the worker sends the iterator state once it has halted iterating. It must hold the state in addition to the last batch, if the queue was full at the time the stop event is sent. - stop_iterating_event: Once this is issued from the main loop, the worker will stop iterating and enter cleanup. During cleanup, the iterator will send the state of the current iterator to the main loop, in addition to possibly the last batch if the batch_queue was full at the time - state_dumped_event: When the main loop issues the stop_iterating_event, it will wait until the state_dumped_event to attempt to get state from the state_queue. It must do this since the worker may take some time to create and send the state. Once received by the main loop, the main loop can safely store the Queue (plus maybe the last batch) as the prefetch buffer, get the worker iterator's state, and terminate the background process + delete associated objects. At this point, calling create_iter() again will bootstrap everything from the stored state and the old iterator will throw an error since it will not iterate anymore (so the caller must call create_iter() again to get a python iterator). """ def __init__( self, base_iterator: StatefulIterator, *, n_batches_to_prefetch: int, prefetch_buffer: list | None = None ): self.base_iterator = base_iterator self.n_batches_to_prefetch = n_batches_to_prefetch if prefetch_buffer is None: prefetch_buffer = [] self.prefetch_buffer = prefetch_buffer self.batch_queue = None self.state_queue = None self.producer = None self.stop_iterating_event = None self.state_dumped_event = None def get_state(self) -> MultiprocessIteratorState: """ This is slightly unusual in effectively destroying the current iterator, its necessary to halt the background process and allow it to write the state to the main loop in order to not lose data """ if self.producer is None: serialized_prefetch_buffer = json.dumps( [b.to_python_dict() for b in self.prefetch_buffer] ) return MultiprocessIteratorState( base_iterator_state=self.base_iterator.get_state(), n_batches_to_prefetch=self.n_batches_to_prefetch, serialized_prefetch_buffer=serialized_prefetch_buffer, ) else: logging.info("Main thread: Sending stop iteration event") self.stop_iterating_event.set() logging.info("Main thread: Waiting for state_dumped event") self.state_dumped_event.wait() self.prefetch_buffer = [] final_batch_received = False while True: try: batch = self.batch_queue.get(timeout=1) if batch.is_final: final_batch_received = True break self.prefetch_buffer.append(batch) except Empty: logging.warning("Main thread: batch_queue is abnormally empty") assert final_batch_received try: base_iterator_state = self.state_queue.get(timeout=1) assert isinstance(base_iterator_state, IteratorState) except Empty: raise ValueError( "Attempted to get the state, but it was unexpectantly missing" ) self.base_iterator = base_iterator_state.build() self.producer.close() self.producer = None self.batch_queue = None self.state_queue = None self.stop_iterating_event = None self.state_dumped_event = None return MultiprocessIteratorState( base_iterator_state=self.base_iterator.get_state(), n_batches_to_prefetch=self.n_batches_to_prefetch, serialized_prefetch_buffer=json.dumps( [b.to_python_dict() for b in self.prefetch_buffer] ), ) def create_iter(self): logging.info("Main thread: Creating MP iterator") # First yield from the stored prefetch buffer. if self.prefetch_buffer is not None: while len(self.prefetch_buffer) > 0: item = self.prefetch_buffer.pop(0) yield item self.prefetch_buffer = None assert ( self.producer is None ), "Cannot create two parallel iterators at once, call get_state() then remake to have two." # using mp context manager avoids excessive CPU loading ctx = mp.get_context("forkserver") self.batch_queue = ctx.Manager().Queue(maxsize=self.n_batches_to_prefetch) # We should only ever one state, which is output at the detection of a stop event self.state_queue = ctx.Manager().Queue(maxsize=1) self.stop_iterating_event = ctx.Event() self.state_dumped_event = ctx.Event() self.producer = mp.Process( name="blt_data_loader", target=start_work_from_state, args=( self.batch_queue, self.state_queue, self.stop_iterating_event, self.state_dumped_event, self.base_iterator.get_state(), ), ) logger.info("Async dataloader started") self.producer.start() while True: if self.producer.exitcode is not None: raise RuntimeError( "Data loader quit unexpectedly, real error has been raised previously" ) try: batch = self.batch_queue.get(timeout=0.1) assert isinstance(batch, Batch) assert ( not batch.is_final ), "is_final should only be used during get_state() being called" yield batch except Empty: pass if self.producer is None: raise ValueError( "Attempted to call this iterator after calling get_state(). You must call create_iter() to make a new iterator instead." )