| 1380 | do=st.sampled_from(hu.device_options)) |
| 1381 | @settings(deadline=None, max_examples=50) |
| 1382 | def test_safe_blobs_queue(self, num_producers, num_consumers, |
| 1383 | capacity, num_blobs, do): |
| 1384 | init_net = core.Net('init_net') |
| 1385 | queue = init_net.CreateBlobsQueue( |
| 1386 | [], 1, capacity=capacity, num_blobs=num_blobs) |
| 1387 | producer_steps = [] |
| 1388 | truth = 0 |
| 1389 | for i in range(num_producers): |
| 1390 | name = 'producer_%d' % i |
| 1391 | net = core.Net(name) |
| 1392 | blobs = [net.ConstantFill([], 1, value=1.0, run_once=False) |
| 1393 | for times in range(num_blobs)] |
| 1394 | status = net.NextName() |
| 1395 | net.SafeEnqueueBlobs([queue] + blobs, blobs + [status]) |
| 1396 | count = (i + 1) * 10 |
| 1397 | step = core.execution_step(name, net, num_iter=count) |
| 1398 | truth += count |
| 1399 | producer_steps.append(step) |
| 1400 | producer_exit_net = core.Net('producer_exit_net') |
| 1401 | producer_exit_net.CloseBlobsQueue([queue], 0) |
| 1402 | producer_step = core.execution_step('producer', [ |
| 1403 | core.execution_step( |
| 1404 | 'producers', producer_steps, concurrent_substeps=True), |
| 1405 | core.execution_step('producer_exit', producer_exit_net)] |
| 1406 | ) |
| 1407 | |
| 1408 | consumer_steps = [] |
| 1409 | counters = [] |
| 1410 | const_1 = init_net.ConstantFill([], 1, value=1.0) |
| 1411 | for i in range(num_consumers): |
| 1412 | name = 'consumer_%d' % i |
| 1413 | net1 = core.Net(name) |
| 1414 | blobs = net1.SafeDequeueBlobs([queue], num_blobs + 1) |
| 1415 | status = blobs[-1] |
| 1416 | |
| 1417 | net2 = core.Net(name + '_counter') |
| 1418 | counter = init_net.ConstantFill([], 1, value=0.0) |
| 1419 | counters.append(counter) |
| 1420 | net2.Add([counter, const_1], counter) |
| 1421 | consumer_steps.append(core.execution_step( |
| 1422 | name, [net1, net2], should_stop_blob=status)) |
| 1423 | consumer_step = core.execution_step( |
| 1424 | 'consumer', consumer_steps, concurrent_substeps=True) |
| 1425 | |
| 1426 | init_step = core.execution_step('init', init_net) |
| 1427 | worker_step = core.execution_step( |
| 1428 | 'worker', [consumer_step, producer_step], concurrent_substeps=True) |
| 1429 | |
| 1430 | plan = core.Plan('test') |
| 1431 | plan.AddStep(init_step) |
| 1432 | plan.AddStep(worker_step) |
| 1433 | |
| 1434 | self.ws.run(plan) |
| 1435 | v = 0 |
| 1436 | for counter in counters: |
| 1437 | v += self.ws.blobs[str(counter)].fetch().tolist() |
| 1438 | self.assertEqual(v, truth) |
| 1439 | |