diff --git a/tests/cli/test_neuron_cache_cli.py b/tests/cli/test_neuron_cache_cli.py index 45629abea..ba83d64ad 100644 --- a/tests/cli/test_neuron_cache_cli.py +++ b/tests/cli/test_neuron_cache_cli.py @@ -59,6 +59,8 @@ def _optimum_neuron_cache_create(self, cache_repo_id: Optional[str] = None, publ except RepositoryNotFoundError: pytest.fail("The repo was not created.") + finally: + delete_repo(repo_id) assert repo_id == load_custom_cache_repo_name_from_hf_home(), f"Saved local Neuron cache name should be equal to {repo_id}." @@ -84,27 +86,6 @@ def test_optimum_neuron_cache_add(self, hub_test): with TemporaryDirectory() as tmpdir: tmpdir = Path(tmpdir) - # TODO: activate those later. - # Without any sequence length, it should fail. - # command = ( - # "optimum-cli neuron cache add -m bert-base-uncased --task text-classification --train_batch_size 16 " - # "--precision bf16 --num_cores 2" - # ).split() - # p = subprocess.Popen(command, stderr=PIPE) - # _, stderr = p.communicate() - # stderr = stderr.decode("utf-8") - # self.assertIn("either sequence_length or encoder_sequence and decoder_sequence_length", stderr) - - # Without both encoder and decoder sequence lengths, it should fail. - # command = ( - # "optimum-cli neuron cache add -m t5-small --task translation --train_batch_size 16 --precision bf16 " - # "--num_cores 2 --encoder_sequence_length 512" - # ).split() - # p = subprocess.Popen(command, stderr=PIPE) - # _, stderr = p.communicate() - # stderr = stderr.decode("utf-8") - # self.assertIn("Both the encoder_sequence and decoder_sequence_length", stderr) - # Create dummy BERT model. bert_model_name = tmpdir / "bert_model" config = BertConfig() diff --git a/tests/conftest.py b/tests/conftest.py index a681ed087..bc05b7043 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -17,7 +17,7 @@ from pathlib import Path import pytest -from huggingface_hub import create_repo, delete_repo, get_token, login, logout +from huggingface_hub import create_repo, delete_repo, get_token, login, logout, HfApi from optimum.neuron.utils.cache_utils import ( delete_custom_cache_repo_name_from_hf_home, @@ -115,6 +115,12 @@ def _hub_test(create_local_cache: bool = False): yield custom_cache_repo_with_seed delete_repo(custom_cache_repo_with_seed, repo_type="model") + + model_repos = HfApi().list_models() + for repo in model_repos: + if repo.id.startswith("optimum-neuron-cache-for-testing-"): + delete_repo(repo.id) + if local_cache_path_with_seed.is_dir(): shutil.rmtree(local_cache_path_with_seed) if orig_token is not None: @@ -165,3 +171,5 @@ def pytest_fixture_setup(fixturedef, request): if getattr(fixturedef.func, "is_dist_fixture", False): dist_fixture_class = fixturedef.func() dist_fixture_class(request) + +