PoC: use isolated contextes for environment variables

easybuild/framework/easyconfig/format/one.py

@@ -39,6 +39,7 @@
 import tempfile
 
 from easybuild.base import fancylogger
+from easybuild.os_hook import OSProxy
 from easybuild.framework.easyconfig.format.format import DEPENDENCY_PARAMETERS, EXCLUDED_KEYS_REPLACE_TEMPLATES
 from easybuild.framework.easyconfig.format.format import FORMAT_DEFAULT_VERSION, GROUPED_PARAMS, LAST_PARAMS
 from easybuild.framework.easyconfig.format.format import SANITY_CHECK_PATHS_DIRS, SANITY_CHECK_PATHS_FILES
@@ -140,7 +141,7 @@ def get_config_dict(self):
         cfg_copy = {}
         for key in cfg:
             # skip special variables like __builtins__, and imported modules (like 'os')
-            if key != '__builtins__' and "'module'" not in str(type(cfg[key])):
+            if key != '__builtins__' and "'module'" not in str(type(cfg[key])) and not isinstance(cfg[key], OSProxy):
                 try:
                     cfg_copy[key] = copy.deepcopy(cfg[key])
                 except Exception as err:

easybuild/main.py

@@ -39,6 +39,7 @@
 * Maxime Boissonneault (Compute Canada)
 * Bart Oldeman (McGill University, Calcul Quebec, Digital Research Alliance of Canada)
 """
+# flake8: noqa: E402
 import copy
 import os
 import stat
@@ -47,9 +48,13 @@
 import traceback
 from datetime import datetime
 
+# if os.environ.get('EB_ISOLATED_CONTEXTS', '').lower() in ('1', 'true', 'yes'):
 # IMPORTANT this has to be the first easybuild import as it customises the logging
 #  expect missing log output when this not the case!
-from easybuild.tools.build_log import EasyBuildError, print_error_and_exit, print_msg, print_warning, stop_logging
+from easybuild import os_hook  # Imported to inject hook that replaces system os with our wrapped version
+os_hook.install_os_hook()
+
+from easybuild.tools.build_log import EasyBuildError, print_error_and_exit, print_msg, print_warning, stop_logging   # noqa: E402
 from easybuild.tools.build_log import EasyBuildExit
 
 from easybuild.framework.easyblock import build_and_install_one, inject_checksums, inject_checksums_to_json

easybuild/os_hook.py

@@ -0,0 +1,136 @@
+import importlib
+import importlib._bootstrap_external
+import importlib.abc
+import importlib.util
+import sys
+import types
+
+
+class ProxyLoader(importlib.abc.Loader):
+    """Loader to create our proxy instead of the real module."""
+    proxy_cls = None  # To be defined in subclasses
+
+    def create_module(self, spec):
+        # Import real module safely
+        sys.meta_path = [f for f in sys.meta_path if not isinstance(f, HookFinder)]
+        real_module = importlib.import_module(spec.name)
+        sys.meta_path.insert(0, HookFinder())
+
+        # Return proxy instead of real module
+        return self.proxy_cls(real_module)
+
+    def exec_module(self, module):
+        """Needs to be defined, can be used to alter the module after creation if needed."""
+
+
+class ModuleProxy(types.ModuleType):
+    """Generic proxy module to intercept attribute access."""
+    overrides = None
+    module_name = None
+
+    def __init__(self, real):
+        super().__init__(self.module_name)
+        self._real = real
+        # self._not_found = set()
+
+    def __getattr__(self, name):
+        # Intercept specific attributes
+        # if name in self.overrides:
+        #     # print(f"Intercepted access to {self.module_name}.{name}, returning override value.")
+        #     pass
+        # else:
+        #     self._not_found.add(name)
+        #     print("NOTFOUND", self.module_name, sorted(self._not_found))
+        return self.overrides.get(name, getattr(self._real, name))
+
+    def __dir__(self):
+        return dir(self._real)
+
+    @classmethod
+    def register_override(cls, name, value):
+        cls.overrides[name] = value
+
+    @classmethod
+    def loader(cls):
+        class Loader(ProxyLoader):
+            proxy_cls = cls
+        return Loader()
+
+
+class SubprocessProxy(ModuleProxy):
+    """Proxy module to intercept subprocess attribute access."""
+    overrides = {}
+    module_name = "subprocess"
+
+
+class OSProxy(ModuleProxy):
+    """Proxy module to intercept os attribute access."""
+    overrides = {}
+    module_name = "os"
+
+
+class PosixProxy(ModuleProxy):
+    """Proxy module to intercept posix attribute access."""
+    overrides = {}
+    module_name = "posix"
+
+
+class PosixpathProxy(ModuleProxy):
+    """Proxy module to intercept posixpath attribute access."""
+    overrides = {}
+    module_name = "posixpath"
+
+
+proxy_map = {
+    "os": OSProxy,
+    "subprocess": SubprocessProxy,
+    "posix": PosixProxy,
+    "posixpath": PosixpathProxy,
+    # "builtins": BuiltinProxy,
+}
+
+
+class HookFinder(importlib.abc.MetaPathFinder):
+    """Meta path finder to intercept imports of 'os' and return our proxy."""
+    def find_spec(self, fullname, path, target=None):
+        if fullname in proxy_map:
+            return importlib.util.spec_from_loader(fullname, proxy_map[fullname].loader())
+        return None
+
+
+def install_os_hook():
+    """Install the os hooking mechanism to intercept imports of 'os' and return our proxy."""
+    if not any(isinstance(f, HookFinder) for f in sys.meta_path):
+        sys.meta_path.insert(0, HookFinder())
+
+    # If already imported, replace in place
+    for name, proxy in proxy_map.items():
+        if name in sys.modules and not isinstance(sys.modules[name], proxy):
+            real_module = sys.modules[name]
+            sys.modules[name] = proxy(real_module)
+
+    # https://stackoverflow.com/questions/79420610/undertanding-python-import-process-importing-custom-os-module
+    # Reload system modules that might have already imported os with a different name, at python initialization
+    # - tempfile imports os as _os and this is happening before we have a chance to install our hook.
+    # - os.path is a separate module (eg posixpath) that imports os into itself and needs to be reloaded to import
+    #   our hook for eg `os.path.expanduser` to work with `os.environ['HOME'] = '...'`
+    # - shutil is used in CUDA sanity check with `shutil.which` to find `cuobjdum`
+    system_modules = [
+        "os", "sys", "tempfile", "posixpath", "shutil", "importlib", "io", "glob",
+    ]
+    for name in system_modules:
+        if name in sys.modules:
+            # print(f"Reloading system module {name} to ensure it imports our os hook.")
+            importlib.reload(sys.modules[name])
+
+    # Needed to override how import paths are resolved in case '' is in sys.path indicating the CWD.
+    # Cannot be reloaded without breaking stuff
+    importlib._bootstrap_external._os = sys.modules["posix"]
+
+    # To ensure we have the contextes module loaded to set all the function overrides
+    importlib.import_module("easybuild.tools.contextes")
+
+    sys.modules["posixpath"].os = sys.modules["os"]
+    sys.modules["posixpath"]._real.os = sys.modules["os"]
+    sys.modules["io"].os = sys.modules["os"]
+    # sys.modules["io"]._real.os = sys.modules["os"]