Repository: huggingface/accelerate

Issue:
patch_environment does not revert previously existing keys
### System Info

```Shell
does not matter here
```


### Information

- [ ] The official example scripts
- [ ] My own modified scripts

### Tasks

- [ ] One of the scripts in the examples/ folder of Accelerate or an officially supported `no_trainer` script in the `examples` folder of the `transformers` repo (such as `run_no_trainer_glue.py`)
- [ ] My own task or dataset (give details below)

### Reproduction

```python
os.environ["FOO"] = "BAR"
with patch_environment(FOO="BAZ"):
     assert os.environ["FOO"] == "BAZ"
assert os.environ["F00"] == "BAR"  # KeyError
```

### Expected behavior

After leaving the context, the initial environment variables should be restored.

I noticed this while going through the code. If this is an issue, LMK and I'll provide a fix + tests. If, in the context that this function is used, it's not a problem, feel free to close the issue.


Discussion:
cc @muellerzr 
This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/accelerate/blob/main/CONTRIBUTING.md) are likely to be ignored.

Resolving patch:
diff --git a/src/accelerate/utils/other.py b/src/accelerate/utils/other.py
--- a/src/accelerate/utils/other.py
+++ b/src/accelerate/utils/other.py
@@ -172,14 +172,22 @@ def patch_environment(**kwargs):
     >>> print(os.environ["FOO"])  # raises KeyError
     ```
     """
+    existing_vars = {}
     for key, value in kwargs.items():
-        os.environ[key.upper()] = str(value)
+        key = key.upper()
+        if key in os.environ:
+            existing_vars[key] = os.environ[key]
+        os.environ[key] = str(value)
 
     yield
 
     for key in kwargs:
-        if key.upper() in os.environ:
-            del os.environ[key.upper()]
+        key = key.upper()
+        if key in existing_vars:
+            # restore previous value
+            os.environ[key] = existing_vars[key]
+        else:
+            os.environ.pop(key, None)
 
 
 def get_pretty_name(obj):


Test patch:
diff --git a/tests/test_utils.py b/tests/test_utils.py
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -103,6 +103,25 @@ def test_patch_environment(self):
         self.assertNotIn("AA", os.environ)
         self.assertNotIn("BB", os.environ)
 
+    def test_patch_environment_key_exists(self):
+        # check that patch_environment correctly restores pre-existing env vars
+        with patch_environment(aa=1, BB=2):
+            self.assertEqual(os.environ.get("AA"), "1")
+            self.assertEqual(os.environ.get("BB"), "2")
+
+            with patch_environment(Aa=10, bb="20", cC=30):
+                self.assertEqual(os.environ.get("AA"), "10")
+                self.assertEqual(os.environ.get("BB"), "20")
+                self.assertEqual(os.environ.get("CC"), "30")
+
+            self.assertEqual(os.environ.get("AA"), "1")
+            self.assertEqual(os.environ.get("BB"), "2")
+            self.assertNotIn("CC", os.environ)
+
+        self.assertNotIn("AA", os.environ)
+        self.assertNotIn("BB", os.environ)
+        self.assertNotIn("CC", os.environ)
+
     def test_can_undo_convert_outputs(self):
         model = RegressionModel()
         model._original_forward = model.forward