#!/usr/bin/env python
# -*- coding: utf-8 -*-
# File: remote.py
# Author: Yuxin Wu <ppwwyyxxc@gmail.com>
from ..utils import logger
try:
import zmq
except ImportError:
logger.warn("Error in 'import zmq'. remote feature won't be available")
__all__ = []
else:
__all__ = ['serve_data', 'RemoteData']
from .base import DataFlow
from .common import RepeatedData
from ..utils import logger
from ..utils.serialize import dumps, loads
def serve_data(ds, addr):
ctx = zmq.Context()
socket = ctx.socket(zmq.PUSH)
socket.set_hwm(10)
socket.bind(addr)
ds = RepeatedData(ds, -1)
try:
ds.reset_state()
logger.info("Serving data at {}".format(addr))
while True:
for dp in ds.get_data():
socket.send(dumps(dp), copy=False)
finally:
socket.setsockopt(zmq.LINGER, 0)
socket.close()
if not ctx.closed:
ctx.destroy(0)
class RemoteData(DataFlow):
def __init__(self, addr):
self.ctx = zmq.Context()
self.socket = self.ctx.socket(zmq.PULL)
self.socket.set_hwm(10)
self.socket.connect(addr)
def get_data(self):
while True:
dp = loads(self.socket.recv(copy=False))
yield dp
if __name__ == '__main__':
import sys
from tqdm import tqdm
from .raw import FakeData
addr = "tcp://127.0.0.1:8877"
if sys.argv[1] == 'serve':
ds = FakeData([(128,244,244,3)], 1000)
serve_data(ds, addr)
else:
ds = RemoteData(addr)
logger.info("Each DP is 73.5MB")
with tqdm(total=10000) as pbar:
for k in ds.get_data():
pbar.update()
