Added monkey patching

Signed-off-by: romit <[email protected]>

Author: romitjain

tests/build/test_launch_script.py

@@ -81,25 +81,19 @@
 }
 
 
-def setup_env(tempdir):
-    os.environ["TRAINING_SCRIPT"] = SCRIPT
-    os.environ["PYTHONPATH"] = "./:$PYTHONPATH"
-    os.environ["TERMINATION_LOG_FILE"] = tempdir + "/termination-log"
+def setup_env(monkeypatch, tempdir):
+    monkeypatch.setenv("TRAINING_SCRIPT", SCRIPT)
+    monkeypatch.setenv("PYTHONPATH", "./:$PYTHONPATH")
+    monkeypatch.setenv("TERMINATION_LOG_FILE", os.path.join(tempdir, "termination-log"))
 
 
-def cleanup_env():
-    os.environ.pop("TRAINING_SCRIPT", None)
-    os.environ.pop("PYTHONPATH", None)
-    os.environ.pop("TERMINATION_LOG_FILE", None)
-
-
-def test_successful_ft():
+def test_successful_ft(monkeypatch):
     """Check if we can bootstrap and fine tune causallm models"""
     with tempfile.TemporaryDirectory() as tempdir:
-        setup_env(tempdir)
+        setup_env(monkeypatch, tempdir)
         TRAIN_KWARGS = {**BASE_KWARGS, **{"output_dir": tempdir}}
         serialized_args = serialize_args(TRAIN_KWARGS)
-        os.environ["SFT_TRAINER_CONFIG_JSON_ENV_VAR"] = serialized_args
+        monkeypatch.setenv("SFT_TRAINER_CONFIG_JSON_ENV_VAR", serialized_args)
 
         assert main() == 0
         _validate_termination_files_when_tuning_succeeds(tempdir)
@@ -108,43 +102,43 @@ def test_successful_ft():
 
 
 @pytest.mark.skipif(True, reason="This test is deprecated so always skipped")
-def test_successful_pt():
+def test_successful_pt(monkeypatch):
     """Check if we can bootstrap and peft tune causallm models"""
     with tempfile.TemporaryDirectory() as tempdir:
-        setup_env(tempdir)
+        setup_env(monkeypatch, tempdir)
         TRAIN_KWARGS = {**BASE_PEFT_KWARGS, **{"output_dir": tempdir}}
         serialized_args = serialize_args(TRAIN_KWARGS)
-        os.environ["SFT_TRAINER_CONFIG_JSON_ENV_VAR"] = serialized_args
+        monkeypatch.setenv("SFT_TRAINER_CONFIG_JSON_ENV_VAR", serialized_args)
 
         assert main() == 0
         _validate_termination_files_when_tuning_succeeds(tempdir)
         checkpoint = os.path.join(tempdir, get_highest_checkpoint(tempdir))
         _validate_training_output(checkpoint, "pt")
 
 
-def test_successful_lora():
+def test_successful_lora(monkeypatch):
     """Check if we can bootstrap and LoRA tune causallm models"""
     with tempfile.TemporaryDirectory() as tempdir:
-        setup_env(tempdir)
+        setup_env(monkeypatch, tempdir)
         TRAIN_KWARGS = {**BASE_LORA_KWARGS, **{"output_dir": tempdir}}
         serialized_args = serialize_args(TRAIN_KWARGS)
-        os.environ["SFT_TRAINER_CONFIG_JSON_ENV_VAR"] = serialized_args
+        monkeypatch.setenv("SFT_TRAINER_CONFIG_JSON_ENV_VAR", serialized_args)
 
         assert main() == 0
         _validate_termination_files_when_tuning_succeeds(tempdir)
         checkpoint = os.path.join(tempdir, get_highest_checkpoint(tempdir))
         _validate_training_output(checkpoint, "lora")
 
 
-def test_lora_save_model_dir_separate_dirs():
+def test_lora_save_model_dir_separate_dirs(monkeypatch):
     """Run LoRA tuning with separate save_model_dir and output_dir.
     Verify model saved to save_model_dir and checkpoints saved to
     output_dir.
     """
     with tempfile.TemporaryDirectory() as tempdir:
         output_dir = os.path.join(tempdir, "output_dir")
         save_model_dir = os.path.join(tempdir, "save_model_dir")
-        setup_env(tempdir)
+        setup_env(monkeypatch, tempdir)
         TRAIN_KWARGS = {
             **BASE_LORA_KWARGS,
             **{
@@ -154,7 +148,7 @@ def test_lora_save_model_dir_separate_dirs():
             },
         }
         serialized_args = serialize_args(TRAIN_KWARGS)
-        os.environ["SFT_TRAINER_CONFIG_JSON_ENV_VAR"] = serialized_args
+        monkeypatch.setenv("SFT_TRAINER_CONFIG_JSON_ENV_VAR", serialized_args)
 
         assert main() == 0
         _validate_termination_files_when_tuning_succeeds(output_dir)
@@ -165,12 +159,12 @@ def test_lora_save_model_dir_separate_dirs():
         assert len(checkpoints) == 1
 
 
-def test_lora_save_model_dir_same_dir_as_output_dir():
+def test_lora_save_model_dir_same_dir_as_output_dir(monkeypatch):
     """Run LoRA tuning with same save_model_dir and output_dir.
     Verify checkpoints, logs, and model saved to path.
     """
     with tempfile.TemporaryDirectory() as tempdir:
-        setup_env(tempdir)
+        setup_env(monkeypatch, tempdir)
         TRAIN_KWARGS = {
             **BASE_LORA_KWARGS,
             **{
@@ -180,7 +174,7 @@ def test_lora_save_model_dir_same_dir_as_output_dir():
             },
         }
         serialized_args = serialize_args(TRAIN_KWARGS)
-        os.environ["SFT_TRAINER_CONFIG_JSON_ENV_VAR"] = serialized_args
+        monkeypatch.setenv("SFT_TRAINER_CONFIG_JSON_ENV_VAR", serialized_args)
 
         assert main() == 0
         # check logs, checkpoint dir, and model exists in path
@@ -195,19 +189,21 @@ def test_lora_save_model_dir_same_dir_as_output_dir():
         assert len(checkpoints) == TRAIN_KWARGS["num_train_epochs"]
 
 
-def test_lora_save_model_dir_same_dir_as_output_dir_save_strategy_no():
+def test_lora_save_model_dir_same_dir_as_output_dir_save_strategy_no(
+    monkeypatch,
+):
     """Run LoRA tuning with same save_model_dir and output_dir and
     save_strategy=no. Verify no checkpoints created, only
     logs and final model.
     """
     with tempfile.TemporaryDirectory() as tempdir:
-        setup_env(tempdir)
+        setup_env(monkeypatch, tempdir)
         TRAIN_KWARGS = {
             **BASE_LORA_KWARGS,
             **{"output_dir": tempdir, "save_model_dir": tempdir, "save_strategy": "no"},
         }
         serialized_args = serialize_args(TRAIN_KWARGS)
-        os.environ["SFT_TRAINER_CONFIG_JSON_ENV_VAR"] = serialized_args
+        monkeypatch.setenv("SFT_TRAINER_CONFIG_JSON_ENV_VAR", serialized_args)
 
         assert main() == 0
         # check that model and logs exists in output_dir
@@ -219,9 +215,9 @@ def test_lora_save_model_dir_same_dir_as_output_dir_save_strategy_no():
         assert len(checkpoints) == 0
 
 
-def test_lora_with_lora_post_process_for_vllm_set_to_true():
+def test_lora_with_lora_post_process_for_vllm_set_to_true(monkeypatch):
     with tempfile.TemporaryDirectory() as tempdir:
-        setup_env(tempdir)
+        setup_env(monkeypatch, tempdir)
         TRAIN_KWARGS = {
             **BASE_LORA_KWARGS,
             **{
@@ -231,7 +227,7 @@ def test_lora_with_lora_post_process_for_vllm_set_to_true():
             },
         }
         serialized_args = serialize_args(TRAIN_KWARGS)
-        os.environ["SFT_TRAINER_CONFIG_JSON_ENV_VAR"] = serialized_args
+        monkeypatch.setenv("SFT_TRAINER_CONFIG_JSON_ENV_VAR", serialized_args)
 
         assert main() == 0
         # check that model and logs exists in output_dir
@@ -255,9 +251,9 @@ def test_lora_with_lora_post_process_for_vllm_set_to_true():
     not _is_package_available("HFResourceScanner"),
     reason="Only runs if HFResourceScanner is installed",
 )
-def test_launch_with_HFResourceScanner_enabled():
+def test_launch_with_HFResourceScanner_enabled(monkeypatch):
     with tempfile.TemporaryDirectory() as tempdir:
-        setup_env(tempdir)
+        setup_env(monkeypatch, tempdir)
         scanner_outfile = os.path.join(tempdir, TrackerConfigs.scanner_output_filename)
         TRAIN_KWARGS = {
             **BASE_LORA_KWARGS,
@@ -271,7 +267,7 @@ def test_launch_with_HFResourceScanner_enabled():
             },
         }
         serialized_args = serialize_args(TRAIN_KWARGS)
-        os.environ["SFT_TRAINER_CONFIG_JSON_ENV_VAR"] = serialized_args
+        monkeypatch.setenv("SFT_TRAINER_CONFIG_JSON_ENV_VAR", serialized_args)
 
         assert main() == 0
         assert os.path.exists(scanner_outfile) is True
@@ -281,14 +277,14 @@ def test_launch_with_HFResourceScanner_enabled():
         assert scanner_res["mem_data"] is not None
 
 
-def test_bad_script_path():
+def test_bad_script_path(monkeypatch):
     """Check for appropriate error for an invalid training script location"""
     with tempfile.TemporaryDirectory() as tempdir:
-        setup_env(tempdir)
+        setup_env(monkeypatch, tempdir)
         TRAIN_KWARGS = {**BASE_LORA_KWARGS, **{"output_dir": tempdir}}
         serialized_args = serialize_args(TRAIN_KWARGS)
-        os.environ["SFT_TRAINER_CONFIG_JSON_ENV_VAR"] = serialized_args
-        os.environ["TRAINING_SCRIPT"] = "/not/here"
+        monkeypatch.setenv("SFT_TRAINER_CONFIG_JSON_ENV_VAR", serialized_args)
+        monkeypatch.setenv("TRAINING_SCRIPT", "/not/here")
 
         with pytest.raises(SystemExit) as pytest_wrapped_e:
             main()
@@ -297,61 +293,61 @@ def test_bad_script_path():
         assert os.stat(tempdir + "/termination-log").st_size > 0
 
 
-def test_blank_env_var():
+def test_blank_env_var(monkeypatch):
     with tempfile.TemporaryDirectory() as tempdir:
-        setup_env(tempdir)
-        os.environ["SFT_TRAINER_CONFIG_JSON_ENV_VAR"] = ""
+        setup_env(monkeypatch, tempdir)
+        monkeypatch.setenv("SFT_TRAINER_CONFIG_JSON_ENV_VAR", "")
         with pytest.raises(SystemExit) as pytest_wrapped_e:
             main()
         assert pytest_wrapped_e.type == SystemExit
         assert pytest_wrapped_e.value.code == USER_ERROR_EXIT_CODE
         assert os.stat(tempdir + "/termination-log").st_size > 0
 
 
-def test_faulty_file_path():
+def test_faulty_file_path(monkeypatch):
     with tempfile.TemporaryDirectory() as tempdir:
-        setup_env(tempdir)
+        setup_env(monkeypatch, tempdir)
         faulty_path = os.path.join(tempdir, "non_existent_file.pkl")
         TRAIN_KWARGS = {
             **BASE_LORA_KWARGS,
             **{"training_data_path": faulty_path, "output_dir": tempdir},
         }
         serialized_args = serialize_args(TRAIN_KWARGS)
-        os.environ["SFT_TRAINER_CONFIG_JSON_ENV_VAR"] = serialized_args
+        monkeypatch.setenv("SFT_TRAINER_CONFIG_JSON_ENV_VAR", serialized_args)
         with pytest.raises(SystemExit) as pytest_wrapped_e:
             main()
         assert pytest_wrapped_e.type == SystemExit
         assert pytest_wrapped_e.value.code == USER_ERROR_EXIT_CODE
         assert os.stat(tempdir + "/termination-log").st_size > 0
 
 
-def test_bad_base_model_path():
+def test_bad_base_model_path(monkeypatch):
     with tempfile.TemporaryDirectory() as tempdir:
-        setup_env(tempdir)
+        setup_env(monkeypatch, tempdir)
         TRAIN_KWARGS = {
             **BASE_LORA_KWARGS,
             **{"model_name_or_path": "/wrong/path"},
             "output_dir": tempdir,
         }
         serialized_args = serialize_args(TRAIN_KWARGS)
-        os.environ["SFT_TRAINER_CONFIG_JSON_ENV_VAR"] = serialized_args
+        monkeypatch.setenv("SFT_TRAINER_CONFIG_JSON_ENV_VAR", serialized_args)
         with pytest.raises(SystemExit) as pytest_wrapped_e:
             main()
         assert pytest_wrapped_e.type == SystemExit
         assert pytest_wrapped_e.value.code == USER_ERROR_EXIT_CODE
         assert os.stat(tempdir + "/termination-log").st_size > 0
 
 
-def test_config_parsing_error():
+def test_config_parsing_error(monkeypatch):
     with tempfile.TemporaryDirectory() as tempdir:
-        setup_env(tempdir)
+        setup_env(monkeypatch, tempdir)
         TRAIN_KWARGS = {
             **BASE_LORA_KWARGS,
             **{"num_train_epochs": "five"},
             "output_dir": tempdir,
         }  # Intentional type error
         serialized_args = serialize_args(TRAIN_KWARGS)
-        os.environ["SFT_TRAINER_CONFIG_JSON_ENV_VAR"] = serialized_args
+        monkeypatch.setenv("SFT_TRAINER_CONFIG_JSON_ENV_VAR", serialized_args)
         with pytest.raises(SystemExit) as pytest_wrapped_e:
             main()
         assert pytest_wrapped_e.type == SystemExit
@@ -376,9 +372,3 @@ def _validate_training_output(base_dir, tuning_technique):
     else:
         assert os.path.exists(base_dir + "/adapter_config.json") is True
         assert os.path.exists(base_dir + "/adapter_model.safetensors") is True
-
-
-def test_cleanup():
-    # This runs to unset env variables that could disrupt other tests
-    cleanup_env()
-    assert True

tests/build/test_utils.py

@@ -55,9 +55,17 @@ def test_process_accelerate_launch_args(job_config):
 
 
 @patch("torch.cuda.device_count", return_value=1)
-def test_accelerate_launch_args_user_set_num_processes_ignored(job_config):
+def test_accelerate_launch_args_user_set_num_processes_ignored(
+    _mock_cuda_count, job_config, monkeypatch
+):
     job_config_copy = copy.deepcopy(job_config)
     job_config_copy["accelerate_launch_args"]["num_processes"] = "3"
+    if "CUDA_VISIBLE_DEVICES" in os.environ:
+        monkeypatch.setenv(
+            "CUDA_VISIBLE_DEVICES", os.environ["CUDA_VISIBLE_DEVICES"]
+        )
+    else:
+        monkeypatch.delenv("CUDA_VISIBLE_DEVICES", raising=False)
     args = process_accelerate_launch_args(job_config_copy)
     # determine number of processes by number of GPUs available
     assert args.num_processes == 1

tests/test_sft_trainer.py

@@ -801,7 +801,7 @@ def test_run_causallm_alora_and_inference(request, target_modules, expected):
         assert "Simply put, the theory of relativity states that \n" in output_inference
 
 
-def test_successful_lora_target_modules_default_from_main():
+def test_successful_lora_target_modules_default_from_main(monkeypatch):
     """Check that if target_modules is not set, or set to None via JSON, the
     default value by model type will be using in LoRA tuning.
     The correct default target modules will be used for model type llama
@@ -818,7 +818,7 @@ def test_successful_lora_target_modules_default_from_main():
             **{"peft_method": "lora", "output_dir": tempdir},
         }
         serialized_args = serialize_args(TRAIN_KWARGS)
-        os.environ["SFT_TRAINER_CONFIG_JSON_ENV_VAR"] = serialized_args
+        monkeypatch.setenv("SFT_TRAINER_CONFIG_JSON_ENV_VAR", serialized_args)
 
         sft_trainer.main()
 

tests/utils/test_config_utils.py

@@ -208,28 +208,30 @@ def test_update_config_can_handle_multiple_config_updates():
     assert config[1].r == 98
 
 
-def test_get_json_config_can_load_from_path():
+def test_get_json_config_can_load_from_path(monkeypatch):
     """Test that the function get_json_config can read
     the json path from env var SFT_TRAINER_CONFIG_JSON_PATH
     """
-    if "SFT_TRAINER_CONFIG_JSON_ENV_VAR" in os.environ:
-        del os.environ["SFT_TRAINER_CONFIG_JSON_ENV_VAR"]
-    os.environ["SFT_TRAINER_CONFIG_JSON_PATH"] = HAPPY_PATH_DUMMY_CONFIG_PATH
+    monkeypatch.delenv("SFT_TRAINER_CONFIG_JSON_ENV_VAR", raising=False)
+    monkeypatch.setenv(
+        "SFT_TRAINER_CONFIG_JSON_PATH", HAPPY_PATH_DUMMY_CONFIG_PATH
+    )
 
     job_config = config_utils.get_json_config()
     assert job_config is not None
     assert job_config["model_name_or_path"] == "bigscience/bloom-560m"
 
 
-def test_get_json_config_can_load_from_envvar():
+def test_get_json_config_can_load_from_envvar(monkeypatch):
     """Test that the function get_json_config can read
     the json path from env var SFT_TRAINER_CONFIG_JSON_ENV_VAR
     """
     config_json = {"model_name_or_path": "foobar"}
     message_bytes = pickle.dumps(config_json)
     base64_bytes = base64.b64encode(message_bytes)
     encoded_json = base64_bytes.decode("ascii")
-    os.environ["SFT_TRAINER_CONFIG_JSON_ENV_VAR"] = encoded_json
+    monkeypatch.delenv("SFT_TRAINER_CONFIG_JSON_PATH", raising=False)
+    monkeypatch.setenv("SFT_TRAINER_CONFIG_JSON_ENV_VAR", encoded_json)
 
     job_config = config_utils.get_json_config()
     assert job_config is not None