Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

make sure to cleanup tmp output_dir for e2e tests #831

Merged
merged 4 commits into from
Nov 6, 2023
Merged

make sure to cleanup tmp output_dir for e2e tests #831

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
bb719b4
make sure to cleanup tmp output_dir for e2e tests
winglian Nov 5, 2023
65ff70b
chore: lint
winglian Nov 6, 2023
53f7aa7
missing dunder-init
winglian Nov 6, 2023
c17d719
use temp_dir kwarg instead
winglian Nov 6, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Empty file added tests/e2e/__init__.py
View file
Open in desktop
Empty file.
11 changes: 6 additions & 5 deletions tests/e2e/test_fused_llama.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@

import logging
import os
import tempfile
import unittest
from pathlib import Path

Expand All @@ -16,6 +15,8 @@
from axolotl.utils.config import normalize_config
from axolotl.utils.dict import DictDefault

from .utils import with_temp_dir

LOG = logging.getLogger("axolotl.tests.e2e")
os.environ["WANDB_DISABLED"] = "true"

Expand All @@ -25,9 +26,9 @@ class TestFusedLlama(unittest.TestCase):
Test case for Llama models using Fused layers
"""

def test_fft_packing(self):
@with_temp_dir
def test_fft_packing(self, temp_dir):
# pylint: disable=duplicate-code
output_dir = tempfile.mkdtemp()
cfg = DictDefault(
{
"base_model": "JackFram/llama-68m",
Expand All @@ -51,7 +52,7 @@ def test_fft_packing(self):
"num_epochs": 2,
"micro_batch_size": 2,
"gradient_accumulation_steps": 1,
"output_dir": output_dir,
"output_dir": temp_dir,
"learning_rate": 0.00001,
"optimizer": "adamw_torch",
"lr_scheduler": "cosine",
Expand All @@ -69,4 +70,4 @@ def test_fft_packing(self):
dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)

train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
assert (Path(output_dir) / "pytorch_model.bin").exists()
assert (Path(temp_dir) / "pytorch_model.bin").exists()
27 changes: 14 additions & 13 deletions tests/e2e/test_lora_llama.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@

import logging
import os
import tempfile
import unittest
from pathlib import Path

Expand All @@ -14,6 +13,8 @@
from axolotl.utils.config import normalize_config
from axolotl.utils.dict import DictDefault

from .utils import with_temp_dir

LOG = logging.getLogger("axolotl.tests.e2e")
os.environ["WANDB_DISABLED"] = "true"

Expand All @@ -23,9 +24,9 @@ class TestLoraLlama(unittest.TestCase):
Test case for Llama models using LoRA
"""

def test_lora(self):
@with_temp_dir
def test_lora(self, temp_dir):
# pylint: disable=duplicate-code
output_dir = tempfile.mkdtemp()
cfg = DictDefault(
{
"base_model": "JackFram/llama-68m",
Expand All @@ -52,7 +53,7 @@ def test_lora(self):
"num_epochs": 2,
"micro_batch_size": 8,
"gradient_accumulation_steps": 1,
"output_dir": output_dir,
"output_dir": temp_dir,
"learning_rate": 0.00001,
"optimizer": "adamw_torch",
"lr_scheduler": "cosine",
Expand All @@ -63,11 +64,11 @@ def test_lora(self):
dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)

train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
assert (Path(output_dir) / "adapter_model.bin").exists()
assert (Path(temp_dir) / "adapter_model.bin").exists()

def test_lora_packing(self):
@with_temp_dir
def test_lora_packing(self, temp_dir):
# pylint: disable=duplicate-code
output_dir = tempfile.mkdtemp()
cfg = DictDefault(
{
"base_model": "JackFram/llama-68m",
Expand Down Expand Up @@ -96,7 +97,7 @@ def test_lora_packing(self):
"num_epochs": 2,
"micro_batch_size": 8,
"gradient_accumulation_steps": 1,
"output_dir": output_dir,
"output_dir": temp_dir,
"learning_rate": 0.00001,
"optimizer": "adamw_torch",
"lr_scheduler": "cosine",
Expand All @@ -107,11 +108,11 @@ def test_lora_packing(self):
dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)

train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
assert (Path(output_dir) / "adapter_model.bin").exists()
assert (Path(temp_dir) / "adapter_model.bin").exists()

def test_lora_gptq(self):
@with_temp_dir
def test_lora_gptq(self, temp_dir):
# pylint: disable=duplicate-code
output_dir = tempfile.mkdtemp()
cfg = DictDefault(
{
"base_model": "TheBlokeAI/jackfram_llama-68m-GPTQ",
Expand Down Expand Up @@ -144,7 +145,7 @@ def test_lora_gptq(self):
"save_steps": 0.5,
"micro_batch_size": 8,
"gradient_accumulation_steps": 1,
"output_dir": output_dir,
"output_dir": temp_dir,
"learning_rate": 0.00001,
"optimizer": "adamw_torch",
"lr_scheduler": "cosine",
Expand All @@ -155,4 +156,4 @@ def test_lora_gptq(self):
dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)

train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
assert (Path(output_dir) / "adapter_model.bin").exists()
assert (Path(temp_dir) / "adapter_model.bin").exists()
19 changes: 10 additions & 9 deletions tests/e2e/test_mistral.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@

import logging
import os
import tempfile
import unittest
from pathlib import Path

Expand All @@ -16,6 +15,8 @@
from axolotl.utils.config import normalize_config
from axolotl.utils.dict import DictDefault

from .utils import with_temp_dir

LOG = logging.getLogger("axolotl.tests.e2e")
os.environ["WANDB_DISABLED"] = "true"

Expand All @@ -25,9 +26,9 @@ class TestMistral(unittest.TestCase):
Test case for Llama models using LoRA
"""

def test_lora(self):
@with_temp_dir
def test_lora(self, temp_dir):
# pylint: disable=duplicate-code
output_dir = tempfile.mkdtemp()
cfg = DictDefault(
{
"base_model": "openaccess-ai-collective/tiny-mistral",
Expand All @@ -54,7 +55,7 @@ def test_lora(self):
"num_epochs": 2,
"micro_batch_size": 2,
"gradient_accumulation_steps": 1,
"output_dir": output_dir,
"output_dir": temp_dir,
"learning_rate": 0.00001,
"optimizer": "adamw_torch",
"lr_scheduler": "cosine",
Expand All @@ -68,11 +69,11 @@ def test_lora(self):
dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)

train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
assert (Path(output_dir) / "adapter_model.bin").exists()
assert (Path(temp_dir) / "adapter_model.bin").exists()

def test_ft(self):
@with_temp_dir
def test_ft(self, temp_dir):
# pylint: disable=duplicate-code
output_dir = tempfile.mkdtemp()
cfg = DictDefault(
{
"base_model": "openaccess-ai-collective/tiny-mistral",
Expand All @@ -93,7 +94,7 @@ def test_ft(self):
"num_epochs": 2,
"micro_batch_size": 2,
"gradient_accumulation_steps": 1,
"output_dir": output_dir,
"output_dir": temp_dir,
"learning_rate": 0.00001,
"optimizer": "adamw_torch",
"lr_scheduler": "cosine",
Expand All @@ -111,4 +112,4 @@ def test_ft(self):
dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)

train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
assert (Path(output_dir) / "pytorch_model.bin").exists()
assert (Path(temp_dir) / "pytorch_model.bin").exists()
19 changes: 10 additions & 9 deletions tests/e2e/test_mistral_samplepack.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,6 @@

import logging
import os
import tempfile
import unittest
from pathlib import Path

Expand All @@ -16,6 +15,8 @@
from axolotl.utils.config import normalize_config
from axolotl.utils.dict import DictDefault

from .utils import with_temp_dir

LOG = logging.getLogger("axolotl.tests.e2e")
os.environ["WANDB_DISABLED"] = "true"

Expand All @@ -25,9 +26,9 @@ class TestMistral(unittest.TestCase):
Test case for Llama models using LoRA
"""

def test_lora_packing(self):
@with_temp_dir
def test_lora_packing(self, temp_dir):
# pylint: disable=duplicate-code
output_dir = tempfile.mkdtemp()
cfg = DictDefault(
{
"base_model": "openaccess-ai-collective/tiny-mistral",
Expand Down Expand Up @@ -55,7 +56,7 @@ def test_lora_packing(self):
"num_epochs": 2,
"micro_batch_size": 2,
"gradient_accumulation_steps": 1,
"output_dir": output_dir,
"output_dir": temp_dir,
"learning_rate": 0.00001,
"optimizer": "adamw_torch",
"lr_scheduler": "cosine",
Expand All @@ -69,11 +70,11 @@ def test_lora_packing(self):
dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)

train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
assert (Path(output_dir) / "adapter_model.bin").exists()
assert (Path(temp_dir) / "adapter_model.bin").exists()

def test_ft_packing(self):
@with_temp_dir
def test_ft_packing(self, temp_dir):
# pylint: disable=duplicate-code
output_dir = tempfile.mkdtemp()
cfg = DictDefault(
{
"base_model": "openaccess-ai-collective/tiny-mistral",
Expand All @@ -95,7 +96,7 @@ def test_ft_packing(self):
"num_epochs": 2,
"micro_batch_size": 2,
"gradient_accumulation_steps": 1,
"output_dir": output_dir,
"output_dir": temp_dir,
"learning_rate": 0.00001,
"optimizer": "adamw_torch",
"lr_scheduler": "cosine",
Expand All @@ -113,4 +114,4 @@ def test_ft_packing(self):
dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)

train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
assert (Path(output_dir) / "pytorch_model.bin").exists()
assert (Path(temp_dir) / "pytorch_model.bin").exists()
16 changes: 11 additions & 5 deletions tests/e2e/test_phi.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,15 +4,17 @@

import logging
import os
import tempfile
import unittest
from pathlib import Path

from axolotl.cli import load_datasets
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 .utils import with_temp_dir

LOG = logging.getLogger("axolotl.tests.e2e")
os.environ["WANDB_DISABLED"] = "true"

Expand All @@ -22,7 +24,8 @@ class TestPhi(unittest.TestCase):
Test case for Llama models using LoRA
"""

def test_ft(self):
@with_temp_dir
def test_ft(self, temp_dir):
# pylint: disable=duplicate-code
cfg = DictDefault(
{
Expand Down Expand Up @@ -52,7 +55,7 @@ def test_ft(self):
"num_epochs": 1,
"micro_batch_size": 1,
"gradient_accumulation_steps": 1,
"output_dir": tempfile.mkdtemp(),
"output_dir": temp_dir,
"learning_rate": 0.00001,
"optimizer": "adamw_bnb_8bit",
"lr_scheduler": "cosine",
Expand All @@ -64,8 +67,10 @@ def test_ft(self):
dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)

train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
assert (Path(temp_dir) / "pytorch_model.bin").exists()

def test_ft_packed(self):
@with_temp_dir
def test_ft_packed(self, temp_dir):
# pylint: disable=duplicate-code
cfg = DictDefault(
{
Expand Down Expand Up @@ -95,7 +100,7 @@ def test_ft_packed(self):
"num_epochs": 1,
"micro_batch_size": 1,
"gradient_accumulation_steps": 1,
"output_dir": tempfile.mkdtemp(),
"output_dir": temp_dir,
"learning_rate": 0.00001,
"optimizer": "adamw_bnb_8bit",
"lr_scheduler": "cosine",
Expand All @@ -107,3 +112,4 @@ def test_ft_packed(self):
dataset_meta = load_datasets(cfg=cfg, cli_args=cli_args)

train(cfg=cfg, cli_args=cli_args, dataset_meta=dataset_meta)
assert (Path(temp_dir) / "pytorch_model.bin").exists()
22 changes: 22 additions & 0 deletions tests/e2e/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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(*args, temp_dir=temp_dir, **kwargs)
finally:
# Clean up the directory after the test
shutil.rmtree(temp_dir)

return wrapper