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 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 testTaskQueueActor(self):
with create_actor_pool(n_process=1, backend='gevent') as pool:
pool.create_actor(MockExecutionActor,
10,
uid=ExecutionActor.default_name())
quota_ref = pool.create_actor(QuotaActor,
30,
uid=MemQuotaActor.default_name())
pool.create_actor(TaskQueueActor,
4,
uid=TaskQueueActor.default_name())
session_id = str(uuid.uuid4())
chunk_keys = [str(uuid.uuid4()).replace('-', '') for _ in range(6)]
with self.run_actor_test(pool) as test_actor:
queue_ref = test_actor.promise_ref(
TaskQueueActor.default_name())
res_times = dict()
def callback_fun(key):
res_times[key] = time.time()
for idx, k in enumerate(chunk_keys):
depth = len(chunk_keys) - idx
queue_ref.enqueue_task(session_id, k, dict(depth=depth), _promise=True) \
.then(functools.partial(callback_fun, k))
gevent.sleep(1)
self.assertEqual(queue_ref.get_allocated_count(), 4)
queue_ref.update_priority(session_id, chunk_keys[-1],
dict(depth=len(chunk_keys)))
quota_ref.release_quota(chunk_keys[0])
queue_ref.release_task(session_id, chunk_keys[0])
quota_ref.release_quota(chunk_keys[1])
queue_ref.release_task(session_id, chunk_keys[1])
gevent.sleep(0.5)
self.assertIn(chunk_keys[-1], res_times)
for k in chunk_keys[:3]:
self.assertLessEqual(res_times[k],
res_times[chunk_keys[-1]] - 0.5)
self.assertIn(k, res_times)
