make allocators and sanitizers work for processes created with multiprocessing's spawn method in dev mode (facebook#2660)

Summary:
Pull Request resolved: facebook#2660

**The first attempt (D30802446) overlooked that fact that the interpreter
wrapper is not an executable (for `execv`) on Mac, and introduced some bugs due
to the refactoring. The attempt 2 addressed the issues, and isolated the effect
of the change to only processes created by multiprocess's spawn method on
Linux.**

#### Problem
Currently, the entrypoint for in-place Python binaries (i.e. built with dev
mode) executes the following steps to load system native dependencies (e.g.
sanitizers and allocators):
- Backup `LD_PRELOAD` set by the caller
- Append system native dependencies to `LD_PRELOAD`
- Inject a prologue in user code which restores `LD_PRELOAD` set by the caller
- `execv` Python interpreter

The steps work as intended for single process Python programs. However, when a
Python program spawns child processes, the child processes will not load native
dependencies, since they simply `execv`'s the vanilla Python interpreter. A few
examples why this is problematic:
- The ASAN runtime library is a system native dependency. Without loading it, a
  child process that loads user native dependencies compiled with ASAN will
  crash during static initialization because it can't find `_asan_init`.
- `jemalloc` is also a system native dependency.

Many if not most ML use cases "bans" dev mode because of these problems. It is
very unfortunate considering the developer efficiency dev mode provides. In
addition, a huge amount of unit tests have to run in a more expensive build
mode because of these problems.

#### Solution
Move the system native dependencies loading logic out of the Python binary
entrypoint into an interpreter wrapper, and set the interpreter as
`sys.executable` in the injected prologue:
- The Python binary entrypoint now uses the interpreter wrapper, which has the
  same command line interface as the Python interpreter, to run the main
  module.
- `multiprocessing`'s `spawn` method now uses the interpreter wrapper to create
  child processes, ensuring system native dependencies get loaded correctly.

#### Alternative Considered
One alternative considered is to simply not removing system native dependencies
from `LD_PRELOAD`, so they are present in the spawned processes. However, this
can cause some linking issues, which were perhaps the reason `LD_PRELOAD` was
restored in the first place.

Reviewed By: fried

fbshipit-source-id: e3f27dc66894a4a3fa13f2d725a4eb43ced41451

Author: Yifu Wang

build.xml

@@ -1016,6 +1016,7 @@
         <include name="com/facebook/buck/maven/build-file.st"/>
         <include name="com/facebook/buck/python/*.py"/>
         <include name="com/facebook/buck/python/run_inplace.py.in"/>
+        <include name="com/facebook/buck/python/run_inplace_interpreter_wrapper.py.in"/>
         <include name="com/facebook/buck/python/run_inplace_lite.py.in"/>
         <include name="com/facebook/buck/parser/function/BuckPyFunction.stg"/>
         <include name="com/facebook/buck/shell/sh_binary_template"/>

src/com/facebook/buck/features/python/BUCK

@@ -69,6 +69,7 @@ java_library_with_plugins(
         "__test_main__.py",
         "compile.py",
         "run_inplace.py.in",
+        "run_inplace_interpreter_wrapper.py.in",
         "run_inplace_lite.py.in",
     ],
     tests = [

src/com/facebook/buck/features/python/PythonInPlaceBinary.java

@@ -18,15 +18,18 @@
 
 import com.facebook.buck.core.build.buildable.context.BuildableContext;
 import com.facebook.buck.core.build.context.BuildContext;
+import com.facebook.buck.core.filesystems.AbsPath;
 import com.facebook.buck.core.filesystems.RelPath;
 import com.facebook.buck.core.model.BuildTarget;
 import com.facebook.buck.core.model.OutputLabel;
 import com.facebook.buck.core.model.TargetConfiguration;
+import com.facebook.buck.core.model.impl.BuildTargetPaths;
 import com.facebook.buck.core.rulekey.AddToRuleKey;
 import com.facebook.buck.core.rules.BuildRule;
 import com.facebook.buck.core.rules.BuildRuleResolver;
 import com.facebook.buck.core.rules.attr.HasRuntimeDeps;
 import com.facebook.buck.core.rules.impl.SymlinkTree;
+import com.facebook.buck.core.sourcepath.ExplicitBuildTargetSourcePath;
 import com.facebook.buck.core.toolchain.tool.Tool;
 import com.facebook.buck.core.toolchain.tool.impl.CommandTool;
 import com.facebook.buck.cxx.toolchain.CxxPlatform;
@@ -39,6 +42,7 @@
 import com.facebook.buck.step.Step;
 import com.facebook.buck.step.fs.MkdirStep;
 import com.facebook.buck.step.isolatedsteps.common.WriteFileIsolatedStep;
+import com.facebook.buck.test.selectors.Nullable;
 import com.facebook.buck.util.Escaper;
 import com.facebook.buck.util.stream.RichStream;
 import com.google.common.base.Joiner;
@@ -57,6 +61,8 @@
 public class PythonInPlaceBinary extends PythonBinary implements HasRuntimeDeps {
 
   private static final String RUN_INPLACE_RESOURCE = "run_inplace.py.in";
+  private static final String RUN_INPLACE_INTERPRETER_WRAPPER_RESOURCE =
+      "run_inplace_interpreter_wrapper.py.in";
   private static final String RUN_INPLACE_LITE_RESOURCE = "run_inplace_lite.py.in";
 
   // TODO(agallagher): Task #8098647: This rule has no steps, so it
@@ -68,8 +74,10 @@ public class PythonInPlaceBinary extends PythonBinary implements HasRuntimeDeps
   //
   // We should upate the Python test rule to account for this.
   private final SymlinkTree linkTree;
+  private final RelPath interpreterWrapperGenPath;
   @AddToRuleKey private final Tool python;
-  @AddToRuleKey private final Supplier<String> script;
+  @AddToRuleKey private final Supplier<String> binScript;
+  @AddToRuleKey private final Supplier<String> interpreterWrapperScript;
 
   PythonInPlaceBinary(
       BuildTarget buildTarget,
@@ -98,18 +106,28 @@ public class PythonInPlaceBinary extends PythonBinary implements HasRuntimeDeps
         legacyOutputPath);
     this.linkTree = linkTree;
     this.python = python;
-    this.script =
-        getScript(
+    this.interpreterWrapperGenPath =
+        getInterpreterWrapperGenPath(
+            buildTarget, projectFilesystem, pexExtension, legacyOutputPath);
+    AbsPath targetRoot =
+        projectFilesystem
+            .resolve(getBinPath(buildTarget, projectFilesystem, pexExtension, legacyOutputPath))
+            .getParent();
+    this.binScript =
+        getBinScript(
+            pythonPlatform,
+            mainModule,
+            targetRoot.relativize(linkTree.getRoot()),
+            targetRoot.relativize(projectFilesystem.resolve(interpreterWrapperGenPath)),
+            packageStyle);
+    this.interpreterWrapperScript =
+        getInterpreterWrapperScript(
             ruleResolver,
             buildTarget.getTargetConfiguration(),
             pythonPlatform,
             cxxPlatform,
-            mainModule,
             components,
-            projectFilesystem
-                .resolve(getBinPath(buildTarget, projectFilesystem, pexExtension, legacyOutputPath))
-                .getParent()
-                .relativize(linkTree.getRoot()),
+            targetRoot.relativize(linkTree.getRoot()),
             preloadLibraries,
             packageStyle);
   }
@@ -123,6 +141,10 @@ private static String getRunInplaceResource() {
     return getNamedResource(RUN_INPLACE_RESOURCE);
   }
 
+  private static String getRunInplaceInterpreterWrapperResource() {
+    return getNamedResource(RUN_INPLACE_INTERPRETER_WRAPPER_RESOURCE);
+  }
+
   private static String getRunInplaceLiteResource() {
     return getNamedResource(RUN_INPLACE_LITE_RESOURCE);
   }
@@ -136,29 +158,64 @@ private static String getNamedResource(String resourceName) {
     }
   }
 
-  private static Supplier<String> getScript(
+  private static RelPath getInterpreterWrapperGenPath(
+      BuildTarget target,
+      ProjectFilesystem filesystem,
+      String extension,
+      boolean legacyOutputPath) {
+    if (!legacyOutputPath) {
+      target = target.withFlavors();
+    }
+    return BuildTargetPaths.getGenPath(
+        filesystem.getBuckPaths(), target, "%s#interpreter_wrapper" + extension);
+  }
+
+  private static Supplier<String> getBinScript(
+      PythonPlatform pythonPlatform,
+      String mainModule,
+      RelPath linkTreeRoot,
+      RelPath interpreterWrapperPath,
+      PackageStyle packageStyle) {
+    return () -> {
+      String linkTreeRootStr = Escaper.escapeAsPythonString(linkTreeRoot.toString());
+      String interpreterWrapperPathStr =
+          Escaper.escapeAsPythonString(interpreterWrapperPath.toString());
+      return new ST(
+              new STGroup(),
+              packageStyle == PackageStyle.INPLACE
+                  ? getRunInplaceResource()
+                  : getRunInplaceLiteResource())
+          .add("PYTHON", pythonPlatform.getEnvironment().getPythonPath())
+          .add("PYTHON_INTERPRETER_FLAGS", pythonPlatform.getInplaceBinaryInterpreterFlags())
+          .add("MODULES_DIR", linkTreeRootStr)
+          .add("MAIN_MODULE", Escaper.escapeAsPythonString(mainModule))
+          .add("INTERPRETER_WRAPPER_REL_PATH", interpreterWrapperPathStr)
+          .render();
+    };
+  }
+
+  @Nullable
+  private static Supplier<String> getInterpreterWrapperScript(
       BuildRuleResolver resolver,
       TargetConfiguration targetConfiguration,
       PythonPlatform pythonPlatform,
       CxxPlatform cxxPlatform,
-      String mainModule,
       PythonPackageComponents components,
       RelPath relativeLinkTreeRoot,
       ImmutableSet<String> preloadLibraries,
       PackageStyle packageStyle) {
     String relativeLinkTreeRootStr = Escaper.escapeAsPythonString(relativeLinkTreeRoot.toString());
     Linker ld = cxxPlatform.getLd().resolve(resolver, targetConfiguration);
+    // Lite mode doesn't need an interpreter wrapper as there's no LD_PRELOADs involved.
+    if (packageStyle != PackageStyle.INPLACE) {
+      return null;
+    }
     return () -> {
       ST st =
-          new ST(
-                  new STGroup(),
-                  packageStyle == PackageStyle.INPLACE
-                      ? getRunInplaceResource()
-                      : getRunInplaceLiteResource())
+          new ST(new STGroup(), getRunInplaceInterpreterWrapperResource())
               .add("PYTHON", pythonPlatform.getEnvironment().getPythonPath())
-              .add("MAIN_MODULE", Escaper.escapeAsPythonString(mainModule))
-              .add("MODULES_DIR", relativeLinkTreeRootStr)
-              .add("PYTHON_INTERPRETER_FLAGS", pythonPlatform.getInplaceBinaryInterpreterFlags());
+              .add("PYTHON_INTERPRETER_FLAGS", pythonPlatform.getInplaceBinaryInterpreterFlags())
+              .add("MODULES_DIR", relativeLinkTreeRootStr);
 
       // Only add platform-specific values when the binary includes native libraries.
       if (components.getNativeLibraries().getComponents().isEmpty()) {
@@ -187,11 +244,26 @@ public ImmutableList<Step> getBuildSteps(
       BuildContext context, BuildableContext buildableContext) {
     RelPath binPath = context.getSourcePathResolver().getCellUnsafeRelPath(getSourcePathToOutput());
     buildableContext.recordArtifact(binPath.getPath());
-    return ImmutableList.of(
-        MkdirStep.of(
-            BuildCellRelativePath.fromCellRelativePath(
-                context.getBuildCellRootPath(), getProjectFilesystem(), binPath.getParent())),
-        WriteFileIsolatedStep.of(script, binPath, /* executable */ true));
+    ImmutableList.Builder<Step> stepsBuilder = new ImmutableList.Builder<Step>();
+    stepsBuilder
+        .add(
+            MkdirStep.of(
+                BuildCellRelativePath.fromCellRelativePath(
+                    context.getBuildCellRootPath(), getProjectFilesystem(), binPath.getParent())))
+        .add(WriteFileIsolatedStep.of(binScript, binPath, /* executable */ true));
+
+    if (interpreterWrapperScript != null) {
+      RelPath interpreterWrapperPath =
+          context
+              .getSourcePathResolver()
+              .getCellUnsafeRelPath(
+                  ExplicitBuildTargetSourcePath.of(getBuildTarget(), interpreterWrapperGenPath));
+      buildableContext.recordArtifact(interpreterWrapperPath.getPath());
+      stepsBuilder.add(
+          WriteFileIsolatedStep.of(
+              interpreterWrapperScript, interpreterWrapperPath, /* executable */ true));
+    }
+    return stepsBuilder.build();
   }
 
   @Override

src/com/facebook/buck/features/python/run_inplace.py.in

@@ -7,11 +7,6 @@ import subprocess
 import sys
 
 main_module = <MAIN_MODULE>
-modules_dir = <MODULES_DIR>
-native_libs_env_var = <NATIVE_LIBS_ENV_VAR>
-native_libs_dir = <NATIVE_LIBS_DIR>
-native_libs_preload_env_var = <NATIVE_LIBS_PRELOAD_ENV_VAR>
-native_libs_preload = <NATIVE_LIBS_PRELOAD>
 
 def try_resolve_possible_symlink(path):
     import ctypes
@@ -63,26 +58,6 @@ if platform.system() == "Windows":
     # does *not* dereference symlinks on windows until, like, 3.8 maybe.
     dirpath = os.path.dirname(try_resolve_possible_symlink(sys.argv[0]))
 
-env_vals_to_restore = {}
-# Update the environment variable for the dynamic loader to the native
-# libraries location.
-if native_libs_dir is not None:
-    old_native_libs_dir = os.environ.get(native_libs_env_var)
-    os.environ[native_libs_env_var] = os.path.join(dirpath, native_libs_dir)
-    env_vals_to_restore[native_libs_env_var] = old_native_libs_dir
-
-# Update the environment variable for the dynamic loader to find libraries
-# to preload.
-if native_libs_preload is not None:
-    old_native_libs_preload = os.environ.get(native_libs_preload_env_var)
-    env_vals_to_restore[native_libs_preload_env_var] = old_native_libs_preload
-
-    # On macos, preloaded libs are found via paths.
-    os.environ[native_libs_preload_env_var] = ":".join(
-        os.path.join(dirpath, native_libs_dir, l)
-        for l in native_libs_preload.split(":")
-    )
-
 # Allow users to decorate the main module. In normal Python invocations this
 # can be done by prefixing the arguments with `-m decoratingmodule`. It's not
 # that easy for par files. The startup script below sets up `sys.path` from
@@ -128,73 +103,18 @@ if os.environ.pop("PYTHONDEBUGWITHPDB", None):
         initial_commands=initial_commands,
     )
 
-# Note: this full block of code will be included as the argument to Python,
-# and will be the first thing that shows up in the process arguments as displayed
-# by programs like ps and top.
-#
-# We include arg0 at the start of this comment just to make it more visible what program
-# is being run in the ps and top output.
-STARTUP = """\
-# {arg0!r}
-# Wrap everything in a private function to prevent globals being captured by
-# the `runpy._run_module_as_main` below.
-def __run():
-    import sys
-
-    # We set the paths beforehand to have a minimal amount of imports before
-    # nuking PWD from sys.path. Otherwise, there can be problems if someone runs
-    # from a directory with a similarly named file, even if their code is properly
-    # namespaced. e.g. if one has foo/bar/contextlib.py and while in foo/bar runs
-    # `buck run foo/bar:bin`, runpy will fail as it tries to import
-    # foo/bar/contextlib.py. You're just out of luck if you have sys.py or os.py
-
-    # Set `argv[0]` to the executing script.
-    assert sys.argv[0] == '-c'
-    sys.argv[0] = {arg0!r}
-
-    # Replace the working directory with location of the modules directory.
-    assert sys.path[0] == ''
-    sys.path[0] = {pythonpath!r}
-
-    import os
-    import runpy
-
-    def setenv(var, val):
-        if val is None:
-            os.environ.pop(var, None)
-        else:
-            os.environ[var] = val
-
-    def restoreenv(d):
-        for k, v in d.items():
-            setenv(k, v)
-
-    restoreenv({env_vals!r})
-    {module_call}
-
-__run()
-""".format(
-    arg0=sys.argv[0],
-    pythonpath=os.path.join(dirpath, modules_dir),
-    env_vals=env_vals_to_restore,
-    main_module=main_module,
-    this_file=__file__,
-    module_call=module_call,
-)
-
-args = [sys.executable, "<PYTHON_INTERPRETER_FLAGS>", "-c", STARTUP]
-
+interpreter_opts = ["<PYTHON_INTERPRETER_FLAGS>"]
 # Default to 'd' warnings, but allow users to control this via PYTHONWARNINGS
 # The -E causes python to ignore all PYTHON* environment vars so we have to
 # pass this down using the command line.
 warnings = os.environ.get("PYTHONWARNINGS", "d").split(",")
 for item in reversed(warnings):
-    args.insert(1, "-W{0}".format(item.strip()))
+    interpreter_opts.insert(0, "-W{0}".format(item.strip()))
 
 # Allow users to disable byte code generation by setting the standard environment var.
 # Same as above, because of -E we have to pass this down using the command line.
 if "PYTHONDONTWRITEBYTECODE" in os.environ:
-    args.insert(1, "-B")
+    interpreter_opts.insert(0, "-B")
 
 # Python 3.7 allows benchmarking import time with this variable. Similar issues to
 # PYTHONDONTWRITEBYTECODE above. If using an earlier version of python... dont set this
@@ -205,30 +125,17 @@ if (
     and platform.python_implementation() == "CPython"
     and (sys.version_info[0], sys.version_info[1]) >= (3, 7)
 ):
-    args[1:1] = ["-X", "importtime"]
+    interpreter_opts[0:0] = ["-X", "importtime"]
 
-if platform.system() == "Windows":
-    # exec on Windows is not true exec - there is only 'spawn' ('CreateProcess').
-    # However, creating processes unnecessarily is painful, so we only do the spawn
-    # path if we have to, which is on Windows. That said, this complicates signal
-    # handling, so we need to set up some signal forwarding logic.
-
-    p = subprocess.Popen(args + sys.argv[1:])
-
-    def handler(signum, frame):
-        # If we're getting this, we need to forward signum to subprocesses
-        if signum == signal.SIGINT:
-            p.send_signal(signal.CTRL_C_EVENT)
-        elif signum == signal.SIGBREAK:
-            p.send_signal(signal.CTRL_BREAK_EVENT)
-        else:
-            # shouldn't happen, we should be killed instead
-            p.terminate()
-
-    signal.signal(signal.SIGINT, handler)
-    signal.signal(signal.SIGBREAK, handler)
-
-    p.wait()
-    sys.exit(p.returncode)
+interpreter_wrapper_path = os.path.join(dirpath, <INTERPRETER_WRAPPER_REL_PATH>)
+if sys.version_info >= (3, 0):
+    import importlib.machinery
+    loader = importlib.machinery.SourceFileLoader("interpreter_wrapper", interpreter_wrapper_path)
+    interpreter_wrapper = loader.load_module()
 else:
-    os.execv(sys.executable, args + sys.argv[1:])
+    # Buck is sunsetting Python2 support. However this is still needed for some
+    # unit tests.
+    import imp
+    interpreter_wrapper = imp.load_source("interpreter_wrapper", interpreter_wrapper_path)
+
+interpreter_wrapper.exec_interpreter(dirpath, interpreter_opts, module_call, sys.argv[1:])