def run_test_ensure_timeout(self): from mars.worker.chunkholder import ensure_chunk promises = [ ensure_chunk(self, str(uuid.uuid4()), str(uuid.uuid4())), ensure_chunk(self, str(uuid.uuid4()), str(uuid.uuid4())), ] promise.all_(promises) \ .then(lambda *_: setattr(self, '_exc_info', None), lambda *exc: setattr(self, '_exc_info', exc)) \ .then(lambda *_: setattr(self, '_finished', True))
def run_test_cache(self): session_id = str(uuid.uuid4()) chunk_holder_ref = self.promise_ref(ChunkHolderActor.default_name()) chunk_store = self._chunk_store data_list = [] for _ in range(9): data_id = str(uuid.uuid4()) data = np.random.randint(0, 32767, (655360, ), np.int16) data_list.append((data_id, data)) def _put_chunk(data_key, data, *_): def _handle_reject(*exc): if issubclass(exc[0], NoDataToSpill): return six.reraise(*exc) try: ref = chunk_store.put(session_id, data_key, data) chunk_holder_ref.register_chunk(session_id, data_key) self.ctx.sleep(0.5) del ref except StoreFull: return chunk_holder_ref.spill_size(calc_data_size(data) * 2, _promise=True) \ .then(partial(_put_chunk, data_key, data), _handle_reject) data_promises = [] for data_id, data in data_list: data_promises.append(promise.finished().then( partial(_put_chunk, data_id, data))) def assert_true(v): assert v last_id = data_list[-1][0] p = promise.all_(data_promises) \ .then(lambda *_: assert_true(chunk_store.contains(session_id, last_id))) \ .then(lambda *_: ensure_chunk(self, session_id, last_id)) \ .then(lambda *_: assert_true(chunk_store.contains(session_id, last_id))) first_id = data_list[0][0] p = p.then(lambda *_: assert_true(not chunk_store.contains(session_id, first_id))) \ .then(lambda *_: ensure_chunk(self, session_id, first_id)) \ .then(lambda *_: assert_true(chunk_store.contains(session_id, first_id))) p = p.then(lambda *_: chunk_holder_ref.unregister_chunk(session_id, first_id)) \ .then(lambda *_: self._plasma_client.evict(128)) \ .then(lambda *_: assert_true(not chunk_store.contains(session_id, first_id))) p = p.then(lambda *_: chunk_holder_ref.unregister_chunk(session_id, last_id)) \ .then(lambda *_: self._plasma_client.evict(128)) \ .then(lambda *_: assert_true(not chunk_store.contains(session_id, last_id))) p.catch(lambda *exc: setattr(self, '_exc_info', exc)) \ .then(lambda *_: setattr(self, '_finished', True))