remove logger in data scripts and callback

llmfoundry/callbacks/run_timeout_callback.py

@@ -17,7 +17,10 @@
 
 def _timeout(timeout: int, mosaicml_logger: Optional[MosaicMLLogger] = None):
     log.error(f'Timeout after {timeout} seconds of inactivity after fit_end.',)
-    if mosaicml_logger is not None:
+    if mosaicml_logger is not None and os.environ.get(
+        'OVERRIDE_EXCEPTHOOK',
+        'false',
+    ).lower() != 'true':
         mosaicml_logger.log_exception(RunTimeoutError(timeout=timeout))
     os.kill(os.getpid(), signal.SIGINT)
 

scripts/data_prep/convert_delta_to_json.py

@@ -27,8 +27,9 @@
 from packaging import version
 from pyspark.sql import SparkSession
 from pyspark.sql.connect.client.core import SparkConnectClient
-from pyspark.sql.connect.client.reattach import \
-    ExecutePlanResponseReattachableIterator
+from pyspark.sql.connect.client.reattach import (
+    ExecutePlanResponseReattachableIterator,
+)
 from pyspark.sql.connect.dataframe import DataFrame
 from pyspark.sql.dataframe import DataFrame as SparkDataFrame
 from pyspark.sql.types import Row
@@ -644,6 +645,9 @@ def fetch_DT(args: Namespace) -> None:
         log.info(f'Elapsed time {time.time() - tik}')
 
     except Exception as e:
-        if mosaicml_logger is not None:
+        if mosaicml_logger is not None and os.environ.get(
+            'OVERRIDE_EXCEPTHOOK',
+            'false',
+        ).lower() != 'true':
             mosaicml_logger.log_exception(e)
         raise e

scripts/data_prep/convert_text_to_mds.py

@@ -624,6 +624,9 @@ def _configure_logging(logging_level: str):
             args_str=_args_str(args),
         )
     except Exception as e:
-        if mosaicml_logger is not None:
+        if mosaicml_logger is not None and os.environ.get(
+            'OVERRIDE_EXCEPTHOOK',
+            'false',
+        ).lower() != 'true':
             mosaicml_logger.log_exception(e)
         raise e

scripts/train/train.py

@@ -394,10 +394,12 @@ def main(cfg: DictConfig) -> Trainer:
             train_cfg.device_train_batch_size,
         )
     except BaseContextualError as e:
-        if mosaicml_logger is not None:
-            e.location = TrainDataLoaderLocation
-            if os.environ.get('OVERRIDE_EXCEPTHOOK', 'false').lower() != 'true':
-                mosaicml_logger.log_exception(e)
+        e.location = TrainDataLoaderLocation
+        if mosaicml_logger is not None and os.environ.get(
+            'OVERRIDE_EXCEPTHOOK',
+            'false',
+        ).lower() != 'true':
+            mosaicml_logger.log_exception(e)
         raise e
 
     if mosaicml_logger is not None:
@@ -428,11 +430,12 @@ def main(cfg: DictConfig) -> Trainer:
             if eval_gauntlet_callback is not None:
                 callbacks.append(eval_gauntlet_callback)
         except BaseContextualError as e:
-            if mosaicml_logger is not None:
-                e.location = EvalDataLoaderLocation
-                if os.environ.get('OVERRIDE_EXCEPTHOOK',
-                                  'false').lower() != 'true':
-                    mosaicml_logger.log_exception(e)
+            e.location = EvalDataLoaderLocation
+            if mosaicml_logger is not None and os.environ.get(
+                'OVERRIDE_EXCEPTHOOK',
+                'false',
+            ).lower() != 'true':
+                mosaicml_logger.log_exception(e)
             raise e
 
     if mosaicml_logger is not None:
@@ -480,8 +483,11 @@ def main(cfg: DictConfig) -> Trainer:
                 non_icl_metrics,
             )
     except BaseContextualError as e:
-        if mosaicml_logger is not None:
-            e.location = EvalDataLoaderLocation
+        e.location = EvalDataLoaderLocation
+        if mosaicml_logger is not None and os.environ.get(
+            'OVERRIDE_EXCEPTHOOK',
+            'false',
+        ).lower() != 'true':
             mosaicml_logger.log_exception(e)
         raise e