fork_exec.py 1.3 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748
  1. ###############################################################################
  2. # Launch a subprocess using forkexec and make sure only the needed fd are
  3. # shared in the two process.
  4. #
  5. # author: Thomas Moreau and Olivier Grisel
  6. #
  7. import os
  8. import sys
  9. if sys.platform == "darwin" and sys.version_info < (3, 3):
  10. FileNotFoundError = OSError
  11. def close_fds(keep_fds): # pragma: no cover
  12. """Close all the file descriptors except those in keep_fds."""
  13. # Make sure to keep stdout and stderr open for logging purpose
  14. keep_fds = set(keep_fds).union([1, 2])
  15. # We try to retrieve all the open fds
  16. try:
  17. open_fds = set(int(fd) for fd in os.listdir('/proc/self/fd'))
  18. except FileNotFoundError:
  19. import resource
  20. max_nfds = resource.getrlimit(resource.RLIMIT_NOFILE)[0]
  21. open_fds = set(fd for fd in range(3, max_nfds))
  22. open_fds.add(0)
  23. for i in open_fds - keep_fds:
  24. try:
  25. os.close(i)
  26. except OSError:
  27. pass
  28. def fork_exec(cmd, keep_fds, env=None):
  29. # copy the environment variables to set in the child process
  30. env = {} if env is None else env
  31. child_env = os.environ.copy()
  32. child_env.update(env)
  33. pid = os.fork()
  34. if pid == 0: # pragma: no cover
  35. close_fds(keep_fds)
  36. os.execve(sys.executable, cmd, child_env)
  37. else:
  38. return pid