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)