async def test_mem_quota_allocation(actor_pool, enable_kill_slot):
from mars.utils import AttributeDict
mock_mem_stat = AttributeDict(dict(total=300, available=50, used=0, free=50))
mock_band_slot_manager_ref = await mo.create_actor(
MockBandSlotManagerActor, uid=BandSlotManagerActor.gen_uid('numa-0'),
address=actor_pool.external_address)
quota_ref = await mo.create_actor(
MemQuotaActor, (actor_pool.external_address, 'numa-0'), 300,
hard_limit=300, refresh_time=0.1,
enable_kill_slot=enable_kill_slot,
uid=MemQuotaActor.gen_uid('cpu-0'),
address=actor_pool.external_address) # type: Union[QuotaActorRef, mo.ActorRef]
with mock.patch('mars.resource.virtual_memory', new=lambda: mock_mem_stat):
time_recs = [time.time()]
async def task_fun():
await quota_ref.request_batch_quota({'req': 100})
await quota_ref.release_quotas(['req'])
time_recs.append(time.time())
task = asyncio.create_task(task_fun())
await asyncio.sleep(0.2)
assert 'req' not in (await quota_ref.dump_data()).allocations
mock_mem_stat['available'] = 150
mock_mem_stat['free'] = 150
await asyncio.wait_for(task, timeout=1)
assert 0.15 < abs(time_recs[0] - time_recs[1]) < 1
assert bool(await mock_band_slot_manager_ref.get_restart_record()) == enable_kill_slot
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)