update connections and diagrams for reduced sockets

IPython/zmq/parallel/client.py

@@ -270,8 +270,9 b' class Client(HasTraits):'
     _control_socket=Instance('zmq.Socket')
     _iopub_socket=Instance('zmq.Socket')
     _notification_socket=Instance('zmq.Socket')
-    _mux_socket=Instance('zmq.Socket')
-    _task_socket=Instance('zmq.Socket')
+    _apply_socket=Instance('zmq.Socket')
+    _mux_ident=Str()
+    _task_ident=Str()
     _task_scheme=Str()
     _balanced_views=Dict()
     _direct_views=Dict()
@@ -401,16 +402,16 b' class Client(HasTraits):'
             self._ids.append(eid)
         self._ids = sorted(self._ids)
         if sorted(self._engines.keys()) != range(len(self._engines)) and \
-                        self._task_scheme == 'pure' and self._task_socket:
+                        self._task_scheme == 'pure' and self._task_ident:
             self._stop_scheduling_tasks()
     
     def _stop_scheduling_tasks(self):
         """Stop scheduling tasks because an engine has been unregistered
         from a pure ZMQ scheduler.
         """
-        
-        self._task_socket.close()
-        self._task_socket = None
+        self._task_ident = ''
+        # self._task_socket.close()
+        # self._task_socket = None
         msg = "An engine has been unregistered, and we are using pure " +\
               "ZMQ task scheduling.  Task farming will be disabled."
         if self.outstanding:
@@ -457,15 +458,18 b' class Client(HasTraits):'
         content = msg.content
         self._config['registration'] = dict(content)
         if content.status == 'ok':
+            self._apply_socket = self._context.socket(zmq.XREP)
+            self._apply_socket.setsockopt(zmq.IDENTITY, self.session.session)
             if content.mux:
-                self._mux_socket = self._context.socket(zmq.XREQ)
-                self._mux_socket.setsockopt(zmq.IDENTITY, self.session.session)
-                connect_socket(self._mux_socket, content.mux)
+                # self._mux_socket = self._context.socket(zmq.XREQ)
+                self._mux_ident = 'mux'
+                connect_socket(self._apply_socket, content.mux)
             if content.task:
                 self._task_scheme, task_addr = content.task
-                self._task_socket = self._context.socket(zmq.XREQ)
-                self._task_socket.setsockopt(zmq.IDENTITY, self.session.session)
-                connect_socket(self._task_socket, task_addr)
+                # self._task_socket = self._context.socket(zmq.XREQ)
+                # self._task_socket.setsockopt(zmq.IDENTITY, self.session.session)
+                connect_socket(self._apply_socket, task_addr)
+                self._task_ident = 'task'
             if content.notification:
                 self._notification_socket = self._context.socket(zmq.SUB)
                 connect_socket(self._notification_socket, content.notification)
@@ -484,7 +488,8 b' class Client(HasTraits):'
                 self._iopub_socket.setsockopt(zmq.IDENTITY, self.session.session)
                 connect_socket(self._iopub_socket, content.iopub)
             self._update_engines(dict(content.engines))
-                
+            # give XREP apply_socket some time to connect
+            time.sleep(0.25)
         else:
             self._connected = False
             raise Exception("Failed to connect!")
@@ -496,7 +501,7 b' class Client(HasTraits):'
     def _unwrap_exception(self, content):
         """unwrap exception, and remap engineid to int."""
         e = error.unwrap_exception(content)
-        print e.traceback
+        # print e.traceback
         if e.engine_info:
             e_uuid = e.engine_info['engine_uuid']
             eid = self._engines[e_uuid]
@@ -540,7 +545,7 b' class Client(HasTraits):'
             
             self._handle_stranded_msgs(eid, uuid)
                 
-        if self._task_socket and self._task_scheme == 'pure':
+        if self._task_ident and self._task_scheme == 'pure':
             self._stop_scheduling_tasks()
     
     def _handle_stranded_msgs(self, eid, uuid):
@@ -725,10 +730,8 b' class Client(HasTraits):'
         """
         if self._notification_socket:
             self._flush_notifications()
-        if self._mux_socket:
-            self._flush_results(self._mux_socket)
-        if self._task_socket:
-            self._flush_results(self._task_socket)
+        if self._apply_socket:
+            self._flush_results(self._apply_socket)
         if self._control_socket:
             self._flush_control(self._control_socket)
         if self._iopub_socket:
@@ -1030,6 +1033,12 b' class Client(HasTraits):'
         args = args if args is not None else []
         kwargs = kwargs if kwargs is not None else {}
         
+        if not self._ids:
+            # flush notification socket if no engines yet
+            any_ids = self.ids
+            if not any_ids:
+                raise error.NoEnginesRegistered("Can't execute without any connected engines.")
+        
         if balanced is None:
             if targets is None:
                 # default to balanced if targets unspecified
@@ -1074,7 +1083,7 b' class Client(HasTraits):'
         for name in ('bound', 'block', 'track'):
             assert loc[name] is not None, "kwarg %r must be specified!"%name
         
-        if self._task_socket is None:
+        if not self._task_ident:
             msg = "Task farming is disabled"
             if self._task_scheme == 'pure':
                 msg += " because the pure ZMQ scheduler cannot handle"
@@ -1106,7 +1115,7 b' class Client(HasTraits):'
         bufs = util.pack_apply_message(f,args,kwargs)
         content = dict(bound=bound)
         
-        msg = self.session.send(self._task_socket, "apply_request", 
+        msg = self.session.send(self._apply_socket, "apply_request", ident=self._task_ident,
                 content=content, buffers=bufs, subheader=subheader, track=track)
         msg_id = msg['msg_id']
         self.outstanding.add(msg_id)
@@ -1130,6 +1139,11 b' class Client(HasTraits):'
         This is a private method, see `apply` for details.
         Not to be called directly!
         """
+        
+        if not self._mux_ident:
+            msg = "Multiplexing is disabled"
+            raise RuntimeError(msg)
+        
         loc = locals()
         for name in ('bound', 'block', 'targets', 'track'):
             assert loc[name] is not None, "kwarg %r must be specified!"%name
@@ -1143,8 +1157,8 b' class Client(HasTraits):'
         msg_ids = []
         trackers = []
         for ident in idents:
-            msg = self.session.send(self._mux_socket, "apply_request", 
-                    content=content, buffers=bufs, ident=ident, subheader=subheader,
+            msg = self.session.send(self._apply_socket, "apply_request", 
+                    content=content, buffers=bufs, ident=[self._mux_ident, ident], subheader=subheader,
                     track=track)
             if track:
                 trackers.append(msg['tracker'])

IPython/zmq/parallel/controller.py

@@ -80,6 +80,7 b' class ControllerFactory(HubFactory):'
         # Multiplexer Queue (in a Process)
         q = mq(zmq.XREP, zmq.XREP, zmq.PUB, 'in', 'out')
         q.bind_in(self.client_info['mux'])
+        q.setsockopt_in(zmq.IDENTITY, 'mux')
         q.bind_out(self.engine_info['mux'])
         q.connect_mon(maybe_inproc)
         q.daemon=True
@@ -88,6 +89,7 b' class ControllerFactory(HubFactory):'
         # Control Queue (in a Process)
         q = mq(zmq.XREP, zmq.XREP, zmq.PUB, 'incontrol', 'outcontrol')
         q.bind_in(self.client_info['control'])
+        q.setsockopt_in(zmq.IDENTITY, 'control')
         q.bind_out(self.engine_info['control'])
         q.connect_mon(maybe_inproc)
         q.daemon=True
@@ -98,6 +100,7 b' class ControllerFactory(HubFactory):'
             q = mq(zmq.XREP, zmq.XREQ, zmq.PUB, 'intask', 'outtask')
             q.setsockopt_out(zmq.HWM, self.hwm)
             q.bind_in(self.client_info['task'][1])
+            q.setsockopt_in(zmq.IDENTITY, 'task')
             q.bind_out(self.engine_info['task'])
             q.connect_mon(maybe_inproc)
             q.daemon=True

IPython/zmq/parallel/scheduler.py

@@ -544,7 +544,8 b' class TaskScheduler(SessionFactory):'
 
 
 def launch_scheduler(in_addr, out_addr, mon_addr, not_addr, config=None,logname='ZMQ', 
-                            log_addr=None, loglevel=logging.DEBUG, scheme='lru'):
+                            log_addr=None, loglevel=logging.DEBUG, scheme='lru',
+                            identity=b'task'):
     from zmq.eventloop import ioloop
     from zmq.eventloop.zmqstream import ZMQStream
     
@@ -552,8 +553,11 b' def launch_scheduler(in_addr, out_addr, mon_addr, not_addr, config=None,logname='
     loop = ioloop.IOLoop()
     print (in_addr, out_addr, mon_addr, not_addr)
     ins = ZMQStream(ctx.socket(zmq.XREP),loop)
+    ins.setsockopt(zmq.IDENTITY, identity)
     ins.bind(in_addr)
+    
     outs = ZMQStream(ctx.socket(zmq.XREP),loop)
+    outs.setsockopt(zmq.IDENTITY, identity)
     outs.bind(out_addr)
     mons = ZMQStream(ctx.socket(zmq.PUB),loop)
     mons.connect(mon_addr)

IPython/zmq/parallel/tests/__init__.py

@@ -12,7 +12,7 b' blackhole = tempfile.TemporaryFile()'
 # nose setup/teardown
 
 def setup():
-    cp = Popen('ipcontrollerz --profile iptest -r --log-level 40'.split(), stdout=blackhole, stderr=STDOUT)
+    cp = Popen('ipcontrollerz --profile iptest -r --log-level 10 --log-to-file'.split(), stdout=blackhole, stderr=STDOUT)
     processes.append(cp)
     time.sleep(.5)
     add_engine()
@@ -22,7 +22,7 b' def setup():'
         c.spin()
 
 def add_engine(profile='iptest'):
-    ep = Popen(['ipenginez']+ ['--profile', profile, '--log-level', '40'], stdout=blackhole, stderr=STDOUT)
+    ep = Popen(['ipenginez']+ ['--profile', profile, '--log-level', '10', '--log-to-file'], stdout=blackhole, stderr=STDOUT)
     # ep.start()
     processes.append(ep)
     return ep

IPython/zmq/parallel/tests/clienttest.py

@@ -1,3 +1,5 b''
+import sys
+import tempfile
 import time
 from signal import SIGINT
 from multiprocessing import Process
@@ -62,7 +64,7 b' class ClusterTestCase(BaseZMQTestCase):'
         while time.time()-tic < timeout and len(self.client.ids) < n:
             time.sleep(0.1)
         
-        assert not self.client.ids < n, "waiting for engines timed out"
+        assert not len(self.client.ids) < n, "waiting for engines timed out"
     
     def connect_client(self):
         """connect a client with my Context, and track its sockets for cleanup"""
@@ -89,12 +91,15 b' class ClusterTestCase(BaseZMQTestCase):'
         self.engines=[]
     
     def tearDown(self):
+        
+        # close fds:
+        for e in filter(lambda e: e.poll() is not None, processes):
+            processes.remove(e)
+        
         self.client.close()
         BaseZMQTestCase.tearDown(self)
-        # [ e.terminate() for e in filter(lambda e: e.poll() is None, self.engines) ]
-        # [ e.wait() for e in self.engines ]
-    #     while len(self.client.ids) > self.base_engine_count:
-    #         time.sleep(.1)
-    #     del self.engines
-    #     BaseZMQTestCase.tearDown(self)
+        # this will be superfluous when pyzmq merges PR #88
+        self.context.term()
+        print tempfile.TemporaryFile().fileno(),
+        sys.stdout.flush()
         
  No newline at end of file

IPython/zmq/parallel/tests/test_client.py

@@ -17,7 +17,6 b' class TestClient(ClusterTestCase):'
         n = len(self.client.ids)
         self.add_engines(3)
         self.assertEquals(len(self.client.ids), n+3)
-        self.assertTrue
     
     def test_segfault_task(self):
         """test graceful handling of engine death (balanced)"""
@@ -179,8 +178,9 b' class TestClient(ClusterTestCase):'
     def test_get_result(self):
         """test getting results from the Hub."""
         c = clientmod.Client(profile='iptest')
-        t = self.client.ids[-1]
+        self.add_engines(1)
         ar = c.apply(wait, (1,), block=False, targets=t)
+        # give the monitor time to notice the message
         time.sleep(.25)
         ahr = self.client.get_result(ar.msg_ids)
         self.assertTrue(isinstance(ahr, AsyncHubResult))