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 test_spawn(self, raises=False): def _task(): self.ctx.sleep(0.1) if raises: raise SystemError ref = self.promise_ref('ServeActor') promise.all_([self.spawn_promised(_task) for _ in range(4)]) \ .then(lambda *_: ref.serve(0, delay=0, _promise=True)) \ .catch(lambda *exc: ref.serve(exc[0].__name__, delay=0, _promise=True)) \ .then(lambda *_: setattr(self, '_finished', True))
def subsequent_all(*_): def func(idx, *_, **kw): return ref.serve(idx, _promise=True, **kw) for idx in range(10): promises.append(func(idx * 2).then(functools.partial(func, idx * 2 + 1))) return promise.all_(promises)
def waitp(self, *promises, **kw): timeout = kw.pop('timeout', 10) if len(promises) > 1: p = promise.all_(promises) else: p = promises[0] p.then(lambda *s: self._test_actor_ref.set_result(s, _tell=True), lambda *exc: self._test_actor_ref.set_result(exc, accept=False, _tell=True)) return self.get_result(timeout)
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))
def register(self, session_id, chunk_keys): import numpy as np cache_ref = self.promise_ref(ChunkHolderActor.default_name()) left_keys = set(chunk_keys) def _put_chunk(chunk_key, data, spill_times=1): try: refs = self._chunk_store.put(session_id, chunk_key, data) cache_ref.register_chunk(session_id, chunk_key) del refs left_keys.remove(chunk_key) except StoreFull: return cache_ref.spill_size(2 * spill_times * calc_data_size(data), _promise=True) \ .then(partial(_put_chunk, chunk_key, data, 2 * spill_times)) promises = [] for idx, chunk_key in enumerate(chunk_keys): data = np.ones((640 * 1024,), dtype=np.int16) * idx promises.append(promise.Promise(done=True) \ .then(partial(_put_chunk, chunk_key, data))) promise.all_(promises).then(lambda *_: setattr(self, '_finished', True))
def _gen_all_promise(*_): prior_promises = [ gen_promise(idx, bool((idx + 1) % 2)) for idx in range(4) ] return promise.all_(prior_promises)
def testAll(self): promises = weakref.WeakValueDictionary() req_queue = Queue() value_list = [] time_unit = 0.1 def test_thread_body(): while True: idx, v, success = req_queue.get() if v is None: break value_list.append(('thread_body', v)) time.sleep(time_unit) promises[idx].step_next([(v, ), dict(_accept=success)]) def gen_promise(value, accept=True): p = promise.Promise() promises[p.id] = p req_queue.put((p.id, value + 1, accept)) return p try: thread = threading.Thread(target=test_thread_body) thread.daemon = True thread.start() value_list = [] promise.all_([]).then(lambda: value_list.append(('all', 0))).wait() self.assertListEqual(value_list, [('all', 0)]) value_list = [] prior_promises = [gen_promise(idx) for idx in range(4)] promise.all_(prior_promises).then(lambda: value_list.append( ('all', 5))).wait() del prior_promises self.assertListEqual(value_list, [('thread_body', 1), ('thread_body', 2), ('thread_body', 3), ('thread_body', 4), ('all', 5)]) value_list = [] prior_promises = [ gen_promise(idx, bool((idx + 1) % 2)) for idx in range(4) ] promise.all_(prior_promises).then( lambda: value_list.append(('all', 5)), lambda *_: value_list.append(('all_catch', 5)), ).wait() del prior_promises expected = [('thread_body', 1), ('thread_body', 2), ('all_catch', 5)] self.assertListEqual(value_list[:len(expected)], expected) time.sleep(0.5) def _gen_all_promise(*_): prior_promises = [ gen_promise(idx, bool((idx + 1) % 2)) for idx in range(4) ] return promise.all_(prior_promises) value_list = [] gen_promise(0) \ .then(lambda *_: value_list.append(('pre_all', 0))) \ .then(_gen_all_promise) \ .then(lambda v: gen_promise(v)) \ .then( lambda: value_list.append(('all', 5)), lambda *_: value_list.append(('all_catch', 5)), ).wait() expected = [('thread_body', 1), ('pre_all', 0), ('thread_body', 1), ('thread_body', 2), ('all_catch', 5)] self.assertListEqual(value_list[:len(expected)], expected) time.sleep(0.5) finally: self.assertEqual(promise.get_active_promise_count(), 0) req_queue.put((None, None, None))
def testDiskReadAndWriteMerger(self): import logging logging.basicConfig(level=logging.DEBUG) test_addr = f'127.0.0.1:{get_next_port()}' options.worker.filemerger.max_file_size = 2400 options.worker.filemerger.concurrency = 16 with self.create_pool(n_process=1, address=test_addr) as pool, \ self.run_actor_test(pool) as test_actor: pool.create_actor(WorkerClusterInfoActor, [test_addr], uid=WorkerClusterInfoActor.default_uid()) pool.create_actor(StatusActor, test_addr, uid=StatusActor.default_uid()) pool.create_actor(EventsActor, uid=EventsActor.default_uid()) disk_file_merger_ref = pool.create_actor( DiskFileMergerActor, uid=DiskFileMergerActor.default_uid()) pool.create_actor(WorkerDaemonActor, uid=WorkerDaemonActor.default_uid()) storage_manager_ref = pool.create_actor( StorageManagerActor, uid=StorageManagerActor.default_uid()) session_id = str(uuid.uuid4()) data_count = 30 data = [ np.random.rand(random.randint(10, 30), random.randint(10, 30)) for _ in range(data_count) ] ser_data = [dataserializer.serialize(d) for d in data] storage_client = test_actor.storage_client handler = storage_client.get_storage_handler( (0, DataStorageDevice.DISK)) for handler._compress in self._get_compress_types(): data_keys = [str(uuid.uuid4()) for _ in range(data_count)] promises = [] for idx in range(data_count): block_data = dataserializer.dumps( data[idx], compress=handler._compress) def _write_data(ser, writer): with writer: writer.write(ser) return writer.filename promises.append( handler.create_bytes_writer(session_id, data_keys[idx], ser_data[idx].total_bytes, packed=True, with_merger_lock=True, _promise=True).then( functools.partial( _write_data, block_data))) promise.all_(promises).then( lambda *_: test_actor.set_result(0), lambda *exc: test_actor.set_result(exc, accept=False)) self.get_result(50) for key in data_keys: self.assertEqual( sorted( storage_manager_ref.get_data_locations( session_id, [key])[0]), [(0, DataStorageDevice.DISK)]) dump_result = disk_file_merger_ref.dump_info() written_files = list(dump_result[2]) for fn in written_files: self.assertTrue(os.path.exists(fn)) data_store = [None] * len(data) promises = [] for idx in range(data_count): def _read_data(reader, idx): with reader: data_store[idx] = dataserializer.loads( reader.read()) promises.append( handler.create_bytes_reader(session_id, data_keys[idx], with_merger_lock=True, packed=True, _promise=True).then( functools.partial( _read_data, idx=idx))) promise.all_(promises).then( lambda *_: test_actor.set_result(0), lambda *exc: test_actor.set_result(exc, accept=False)) self.get_result(50) for true_data, read_data in zip(data, data_store): assert_allclose(true_data, read_data) data_store = [None] * len(data) promises = [] for idx in range(data_count): def _read_data(reader, idx): with reader: data_store[idx] = dataserializer.deserialize( reader.read()) promises.append( handler.create_bytes_reader(session_id, data_keys[idx], _promise=True).then( functools.partial( _read_data, idx=idx))) promise.all_(promises).then( lambda *_: test_actor.set_result(0), lambda *exc: test_actor.set_result(exc, accept=False)) self.get_result(50) for true_data, read_data in zip(data, data_store): assert_allclose(true_data, read_data) storage_client.delete(session_id, data_keys) pool.sleep(0.1) for fn in written_files: self.assertFalse(os.path.exists(fn))