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)