track system metrics + total eval time + stdout

lm_eval/__main__.py

@@ -201,6 +201,9 @@ def cli_evaluate(args: Union[argparse.Namespace, None] = None) -> None:
         # we allow for args to be passed externally, else we parse them ourselves
         args = parse_eval_args()
 
+    if args.wandb_args:
+        wandb_logger = WandbLogger(args)
+
     eval_logger = utils.eval_logger
     eval_logger.setLevel(getattr(logging, f"{args.verbosity}"))
     eval_logger.info(f"Verbosity set to {args.verbosity}")
@@ -318,12 +321,10 @@ def cli_evaluate(args: Union[argparse.Namespace, None] = None) -> None:
         # Add W&B logging
         if args.wandb_args:
             try:
-                wandb_logger = WandbLogger(results, args)
+                wandb_logger.post_init(results)
                 wandb_logger.log_eval_result()
                 if args.log_samples:
                     wandb_logger.log_eval_samples(samples)
-                # Tear down wandb run once all the logging is done.
-                wandb_logger.run.finish()
             except Exception as e:
                 eval_logger.info(f"Logging to Weights and Biases failed due to {e}")
 
@@ -352,6 +353,10 @@ def cli_evaluate(args: Union[argparse.Namespace, None] = None) -> None:
         if "groups" in results:
             print(make_table(results, "groups"))
 
+        if args.wandb_args:
+            # Tear down wandb run once all the logging is done.
+            wandb_logger.run.finish()
+
 
 if __name__ == "__main__":
     cli_evaluate()

lm_eval/logging_utils.py

@@ -78,21 +78,17 @@ def get_wandb_printer() -> Literal["Printer"]:
 
 
 class WandbLogger:
-    def __init__(self, results: Dict[str, Any], args: Any) -> None:
+    def __init__(self, args: Any) -> None:
         """Initialize the WandbLogger.
 
         Args:
             results (Dict[str, Any]): The results dictionary.
             args (Any): Arguments for configuration.
         """
-        self.results: Dict[str, Any] = copy.deepcopy(results)
         self.wandb_args: Dict[str, Any] = utils.simple_parse_args_string(
             args.wandb_args
         )
 
-        self.task_names: List[str] = list(results.get("results", {}).keys())
-        self.group_names: List[str] = list(results.get("groups", {}).keys())
-
         # initialize a W&B run
         if wandb.run is None:
             self.run = wandb.init(**self.wandb_args)
@@ -101,6 +97,11 @@ def __init__(self, results: Dict[str, Any], args: Any) -> None:
 
         self.printer = get_wandb_printer()
 
+    def post_init(self, results: Dict[str, Any]) -> None:
+        self.results: Dict[str, Any] = copy.deepcopy(results)
+        self.task_names: List[str] = list(results.get("results", {}).keys())
+        self.group_names: List[str] = list(results.get("groups", {}).keys())
+
     def _get_config(self) -> Dict[str, Any]:
         """Get configuration parameters."""
         self.task_configs = self.results.get("configs", {})