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))
