async def actor_pool():
pool = await mo.create_actor_pool('127.0.0.1', n_process=0)
async with pool:
session_id = 'test_session'
await MockClusterAPI.create(pool.external_address)
queue_ref = await mo.create_actor(
MockSubtaskQueueingActor,
uid=SubtaskQueueingActor.gen_uid(session_id),
address=pool.external_address)
slots_ref = await mo.create_actor(
GlobalSlotManagerActor,
uid=GlobalSlotManagerActor.default_uid(),
address=pool.external_address)
task_manager_ref = await mo.create_actor(
MockTaskManagerActor,
uid=TaskManagerActor.gen_uid(session_id),
address=pool.external_address)
execution_ref = await mo.create_actor(
MockSubtaskExecutionActor,
uid=SubtaskExecutionActor.default_uid(),
address=pool.external_address)
submitter_ref = await mo.create_actor(
SubtaskManagerActor,
session_id,
uid=SubtaskManagerActor.gen_uid(session_id),
address=pool.external_address)
yield pool, session_id, execution_ref, submitter_ref, queue_ref, task_manager_ref
await mo.destroy_actor(slots_ref)
async def actor_pool():
pool = await mo.create_actor_pool('127.0.0.1', n_process=0)
async with pool:
session_id = 'test_session'
await MockClusterAPI.create(pool.external_address)
# create assigner actor
await mo.create_actor(MockAssignerActor,
uid=AssignerActor.gen_uid(session_id),
address=pool.external_address)
# create queueing actor
manager_ref = await mo.create_actor(
MockSubtaskManagerActor,
uid=SubtaskManagerActor.gen_uid(session_id),
address=pool.external_address)
# create slots actor
slots_ref = await mo.create_actor(
MockSlotsActor,
uid=GlobalSlotManagerActor.default_uid(),
address=pool.external_address)
# create queueing actor
queueing_ref = await mo.create_actor(
SubtaskQueueingActor,
session_id,
1,
uid=SubtaskQueueingActor.gen_uid(session_id),
address=pool.external_address)
yield pool, session_id, queueing_ref, slots_ref, manager_ref
await mo.destroy_actor(queueing_ref)
async def test_report_usage(actor_pool):
pool, slot_manager_ref = actor_pool
await slot_manager_ref.acquire_free_slot(('session_id', 'subtask_id'))
await asyncio.sleep(1.3)
global_slot_ref = await mo.actor_ref(
uid=GlobalSlotManagerActor.default_uid(),
address=pool.external_address)
_band, session_id, subtask_id, slots = await global_slot_ref.get_result()
assert slots == pytest.approx(1.0)
assert session_id == 'session_id'
assert subtask_id == 'subtask_id'
async def actor_pool():
pool = await mo.create_actor_pool('127.0.0.1', n_process=0)
async with pool:
session_id = 'test_session'
await MockClusterAPI.create(pool.external_address)
await MockSessionAPI.create(pool.external_address,
session_id=session_id)
global_slot_ref = await mo.create_actor(
GlobalSlotManagerActor,
uid=GlobalSlotManagerActor.default_uid(),
address=pool.external_address)
yield pool, session_id, global_slot_ref
await mo.destroy_actor(global_slot_ref)
async def test_schedule_error(actor_pools):
sv_pool, worker_pool, session_id, task_manager_ref = actor_pools
global_slot_ref = await mo.actor_ref(GlobalSlotManagerActor.default_uid(),
address=sv_pool.external_address)
scheduling_api = await SchedulingAPI.create(session_id,
sv_pool.external_address)
def _remote_fun():
raise ValueError
a = mr.spawn(_remote_fun)
subtask = _gen_subtask(a, session_id)
subtask.expect_bands = [(worker_pool.external_address, 'numa-0')]
await scheduling_api.add_subtasks([subtask])
with pytest.raises(ValueError):
await task_manager_ref.wait_subtask_result(subtask.subtask_id)
assert (await global_slot_ref.get_used_slots())['numa-0'] == 0
async def actor_pool(request):
start_method = os.environ.get('POOL_START_METHOD', 'forkserver') \
if sys.platform != 'win32' else None
n_slots = request.param
pool = await mo.create_actor_pool(f'127.0.0.1:{get_next_port()}',
n_process=n_slots,
labels=[None] + ['numa-0'] * n_slots,
subprocess_start_method=start_method)
async with pool:
global_slots_ref = await mo.create_actor(
MockGlobalSlotManagerActor,
uid=GlobalSlotManagerActor.default_uid(),
address=pool.external_address)
slot_manager_ref = await mo.create_actor(
BandSlotManagerActor,
'numa-0',
n_slots,
global_slots_ref,
uid=BandSlotManagerActor.gen_uid('numa-0'),
address=pool.external_address)
yield pool, slot_manager_ref
async def test_schedule_success(actor_pools):
sv_pool, worker_pool, session_id, task_manager_ref = actor_pools
global_slot_ref = await mo.actor_ref(GlobalSlotManagerActor.default_uid(),
address=sv_pool.external_address)
scheduling_api = await SchedulingAPI.create(session_id,
sv_pool.external_address)
storage_api = await StorageAPI.create(session_id,
worker_pool.external_address)
a = mt.ones((10, 10), chunk_size=10)
b = a + 1
subtask = _gen_subtask(b, session_id)
subtask.expect_bands = [(worker_pool.external_address, 'numa-0')]
await scheduling_api.add_subtasks([subtask], [(0, )])
await task_manager_ref.wait_subtask_result(subtask.subtask_id)
result_key = next(subtask.chunk_graph.iter_indep(reverse=True)).key
result = await storage_api.get(result_key)
np.testing.assert_array_equal(np.ones((10, 10)) + 1, result)
assert (await global_slot_ref.get_used_slots())['numa-0'] == 0
async def test_schedule_cancel(actor_pools):
sv_pool, worker_pool, session_id, task_manager_ref = actor_pools
global_slot_ref = await mo.actor_ref(GlobalSlotManagerActor.default_uid(),
address=sv_pool.external_address)
scheduling_api = await SchedulingAPI.create(session_id,
sv_pool.external_address)
def _remote_fun(secs):
time.sleep(secs)
return secs
async def _waiter_fun(subtask_id):
await task_manager_ref.wait_subtask_result(subtask_id)
await scheduling_api.finish_subtasks([subtask_id])
subtasks = []
wait_tasks = []
for task_id in range(6):
a = mr.spawn(_remote_fun, args=(1 - 0.01 * task_id, ))
subtask = _gen_subtask(a, session_id)
subtask.subtask_id = f'test_schedule_queue_subtask_{task_id}'
subtask.expect_bands = [(worker_pool.external_address, 'numa-0')]
subtask.priority = (4 - task_id, )
wait_tasks.append(asyncio.create_task(_waiter_fun(subtask.subtask_id)))
subtasks.append(subtask)
await scheduling_api.add_subtasks(subtasks)
await asyncio.gather(*wait_tasks[:2])
await scheduling_api.cancel_subtasks(
[subtask.subtask_id for subtask in subtasks], kill_timeout=0.1)
for wait_task in wait_tasks[2:]:
with pytest.raises(asyncio.CancelledError):
await wait_task
assert (await global_slot_ref.get_used_slots())['numa-0'] == 0
async def test_schedule_queue(actor_pools):
sv_pool, worker_pool, session_id, task_manager_ref = actor_pools
global_slot_ref = await mo.actor_ref(GlobalSlotManagerActor.default_uid(),
address=sv_pool.external_address)
scheduling_api = await SchedulingAPI.create(session_id,
sv_pool.external_address)
finish_ids, finish_time = [], []
def _remote_fun(secs):
time.sleep(secs)
return secs
async def _waiter_fun(subtask_id):
await task_manager_ref.wait_subtask_result(subtask_id)
await scheduling_api.finish_subtasks([subtask_id])
finish_ids.append(subtask_id)
finish_time.append(time.time())
subtasks = []
wait_tasks = []
for task_id in range(6):
a = mr.spawn(_remote_fun, args=(0.5 + 0.01 * task_id, ))
subtask = _gen_subtask(a, session_id)
subtask.subtask_id = f'test_schedule_queue_subtask_{task_id}'
subtask.expect_bands = [(worker_pool.external_address, 'numa-0')]
subtask.priority = (4 - task_id, )
wait_tasks.append(asyncio.create_task(_waiter_fun(subtask.subtask_id)))
subtasks.append(subtask)
await scheduling_api.add_subtasks(subtasks)
await scheduling_api.update_subtask_priority(subtasks[-1].subtask_id,
(6, ))
await asyncio.gather(*wait_tasks)
assert (await global_slot_ref.get_used_slots())['numa-0'] == 0
