Produce data from ZMQ PULL socket(s). It is the receiver-side counterpart of :func:`send_dataflow_zmq`, which uses :mod:`tensorpack.utils.serialize` for serialization. See http://tensorpack.readthedocs.io/tutorial/efficient-dataflow.html#distributed-dataflow Attributes:
| 85 | |
| 86 | |
| 87 | class RemoteDataZMQ(DataFlow): |
| 88 | """ |
| 89 | Produce data from ZMQ PULL socket(s). |
| 90 | It is the receiver-side counterpart of :func:`send_dataflow_zmq`, which uses :mod:`tensorpack.utils.serialize` |
| 91 | for serialization. |
| 92 | See http://tensorpack.readthedocs.io/tutorial/efficient-dataflow.html#distributed-dataflow |
| 93 | |
| 94 | Attributes: |
| 95 | cnt1, cnt2 (int): number of data points received from addr1 and addr2 |
| 96 | """ |
| 97 | def __init__(self, addr1, addr2=None, hwm=50, bind=True): |
| 98 | """ |
| 99 | Args: |
| 100 | addr1,addr2 (str): addr of the zmq endpoint to connect to. |
| 101 | Use both if you need two protocols (e.g. both IPC and TCP). |
| 102 | I don't think you'll ever need 3. |
| 103 | hwm (int): ZMQ high-water mark (buffer size) |
| 104 | bind (bool): whether to connect or bind the endpoint |
| 105 | """ |
| 106 | assert addr1 |
| 107 | self._addr1 = addr1 |
| 108 | self._addr2 = addr2 |
| 109 | self._hwm = int(hwm) |
| 110 | self._guard = DataFlowReentrantGuard() |
| 111 | self._bind = bind |
| 112 | |
| 113 | def reset_state(self): |
| 114 | self.cnt1 = 0 |
| 115 | self.cnt2 = 0 |
| 116 | |
| 117 | def bind_or_connect(self, socket, addr): |
| 118 | if self._bind: |
| 119 | socket.bind(addr) |
| 120 | else: |
| 121 | socket.connect(addr) |
| 122 | |
| 123 | def __iter__(self): |
| 124 | with self._guard: |
| 125 | try: |
| 126 | ctx = zmq.Context() |
| 127 | if self._addr2 is None: |
| 128 | socket = ctx.socket(zmq.PULL) |
| 129 | socket.set_hwm(self._hwm) |
| 130 | self.bind_or_connect(socket, self._addr1) |
| 131 | |
| 132 | while True: |
| 133 | dp = loads(socket.recv(copy=False)) |
| 134 | yield dp |
| 135 | self.cnt1 += 1 |
| 136 | else: |
| 137 | socket1 = ctx.socket(zmq.PULL) |
| 138 | socket1.set_hwm(self._hwm) |
| 139 | self.bind_or_connect(socket1, self._addr1) |
| 140 | |
| 141 | socket2 = ctx.socket(zmq.PULL) |
| 142 | socket2.set_hwm(self._hwm) |
| 143 | self.bind_or_connect(socket2, self._addr2) |
| 144 |