Add configurable engine log level

deepspeed/runtime/config.py

@@ -637,6 +637,10 @@ def get_dataloader_drop_last(param_dict):
     return get_scalar_param(param_dict, DATALOADER_DROP_LAST, DATALOADER_DROP_LAST_DEFAULT)
 
 
+def get_log_level(param_dict):
+    return get_scalar_param(param_dict, LOG_LEVEL, LOG_LEVEL_DEFAULT)
+
+
 '''Write deepspeed config files by modifying basic templates.
 Can be used for quickly changing parameters via command line parameters.'''
 
@@ -873,6 +877,8 @@ def _initialize_params(self, param_dict):
 
         self.dataloader_drop_last = get_dataloader_drop_last(param_dict)
 
+        self.log_level = get_log_level(param_dict)
+
         self.nebula_config = DeepSpeedNebulaConfig(param_dict)
         self.datastates_config = DeepSpeedDataStatesConfig(param_dict)
         self.checkpoint_config = get_checkpoint_config(param_dict)

deepspeed/runtime/constants.py

@@ -510,3 +510,6 @@ class ValidationMode:
 #########################################
 USE_DATA_BEFORE_EXPERT_PARALLEL = "use_data_before_expert_parallelism"
 USE_DATA_BEFORE_EXPERT_PARALLEL_DEFAULT = False
+
+LOG_LEVEL = "log_level"
+LOG_LEVEL_DEFAULT = None

deepspeed/runtime/engine.py

@@ -116,7 +116,7 @@
 from ..git_version_info import version
 
 from deepspeed.profiling.flops_profiler.profiler import FlopsProfiler
-from deepspeed.utils.logging import print_json_dist, print_configuration
+from deepspeed.utils.logging import print_json_dist, print_configuration, set_log_level_from_string
 
 from deepspeed.accelerator import get_accelerator
 
@@ -277,6 +277,8 @@ def __init__(self,
         self._do_args_sanity_check(args)
         self._configure_with_arguments(args, mpu)
         self._do_sanity_check()
+        if self.log_level() is not None:
+            set_log_level_from_string(self.log_level())
         self._configure_expert_parallel(model)
         if self.autotp_size() > 1:
             self._configure_tensor_parallel(model, self.tensor_parallel_config())
@@ -1228,6 +1230,9 @@ def use_node_local_storage(self):
     def load_universal_checkpoint(self):
         return self._config.load_universal_checkpoint
 
+    def log_level(self):
+        return self._config.log_level
+
     @property
     def communication_data_type(self):
         res = self._config.communication_data_type

tests/unit/v1/test_engine_log_level.py

@@ -0,0 +1,80 @@
+# Copyright (c) DeepSpeed Team.
+# SPDX-License-Identifier: Apache-2.0
+
+# DeepSpeed Team
+
+import logging
+import types
+
+import pytest
+
+import deepspeed
+from deepspeed.runtime.engine import DeepSpeedEngine
+from deepspeed.runtime.config import get_log_level
+from deepspeed.runtime.constants import LOG_LEVEL_DEFAULT
+from deepspeed.utils import logging as ds_logging
+from unit.common import DistributedTest
+from unit.simple_model import SimpleModel
+
+
+class TestLogLevelConfig:
+
+    def test_default_is_none(self):
+        assert get_log_level({}) is None
+        assert LOG_LEVEL_DEFAULT is None
+
+    @pytest.mark.parametrize("level", ["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"])
+    def test_explicit_value(self, level):
+        assert get_log_level({"log_level": level}) == level
+
+    def test_engine_getter_returns_config_value(self):
+        engine = types.SimpleNamespace(_config=types.SimpleNamespace(log_level="ERROR"))
+        assert DeepSpeedEngine.log_level(engine) == "ERROR"
+
+
+class TestLogLevelEndToEnd(DistributedTest):
+    world_size = 1
+
+    def _config(self, log_level=None):
+        config = {
+            "train_batch_size": 1,
+            "optimizer": {
+                "type": "Adam",
+                "params": {
+                    "lr": 1e-3,
+                    "torch_adam": True
+                }
+            },
+        }
+        if log_level is not None:
+            config["log_level"] = log_level
+        return config
+
+    def _init(self, log_level=None):
+        model = SimpleModel(hidden_dim=8)
+        engine, _, _, _ = deepspeed.initialize(config=self._config(log_level),
+                                               model=model,
+                                               model_parameters=model.parameters())
+        return engine
+
+    def test_init_with_error_log_level(self):
+        saved = ds_logging.logger.level
+        ds_logging.logger.setLevel(logging.DEBUG)
+        try:
+            engine = self._init("ERROR")
+            assert engine.log_level() == "ERROR"
+            assert ds_logging.logger.getEffectiveLevel() == logging.ERROR
+        finally:
+            ds_logging.logger.setLevel(saved)
+
+    def test_omit_log_level_defaults_to_warning(self):
+        saved = ds_logging.logger.level
+        try:
+            ds_logging.logger.setLevel(logging.WARNING)
+            for handler in ds_logging.logger.handlers:
+                handler.setLevel(logging.WARNING)
+            engine = self._init()
+            assert engine.log_level() is None
+            assert ds_logging.logger.getEffectiveLevel() == logging.WARNING
+        finally:
+            ds_logging.logger.setLevel(saved)