chore: lint

tests/e2e/test_fused_llama.py

@@ -14,7 +14,8 @@ from axolotl.common.cli import TrainerCliArgs
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
-from tests.utils import with_temp_dir
+
+from .utils import with_temp_dir  # pylint: disable=relative-beyond-top-level
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"

tests/e2e/test_lora_llama.py

@@ -12,7 +12,8 @@ from axolotl.common.cli import TrainerCliArgs
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
-from tests.utils import with_temp_dir
+
+from .utils import with_temp_dir  # pylint: disable=relative-beyond-top-level
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"

tests/e2e/test_mistral.py

@@ -14,7 +14,8 @@ from axolotl.common.cli import TrainerCliArgs
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
-from tests.utils import with_temp_dir
+
+from .utils import with_temp_dir  # pylint: disable=relative-beyond-top-level
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"

tests/e2e/test_mistral_samplepack.py

@@ -14,7 +14,8 @@ from axolotl.common.cli import TrainerCliArgs
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
-from tests.utils import with_temp_dir
+
+from .utils import with_temp_dir  # pylint: disable=relative-beyond-top-level
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"

tests/e2e/test_phi.py

@@ -12,7 +12,8 @@ from axolotl.common.cli import TrainerCliArgs
 from axolotl.train import train
 from axolotl.utils.config import normalize_config
 from axolotl.utils.dict import DictDefault
-from tests.utils import with_temp_dir
+
+from .utils import with_temp_dir  # pylint: disable=relative-beyond-top-level
 
 LOG = logging.getLogger("axolotl.tests.e2e")
 os.environ["WANDB_DISABLED"] = "true"

tests/e2e/utils.py

@@ -0,0 +1,22 @@
+"""
+helper utils for tests
+"""
+
+import shutil
+import tempfile
+from functools import wraps
+
+
+def with_temp_dir(test_func):
+    @wraps(test_func)
+    def wrapper(*args, **kwargs):
+        # Create a temporary directory
+        temp_dir = tempfile.mkdtemp()
+        try:
+            # Pass the temporary directory to the test function
+            test_func(temp_dir, *args, **kwargs)
+        finally:
+            # Clean up the directory after the test
+            shutil.rmtree(temp_dir)
+
+    return wrapper