def testDiskReadAndWritePacked(self, *_):
test_addr = f'127.0.0.1:{get_next_port()}'
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())
pool.create_actor(WorkerDaemonActor,
uid=WorkerDaemonActor.default_uid())
storage_manager_ref = pool.create_actor(
StorageManagerActor, uid=StorageManagerActor.default_uid())
session_id = str(uuid.uuid4())
data1 = np.random.random((10, 10))
ser_data1 = dataserializer.serialize(data1)
storage_client = test_actor.storage_client
handler = storage_client.get_storage_handler(
(0, DataStorageDevice.DISK))
for handler._compress in self._get_compress_types():
data_key1 = str(uuid.uuid4())
storage_client.delete(session_id, [data_key1])
self.rm_spill_dirs()
block_data1 = dataserializer.dumps(data1,
compress=handler._compress)
def _write_data(ser, writer):
with writer:
writer.write(ser)
return writer.filename
handler.create_bytes_writer(session_id, data_key1, ser_data1.total_bytes,
packed=True, _promise=True) \
.then(functools.partial(_write_data, block_data1)) \
.then(test_actor.set_result,
lambda *exc: test_actor.set_result(exc, accept=False))
file_name = self.get_result(5)
self.assertEqual(
sorted(
storage_manager_ref.get_data_locations(
session_id, [data_key1])[0]),
[(0, DataStorageDevice.DISK)])
self.assertTrue(os.path.exists(file_name))
def _read_data(reader):
with reader:
return dataserializer.loads(reader.read())
handler.create_bytes_reader(session_id, data_key1, packed=True, _promise=True) \
.then(_read_data) \
.then(functools.partial(test_actor.set_result),
lambda *exc: test_actor.set_result(exc, accept=False))
assert_allclose(self.get_result(5), data1)
def testDiskLoad(self, *_):
test_addr = f'127.0.0.1:{get_next_port()}'
with self.create_pool(n_process=1, address=test_addr) as pool, \
self.run_actor_test(pool) as test_actor:
pool.create_actor(WorkerDaemonActor, uid=WorkerDaemonActor.default_uid())
storage_manager_ref = pool.create_actor(
StorageManagerActor, uid=StorageManagerActor.default_uid())
pool.create_actor(QuotaActor, 1024 ** 2, uid=MemQuotaActor.default_uid())
pool.create_actor(InProcHolderActor)
pool.create_actor(PlasmaKeyMapActor, uid=PlasmaKeyMapActor.default_uid())
pool.create_actor(SharedHolderActor, uid=SharedHolderActor.default_uid())
data1 = np.random.random((10, 10))
data2 = np.random.random((10, 10))
ser_data1 = dataserializer.serialize(data1)
session_id = str(uuid.uuid4())
data_key1 = str(uuid.uuid4())
data_key2 = str(uuid.uuid4())
storage_client = test_actor.storage_client
handler = storage_client.get_storage_handler((0, DataStorageDevice.DISK))
# load from bytes io
shared_handler = storage_client.get_storage_handler((0, DataStorageDevice.SHARED_MEMORY))
with shared_handler.create_bytes_writer(
session_id, data_key1, ser_data1.total_bytes) as writer:
ser_data1.write_to(writer)
handler.load_from_bytes_io(session_id, [data_key1], shared_handler) \
.then(lambda *_: test_actor.set_result(None),
lambda *exc: test_actor.set_result(exc, accept=False))
self.get_result(5)
self.assertEqual(sorted(storage_manager_ref.get_data_locations(session_id, [data_key1])[0]),
[(0, DataStorageDevice.SHARED_MEMORY), (0, DataStorageDevice.DISK)])
shared_handler.delete(session_id, [data_key1])
handler.delete(session_id, [data_key1])
# load from object io
ref_data2 = weakref.ref(data2)
proc_handler = storage_client.get_storage_handler((0, DataStorageDevice.PROC_MEMORY))
proc_handler.put_objects(session_id, [data_key2], [data2])
del data2
handler.load_from_object_io(session_id, [data_key2], proc_handler) \
.then(lambda *_: test_actor.set_result(None),
lambda *exc: test_actor.set_result(exc, accept=False))
self.get_result(5)
self.assertEqual(sorted(storage_manager_ref.get_data_locations(session_id, [data_key2])[0]),
[(0, DataStorageDevice.PROC_MEMORY), (0, DataStorageDevice.DISK)])
proc_handler.delete(session_id, [data_key2])
self.assertIsNone(ref_data2())
handler.delete(session_id, [data_key2])
def testSharedPutAndGet(self, *_):
test_addr = '127.0.0.1:%d' % get_next_port()
with self.create_pool(n_process=1, address=test_addr) as pool, \
self.run_actor_test(pool) as test_actor:
pool.create_actor(WorkerDaemonActor,
uid=WorkerDaemonActor.default_uid())
storage_manager_ref = pool.create_actor(
StorageManagerActor, uid=StorageManagerActor.default_uid())
pool.create_actor(PlasmaKeyMapActor,
uid=PlasmaKeyMapActor.default_uid())
pool.create_actor(SharedHolderActor,
uid=SharedHolderActor.default_uid())
data1 = np.random.random((10, 10))
data2 = np.random.random((10, 10))
ser_data2 = dataserializer.serialize(data2)
bytes_data2 = ser_data2.to_buffer()
session_id = str(uuid.uuid4())
data_key1 = str(uuid.uuid4())
data_key2 = str(uuid.uuid4())
storage_client = test_actor.storage_client
handler = storage_client.get_storage_handler(
(0, DataStorageDevice.SHARED_MEMORY))
handler.put_objects(session_id, [data_key1], [data1])
self.assertEqual(
sorted(
storage_manager_ref.get_data_locations(
session_id, [data_key1])[0]),
[(0, DataStorageDevice.SHARED_MEMORY)])
assert_allclose(data1,
handler.get_objects(session_id, [data_key1])[0])
handler.delete(session_id, [data_key1])
self.assertEqual(
list(
storage_manager_ref.get_data_locations(
session_id, [data_key1])[0]), [])
with self.assertRaises(KeyError):
handler.get_objects(session_id, [data_key1])
handler.put_objects(session_id, [data_key2], [ser_data2],
serialize=True)
assert_allclose(data2,
handler.get_objects(session_id, [data_key2])[0])
handler.delete(session_id, [data_key2])
handler.put_objects(session_id, [data_key2], [bytes_data2],
serialize=True)
assert_allclose(data2,
handler.get_objects(session_id, [data_key2])[0])
handler.delete(session_id, [data_key2])
def testLoadStoreInOtherProcess(self):
test_addr = '127.0.0.1:%d' % get_next_port()
with self.create_pool(n_process=3,
address=test_addr,
distributor=MarsDistributor(3)) as pool:
pool.create_actor(WorkerDaemonActor,
uid=WorkerDaemonActor.default_uid())
pool.create_actor(StorageManagerActor,
uid=StorageManagerActor.default_uid())
pool.create_actor(DispatchActor, uid=DispatchActor.default_uid())
pool.create_actor(QuotaActor,
1024**2,
uid=MemQuotaActor.default_uid())
pool.create_actor(PlasmaKeyMapActor,
uid=PlasmaKeyMapActor.default_uid())
pool.create_actor(SharedHolderActor,
self.plasma_storage_size,
uid=SharedHolderActor.default_uid())
pool.create_actor(InProcHolderActor, uid='w:1:InProcHolderActor1')
pool.create_actor(InProcHolderActor, uid='w:2:InProcHolderActor2')
pool.create_actor(IORunnerActor,
lock_free=True,
dispatched=False,
uid=IORunnerActor.gen_uid(1))
test_ref = pool.create_actor(OtherProcessTestActor,
uid='w:0:OtherProcTest')
def _get_result():
start_time = time.time()
while test_ref.get_result() is None:
pool.sleep(0.5)
if time.time() - start_time > 10:
raise TimeoutError
test_ref.run_copy_test((0, DataStorageDevice.SHARED_MEMORY),
(1, DataStorageDevice.PROC_MEMORY),
_tell=True)
_get_result()
test_ref.run_copy_test((1, DataStorageDevice.PROC_MEMORY),
(0, DataStorageDevice.SHARED_MEMORY),
_tell=True)
_get_result()
test_ref.run_copy_test((1, DataStorageDevice.PROC_MEMORY),
(2, DataStorageDevice.PROC_MEMORY),
_tell=True)
_get_result()
def testSharedLoadFromObjects(self, *_):
test_addr = '127.0.0.1:%d' % get_next_port()
with self.create_pool(n_process=1, address=test_addr) as pool, \
self.run_actor_test(pool) as test_actor:
pool.create_actor(WorkerDaemonActor,
uid=WorkerDaemonActor.default_uid())
storage_manager_ref = pool.create_actor(
StorageManagerActor, uid=StorageManagerActor.default_uid())
pool.create_actor(QuotaActor,
1024**2,
uid=MemQuotaActor.default_uid())
pool.create_actor(InProcHolderActor)
pool.create_actor(PlasmaKeyMapActor,
uid=PlasmaKeyMapActor.default_uid())
pool.create_actor(SharedHolderActor,
uid=SharedHolderActor.default_uid())
data1 = np.random.random((10, 10))
session_id = str(uuid.uuid4())
data_key1 = str(uuid.uuid4())
storage_client = test_actor.storage_client
handler = storage_client.get_storage_handler(
(0, DataStorageDevice.SHARED_MEMORY))
# load from object io
ref_data1 = weakref.ref(data1)
proc_handler = storage_client.get_storage_handler(
(0, DataStorageDevice.PROC_MEMORY))
proc_handler.put_objects(session_id, [data_key1], [data1])
del data1
handler.load_from_object_io(session_id, [data_key1], proc_handler) \
.then(lambda *_: test_actor.set_result(None),
lambda *exc: test_actor.set_result(exc, accept=False))
self.get_result(5)
self.assertEqual(
sorted(
storage_manager_ref.get_data_locations(
session_id, [data_key1])[0]),
[(0, DataStorageDevice.PROC_MEMORY),
(0, DataStorageDevice.SHARED_MEMORY)])
proc_handler.delete(session_id, [data_key1])
self.assertIsNone(ref_data1())
handler.delete(session_id, [data_key1])
def _start_shared_holder_pool(self):
test_addr = '127.0.0.1:%d' % get_next_port()
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(WorkerDaemonActor, uid=WorkerDaemonActor.default_uid())
pool.create_actor(StorageManagerActor, uid=StorageManagerActor.default_uid())
pool.create_actor(PlasmaKeyMapActor, uid=PlasmaKeyMapActor.default_uid())
pool.create_actor(SharedHolderActor, self.plasma_storage_size,
uid=SharedHolderActor.default_uid())
yield pool, test_actor
def testCudaMemPutAndGet(self):
test_addr = f'127.0.0.1:{get_next_port()}'
with self.create_pool(n_process=1, address=test_addr) as pool, \
self.run_actor_test(pool) as test_actor:
pool.create_actor(WorkerDaemonActor, uid=WorkerDaemonActor.default_uid())
storage_manager_ref = pool.create_actor(
StorageManagerActor, uid=StorageManagerActor.default_uid())
pool.create_actor(QuotaActor, 1024 ** 2, uid=MemQuotaActor.default_uid())
pool.create_actor(CudaHolderActor)
test_data = np.random.random((10, 10))
test_suites = [
(test_data, cp.ndarray, cp.asnumpy, assert_allclose),
(pd.Series(test_data.flatten()), cudf.Series,
lambda o: o.to_pandas(), pd.testing.assert_series_equal),
(pd.DataFrame(dict(col=test_data.flatten())), cudf.DataFrame,
lambda o: o.to_pandas(), pd.testing.assert_frame_equal),
]
for data, cuda_type, move_to_mem, assert_obj_equal in test_suites:
ser_data = dataserializer.serialize(data)
session_id = str(uuid.uuid4())
data_key1 = str(uuid.uuid4())
data_key2 = str(uuid.uuid4())
storage_client = test_actor.storage_client
handler = storage_client.get_storage_handler((0, DataStorageDevice.CUDA))
handler.put_objects(session_id, [data_key1], [data])
self.assertEqual(sorted(storage_manager_ref.get_data_locations(session_id, [data_key1])[0]),
[(0, DataStorageDevice.CUDA)])
self.assertIsInstance(handler.get_objects(session_id, [data_key1])[0], cuda_type)
assert_obj_equal(data, move_to_mem(handler.get_objects(session_id, [data_key1])[0]))
handler.delete(session_id, [data_key1])
self.assertEqual(sorted(storage_manager_ref.get_data_locations(session_id, [data_key1])[0]), [])
with self.assertRaises(KeyError):
handler.get_objects(session_id, [data_key1])
handler.put_objects(session_id, [data_key2], [ser_data], serialize=True)
self.assertIsInstance(handler.get_objects(session_id, [data_key2])[0], cuda_type)
assert_obj_equal(data, move_to_mem(handler.get_objects(session_id, [data_key2])[0]))
handler.delete(session_id, [data_key2])
def testClientSpill(self, *_):
test_addr = '127.0.0.1:%d' % get_next_port()
with self.create_pool(n_process=1, address=test_addr) as pool:
pool.create_actor(WorkerDaemonActor,
uid=WorkerDaemonActor.default_uid())
storage_manager_ref = pool.create_actor(
StorageManagerActor, uid=StorageManagerActor.default_uid())
pool.create_actor(DispatchActor, uid=DispatchActor.default_uid())
pool.create_actor(IORunnerActor)
pool.create_actor(QuotaActor,
1024**2,
uid=MemQuotaActor.default_uid())
pool.create_actor(InProcHolderActor)
pool.create_actor(PlasmaKeyMapActor,
uid=PlasmaKeyMapActor.default_uid())
pool.create_actor(SharedHolderActor,
self.plasma_storage_size,
uid=SharedHolderActor.default_uid())
session_id = str(uuid.uuid4())
data_list = [
np.random.randint(0, 32767, (655360, ), np.int16)
for _ in range(20)
]
data_keys = [str(uuid.uuid4()) for _ in range(20)]
with self.run_actor_test(pool) as test_actor:
storage_client = test_actor.storage_client
idx = 0
shared_handler = storage_client.get_storage_handler(
(0, DataStorageDevice.SHARED_MEMORY))
proc_handler = storage_client.get_storage_handler(
(0, DataStorageDevice.PROC_MEMORY))
def _fill_data():
i = 0
for i, (key,
data) in enumerate(zip(data_keys[idx:],
data_list)):
try:
shared_handler.put_objects(session_id, [key],
[data])
except StorageFull:
break
return i + idx
idx = _fill_data()
# test copying non-existing keys
storage_client.copy_to(session_id, ['non-exist-key'], [DataStorageDevice.SHARED_MEMORY]) \
.then(lambda *_: test_actor.set_result(None),
lambda *exc: test_actor.set_result(exc, accept=False))
with self.assertRaises(KeyError):
self.get_result(5)
# test copying into containing locations
storage_client.copy_to(session_id, [data_keys[0]], [DataStorageDevice.SHARED_MEMORY]) \
.then(lambda *_: test_actor.set_result(None),
lambda *exc: test_actor.set_result(exc, accept=False))
self.get_result(5)
self.assertEqual(
sorted(
storage_manager_ref.get_data_locations(
session_id, [data_keys[0]])[0]),
[(0, DataStorageDevice.SHARED_MEMORY)])
# test unsuccessful copy when no data at target
def _mock_load_from(*_, **__):
return promise.finished(*build_exc_info(SystemError),
_accept=False)
with patch_method(StorageHandler.load_from, _mock_load_from), \
self.assertRaises(SystemError):
storage_client.copy_to(session_id, [data_keys[0]], [DataStorageDevice.DISK]) \
.then(lambda *_: test_actor.set_result(None),
lambda *exc: test_actor.set_result(exc, accept=False))
self.get_result(5)
# test successful copy for multiple objects
storage_client.delete(session_id, [data_keys[idx - 1]])
ref_data = weakref.ref(data_list[idx])
ref_data2 = weakref.ref(data_list[idx + 1])
proc_handler.put_objects(session_id, data_keys[idx:idx + 2],
data_list[idx:idx + 2])
data_list[idx:idx + 2] = [None, None]
storage_client.copy_to(session_id, data_keys[idx:idx + 2],
[DataStorageDevice.SHARED_MEMORY, DataStorageDevice.DISK]) \
.then(lambda *_: test_actor.set_result(None),
lambda *exc: test_actor.set_result(exc, accept=False))
self.get_result(5)
proc_handler.delete(session_id, data_keys[idx:idx + 2])
self.assertEqual(
storage_manager_ref.get_data_locations(
session_id, data_keys[idx:idx + 2]),
[{(0, DataStorageDevice.SHARED_MEMORY)},
{(0, DataStorageDevice.DISK)}])
self.assertIsNone(ref_data())
self.assertIsNone(ref_data2())
# test copy with spill
idx += 2
proc_handler.put_objects(session_id, [data_keys[idx]],
[data_list[idx]])
storage_client.copy_to(session_id, [data_keys[idx]], [DataStorageDevice.SHARED_MEMORY]) \
.then(lambda *_: test_actor.set_result(None),
lambda *exc: test_actor.set_result(exc, accept=False))
self.get_result(5)
self.assertEqual(
sorted(
storage_manager_ref.get_data_locations(
session_id, [data_keys[idx]])[0]),
[(0, DataStorageDevice.PROC_MEMORY),
(0, DataStorageDevice.SHARED_MEMORY)])
def testClientPutAndGet(self):
test_addr = '127.0.0.1:%d' % get_next_port()
with self.create_pool(n_process=1, address=test_addr) as pool:
pool.create_actor(WorkerDaemonActor,
uid=WorkerDaemonActor.default_uid())
pool.create_actor(StorageManagerActor,
uid=StorageManagerActor.default_uid())
pool.create_actor(DispatchActor, uid=DispatchActor.default_uid())
pool.create_actor(IORunnerActor)
pool.create_actor(PlasmaKeyMapActor,
uid=PlasmaKeyMapActor.default_uid())
pool.create_actor(SharedHolderActor,
self.plasma_storage_size,
uid=SharedHolderActor.default_uid())
pool.create_actor(InProcHolderActor, uid='w:1:InProcHolderActor')
session_id = str(uuid.uuid4())
data_list = [
np.random.randint(0, 32767, (655360, ), np.int16)
for _ in range(20)
]
data_keys = [str(uuid.uuid4()) for _ in range(20)]
data_dict = dict(zip(data_keys, data_list))
with self.run_actor_test(pool) as test_actor:
storage_client = test_actor.storage_client
# check batch object put with size exceeds
storage_client.put_objects(session_id, data_keys, data_list,
[DataStorageDevice.SHARED_MEMORY, DataStorageDevice.PROC_MEMORY]) \
.then(functools.partial(test_actor.set_result),
lambda *exc: test_actor.set_result(exc, accept=False))
self.get_result(5)
locations = storage_client.get_data_locations(
session_id, data_keys)
loc_to_keys = defaultdict(list)
for key, location in zip(data_keys, locations):
self.assertEqual(len(location), 1)
loc_to_keys[list(location)[0][-1]].append(key)
self.assertGreater(
len(loc_to_keys[DataStorageDevice.PROC_MEMORY]), 1)
self.assertGreater(
len(loc_to_keys[DataStorageDevice.SHARED_MEMORY]), 1)
# check get object with all cases
with self.assertRaises(IOError):
first_shared_key = loc_to_keys[
DataStorageDevice.SHARED_MEMORY][0]
storage_client.get_object(session_id,
first_shared_key,
[DataStorageDevice.PROC_MEMORY],
_promise=False)
shared_objs = storage_client.get_objects(
session_id, [first_shared_key],
[DataStorageDevice.SHARED_MEMORY],
_promise=False)
self.assertEqual(len(shared_objs), 1)
assert_allclose(shared_objs[0], data_dict[first_shared_key])
storage_client.get_object(session_id, first_shared_key,
[DataStorageDevice.PROC_MEMORY], _promise=True) \
.then(functools.partial(test_actor.set_result),
lambda *exc: test_actor.set_result(exc, accept=False))
assert_allclose(self.get_result(5),
data_dict[first_shared_key])
storage_client.delete(session_id, data_keys)
time.sleep(0.5)
ref = weakref.ref(data_dict[data_keys[0]])
storage_client.put_objects(session_id, data_keys[:1], [ref()],
[DataStorageDevice.SHARED_MEMORY])
data_list[:] = []
data_dict.clear()
self.assertIsNone(ref())
def testWebApi(self):
service_ep = 'http://127.0.0.1:' + self.web_port
timeout = 120 if 'CI' in os.environ else -1
with new_session(service_ep) as sess:
self.assertEqual(sess.count_workers(), 1)
a = mt.ones((100, 100), chunk_size=30)
b = mt.ones((100, 100), chunk_size=30)
c = a.dot(b)
value = sess.run(c, timeout=timeout)
assert_array_equal(value, np.ones((100, 100)) * 100)
# check resubmission
value2 = sess.run(c, timeout=timeout)
assert_array_equal(value, value2)
# check when local compression libs are missing
from mars.serialize import dataserializer
try:
a = mt.ones((10, 10), chunk_size=30)
b = mt.ones((10, 10), chunk_size=30)
c = a.dot(b)
value = sess.run(c, timeout=timeout)
assert_array_equal(value, np.ones((10, 10)) * 10)
dataserializer.decompressors[dataserializer.CompressType.LZ4] = None
dataserializer.decompressobjs[dataserializer.CompressType.LZ4] = None
dataserializer.compress_openers[dataserializer.CompressType.LZ4] = None
assert_array_equal(sess.fetch(c), np.ones((10, 10)) * 10)
finally:
dataserializer.decompressors[dataserializer.CompressType.LZ4] = dataserializer.lz4_decompress
dataserializer.decompressobjs[dataserializer.CompressType.LZ4] = dataserializer.lz4_decompressobj
dataserializer.compress_openers[dataserializer.CompressType.LZ4] = dataserializer.lz4_open
# check serialization by pickle
try:
sess._sess._serial_type = SerialType.PICKLE
a = mt.ones((10, 10), chunk_size=30)
b = mt.ones((10, 10), chunk_size=30)
c = a.dot(b)
value = sess.run(c, timeout=timeout)
assert_array_equal(value, np.ones((10, 10)) * 10)
finally:
sess._sess._serial_type = SerialType.ARROW
va = np.random.randint(0, 10000, (100, 100))
vb = np.random.randint(0, 10000, (100, 100))
a = mt.array(va, chunk_size=30)
b = mt.array(vb, chunk_size=30)
c = a.dot(b)
value = sess.run(c, timeout=timeout)
assert_array_equal(value, va.dot(vb))
graphs = sess.get_graph_states()
# make sure status got uploaded
time.sleep(1.5)
# check web UI requests
res = requests.get(service_ep)
self.assertEqual(res.status_code, 200)
res = requests.get('%s/scheduler' % (service_ep,))
self.assertEqual(res.status_code, 200)
res = requests.get('%s/scheduler/127.0.0.1:%s' % (service_ep, self.scheduler_port))
self.assertEqual(res.status_code, 200)
res = requests.get('%s/worker' % (service_ep,))
self.assertEqual(res.status_code, 200)
res = requests.get('%s/worker/127.0.0.1:%s' % (service_ep, self.worker_port))
self.assertEqual(res.status_code, 200)
res = requests.get('%s/worker/127.0.0.1:%s/timeline' % (service_ep, self.worker_port))
self.assertEqual(res.status_code, 200)
res = requests.get('%s/session' % (service_ep,))
self.assertEqual(res.status_code, 200)
task_id = next(iter(graphs.keys()))
res = requests.get('%s/session/%s/graph/%s' % (service_ep, sess._session_id, task_id))
self.assertEqual(res.status_code, 200)
res = requests.get('%s/session/%s/graph/%s/running_nodes' % (service_ep, sess._session_id, task_id))
self.assertEqual(res.status_code, 200)
from mars.web.task_pages import PROGRESS_APP_NAME
res = requests.get('%s/%s?session_id=%s&task_id=%s'
% (service_ep, PROGRESS_APP_NAME, sess._session_id, task_id))
self.assertEqual(res.status_code, 200)
from mars.web.worker_pages import TIMELINE_APP_NAME
res = requests.get('%s/%s?endpoint=127.0.0.1:%s'
% (service_ep, TIMELINE_APP_NAME, self.worker_port))
self.assertEqual(res.status_code, 200)
# make sure all chunks freed when session quits
from mars.worker.storage import StorageManagerActor
actor_client = new_client()
storage_manager_ref = actor_client.actor_ref(StorageManagerActor.default_uid(),
address='127.0.0.1:' + str(self.worker_port))
self.assertSetEqual(set(storage_manager_ref.dump_keys()), set())
def testDiskReadAndWrite(self, *_):
test_addr = f'127.0.0.1:{get_next_port()}'
with self.create_pool(n_process=1, address=test_addr) as pool, \
self.run_actor_test(pool) as test_actor:
pool.create_actor(WorkerDaemonActor,
uid=WorkerDaemonActor.default_uid())
storage_manager_ref = pool.create_actor(
StorageManagerActor, uid=StorageManagerActor.default_uid())
data1 = np.random.random((10, 10))
ser_data1 = dataserializer.serialize(data1)
data2 = np.random.random((10, 10))
ser_data2 = dataserializer.serialize(data2)
session_id = str(uuid.uuid4())
storage_client = test_actor.storage_client
handler = storage_client.get_storage_handler(
(0, DataStorageDevice.DISK))
for handler._compress in self._get_compress_types():
data_key1 = str(uuid.uuid4())
data_key2 = (str(uuid.uuid4()), 'subkey')
storage_client.delete(session_id, [data_key1])
storage_client.delete(session_id, [data_key2])
self.rm_spill_dirs()
def _write_data(ser, writer):
self.assertEqual(writer.nbytes, ser.total_bytes)
with writer:
ser.write_to(writer)
return writer.filename
def _read_data(reader):
with reader:
return dataserializer.deserialize(reader.read())
# test normal file write
handler.create_bytes_writer(session_id, data_key1, ser_data1.total_bytes, _promise=True) \
.then(functools.partial(_write_data, ser_data1)) \
.then(test_actor.set_result,
lambda *exc: test_actor.set_result(exc, accept=False))
file_name = self.get_result(5)
self.assertTrue(os.path.exists(file_name))
self.assertEqual(
sorted(
storage_manager_ref.get_data_locations(
session_id, [data_key1])[0]),
[(0, DataStorageDevice.DISK)])
# test write existing (this should produce an error)
handler.create_bytes_writer(session_id, data_key1, ser_data1.total_bytes, _promise=True) \
.then(functools.partial(_write_data, ser_data1)) \
.then(test_actor.set_result,
lambda *exc: test_actor.set_result(exc, accept=False))
with self.assertRaises(StorageDataExists):
self.get_result(5)
# test writing with unreferenced file
storage_manager_ref.unregister_data(
session_id, [data_key1], (0, DataStorageDevice.DISK))
handler.create_bytes_writer(session_id, data_key1, ser_data1.total_bytes, _promise=True) \
.then(functools.partial(_write_data, ser_data1)) \
.then(test_actor.set_result,
lambda *exc: test_actor.set_result(exc, accept=False))
file_name = self.get_result(5)
self.assertTrue(os.path.exists(file_name))
self.assertEqual(
sorted(
storage_manager_ref.get_data_locations(
session_id, [data_key1])[0]),
[(0, DataStorageDevice.DISK)])
# test reading and verifying written data
handler.create_bytes_reader(session_id, data_key1, _promise=True) \
.then(_read_data) \
.then(test_actor.set_result,
lambda *exc: test_actor.set_result(exc, accept=False))
assert_allclose(self.get_result(5), data1)
# test unregistering data
handler.delete(session_id, [data_key1])
while os.path.exists(file_name):
test_actor.ctx.sleep(0.05)
self.assertFalse(os.path.exists(file_name))
# test reading and writing with tuple keys
handler.create_bytes_writer(session_id, data_key2, ser_data2.total_bytes, _promise=True) \
.then(functools.partial(_write_data, ser_data2)) \
.then(test_actor.set_result,
lambda *exc: test_actor.set_result(exc, accept=False))
self.get_result(5)
self.assertEqual(
sorted(
storage_manager_ref.get_data_locations(
session_id, [data_key2])[0]),
[(0, DataStorageDevice.DISK)])
handler.create_bytes_reader(session_id, data_key2, _promise=True) \
.then(_read_data) \
.then(functools.partial(test_actor.set_result),
lambda *exc: test_actor.set_result(exc, accept=False))
assert_allclose(self.get_result(5), data2)
def testWebApi(self):
service_ep = 'http://127.0.0.1:' + self.web_port
timeout = 120 if 'CI' in os.environ else -1
with new_session(service_ep) as sess:
session_id = sess._session_id
self.assertEqual(sess.count_workers(), 1)
a = mt.ones((100, 100), chunk_size=30)
b = mt.ones((100, 100), chunk_size=30)
c = a.dot(b)
value = sess.run(c, timeout=timeout)
np.testing.assert_array_equal(value, np.ones((100, 100)) * 100)
# check resubmission
value2 = sess.run(c, timeout=timeout)
np.testing.assert_array_equal(value, value2)
# check when local compression libs are missing
from mars.serialize import dataserializer
try:
a = mt.ones((10, 10), chunk_size=30)
b = mt.ones((10, 10), chunk_size=30)
c = a.dot(b)
value = sess.run(c, timeout=timeout)
np.testing.assert_array_equal(value, np.ones((10, 10)) * 10)
dataserializer.decompressors[
dataserializer.CompressType.LZ4] = None
dataserializer.decompressobjs[
dataserializer.CompressType.LZ4] = None
dataserializer.compress_openers[
dataserializer.CompressType.LZ4] = None
np.testing.assert_array_equal(sess.fetch(c),
np.ones((10, 10)) * 10)
finally:
dataserializer.decompressors[
dataserializer.CompressType.
LZ4] = dataserializer.lz4_decompress
dataserializer.decompressobjs[
dataserializer.CompressType.
LZ4] = dataserializer.lz4_decompressobj
dataserializer.compress_openers[
dataserializer.CompressType.LZ4] = dataserializer.lz4_open
# check serialization by pickle
try:
sess._sess._serial_type = SerialType.PICKLE
a = mt.ones((10, 10), chunk_size=30)
b = mt.ones((10, 10), chunk_size=30)
c = a.dot(b)
value = sess.run(c, timeout=timeout)
np.testing.assert_array_equal(value, np.ones((10, 10)) * 10)
raw = pd.DataFrame(np.random.rand(10, 5),
columns=list('ABCDE'),
index=pd.RangeIndex(10, 0, -1))
data = md.DataFrame(raw).astype({'E': 'arrow_string'})
ret_data = data.execute(session=sess).fetch(session=sess)
self.assertEqual(ret_data.dtypes['E'], np.dtype('O'))
pd.testing.assert_frame_equal(ret_data.astype({'E': 'float'}),
raw,
check_less_precise=True)
raw = pd.Series(np.random.rand(10),
index=pd.RangeIndex(10, 0, -1),
name='r')
data = md.Series(raw).astype('Arrow[string]')
ret_data = data.execute(session=sess).fetch(session=sess)
self.assertEqual(ret_data.dtype, np.dtype('O'))
pd.testing.assert_series_equal(ret_data.astype('float'), raw)
finally:
sess._sess._serial_type = SerialType.ARROW
va = np.random.randint(0, 10000, (100, 100))
vb = np.random.randint(0, 10000, (100, 100))
a = mt.array(va, chunk_size=30)
b = mt.array(vb, chunk_size=30)
c = a.dot(b)
value = sess.run(c, timeout=timeout)
np.testing.assert_array_equal(value, va.dot(vb))
# test fetch log
def f():
print('test')
r = mr.spawn(f).execute(session=sess, timeout=timeout)
self.assertEqual(str(r.fetch_log()).strip(), 'test')
self.assertEqual(str(r.fetch_log(offsets=0)).strip(), 'test')
self.assertEqual(str(r.fetch_log()).strip(), '')
self.assertEqual(
str(r.fetch_log(offsets='-0.003k', sizes=2)).strip(), 'st')
graphs = sess.get_graph_states()
# make sure status got uploaded
time.sleep(1.5)
# check web UI requests
res = requests.get(service_ep)
self.assertEqual(res.status_code, 200)
res = requests.get(f'{service_ep}/scheduler')
self.assertEqual(res.status_code, 200)
res = requests.get(
f'{service_ep}/scheduler/127.0.0.1:{self.scheduler_port}')
self.assertEqual(res.status_code, 200)
res = requests.get(f'{service_ep}/worker')
self.assertEqual(res.status_code, 200)
res = requests.get(
f'{service_ep}/worker/127.0.0.1:{self.worker_port}')
self.assertEqual(res.status_code, 200)
res = requests.get(
f'{service_ep}/worker/127.0.0.1:{self.worker_port}/timeline')
self.assertEqual(res.status_code, 200)
res = requests.get(f'{service_ep}/session')
self.assertEqual(res.status_code, 200)
task_id = next(iter(graphs.keys()))
res = requests.get(
f'{service_ep}/session/{session_id}/graph/{task_id}')
self.assertEqual(res.status_code, 200)
res = requests.get(
f'{service_ep}/session/{session_id}/graph/{task_id}/running_nodes'
)
self.assertEqual(res.status_code, 200)
from mars.web.task_pages import PROGRESS_APP_NAME
res = requests.get(
f'{service_ep}/{PROGRESS_APP_NAME}?session_id={session_id}&task_id={task_id}'
)
self.assertEqual(res.status_code, 200)
from mars.web.worker_pages import TIMELINE_APP_NAME
res = requests.get(
f'{service_ep}/{TIMELINE_APP_NAME}?endpoint=127.0.0.1:{self.worker_port}'
)
self.assertEqual(res.status_code, 200)
# make sure all chunks freed when session quits
from mars.worker.storage import StorageManagerActor
actor_client = new_client()
storage_manager_ref = actor_client.actor_ref(
StorageManagerActor.default_uid(),
address='127.0.0.1:' + str(self.worker_port))
self.assertSetEqual(set(storage_manager_ref.dump_keys()), set())
def testSharedHolderSpill(self):
with self._start_shared_holder_pool() as (pool, test_actor):
pool.create_actor(DispatchActor, uid=DispatchActor.default_uid())
pool.create_actor(MockIORunnerActor,
uid=MockIORunnerActor.default_uid())
manager_ref = pool.actor_ref(StorageManagerActor.default_uid())
shared_holder_ref = pool.actor_ref(SharedHolderActor.default_uid())
mock_runner_ref = pool.actor_ref(MockIORunnerActor.default_uid())
status_ref = pool.actor_ref(StatusActor.default_uid())
storage_client = test_actor.storage_client
shared_handler = storage_client.get_storage_handler(
(0, DataStorageDevice.SHARED_MEMORY))
cache_allocations = status_ref.get_cache_allocations()
self.assertGreater(cache_allocations['total'], 0)
session_id = str(uuid.uuid4())
data_list = [
np.random.randint(0, 32767, (655360, ), np.int16)
for _ in range(20)
]
key_list = [str(uuid.uuid4()) for _ in range(20)]
self._fill_shared_storage(session_id, key_list, data_list)
data_size = manager_ref.get_data_sizes(session_id,
[key_list[0]])[0]
# spill huge sizes
with self.assertRaises(SpillSizeExceeded):
self.waitp(
shared_handler.spill_size(self.plasma_storage_size * 2), )
# spill size of two data chunks
keys_before = [tp[1] for tp in shared_holder_ref.dump_keys()]
pin_token = str(uuid.uuid4())
shared_holder_ref.pin_data_keys(session_id, key_list[1:2],
pin_token)
expect_spills = key_list[2:4]
shared_holder_ref.lift_data_keys(session_id, [key_list[0]])
shared_handler.spill_size(data_size * 2) \
.then(lambda *_: test_actor.set_result(None),
lambda *exc: test_actor.set_result(exc, accept=False))
pool.sleep(0.5)
# when the key is in spill (here we trigger it manually in mock),
# it cannot be spilled
with self.assertRaises(PinDataKeyFailed):
shared_holder_ref.pin_data_keys(session_id, key_list[2:3],
str(uuid.uuid4()))
for k in key_list[2:6]:
mock_runner_ref.submit_item(session_id, k)
self.get_result(5)
shared_holder_ref.unpin_data_keys(session_id, key_list[1:2],
pin_token)
keys_after = [tp[1] for tp in shared_holder_ref.dump_keys()]
self.assertSetEqual(
set(keys_before) - set(keys_after), set(expect_spills))
# spill size of a single chunk, should return immediately
keys_before = [tp[1] for tp in shared_holder_ref.dump_keys()]
shared_handler.spill_size(data_size) \
.then(lambda *_: test_actor.set_result(None),
lambda *exc: test_actor.set_result(exc, accept=False))
self.get_result(5)
keys_after = [tp[1] for tp in shared_holder_ref.dump_keys()]
self.assertSetEqual(set(keys_before), set(keys_after))
# when all pinned, nothing can be spilled
# and spill_size() should raises an error
pin_token = str(uuid.uuid4())
shared_holder_ref.pin_data_keys(session_id, key_list, pin_token)
shared_handler.spill_size(data_size * 3) \
.then(lambda *_: test_actor.set_result(None),
lambda *exc: test_actor.set_result(exc, accept=False))
with self.assertRaises(NoDataToSpill):
self.get_result(5)
shared_holder_ref.unpin_data_keys(session_id, key_list, pin_token)
# when some errors raise when spilling,
# spill_size() should report it
mock_runner_ref.clear_submissions()
shared_handler.spill_size(data_size * 3) \
.then(lambda *_: test_actor.set_result(None),
lambda *exc: test_actor.set_result(exc, accept=False))
pool.sleep(0.5)
spill_keys = mock_runner_ref.get_request_keys()
mock_runner_ref.submit_item(session_id, spill_keys[0],
build_exc_info(SystemError))
for k in spill_keys[1:]:
mock_runner_ref.submit_item(session_id, k)
with self.assertRaises(SystemError):
self.get_result(5)
def testSharedReadAndWrite(self, *_):
test_addr = '127.0.0.1:%d' % get_next_port()
io_size = dataserializer.HEADER_LENGTH * 2
with self.create_pool(n_process=1, address=test_addr) as pool, \
self.run_actor_test(pool) as test_actor:
pool.create_actor(WorkerDaemonActor, uid=WorkerDaemonActor.default_uid())
storage_manager_ref = pool.create_actor(
StorageManagerActor, uid=StorageManagerActor.default_uid())
pool.create_actor(PlasmaKeyMapActor, uid=PlasmaKeyMapActor.default_uid())
pool.create_actor(SharedHolderActor, uid=SharedHolderActor.default_uid())
data1 = np.random.random((100, 100))
ser_data1 = dataserializer.serialize(data1)
session_id = str(uuid.uuid4())
data_key1 = str(uuid.uuid4())
storage_client = test_actor.storage_client
handler = storage_client.get_storage_handler((0, DataStorageDevice.SHARED_MEMORY))
def _write_data(ser, writer):
self.assertEqual(writer.nbytes, ser_data1.total_bytes)
with writer:
ser.write_to(writer)
handler.create_bytes_writer(session_id, data_key1, ser_data1.total_bytes, _promise=True) \
.then(functools.partial(_write_data, ser_data1)) \
.then(lambda *_: test_actor.set_result(None),
lambda *exc: test_actor.set_result(exc, accept=False))
self.get_result(5)
self.assertEqual(sorted(storage_manager_ref.get_data_locations(session_id, [data_key1])[0]),
[(0, DataStorageDevice.SHARED_MEMORY)])
handler.delete(session_id, [data_key1])
def _write_data(ser, writer):
with writer:
for start in range(0, len(ser), io_size):
writer.write(ser[start:start + io_size])
handler.create_bytes_writer(session_id, data_key1, ser_data1.total_bytes, _promise=True) \
.then(functools.partial(_write_data, ser_data1.to_buffer())) \
.then(lambda *_: test_actor.set_result(None),
lambda *exc: test_actor.set_result(exc, accept=False))
self.get_result(5)
self.assertEqual(sorted(storage_manager_ref.get_data_locations(session_id, [data_key1])[0]),
[(0, DataStorageDevice.SHARED_MEMORY)])
def _read_data_all(reader):
with reader:
return dataserializer.deserialize(reader.read())
handler.create_bytes_reader(session_id, data_key1, _promise=True) \
.then(_read_data_all) \
.then(functools.partial(test_actor.set_result),
lambda *exc: test_actor.set_result(exc, accept=False))
assert_allclose(self.get_result(5), data1)
def _read_data_batch(reader):
bio = BytesIO()
with reader:
while True:
buf = reader.read(io_size)
if buf:
bio.write(buf)
else:
break
return dataserializer.deserialize(bio.getvalue())
handler.create_bytes_reader(session_id, data_key1, _promise=True) \
.then(_read_data_batch) \
.then(functools.partial(test_actor.set_result),
lambda *exc: test_actor.set_result(exc, accept=False))
assert_allclose(self.get_result(5), data1)
handler.delete(session_id, [data_key1])
def testSharedSpill(self, *_):
test_addr = '127.0.0.1:%d' % get_next_port()
with self.create_pool(n_process=1, address=test_addr) as pool, \
self.run_actor_test(pool) as test_actor:
pool.create_actor(WorkerDaemonActor, uid=WorkerDaemonActor.default_uid())
storage_manager_ref = pool.create_actor(
StorageManagerActor, uid=StorageManagerActor.default_uid())
pool.create_actor(PlasmaKeyMapActor, uid=PlasmaKeyMapActor.default_uid())
holder_ref = pool.create_actor(
SharedHolderActor, self.plasma_storage_size,
uid=SharedHolderActor.default_uid())
session_id = str(uuid.uuid4())
data_list = [np.random.randint(0, 32767, (655360,), np.int16)
for _ in range(20)]
data_keys = [str(uuid.uuid4()) for _ in range(20)]
storage_client = test_actor.storage_client
handler = storage_client.get_storage_handler((0, DataStorageDevice.SHARED_MEMORY))
idx = 0
def _fill_data():
i = 0
for i, (key, data) in enumerate(zip(data_keys[idx:], data_list)):
try:
handler.put_objects(session_id, [key], [data])
except StorageFull:
break
return i + idx
def _do_spill():
data_size = storage_manager_ref.get_data_sizes(session_id, [data_keys[0]])[0]
handler.spill_size(2 * data_size) \
.then(lambda *_: test_actor.set_result(None),
lambda *exc: test_actor.set_result(exc, accept=False))
self.get_result(5)
# test lift data key
idx = _fill_data()
handler.lift_data_keys(session_id, [data_keys[0]])
_do_spill()
self.assertEqual(list(storage_manager_ref.get_data_locations(session_id, [data_keys[0]])[0]),
[(0, DataStorageDevice.SHARED_MEMORY)])
self.assertEqual(list(storage_manager_ref.get_data_locations(session_id, [data_keys[1]])[0]),
[(0, DataStorageDevice.DISK)])
handler.put_objects(session_id, [data_keys[idx]], [data_list[idx]])
self.assertEqual(list(storage_manager_ref.get_data_locations(session_id, [data_keys[idx]])[0]),
[(0, DataStorageDevice.SHARED_MEMORY)])
idx += 1
# test pin data key
idx = _fill_data()
holder_ref.lift_data_keys(session_id, [data_keys[0]], last=False)
pin_token = str(uuid.uuid4())
pinned_keys = handler.pin_data_keys(session_id, (data_keys[0],), pin_token)
self.assertIn(data_keys[0], pinned_keys)
_do_spill()
self.assertEqual(list(storage_manager_ref.get_data_locations(session_id, [data_keys[0]])[0]),
[(0, DataStorageDevice.SHARED_MEMORY)])
self.assertEqual(list(storage_manager_ref.get_data_locations(session_id, [data_keys[1]])[0]),
[(0, DataStorageDevice.DISK)])
handler.put_objects(session_id, [data_keys[idx]], [data_list[idx]])
self.assertEqual(list(storage_manager_ref.get_data_locations(session_id, [data_keys[idx]])[0]),
[(0, DataStorageDevice.SHARED_MEMORY)])
idx += 1
# test unpin data key
idx = _fill_data()
handler.unpin_data_keys(session_id, (data_keys[0],), pin_token)
_do_spill()
self.assertEqual(list(storage_manager_ref.get_data_locations(session_id, [data_keys[0]])[0]),
[(0, DataStorageDevice.DISK)])
def testSharedLoadFromBytes(self, *_):
import logging
logging.basicConfig(level=logging.DEBUG)
test_addr = '127.0.0.1:%d' % get_next_port()
with self.create_pool(n_process=1, address=test_addr) as pool, \
self.run_actor_test(pool) as test_actor:
pool.create_actor(WorkerDaemonActor, uid=WorkerDaemonActor.default_uid())
storage_manager_ref = pool.create_actor(
StorageManagerActor, uid=StorageManagerActor.default_uid())
pool.create_actor(QuotaActor, 1024 ** 2, uid=MemQuotaActor.default_uid())
pool.create_actor(InProcHolderActor)
pool.create_actor(PlasmaKeyMapActor, uid=PlasmaKeyMapActor.default_uid())
pool.create_actor(SharedHolderActor, uid=SharedHolderActor.default_uid())
data1 = np.random.random((10, 10))
ser_data1 = dataserializer.serialize(data1)
session_id = str(uuid.uuid4())
data_key1 = str(uuid.uuid4())
storage_client = test_actor.storage_client
handler = storage_client.get_storage_handler((0, DataStorageDevice.SHARED_MEMORY))
# load from bytes io
disk_handler = storage_client.get_storage_handler((0, DataStorageDevice.DISK))
with disk_handler.create_bytes_writer(
session_id, data_key1, ser_data1.total_bytes) as writer:
ser_data1.write_to(writer)
handler.load_from_bytes_io(session_id, [data_key1], disk_handler) \
.then(lambda *_: test_actor.set_result(None),
lambda *exc: test_actor.set_result(exc, accept=False))
self.get_result(5)
self.assertEqual(sorted(storage_manager_ref.get_data_locations(session_id, [data_key1])[0]),
[(0, DataStorageDevice.SHARED_MEMORY), (0, DataStorageDevice.DISK)])
disk_handler.delete(session_id, [data_key1])
handler.delete(session_id, [data_key1])
# load from bytes io till no capacity
data_list = [np.random.randint(0, 32767, (655360,), np.int16)
for _ in range(20)]
data_keys = [str(uuid.uuid4()) for _ in range(20)]
for key, data in zip(data_keys, data_list):
ser_data = dataserializer.serialize(data)
with disk_handler.create_bytes_writer(
session_id, key, ser_data.total_bytes) as writer:
ser_data.write_to(writer)
handler.load_from_bytes_io(session_id, data_keys, disk_handler) \
.then(lambda *_: test_actor.set_result(None),
lambda *exc: test_actor.set_result(exc, accept=False))
affected_keys = set()
try:
self.get_result(5)
except StorageFull as ex:
affected_keys.update(ex.affected_keys)
storage_client.delete(session_id, data_keys, [DataStorageDevice.DISK])
self.assertLess(len(affected_keys), len(data_keys))
self.assertGreater(len(affected_keys), 1)
for k, size in zip(data_keys, storage_client.get_data_sizes(session_id, data_keys)):
if k in affected_keys:
self.assertIsNone(size)
else:
self.assertIsNotNone(size)
def post_create(self):
super().post_create()
self._manager_ref = self.ctx.actor_ref(
StorageManagerActor.default_uid())
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))
def testClientReadAndWrite(self):
test_addr = '127.0.0.1:%d' % get_next_port()
with self.create_pool(n_process=1, address=test_addr) as pool:
options.worker.lock_free_fileio = True
pool.create_actor(WorkerDaemonActor,
uid=WorkerDaemonActor.default_uid())
pool.create_actor(StorageManagerActor,
uid=StorageManagerActor.default_uid())
pool.create_actor(DispatchActor, uid=DispatchActor.default_uid())
pool.create_actor(IORunnerActor)
pool.create_actor(PlasmaKeyMapActor,
uid=PlasmaKeyMapActor.default_uid())
pool.create_actor(SharedHolderActor,
self.plasma_storage_size,
uid=SharedHolderActor.default_uid())
data1 = np.random.random((10, 10))
ser_data1 = dataserializer.serialize(data1)
session_id = str(uuid.uuid4())
data_key1 = str(uuid.uuid4())
data_key2 = str(uuid.uuid4())
with self.run_actor_test(pool) as test_actor:
storage_client = test_actor.storage_client
file_names = []
def _write_data(ser, writer):
file_names.append(writer.filename)
self.assertEqual(writer.nbytes, ser_data1.total_bytes)
with writer:
ser.write_to(writer)
# test creating non-promised writer and write
with storage_client.create_writer(session_id,
data_key1,
ser_data1.total_bytes,
(DataStorageDevice.DISK, ),
_promise=False) as writer:
_write_data(ser_data1, writer)
self.assertTrue(os.path.exists(file_names[0]))
self.assertEqual(
sorted(
storage_client.get_data_locations(
session_id, [data_key1])[0]),
[(0, DataStorageDevice.DISK)])
storage_client.delete(session_id, [data_key1])
# test creating promised writer and write
file_names[:] = []
self.waitp(
storage_client.create_writer(
session_id, data_key2, ser_data1.total_bytes,
(DataStorageDevice.DISK, )).then(
functools.partial(_write_data, ser_data1)))
self.assertTrue(os.path.exists(file_names[0]))
self.assertEqual(
sorted(
storage_client.get_data_locations(
session_id, [data_key2])[0]),
[(0, DataStorageDevice.DISK)])
def _read_data(reader):
with reader:
return dataserializer.deserialize(reader.read())
# test creating reader when data exist in location
result = self.waitp(
storage_client.create_reader(
session_id, data_key2,
(DataStorageDevice.DISK, )).then(_read_data))[0]
assert_allclose(result, data1)
# test creating reader when no data in location (should raise)
with self.assertRaises(IOError):
storage_client.create_reader(
session_id,
data_key2, (DataStorageDevice.SHARED_MEMORY, ),
_promise=False)
# test creating reader when copy needed
self.waitp(
storage_client.create_reader(
session_id, data_key2,
(DataStorageDevice.SHARED_MEMORY, )).then(_read_data))
self.assertEqual(
sorted(
storage_client.get_data_locations(
session_id, [data_key2])[0]),
[(0, DataStorageDevice.SHARED_MEMORY),
(0, DataStorageDevice.DISK)])
storage_client.delete(session_id, [data_key2])
while os.path.exists(file_names[0]):
test_actor.ctx.sleep(0.05)
self.assertFalse(os.path.exists(file_names[0]))
def testStorageManager(self):
test_addr = '127.0.0.1:%d' % get_next_port()
with self.create_pool(n_process=1, address=test_addr) as pool:
session_id = str(uuid.uuid4())
data_key1 = str(uuid.uuid4())
data_key2 = str(uuid.uuid4())
pool.create_actor(WorkerDaemonActor,
uid=WorkerDaemonActor.default_uid())
manager_ref = pool.create_actor(
StorageManagerActor, uid=StorageManagerActor.default_uid())
self.assertEqual(
list(
manager_ref.get_data_locations(session_id,
['NON_EXIST'])[0]), [])
manager_ref.register_data(session_id, [data_key1],
(0, DataStorageDevice.SHARED_MEMORY),
[1024],
shapes=[(16, 8)])
manager_ref.register_data(session_id, [data_key1],
(1, DataStorageDevice.PROC_MEMORY),
[1024],
shapes=[(16, 8)])
manager_ref.register_data(session_id, [data_key1],
(0, DataStorageDevice.DISK), [2048])
self.assertEqual([(0, DataStorageDevice.SHARED_MEMORY),
(0, DataStorageDevice.DISK),
(1, DataStorageDevice.PROC_MEMORY)],
sorted(
manager_ref.get_data_locations(
session_id, [data_key1])[0]))
self.assertEqual(
2048,
manager_ref.get_data_sizes(session_id, [data_key1])[0])
self.assertEqual(
(16, 8),
manager_ref.get_data_shapes(session_id, [data_key1])[0])
manager_ref.register_data(session_id, [data_key2],
(0, DataStorageDevice.SHARED_MEMORY),
[1024])
manager_ref.register_data(session_id, [data_key2],
(1, DataStorageDevice.PROC_MEMORY),
[1024])
self.assertEqual([(0, DataStorageDevice.SHARED_MEMORY),
(1, DataStorageDevice.PROC_MEMORY)],
sorted(
manager_ref.get_data_locations(
session_id, [data_key2])[0]))
manager_ref.unregister_data(session_id, [data_key2],
(0, DataStorageDevice.SHARED_MEMORY))
self.assertEqual([(1, DataStorageDevice.PROC_MEMORY)],
sorted(
manager_ref.get_data_locations(
session_id, [data_key2])[0]))
self.assertEqual(
1024,
manager_ref.get_data_sizes(session_id, [data_key2])[0])
self.assertEqual([data_key1],
list(
manager_ref.filter_exist_keys(
session_id,
[data_key1, data_key2, 'non-exist'],
[(0, DataStorageDevice.SHARED_MEMORY)])))
manager_ref.unregister_data(session_id, [data_key2],
(1, DataStorageDevice.PROC_MEMORY))
self.assertEqual(
list(
manager_ref.get_data_locations(session_id,
[data_key2])[0]), [])
self.assertIsNone(
manager_ref.get_data_sizes(session_id, [data_key2])[0])
self.assertIsNone(
manager_ref.get_data_shapes(session_id, data_key2)[0])
manager_ref.register_data(session_id, [data_key2],
(1, DataStorageDevice.PROC_MEMORY),
[1024])
manager_ref.handle_process_down([1])
self.assertEqual([(0, DataStorageDevice.SHARED_MEMORY),
(0, DataStorageDevice.DISK)],
sorted(
manager_ref.get_data_locations(
session_id, [data_key1])[0]))
self.assertEqual(
list(
manager_ref.get_data_locations(session_id,
[data_key2])[0]), [])
self.assertIsNone(
manager_ref.get_data_sizes(session_id, [data_key2])[0])
self.assertIsNone(
manager_ref.get_data_shapes(session_id, [data_key2])[0])