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