def stage_out_to_operation_folder(working_dir, operation, simulator_gid):
# type: (Storage, Operation, typing.Union[uuid.UUID, str]) -> (list, Operation, str)
encrypted_files = HPCSchedulerClient._stage_out_results(
working_dir, simulator_gid)
simulation_results = list()
metric_encrypted_file = None
metric_vm_encrypted_file = None
for encrypted_file in encrypted_files:
if os.path.basename(encrypted_file).startswith(
DatatypeMeasureH5.file_name_base()):
metric_encrypted_file = encrypted_file
elif os.path.basename(encrypted_file).startswith(
MEASURE_METRICS_MODEL_CLASS):
metric_vm_encrypted_file = encrypted_file
else:
simulation_results.append(encrypted_file)
encryption_handler = StorageInterface.get_encryption_handler(
simulator_gid)
metric_op, metric_file = HPCSchedulerClient._handle_metric_results(
metric_encrypted_file, metric_vm_encrypted_file, operation,
encryption_handler)
project = dao.get_project_by_id(operation.fk_launched_in)
operation_dir = HPCSchedulerClient.storage_interface.get_project_folder(
project.name, str(operation.id))
h5_filenames = encryption_handler.decrypt_files_to_dir(
simulation_results, operation_dir)
encryption_handler.cleanup_encryption_handler()
LOGGER.info("Decrypted h5: {}".format(h5_filenames))
LOGGER.info("Metric op: {}".format(metric_op))
LOGGER.info("Metric file: {}".format(metric_file))
return h5_filenames, metric_op, metric_file
def _configure_job(simulator_gid, available_space, is_group_launch,
operation_id):
# type: (str, int, bool, int) -> (dict, list)
bash_entrypoint = os.path.join(
os.environ[HPCSchedulerClient.TVB_BIN_ENV_KEY],
HPCSettings.HPC_LAUNCHER_SH_SCRIPT)
base_url = TvbProfile.current.web.BASE_URL
inputs_in_container = os.path.join(
HPCSchedulerClient.CONTAINER_INPUT_FOLDER,
StorageInterface.get_encryption_handler(
simulator_gid).current_enc_dirname)
# Build job configuration JSON
my_job = {
HPCSettings.UNICORE_EXE_KEY:
os.path.basename(bash_entrypoint),
HPCSettings.UNICORE_ARGS_KEY: [
simulator_gid, available_space, is_group_launch, base_url,
inputs_in_container, HPCSchedulerClient.HOME_FOLDER_MOUNT,
operation_id
],
HPCSettings.UNICORE_RESOURCER_KEY: {
"CPUs": "1"
}
}
if HPCSchedulerClient.CSCS_PROJECT in os.environ:
my_job[HPCSettings.UNICORE_PROJECT_KEY] = os.environ[
HPCSchedulerClient.CSCS_PROJECT]
return my_job, bash_entrypoint
def do_operation_launch(simulator_gid,
available_disk_space,
is_group_launch,
base_url,
operation_id,
plain_dir='/root/plain'):
try:
log.info("Preparing HPC launch for simulation with id={}".format(
simulator_gid))
populate_datatypes_registry()
log.info("Current TVB profile has HPC run=: {}".format(
TvbProfile.current.hpc.IS_HPC_RUN))
encryption_handler = StorageInterface.get_encryption_handler(
simulator_gid)
_request_passfile(simulator_gid, operation_id, base_url,
encryption_handler.get_password_file())
encryption_handler.decrypt_results_to_dir(plain_dir)
log.info("Current wdir is: {}".format(plain_dir))
view_model = h5.load_view_model(simulator_gid, plain_dir)
adapter_instance = HPCSimulatorAdapter(plain_dir, is_group_launch)
_update_operation_status(STATUS_STARTED, simulator_gid, operation_id,
base_url)
adapter_instance._prelaunch(None, view_model, available_disk_space)
_encrypt_results(adapter_instance, encryption_handler)
_update_operation_status(STATUS_FINISHED, simulator_gid, operation_id,
base_url)
except Exception as excep:
log.error("Could not execute operation {}".format(str(sys.argv[1])))
log.exception(excep)
_update_operation_status(STATUS_ERROR, simulator_gid, operation_id,
base_url)
raise excep
def _do_operation_launch(self, op, sim_gid, mocker, is_pse=False):
# Prepare encrypted dir
self.dir_gid = sim_gid
self.encryption_handler = StorageInterface.get_encryption_handler(
self.dir_gid)
job_encrypted_inputs = HPCSchedulerClient()._prepare_input(
op, self.dir_gid)
self.encryption_handler.encrypt_inputs(job_encrypted_inputs)
encrypted_dir = self.encryption_handler.get_encrypted_dir()
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(self.dir_gid, 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 encryption_config(self, simulator_gid, operation_id):
self.logger.info("Received a request for passfile with gid: {}".format(simulator_gid))
if cherrypy.request.method != 'GET':
raise cherrypy.HTTPError(HTTPStatus.METHOD_NOT_ALLOWED)
self._validate_request_params(simulator_gid, operation_id)
file_path = StorageInterface.get_encryption_handler(simulator_gid).get_password_file()
return serve_file(file_path, "application/x-download", "attachment", os.path.basename(file_path))
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))
encryption_handler = StorageInterface.get_encryption_handler(
simulator_gid)
LOGGER.info("Encrypt job inputs for operation: {}".format(
operation.id))
job_encrypted_inputs = encryption_handler.encrypt_inputs(
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
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_encryption_handler(
simulator_gid).get_encrypted_dir(),
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.encryption_handler = self.storage_interface.get_encryption_handler(
self.dir_gid)
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.encryption_handler.encrypt_inputs(
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.encryption_handler.encrypt_inputs(job_inputs)
encrypted_dir = self.encryption_handler.get_encrypted_dir()
# Unencrypt data
out_dir = os.path.join(str(tmpdir), 'output')
os.mkdir(out_dir)
self.encryption_handler.decrypt_results_to_dir(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.encryption_handler.encrypt_inputs(job_inputs)
# Unencrypt data
out_dir = os.path.join(str(tmpdir), 'output')
os.mkdir(out_dir)
self.encryption_handler.decrypt_files_to_dir([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.dir_gid = sim_gid
self.encryption_handler = StorageInterface.get_encryption_handler(
self.dir_gid)
job_encrypted_inputs = HPCSchedulerClient()._prepare_input(
op, self.dir_gid)
self.encryption_handler.encrypt_inputs(job_encrypted_inputs)
encrypted_dir = self.encryption_handler.get_encrypted_dir()
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(self.dir_gid, 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()
prj_db_db = h5.store_complete(proj, op.id, op.project.name)
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.encryption_handler.get_encrypted_dir()
passfile = self.encryption_handler.get_password_file()
self.storage_interface.remove_files([encrypted_dir, passfile])
self.clean_database()
