1 #

     2 # Module for starting a process object using os.fork() or CreateProcess()

     3 #

     4 # multiprocessing/forking.py

     5 #

     6 # Copyright (c) 2006-2008, R Oudkerk --- see COPYING.txt

     7 #

     8 

     9 import os

    10 import sys

    11 import signal

    12 

    13 from multiprocessing import util, process

    14 

    15 __all__ = ['Popen', 'assert_spawning', 'exit', 'duplicate', 'close', 'ForkingPickler']

    16 

    17 #

    18 # Check that the current thread is spawning a child process

    19 #

    20 

    21 def assert_spawning(self):

    22     if not Popen.thread_is_spawning():

    23         raise RuntimeError(

    24             '%s objects should only be shared between processes'

    25             ' through inheritance' % type(self).__name__

    26             )

    27 

    28 #

    29 # Try making some callable types picklable

    30 #

    31 

    32 from pickle import Pickler

    33 class ForkingPickler(Pickler):

    34     dispatch = Pickler.dispatch.copy()

    35 

    36     @classmethod

    37     def register(cls, type, reduce):

    38         def dispatcher(self, obj):

    39             rv = reduce(obj)

    40             self.save_reduce(obj=obj, *rv)

    41         cls.dispatch[type] = dispatcher

    42 

    43 def _reduce_method(m):

    44     if m.im_self is None:

    45         return getattr, (m.im_class, m.im_func.func_name)

    46     else:

    47         return getattr, (m.im_self, m.im_func.func_name)

    48 ForkingPickler.register(type(ForkingPickler.save), _reduce_method)

    49 

    50 def _reduce_method_descriptor(m):

    51     return getattr, (m.__objclass__, m.__name__)

    52 ForkingPickler.register(type(list.append), _reduce_method_descriptor)

    53 ForkingPickler.register(type(int.__add__), _reduce_method_descriptor)

    54 

    55 #def _reduce_builtin_function_or_method(m):

    56 #    return getattr, (m.__self__, m.__name__)

    57 #ForkingPickler.register(type(list().append), _reduce_builtin_function_or_method)

    58 #ForkingPickler.register(type(int().__add__), _reduce_builtin_function_or_method)

    59 

    60 try:

    61     from functools import partial

    62 except ImportError:

    63     pass

    64 else:

    65     def _reduce_partial(p):

    66         return _rebuild_partial, (p.func, p.args, p.keywords or {})

    67     def _rebuild_partial(func, args, keywords):

    68         return partial(func, *args, **keywords)

    69     ForkingPickler.register(partial, _reduce_partial)

    70 

    71 #

    72 # Unix

    73 #

    74 

    75 if sys.platform != 'win32':

    76     import time

    77 

    78     exit = os._exit

    79     duplicate = os.dup

    80     close = os.close

    81 

    82     #

    83     # We define a Popen class similar to the one from subprocess, but

    84     # whose constructor takes a process object as its argument.

    85     #

    86 

    87     class Popen(object):

    88 

    89         def __init__(self, process_obj):

    90             sys.stdout.flush()

    91             sys.stderr.flush()

    92             self.returncode = None

    93 

    94             self.pid = os.fork()

    95             if self.pid == 0:

    96                 if 'random' in sys.modules:

    97                     import random

    98                     random.seed()

    99                 code = process_obj._bootstrap()

   100                 sys.stdout.flush()

   101                 sys.stderr.flush()

   102                 os._exit(code)

   103 

   104         def poll(self, flag=os.WNOHANG):

   105             if self.returncode is None:

   106                 pid, sts = os.waitpid(self.pid, flag)

   107                 if pid == self.pid:

   108                     if os.WIFSIGNALED(sts):

   109                         self.returncode = -os.WTERMSIG(sts)

   110                     else:

   111                         assert os.WIFEXITED(sts)

   112                         self.returncode = os.WEXITSTATUS(sts)

   113             return self.returncode

   114 

   115         def wait(self, timeout=None):

   116             if timeout is None:

   117                 return self.poll(0)

   118             deadline = time.time() + timeout

   119             delay = 0.0005

   120             while 1:

   121                 res = self.poll()

   122                 if res is not None:

   123                     break

   124                 remaining = deadline - time.time()

   125                 if remaining <= 0:

   126                     break

   127                 delay = min(delay * 2, remaining, 0.05)

   128                 time.sleep(delay)

   129             return res

   130 

   131         def terminate(self):

   132             if self.returncode is None:

   133                 try:

   134                     os.kill(self.pid, signal.SIGTERM)

   135                 except OSError, e:

   136                     if self.wait(timeout=0.1) is None:

   137                         raise

   138 

   139         @staticmethod

   140         def thread_is_spawning():

   141             return False

   142 

   143 #

   144 # Windows

   145 #

   146 

   147 else:

   148     import thread

   149     import msvcrt

   150     import _subprocess

   151     import time

   152 

   153     from ._multiprocessing import win32, Connection, PipeConnection

   154     from .util import Finalize

   155 

   156     #try:

   157     #    from cPickle import dump, load, HIGHEST_PROTOCOL

   158     #except ImportError:

   159     from pickle import load, HIGHEST_PROTOCOL

   160 

   161     def dump(obj, file, protocol=None):

   162         ForkingPickler(file, protocol).dump(obj)

   163 

   164     #

   165     #

   166     #

   167 

   168     TERMINATE = 0x10000

   169     WINEXE = (sys.platform == 'win32' and getattr(sys, 'frozen', False))

   170 

   171     exit = win32.ExitProcess

   172     close = win32.CloseHandle

   173 

   174     #

   175     # _python_exe is the assumed path to the python executable.

   176     # People embedding Python want to modify it.

   177     #

   178 

   179     if sys.executable.lower().endswith('pythonservice.exe'):

   180         _python_exe = os.path.join(sys.exec_prefix, 'python.exe')

   181     else:

   182         _python_exe = sys.executable

   183 

   184     def set_executable(exe):

   185         global _python_exe

   186         _python_exe = exe

   187 

   188     #

   189     #

   190     #

   191 

   192     def duplicate(handle, target_process=None, inheritable=False):

   193         if target_process is None:

   194             target_process = _subprocess.GetCurrentProcess()

   195         return _subprocess.DuplicateHandle(

   196             _subprocess.GetCurrentProcess(), handle, target_process,

   197             0, inheritable, _subprocess.DUPLICATE_SAME_ACCESS

   198             ).Detach()

   199 

   200     #

   201     # We define a Popen class similar to the one from subprocess, but

   202     # whose constructor takes a process object as its argument.

   203     #

   204 

   205     class Popen(object):

   206         '''

   207         Start a subprocess to run the code of a process object

   208         '''

   209         _tls = thread._local()

   210 

   211         def __init__(self, process_obj):

   212             # create pipe for communication with child

   213             rfd, wfd = os.pipe()

   214 

   215             # get handle for read end of the pipe and make it inheritable

   216             rhandle = duplicate(msvcrt.get_osfhandle(rfd), inheritable=True)

   217             os.close(rfd)

   218 

   219             # start process

   220             cmd = get_command_line() + [rhandle]

   221             cmd = ' '.join('"%s"' % x for x in cmd)

   222             hp, ht, pid, tid = _subprocess.CreateProcess(

   223                 _python_exe, cmd, None, None, 1, 0, None, None, None

   224                 )

   225             ht.Close()

   226             close(rhandle)

   227 

   228             # set attributes of self

   229             self.pid = pid

   230             self.returncode = None

   231             self._handle = hp

   232 

   233             # send information to child

   234             prep_data = get_preparation_data(process_obj._name)

   235             to_child = os.fdopen(wfd, 'wb')

   236             Popen._tls.process_handle = int(hp)

   237             try:

   238                 dump(prep_data, to_child, HIGHEST_PROTOCOL)

   239                 dump(process_obj, to_child, HIGHEST_PROTOCOL)

   240             finally:

   241                 del Popen._tls.process_handle

   242                 to_child.close()

   243 

   244         @staticmethod

   245         def thread_is_spawning():

   246             return getattr(Popen._tls, 'process_handle', None) is not None

   247 

   248         @staticmethod

   249         def duplicate_for_child(handle):

   250             return duplicate(handle, Popen._tls.process_handle)

   251 

   252         def wait(self, timeout=None):

   253             if self.returncode is None:

   254                 if timeout is None:

   255                     msecs = _subprocess.INFINITE

   256                 else:

   257                     msecs = max(0, int(timeout * 1000 + 0.5))

   258 

   259                 res = _subprocess.WaitForSingleObject(int(self._handle), msecs)

   260                 if res == _subprocess.WAIT_OBJECT_0:

   261                     code = _subprocess.GetExitCodeProcess(self._handle)

   262                     if code == TERMINATE:

   263                         code = -signal.SIGTERM

   264                     self.returncode = code

   265 

   266             return self.returncode

   267 

   268         def poll(self):

   269             return self.wait(timeout=0)

   270 

   271         def terminate(self):

   272             if self.returncode is None:

   273                 try:

   274                     _subprocess.TerminateProcess(int(self._handle), TERMINATE)

   275                 except WindowsError:

   276                     if self.wait(timeout=0.1) is None:

   277                         raise

   278 

   279     #

   280     #

   281     #

   282 

   283     def is_forking(argv):

   284         '''

   285         Return whether commandline indicates we are forking

   286         '''

   287         if len(argv) >= 2 and argv[1] == '--multiprocessing-fork':

   288             assert len(argv) == 3

   289             return True

   290         else:

   291             return False

   292 

   293 

   294     def freeze_support():

   295         '''

   296         Run code for process object if this in not the main process

   297         '''

   298         if is_forking(sys.argv):

   299             main()

   300             sys.exit()

   301 

   302 

   303     def get_command_line():

   304         '''

   305         Returns prefix of command line used for spawning a child process

   306         '''

   307         if process.current_process()._identity==() and is_forking(sys.argv):

   308             raise RuntimeError('''

   309             Attempt to start a new process before the current process

   310             has finished its bootstrapping phase.

   311 

   312             This probably means that you are on Windows and you have

   313             forgotten to use the proper idiom in the main module:

   314 

   315                 if __name__ == '__main__':

   316                     freeze_support()

   317                     ...

   318 

   319             The "freeze_support()" line can be omitted if the program

   320             is not going to be frozen to produce a Windows executable.''')

   321 

   322         if getattr(sys, 'frozen', False):

   323             return [sys.executable, '--multiprocessing-fork']

   324         else:

   325             prog = 'from multiprocessing.forking import main; main()'

   326             return [_python_exe, '-c', prog, '--multiprocessing-fork']

   327 

   328 

   329     def main():

   330         '''

   331         Run code specifed by data received over pipe

   332         '''

   333         assert is_forking(sys.argv)

   334 

   335         handle = int(sys.argv[-1])

   336         fd = msvcrt.open_osfhandle(handle, os.O_RDONLY)

   337         from_parent = os.fdopen(fd, 'rb')

   338 

   339         process.current_process()._inheriting = True

   340         preparation_data = load(from_parent)

   341         prepare(preparation_data)

   342         self = load(from_parent)

   343         process.current_process()._inheriting = False

   344 

   345         from_parent.close()

   346 

   347         exitcode = self._bootstrap()

   348         exit(exitcode)

   349 

   350 

   351     def get_preparation_data(name):

   352         '''

   353         Return info about parent needed by child to unpickle process object

   354         '''

   355         from .util import _logger, _log_to_stderr

   356 

   357         d = dict(

   358             name=name,

   359             sys_path=sys.path,

   360             sys_argv=sys.argv,

   361             log_to_stderr=_log_to_stderr,

   362             orig_dir=process.ORIGINAL_DIR,

   363             authkey=process.current_process().authkey,

   364             )

   365 

   366         if _logger is not None:

   367             d['log_level'] = _logger.getEffectiveLevel()

   368 

   369         if not WINEXE:

   370             main_path = getattr(sys.modules['__main__'], '__file__', None)

   371             if not main_path and sys.argv[0] not in ('', '-c'):

   372                 main_path = sys.argv[0]

   373             if main_path is not None:

   374                 if not os.path.isabs(main_path) and \

   375                                           process.ORIGINAL_DIR is not None:

   376                     main_path = os.path.join(process.ORIGINAL_DIR, main_path)

   377                 d['main_path'] = os.path.normpath(main_path)

   378 

   379         return d

   380 

   381     #

   382     # Make (Pipe)Connection picklable

   383     #

   384 

   385     def reduce_connection(conn):

   386         if not Popen.thread_is_spawning():

   387             raise RuntimeError(

   388                 'By default %s objects can only be shared between processes\n'

   389                 'using inheritance' % type(conn).__name__

   390                 )

   391         return type(conn), (Popen.duplicate_for_child(conn.fileno()),

   392                             conn.readable, conn.writable)

   393 

   394     ForkingPickler.register(Connection, reduce_connection)

   395     ForkingPickler.register(PipeConnection, reduce_connection)

   396 

   397 #

   398 # Prepare current process

   399 #

   400 

   401 old_main_modules = []

   402 

   403 def prepare(data):

   404     '''

   405     Try to get current process ready to unpickle process object

   406     '''

   407     old_main_modules.append(sys.modules['__main__'])

   408 

   409     if 'name' in data:

   410         process.current_process().name = data['name']

   411 

   412     if 'authkey' in data:

   413         process.current_process()._authkey = data['authkey']

   414 

   415     if 'log_to_stderr' in data and data['log_to_stderr']:

   416         util.log_to_stderr()

   417 

   418     if 'log_level' in data:

   419         util.get_logger().setLevel(data['log_level'])

   420 

   421     if 'sys_path' in data:

   422         sys.path = data['sys_path']

   423 

   424     if 'sys_argv' in data:

   425         sys.argv = data['sys_argv']

   426 

   427     if 'dir' in data:

   428         os.chdir(data['dir'])

   429 

   430     if 'orig_dir' in data:

   431         process.ORIGINAL_DIR = data['orig_dir']

   432 

   433     if 'main_path' in data:

   434         main_path = data['main_path']

   435         main_name = os.path.splitext(os.path.basename(main_path))[0]

   436         if main_name == '__init__':

   437             main_name = os.path.basename(os.path.dirname(main_path))

   438 

   439         if main_name != 'ipython':

   440             import imp

   441 

   442             if main_path is None:

   443                 dirs = None

   444             elif os.path.basename(main_path).startswith('__init__.py'):

   445                 dirs = [os.path.dirname(os.path.dirname(main_path))]

   446             else:

   447                 dirs = [os.path.dirname(main_path)]

   448 

   449             assert main_name not in sys.modules, main_name

   450             file, path_name, etc = imp.find_module(main_name, dirs)

   451             try:

   452                 # We would like to do "imp.load_module('__main__', ...)"

   453                 # here.  However, that would cause 'if __name__ ==

   454                 # "__main__"' clauses to be executed.

   455                 main_module = imp.load_module(

   456                     '__parents_main__', file, path_name, etc

   457                     )

   458             finally:

   459                 if file:

   460                     file.close()

   461 

   462             sys.modules['__main__'] = main_module

   463             main_module.__name__ = '__main__'

   464 

   465             # Try to make the potentially picklable objects in

   466             # sys.modules['__main__'] realize they are in the main

   467             # module -- somewhat ugly.

   468             for obj in main_module.__dict__.values():

   469                 try:

   470                     if obj.__module__ == '__parents_main__':

   471                         obj.__module__ = '__main__'

   472                 except Exception:

   473                     pass