async def create_actors(actor_pools):
worker_pool_1, worker_pool_2 = actor_pools
if sys.platform == 'darwin':
plasma_dir = '/tmp'
else:
plasma_dir = '/dev/shm'
plasma_setup_params = dict(store_memory=5 * 1024 * 1024,
plasma_directory=plasma_dir,
check_dir_size=False)
storage_configs = {
"plasma": plasma_setup_params,
}
manager_ref1 = await mo.create_actor(
StorageManagerActor,
storage_configs,
uid=StorageManagerActor.default_uid(),
address=worker_pool_1.external_address)
manager_ref2 = await mo.create_actor(
StorageManagerActor,
storage_configs,
uid=StorageManagerActor.default_uid(),
address=worker_pool_2.external_address)
yield worker_pool_1.external_address, worker_pool_2.external_address
await mo.destroy_actor(manager_ref1)
await mo.destroy_actor(manager_ref2)
async def create_actors(actor_pool):
storage_configs = _build_storage_config()
manager_ref = await mo.create_actor(StorageManagerActor,
storage_configs,
uid=StorageManagerActor.default_uid(),
address=actor_pool.external_address)
sub_processes = list(actor_pool.sub_processes)
yield actor_pool.external_address, sub_processes[0], sub_processes[1]
await mo.destroy_actor(manager_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')
await mo.create_actor(mock_sender,
uid=mock_sender.default_uid(),
address=worker_address_1,
allocate_strategy=strategy)
await mo.create_actor(mock_receiver,
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)
storage_manager2 = await mo.actor_ref(
uid=StorageManagerActor.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 storage_manager2.get_quota(StorageLevel.MEMORY))[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 storage_manager2.get_quota(StorageLevel.MEMORY))[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_spill(create_actors):
worker_address = create_actors
storage_handler = await mo.actor_ref(uid=StorageHandlerActor.default_uid(),
address=worker_address)
storage_manager = await mo.actor_ref(uid=StorageManagerActor.default_uid(),
address=worker_address)
init_params = await storage_manager.get_client_params()
plasma_init_params = init_params['plasma']
plasma_handler = PlasmaStorage(**plasma_init_params)
memory_quota = await mo.actor_ref(StorageQuotaActor,
StorageLevel.MEMORY,
MEMORY_SIZE,
address=worker_address,
uid=StorageQuotaActor.gen_uid(
StorageLevel.MEMORY))
# fill to trigger spill
session_id = 'mock_session'
data_list = []
key_list = []
for i in range(10):
data = np.random.randint(0, 10000, (8000, ), np.int16)
key = f'mock_key_{i}'
await storage_handler.put(session_id, key, data, StorageLevel.MEMORY)
used = (await memory_quota.get_quota())[1]
assert used < MEMORY_SIZE
data_list.append(data)
key_list.append(key)
memory_object_list = await storage_handler.list(StorageLevel.MEMORY)
disk_object_list = await storage_handler.list(StorageLevel.DISK)
assert len(memory_object_list) == 2
assert len(disk_object_list) == 8
for key, data in zip(key_list, data_list):
get_data = await storage_handler.get(session_id, key)
np.testing.assert_array_equal(data, get_data)
plasma_list = await plasma_handler.list()
assert len(plasma_list) == len(memory_object_list)
async def create_actors(actor_pool):
if sys.platform == 'darwin':
plasma_dir = '/tmp'
else:
plasma_dir = '/dev/shm'
plasma_setup_params = dict(store_memory=MEMORY_SIZE,
plasma_directory=plasma_dir,
check_dir_size=False)
tempdir = tempfile.mkdtemp()
disk_setup_params = dict(root_dirs=tempdir, level='disk')
storage_configs = {
"plasma": plasma_setup_params,
"filesystem": disk_setup_params
}
manager_ref = await mo.create_actor(StorageManagerActor,
storage_configs,
uid=StorageManagerActor.default_uid(),
address=actor_pool.external_address)
yield actor_pool.external_address
await mo.destroy_actor(manager_ref)