def testOneEachFeeding(self):
# One feeder, one consumer
servers = self._create_local_cluster(consumer=1, feeder=1)
coord = coordinator.Coordinator()
feeder_thread = FeederThread(self, coord, servers, 'feeder', 0)
feeder_thread.add_remote_device('/job:consumer/task:0')
feeder_thread.start()
with ops.Graph().as_default():
with ops.device('/job:consumer/task:0'):
feeder = feeder_lib.Feeder(
dtypes=[dtypes_lib.string, dtypes_lib.string],
shapes=[[], []],
capacity=1)
out_a, out_b = feeder.get_fed_tensors()
with session_lib.Session(servers['consumer'][0].target) as session:
a, b = session.run([out_a, out_b])
self.assertEquals(b'a0', a)
self.assertEquals(b'b0', b)
a = session.run(out_a) # Omit b!
self.assertEquals(b'a1', a)
coord.request_stop()
coord.join()
feeder_thread.join()
def testFeederSeparateThread(self):
# Start a feeder on a seperate thread, but with a shared local queue
servers = self._create_local_cluster(worker=1)
coord = coordinator.Coordinator()
feed_thread = FeederThread(self, coord, servers, 'worker', 0)
feed_thread.start()
with ops.Graph().as_default():
with ops.device('/job:worker/task:0'):
feeder = feeder_lib.Feeder(
dtypes=[dtypes_lib.string, dtypes_lib.string],
shapes=[[], []],
capacity=1)
out_a, out_b = feeder.get_fed_tensors()
with session_lib.Session(servers['worker'][0].target) as session:
a, b = session.run([out_a, out_b])
self.assertEquals(b'a0', a)
self.assertEquals(b'b0', b)
a = session.run(out_a) # Omit b!
self.assertEquals(b'a1', a)
coord.request_stop()
coord.join()
feed_thread.join()
def testFeederActsLikeQueue(self):
# Tests that a feeder acts like a queue
feeder = feeder_lib.Feeder(
dtypes=[dtypes_lib.string, dtypes_lib.string],
shapes=[[], []],
capacity=10)
feeder.set_many_fed_tensors([
constant_op.constant(['a0', 'a1', 'a2']),
constant_op.constant(['b0', 'b1', 'b2'])
])
out_a, out_b = feeder.get_fed_tensors()
with self.test_session() as session:
coord = coordinator.Coordinator()
queue_runner_impl.start_queue_runners(session, coord=coord)
a, b = session.run([out_a, out_b])
self.assertEquals(b'a0', a)
self.assertEquals(b'b0', b)
a = session.run(out_a) # Omit b!
self.assertEquals(b'a1', a)
a, b = session.run([out_a, out_b])
self.assertEquals(b'a2', a)
self.assertEquals(b'b2', b) # queued together
a, b = session.run([out_a, out_b]) # loops around
self.assertEquals(b'a0', a)
self.assertEquals(b'b0', b) # queued together
coord.request_stop()
coord.join()
def testAddRemoteReplicas(self):
with ops.Graph().as_default():
for idx in range(3):
with ops.name_scope('replica_%d' % idx):
feeder = feeder_lib.Feeder(
dtypes=[dtypes_lib.string, dtypes_lib.string],
shapes=[[], []],
capacity=10)
feeder.add_remote_replicas(
'consumer',
replica_count=3,
feeder_task_num=idx,
replicas_per_feeder=2,
base_device_spec='/device:cpu:0')
# Examine ops...
op_types_by_scope_and_device = collections.defaultdict(
lambda: collections.defaultdict(collections.Counter))
for op in ops.get_default_graph().get_operations():
scope = '/'.join(op.name.split('/')[:-1])
dev = op.device
op_types_by_scope_and_device[scope][dev][op.type] += 1
expected_ops = collections.Counter({
'QueueEnqueueV2': 1,
'FIFOQueueV2': 1
})
expected_enq_devices = [('replica_0', [
'/job:consumer/replica:0/device:cpu:0',
'/job:consumer/replica:1/device:cpu:0',
]),
('replica_1', [
'/job:consumer/replica:2/device:cpu:0',
'/job:consumer/replica:0/device:cpu:0',
]),
('replica_2', [
'/job:consumer/replica:1/device:cpu:0',
'/job:consumer/replica:2/device:cpu:0',
])]
for scope, devs in expected_enq_devices:
for dev in devs:
self.assertEqual(expected_ops,
op_types_by_scope_and_device[scope][dev])
def __init__(self, test_case, coord, servers, job, task_num, prefix=''):
self.graph = ops.Graph()
self.coord = coord
self.server = servers[job][task_num]
self.remote_devices = []
# Just because we do tf.session(X) doesn't mean ops will located
# on the X task; wrapping all feeder creation/interaction in an
# extra tf.device(X) ensures that any ops that don't provider
# their own tf.device() wrapper will be placed on the correct "local"
# feeder task. A session can and does put ops that have no device
# assignment onto any of the tasks it knows about, not just the
# task passed as its target= argument!
self.device = '/job:%s/task:%d' % (job, task_num)
self.prefix = prefix
self.thread = test_case.checkedThread(target=self._feed_thread)
with self.graph.as_default(), ops.device(self.device):
self.feeder = feeder_lib.Feeder(
[dtypes_lib.string, dtypes_lib.string], [[], []], capacity=1)
self.feeder.set_many_fed_tensors(self._get_feed_values())
def _run_consumer(task, expected_keys):
server = servers['consumer'][task]
# Runs until everything in expected_keys has been seen at least once;
# fails if any prefix not in expected_keys shows up
with ops.Graph().as_default(), ops.device('/job:consumer/task:%d' %
task):
feeder = feeder_lib.Feeder(
dtypes=[dtypes_lib.string, dtypes_lib.string],
shapes=[[], []],
capacity=1)
out_a, out_b = feeder.get_fed_tensors()
counts = collections.Counter()
with session_lib.Session(server.target) as sess:
while True:
a, b = sess.run([out_a, out_b])
counts[a[:-1]] += 1
counts[b[:-1]] += 1
self.assertTrue(a[:-1] in expected_keys)
self.assertTrue(b[:-1] in expected_keys)
if all(counts[k] > 0 for k in expected_keys):
return
