File size: 12,127 Bytes
d9115cb | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256 257 258 259 260 261 262 263 264 265 266 267 268 269 270 271 272 273 274 275 276 277 278 279 280 281 282 283 284 | 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"
)
|