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