correctly set log level for benchmark runs (meta-pytorch#3494)

Summary:
Pull Request resolved: meta-pytorch#3494

# context
* loglevel is not correctly set in train pipeline benchmark due to the multiprocess setup
* the log level is only set in the main process but not correctly set in the forked/spawn processes
* this diff add the `loglevel` argument into the RunConfig so that in every runner funcion can call `set_logger_level`
* also directly pass the error message on yaml or json parser failure, which previously just warn silently and the warning message is buried in lengthy logs.
* with loglevel=info we can now see the planner info: P2014482201

Reviewed By: spmex

Differential Revision: D85829837

fbshipit-source-id: 9719baf4307972a1794bf8870cd5c2df8add4436

Author: TroyGarden

torchrec/distributed/benchmark/base.py

@@ -425,17 +425,11 @@ def _load_config_file(
             if not config_path:
                 return {}
 
-            try:
-                with open(config_path, "r") as f:
-                    if is_json:
-                        return json.load(f) or {}
-                    else:
-                        return yaml.safe_load(f) or {}
-            except Exception as e:
-                logger.error(
-                    f"Failed to load config because {e}. Proceeding without it."
-                )
-                return {}
+            with open(config_path, "r") as f:
+                if is_json:
+                    return json.load(f) or {}
+                else:
+                    return yaml.safe_load(f) or {}
 
         @functools.wraps(func)
         def wrapper() -> Any:  # pyre-ignore [3]
@@ -479,7 +473,12 @@ def wrapper() -> Any:  # pyre-ignore [3]
             # Merge the two dictionaries, JSON overrides YAML
             merged_defaults = {**yaml_defaults, **json_defaults}
 
-            seen_args = set()  # track all --<name> we've added
+            # track all --<name> we've added
+            seen_args = {
+                "json_config",
+                "yaml_config",
+                "loglevel",
+            }
 
             for _name, param in sig.parameters.items():
                 cls = param.annotation
@@ -548,7 +547,12 @@ def wrapper() -> Any:  # pyre-ignore [3]
                     logger.info(config_instance)
 
             loglevel = logging._nameToLevel[args.loglevel.upper()]
-            logger.setLevel(loglevel)
+            # Set loglevel for all existing loggers
+            for existing_logger_name in logging.root.manager.loggerDict:
+                existing_logger = logging.getLogger(existing_logger_name)
+                existing_logger.setLevel(loglevel)
+            # Also set the root logger
+            logging.root.setLevel(loglevel)
 
             return func(**kwargs)
 
@@ -857,6 +861,7 @@ class BenchFuncConfig:
     export_stacks: bool = False
     all_rank_traces: bool = False
     memory_snapshot: bool = False
+    loglevel: str = "WARNING"
 
     # pyre-ignore [2]
     def benchmark_func_kwargs(self, **kwargs_to_override) -> Dict[str, Any]:
@@ -873,6 +878,10 @@ def benchmark_func_kwargs(self, **kwargs_to_override) -> Dict[str, Any]:
             "memory_snapshot": self.memory_snapshot,
         } | kwargs_to_override
 
+    def set_log_level(self) -> None:
+        loglevel = logging._nameToLevel[self.loglevel.upper()]
+        logging.root.setLevel(loglevel)
+
 
 def benchmark_func(
     name: str,

torchrec/distributed/benchmark/benchmark_train_pipeline.py

@@ -129,7 +129,7 @@ def runner(
         torch.cuda.is_available() and torch.cuda.device_count() >= world_size
     ), "CUDA not available or insufficient GPUs for the requested world_size"
 
-    torch.autograd.set_detect_anomaly(True)
+    run_option.set_log_level()
     with MultiProcessContext(
         rank=rank,
         world_size=world_size,