tmp outputs dir at tests

KompleteAI · Oct 17, 2023 · b152c19 · b152c19
1 parent 77b8cfa
commit b152c19
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 4 deletions.
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -110,9 +110,9 @@ def llama_lm_collator(llama_tokenizer: PreTrainedTokenizer) -> LMCollator:
 
 
 @pytest.fixture(scope="session")
-def training_arguments() -> TrainingArguments:
+def training_arguments(path_to_outputs: str) -> TrainingArguments:
     arguments = TrainingArguments(
-        output_dir="./outputs/",
+        output_dir=path_to_outputs,
         per_device_train_batch_size=2,
         gradient_accumulation_steps=2,
         warmup_steps=50,
@@ -191,3 +191,9 @@ def path_to_fused_model_local_path(tmp_path_factory: TempPathFactory) -> str:
 def path_to_download_result(tmp_path_factory: TempPathFactory) -> str:
     path = tmp_path_factory.mktemp("tmp") / "data.jsonl"
     return os.path.abspath(path)
+
+
+@pytest.fixture(scope="session")
+def path_to_outputs(tmp_path_factory: TempPathFactory) -> str:
+    path = tmp_path_factory.mktemp("tmp") / "outputs/"
+    return os.path.abspath(path)
diff --git a/tests/unit/experiments/test_base.py b/tests/unit/experiments/test_base.py
@@ -19,7 +19,7 @@ def test_base_experiment_init(monkeypatch: MonkeyPatch, path_to_train_dummy_data
         Experiment(config=config)
 
 
-def test_base_experiment_train(monkeypatch: MonkeyPatch, path_to_train_prepared_dummy_data: str):
+def test_base_experiment_train(monkeypatch: MonkeyPatch, path_to_train_prepared_dummy_data: str, path_to_outputs: str):
     os.environ["TOKENIZERS_PARALLELISM"] = "false"
     config = HuggingFaceConfig(
         push_to_hub=False,
@@ -29,6 +29,7 @@ def test_base_experiment_train(monkeypatch: MonkeyPatch, path_to_train_prepared_
         save_total_limit=0,
         max_steps=2,
         tokenizer_name_or_path=LLAMA_TOKENIZER_DIR,
+        output_dir=path_to_outputs,
     )
 
     with patch_from_pretrained_auto_causal_lm(monkeypatch=monkeypatch):

diff --git a/tests/unit/run/test_train.py b/tests/unit/run/test_train.py
@@ -6,7 +6,7 @@
 from tests.helpers.patches import patch_from_pretrained_auto_causal_lm, patch_trainer_train
 
 
-def test_train(monkeypatch: MonkeyPatch, path_to_train_prepared_dummy_data: str):
+def test_train(monkeypatch: MonkeyPatch, path_to_train_prepared_dummy_data: str, path_to_outputs: str):
     config = HuggingFaceConfig(
         push_to_hub=False,
         deepspeed_stage=0,
@@ -15,6 +15,7 @@ def test_train(monkeypatch: MonkeyPatch, path_to_train_prepared_dummy_data: str)
         save_total_limit=0,
         max_steps=2,
         tokenizer_name_or_path=LLAMA_TOKENIZER_DIR,
+        output_dir=path_to_outputs,
     )
     with patch_from_pretrained_auto_causal_lm(monkeypatch=monkeypatch):
         with patch_trainer_train(monkeypatch=monkeypatch):

diff --git a/tests/unit/trainers/test_registry.py b/tests/unit/trainers/test_registry.py
@@ -15,6 +15,7 @@ def test_get_lm_trainer(
     training_arguments: TrainingArguments,
     llama_lm_collator: LMCollator,
     soda_dataset: SodaDataset,
+    path_to_outputs: str,
 ):
     trainer_cls = trainers_registry.get(key=enums.Trainers.lm)
     trainer = trainer_cls(