import datetime import logging import os import pickle from typing import Any, Dict, List, Optional, Union import pyarrow as pa import pyarrow.parquet as pq import torch from google.cloud import bigquery from lightning import LightningModule, Trainer from lightning.pytorch.callbacks import BasePredictionWriter from src.models.components.interfaces import ModelOutput from src.utils.decorators import RetriesFailedException, retry from src.utils.tensor_utils import merge_list_of_keyed_tensors_to_single_tensor log = logging.getLogger(__name__) class BaseBufferedWriter(BasePredictionWriter): def __init__( self, # Update this based on data size. The goal is to limit the number of writes to BQ without exceeding the memory of your machine. flush_frequency: int = 5000, write_interval: str = "batch", schema: Optional[List[Union[bigquery.SchemaField, pa.Field]]] = None, prediction_key_name: Optional[str] = None, prediction_name: Optional[str] = None, ): """ Args: flush_frequency: Number of rows to accumulate before flushing write_interval: "batch" or "epoch". schema: (Optional) Schema for the output data. prediction_key_name: (Optional) The key to key the predictions by e.g. user_id, item_id, etc. If not provided, it will be the module's responsibility to either hardcode or set it through the module's attributes. prediction_name: (Optional) The name for the prediction ie. label, logits, etc. If not provided, it will be the module's responsibility to either hardcode or set it through the module's attributes. """ super().__init__(write_interval) self.flush_frequency = flush_frequency # Buffer to accumulate rows before writing self.rows_buffer: List[dict] = [] self.schema = schema self.global_rank = None self.prediction_key_name = prediction_key_name self.prediction_name = prediction_name def setup(self, trainer: Trainer, pl_module: LightningModule, stage: str) -> None: self.global_rank = trainer.global_rank if trainer.global_rank else 0 log.info(f"Rank {self.global_rank} initialized for inference.") # If the module does not have the prediction_key_name or prediction_name attributes, # we don't do anything as it might be using the previous interface. # If the module has the attributes and they are set, we also don't change it as they # might be hardcoded in the module. # We only set the prediction_key_name or prediction_name if they are not set in the module but are # passed in the callback. This allows us to control the prediction_key_name and prediction_name # from the callback config. # TODO (lneves) Deprecate the way folks do and have all inference pipelines to use this. if ( hasattr(pl_module, "prediction_key_name") and pl_module.prediction_key_name is None and self.prediction_key_name is not None ): pl_module.prediction_key_name = self.prediction_key_name if ( hasattr(pl_module, "prediction_name") and pl_module.prediction_name is None and self.prediction_name is not None ): pl_module.prediction_name = self.prediction_name def flush_buffer(self) -> None: """Flush the buffer and then clear it.""" if self.rows_buffer: self._flush_buffer() self.rows_buffer.clear() else: log.info("Buffer is empty, nothing to flush.") def _flush_buffer(self) -> None: """Override this method to implement the logic for flushing the buffer.""" raise NotImplementedError( "You need to implement the `_flush_buffer` method in your subclass." ) def handle_batch(self, model_output: ModelOutput) -> None: """ Handles a batch of predictions by appending them to the rows buffer and flushing the buffer if it is full. Args: model_output: The ModelOutput object containing the predictions. """ if model_output is None: log.warning( f"Rank {self.global_rank} received an empty model output. Skipping this batch. This is expected if the batch is a dummy batch." ) return None rows = model_output.list_of_row_format if len(rows) + len(self.rows_buffer) < self.flush_frequency: self.rows_buffer.extend(rows) else: rows_left = self.flush_frequency - len(self.rows_buffer) self.rows_buffer.extend(rows[:rows_left]) self.flush_buffer() self.rows_buffer.extend(rows[rows_left:]) def write_on_batch_end( self, trainer: Trainer, pl_module: LightningModule, prediction: ModelOutput, batch_indices: List[int], batch: Any, batch_idx: int, dataloader_idx: int, ) -> None: """ Called at the end of each prediction batch. We'll accumulate rows in the buffer and only write to BigQuery once we reach the flush_frequency. """ self.handle_batch(prediction) def write_on_epoch_end( self, trainer: Trainer, pl_module: LightningModule, predictions: List[ModelOutput], batch_indices: List[List[int]], ) -> None: """ Called at the end of a prediction epoch. We'll continue to buffer predictions from all batches in this epoch and flush if the buffer exceeds flush_frequency. """ # predictions is typically a list of tensors (one per batch). for batch_pred in predictions: self.handle_batch(batch_pred) # At the very end of the epoch, flush any remaining rows self.flush_buffer() def on_predict_end( self, trainer: Trainer, pl_module: LightningModule, ) -> None: """ Called at the end of the prediction process. We'll flush any remaining rows in the buffer. """ self.flush_buffer() log.info(f"Rank {self.global_rank} finished writing predictions.") # TODO (clark): technically write_on_epoch_end should handle this correctly # but if we dont't do this as well here, the number of rows in the final BQ table will # always be a multiplier of flush_frequency # and we will lose some rows if the last batch is smaller than flush_frequency # this is an indicator of something not working fully as expected # i'll investigate this issue later class LocalPickleWriter(BaseBufferedWriter): """ Callback to write predictions to local pickle files during inference. """ def __init__( self, output_dir: str, flush_frequency: int = 1000, write_interval: str = "batch", should_merge_files_on_main: bool = True, should_merge_list_of_keyed_tensors_to_single_tensor: bool = True, post_processing_functions: Optional[List[Dict[str, callable]]] = [], **kwargs, ): """ Args: output_dir: Directory to save the pickle files. flush_frequency: Number of rows to accumulate before writing to a pickle file. write_interval: "batch" or "epoch". should_merge_files_on_main: If True, merge all files on the main process after writing. should_merge_list_of_keyed_tensors_to_single_tensor: If True, merge list of keyed tensors to a single tensor. post_processing_functions: List of ordered post-processing functions to apply to the files. """ super().__init__( write_interval=write_interval, flush_frequency=flush_frequency, **kwargs ) self.output_dir = output_dir os.makedirs(self.output_dir, exist_ok=True) self.should_merge_files_on_main = should_merge_files_on_main self.should_merge_list_of_keyed_tensors_to_single_tensor = ( should_merge_list_of_keyed_tensors_to_single_tensor ) self.post_processing_functions = post_processing_functions def _create_file_path(self) -> str: """Create a file path for the pickle file.""" return f"predictions_{self.global_rank}_{datetime.datetime.now(datetime.timezone.utc).strftime('%Y%m%dT%H%M%S%f')[:-3]}.pkl" def _local_file_path(self, file_path: Optional[str] = None) -> str: """Create a local file path for the pickle file.""" return ( f"{self.output_dir}/{file_path if file_path else self._create_file_path()}" ) @retry() def _flush_buffer(self): """Flush the buffer to a pickle file.""" file_path = self._create_file_path() with open(self._local_file_path(file_path=file_path), "wb") as f: pickle.dump(self.rows_buffer, f) log.info( f"Global Rank: {self.global_rank} wrote {len(self.rows_buffer)} rows to {self._local_file_path(file_path=file_path)}." ) @retry() def on_predict_end( self, trainer: Trainer, pl_module: LightningModule, ) -> None: super().on_predict_end(trainer, pl_module) if self.should_merge_files_on_main: # if we use multiple workers, we need to wait for all of them to finish writing # before merging the files if trainer.global_rank != None: torch.distributed.barrier() if self.global_rank == 0: log.info("Merging pickle files on main process.") self._merge_files() # other processes can continue after merging if trainer.global_rank != None: torch.distributed.barrier() # conducting post-processing functions on the files for process_func in self.post_processing_functions: all_files = [f for f in os.listdir(self.output_dir)] for file in all_files: file_path = os.path.join(self.output_dir, file) if process_func.get("main_only", False): if self.global_rank == 0: process_func["function"](file_path) else: process_func["function"](file_path) if trainer.global_rank != None: torch.distributed.barrier() def _merge_files(self): """Merge all pickle files in the output directory into a single file.""" all_files = [f for f in os.listdir(self.output_dir) if f.endswith(".pkl")] merged_data = [] for file in all_files: with open(os.path.join(self.output_dir, file), "rb") as f: merged_data.extend(pickle.load(f)) os.remove(os.path.join(self.output_dir, file)) with open(os.path.join(self.output_dir, "merged_predictions.pkl"), "wb") as f: pickle.dump(merged_data, f) log.info(f"Merged {len(merged_data)} rows into merged_predictions.pkl.") if self.should_merge_list_of_keyed_tensors_to_single_tensor: merged_data_tensor = merge_list_of_keyed_tensors_to_single_tensor( data=merged_data, index_key=self.prediction_key_name, value_key=self.prediction_name, ) torch.save( merged_data_tensor.cpu(), os.path.join(self.output_dir, "merged_predictions_tensor.pt"), ) log.info( f"Merged {len(merged_data_tensor)} rows into merged_predictions_tensor.pt. as pytorch tensor" )