async def test_parallel_allocate_idle_label():
start_method = os.environ.get('POOL_START_METHOD', 'forkserver') \
if sys.platform != 'win32' else None
pool = await create_actor_pool('127.0.0.1',
pool_cls=MainActorPool,
n_process=2,
subprocess_start_method=start_method,
labels=[None, 'my_label', 'my_label'])
class _Actor(Actor):
def get_pid(self):
return os.getpid()
ctx = get_context()
strategy = IdleLabel('my_label', 'tests')
tasks = [
ctx.create_actor(_Actor,
allocate_strategy=strategy,
address=pool.external_address),
ctx.create_actor(_Actor,
allocate_strategy=strategy,
address=pool.external_address),
]
refs = await asyncio.gather(*tasks)
# outputs identical process ids, while the result should be different
assert len({await ref.get_pid() for ref in refs}) == 2
async def test_slot_kill(actor_pool: ActorPoolType):
pool, slot_manager_ref = actor_pool
strategy = IdleLabel('numa-0', 'task_actor')
task_ref = await mo.create_actor(TaskActor, {},
allocate_strategy=strategy,
address=pool.external_address)
assert await mo.actor_ref(BandSlotControlActor.gen_uid('numa-0', 0),
address=pool.external_address)
delayed_task = asyncio.create_task(task_ref.queued_call('key', 10))
await asyncio.sleep(0.1)
# check if process hosting the actor is closed
kill_task = asyncio.create_task(slot_manager_ref.kill_slot(0))
await asyncio.sleep(0)
kill_task2 = asyncio.create_task(slot_manager_ref.kill_slot(0))
with pytest.raises(ServerClosed):
await delayed_task
# check if slot actor is restored
await kill_task
# check if secondary task makes no change
await kill_task2
assert await mo.actor_ref(BandSlotControlActor.gen_uid('numa-0', 0),
address=pool.external_address)
assert await mo.actor_ref(task_ref)
async def test_cancel_transfer(create_actors, mock_sender, mock_receiver):
worker_address_1, worker_address_2 = create_actors
strategy = IdleLabel('io', 'mock_sender')
quota_refs = {StorageLevel.MEMORY: await mo.actor_ref(
StorageQuotaActor, StorageLevel.MEMORY, 5 * 1024 * 1024,
address=worker_address_2, uid=StorageQuotaActor.gen_uid(StorageLevel.MEMORY))}
await mo.create_actor(
mock_sender, uid=mock_sender.default_uid(),
address=worker_address_1, allocate_strategy=strategy)
await mo.create_actor(
mock_receiver, quota_refs, uid=mock_receiver.default_uid(),
address=worker_address_2, allocate_strategy=strategy)
data1 = np.random.rand(10, 10)
storage_handler1 = await mo.actor_ref(
uid=StorageHandlerActor.default_uid(),
address=worker_address_1)
storage_handler2 = await mo.actor_ref(
uid=StorageHandlerActor.default_uid(),
address=worker_address_2)
await storage_handler1.put('mock', 'data_key1',
data1, StorageLevel.MEMORY)
sender_actor = await mo.actor_ref(address=worker_address_1,
uid=mock_sender.default_uid())
used_before = (await quota_refs[StorageLevel.MEMORY].get_quota())[1]
send_task = asyncio.create_task(sender_actor.send_data(
'mock', 'data_key1', worker_address_2, StorageLevel.MEMORY))
await asyncio.sleep(0.5)
send_task.cancel()
try:
await send_task
except asyncio.CancelledError:
pass
used = (await quota_refs[StorageLevel.MEMORY].get_quota())[1]
assert used == used_before
with pytest.raises(DataNotExist):
await storage_handler2.get('mock', 'data_key1')
send_task = asyncio.create_task(sender_actor.send_data(
'mock', 'data_key1', worker_address_2, StorageLevel.MEMORY))
await send_task
get_data = await storage_handler2.get('mock', 'data_key1')
np.testing.assert_array_equal(data1, get_data)
async def test_logging_config(logging_conf):
start_method = os.environ.get('POOL_START_METHOD', 'forkserver') \
if sys.platform != 'win32' else None
pool = await create_actor_pool('127.0.0.1',
pool_cls=MainActorPool,
n_process=1,
subprocess_start_method=start_method,
labels=[None, 'my_label'],
logging_conf=logging_conf)
class _Actor(Actor):
def get_logger_level(self):
logger = logging.getLogger(__name__)
return logger.getEffectiveLevel()
async with pool:
ctx = get_context()
strategy = IdleLabel('my_label', 'tests')
ref = await ctx.create_actor(_Actor,
allocate_strategy=strategy,
address=pool.external_address)
assert await ref.get_logger_level() == logging.DEBUG
async def test_slot_restart(actor_pool: ActorPoolType):
pool, slot_manager_ref = actor_pool
strategy = IdleLabel('numa-0', 'task_actor')
task_refs = []
for idx in range(3):
ref = await mo.create_actor(TaskActor, {},
slot_id=idx,
allocate_strategy=strategy,
address=pool.external_address)
await ref.queued_call('idx', idx)
task_refs.append(ref)
await slot_manager_ref.acquire_free_slot(('session_id', 'subtask_id1'))
slot_id2 = await slot_manager_ref.acquire_free_slot(
('session_id', 'subtask_id2'))
await slot_manager_ref.release_free_slot(slot_id2)
async def record_finish_time(coro):
await coro
return time.time()
restart_task1 = asyncio.create_task(
record_finish_time(slot_manager_ref.restart_free_slots()))
await asyncio.sleep(0)
restart_task2 = asyncio.create_task(
record_finish_time(slot_manager_ref.restart_free_slots()))
acquire_task = asyncio.create_task(
record_finish_time(
slot_manager_ref.acquire_free_slot(('session_id', 'subtask_id3'))))
await asyncio.gather(restart_task1, restart_task2, acquire_task)
# check only slots with running records are restarted
assert len(await task_refs[0].get_call_logs()) > 0
assert len(await task_refs[1].get_call_logs()) == 0
assert len(await task_refs[2].get_call_logs()) > 0
assert abs(restart_task1.result() - acquire_task.result()) < 0.1
def test_idle_label():
strategy = IdleLabel('test', 'my_mark')
addresses = config.get_external_addresses(label='test')
assert len(addresses) == 2
allocated = {
addresses[0]: {
create_actor_ref(addresses[0], b'id1'): (strategy, None)
}
}
assert strategy.get_allocated_address(config, allocated) == addresses[1]
strategy2 = IdleLabel('test', 'my_mark')
allocated = {
addresses[0]: {
create_actor_ref(addresses[0], b'id1'): (strategy, None),
create_actor_ref(addresses[0], b'id2'): (RandomLabel('test'), None)
},
addresses[1]: {
create_actor_ref(addresses[1], b'id3'): (strategy2, None)
}
}
with pytest.raises(NoIdleSlot):
strategy2.get_allocated_address(config, allocated)
async def test_main_actor_pool():
config = ActorPoolConfig()
my_label = 'computation'
main_address = f'127.0.0.1:{get_next_port()}'
_add_pool_conf(config, 0, 'main', 'unixsocket:///0', main_address)
_add_pool_conf(config,
1,
my_label,
'unixsocket:///1',
f'127.0.0.1:{get_next_port()}',
env={'my_env': '1'})
_add_pool_conf(config, 2, my_label, 'unixsocket:///2',
f'127.0.0.1:{get_next_port()}')
strategy = IdleLabel(my_label, 'my_test')
async with await MainActorPool.create({'actor_pool_config':
config}) as pool:
create_actor_message = CreateActorMessage(new_message_id(), TestActor,
b'test', tuple(), dict(),
MainPool())
message = await pool.create_actor(create_actor_message)
actor_ref = message.result
assert actor_ref.address == main_address
create_actor_message1 = CreateActorMessage(new_message_id(), TestActor,
b'test1', tuple(), dict(),
strategy)
message1 = await pool.create_actor(create_actor_message1)
actor_ref1 = message1.result
assert actor_ref1.address in config.get_external_addresses(my_label)
create_actor_message2 = CreateActorMessage(new_message_id(), TestActor,
b'test2', tuple(), dict(),
strategy)
message2 = await pool.create_actor(create_actor_message2)
actor_ref2 = message2.result
assert actor_ref2.address in config.get_external_addresses(my_label)
assert actor_ref2.address != actor_ref1.address
create_actor_message3 = CreateActorMessage(new_message_id(), TestActor,
b'test3', tuple(), dict(),
strategy)
message3 = await pool.create_actor(create_actor_message3)
# no slot to allocate the same label
assert isinstance(message3.error, NoIdleSlot)
has_actor_message = HasActorMessage(
new_message_id(), create_actor_ref(main_address, b'test2'))
assert (await pool.has_actor(has_actor_message)).result is True
actor_ref_message = ActorRefMessage(
new_message_id(), create_actor_ref(main_address, b'test2'))
assert (await pool.actor_ref(actor_ref_message)).result == actor_ref2
# tell
tell_message = TellMessage(new_message_id(), actor_ref1,
('add', 0, (2, ), dict()))
message = await pool.tell(tell_message)
assert message.result is None
# send
send_message = SendMessage(new_message_id(), actor_ref1,
('add', 0, (4, ), dict()))
assert (await pool.send(send_message)).result == 6
# test error message
# type mismatch
send_message = SendMessage(new_message_id(), actor_ref1,
('add', 0, ('3', ), dict()))
result = await pool.send(send_message)
assert isinstance(result.error, TypeError)
# send and tell to main process
tell_message = TellMessage(new_message_id(), actor_ref,
('add', 0, (2, ), dict()))
message = await pool.tell(tell_message)
assert message.result is None
send_message = SendMessage(new_message_id(), actor_ref,
('add', 0, (4, ), dict()))
assert (await pool.send(send_message)).result == 6
# send and cancel
send_message = SendMessage(new_message_id(), actor_ref1,
('sleep', 0, (20, ), dict()))
result_task = asyncio.create_task(pool.send(send_message))
start = time.time()
cancel_message = CancelMessage(new_message_id(), actor_ref1.address,
send_message.message_id)
cancel_task = asyncio.create_task(pool.cancel(cancel_message))
result = await asyncio.wait_for(cancel_task, 3)
assert result.message_type == MessageType.result
assert result.result is True
result = await result_task
assert time.time() - start < 3
assert result.message_type == MessageType.result
assert result.result == 7
# destroy
destroy_actor_message = DestroyActorMessage(new_message_id(),
actor_ref1)
message = await pool.destroy_actor(destroy_actor_message)
assert message.result == actor_ref1.uid
tell_message = TellMessage(new_message_id(), actor_ref1,
('add', 0, (2, ), dict()))
message = await pool.tell(tell_message)
assert isinstance(message, ErrorMessage)
# destroy via connecting to sub pool directly
async with await pool.router.get_client(
config.get_external_addresses()[-1]) as client:
destroy_actor_message = DestroyActorMessage(
new_message_id(), actor_ref2)
await client.send(destroy_actor_message)
result = await client.recv()
assert result.result == actor_ref2.uid
# test sync config
config.add_pool_conf(3, 'sub', 'unixsocket:///3',
f'127.0.0.1:{get_next_port()}')
sync_config_message = ControlMessage(new_message_id(),
pool.external_address,
ControlMessageType.sync_config,
config)
message = await pool.handle_control_command(sync_config_message)
assert message.result is True
# test get config
get_config_message = ControlMessage(new_message_id(),
config.get_external_addresses()[1],
ControlMessageType.get_config,
None)
message = await pool.handle_control_command(get_config_message)
config2 = message.result
assert config.as_dict() == config2.as_dict()
assert pool.stopped
async def test_cancel_transfer(create_actors, mock_sender, mock_receiver):
worker_address_1, worker_address_2 = create_actors
quota_refs = {
StorageLevel.MEMORY:
await
mo.actor_ref(StorageQuotaActor,
StorageLevel.MEMORY,
5 * 1024 * 1024,
address=worker_address_2,
uid=StorageQuotaActor.gen_uid('numa-0',
StorageLevel.MEMORY))
}
data_manager_ref = await mo.actor_ref(uid=DataManagerActor.default_uid(),
address=worker_address_1)
storage_handler1 = await mo.actor_ref(
uid=StorageHandlerActor.gen_uid('numa-0'), address=worker_address_1)
storage_handler2 = await mo.actor_ref(
uid=StorageHandlerActor.gen_uid('numa-0'), address=worker_address_2)
sender_actor = await mo.create_actor(mock_sender,
data_manager_ref=data_manager_ref,
uid=mock_sender.default_uid(),
address=worker_address_1,
allocate_strategy=IdleLabel(
'io', 'mock_sender'))
await mo.create_actor(mock_receiver,
quota_refs,
uid=mock_receiver.default_uid(),
address=worker_address_2,
allocate_strategy=IdleLabel('io', 'mock_receiver'))
data1 = np.random.rand(10, 10)
await storage_handler1.put('mock', 'data_key1', data1, StorageLevel.MEMORY)
data2 = pd.DataFrame(np.random.rand(100, 100))
await storage_handler1.put('mock', 'data_key2', data2, StorageLevel.MEMORY)
used_before = (await quota_refs[StorageLevel.MEMORY].get_quota())[1]
send_task = asyncio.create_task(
sender_actor.send_batch_data('mock', ['data_key1'], worker_address_2,
StorageLevel.MEMORY))
await asyncio.sleep(0.5)
send_task.cancel()
with pytest.raises(asyncio.CancelledError):
await send_task
used = (await quota_refs[StorageLevel.MEMORY].get_quota())[1]
assert used == used_before
with pytest.raises(DataNotExist):
await storage_handler2.get('mock', 'data_key1')
send_task = asyncio.create_task(
sender_actor.send_batch_data('mock', ['data_key1'], worker_address_2,
StorageLevel.MEMORY))
await send_task
get_data = await storage_handler2.get('mock', 'data_key1')
np.testing.assert_array_equal(data1, get_data)
# cancel when fetch the same data Simultaneously
if mock_sender is MockSenderManagerActor:
send_task1 = asyncio.create_task(
sender_actor.send_batch_data('mock', ['data_key2'],
worker_address_2,
StorageLevel.MEMORY))
send_task2 = asyncio.create_task(
sender_actor.send_batch_data('mock', ['data_key2'],
worker_address_2,
StorageLevel.MEMORY))
await asyncio.sleep(0.5)
send_task1.cancel()
with pytest.raises(asyncio.CancelledError):
await send_task1
await send_task2
get_data2 = await storage_handler2.get('mock', 'data_key2')
pd.testing.assert_frame_equal(get_data2, data2)