Python zmq.Context() Examples

def run(self):
            enable_death_signal(_warn=self.identity == b'0')
            ctx = zmq.Context()

            # recv jobs
            socket = ctx.socket(zmq.PULL)
            socket.setsockopt(zmq.IDENTITY, self.identity)
            socket.set_hwm(self.hwm * self.batch_size)
            socket.connect(self.input_pipe)

            # send results
            out_socket = ctx.socket(zmq.PUSH)
            out_socket.set_hwm(max(self.hwm, 5))
            out_socket.connect(self.result_pipe)

            batch = []
            while True:
                dp = loads(socket.recv(copy=False))
                dp = self.map_func(dp)
                if dp is not None:
                    batch.append(dp)
                    if len(batch) == self.batch_size:
                        dp = BatchData.aggregate_batch(batch)
                        out_socket.send(dumps(dp), copy=False)
                        del batch[:] 

def multi_skts_method_other_args(self):
        @socket(zmq.PUB)
        @socket(zmq.SUB)
        def f(foo, pub, sub, bar=None):
            assert isinstance(pub, zmq.Socket), pub
            assert isinstance(sub, zmq.Socket), sub

            assert foo == 'mock'
            assert bar == 'fake'

            assert pub.context is zmq.Context.instance()
            assert sub.context is zmq.Context.instance()

            assert pub.type is zmq.PUB
            assert sub.type is zmq.SUB

        f('mock', bar='fake') 

def reset_state(self):
        _MultiProcessZMQDataFlow.reset_state(self)
        self._guard = DataFlowReentrantGuard()

        job_pipe = _get_pipe_name("dataflow_MaB_job")
        result_pipe = _get_pipe_name("dataflow_MaB_result")

        self.context = zmq.Context()
        self.socket = self.context.socket(zmq.PULL)
        self.socket.set_hwm(max(5, self.buffer_size // self.batch_size))
        _bind_guard(self.socket, result_pipe)

        dispatcher = MultiProcessMapAndBatchDataZMQ._Dispatcher(self.ds, job_pipe, self.buffer_size)

        self._proc_ids = [u'{}'.format(k).encode('utf-8') for k in range(self.num_proc)]
        worker_hwm = max(3, self.buffer_size // self.num_proc // self.batch_size)
        self._procs = [MultiProcessMapAndBatchDataZMQ._Worker(
            self._proc_ids[k], self.map_func, job_pipe, result_pipe, worker_hwm, self.batch_size)
            for k in range(self.num_proc)]

        self._procs.append(dispatcher)
        self._start_processes() 

def run(self):
            enable_death_signal(_warn=self.idx == 0)
            self.ds.reset_state()
            itr = _repeat_iter(lambda: self.ds)

            context = zmq.Context()
            socket = context.socket(zmq.PUSH)
            socket.set_hwm(self.hwm)
            socket.connect(self.conn_name)
            try:
                while True:
                    try:
                        dp = next(itr)
                        socket.send(dumps(dp), copy=False)
                    except Exception:
                        dp = _ExceptionWrapper(sys.exc_info()).pack()
                        socket.send(dumps(dp), copy=False)
                        raise
            # sigint could still propagate here, e.g. when nested
            except KeyboardInterrupt:
                pass
            finally:
                socket.close(0)
                context.destroy(0) 

def __init__(self, max_nprocs):
        import subprocess
        import sys
        import os
        import zmq

        # Since the output terminals are used for lots of debug output etc., we use
        # ZeroMQ to communicate with the workers.
        zctx = zmq.Context()
        socket = zctx.socket(zmq.REQ)
        port = socket.bind_to_random_port("tcp://*")
        cmd = 'import %s as mod; mod._mpi_worker("tcp://127.0.0.1:%d")' % (__name__, port)
        env = dict(os.environ)
        env['PYTHONPATH'] = ':'.join(sys.path)
        self.child = subprocess.Popen(['mpiexec', '-np', str(max_nprocs), sys.executable,
                                       '-c', cmd], env=env)
        self.socket = socket 

def __init__(self, ip_address, port_range):
        """
        Parameters
        ----------

        ip_address: str
           IP address of the client (where Parsl runs)
        port_range: tuple(int, int)
           Port range for the comms between client and interchange

        """
        self.context = zmq.Context()
        self.zmq_socket = self.context.socket(zmq.DEALER)
        self.zmq_socket.set_hwm(0)
        self.port = self.zmq_socket.bind_to_random_port("tcp://{}".format(ip_address),
                                                        min_port=port_range[0],
                                                        max_port=port_range[1]) 

def reset_state(self):
        _MultiProcessZMQDataFlow.reset_state(self)
        _ParallelMapData.reset_state(self)
        self._guard = DataFlowReentrantGuard()

        self.context = zmq.Context()
        self.socket = self.context.socket(zmq.DEALER)
        self.socket.set_hwm(self._buffer_size * 2)
        pipename = _get_pipe_name('dataflow-map')
        _bind_guard(self.socket, pipename)

        self._proc_ids = [u'{}'.format(k).encode('utf-8') for k in range(self.num_proc)]
        worker_hwm = int(self._buffer_size * 2 // self.num_proc)
        self._procs = [self._create_worker(self._proc_ids[k], pipename, worker_hwm)
                       for k in range(self.num_proc)]

        self._start_processes()
        self._fill_buffer()     # pre-fill the bufer 

def __init__(self, ip_address, port_range):
        """
        Parameters
        ----------

        ip_address: str
           IP address of the client (where Parsl runs)
        port_range: tuple(int, int)
           Port range for the comms between client and interchange

        """
        self.context = zmq.Context()
        self.zmq_socket = self.context.socket(zmq.DEALER)
        self.zmq_socket.set_hwm(0)
        self.port = self.zmq_socket.bind_to_random_port("tcp://{}".format(ip_address),
                                                        min_port=port_range[0],
                                                        max_port=port_range[1])
        self.poller = zmq.Poller()
        self.poller.register(self.zmq_socket, zmq.POLLOUT) 

def __init__(self, ip_address, port_range):
        """
        Parameters
        ----------

        ip_address: str
           IP address of the client (where Parsl runs)
        port_range: tuple(int, int)
           Port range for the comms between client and interchange

        """
        self.context = zmq.Context()
        self.results_receiver = self.context.socket(zmq.DEALER)
        self.results_receiver.set_hwm(0)
        self.port = self.results_receiver.bind_to_random_port("tcp://{}".format(ip_address),
                                                              min_port=port_range[0],
                                                              max_port=port_range[1]) 

def run(self):
        context = zmq.Context()
        receiver = context.socket(zmq.PULL)
        receiver.connect(self.worker_address)

        sink = context.socket(zmq.PUSH)
        sink.connect(self.sink_address)

        input_fn = self.input_fn_builder(receiver)

        self.logger.info('ready and listening')
        start_t = time.perf_counter()
        for r in self.estimator.predict(input_fn, yield_single_examples=False):
            # logger.info('new result!')
            send_ndarray(sink, r['client_id'], r['encodes'])
            time_used = time.perf_counter() - start_t
            start_t = time.perf_counter()
            self.logger.info('job %s\tsamples: %4d\tdone: %.2fs' %
                             (r['client_id'], r['encodes'].shape[0], time_used))

        receiver.close()
        sink.close()
        context.term()
        self.logger.info('terminated!') 

def _get_context(self, *args, **kwargs):
        '''
        Find the ``zmq.Context`` from ``args`` and ``kwargs`` at call time.

        First, if there is an keyword argument named ``context`` and it is a
        ``zmq.Context`` instance , we will take it.

        Second, we check all the ``args``, take the first ``zmq.Context``
        instance.

        Finally, we will provide default Context -- ``zmq.Context.instance``

        :return: a ``zmq.Context`` instance
        '''
        if self.context_name in kwargs:
            ctx = kwargs[self.context_name]

            if isinstance(ctx, zmq.Context):
                return ctx

        for arg in args:
            if isinstance(arg, zmq.Context):
                return arg
        # not specified by any decorator
        return zmq.Context.instance() 

def test_gc(self):
        """test close&term by garbage collection alone"""
        if PYPY:
            raise SkipTest("GC doesn't work ")
            
        # test credit @dln (GH #137):
        def gcf():
            def inner():
                ctx = self.Context()
                s = ctx.socket(zmq.PUSH)
            inner()
            gc.collect()
        t = Thread(target=gcf)
        t.start()
        t.join(timeout=1)
        self.assertFalse(t.is_alive(), "Garbage collection should have cleaned up context") 

def test_poll_base_socket(self):
        @asyncio.coroutine
        def test():
            ctx = zmq.Context()
            url = 'inproc://test'
            a = ctx.socket(zmq.PUSH)
            b = ctx.socket(zmq.PULL)
            self.sockets.extend([a, b])
            a.bind(url)
            b.connect(url)

            poller = zaio.Poller()
            poller.register(b, zmq.POLLIN)

            f = poller.poll(timeout=1000)
            assert not f.done()
            a.send_multipart([b'hi', b'there'])
            evt = yield from f
            self.assertEqual(evt, [(b, zmq.POLLIN)])
            recvd = b.recv_multipart()
            self.assertEqual(recvd, [b'hi', b'there'])
        self.loop.run_until_complete(test()) 

def test_multi_skts_with_name():
    @socket('foo', zmq.PUSH)
    @socket('bar', zmq.SUB)
    @socket('baz', zmq.PUB)
    def test(foo, bar, baz):
        assert isinstance(foo, zmq.Socket), foo
        assert isinstance(bar, zmq.Socket), bar
        assert isinstance(baz, zmq.Socket), baz

        assert foo.context is zmq.Context.instance()
        assert bar.context is zmq.Context.instance()
        assert baz.context is zmq.Context.instance()

        assert foo.type == zmq.PUSH
        assert bar.type == zmq.SUB
        assert baz.type == zmq.PUB
    test() 

def test_multi_skts_single_ctx():
    @context()
    @socket(zmq.PUB)
    @socket(zmq.SUB)
    @socket(zmq.PUSH)
    def test(ctx, pub, sub, push):
        assert isinstance(ctx, zmq.Context), ctx
        assert isinstance(pub, zmq.Socket), pub
        assert isinstance(sub, zmq.Socket), sub
        assert isinstance(push, zmq.Socket), push

        assert pub.context is ctx
        assert sub.context is ctx
        assert push.context is ctx

        assert pub.type == zmq.PUB
        assert sub.type == zmq.SUB
        assert push.type == zmq.PUSH
    test() 

def test_multi_skts():
    @socket(zmq.PUB)
    @socket(zmq.SUB)
    @socket(zmq.PUSH)
    def test(pub, sub, push):
        assert isinstance(pub, zmq.Socket), pub
        assert isinstance(sub, zmq.Socket), sub
        assert isinstance(push, zmq.Socket), push

        assert pub.context is zmq.Context.instance()
        assert sub.context is zmq.Context.instance()
        assert push.context is zmq.Context.instance()

        assert pub.type == zmq.PUB
        assert sub.type == zmq.SUB
        assert push.type == zmq.PUSH
    test() 

def test_ctx_multi_thread():
    @context()
    @context()
    def f(foo, bar):
        assert isinstance(foo, zmq.Context), foo
        assert isinstance(bar, zmq.Context), bar

        assert len(set(map(id, [foo, bar]))) == 2, set(map(id, [foo, bar]))

    threads = [threading.Thread(target=f) for i in range(8)]
    [t.start() for t in threads]
    [t.join() for t in threads]


##############################################
#  Test cases for @socket
############################################## 

def test_shadow_pyczmq(self):
        try:
            from pyczmq import zctx, zsocket, zstr
        except Exception:
            raise SkipTest("Requires pyczmq")
        
        ctx = zctx.new()
        a = zsocket.new(ctx, zmq.PUSH)
        zsocket.bind(a, "inproc://a")
        ctx2 = self.Context.shadow_pyczmq(ctx)
        b = ctx2.socket(zmq.PULL)
        b.connect("inproc://a")
        zstr.send(a, b'hi')
        rcvd = self.recv(b)
        self.assertEqual(rcvd, b'hi')
        b.close() 

def test_shadow(self):
        ctx = self.Context()
        ctx2 = self.Context.shadow(ctx.underlying)
        self.assertEqual(ctx.underlying, ctx2.underlying)
        s = ctx.socket(zmq.PUB)
        s.close()
        del ctx2
        self.assertFalse(ctx.closed)
        s = ctx.socket(zmq.PUB)
        ctx2 = self.Context.shadow(ctx.underlying)
        s2 = ctx2.socket(zmq.PUB)
        s.close()
        s2.close()
        ctx.term()
        self.assertRaisesErrno(zmq.EFAULT, ctx2.socket, zmq.PUB)
        del ctx2 

def test_cyclic_destroy(self):
        """ctx.destroy should succeed when cyclic ref prevents gc"""
        # test credit @dln (GH #137):
        class CyclicReference(object):
            def __init__(self, parent=None):
                self.parent = parent
            
            def crash(self, sock):
                self.sock = sock
                self.child = CyclicReference(self)
        
        def crash_zmq():
            ctx = self.Context()
            sock = ctx.socket(zmq.PULL)
            c = CyclicReference()
            c.crash(sock)
            ctx.destroy()
        
        crash_zmq() 

def tearDown(self):
        contexts = set([self.context])
        while self.sockets:
            sock = self.sockets.pop()
            contexts.add(sock.context) # in case additional contexts are created
            sock.close(0)
        for ctx in contexts:
            t = Thread(target=ctx.term)
            t.daemon = True
            t.start()
            t.join(timeout=2)
            if t.is_alive():
                # reset Context.instance, so the failure to term doesn't corrupt subsequent tests
                zmq.sugar.context.Context._instance = None
                raise RuntimeError("context could not terminate, open sockets likely remain in test")
        super(BaseZMQTestCase, self).tearDown() 

def _mpi_worker(addr):
    import importlib
    import zmq
    from pickle import loads, dumps

    rank = MPI.COMM_WORLD.Get_rank()
    if rank == 0:
        zctx = zmq.Context()
        socket = zctx.socket(zmq.REP)
        socket.connect(addr)

    while True:
        if rank == 0:
            pickled_msg = socket.recv()
        else:
            pickled_msg = None
        pickled_msg = MPI.COMM_WORLD.bcast(pickled_msg, root=0)
        msg = loads(pickled_msg)
        if msg == 'stop':
            if rank == 0:
                socket.send_pyobj('')
            break
        else:
            module_name, func_name = msg
            mod = importlib.import_module(module_name)
            func = getattr(mod, func_name)
            status = func(_return_status=True)
            if rank == 0:
                socket.send_pyobj(status)
    # All processes wait until they can terminate
    MPI.COMM_WORLD.barrier() 

def test_func_return():
    @context()
    def f(ctx):
        assert isinstance(ctx, zmq.Context), ctx
        return 'something'

    assert f() == 'something' 

def reset_state(self):
        super(MultiProcessRunnerZMQ, self).reset_state()
        self._guard = DataFlowReentrantGuard()
        self.context = zmq.Context()
        self.socket = self.context.socket(zmq.PULL)
        self.socket.set_hwm(self._hwm)
        pipename = _get_pipe_name('dataflow')
        _bind_guard(self.socket, pipename)

        self._procs = [MultiProcessRunnerZMQ._Worker(self.ds, pipename, self._hwm, idx)
                       for idx in range(self.num_proc)]
        self._start_processes() 

def test_copy(self):
        c1 = self.Context()
        c2 = copy.copy(c1)
        c2b = copy.deepcopy(c1)
        c3 = copy.deepcopy(c2)
        self.assert_(c2._shadow)
        self.assert_(c3._shadow)
        self.assertEqual(c1.underlying, c2.underlying)
        self.assertEqual(c1.underlying, c3.underlying)
        self.assertEqual(c1.underlying, c2b.underlying)
        s = c3.socket(zmq.PUB)
        s.close()
        c1.term() 

def test_destroy_no_sockets(self):
        ctx = self.Context()
        s = ctx.socket(zmq.PUB)
        s.bind_to_random_port('tcp://127.0.0.1')
        s.close()
        ctx.destroy()
        assert s.closed
        assert ctx.closed 

def test_term_thread(self):
        """ctx.term should not crash active threads (#139)"""
        ctx = self.Context()
        evt = Event()
        evt.clear()

        def block():
            s = ctx.socket(zmq.REP)
            s.bind_to_random_port('tcp://127.0.0.1')
            evt.set()
            try:
                s.recv()
            except zmq.ZMQError as e:
                self.assertEqual(e.errno, zmq.ETERM)
                return
            finally:
                s.close()
            self.fail("recv should have been interrupted with ETERM")
        t = Thread(target=block)
        t.start()
        
        evt.wait(1)
        self.assertTrue(evt.is_set(), "sync event never fired")
        time.sleep(0.01)
        ctx.term()
        t.join(timeout=1)
        self.assertFalse(t.is_alive(), "term should have interrupted s.recv()") 

def multi_skts_method(self, ctx, pub, sub, foo='bar'):
        assert isinstance(self, TestMethodDecorators), self
        assert isinstance(ctx, zmq.Context), ctx
        assert isinstance(pub, zmq.Socket), pub
        assert isinstance(sub, zmq.Socket), sub
        assert foo == 'bar'

        assert pub.context is ctx
        assert sub.context is ctx

        assert pub.type is zmq.PUB
        assert sub.type is zmq.SUB 

def test_instance_threadsafe(self):
        self.context.term() # clear default context

        q = Queue()
        # slow context initialization,
        # to ensure that we are both trying to create one at the same time
        class SlowContext(self.Context):
            def __init__(self, *a, **kw):
                time.sleep(1)
                super(SlowContext, self).__init__(*a, **kw)

        def f():
            q.put(SlowContext.instance())

        # call ctx.instance() in several threads at once
        N = 16
        threads = [ Thread(target=f) for i in range(N) ]
        [ t.start() for t in threads ]
        # also call it in the main thread (not first)
        ctx = SlowContext.instance()
        assert isinstance(ctx, SlowContext)
        # check that all the threads got the same context
        for i in range(N):
            thread_ctx = q.get(timeout=5)
            assert thread_ctx is ctx
        # cleanup
        ctx.term()
        [ t.join(timeout=5) for t in threads ] 

def test_many_sockets(self):
        """opening and closing many sockets shouldn't cause problems"""
        ctx = self.Context()
        for i in range(16):
            sockets = [ ctx.socket(zmq.REP) for i in range(65) ]
            [ s.close() for s in sockets ]
            # give the reaper a chance
            time.sleep(1e-2)
        ctx.term()