Rank 0-only logging (#2608)

Co-authored-by: Wing Lian <wing@axolotl.ai>

tests/e2e/multigpu/solo/test_flex.py

@@ -2,7 +2,6 @@
 E2E tests for multigpu lora tinyllama
 """
 
-import logging
 import os
 from pathlib import Path
 
@@ -14,10 +13,11 @@ from transformers.testing_utils import get_torch_dist_unique_port
 from transformers.utils import is_torch_bf16_gpu_available
 
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from tests.e2e.utils import check_tensorboard, require_torch_2_6_0
 
-LOG = logging.getLogger("axolotl.tests.e2e.multigpu")
+LOG = get_logger("axolotl.tests.e2e.multigpu")
 os.environ["WANDB_DISABLED"] = "true"
 
 AXOLOTL_ROOT = Path(__file__).parent.parent.parent.parent

tests/e2e/multigpu/test_eval.py

@@ -2,7 +2,6 @@
 E2E tests for multigpu eval
 """
 
-import logging
 import os
 from pathlib import Path
 
@@ -11,10 +10,11 @@ from accelerate.test_utils import execute_subprocess_async
 from transformers.testing_utils import get_torch_dist_unique_port
 
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from ..utils import check_tensorboard
 
-LOG = logging.getLogger("axolotl.tests.e2e.multigpu")
+LOG = get_logger("axolotl.tests.e2e.multigpu")
 os.environ["WANDB_DISABLED"] = "true"
 
 AXOLOTL_ROOT = Path(__file__).parent.parent.parent.parent

tests/e2e/multigpu/test_gemma3.py

@@ -2,7 +2,6 @@
 E2E tests for multigpu lora tinyllama
 """
 
-import logging
 import os
 from pathlib import Path
 
@@ -13,10 +12,11 @@ from huggingface_hub import snapshot_download
 from transformers.testing_utils import get_torch_dist_unique_port
 
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from tests.e2e.utils import check_tensorboard
 
-LOG = logging.getLogger("axolotl.tests.e2e.multigpu")
+LOG = get_logger("axolotl.tests.e2e.multigpu")
 os.environ["WANDB_DISABLED"] = "true"
 
 AXOLOTL_ROOT = Path(__file__).parent.parent.parent.parent

tests/e2e/multigpu/test_llama.py

@@ -2,7 +2,6 @@
 E2E tests for multigpu lora tinyllama
 """
 
-import logging
 import os
 from pathlib import Path
 
@@ -15,10 +14,11 @@ from packaging import version
 from transformers.testing_utils import get_torch_dist_unique_port
 
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from tests.e2e.utils import check_tensorboard, require_torch_2_6_0
 
-LOG = logging.getLogger("axolotl.tests.e2e.multigpu")
+LOG = get_logger("axolotl.tests.e2e.multigpu")
 os.environ["WANDB_DISABLED"] = "true"
 
 AXOLOTL_ROOT = Path(__file__).parent.parent.parent.parent

tests/e2e/multigpu/test_qwen2.py

@@ -2,7 +2,6 @@
 E2E tests for multigpu qwen2
 """
 
-import logging
 import os
 from pathlib import Path
 
@@ -12,8 +11,9 @@ from accelerate.test_utils import execute_subprocess_async
 from transformers.testing_utils import get_torch_dist_unique_port
 
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
-LOG = logging.getLogger("axolotl.tests.e2e.multigpu")
+LOG = get_logger("axolotl.tests.e2e.multigpu")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/multigpu/test_ray.py

@@ -2,7 +2,6 @@
 E2E tests for multigpu post-training use Ray Train
 """
 
-import logging
 import os
 from pathlib import Path
 
@@ -11,10 +10,11 @@ import yaml
 from accelerate.test_utils import execute_subprocess_async
 
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from tests.e2e.utils import check_tensorboard, require_torch_lt_2_6_0
 
-LOG = logging.getLogger(__name__)
+LOG = get_logger(__name__)
 os.environ["WANDB_DISABLED"] = "true"
 
 AXOLOTL_ROOT = Path(__file__).parent.parent.parent.parent

tests/e2e/patched/test_4d_multipack_llama.py

@@ -2,7 +2,6 @@
 E2E tests for multipack fft llama using 4d attention masks
 """
 
-import logging
 import os
 import unittest
 
@@ -11,10 +10,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from ..utils import check_model_output_exists, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/patched/test_fa_xentropy.py

@@ -2,7 +2,6 @@
 E2E tests for lora llama
 """
 
-import logging
 import os
 
 import pytest
@@ -13,10 +12,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from ..utils import check_model_output_exists, check_tensorboard
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/patched/test_falcon_samplepack.py

@@ -2,7 +2,6 @@
 E2E tests for falcon
 """
 
-import logging
 import os
 import unittest
 
@@ -13,10 +12,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from ..utils import check_model_output_exists, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/patched/test_fused_llama.py

@@ -2,7 +2,6 @@
 E2E tests for lora llama
 """
 
-import logging
 import os
 import unittest
 
@@ -14,10 +13,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from ..utils import check_model_output_exists, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/patched/test_llama_s2_attention.py

@@ -2,7 +2,6 @@
 E2E tests for llama w/ S2 attn
 """
 
-import logging
 import os
 import unittest
 
@@ -13,10 +12,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from ..utils import check_model_output_exists, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/patched/test_lora_llama_multipack.py

@@ -2,7 +2,6 @@
 E2E tests for lora llama
 """
 
-import logging
 import os
 import unittest
 
@@ -14,10 +13,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from ..utils import check_model_output_exists, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/patched/test_mistral_samplepack.py

@@ -2,7 +2,6 @@
 E2E tests for lora llama
 """
 
-import logging
 import os
 import unittest
 
@@ -11,10 +10,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from ..utils import check_model_output_exists, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/patched/test_mixtral_samplepack.py

@@ -2,7 +2,6 @@
 E2E tests for mixtral
 """
 
-import logging
 import os
 import unittest
 
@@ -11,10 +10,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from ..utils import check_model_output_exists, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/patched/test_phi_multipack.py

@@ -2,7 +2,6 @@
 E2E tests for lora llama
 """
 
-import logging
 import os
 import unittest
 
@@ -11,10 +10,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from ..utils import check_model_output_exists, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/patched/test_resume.py

@@ -2,7 +2,6 @@
 E2E tests for resuming training
 """
 
-import logging
 import os
 import re
 import subprocess
@@ -14,10 +13,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from ..utils import check_model_output_exists, most_recent_subdir, require_torch_2_6_0
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/patched/test_unsloth_qlora.py

@@ -2,7 +2,6 @@
 e2e tests for unsloth qlora
 """
 
-import logging
 import os
 
 import pytest
@@ -12,10 +11,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from ..utils import check_model_output_exists, check_tensorboard
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/solo/test_flex.py

@@ -2,7 +2,6 @@
 E2E tests for packed training w/ flex attention
 """
 
-import logging
 import os
 import unittest
 
@@ -13,10 +12,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from ..utils import check_tensorboard, require_torch_2_6_0, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/solo/test_relora_llama.py

@@ -2,7 +2,6 @@
 E2E tests for relora llama
 """
 
-import logging
 import os
 import unittest
 from pathlib import Path
@@ -12,10 +11,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from ..utils import check_model_output_exists, check_tensorboard, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/test_deepseekv3.py

@@ -2,7 +2,6 @@
 E2E tests for deepseekv3
 """
 
-import logging
 import os
 from pathlib import Path
 
@@ -13,10 +12,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from tests.hf_offline_utils import enable_hf_offline
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/test_dpo.py

@@ -2,7 +2,6 @@
 E2E tests for lora llama
 """
 
-import logging
 import os
 import unittest
 from pathlib import Path
@@ -14,10 +13,11 @@ from axolotl.common.datasets import load_preference_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from .utils import check_model_output_exists, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/test_embeddings_lr.py

@@ -2,7 +2,6 @@
 E2E tests for llama pretrain
 """
 
-import logging
 import os
 import unittest
 
@@ -11,10 +10,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from .utils import check_model_output_exists, check_tensorboard, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/test_falcon.py

@@ -2,7 +2,6 @@
 E2E tests for falcon
 """
 
-import logging
 import os
 import unittest
 
@@ -13,10 +12,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from .utils import check_model_output_exists, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/test_gemma2.py

@@ -2,7 +2,6 @@
 E2E tests for gemma2
 """
 
-import logging
 import os
 from pathlib import Path
 
@@ -13,8 +12,9 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/test_gemma3_text.py

@@ -2,7 +2,6 @@
 E2E tests for gemma3_text
 """
 
-import logging
 import os
 from pathlib import Path
 
@@ -13,8 +12,9 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/test_llama.py

@@ -2,7 +2,6 @@
 E2E tests for llama
 """
 
-import logging
 import os
 
 from axolotl.cli.args import TrainerCliArgs
@@ -10,10 +9,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from tests.e2e.utils import check_model_output_exists
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/test_llama_pretrain.py

@@ -2,7 +2,6 @@
 E2E tests for llama pretrain
 """
 
-import logging
 import os
 
 import pytest
@@ -12,10 +11,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from .utils import check_model_output_exists, check_tensorboard
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/test_llama_vision.py

@@ -2,7 +2,6 @@
 E2E tests for lora llama
 """
 
-import logging
 import os
 import unittest
 
@@ -11,10 +10,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from .utils import check_model_output_exists, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/test_lora_llama.py

@@ -2,7 +2,6 @@
 E2E tests for lora llama
 """
 
-import logging
 import os
 import unittest
 
@@ -11,10 +10,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from .utils import check_model_output_exists, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/test_mamba.py

@@ -2,7 +2,6 @@
 E2E tests for lora llama
 """
 
-import logging
 import os
 import unittest
 
@@ -13,10 +12,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from .utils import check_model_output_exists, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/test_mistral.py

@@ -2,7 +2,6 @@
 E2E tests for lora llama
 """
 
-import logging
 import os
 import unittest
 
@@ -13,10 +12,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from .utils import check_model_output_exists, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/test_mixtral.py

@@ -2,7 +2,6 @@
 E2E tests for mixtral
 """
 
-import logging
 import os
 import unittest
 
@@ -14,10 +13,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from .utils import check_model_output_exists, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/test_optimizers.py

@@ -2,7 +2,6 @@
 E2E tests for custom optimizers using Llama
 """
 
-import logging
 import os
 import unittest
 
@@ -11,10 +10,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from .utils import check_model_output_exists, require_torch_2_5_1, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/test_packing_loss.py

@@ -2,7 +2,6 @@
 E2E tests for packed training
 """
 
-import logging
 import os
 import unittest
 
@@ -13,10 +12,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from .utils import check_tensorboard, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/test_phi.py

@@ -2,7 +2,6 @@
 E2E tests for lora llama
 """
 
-import logging
 import os
 import unittest
 
@@ -11,10 +10,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from .utils import check_model_output_exists, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/test_process_reward_model_smollm2.py

@@ -2,7 +2,6 @@
 E2E tests for process reward model w/ lora llama
 """
 
-import logging
 import os
 import unittest
 
@@ -11,10 +10,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from .utils import check_model_output_exists, check_tensorboard, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/test_qwen.py

@@ -2,7 +2,6 @@
 E2E tests for qwen
 """
 
-import logging
 import os
 from pathlib import Path
 
@@ -12,8 +11,9 @@ from accelerate.test_utils import execute_subprocess_async
 from transformers.testing_utils import get_torch_dist_unique_port
 
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
-LOG = logging.getLogger("axolotl.tests.qwen")
+LOG = get_logger("axolotl.tests.qwen")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/test_reward_model_smollm2.py

@@ -2,7 +2,6 @@
 E2E tests for reward model lora llama
 """
 
-import logging
 import os
 import unittest
 
@@ -11,10 +10,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from .utils import check_model_output_exists, check_tensorboard, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"
 
 

tests/e2e/test_schedulers.py

@@ -2,7 +2,6 @@
 E2E tests for custom schedulers using Llama
 """
 
-import logging
 import os
 import unittest
 
@@ -11,10 +10,11 @@ from axolotl.common.datasets import load_datasets
 from axolotl.train import train
 from axolotl.utils.config import normalize_config, validate_config
 from axolotl.utils.dict import DictDefault
+from axolotl.utils.logging import get_logger
 
 from .utils import check_model_output_exists, with_temp_dir
 
-LOG = logging.getLogger("axolotl.tests.e2e")
+LOG = get_logger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"