Only unset HOME when running cmake as root (#3507)

Summary:
This hack is required to work around pytorch/test-infra#5091, which runs some CI jobs as root, which buck2 doesn't like.

But we saw in #3502 that this can break things for some normal users.

Reduce the blast radius of this hack, only modifying HOME when actually running as root.

Mitigates #3502

Pull Request resolved: #3507

Test Plan:
`./install_requirements.sh` succeeded locally.

The build-wheels jobs for this PR do not break during the buck2 phase, and show that they're unsetting HOME.

https://github.com/pytorch/executorch/actions/runs/8946028682/job/24575999986?pr=3507#step:14:118
```
2024-05-03T23:32:08.9616508Z temporarily unsetting HOME while running as root
```

https://github.com/pytorch/executorch/actions/runs/8946028682/job/24575999986?pr=3507#step:14:465
```
2024-05-03T23:32:08.9914557Z restored HOME
```

Reviewed By: larryliu0820

Differential Revision: D56958571

Pulled By: dbort

fbshipit-source-id: c7c6abdd52361af8253ce068002e3c23dee16f6b

setup.py

@@ -45,6 +45,7 @@
 # other computer software, distribute, and sublicense such enhancements or
 # derivative works thereof, in binary and source code form.
 
+import contextlib
 import os
 import re
 import sys
@@ -382,6 +383,31 @@ def run(self):
             self.copy_file(src, dst, preserve_mode=False)
 
 
+class Buck2EnvironmentFixer(contextlib.AbstractContextManager):
+    """Removes HOME from the environment when running as root.
+
+    This script is sometimes run as root in docker containers. buck2 doesn't
+    allow running as root unless $HOME is owned by root or is not set.
+
+    TODO(pytorch/test-infra#5091): Remove this once the CI jobs stop running as
+    root.
+    """
+
+    def __init__(self):
+        self.saved_env = {}
+
+    def __enter__(self):
+        if os.geteuid() == 0 and "HOME" in os.environ:
+            log.info("temporarily unsetting HOME while running as root")
+            self.saved_env["HOME"] = os.environ.pop("HOME")
+        return self
+
+    def __exit__(self, *args, **kwargs):
+        if "HOME" in self.saved_env:
+            log.info("restored HOME")
+            os.environ["HOME"] = self.saved_env["HOME"]
+
+
 # TODO(dbort): For editable wheels, may need to update get_source_files(),
 # get_outputs(), and get_output_mapping() to satisfy
 # https://setuptools.pypa.io/en/latest/userguide/extension.html#setuptools.command.build.SubCommand.get_output_mapping
@@ -490,17 +516,13 @@ def run(self):
         if not self.dry_run:
             # Dry run should log the command but not actually run it.
             (Path(cmake_cache_dir) / "CMakeCache.txt").unlink(missing_ok=True)
-        try:
-            # This script is sometimes run as root in docker containers. buck2
-            # doesn't allow running as root unless $HOME is owned by root or
-            # does not exist. So temporarily undefine it while configuring
-            # cmake, which runs buck2 to get some source lists.
-            old_home = os.environ.pop("HOME", None)
+        with Buck2EnvironmentFixer():
+            # The context manager may patch the environment while running this
+            # cmake command, which happens to run buck2 to get some source
+            # lists.
+
             # Generate the build system files.
             self.spawn(["cmake", "-S", repo_root, "-B", cmake_cache_dir, *cmake_args])
-        finally:
-            if old_home is not None:
-                os.environ["HOME"] = old_home
 
         # Build the system.
         self.spawn(["cmake", "--build", cmake_cache_dir, *build_args])