Made it possible to run multiple workerpools on the same machine

debian/changelog

@@ -1,4 +1,8 @@
   [Michele Simionato]
+  * Internal: the `host_cores` entries of the `[zworkers]` section can now
+    specify a port, as in `127.0.0.1:1910 4`, so that several workerpools
+    can run on the same host; `oq workers status` shows each pool as
+    `host:port`, and stopping the workerpools no longer leaks semaphores.
   * Fixed a bug affecting previous versions of the engine: with few sites
     only the ruptures of the first atomic group of a task were stored,
     so the disaggregation missed the ruptures of the other atomic groups

openquake/baselib/parallel.py

@@ -219,6 +219,18 @@ def count(word):
 host_cores = config.zworkers.host_cores.split(',')
 
 
+def parse_hostcore(hostcore, default_port):
+    """
+    :param hostcore: a string "host cores" or "host:port cores"
+    :param default_port: the ctrl_port to use if the port is missing
+    :returns: a triple (host, port, cores) with port an integer
+    """
+    hostport, cores = hostcore.split()
+    host, sep, port = hostport.partition(':')
+    # without an explicit port the default ctrl_port is used
+    return host, (int(port) if sep else int(default_port)), cores
+
+
 def calc_dir(job_id_or_fname):
     """
     :param job_id_or_fname: job ID (integer) or .hdf5 pathname (string)
@@ -255,8 +267,8 @@ def processpool_submit(self, func, args, monitor):
 @submit.add('zmq', 'slurm')
 def zmq_submit(self, func, args, monitor):
     idx = self.task_no % len(host_cores)
-    host = host_cores[idx].split()[0]
-    port = int(config.zworkers.ctrl_port)
+    host, port, _cores = parse_hostcore(
+        host_cores[idx], config.zworkers.ctrl_port)
     dest = 'tcp://%s:%d' % (host, port)
     logging.debug('Sending to %s', dest)
     with Socket(dest, zmq.REQ, 'connect', timeout=300) as sock:
@@ -1075,9 +1087,9 @@ def _task_ended(self, res, finished):
                 # measure the memory only for 1 task out of 10
                 # with 8 nodes the time to get the memory is 0.01 secs
                 for line in host_cores:
-                    host, _cores = line.split()
-                    addr = 'tcp://%s:%s' % (
-                        host, config.zworkers.ctrl_port)
+                    host, port, _cores = parse_hostcore(
+                        line, config.zworkers.ctrl_port)
+                    addr = 'tcp://%s:%s' % (host, port)
                     with Socket(addr, zmq.REQ, 'connect') as sock:
                         mem_gb += sock.send('memory_gb')
         elif self._shared:

openquake/baselib/tests/workerpool_test.py

@@ -16,40 +16,75 @@
 # You should have received a copy of the GNU Affero General Public License
 # along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.
 
-import time
+import os
 import unittest
-from openquake.baselib import config
-from openquake.baselib.workerpool import WorkerMaster
-from openquake.baselib.parallel import Starmap
-from openquake.baselib.general import socket_ready
+from openquake.baselib import config, parallel, zeromq as z
+from openquake.baselib.workerpool import WorkerMaster, ssh_args
+from openquake.baselib.parallel import Starmap, parse_hostcore
 
 
-def double(x):
-    return 2 * x
+def whoami(x):
+    # the parent of a worker process is the workerpool process
+    return os.getppid()
+
+
+class ParseHostcoreTestCase(unittest.TestCase):
+    def test_default_port(self):
+        self.assertEqual(parse_hostcore('10.0.0.1 4', 1909),
+                         ('10.0.0.1', 1909, '4'))
+
+    def test_explicit_port(self):
+        self.assertEqual(parse_hostcore('127.0.0.1:1910 -1', '1909'),
+                         ('127.0.0.1', 1910, '-1'))
+
+    def test_ssh_args(self):
+        host_cores = [('127.0.0.1', 1910, '2'), ('127.0.0.1', 1909, '3')]
+        got = [(host, port, cores)
+               for host, port, cores, _args in ssh_args(host_cores)]
+        self.assertEqual(got, [('127.0.0.1', 1910, '2'),
+                               ('127.0.0.1', 1909, '3')])
 
 
 class WorkerPoolTestCase(unittest.TestCase):
+    # two workerpools on the same host, on the two ports after ctrl_port
     @classmethod
     def setUpClass(cls):
         cls.z = config.zworkers.copy()
-        cls.z['host_cores'] = '127.0.0.1 4'
-        hostport = '127.0.0.1', int(cls.z['ctrl_port']) + 1
-        if not socket_ready(hostport):
-            raise unittest.SkipTest('The task streamer is off')
+        base = int(cls.z['ctrl_port'])
+        cls.ports = [base, base + 1]
+        cls.z['host_cores'] = '127.0.0.1:%d 2,127.0.0.1:%d 2' % tuple(
+            cls.ports)
+        cls.host_cores = parallel.host_cores[:]
+        parallel.host_cores[:] = cls.z['host_cores'].split(',')
         cls.master = WorkerMaster(cls.z)
         cls.master.start()
+        cls.status = cls.master.wait(30)
 
     def test(self):
         iterargs = ((i,) for i in range(10))
-        smap = Starmap(double, iterargs, distribute='zmq')
-        self.assertEqual(sum(res for res in smap), 90)
-        # sum[0, 2, 4, 6, 8, 10, 12, 14, 16, 18]
+        smap = Starmap(whoami, iterargs, distribute='zmq')
+        # the tasks alternate between the two workerpools
+        pids = {pid for pid in smap}
+        # ask the workerpools for their own PIDs: on Windows the Popen
+        # pid is the venv launcher, not the process running the pool
+        expected = set()
+        for port in self.ports:
+            url = 'tcp://127.0.0.1:%d' % port
+            with z.Socket(url, z.zmq.REQ, 'connect') as sock:
+                expected.add(sock.send('getpid'))
+        self.assertEqual(pids, expected)
 
     def test_status(self):
-        time.sleep(1)  # wait a bit for the workerpool to start
-        self.assertEqual(self.master.status(), [('127.0.0.1', 'running')])
+        self.assertEqual(
+            self.status,
+            [('127.0.0.1:%d' % port, 0, 2) for port in self.ports])
 
     @classmethod
     def tearDownClass(cls):
+        procs = list(cls.master.popens)
         cls.master.stop()
-        config.zworkers = cls.z
+        parallel.host_cores[:] = cls.host_cores
+        # the workerpools must exit on their own after the stop command
+        codes = [p.returncode for p in procs]
+        if codes != [0, 0]:  # explicit, so that it survives python -O
+            raise AssertionError('exit codes %s' % codes)

openquake/baselib/workerpool.py

@@ -60,22 +60,22 @@ def get_zworkers(job_id):
     return {}
 
 
-def ssh_args(zworkers):
+def ssh_args(host_cores):
     """
-    :yields: triples (hostIP, num_cores, [ssh remote python command])
+    :param host_cores: list of triples (host, ctrl_port, num_cores)
+    :yields: quadruples (hostIP, ctrl_port, num_cores, [ssh remote python
+             command])
     """
     user = getpass.getuser()
     if os.environ.get('OQ_WORKERPOOL_ERROR'):
         # can be set in CI to emulate a bug in ssh_args
         raise RuntimeError('Testing a bug in ssh_args')
-    if zworkers['host_cores'].strip():
-        for hostcores in zworkers['host_cores'].split(','):
-            host, cores = hostcores.split()
-            if host == '127.0.0.1':  # localhost
-                yield host, cores, [sys.executable]
-            else:
-                yield host, cores, [
-                    'ssh', '-f', '-T', user + '@' + host, sys.executable]
+    for host, port, cores in host_cores:
+        if host == '127.0.0.1':  # localhost
+            yield host, port, cores, [sys.executable]
+        else:
+            yield host, port, cores, [
+                'ssh', '-f', '-T', user + '@' + host, sys.executable]
 
 
 class WorkerMaster(object):
@@ -90,21 +90,31 @@ def __init__(self, zworkers, receiver_ports=None):
         # NB: receiver_ports is not used but needed for compliance
         self.ctrl_port = int(self.zworkers['ctrl_port'])
         self.host_cores = (
-            [hc.split() for hc in self.zworkers['host_cores'].split(',')]
+            [parallel.parse_hostcore(hc, self.ctrl_port)
+             for hc in self.zworkers['host_cores'].split(',')]
             if self.zworkers['host_cores'] else [])
         self.popens = []
 
+    def _running(self):
+        """
+        :yields: (host, port, ctrl_url) for the workerpools already running
+        """
+        for host, port, _ in self.host_cores:
+            if general.socket_ready((host, port)):
+                yield host, port, 'tcp://%s:%s' % (host, port)
+
     def start(self):
         """
         Start multiple workerpools on remote servers via ssh and/or a single
         workerpool on localhost.
         """
         starting = []
-        for host, cores, args in ssh_args(self.zworkers):
-            if general.socket_ready((host, self.ctrl_port)):
-                print('%s:%s already running' % (host, self.ctrl_port))
+        for host, port, cores, args in ssh_args(self.host_cores):
+            if general.socket_ready((host, port)):
+                print('%s:%s already running' % (host, port))
                 continue
-            args += ['-m', 'openquake.baselib.workerpool', cores]
+            args += ['-m', 'openquake.baselib.workerpool', cores,
+                     '--ctrl-port', str(port)]
             if host != '127.0.0.1':
                 print('%s: if it hangs, check the ssh keys' % ' '.join(args))
             self.popens.append(subprocess.Popen(args))
@@ -116,23 +126,25 @@ def stop(self):
         Send a "stop" command to all worker pools
         """
         stopped = []
-        for host, _ in self.host_cores:
-            if not general.socket_ready((host, self.ctrl_port)):
-                continue
-            ctrl_url = 'tcp://%s:%s' % (host, self.ctrl_port)
+        for host, _port, ctrl_url in self._running():
             with z.Socket(ctrl_url, z.zmq.REQ, 'connect') as sock:
                 try:
                     print(sock.send('stop'))
                 except z.TimeoutError:
                     pass  # workerpool dead for some reason
                 stopped.append(host)
         for popen in self.popens:
-            popen.terminate()
-            # since we are not consuming any output from the spawned process
-            # we must call wait() after terminate() to have Popen()
-            # fully deallocate the process file descriptors, otherwise
-            # zombies will arise
-            popen.wait()
+            # the workerpool needs time to shut down its process pool,
+            # otherwise its semaphores leak: terminate only if it hangs
+            try:
+                popen.wait(timeout=10)
+            except subprocess.TimeoutExpired:
+                popen.terminate()
+                # since we are not consuming any output from the spawned
+                # process we must call wait() after terminate() to have
+                # Popen() fully deallocate the process file descriptors,
+                # otherwise zombies will arise
+                popen.wait()
         self.popens = []
         return 'stopped %s' % stopped
 
@@ -142,7 +154,7 @@ def kill(self):
         in case of hard out of memory situations
         """
         killed = []
-        for host, cores, args in ssh_args(self.zworkers):
+        for host, _port, _cores, args in ssh_args(self.host_cores):
             args = args[:-1] + ['killall', '-r', 'oq-zworker|multiprocessing']
             print(' '.join(args))
             subprocess.run(args)
@@ -151,25 +163,22 @@ def kill(self):
 
     def status(self):
         """
-        :returns: a list [(host, running, total), ...]
+        :returns: a list [(host:port, running, total), ...]
         """
         executing = []
-        for host, _cores in self.host_cores:
-            if not general.socket_ready((host, self.ctrl_port)):
-                continue
-            ctrl_url = 'tcp://%s:%s' % (host, self.ctrl_port)
+        for host, port, ctrl_url in self._running():
             with z.Socket(ctrl_url, z.zmq.REQ, 'connect') as sock:
                 running = len(sock.send('get_executing').split())
                 total = sock.send('get_num_workers')
-                executing.append((host, running, total))
+                executing.append(('%s:%s' % (host, port), running, total))
         return executing
 
     def send_jobs(self):
         """
         Send an asynchronous "run_jobs" command to the first WorkerPool
         """
-        host, _cores = self.host_cores[0]
-        ctrl_url = 'tcp://%s:%s' % (host, self.ctrl_port)
+        host, port, _cores = self.host_cores[0]
+        ctrl_url = 'tcp://%s:%s' % (host, port)
         with z.Socket(ctrl_url, z.zmq.REQ, 'connect') as sock:
             return sock.send('run_jobs')
 
@@ -192,10 +201,7 @@ def restart(self):
         """
         Stop and start again
         """
-        for host, _ in self.host_cores:
-            if not general.socket_ready((host, self.ctrl_port)):
-                continue
-            ctrl_url = 'tcp://%s:%s' % (host, self.ctrl_port)
+        for _host, _port, ctrl_url in self._running():
             with z.Socket(ctrl_url, z.zmq.REQ, 'connect', timeout=120) as sock:
                 sock.send('restart')
         return 'restarted'
@@ -217,8 +223,8 @@ def debug(self):
             task_no = 0
             with z.Socket(receiver, z.zmq.PULL, 'bind') as pull:
                 mon.backurl = 'tcp://%s:%s' % (rec_host, pull.port)
-                for host, _ in self.host_cores:
-                    url = 'tcp://%s:%d' % (host, self.ctrl_port)
+                for host, port, _ in self.host_cores:
+                    url = 'tcp://%s:%s' % (host, port)
                     print('Sending to', url)
                     with z.Socket(url, z.zmq.REQ, 'connect') as sock:
                         for i in range(2):
@@ -344,7 +350,7 @@ def stop(self):
         return 'WorkerPool on %s stopped' % self.hostname
 
 
-def workerpool(num_workers: int=-1, job_id: int=0):
+def workerpool(num_workers: int=-1, job_id: int=0, *, ctrl_port: int=0):
     """
     Start a workerpool with the given number of workers.
     """
@@ -353,12 +359,15 @@ def workerpool(num_workers: int=-1, job_id: int=0):
     signal.signal(signal.SIGINT, signal.SIG_IGN)
     if hasattr(signal, 'SIGHUP'):
         signal.signal(signal.SIGHUP, signal.SIG_IGN)
+    # 0 means use the ctrl_port in the configuration
+    port = ctrl_port or int(config.zworkers['ctrl_port'])
     # NB: unexpected errors will appear in the worker service log
-    wpool = WorkerPool(int(config.zworkers['ctrl_port']), num_workers, job_id)
+    wpool = WorkerPool(port, num_workers, job_id)
     wpool.start()
 
 workerpool.num_workers = dict(help='number of cores to use')
 workerpool.job_id = dict(help='associated job, if any')
+workerpool.ctrl_port = dict(help='ZMQ control port (0 means from config)')
 
 
 if __name__ == '__main__':

openquake/engine/openquake.cfg

@@ -92,6 +92,8 @@ authkey = changeme
 calc_timeout =
 
 [zworkers]
+# comma-separated entries "host cores" or "host:port cores"; the port
+# defaults to ctrl_port and allows several workerpools on the same host
 host_cores = 127.0.0.1 -1
 ctrl_port = 1909