Esempio n. 1
0
 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
Esempio n. 2
0
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()
Esempio n. 3
0
    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)
Esempio n. 4
0
    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)