Fix: Removal of transformers logger and addition of python native logger

build/accelerate_launch.py

@@ -61,9 +61,6 @@ def get_base_model_from_adapter_config(adapter_config):
 
 
 def main():
-    LOGLEVEL = os.environ.get("LOG_LEVEL", "WARNING").upper()
-    logging.basicConfig(level=LOGLEVEL)
-
     if not os.getenv("TERMINATION_LOG_FILE"):
         os.environ["TERMINATION_LOG_FILE"] = ERROR_LOG
 
@@ -80,6 +77,14 @@ def main():
             or 'SFT_TRAINER_CONFIG_JSON_ENV_VAR'."
             )
 
+        # Configure log_level of python native logger.
+        LOGLEVEL = None
+        if "log_level" in job_config and job_config["log_level"]:
+            LOGLEVEL = job_config["log_level"].upper()
+            logging.basicConfig(level=LOGLEVEL)
+        else:
+            LOGLEVEL = os.environ.get("LOG_LEVEL", "WARNING").upper()
+            logging.basicConfig(level=LOGLEVEL)
         args = process_accelerate_launch_args(job_config)
         logging.debug("accelerate launch parsed args: %s", args)
     except FileNotFoundError as e:
@@ -110,6 +115,10 @@ def main():
             updated_args = serialize_args(job_config)
             os.environ["SFT_TRAINER_CONFIG_JSON_ENV_VAR"] = updated_args
 
+            # Configure for Image to get log level as the code after
+            # launch_command only takes log level from env var LOG_LEVEL and not CLI
+            os.environ["LOG_LEVEL"] = LOGLEVEL
+
             launch_command(args)
         except subprocess.CalledProcessError as e:
             # If the subprocess throws an exception, the base exception is hidden in the

tests/test_sft_trainer.py

@@ -18,6 +18,7 @@
 # Standard
 import copy
 import json
+import logging
 import os
 import tempfile
 
@@ -780,3 +781,61 @@ def test_pretokenized_dataset_wrong_format():
         # is essentially swallowing a KeyError here.
         with pytest.raises(ValueError):
             sft_trainer.train(MODEL_ARGS, data_args, train_args, PEFT_PT_ARGS)
+
+
+def test_set_log_level_for_logger_default():
+    """
+    Ensure that the correct log level is being set
+    for python native logger and transformers logger
+    """
+
+    # Set env var TRANSFORMERS_VERBOSITY as None and test
+    os.unsetenv("TRANSFORMERS_VERBOSITY")
+    os.unsetenv("LOG_LEVEL")
+    os.unsetenv("SFT_TRAINER_CONFIG_JSON_ENV_VAR")
+    os.unsetenv("SFT_TRAINER_CONFIG_JSON_PATH")
+    train_args = copy.deepcopy(TRAIN_ARGS)
+
+    # TEST IF NO ENV VAR ARE SET AND NO CLI ARGUMENT IS PASSED
+    training_args, logger = sft_trainer.set_log_level(train_args)
+    assert logger.level == logging.WARNING
+    assert training_args.log_level == "warning"
+
+
+def test_set_log_level_for_logger_with_env_var():
+    """
+    Ensure that the correct log level is being set
+    for python native logger and transformers logger
+    """
+
+    os.unsetenv("TRANSFORMERS_VERBOSITY")
+    os.unsetenv("LOG_LEVEL")
+    os.unsetenv("SFT_TRAINER_CONFIG_JSON_ENV_VAR")
+    os.unsetenv("SFT_TRAINER_CONFIG_JSON_PATH")
+    train_args = copy.deepcopy(TRAIN_ARGS)
+
+    # TEST IF LOG_LEVEL ENV VAR IS SET AND NO CLI ARGUMENT IS PASSED
+    os.environ["LOG_LEVEL"] = "info"
+    training_args, logger = sft_trainer.set_log_level(train_args)
+    assert logger.level == logging.INFO
+    assert training_args.log_level == "info"
+
+
+def test_set_log_level_for_logger_with_env_var_and_cli():
+    """
+    Ensure that the correct log level is being set
+    for python native logger and transformers logger
+    """
+
+    os.unsetenv("TRANSFORMERS_VERBOSITY")
+    os.unsetenv("LOG_LEVEL")
+    os.unsetenv("SFT_TRAINER_CONFIG_JSON_ENV_VAR")
+    os.unsetenv("SFT_TRAINER_CONFIG_JSON_PATH")
+    train_args = copy.deepcopy(TRAIN_ARGS)
+
+    # TEST IF LOG_LEVEL ENV VAR IS SET AND --log_level CLI ARGUMENT IS PASSED
+    os.environ["LOG_LEVEL"] = "info"
+    train_args.log_level = "error"
+    training_args, logger = sft_trainer.set_log_level(train_args)
+    assert logger.level == logging.ERROR
+    assert training_args.log_level == "error"

tuning/config/configs.py

@@ -136,6 +136,13 @@ class TrainingArguments(transformers.TrainingArguments):
             + "Requires additional configs, see tuning.configs/tracker_configs.py"
         },
     )
+    log_level: str = field(
+        default="passive",
+        metadata={
+            "help": "The log level to adopt during training. \
+            Possible values are 'debug', 'info', 'warning', 'error' and 'critical'"
+        },
+    )
 
 
 @dataclass

tuning/sft_trainer.py

@@ -16,6 +16,8 @@
 from typing import Dict, List, Optional, Union
 import dataclasses
 import json
+import logging
+import os
 import sys
 import time
 import traceback
@@ -33,7 +35,7 @@
     LlamaTokenizerFast,
     TrainerCallback,
 )
-from transformers.utils import is_accelerate_available, logging
+from transformers.utils import is_accelerate_available
 from trl import SFTConfig, SFTTrainer
 import fire
 import transformers
@@ -111,8 +113,6 @@ def train(
             fused_lora and fast_kernels must used together (may change in future). \
     """
 
-    logger = logging.get_logger("sft_trainer")
-
     # Validate parameters
     if (not isinstance(train_args.num_train_epochs, (float, int))) or (
         train_args.num_train_epochs <= 0
@@ -218,9 +218,9 @@ def train(
         )
 
     max_seq_length = min(train_args.max_seq_length, tokenizer.model_max_length)
-    logger.info("Max sequence length is %s", max_seq_length)
+    logging.info("Max sequence length is %s", max_seq_length)
     if train_args.max_seq_length > tokenizer.model_max_length:
-        logger.warning(
+        logging.warning(
             "max_seq_length %s exceeds tokenizer.model_max_length \
             %s, using tokenizer.model_max_length %s",
             train_args.max_seq_length,
@@ -231,16 +231,16 @@ def train(
     # TODO: we need to change this, perhaps follow what open instruct does?
     special_tokens_dict = {}
     if tokenizer.pad_token is None:
-        logger.warning("PAD token set to default, missing in tokenizer")
+        logging.warning("PAD token set to default, missing in tokenizer")
         special_tokens_dict["pad_token"] = configs.DEFAULT_PAD_TOKEN
     if tokenizer.eos_token is None:
-        logger.warning("EOS token set to default, missing in tokenizer")
+        logging.warning("EOS token set to default, missing in tokenizer")
         special_tokens_dict["eos_token"] = configs.DEFAULT_EOS_TOKEN
     if tokenizer.bos_token is None:
-        logger.warning("BOS token set to default, missing in tokenizer")
+        logging.warning("BOS token set to default, missing in tokenizer")
         special_tokens_dict["bos_token"] = configs.DEFAULT_BOS_TOKEN
     if tokenizer.unk_token is None:
-        logger.warning("UNK token set to default, missing in tokenizer")
+        logging.warning("UNK token set to default, missing in tokenizer")
         special_tokens_dict["unk_token"] = configs.DEFAULT_UNK_TOKEN
 
     # TODO: lower priority but understand if resizing impacts inference quality and why its needed.
@@ -254,11 +254,11 @@ def train(
 
     # Configure the collator and validate args related to packing prior to formatting the dataset
     if train_args.packing:
-        logger.info("Packing is set to True")
+        logging.info("Packing is set to True")
         data_collator = None
         packing = True
     else:
-        logger.info("Packing is set to False")
+        logging.info("Packing is set to False")
         packing = False
 
     # Validate if data args are set properly
@@ -323,7 +323,7 @@ def train(
                     tracker.track(metric=v, name=k, stage="additional_metrics")
                     tracker.set_params(params=exp_metadata, name="experiment_metadata")
             except ValueError as e:
-                logger.error(
+                logging.error(
                     "Exception while saving additional metrics and metadata %s",
                     repr(e),
                 )
@@ -461,12 +461,46 @@ def parse_arguments(parser, json_config=None):
     )
 
 
-def main(**kwargs):  # pylint: disable=unused-argument
-    logger = logging.get_logger("__main__")
+def set_log_level(parsed_training_args):
+    """Set log level of python native logger and TF logger via argument from CLI or env variable.
+
+    Args:
+        parsed_training_args
+            Training arguments for training model.
+    """
+
+    # Clear any existing handlers if necessary
+    for handler in logging.root.handlers[:]:
+        logging.root.removeHandler(handler)
+
+    # Configure Python native logger log level
+    # If CLI arg is passed, assign same log level to python native logger
+    if parsed_training_args.log_level != "passive":
+        logging.basicConfig(level=parsed_training_args.log_level.upper())
+    else:
+        # Assign value of either env var LOG_LEVEL or warning
+        LOGLEVEL = os.environ.get("LOG_LEVEL", "WARNING").upper()
+        logging.basicConfig(level=LOGLEVEL)
+        train_logger = logging.getLogger()
+
+    # Check if env var TRANSFORMERS_VERBOSITY is not set.
+    # Else if env var is already set then, log level of transformers is automatically set.
+    if os.environ.get("TRANSFORMERS_VERBOSITY") is None:
+
+        # Check if "--log_level" CLI argument is not used (passive/warning is the default log level)
+        if parsed_training_args.log_level == "passive":
+            LOGLEVEL = os.environ.get("LOG_LEVEL", "WARNING").upper()
+
+            # Set log_level in TrainingArguments
+            parsed_training_args.log_level = LOGLEVEL.lower()
 
+    train_logger = logging.getLogger()
+    return parsed_training_args, train_logger
+
+
+def main(**kwargs):  # pylint: disable=unused-argument
     parser = get_parser()
     job_config = get_json_config()
-    logger.debug("Input args parsed: %s", job_config)
     # accept arguments via command-line or JSON
     try:
         (
@@ -481,7 +515,11 @@ def main(**kwargs):  # pylint: disable=unused-argument
             fusedops_kernels_config,
             exp_metadata,
         ) = parse_arguments(parser, job_config)
-        logger.debug(
+
+        # Function to set log level for python native logger and transformers training logger
+        training_args, _ = set_log_level(training_args)
+
+        logging.debug(
             "Input args parsed: \
             model_args %s, data_args %s, training_args %s, trainer_controller_args %s, \
             tune_config %s, file_logger_config, %s aim_config %s, \
@@ -511,12 +549,12 @@ def main(**kwargs):  # pylint: disable=unused-argument
         try:
             metadata = json.loads(exp_metadata)
             if metadata is None or not isinstance(metadata, Dict):
-                logger.warning(
+                logging.warning(
                     "metadata cannot be converted to simple k:v dict ignoring"
                 )
                 metadata = None
         except ValueError as e:
-            logger.error(
+            logging.error(
                 "failed while parsing extra metadata. pass a valid json %s", repr(e)
             )
 
@@ -539,27 +577,27 @@ def main(**kwargs):  # pylint: disable=unused-argument
             fusedops_kernels_config=fusedops_kernels_config,
         )
     except (MemoryError, OutOfMemoryError) as e:
-        logger.error(traceback.format_exc())
+        logging.error(traceback.format_exc())
         write_termination_log(f"OOM error during training. {e}")
         sys.exit(INTERNAL_ERROR_EXIT_CODE)
     except FileNotFoundError as e:
-        logger.error(traceback.format_exc())
+        logging.error(traceback.format_exc())
         write_termination_log("Unable to load file: {}".format(e))
         sys.exit(USER_ERROR_EXIT_CODE)
     except HFValidationError as e:
-        logger.error(traceback.format_exc())
+        logging.error(traceback.format_exc())
         write_termination_log(
             f"There may be a problem with loading the model. Exception: {e}"
         )
         sys.exit(USER_ERROR_EXIT_CODE)
     except (TypeError, ValueError, EnvironmentError) as e:
-        logger.error(traceback.format_exc())
+        logging.error(traceback.format_exc())
         write_termination_log(
             f"Exception raised during training. This may be a problem with your input: {e}"
         )
         sys.exit(USER_ERROR_EXIT_CODE)
     except Exception as e:  # pylint: disable=broad-except
-        logger.error(traceback.format_exc())
+        logging.error(traceback.format_exc())
         write_termination_log(f"Unhandled exception during training: {e}")
         sys.exit(INTERNAL_ERROR_EXIT_CODE)
 

tuning/trackers/aimstack_tracker.py

@@ -14,11 +14,11 @@
 
 # Standard
 import json
+import logging
 import os
 
 # Third Party
 from aim.hugging_face import AimCallback  # pylint: disable=import-error
-from transformers.utils import logging
 
 # Local
 from .tracker import Tracker
@@ -97,7 +97,10 @@ def __init__(self, tracker_config: AimConfig):
             information about the repo or the server and port where aim db is present.
         """
         super().__init__(name="aim", tracker_config=tracker_config)
-        self.logger = logging.get_logger("aimstack_tracker")
+        # Configure log level
+        LOGLEVEL = os.environ.get("LOG_LEVEL", "WARNING").upper()
+        logging.basicConfig(level=LOGLEVEL)
+        self.logger = logging.getLogger("aimstack_tracker")
 
     def get_hf_callback(self):
         """Returns the aim.hugging_face.AimCallback object associated with this tracker.

tuning/trackers/filelogging_tracker.py

@@ -15,11 +15,11 @@
 # Standard
 from datetime import datetime
 import json
+import logging
 import os
 
 # Third Party
 from transformers import TrainerCallback
-from transformers.utils import logging
 
 # Local
 from .tracker import Tracker
@@ -80,7 +80,10 @@ def __init__(self, tracker_config: FileLoggingTrackerConfig):
                 which contains the location of file where logs are recorded.
         """
         super().__init__(name="file_logger", tracker_config=tracker_config)
-        self.logger = logging.get_logger("file_logging_tracker")
+        # Configure log level
+        LOGLEVEL = os.environ.get("LOG_LEVEL", "WARNING").upper()
+        logging.basicConfig(level=LOGLEVEL)
+        self.logger = logging.getLogger("file_logging_tracker")
 
     def get_hf_callback(self):
         """Returns the FileLoggingCallback object associated with this tracker.

tuning/trackers/tracker_factory.py

@@ -14,16 +14,20 @@
 
 # Standard
 import dataclasses
+import logging
+import os
 
 # Third Party
-from transformers.utils import logging
 from transformers.utils.import_utils import _is_package_available
 
 # Local
 from .filelogging_tracker import FileLoggingTracker
 from tuning.config.tracker_configs import FileLoggingTrackerConfig, TrackerConfigFactory
 
-logger = logging.get_logger("tracker_factory")
+# Configure log level
+LOGLEVEL = os.environ.get("LOG_LEVEL", "WARNING").upper()
+logging.basicConfig(level=LOGLEVEL)
+logger = logging.getLogger("tracker_factory")
 
 
 # Information about all registered trackers