Setting logging.raiseExceptions=True so that DATA_LOSS errors that occur in debug/info/other logs raise exceptions immediately.

Author: lukebaumann

axlearn/common/launch_trainer.py

@@ -4,6 +4,7 @@
 
 import contextlib
 import json
+import logging as py_logging
 import os
 from typing import Any, Optional
 
@@ -16,6 +17,8 @@
 from axlearn.common.utils import MeshShape, get_data_dir, infer_mesh_shape
 from axlearn.experiments import TrainerConfigFn, get_named_trainer_config
 
+logging.raiseException = True
+
 # Trainer-specific flags.
 flags.DEFINE_string(
     "module",