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