import collections import json import logging import os import pprint import time import tempfile from typing import Dict, Optional import ray.train from ray.data._internal.stats import Timer import torch from logger_utils import ContextLoggerAdapter from benchmark_factory import BenchmarkFactory logger = ContextLoggerAdapter(logging.getLogger(__name__)) class TrainLoopRunner: """Generic runner that sets up the training loop scaffolding. Collects perf metrics and handles periodic checkpointing and validation. """ def __init__(self, factory: BenchmarkFactory): self.factory = factory self.benchmark_config = factory.benchmark_config self._setup() # Training progress state. self._train_batch_idx: int = 0 self._train_epoch_idx: int = 0 self._global_rows_processed_this_epoch: int = 0 # Performance metrics self._metrics = collections.defaultdict(lambda: Timer()) checkpoint = ray.train.get_checkpoint() if checkpoint: self._restore_from_checkpoint(checkpoint) # Methods for subclasses to implement. def _setup(self): """Subclasses should override this to setup the model, optimizer, etc. The attributes initialized in this method should only be used in the other overridden methods.""" pass def _cleanup(self): """Subclasses can override this to cleanup any resources.""" pass def _train_step(self, train_dataloader): """Subclasses should override this to implement the training step. A training step represents a single forward and backward pass on a batch of data. """ raise NotImplementedError def _validate_step(self, val_dataloader): """Subclasses should override this to implement the validation step. A validation step represents a single forward pass on a batch of data.""" raise NotImplementedError def _save_training_state(self, local_dir: str): """Subclasses should override this to save the training state. This should reference the model and optimizer state initialized in the `_setup` method.""" pass def _load_training_state(self, local_dir: str): """Subclasses should override this to load the training state. This should reference the model and optimizer state initialized in the `_setup` method.""" pass def _restore_from_checkpoint(self, checkpoint: ray.train.Checkpoint): logger.info( f"Restoring from checkpoint: {checkpoint} for worker " f"{ray.train.get_context().get_world_rank()}" ) with tempfile.TemporaryDirectory( dir="/mnt/local_storage" ) as temp_checkpoint_dir: download_start = time.perf_counter() checkpoint.to_directory(temp_checkpoint_dir) download_time = time.perf_counter() - download_start load_start = time.perf_counter() self._load_checkpoint(temp_checkpoint_dir) load_time = time.perf_counter() - load_start self._metrics["checkpoint/download"].add(download_time) self._metrics["checkpoint/load"].add(load_time) def _wrap_dataloader(self, dataloader, train: bool = True): dataloader_iter = iter(dataloader) prefix = "train" if train else "validation" def dataloader_with_timers(): try: with self._metrics[f"{prefix}/iter_first_batch"].timer(): batch = next(dataloader_iter) if train: self._train_batch_idx += 1 except StopIteration: return while True: yield batch try: with self._metrics[f"{prefix}/iter_batch"].timer(): batch = next(dataloader_iter) if train: self._train_batch_idx += 1 except StopIteration: return return dataloader_with_timers() @property def _num_batches_to_skip(self) -> int: """Calculate the number of batches to skip based on the number of rows already processed in this epoch.""" global_batch_size = ( self.benchmark_config.dataloader_config.train_batch_size * ray.train.get_context().get_world_size() ) return self._global_rows_processed_this_epoch // global_batch_size def _train_epoch(self): """Subclasses can override the entrire `_train_epoch` method for more training logic customization.""" if ray.train.get_context().get_world_rank() == 0: logger.info(f"Training starting @ epoch={self._train_epoch_idx}") train_dataloader = self.factory.get_train_dataloader() train_dataloader = self._wrap_dataloader(train_dataloader, train=True) # Skip through batches if we restored to a middle of the epoch. # TODO: Compare this baseline to the data checkpointing approach once we have it. if self._num_batches_to_skip: if ray.train.get_context().get_world_rank() == 0: logger.info(f"Skipping {self._num_batches_to_skip} batches...") # Zero before the skip loop drives the wrapper, which would # otherwise double-count against the value restored from the # checkpoint. After the skip, _train_batch_idx is rebuilt to # _num_batches_to_skip — matching the restored value. self._train_batch_idx = 0 for _ in range(self._num_batches_to_skip): with self._metrics["train/iter_skip_batch"].timer(): next(train_dataloader) for batch in train_dataloader: with self._metrics["train/step"].timer(): if not self.benchmark_config.skip_train_step: self._train_step(batch) if self.benchmark_config.train_step_sleep_s > 0: time.sleep(self.benchmark_config.train_step_sleep_s) # TODO: This is slightly off if the last batch is a partial batch (if drop_last=False) global_batch_size = ( self.benchmark_config.dataloader_config.train_batch_size * ray.train.get_context().get_world_size() ) self._metrics["train/rows_processed"].add(global_batch_size) self._global_rows_processed_this_epoch += global_batch_size if self._should_checkpoint_during_epoch(): self._checkpoint() if self._should_validate_during_epoch(): validation_metrics = self._validate() self._checkpoint(validation_metrics) if self._should_log_metrics(): logger.info(pprint.pformat(self.get_metrics(), indent=2)) if ( self.benchmark_config.max_train_batches > 0 and self._train_batch_idx >= self.benchmark_config.max_train_batches ): break self._train_epoch_idx += 1 self._train_batch_idx = 0 self._global_rows_processed_this_epoch = 0 def _validate_epoch(self) -> Dict[str, float]: if ray.train.get_context().get_world_rank() == 0: logger.info( f"Validation starting @ epoch={self._train_epoch_idx}, " f"batch={self._train_batch_idx}" ) val_dataloader = self.factory.get_val_dataloader() val_dataloader = self._wrap_dataloader(val_dataloader, train=False) total_loss = torch.tensor(0.0).to(ray.train.torch.get_device()) num_rows = 0 for batch in val_dataloader: with self._metrics["validation/step"].timer(): if not self.benchmark_config.skip_validation_step: total_loss += self._validate_step(batch) num_rows += self.benchmark_config.dataloader_config.validation_batch_size self._metrics["validation/rows_processed"].add( self.benchmark_config.dataloader_config.validation_batch_size ) assert num_rows > 0, "Validation dataset yielded no batches." return {"validation/loss": total_loss.item() / num_rows} def _should_checkpoint_during_epoch(self) -> bool: """Handles the checkpoint_every_n_steps logic.""" return ( self.benchmark_config.checkpoint_every_n_steps > 0 and self._train_batch_idx % self.benchmark_config.checkpoint_every_n_steps == 0 ) def _should_validate_during_epoch(self) -> bool: """Handles the validate_every_n_steps logic.""" return ( self.benchmark_config.validate_every_n_steps > 0 and self._train_batch_idx % self.benchmark_config.validate_every_n_steps == 0 ) def _should_log_metrics(self) -> bool: """Handles the log_metrics_every_n_steps logic.""" return ( self.benchmark_config.log_metrics_every_n_steps > 0 and self._train_batch_idx % self.benchmark_config.log_metrics_every_n_steps == 0 ) def _validate(self) -> Dict[str, float]: with self._metrics["validation/epoch"].timer(): validation_metrics = self._validate_epoch() return validation_metrics def _checkpoint(self, metrics: Optional[Dict[str, float]] = None): with tempfile.TemporaryDirectory( dir="/mnt/local_storage" ) as temp_checkpoint_dir: with self._metrics["checkpoint/save"].timer(): self._save_checkpoint(temp_checkpoint_dir) with self._metrics["checkpoint/report"].timer(): self._report_checkpoint( metrics=metrics or {}, checkpoint=ray.train.Checkpoint.from_directory(temp_checkpoint_dir), ) def _load_checkpoint(self, local_dir: str): self._load_training_state(local_dir) run_state = torch.load(os.path.join(local_dir, "run_state.pt")) self._train_epoch_idx = run_state["epoch"] self._train_batch_idx = run_state["batch_idx"] self._global_rows_processed_this_epoch = run_state[ "global_rows_processed_this_epoch" ] with open(os.path.join(local_dir, "metrics.json"), "r") as f: metrics_json = json.load(f) for k, v in metrics_json.items(): self._metrics[k].from_dict(v) if ray.train.get_context().get_world_rank() == 0: logger.info( f"Restored to epoch={self._train_epoch_idx}, " f"train_batch_idx={self._train_batch_idx} from checkpoint: " f"{ray.train.get_checkpoint()}" ) def _save_checkpoint(self, local_dir: str): logger.info( f"Saving checkpoint @ epoch={self._train_epoch_idx}, " f"train_batch_idx={self._train_batch_idx}" ) self._save_training_state(local_dir) if ray.train.get_context().get_world_rank() == 0: run_state = { "epoch": self._train_epoch_idx, "batch_idx": self._train_batch_idx, "global_rows_processed_this_epoch": self._global_rows_processed_this_epoch, } torch.save(run_state, os.path.join(local_dir, "run_state.pt")) metrics_json = {k: v.as_dict() for k, v in self._metrics.items()} with open(os.path.join(local_dir, "metrics.json"), "w") as f: json.dump(metrics_json, f) def _report_checkpoint(self, metrics, checkpoint): logger.info( f"Uploading checkpoint @ epoch={self._train_epoch_idx}, " f"train_batch_idx={self._train_batch_idx}" ) checkpoint_dir_name = ( f"checkpoint_epoch={self._train_epoch_idx}_batch={self._train_batch_idx}" ) ray.train.report( metrics, checkpoint=checkpoint, checkpoint_dir_name=checkpoint_dir_name, ) def run(self): starting_epoch = self._train_epoch_idx for _ in range(starting_epoch, self.benchmark_config.num_epochs): with self._metrics["train/epoch"].timer(): self._train_epoch() if not self.benchmark_config.skip_validation_at_epoch_end: validation_metrics = self._validate() self._checkpoint(validation_metrics) if ray.train.get_context().get_world_rank() == 0: logger.info(pprint.pformat(self.get_metrics(), indent=2)) self._cleanup() def get_metrics(self, dataset_creation_time: float = 0.0) -> Dict[str, float]: # TODO: These metrics should be aggregated across training workers. metrics = {} for key, metric in self._metrics.items(): metrics.update( { f"{key}-avg": metric.avg(), f"{key}-min": metric.min(), f"{key}-max": metric.max(), f"{key}-total": metric.get(), } ) metrics["train/dataset_creation_time"] = dataset_creation_time metrics["validation/dataset_creation_time"] = dataset_creation_time # Throughput # TODO: Ray Data can provide these throughput metrics automatically. train_time = ( metrics["train/dataset_creation_time"] + self._metrics["train/step"].get() # Include the time it takes to get the first batch. + self._metrics["train/iter_first_batch"].get() + self._metrics["train/iter_batch"].get() ) if train_time > 0: metrics["train/global_throughput"] = ( self._metrics["train/rows_processed"].get() / train_time ) validation_time = ( metrics["validation/dataset_creation_time"] + self._metrics["validation/step"].get() # Include the time it takes to get the first batch. + self._metrics["validation/iter_first_batch"].get() + self._metrics["validation/iter_batch"].get() ) if validation_time > 0: metrics["validation/global_throughput"] = ( self._metrics["validation/rows_processed"].get() / validation_time ) # Extra time that each worker spends to restore from checkpoint, # which includes downloading the checkpoint, loading the checkpoint, # and skipping through batches that were already processed. restoration_time = ( self._metrics["checkpoint/download"].get() + self._metrics["checkpoint/load"].get() + self._metrics["train/iter_skip_batch"].get() ) if restoration_time > 0: metrics["checkpoint/restoration_time"] = restoration_time # Dataloader metrics (ex: Ray Data stats) metrics.update(self.factory.get_dataloader_metrics()) return metrics class VanillaTorchRunner(TrainLoopRunner): """A simple runner that uses a PyTorch model, optimizer, and loss function.""" def _setup(self): model = self.factory.get_model() self.model = ray.train.torch.prepare_model(model) self.loss_fn = self.factory.get_loss_fn() self.optimizer = torch.optim.Adam(self.model.parameters(), lr=1e-3) def _train_step(self, batch): self.model.train() input_batch, labels = batch self.model.train() self.optimizer.zero_grad() out = self.model(input_batch) loss = self.loss_fn(out, labels) loss.backward() self.optimizer.step() def _validate_step(self, batch): self.model.eval() input_batch, labels = batch with torch.no_grad(): out = self.model(input_batch) loss = self.loss_fn(out, labels) return loss def _save_training_state(self, local_dir: str): # Standard DDP checkpointing. if ray.train.get_context().get_world_rank() == 0: torch.save(self.model.state_dict(), os.path.join(local_dir, "model.pt")) torch.save( self.optimizer.state_dict(), os.path.join(local_dir, "optimizer.pt") ) def _load_training_state(self, local_dir: str): self.model.load_state_dict( torch.load(os.path.join(local_dir, "model.pt"), map_location="cpu") ) self.optimizer.load_state_dict( torch.load(os.path.join(local_dir, "optimizer.pt"), map_location="cpu") )