[dataset] add ms_logger_context (#4428)

Author: Jintao-Huang

swift/hub/hub.py

@@ -1,4 +1,5 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
+import logging
 import os
 import tempfile
 from contextlib import contextmanager
@@ -11,11 +12,11 @@
 from huggingface_hub.hf_api import api, future_compatible
 from requests.exceptions import HTTPError
 from transformers import trainer
-from transformers.utils import logging, strtobool
+from transformers.utils import strtobool
 
-from swift.utils.env import use_hf_hub
+from swift.utils import get_logger, ms_logger_context, use_hf_hub
 
-logger = logging.get_logger(__name__)
+logger = get_logger()
 
 
 class HubOperation:
@@ -287,15 +288,15 @@ def load_dataset(cls,
         cls.try_login(token)
         if revision is None or revision == 'main':
             revision = 'master'
-
-        return MsDataset.load(
-            dataset_id,
-            subset_name=subset_name,
-            split=split,
-            version=revision,
-            download_mode=download_mode,
-            use_streaming=streaming,
-        )
+        with ms_logger_context(logging.ERROR):
+            return MsDataset.load(
+                dataset_id,
+                subset_name=subset_name,
+                split=split,
+                version=revision,
+                download_mode=download_mode,
+                use_streaming=streaming,
+            )
 
     @classmethod
     def download_model(cls,

swift/trainers/mixin.py

@@ -1,6 +1,7 @@
 # Copyright (c) Alibaba, Inc. and its affiliates.
 # Part of the implementation is borrowed from huggingface/transformers.
 import inspect
+import logging
 import os
 import shutil
 import time
@@ -33,7 +34,7 @@
 from swift.llm import BatchSamplerShard, DataLoaderDispatcher, DataLoaderShard, Template
 from swift.plugin import MeanMetric, compute_acc, extra_tuners
 from swift.tuners import SwiftModel
-from swift.utils import get_logger, is_mp_ddp, seed_worker, use_torchacc
+from swift.utils import get_logger, is_mp_ddp, ms_logger_context, seed_worker, use_torchacc
 from swift.utils.torchacc_utils import ta_trim_graph
 from ..utils.torch_utils import get_device_count
 from .arguments import TrainingArguments
@@ -68,8 +69,7 @@ def __init__(self,
             logger.warning('Using IterableDataset, setting args.dataloader_num_workers to 1.')
 
         if args.check_model and hasattr(model, 'model_dir'):
-            from swift.utils.logger import ms_logger_ignore_error
-            with ms_logger_ignore_error():
+            with ms_logger_context(logging.CRITICAL):
                 check_local_model_is_latest(
                     model.model_dir, user_agent={
                         'invoked_by': 'local_trainer',

swift/utils/__init__.py

@@ -7,7 +7,7 @@
                            is_unsloth_available, is_vllm_ascend_available, is_vllm_available, is_wandb_available,
                            is_xtuner_available)
 from .io_utils import JsonlWriter, append_to_jsonl, download_ms_file, get_file_mm_type, read_from_jsonl, write_to_jsonl
-from .logger import get_logger
+from .logger import get_logger, ms_logger_context
 from .np_utils import get_seed, stat_array, transform_jsonl_to_df
 from .tb_utils import TB_COLOR, TB_COLOR_SMOOTH, plot_images, read_tensorboard_file, tensorboard_smoothing
 from .torch_utils import (Serializer, activate_parameters, check_shared_disk, find_all_linears, find_embedding,

swift/utils/logger.py

@@ -111,10 +111,10 @@ def get_logger(log_file: Optional[str] = None, log_level: Optional[int] = None,
 
 
 @contextmanager
-def ms_logger_ignore_error():
+def ms_logger_context(log_leval):
     ms_logger = get_ms_logger()
     origin_log_level = ms_logger.level
-    ms_logger.setLevel(logging.CRITICAL)
+    ms_logger.setLevel(log_leval)
     try:
         yield
     finally: