def start_send_recv_pool():
with start_transfer_test_pool(
address=send_pool_addr, plasma_size=self.plasma_storage_size) as sp:
sp.create_actor(SenderActor, uid=SenderActor.default_name())
with start_transfer_test_pool(
address=recv_pool_addr, plasma_size=self.plasma_storage_size) as rp:
rp.create_actor(MockReceiverActor, uid=ReceiverActor.default_name())
yield sp, rp
def testSender(self):
send_pool_addr = 'localhost:%d' % get_next_port()
recv_pool_addr = 'localhost:%d' % get_next_port()
recv_pool_addr2 = 'localhost:%d' % get_next_port()
options.worker.spill_directory = tempfile.mkdtemp(
prefix='mars_test_sender_')
session_id = str(uuid.uuid4())
mock_data = np.array([1, 2, 3, 4])
chunk_key1 = str(uuid.uuid4())
chunk_key2 = str(uuid.uuid4())
@contextlib.contextmanager
def start_send_recv_pool():
with start_transfer_test_pool(
address=send_pool_addr,
plasma_size=self.plasma_storage_size) as sp:
sp.create_actor(SenderActor, uid=SenderActor.default_uid())
with start_transfer_test_pool(
address=recv_pool_addr,
plasma_size=self.plasma_storage_size) as rp:
rp.create_actor(MockReceiverActor,
uid=ReceiverActor.default_uid())
yield sp, rp
with start_send_recv_pool() as (send_pool, recv_pool):
sender_ref = send_pool.actor_ref(SenderActor.default_uid())
receiver_ref = recv_pool.actor_ref(ReceiverActor.default_uid())
with self.run_actor_test(send_pool) as test_actor:
storage_client = test_actor.storage_client
# send when data missing
sender_ref_p = test_actor.promise_ref(sender_ref)
sender_ref_p.send_data(session_id, str(uuid.uuid4()), recv_pool_addr, _promise=True) \
.then(lambda *s: test_actor.set_result(s)) \
.catch(lambda *exc: test_actor.set_result(exc, accept=False))
with self.assertRaises(DependencyMissing):
self.get_result(5)
# send data in spill
serialized = dataserializer.serialize(mock_data)
self.waitp(
storage_client.create_writer(
session_id, chunk_key1, serialized.total_bytes,
[DataStorageDevice.DISK
]).then(lambda writer: promise.finished().then(
lambda *_: writer.write(serialized)).then(
lambda *_: writer.close())))
sender_ref_p.send_data(session_id, chunk_key1, recv_pool_addr, _promise=True) \
.then(lambda *s: test_actor.set_result(s)) \
.catch(lambda *exc: test_actor.set_result(exc, accept=False))
self.get_result(5)
assert_array_equal(
mock_data,
receiver_ref.get_result_data(session_id, chunk_key1))
storage_client.delete(session_id, chunk_key1)
# send data in plasma store
self.waitp(
storage_client.put_object(
session_id, chunk_key1, mock_data,
[DataStorageDevice.SHARED_MEMORY]))
sender_ref_p.send_data(session_id, chunk_key1, recv_pool_addr, _promise=True) \
.then(lambda *s: test_actor.set_result(s)) \
.catch(lambda *exc: test_actor.set_result(exc, accept=False))
self.get_result(5)
assert_array_equal(
mock_data,
receiver_ref.get_result_data(session_id, chunk_key1))
# send data to multiple targets
with start_transfer_test_pool(
address=recv_pool_addr2,
plasma_size=self.plasma_storage_size) as rp2:
recv_ref2 = rp2.create_actor(
MockReceiverActor, uid=ReceiverActor.default_uid())
self.waitp(
sender_ref_p.send_data(
session_id,
chunk_key1, [recv_pool_addr, recv_pool_addr2],
_promise=True))
# send data to already transferred / transferring
sender_ref_p.send_data(session_id, chunk_key1,
[recv_pool_addr, recv_pool_addr2], _promise=True) \
.then(lambda *s: test_actor.set_result(s)) \
.catch(lambda *exc: test_actor.set_result(exc, accept=False))
self.get_result(5)
assert_array_equal(
mock_data,
recv_ref2.get_result_data(session_id, chunk_key1))
# send data to non-exist endpoint which causes error
self.waitp(
storage_client.put_object(
session_id, chunk_key2, mock_data,
[DataStorageDevice.SHARED_MEMORY]))
sender_ref_p.send_data(session_id, chunk_key2, recv_pool_addr2, _promise=True) \
.then(lambda *s: test_actor.set_result(s)) \
.catch(lambda *exc: test_actor.set_result(exc, accept=False))
with self.assertRaises(BrokenPipeError):
self.get_result(5)
def mocked_receive_data_part(*_):
raise ChecksumMismatch
with patch_method(MockReceiverActor.receive_data_part,
new=mocked_receive_data_part):
sender_ref_p.send_data(session_id, chunk_key2, recv_pool_addr, _promise=True) \
.then(lambda *s: test_actor.set_result(s)) \
.catch(lambda *exc: test_actor.set_result(exc, accept=False))
with self.assertRaises(ChecksumMismatch):
self.get_result(5)
def testSender(self):
send_pool_addr = 'localhost:%d' % get_next_port()
recv_pool_addr = 'localhost:%d' % get_next_port()
recv_pool_addr2 = 'localhost:%d' % get_next_port()
options.worker.spill_directory = os.path.join(
tempfile.gettempdir(), 'mars_spill_%d_%d' % (os.getpid(), id(run_transfer_worker)))
session_id = str(uuid.uuid4())
mock_data = np.array([1, 2, 3, 4])
chunk_key1 = str(uuid.uuid4())
chunk_key2 = str(uuid.uuid4())
@contextlib.contextmanager
def start_send_recv_pool():
with start_transfer_test_pool(
address=send_pool_addr, plasma_size=self.plasma_storage_size) as sp:
sp.create_actor(SenderActor, uid=SenderActor.default_name())
with start_transfer_test_pool(
address=recv_pool_addr, plasma_size=self.plasma_storage_size) as rp:
rp.create_actor(MockReceiverActor, uid=ReceiverActor.default_name())
yield sp, rp
with start_send_recv_pool() as (send_pool, recv_pool):
chunk_holder_ref = send_pool.actor_ref(ChunkHolderActor.default_name())
sender_ref = send_pool.actor_ref(SenderActor.default_name())
receiver_ref = recv_pool.actor_ref(ReceiverActor.default_name())
sender_mapper_ref = send_pool.actor_ref(PlasmaKeyMapActor.default_name())
store = PlasmaChunkStore(self._plasma_client, sender_mapper_ref)
with self.run_actor_test(send_pool) as test_actor:
# send when data missing
sender_ref_p = test_actor.promise_ref(sender_ref)
sender_ref_p.send_data(session_id, str(uuid.uuid4()), recv_pool_addr, _promise=True) \
.then(lambda *s: test_actor.set_result(s, destroy=False)) \
.catch(lambda *exc: test_actor.set_result(exc, accept=False, destroy=False))
with self.assertRaises(DependencyMissing):
self.get_result(5)
# send data in spill
write_spill_file(chunk_key1, mock_data)
sender_ref_p.send_data(session_id, chunk_key1, recv_pool_addr, _promise=True) \
.then(lambda *s: test_actor.set_result(s, destroy=False)) \
.catch(lambda *exc: test_actor.set_result(exc, accept=False, destroy=False))
self.get_result(5)
assert_array_equal(mock_data, receiver_ref.get_result_data(session_id, chunk_key1))
os.unlink(build_spill_file_name(chunk_key1))
# send data in plasma store
store.put(session_id, chunk_key1, mock_data)
chunk_holder_ref.register_chunk(session_id, chunk_key1)
sender_ref_p.send_data(session_id, chunk_key1, recv_pool_addr, _promise=True) \
.then(lambda *s: test_actor.set_result(s, destroy=False)) \
.catch(lambda *exc: test_actor.set_result(exc, accept=False, destroy=False))
self.get_result(5)
assert_array_equal(mock_data, receiver_ref.get_result_data(session_id, chunk_key1))
# send data to multiple targets
with start_transfer_test_pool(
address=recv_pool_addr2, plasma_size=self.plasma_storage_size) as rp2:
recv_ref2 = rp2.create_actor(MockReceiverActor, uid=ReceiverActor.default_name())
sender_ref_p.send_data(session_id, chunk_key1,
[recv_pool_addr, recv_pool_addr2], _promise=True)
# send data to already transferred / transferring
sender_ref_p.send_data(session_id, chunk_key1,
[recv_pool_addr, recv_pool_addr2], _promise=True) \
.then(lambda *s: test_actor.set_result(s, destroy=False)) \
.catch(lambda *exc: test_actor.set_result(exc, accept=False, destroy=False))
self.get_result(5)
assert_array_equal(mock_data, recv_ref2.get_result_data(session_id, chunk_key1))
# send data to non-exist endpoint which causes error
store.put(session_id, chunk_key2, mock_data)
chunk_holder_ref.register_chunk(session_id, chunk_key2)
sender_ref_p.send_data(session_id, chunk_key2, recv_pool_addr2, _promise=True) \
.then(lambda *s: test_actor.set_result(s, destroy=False)) \
.catch(lambda *exc: test_actor.set_result(exc, accept=False, destroy=False))
with self.assertRaises(BrokenPipeError):
self.get_result(5)
def mocked_receive_data_part(*_):
raise ChecksumMismatch
with patch_method(MockReceiverActor.receive_data_part, new=mocked_receive_data_part):
sender_ref_p.send_data(session_id, chunk_key2, recv_pool_addr, _promise=True) \
.then(lambda *s: test_actor.set_result(s, destroy=False)) \
.catch(lambda *exc: test_actor.set_result(exc, accept=False, destroy=False))
with self.assertRaises(ChecksumMismatch):
self.get_result(5)
def testSender(self):
send_pool_addr = f'localhost:{get_next_port()}'
recv_pool_addr = f'localhost:{get_next_port()}'
recv_pool_addr2 = f'localhost:{get_next_port()}'
options.worker.spill_directory = tempfile.mkdtemp(
prefix='mars_test_sender_')
session_id = str(uuid.uuid4())
mock_data = np.array([1, 2, 3, 4])
chunk_key1 = str(uuid.uuid4())
chunk_key2 = str(uuid.uuid4())
chunk_key3 = str(uuid.uuid4())
chunk_key4 = str(uuid.uuid4())
chunk_key5 = str(uuid.uuid4())
chunk_key6 = str(uuid.uuid4())
@contextlib.contextmanager
def start_send_recv_pool():
with start_transfer_test_pool(
address=send_pool_addr,
plasma_size=self.plasma_storage_size) as sp:
sp.create_actor(SenderActor, uid=SenderActor.default_uid())
with start_transfer_test_pool(
address=recv_pool_addr,
plasma_size=self.plasma_storage_size) as rp:
rp.create_actor(MockReceiverWorkerActor,
uid=ReceiverWorkerActor.default_uid())
yield sp, rp
with start_send_recv_pool() as (send_pool, recv_pool), \
self.run_actor_test(send_pool) as test_actor:
sender_ref = send_pool.actor_ref(SenderActor.default_uid())
storage_client = test_actor.storage_client
sender_ref_p = test_actor.promise_ref(sender_ref)
# SCENARIO 1: send when data missing
with self.assertRaises(DependencyMissing):
self.waitp(
sender_ref_p.send_data(session_id, ['non_exist'],
[recv_pool_addr],
_promise=True))
# SCENARIO 2: send data to non-exist endpoint which causes error
self.waitp(
storage_client.put_objects(session_id, [chunk_key1],
[mock_data],
[DataStorageDevice.SHARED_MEMORY]))
with self.assertRaises((BrokenPipeError, ConnectionRefusedError)):
self.waitp(
sender_ref_p.send_data(session_id, [chunk_key1],
[recv_pool_addr2],
_promise=True))
with start_transfer_test_pool(
address=recv_pool_addr2,
plasma_size=self.plasma_storage_size) as rp2:
mock_recv_ref2 = rp2.create_actor(
MockReceiverWorkerActor,
uid=ReceiverWorkerActor.default_uid())
# SCENARIO 3: send data to multiple targets
self.waitp(
storage_client.put_objects(
session_id, [chunk_key2], [mock_data],
[DataStorageDevice.SHARED_MEMORY]))
self.waitp(
sender_ref_p.send_data(session_id, [chunk_key2],
[recv_pool_addr, recv_pool_addr2],
block_size=128,
_promise=True))
# send data to already transferred / transferring
self.waitp(
sender_ref_p.send_data(session_id, [chunk_key2],
[recv_pool_addr, recv_pool_addr2],
_promise=True))
assert_array_equal(
mock_data,
mock_recv_ref2.get_result_data(session_id, chunk_key2))
# SCENARIO 4: send multiple data at one time
self.waitp(
storage_client.put_objects(
session_id, [chunk_key3, chunk_key4], [mock_data] * 2,
[DataStorageDevice.SHARED_MEMORY]))
self.waitp(
sender_ref_p.send_data(session_id,
[chunk_key3, chunk_key4],
[recv_pool_addr2],
_promise=True))
assert_array_equal(
mock_data,
mock_recv_ref2.get_result_data(session_id, chunk_key3))
assert_array_equal(
mock_data,
mock_recv_ref2.get_result_data(session_id, chunk_key4))
# SCENARIO 5: send chunks already under transfer
self.waitp(
storage_client.put_objects(
session_id, [chunk_key5], [mock_data],
[DataStorageDevice.SHARED_MEMORY]))
mock_recv_ref2.set_receive_delay_key(session_id, chunk_key5, 1)
sender_ref_p.send_data(session_id, [chunk_key2, chunk_key5],
[recv_pool_addr2],
_promise=True)
self.waitp(
sender_ref_p.send_data(session_id, [chunk_key5],
[recv_pool_addr2],
_promise=True))
assert_array_equal(
mock_data,
mock_recv_ref2.get_result_data(session_id, chunk_key5))
# SCENARIO 6: send chunks already under transfer
self.waitp(
storage_client.put_objects(
session_id, [chunk_key6], [mock_data],
[DataStorageDevice.SHARED_MEMORY]))
mock_recv_ref2.set_receive_error_key(session_id, chunk_key6)
with self.assertRaises(ValueError):
self.waitp(
sender_ref_p.send_data(session_id, [chunk_key6],
[recv_pool_addr2],
_promise=True))