def testHolder(self):
pool_address = '127.0.0.1:%d' % get_next_port()
with create_actor_pool(n_process=1,
backend='gevent',
address=pool_address) as pool:
pool.create_actor(PlasmaKeyMapActor,
uid=PlasmaKeyMapActor.default_name())
pool.create_actor(WorkerClusterInfoActor,
schedulers=[pool_address],
uid=WorkerClusterInfoActor.default_name())
pool.create_actor(KVStoreActor, uid=KVStoreActor.default_name())
pool.create_actor(DispatchActor, uid=DispatchActor.default_name())
pool.create_actor(QuotaActor,
1024 * 1024 * 10,
uid=MemQuotaActor.default_name())
cache_ref = pool.create_actor(ChunkHolderActor,
self.plasma_storage_size,
uid=ChunkHolderActor.default_name())
pool.create_actor(SpillActor)
try:
test_ref = pool.create_actor(CacheTestActor)
test_ref.run_test_cache()
while not test_ref.get_exc_info()[0]:
pool.sleep(0.1)
exc_info = test_ref.get_exc_info()[1]
if exc_info:
six.reraise(*exc_info)
finally:
pool.destroy_actor(cache_ref)
def testEnsureTimeout(self, *_):
from mars.errors import PromiseTimeout
pool_address = '127.0.0.1:%d' % get_next_port()
with create_actor_pool(n_process=1,
backend='gevent',
address=pool_address) as pool:
pool.create_actor(PlasmaKeyMapActor,
uid=PlasmaKeyMapActor.default_name())
pool.create_actor(WorkerClusterInfoActor,
schedulers=[pool_address],
uid=WorkerClusterInfoActor.default_name())
pool.create_actor(KVStoreActor, uid=KVStoreActor.default_name())
pool.create_actor(DispatchActor, uid=DispatchActor.default_name())
pool.create_actor(QuotaActor,
1024 * 1024 * 10,
uid=MemQuotaActor.default_name())
pool.create_actor(SpillActor, uid=SpillActor.default_name())
cache_ref = pool.create_actor(ChunkHolderActor,
self.plasma_storage_size,
uid=ChunkHolderActor.default_name())
try:
options.worker.prepare_data_timeout = 2
test_ref = pool.create_actor(CacheTestActor)
test_ref.run_test_ensure_timeout()
while not test_ref.get_exc_info()[0]:
pool.sleep(0.1)
exc_info = test_ref.get_exc_info()[1]
self.assertIsNotNone(exc_info)
self.assertIsInstance(exc_info[1], PromiseTimeout)
finally:
options.worker.prepare_data_timeout = 600
pool.destroy_actor(cache_ref)
def start_transfer_test_pool(**kwargs):
address = kwargs.pop('address')
plasma_size = kwargs.pop('plasma_size')
with create_actor_pool(n_process=1,
backend='gevent',
address=address,
**kwargs) as pool:
pool.create_actor(SchedulerClusterInfoActor,
schedulers=[address],
uid=SchedulerClusterInfoActor.default_name())
pool.create_actor(WorkerClusterInfoActor,
schedulers=[address],
uid=WorkerClusterInfoActor.default_name())
pool.create_actor(PlasmaKeyMapActor,
uid=PlasmaKeyMapActor.default_name())
pool.create_actor(ChunkMetaActor, uid=ChunkMetaActor.default_name())
pool.create_actor(DispatchActor, uid=DispatchActor.default_name())
pool.create_actor(QuotaActor,
1024 * 1024 * 20,
uid=MemQuotaActor.default_name())
chunk_holder_ref = pool.create_actor(
ChunkHolderActor, plasma_size, uid=ChunkHolderActor.default_name())
pool.create_actor(SpillActor)
pool.create_actor(StatusActor, address, uid=StatusActor.default_name())
yield pool
chunk_holder_ref.destroy()
def testStatus(self):
pool_address = '127.0.0.1:%d' % get_next_port()
old_spill_dir = options.worker.spill_directory
dir_name = options.worker.spill_directory = tempfile.mkdtemp(prefix='temp-mars-spill-')
try:
with create_actor_pool(n_process=1, backend='gevent', address=pool_address) as pool:
pool.create_actor(SchedulerClusterInfoActor, schedulers=[pool_address],
uid=SchedulerClusterInfoActor.default_name())
pool.create_actor(WorkerClusterInfoActor, schedulers=[pool_address],
uid=WorkerClusterInfoActor.default_name())
resource_ref = pool.create_actor(ResourceActor, uid=ResourceActor.default_name())
pool.create_actor(ChunkHolderActor, self.plasma_storage_size,
uid=ChunkHolderActor.default_name())
status_ref = pool.create_actor(StatusActor, pool_address,
uid=StatusActor.default_name())
status_ref.enable_status_upload()
status_ref.update_slots(dict(cpu=4))
status_ref.update_stats(dict(min_est_finish_time=10))
def delay_read():
gevent.sleep(1.5)
return resource_ref.get_workers_meta()
gl = gevent.spawn(delay_read)
gl.join()
v = gl.value
self.assertIsNotNone(v)
pool.destroy_actor(status_ref)
finally:
options.worker.spill_directory = old_spill_dir
shutil.rmtree(dir_name)
def testPrepareQuota(self, *_):
pinned = [True]
def _mock_pin(graph_key, chunk_keys):
from mars.errors import PinChunkFailed
if pinned[0]:
raise PinChunkFailed
return chunk_keys
ChunkHolderActor.pin_chunks.side_effect = _mock_pin
pool_address = '127.0.0.1:%d' % get_next_port()
session_id = str(uuid.uuid4())
mock_data = np.array([1, 2, 3, 4])
with create_actor_pool(n_process=1, backend='gevent', address=pool_address) as pool:
self.create_standard_actors(pool, pool_address, with_daemon=False, with_status=False)
pool.create_actor(MockSenderActor, mock_data, 'in', uid='w:mock_sender')
cluster_info_ref = pool.actor_ref(WorkerClusterInfoActor.default_name())
chunk_meta_client = ChunkMetaClient(pool, cluster_info_ref)
import mars.tensor as mt
from mars.tensor.expressions.fetch import TensorFetch
arr = mt.ones((4,), chunk_size=4)
arr_add = mt.array(mock_data)
result_tensor = arr + arr_add
graph = result_tensor.build_graph(compose=False, tiled=True)
modified_chunk = arr_add.chunks[0]
arr_add.chunks[0]._op = TensorFetch(
dtype=modified_chunk.dtype, _outputs=[weakref.ref(o) for o in modified_chunk.op.outputs],
_key=modified_chunk.op.key)
chunk_meta_client.set_chunk_meta(session_id, modified_chunk.key, size=mock_data.nbytes,
shape=mock_data.shape, workers=('0.0.0.0:1234', pool_address))
with self.run_actor_test(pool) as test_actor:
graph_key = str(uuid.uuid4())
execution_ref = test_actor.promise_ref(ExecutionActor.default_name())
start_time = time.time()
execution_ref.enqueue_graph(
session_id, graph_key, serialize_graph(graph),
dict(chunks=[result_tensor.chunks[0].key]), None, _promise=True) \
.then(lambda *_: test_actor.set_result(time.time())) \
.catch(lambda *exc: test_actor.set_result(exc, False))
def _delay_fun():
time.sleep(1)
pinned[0] = False
threading.Thread(target=_delay_fun).start()
finish_time = self.get_result()
self.assertGreaterEqual(finish_time, start_time + 1)
def testMemQuotaAllocation(self):
from mars import resource
from mars.utils import AttributeDict
mock_mem_stat = AttributeDict(
dict(total=300, available=50, used=0, free=50))
local_pool_addr = 'localhost:%d' % get_next_port()
with create_actor_pool(n_process=1, backend='gevent', address=local_pool_addr) as pool, \
patch_method(resource.virtual_memory, new=lambda: mock_mem_stat):
pool.create_actor(WorkerClusterInfoActor,
schedulers=[local_pool_addr],
uid=WorkerClusterInfoActor.default_name())
pool.create_actor(StatusActor,
local_pool_addr,
uid=StatusActor.default_name())
pool.create_actor(DispatchActor, uid=DispatchActor.default_name())
pool.create_actor(ProcessHelperActor,
uid=ProcessHelperActor.default_name())
quota_ref = pool.create_actor(MemQuotaActor,
300,
refresh_time=0.1,
uid=MemQuotaActor.default_name())
time_recs = []
with self.run_actor_test(pool) as test_actor:
ref = test_actor.promise_ref(quota_ref)
time_recs.append(time.time())
def actual_exec(x):
ref.release_quota(x)
time_recs.append(time.time())
test_actor.set_result(None)
ref.request_quota('req', 100, _promise=True) \
.then(functools.partial(actual_exec, 'req'))
pool.sleep(0.5)
mock_mem_stat['available'] = 150
mock_mem_stat['free'] = 150
self.get_result(2)
self.assertGreater(abs(time_recs[0] - time_recs[1]), 0.4)
def create_standard_actors(cls, pool, address, quota_size=None, with_daemon=True,
with_status=True, with_resource=False):
quota_size = quota_size or (1024 * 1024)
pool.create_actor(SchedulerClusterInfoActor, schedulers=[address],
uid=SchedulerClusterInfoActor.default_name())
pool.create_actor(WorkerClusterInfoActor, schedulers=[address],
uid=WorkerClusterInfoActor.default_name())
pool.create_actor(PlasmaKeyMapActor, uid=PlasmaKeyMapActor.default_name())
if with_resource:
pool.create_actor(ResourceActor, uid=ResourceActor.default_name())
if with_daemon:
pool.create_actor(WorkerDaemonActor, uid=WorkerDaemonActor.default_name())
if with_status:
pool.create_actor(StatusActor, address, uid=StatusActor.default_name())
pool.create_actor(
ChunkHolderActor, cls.plasma_storage_size, uid=ChunkHolderActor.default_name())
pool.create_actor(ChunkMetaActor, uid=ChunkMetaActor.default_name())
pool.create_actor(TaskQueueActor, uid=TaskQueueActor.default_name())
pool.create_actor(DispatchActor, uid=DispatchActor.default_name())
pool.create_actor(QuotaActor, quota_size, uid=MemQuotaActor.default_name())
pool.create_actor(ExecutionActor, uid=ExecutionActor.default_name())
def testQuota(self):
local_pool_addr = 'localhost:%d' % get_next_port()
with create_actor_pool(n_process=1,
backend='gevent',
address=local_pool_addr) as pool:
pool.create_actor(WorkerClusterInfoActor,
schedulers=[local_pool_addr],
uid=WorkerClusterInfoActor.default_name())
pool.create_actor(StatusActor,
local_pool_addr,
uid=StatusActor.default_name())
quota_ref = pool.create_actor(QuotaActor,
300,
uid=QuotaActor.default_name())
quota_ref.process_quota('non_exist')
quota_ref.hold_quota('non_exist')
quota_ref.release_quota('non_exist')
with self.assertRaises(ValueError):
quota_ref.request_quota('ERROR', 1000)
self.assertTrue(quota_ref.request_quota('0', 100))
self.assertTrue(quota_ref.request_quota('0', 50))
self.assertTrue(quota_ref.request_quota('0', 200))
quota_ref.process_quota('0')
self.assertIn('0', quota_ref.dump_data().proc_sizes)
quota_ref.alter_allocation('0', 190, new_key=('0', 0))
self.assertEqual(quota_ref.dump_data().allocations[('0', 0)], 190)
quota_ref.hold_quota(('0', 0))
self.assertIn(('0', 0), quota_ref.dump_data().hold_sizes)
quota_ref.alter_allocation(('0', 0), new_key=('0', 1))
self.assertEqual(quota_ref.dump_data().allocations[('0', 1)], 190)
with self.run_actor_test(pool) as test_actor:
ref = test_actor.promise_ref(QuotaActor.default_name())
ref.request_quota('1', 150, _promise=True) \
.catch(lambda *exc: test_actor.set_result(exc, accept=False))
self.assertFalse(quota_ref.request_quota('2', 50))
self.assertFalse(quota_ref.request_quota('3', 200))
self.assertFalse(quota_ref.request_quota('3', 180))
self.assertNotIn('2', quota_ref.dump_data().allocations)
ref.cancel_requests(('1', ),
reject_exc=build_exc_info(ValueError))
with self.assertRaises(ValueError):
self.get_result(5)
self.assertNotIn('1', quota_ref.dump_data().requests)
self.assertIn('2', quota_ref.dump_data().allocations)
self.assertNotIn('3', quota_ref.dump_data().allocations)
quota_ref.release_quotas([('0', 1)])
self.assertIn('3', quota_ref.dump_data().allocations)
self.assertFalse(quota_ref.request_quota('4', 180))
quota_ref.alter_allocations(['3'], [50])
self.assertIn('4', quota_ref.dump_data().allocations)
def testFetchRemoteData(self):
pool_address = '127.0.0.1:%d' % get_next_port()
session_id = str(uuid.uuid4())
mock_data = np.array([1, 2, 3, 4])
with create_actor_pool(n_process=1, backend='gevent',
address=pool_address, distributor=MarsDistributor(2, 'w:0:')) as pool:
self.create_standard_actors(pool, pool_address, with_daemon=False, with_status=False,
with_resource=True)
pool.create_actor(CpuCalcActor)
pool.create_actor(MockSenderActor, mock_data, 'in', uid='w:mock_sender')
cluster_info_ref = pool.actor_ref(WorkerClusterInfoActor.default_name())
chunk_meta_client = ChunkMetaClient(pool, cluster_info_ref)
import mars.tensor as mt
from mars.tensor.expressions.fetch import TensorFetch
arr = mt.ones((4,), chunk_size=4)
arr_add = mt.array(mock_data)
result_tensor = arr + arr_add
graph = result_tensor.build_graph(compose=False, tiled=True)
modified_chunk = arr_add.chunks[0]
arr_add.chunks[0]._op = TensorFetch(
dtype=modified_chunk.dtype, _outputs=[weakref.ref(o) for o in modified_chunk.op.outputs],
_key=modified_chunk.op.key)
with self.run_actor_test(pool) as test_actor:
graph_key = str(uuid.uuid4())
execution_ref = test_actor.promise_ref(ExecutionActor.default_name())
execution_ref.enqueue_graph(session_id, graph_key, serialize_graph(graph),
dict(chunks=[result_tensor.chunks[0].key]), None, _promise=True) \
.then(lambda *_: execution_ref.start_execution(session_id, graph_key, _promise=True)) \
.then(lambda *_: test_actor.set_result(None)) \
.catch(lambda *exc: test_actor.set_result(exc, False))
with self.assertRaises(DependencyMissing):
self.get_result()
chunk_meta_client.set_chunk_meta(session_id, modified_chunk.key, size=mock_data.nbytes,
shape=mock_data.shape, workers=('0.0.0.0:1234',))
with self.run_actor_test(pool) as test_actor:
graph_key = str(uuid.uuid4())
execution_ref = test_actor.promise_ref(ExecutionActor.default_name())
execution_ref.enqueue_graph(session_id, graph_key, serialize_graph(graph),
dict(chunks=[result_tensor.chunks[0].key]), None, _promise=True) \
.then(lambda *_: execution_ref.start_execution(session_id, graph_key, _promise=True)) \
.then(lambda *_: test_actor.set_result(None)) \
.catch(lambda *exc: test_actor.set_result(exc, False))
with self.assertRaises(DependencyMissing):
self.get_result()
chunk_meta_client.set_chunk_meta(session_id, modified_chunk.key, size=mock_data.nbytes,
shape=mock_data.shape, workers=('0.0.0.0:1234', pool_address))
with self.run_actor_test(pool) as test_actor:
def _validate(_):
data = test_actor._chunk_store.get(session_id, result_tensor.chunks[0].key)
assert_array_equal(data, mock_data + np.ones((4,)))
graph_key = str(uuid.uuid4())
execution_ref = test_actor.promise_ref(ExecutionActor.default_name())
execution_ref.enqueue_graph(session_id, graph_key, serialize_graph(graph),
dict(chunks=[result_tensor.chunks[0].key]), None, _promise=True) \
.then(lambda *_: execution_ref.start_execution(session_id, graph_key, _promise=True)) \
.then(_validate) \
.then(lambda *_: test_actor.set_result(None)) \
.catch(lambda *exc: test_actor.set_result(exc, False))
self.get_result()