(test_case, x0, src, dst)
| 27 | @unittest.skipIf(os.getenv("ONEFLOW_TEST_CPU_ONLY"), "only test cpu cases") |
| 28 | class TestComm(flow.unittest.TestCase): |
| 29 | def _test_send_recv(test_case, x0, src, dst): |
| 30 | rank = flow.env.get_rank() |
| 31 | if rank == src: |
| 32 | x1 = x0 |
| 33 | flow.comm.send(x1, dst) |
| 34 | |
| 35 | x2 = x0 |
| 36 | flow.comm.send(x2, dst) |
| 37 | elif rank == dst: |
| 38 | x1 = flow.comm.recv(src) |
| 39 | test_case.assertTrue(np.array_equal(x1.numpy(), x0.numpy())) |
| 40 | test_case.assertEqual(x1.device, x0.device) |
| 41 | |
| 42 | x2 = flow.zeros_like(x0) |
| 43 | flow.comm.recv(src, out=x2) |
| 44 | test_case.assertTrue(np.array_equal(x2.numpy(), x0.numpy())) |
| 45 | test_case.assertEqual(x2.device, x0.device) |
| 46 | else: |
| 47 | # do nothing |
| 48 | pass |
| 49 | |
| 50 | @flow.unittest.skip_unless_1n2d() |
| 51 | def test_send_recv_2_devices(test_case): |
no test coverage detected