def post_create(self):
super().post_create()
self._dispatch_ref = self.ctx.actor_ref(DispatchActor.default_uid())
self._dispatch_ref.register_free_slot(self.uid, 'receiver')
self._receiver_manager_ref = self.ctx.actor_ref(ReceiverManagerActor.default_uid())
if not self.ctx.has_actor(self._receiver_manager_ref):
self._receiver_manager_ref = None
def start_transfer_test_pool(**kwargs):
address = kwargs.pop('address')
plasma_size = kwargs.pop('plasma_size')
with create_actor_pool(n_process=1, backend='gevent', address=address, **kwargs) as pool:
pool.create_actor(SchedulerClusterInfoActor, [address],
uid=SchedulerClusterInfoActor.default_uid())
pool.create_actor(WorkerClusterInfoActor, [address],
uid=WorkerClusterInfoActor.default_uid())
pool.create_actor(PlasmaKeyMapActor, uid=PlasmaKeyMapActor.default_uid())
pool.create_actor(StorageManagerActor, uid=StorageManagerActor.default_uid())
pool.create_actor(ChunkMetaActor, uid=ChunkMetaActor.default_uid())
pool.create_actor(DispatchActor, uid=DispatchActor.default_uid())
pool.create_actor(QuotaActor, 1024 * 1024 * 20, uid=MemQuotaActor.default_uid())
shared_holder_ref = pool.create_actor(SharedHolderActor,
plasma_size, uid=SharedHolderActor.default_uid())
pool.create_actor(StatusActor, address, uid=StatusActor.default_uid())
pool.create_actor(IORunnerActor)
pool.create_actor(StorageClientActor, uid=StorageClientActor.default_uid())
pool.create_actor(InProcHolderActor)
pool.create_actor(ReceiverManagerActor, uid=ReceiverManagerActor.default_uid())
try:
yield pool
finally:
shared_holder_ref.destroy()
def testReceiverWorker(self):
pool_addr = f'localhost:{get_next_port()}'
options.worker.spill_directory = tempfile.mkdtemp(
prefix='mars_test_receiver_')
session_id = str(uuid.uuid4())
mock_data = np.array([1, 2, 3, 4])
serialized_arrow_data = dataserializer.serialize(mock_data)
data_size = serialized_arrow_data.total_bytes
dumped_mock_data = dataserializer.dumps(mock_data)
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())
chunk_key7 = str(uuid.uuid4())
chunk_key8 = str(uuid.uuid4())
chunk_key9 = str(uuid.uuid4())
with start_transfer_test_pool(address=pool_addr, plasma_size=self.plasma_storage_size) as pool, \
self.run_actor_test(pool) as test_actor:
storage_client = test_actor.storage_client
receiver_ref = test_actor.promise_ref(
pool.create_actor(ReceiverWorkerActor, uid=str(uuid.uuid4())))
receiver_manager_ref = test_actor.promise_ref(
ReceiverManagerActor.default_uid())
# SCENARIO 1: create two writers and write with chunks
self.waitp(
receiver_ref.create_data_writers(session_id,
[chunk_key1, chunk_key2],
[data_size] * 2,
test_actor,
_promise=True))
receiver_ref.receive_data_part(
session_id, [chunk_key1, chunk_key2], [True, False],
dumped_mock_data,
dumped_mock_data[:len(dumped_mock_data) // 2])
self.assertEqual(receiver_ref.check_status(session_id, chunk_key1),
ReceiveStatus.RECEIVED)
self.assertEqual(receiver_ref.check_status(session_id, chunk_key2),
ReceiveStatus.RECEIVING)
receiver_ref.receive_data_part(
session_id, [chunk_key2], [True],
dumped_mock_data[len(dumped_mock_data) // 2:])
self.assertEqual(receiver_ref.check_status(session_id, chunk_key2),
ReceiveStatus.RECEIVED)
assert_array_equal(
storage_client.get_object(session_id,
chunk_key1,
[DataStorageDevice.SHARED_MEMORY],
_promise=False), mock_data)
assert_array_equal(
storage_client.get_object(session_id,
chunk_key2,
[DataStorageDevice.SHARED_MEMORY],
_promise=False), mock_data)
# SCENARIO 2: one of the writers failed to create,
# will test both existing and non-existing keys
old_create_writer = StorageClient.create_writer
def _create_writer_with_fail(self, session_id, chunk_key, *args,
**kwargs):
if chunk_key == fail_key:
if kwargs.get('_promise', True):
return promise.finished(*build_exc_info(ValueError),
**dict(_accept=False))
else:
raise ValueError
return old_create_writer(self, session_id, chunk_key, *args,
**kwargs)
with patch_method(StorageClient.create_writer, new=_create_writer_with_fail), \
self.assertRaises(ValueError):
fail_key = chunk_key4
self.waitp(
receiver_ref.create_data_writers(
session_id, [chunk_key3, chunk_key4, chunk_key5],
[data_size] * 3,
test_actor,
ensure_cached=False,
_promise=True))
self.assertEqual(receiver_ref.check_status(session_id, chunk_key3),
ReceiveStatus.NOT_STARTED)
self.assertEqual(receiver_ref.check_status(session_id, chunk_key4),
ReceiveStatus.NOT_STARTED)
self.assertEqual(receiver_ref.check_status(session_id, chunk_key5),
ReceiveStatus.NOT_STARTED)
with patch_method(StorageClient.create_writer,
new=_create_writer_with_fail):
fail_key = chunk_key2
self.waitp(
receiver_ref.create_data_writers(session_id,
[chunk_key2, chunk_key3],
[data_size] * 2,
test_actor,
ensure_cached=False,
_promise=True))
# SCENARIO 3: transfer timeout
receiver_manager_ref.register_pending_keys(session_id,
[chunk_key6])
self.waitp(
receiver_ref.create_data_writers(session_id, [chunk_key6],
[data_size],
test_actor,
timeout=1,
_promise=True))
with self.assertRaises(TimeoutError):
self.waitp(
receiver_manager_ref.add_keys_callback(session_id,
[chunk_key6],
_promise=True))
# SCENARIO 4: cancelled transfer (both before and during transfer)
receiver_manager_ref.register_pending_keys(session_id,
[chunk_key7])
self.waitp(
receiver_ref.create_data_writers(session_id, [chunk_key7],
[data_size],
test_actor,
timeout=1,
_promise=True))
receiver_ref.cancel_receive(session_id, [chunk_key2, chunk_key7])
with self.assertRaises(KeyError):
receiver_ref.receive_data_part(
session_id, [chunk_key7], [False],
dumped_mock_data[:len(dumped_mock_data) // 2])
with self.assertRaises(KeyError):
self.waitp(
receiver_manager_ref.add_keys_callback(session_id,
[chunk_key7],
_promise=True))
# SCENARIO 5: sender halt and receiver is notified (reusing previous unsuccessful key)
receiver_manager_ref.register_pending_keys(session_id,
[chunk_key7])
mock_ref = pool.actor_ref(test_actor.uid, address='MOCK_ADDR')
self.waitp(
receiver_ref.create_data_writers(session_id, [chunk_key7],
[data_size],
mock_ref,
timeout=1,
_promise=True))
receiver_ref.notify_dead_senders(['MOCK_ADDR'])
with self.assertRaises(WorkerDead):
self.waitp(
receiver_manager_ref.add_keys_callback(session_id,
[chunk_key7],
_promise=True))
# SCENARIO 6: successful transfer without promise
receiver_ref.create_data_writers(session_id, [chunk_key8],
[data_size],
mock_ref,
use_promise=False)
receiver_ref.receive_data_part(session_id, [chunk_key8], [True],
dumped_mock_data)
self.assertEqual(receiver_ref.check_status(session_id, chunk_key8),
ReceiveStatus.RECEIVED)
assert_array_equal(
storage_client.get_object(session_id,
chunk_key8,
[DataStorageDevice.SHARED_MEMORY],
_promise=False), mock_data)
# SCENARIO 7: failed transfer without promise
with patch_method(StorageClient.create_writer, new=_create_writer_with_fail), \
self.assertRaises(ValueError):
fail_key = chunk_key9
receiver_ref.create_data_writers(session_id, [chunk_key9],
[data_size],
mock_ref,
use_promise=False)
def testReceiverManager(self):
pool_addr = f'localhost:{get_next_port()}'
session_id = str(uuid.uuid4())
mock_data = np.array([1, 2, 3, 4])
serialized_data = dataserializer.dumps(mock_data)
data_size = len(serialized_data)
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())
chunk_key7 = str(uuid.uuid4())
with start_transfer_test_pool(address=pool_addr, plasma_size=self.plasma_storage_size) as pool, \
self.run_actor_test(pool) as test_actor:
mock_receiver_ref = pool.create_actor(MockReceiverWorkerActor,
uid=str(uuid.uuid4()))
storage_client = test_actor.storage_client
receiver_manager_ref = test_actor.promise_ref(
ReceiverManagerActor.default_uid())
# SCENARIO 1: test transferring existing keys
self.waitp(
storage_client.create_writer(
session_id, chunk_key1, data_size,
[DataStorageDevice.DISK
]).then(lambda writer: promise.finished().then(
lambda *_: writer.write(serialized_data)).then(
lambda *_: writer.close())))
result = self.waitp(
receiver_manager_ref.create_data_writers(session_id,
[chunk_key1],
[data_size],
test_actor,
_promise=True))
self.assertEqual(result[0].uid, mock_receiver_ref.uid)
self.assertEqual(result[1][0], ReceiveStatus.RECEIVED)
# test adding callback for transferred key (should return immediately)
result = self.waitp(
receiver_manager_ref.add_keys_callback(session_id,
[chunk_key1],
_promise=True))
self.assertTupleEqual(result, ())
receiver_manager_ref.register_pending_keys(
session_id, [chunk_key1, chunk_key2])
self.assertEqual(
receiver_manager_ref.filter_receiving_keys(
session_id, [chunk_key1, chunk_key2, 'non_exist']),
[chunk_key2])
# SCENARIO 2: test transferring new keys and wait on listeners
result = self.waitp(
receiver_manager_ref.create_data_writers(
session_id, [chunk_key2, chunk_key3], [data_size] * 2,
test_actor,
_promise=True))
self.assertEqual(result[0].uid, mock_receiver_ref.uid)
self.assertIsNone(result[1][0])
# transfer with transferring keys will report RECEIVING
result = self.waitp(
receiver_manager_ref.create_data_writers(session_id,
[chunk_key2],
[data_size],
test_actor,
_promise=True))
self.assertEqual(result[1][0], ReceiveStatus.RECEIVING)
# add listener and finish transfer
receiver_manager_ref.add_keys_callback(session_id, [chunk_key1, chunk_key2], _promise=True) \
.then(lambda *s: test_actor.set_result(s))
mock_receiver_ref.receive_data_part(session_id, [chunk_key2],
[True], serialized_data)
mock_receiver_ref.receive_data_part(session_id, [chunk_key3],
[True], serialized_data)
self.get_result(5)
# SCENARIO 3: test listening on multiple transfers
receiver_manager_ref.create_data_writers(
session_id, [chunk_key4, chunk_key5], [data_size] * 2, test_actor, _promise=True) \
.then(lambda *s: test_actor.set_result(s))
self.get_result(5)
# add listener
receiver_manager_ref.add_keys_callback(session_id, [chunk_key4, chunk_key5], _promise=True) \
.then(lambda *s: test_actor.set_result(s))
mock_receiver_ref.receive_data_part(session_id, [chunk_key4],
[True], serialized_data)
# when some chunks are not transferred, promise will not return
with self.assertRaises(TimeoutError):
self.get_result(0.5)
mock_receiver_ref.receive_data_part(session_id, [chunk_key5],
[True], serialized_data)
self.get_result(5)
# SCENARIO 4: test listening on transfer with errors
self.waitp(
receiver_manager_ref.create_data_writers(session_id,
[chunk_key6],
[data_size],
test_actor,
_promise=True))
receiver_manager_ref.add_keys_callback(session_id, [chunk_key6], _promise=True) \
.then(lambda *s: test_actor.set_result(s)) \
.catch(lambda *exc: test_actor.set_result(exc, accept=False))
mock_receiver_ref.cancel_receive(session_id, [chunk_key6])
with self.assertRaises(ExecutionInterrupted):
self.get_result(5)
# SCENARIO 5: test creating writers without promise
ref, statuses = receiver_manager_ref.create_data_writers(
session_id, [chunk_key7], [data_size],
test_actor,
use_promise=False)
self.assertIsNone(statuses[0])
self.assertEqual(ref.uid, mock_receiver_ref.uid)
# SCENARIO 6: test transferring lost keys
storage_client.delete(session_id, [chunk_key1])
result = self.waitp(
receiver_manager_ref.create_data_writers(session_id,
[chunk_key1],
[data_size],
test_actor,
_promise=True))
self.assertEqual(result[0].uid, mock_receiver_ref.uid)
self.assertIsNone(result[1][0])
# add listener and finish transfer
receiver_manager_ref.add_keys_callback(session_id, [chunk_key1], _promise=True) \
.then(lambda *s: test_actor.set_result(s))
mock_receiver_ref.receive_data_part(session_id, [chunk_key1],
[True], serialized_data)
self.get_result(5)