def _launch_job_with_pyunicore(operation, simulator_gid, is_group_launch):
# type: (Operation, str, bool) -> Job
LOGGER.info("Prepare job inputs for operation: {}".format(
operation.id))
job_plain_inputs = HPCSchedulerClient._prepare_input(
operation, simulator_gid)
available_space = HPCSchedulerClient.compute_available_disk_space(
operation)
LOGGER.info("Prepare job configuration for operation: {}".format(
operation.id))
job_config, job_script = HPCSchedulerClient._configure_job(
simulator_gid, available_space, is_group_launch, operation.id)
LOGGER.info("Prepare encryption for operation: {}".format(
operation.id))
storage_interface = StorageInterface()
LOGGER.info("Encrypt job inputs for operation: {}".format(
operation.id))
job_encrypted_inputs = storage_interface.encrypt_inputs(
simulator_gid, job_plain_inputs)
# use "DAINT-CSCS" -- change if another supercomputer is prepared for usage
LOGGER.info("Prepare unicore client for operation: {}".format(
operation.id))
site_client = HPCSchedulerClient._build_unicore_client(
os.environ[HPCSchedulerClient.CSCS_LOGIN_TOKEN_ENV_KEY],
unicore_client._HBP_REGISTRY_URL,
TvbProfile.current.hpc.HPC_COMPUTE_SITE)
LOGGER.info("Submit job for operation: {}".format(operation.id))
job = HPCSchedulerClient._create_job_with_pyunicore(
pyunicore_client=site_client,
job_description=job_config,
job_script=job_script,
inputs=job_encrypted_inputs)
LOGGER.info("Job url {} for operation: {}".format(
job.resource_url, operation.id))
op_identifier = OperationProcessIdentifier(operation_id=operation.id,
job_id=job.resource_url)
dao.store_entity(op_identifier)
LOGGER.info("Job mount point: {}".format(
job.working_dir.properties[HPCSettings.JOB_MOUNT_POINT_KEY]))
return job
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()
