make allocators and sanitizers work for processes created with multiprocessing's spawn method in dev mode

Summary:
#### 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.

For an earlier discussion, see [this post](https://fb.workplace.com/groups/fbpython/permalink/2897630276944987/).

#### 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
causes some linking issues, which were perhaps the reason `LD_PRELOAD` was
restored in the first place: in-place Python binaries have access to binaries
install on devservers that are not built with the target platform (e.g.
`/bin/sh` which is used by some Python standard libraries). These binaries does
not link properly with the system native dependencies.

#### References
An old RFC for this change: D16210828
The counterpart for opt mode: D16350169

fbshipit-source-id: ab83ba439a66f369794c2febea6993fdaa934e6f

Author: Yifu Wang

build.xml

@@ -1011,6 +1011,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.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.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,7 @@
 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_RESOURCE = "run_inplace_interpreter.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 +73,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 interpreterGenPath;
   @AddToRuleKey private final Tool python;
-  @AddToRuleKey private final Supplier<String> script;
+  @AddToRuleKey private final Supplier<String> binScript;
+  @AddToRuleKey private final Supplier<String> interpreterScript;
 
   PythonInPlaceBinary(
       BuildTarget buildTarget,
@@ -98,18 +105,27 @@ public class PythonInPlaceBinary extends PythonBinary implements HasRuntimeDeps
         legacyOutputPath);
     this.linkTree = linkTree;
     this.python = python;
-    this.script =
-        getScript(
+    this.interpreterGenPath =
+        getInterpreterGenPath(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(interpreterGenPath)),
+            packageStyle);
+    this.interpreterScript =
+        getInterpreterScript(
             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 +139,10 @@ private static String getRunInplaceResource() {
     return getNamedResource(RUN_INPLACE_RESOURCE);
   }
 
+  private static String getRunInplaceInterpreterResource() {
+    return getNamedResource(RUN_INPLACE_INTERPRETER_RESOURCE);
+  }
+
   private static String getRunInplaceLiteResource() {
     return getNamedResource(RUN_INPLACE_LITE_RESOURCE);
   }
@@ -136,29 +156,63 @@ private static String getNamedResource(String resourceName) {
     }
   }
 
-  private static Supplier<String> getScript(
+  private static RelPath getInterpreterGenPath(
+      BuildTarget target,
+      ProjectFilesystem filesystem,
+      String extension,
+      boolean legacyOutputPath) {
+    if (!legacyOutputPath) {
+      target = target.withFlavors();
+    }
+    return BuildTargetPaths.getGenPath(
+        filesystem.getBuckPaths(), target, "%s#interpreter" + extension);
+  }
+
+  private static Supplier<String> getBinScript(
+      PythonPlatform pythonPlatform,
+      String mainModule,
+      RelPath linkTreeRoot,
+      RelPath interpreterPath,
+      PackageStyle packageStyle) {
+    return () -> {
+      String linkTreeRootStr = Escaper.escapeAsPythonString(linkTreeRoot.toString());
+      String interpreterPathStr = Escaper.escapeAsPythonString(interpreterPath.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_REL_PATH", interpreterPathStr)
+          .render();
+    };
+  }
+
+  @Nullable
+  private static Supplier<String> getInterpreterScript(
       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 a par-style interpreter 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(), getRunInplaceInterpreterResource())
               .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 +241,24 @@ 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 (interpreterScript != null) {
+      RelPath interpreterPath =
+          context
+              .getSourcePathResolver()
+              .getCellUnsafeRelPath(
+                  ExplicitBuildTargetSourcePath.of(getBuildTarget(), interpreterGenPath));
+      stepsBuilder.add(
+          WriteFileIsolatedStep.of(interpreterScript, interpreterPath, /* 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,126 +58,10 @@ 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
-# within the Python interpreter. Enable decorating the main module after
-# `sys.path` has been setup by setting the PAR_MAIN_OVERRIDE environment
-# variable.
-decorate_main_module = os.environ.pop("PAR_MAIN_OVERRIDE", None)
-if decorate_main_module:
-    # Pass the original main module as environment variable for the process.
-    # Allowing the decorating module to pick it up.
-    os.environ["PAR_MAIN_ORIGINAL"] = main_module
-    main_module = decorate_main_module
-
-module_call = "runpy._run_module_as_main({main_module!r}, False)".format(
-    main_module=main_module
-)
-
-# Allow users to run the main module under pdb. Encode the call into the
-# startup script, because pdb does not support the -c argument we use to invoke
-# our startup wrapper.
-#
-# Note: use pop to avoid leaking the environment variable to the child process.
-if os.environ.pop("PYTHONDEBUGWITHPDB", None):
-    # Support passing initial commands to pdb. We cannot pass the -c argument
-    # to pdb. Instead, allow users to pass initial commands through the
-    # PYTHONPDBINITIALCOMMANDS env var, separated by the | character.
-    initial_commands = []
-    if "PYTHONPDBINITIALCOMMANDS" in os.environ:
-        # Note: use pop to avoid leaking the environment variable to the child
-        # process.
-        initial_commands_string = os.environ.pop("PYTHONPDBINITIALCOMMANDS", None)
-        initial_commands = initial_commands_string.split("|")
-
-    # Note: indentation of this block of code is important as it gets included
-    # in the bigger block below.
-    module_call = """
-    from pdb import Pdb
-    pdb = Pdb()
-    pdb.rcLines.extend({initial_commands!r})
-    pdb.runcall(runpy._run_module_as_main, {main_module!r}, False)
-""".format(
-        main_module=main_module,
-        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]
+# Run the main module with the interpreter wrapper, which loads native libaries
+# and adds neccessary prologue.
+interpreter_path = os.path.join(dirpath, <INTERPRETER_REL_PATH>)
+args = [interpreter_path, "<PYTHON_INTERPRETER_FLAGS>", "-m", main_module] + sys.argv[1:]
 
 # 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
@@ -213,7 +92,7 @@ if platform.system() == "Windows":
     # 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:])
+    p = subprocess.Popen(args)
 
     def handler(signum, frame):
         # If we're getting this, we need to forward signum to subprocesses
@@ -231,4 +110,4 @@ if platform.system() == "Windows":
     p.wait()
     sys.exit(p.returncode)
 else:
-    os.execv(sys.executable, args + sys.argv[1:])
+    os.execv(interpreter_path, args)