Source code for nose2.plugins.mp

import logging
import multiprocessing
import select
import unittest
import collections

import os
import sys
import six

import multiprocessing.connection as connection
from nose2 import events, loader, result, runner, session, util

log = logging.getLogger(__name__)


class MultiProcess(events.Plugin):
[docs] configSection = 'multiprocess' def __init__(self): self.addArgument(self.setProcs, 'N', 'processes', '# o procs') self.testRunTimeout = self.config.as_float('test-run-timeout', 60.0) self.procs = self.config.as_int( 'processes', multiprocessing.cpu_count()) self.setAddress(self.config.as_str('bind_address', None)) self.cases = {} def setProcs(self, num): self.procs = int(num[0]) # FIXME merge n fix self.register() def setAddress(self, address): if address is None or address.strip() == '': address = [] else: address = [x.strip() for x in address.split(':')[:2]] #Background: On Windows, select.select only works on sockets. So the #ability to select a bindable address and optionally port for the mp #plugin was added. Pipes should support a form of select, but this #would require using pywin32. There are altnernatives but all have #some kind of downside. An alternative might be creating a connection #like object using a shared queue for incomings events. self.bind_host = None self.bind_port = 0 if sys.platform == "win32" or address: self.bind_host = '127.116.157.163' if address and address[0]: self.bind_host = address[0] self.bind_port = 0 if len(address) >= 2: self.bind_port = int(address[1]) def pluginsLoaded(self, event): self.addMethods('registerInSubprocess', 'startSubprocess', 'stopSubprocess') def startTestRun(self, event): event.executeTests = self._runmp def beforeInteraction(self, event): # prevent interactive plugins from running event.handled = True return False def _runmp(self, test, result): flat = list(self._flatten(test)) procs = self._startProcs() # send one initial task to each process for proc, conn in procs: if not flat: break caseid = flat.pop(0) conn.send(caseid) rdrs = [conn for proc, conn in procs if proc.is_alive()] while flat or rdrs: ready, _, _ = select.select(rdrs, [], [], self.testRunTimeout) for conn in ready: # XXX proc could be dead try: remote_events = conn.recv() except EOFError: # probably dead/12 log.warning("Subprocess connection closed unexpectedly") continue # XXX or die? if remote_events is None: # XXX proc is done, how to mark it dead? rdrs.remove(conn) continue # replay events testid, events = remote_events log.debug("Received results for %s", testid) for (hook, event) in events: log.debug("Received %s(%s)", hook, event) self._localize(event) getattr(self.session.hooks, hook)(event) # send a new test to the worker if there is one left if not flat: # if there isn't send None - it's the 'done' flag conn.send(None) continue caseid = flat.pop(0) conn.send(caseid) for _, conn in procs: conn.close() # ensure we wait until all processes are done before # exiting, to allow plugins running there to finalize for proc, _ in procs: proc.join() def _prepConns(self): """ If the bind_host is not none, return: (multiprocessing.connection.Listener, (address, port, authkey)) else: (parent_connection, child_connection) For the former case: accept must be called on the listener. In order to get a Connection object for the socket. """ if self.bind_host is not None: #prevent "accidental" wire crossing authkey = os.urandom(20) address = (self.bind_host, self.bind_port) listener = connection.Listener(address, authkey=authkey) return (listener, listener.address + (authkey,)) else: return multiprocessing.Pipe() def _acceptConns(self, parent_conn): """ When listener is is a connection.Listener instance: accept the next incoming connection. However, a timeout mechanism is needed. Since, this functionality was added to support mp over inet sockets, will assume a Socket based listen and will accept the private _socket member to get a low_level socket to do a select on. """ if isinstance(parent_conn, connection.Listener): #ick private interface rdrs = [parent_conn._listener._socket] readable, _, _ = select.select(rdrs, [], [], self.testRunTimeout) if readable: return parent_conn.accept() else: raise RuntimeError('MP: Socket Connection Failed') else: return parent_conn def _startProcs(self): # XXX create session export session_export = self._exportSession() procs = [] for i in range(0, self.procs): parent_conn, child_conn = self._prepConns() proc = multiprocessing.Process( target=procserver, args=(session_export, child_conn)) proc.daemon = True proc.start() parent_conn = self._acceptConns(parent_conn) procs.append((proc, parent_conn)) return procs def _flatten(self, suite): # XXX # examine suite tests to find out if they have class # or module fixtures and group them that way into names # of test classes or modules # ALSO record all test cases in self.cases mods = {} classes = {} stack = [suite] while stack: suite = stack.pop() for test in suite: if isinstance(test, unittest.TestSuite): stack.append(test) else: testid = util.test_name(test) self.cases[testid] = test if util.has_module_fixtures(test): mods.setdefault(test.__class__.__module__, []).append( testid) elif util.has_class_fixtures(test): classes.setdefault( "%s.%s" % (test.__class__.__module__, test.__class__.__name__), []).append(testid) else: yield testid for cls in sorted(classes.keys()): yield cls for mod in sorted(mods.keys()): yield mod def _localize(self, event): # XXX set loader, case, result etc to local ones, if present in event # (event case will be just the id) # (traceback in exc_info if any won't be real!) if hasattr(event, 'result'): event.result = self.session.testResult if hasattr(event, 'loader'): event.loader = self.session.testLoader if hasattr(event, 'runner'): event.runner = self.session.testRunner if hasattr(event, 'test') and isinstance(event.test, six.string_types): # remote event.case is the test id try: event.test = self.cases[event.test] except KeyError: event.test = self.session.testLoader.failedLoadTests( 'test_not_found', RuntimeError("Unable to locate test case for %s in " "main process" % event.test))._tests[0] def _exportSession(self): # argparse isn't pickleable # no plugin instances # no hooks export = {'config': self.session.config, 'verbosity': self.session.verbosity, 'startDir': self.session.startDir, 'topLevelDir': self.session.topLevelDir, 'logLevel': self.session.logLevel, # XXX classes or modules? 'pluginClasses': []} # XXX fire registerInSubprocess -- add those plugin classes # (classes must be pickleable!) event = RegisterInSubprocessEvent() # FIXME should be own event type self.session.hooks.registerInSubprocess(event) export['pluginClasses'].extend(event.pluginClasses) return export def procserver(session_export, conn):
# init logging system rlog = multiprocessing.log_to_stderr() rlog.setLevel(session_export['logLevel']) # make a real session from the "session" we got ssn = session.Session() ssn.config = session_export['config'] ssn.hooks = RecordingPluginInterface() ssn.verbosity = session_export['verbosity'] ssn.startDir = session_export['startDir'] ssn.topLevelDir = session_export['topLevelDir'] ssn.prepareSysPath() loader_ = loader.PluggableTestLoader(ssn) ssn.testLoader = loader_ result_ = result.PluggableTestResult(ssn) ssn.testResult = result_ runner_ = runner.PluggableTestRunner(ssn) # needed?? ssn.testRunner = runner_ # load and register plugins ssn.plugins = [ plugin(session=ssn) for plugin in session_export['pluginClasses']] rlog.debug("Plugins loaded: %s", ssn.plugins) for plugin in ssn.plugins: plugin.register() rlog.debug("Registered %s in subprocess", plugin) if isinstance(conn, collections.Sequence): conn = connection.Client(conn[:2], authkey=conn[2]) event = SubprocessEvent(loader_, result_, runner_, ssn.plugins, conn) res = ssn.hooks.startSubprocess(event) if event.handled and not res: conn.send(None) conn.close() ssn.hooks.stopSubprocess(event) return # receive and run tests executor = event.executeTests for testid in gentests(conn): if testid is None: break # XXX to handle weird cases like layers, need to # deal with the case that testid is something other # than a simple string. test = event.loader.loadTestsFromName(testid) # xxx try/except? rlog.debug("Execute test %s (%s)", testid, test) executor(test, event.result) events = [e for e in ssn.hooks.flush()] conn.send((testid, events)) rlog.debug("Log for %s returned", testid) conn.send(None) conn.close() ssn.hooks.stopSubprocess(event) # test generator def gentests(conn): while True: try: testid = conn.recv() if testid is None: return yield testid except EOFError: return # custom event classes class SubprocessEvent(events.Event):
[docs] """Event fired at start and end of subprocess execution. .. attribute :: loader Test loader instance .. attribute :: result Test result .. attribute :: plugins List of plugins loaded in the subprocess. .. attribute :: connection The :class:`multiprocessing.Connection` instance that the subprocess uses for communication with the main process. .. attribute :: executeTests Callable that will be used to execute tests. Plugins may set this attribute to wrap or otherwise change test execution. The callable must match the signature:: def execute(suite, result): ... """ def __init__(self, loader, result, runner, plugins, connection, **metadata): self.loader = loader self.result = result self.runner = runner self.plugins = plugins self.connection = connection self.executeTests = lambda test, result: test(result) super(SubprocessEvent, self).__init__(**metadata) class RegisterInSubprocessEvent(events.Event):
[docs] """Event fired to notify plugins that multiprocess testing will occur .. attribute :: pluginClasses Add a plugin class to this list to cause the plugin to be instantiated in each test-running subprocess. The most common thing to do, for plugins that need to run in subprocesses, is:: def registerInSubprocess(self, event): event.pluginClasses.append(self.__class__) """ def __init__(self, **metadata): self.pluginClasses = [] super(RegisterInSubprocessEvent, self).__init__(**metadata) # custom hook system that records calls and events class RecordingHook(events.Hook):
def __init__(self, method, interface): super(RecordingHook, self).__init__(method) self.interface = interface def __call__(self, event): res = super(RecordingHook, self).__call__(event) self.interface.log(self.method, event) return res class RecordingPluginInterface(events.PluginInterface): hookClass = RecordingHook noLogMethods = set( ['getTestCaseNames', 'startSubprocess', 'stopSubprocess', 'registerInSubprocess', 'moduleLoadedSuite']) def __init__(self): super(RecordingPluginInterface, self).__init__() self.events = [] def log(self, method, event): self.events.append((method, event)) def flush(self): events = self.events[:] self.events = [] return events def register(self, method, plugin): """Register a plugin for a method. :param method: A method name :param plugin: A plugin instance """ self._hookForMethod(method).append(plugin) def __getattr__(self, attr): if attr.startswith('__'): raise AttributeError('No %s in %s' % (attr, self)) return self._hookForMethod(attr) def _hookForMethod(self, method): # return recording hook for most hooks, normal hook for those # (like test loading and subprocess events) that we don't want # to send back to the main process. try: return self.hooks[method] except KeyError: if method in self.noLogMethods or method.startswith('loadTest'): hook = events.Hook(method) else: hook = self.hookClass(method, self) self.hooks[method] = hook return hook