Rank 0-only logging (#2608)

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

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"