reset environment variables after each test

Signed-off-by: Joshua Rosenkranz <jmrosenk@us.ibm.com>

Author: JRosenkranz

tests/models/test_decoders.py

@@ -9,8 +9,7 @@
 from aiu_fms_testing_utils.utils.aiu_setup import dprint
 import os
 
-if "HF_HOME" not in os.environ:
-    os.environ["HF_HOME"] = "/tmp/models/hf_cache"
+ORIGINAL_HF_HOME = os.environ.get("HF_HOME", None)
 
 # Add models to test here
 LLAMA_3p1_8B_INSTRUCT = "meta-llama/Llama-3.1-8B-Instruct"
@@ -72,6 +71,11 @@ def reset_compiler():
     yield # run the test
     torch.compiler.reset()
     torch._dynamo.reset()
+    os.environ.pop('COMPILATION_MODE', None)
+    if ORIGINAL_HF_HOME is None:
+        os.environ.pop('HF_HOME', None)
+    else:
+        os.environ['HF_HOME'] = ORIGINAL_HF_HOME
 
 def __prepare_inputs(batch_size, seq_length, tokenizer, seed=0):
     prompts_and_sizes = sample_sharegpt_requests(SHARE_GPT_DATASET_PATH, batch_size, tokenizer, int(seq_length / 2), seq_length, seed)
@@ -116,6 +120,10 @@ def __load_validation_info(model_path, batch_size, seq_length, max_new_tokens, t
 @pytest.mark.parametrize("model_path,batch_size,seq_length,max_new_tokens", common_shapes)
 def test_common_shapes(model_path, batch_size, seq_length, max_new_tokens):
     os.environ["COMPILATION_MODE"] = "offline_decoder"
+    
+    if "HF_HOME" not in os.environ:
+        os.environ["HF_HOME"] = "/tmp/models/hf_cache"
+
     dprint(f"testing model={model_path}, batch_size={batch_size}, seq_length={seq_length}, max_new_tokens={max_new_tokens}, micro_model={USE_MICRO_MODELS}")
 
     if USE_MICRO_MODELS:

tests/models/test_encoders.py

@@ -9,8 +9,7 @@
 from aiu_fms_testing_utils.utils.aiu_setup import dprint
 import os
 
-if "HF_HOME" not in os.environ:
-    os.environ["HF_HOME"] = "/tmp/models/hf_cache"
+ORIGINAL_HF_HOME = os.environ.get("HF_HOME", None)
 
 # Add models to test here
 ROBERTA_SQUAD_V2 = "deepset/roberta-base-squad2"
@@ -49,13 +48,21 @@ def reset_compiler():
     yield # run the test
     torch.compiler.reset()
     torch._dynamo.reset()
+    os.environ.pop('COMPILATION_MODE', None)
+    if ORIGINAL_HF_HOME is None:
+        os.environ.pop('HF_HOME', None)
+    else:
+        os.environ['HF_HOME'] = ORIGINAL_HF_HOME
 
 encoder_paths = ["deepset/roberta-base-squad2"]
 common_encoder_shapes = list(itertools.product(encoder_paths, common_batch_sizes, common_seq_lengths))
 
 @pytest.mark.parametrize("model_path,batch_size,seq_length", common_encoder_shapes)
 def test_common_shapes(model_path, batch_size, seq_length):
     os.environ["COMPILATION_MODE"] = "offline"
+
+    if "HF_HOME" not in os.environ:
+        os.environ["HF_HOME"] = "/tmp/models/hf_cache"
 
     dprint(f"testing model={model_path}, batch_size={batch_size}, seq_length={seq_length}")