def _stage_out_results(working_dir, simulator_gid):
# type: (Storage, typing.Union[uuid.UUID, str]) -> list
output_subfolder = HPCSchedulerClient.CSCS_DATA_FOLDER + '/' + HPCSchedulerClient.OUTPUT_FOLDER
output_list = HPCSchedulerClient._listdir(working_dir,
output_subfolder)
LOGGER.info("Output list {}".format(output_list))
storage_interface = StorageInterface()
encrypted_dir = os.path.join(
storage_interface.get_encrypted_dir(simulator_gid),
HPCSchedulerClient.OUTPUT_FOLDER)
encrypted_files = HPCSchedulerClient._stage_out_outputs(
encrypted_dir, output_list)
# Clean data uploaded on CSCS
LOGGER.info("Clean uploaded files and results")
working_dir.rmdir(HPCSchedulerClient.CSCS_DATA_FOLDER)
LOGGER.info(encrypted_files)
return encrypted_files
class TestHPCSchedulerClient(BaseTestCase):
def setup_method(self):
self.storage_interface = StorageInterface()
self.dir_gid = '123'
self.clean_database()
self.test_user = TestFactory.create_user()
self.test_project = TestFactory.create_project(self.test_user)
def _prepare_dummy_files(self, tmpdir):
dummy_file1 = os.path.join(str(tmpdir), 'dummy1.txt')
open(dummy_file1, 'a').close()
dummy_file2 = os.path.join(str(tmpdir), 'dummy2.txt')
open(dummy_file2, 'a').close()
job_inputs = [dummy_file1, dummy_file2]
return job_inputs
def test_encrypt_inputs(self, tmpdir):
job_inputs = self._prepare_dummy_files(tmpdir)
job_encrypted_inputs = self.storage_interface.encrypt_inputs(
self.dir_gid, job_inputs)
# Encrypted folder has 2 more files are more then plain folder
assert len(job_encrypted_inputs) == len(job_inputs)
def test_decrypt_results(self, tmpdir):
# Prepare encrypted dir
job_inputs = self._prepare_dummy_files(tmpdir)
self.storage_interface.encrypt_inputs(self.dir_gid, job_inputs)
encrypted_dir = self.storage_interface.get_encrypted_dir(self.dir_gid)
# Unencrypt data
out_dir = os.path.join(str(tmpdir), 'output')
os.mkdir(out_dir)
self.storage_interface.decrypt_results_to_dir(self.dir_gid, out_dir)
list_plain_dir = os.listdir(out_dir)
assert len(list_plain_dir) == len(os.listdir(encrypted_dir))
assert 'dummy1.txt' in list_plain_dir
assert 'dummy2.txt' in list_plain_dir
def test_decrypt_files(self, tmpdir):
# Prepare encrypted dir
job_inputs = self._prepare_dummy_files(tmpdir)
enc_files = self.storage_interface.encrypt_inputs(
self.dir_gid, job_inputs)
# Unencrypt data
out_dir = os.path.join(str(tmpdir), 'output')
os.mkdir(out_dir)
self.storage_interface.decrypt_files_to_dir(self.dir_gid,
[enc_files[1]], out_dir)
list_plain_dir = os.listdir(out_dir)
assert len(list_plain_dir) == 1
assert os.path.basename(enc_files[0]).replace('.aes',
'') not in list_plain_dir
assert os.path.basename(enc_files[1]).replace('.aes',
'') in list_plain_dir
def test_do_operation_launch(self, simulator_factory, operation_factory,
mocker):
# Prepare encrypted dir
op = operation_factory(test_user=self.test_user,
test_project=self.test_project)
sim_folder, sim_gid = simulator_factory(op=op)
self._do_operation_launch(op, sim_gid, mocker)
def _do_operation_launch(self, op, sim_gid, mocker, is_pse=False):
# Prepare encrypted dir
self.storage_interface = StorageInterface()
self.dir_gid = sim_gid
job_encrypted_inputs = HPCSchedulerClient()._prepare_input(op, sim_gid)
self.storage_interface.encrypt_inputs(sim_gid, job_encrypted_inputs)
encrypted_dir = self.storage_interface.get_encrypted_dir(self.dir_gid)
mocker.patch('tvb.core.operation_hpc_launcher._request_passfile',
_request_passfile_dummy)
mocker.patch(
'tvb.core.operation_hpc_launcher._update_operation_status',
_update_operation_status)
# Call do_operation_launch similarly to CSCS env
plain_dir = self.storage_interface.get_project_folder(
self.test_project.name, 'plain')
do_operation_launch(sim_gid.hex, 1000, is_pse, '', op.id, plain_dir)
assert len(os.listdir(encrypted_dir)) == 7
output_path = os.path.join(encrypted_dir,
HPCSchedulerClient.OUTPUT_FOLDER)
assert os.path.exists(output_path)
expected_files = 2
if is_pse:
expected_files = 3
assert len(os.listdir(output_path)) == expected_files
return output_path
def test_do_operation_launch_pse(self, simulator_factory,
operation_factory, mocker):
op = operation_factory(test_user=self.test_user,
test_project=self.test_project)
sim_folder, sim_gid = simulator_factory(op=op)
self._do_operation_launch(op, sim_gid, mocker, is_pse=True)
def test_prepare_inputs(self, operation_factory, simulator_factory):
op = operation_factory(test_user=self.test_user,
test_project=self.test_project)
sim_folder, sim_gid = simulator_factory(op=op)
hpc_client = HPCSchedulerClient()
input_files = hpc_client._prepare_input(op, sim_gid)
assert len(input_files) == 6
def test_prepare_inputs_with_surface(self, operation_factory,
simulator_factory):
op = operation_factory(test_user=self.test_user,
test_project=self.test_project)
sim_folder, sim_gid = simulator_factory(op=op, with_surface=True)
hpc_client = HPCSchedulerClient()
input_files = hpc_client._prepare_input(op, sim_gid)
assert len(input_files) == 9
def test_prepare_inputs_with_eeg_monitor(self, operation_factory,
simulator_factory,
surface_index_factory,
sensors_index_factory,
region_mapping_index_factory,
connectivity_index_factory):
surface_idx, surface = surface_index_factory(cortical=True)
sensors_idx, sensors = sensors_index_factory()
proj = ProjectionSurfaceEEG(sensors=sensors,
sources=surface,
projection_data=numpy.ones(3))
op = operation_factory()
storage_path = StorageInterface().get_project_folder(
op.project.name, str(op.id))
prj_db_db = h5.store_complete(proj, storage_path)
prj_db_db.fk_from_operation = op.id
dao.store_entity(prj_db_db)
connectivity = connectivity_index_factory(76, op)
rm_index = region_mapping_index_factory(conn_gid=connectivity.gid,
surface_gid=surface_idx.gid)
eeg_monitor = EEGViewModel(projection=proj.gid, sensors=sensors.gid)
eeg_monitor.region_mapping = rm_index.gid
sim_folder, sim_gid = simulator_factory(op=op,
monitor=eeg_monitor,
conn_gid=connectivity.gid)
hpc_client = HPCSchedulerClient()
input_files = hpc_client._prepare_input(op, sim_gid)
assert len(input_files) == 11
def test_stage_out_to_operation_folder(self, mocker, operation_factory,
simulator_factory,
pse_burst_configuration_factory):
burst = pse_burst_configuration_factory(self.test_project)
op = operation_factory(test_user=self.test_user,
test_project=self.test_project)
op.fk_operation_group = burst.fk_operation_group
dao.store_entity(op)
sim_folder, sim_gid = simulator_factory(op=op)
burst.simulator_gid = sim_gid.hex
dao.store_entity(burst)
output_path = self._do_operation_launch(op,
sim_gid,
mocker,
is_pse=True)
def _stage_out_dummy(dir, sim_gid):
return [
os.path.join(output_path, enc_file)
for enc_file in os.listdir(output_path)
]
mocker.patch.object(HPCSchedulerClient, '_stage_out_results',
_stage_out_dummy)
sim_results_files, metric_op, metric_file = HPCSchedulerClient.stage_out_to_operation_folder(
None, op, sim_gid)
assert op.id != metric_op.id
assert os.path.exists(metric_file)
assert len(sim_results_files) == 1
assert os.path.exists(sim_results_files[0])
def teardown_method(self):
encrypted_dir = self.storage_interface.get_encrypted_dir(self.dir_gid)
if os.path.isdir(encrypted_dir):
shutil.rmtree(encrypted_dir)
passfile = self.storage_interface.get_password_file(self.dir_gid)
if os.path.exists(passfile):
os.remove(passfile)
self.storage_interface.remove_project_structure(self.test_project.name)
self.clean_database()
