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)
encryption_handler = EncryptionHandler(simulator_gid)
simulation_results = list()
metric_op = None
metric_file = None
for encrypted_file in encrypted_files:
if os.path.basename(encrypted_file).startswith(
DatatypeMeasureH5.file_name_base()):
metric_op_dir, metric_op = BurstService.prepare_metrics_operation(
operation)
metric_files = encryption_handler.decrypt_files_to_dir(
[encrypted_file], metric_op_dir)
metric_file = metric_files[0]
else:
simulation_results.append(encrypted_file)
project = dao.get_project_by_id(operation.fk_launched_in)
operation_dir = HPCSchedulerClient.file_handler.get_project_folder(
project, str(operation.id))
h5_filenames = EncryptionHandler(simulator_gid).decrypt_files_to_dir(
simulation_results, operation_dir)
return h5_filenames, metric_op, metric_file
def _do_operation_launch(self, op, sim_gid, mocker, is_pse=False):
# Prepare encrypted dir
self.encryption_handler = EncryptionHandler(sim_gid)
job_encrypted_inputs = HPCSchedulerClient()._prepare_input(op, sim_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.files_helper.get_project_folder(
self.test_project, '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,
HPCSimulatorAdapter.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 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)
encryption_handler = EncryptionHandler(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(
TimeseriesMetricsAdapterModel.__name__):
metric_vm_encrypted_file = encrypted_file
else:
simulation_results.append(encrypted_file)
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.file_handler.get_project_folder(
project, str(operation.id))
h5_filenames = EncryptionHandler(simulator_gid).decrypt_files_to_dir(
simulation_results, operation_dir)
return h5_filenames, metric_op, metric_file
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))
encyrption_handler = EncryptionHandler(simulator_gid)
_request_passfile(
simulator_gid, operation_id, base_url,
os.path.dirname(encyrption_handler.get_password_file()))
encyrption_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, None,
available_disk_space)
_encrypt_results(adapter_instance, encyrption_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)
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)
encryption_handler = EncryptionHandler(simulator_gid)
file_path = encryption_handler.get_password_file()
return serve_file(file_path, "application/x-download", "attachment",
os.path.basename(file_path))
def _stage_out_results(working_dir, simulator_gid):
# type: (Storage, typing.Union[uuid.UUID, str]) -> list
output_subfolder = HPCSchedulerClient.CSCS_DATA_FOLDER + '/' + HPCSimulatorAdapter.OUTPUT_FOLDER
output_list = HPCSchedulerClient._listdir(working_dir,
output_subfolder)
encryption_handler = EncryptionHandler(simulator_gid)
encrypted_dir = os.path.join(encryption_handler.get_encrypted_dir(),
HPCSimulatorAdapter.OUTPUT_FOLDER)
encrypted_files = HPCSchedulerClient._stage_out_outputs(
encrypted_dir, output_list)
LOGGER.info(working_dir.properties)
LOGGER.info(working_dir.listdir())
return encrypted_files
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)
encryption_handler = EncryptionHandler(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)
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.file_handler.get_project_folder(
project, str(operation.id))
h5_filenames = encryption_handler.decrypt_files_to_dir(
simulation_results, operation_dir)
encryption_handler.cleanup()
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,
EncryptionHandler(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 _project_key(project_id):
password_encryption_key = os.environ[
DataEncryptionHandler.
CRYPTO_PASS] if DataEncryptionHandler.CRYPTO_PASS in os.environ else None
if password_encryption_key is None:
raise TVBException("Password encryption key is not defined.")
project_keys_folder = os.path.join(TvbProfile.current.TVB_STORAGE,
DataEncryptionHandler.KEYS_FOLDER)
DataEncryptionHandler.fie_helper.check_created(project_keys_folder)
encrypted_project_key = DataEncryptionHandler.project_key_path(
project_id)
if os.path.exists(encrypted_project_key):
with open(encrypted_project_key, "rb") as fIn:
inputFileSize = stat(encrypted_project_key).st_size
pass_stream = BytesIO()
pyAesCrypt.decryptStream(fIn, pass_stream,
password_encryption_key, 64 * 1024,
inputFileSize)
project_key = pass_stream.getvalue().decode()
pass_stream.close()
return project_key
project_key = EncryptionHandler.generate_random_password(64)
with open(encrypted_project_key, "wb") as fOut:
pass_stream = BytesIO(str.encode(project_key))
pyAesCrypt.encryptStream(pass_stream, fOut,
password_encryption_key, 64 * 1024)
pass_stream.close()
return project_key
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 = EncryptionHandler(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))
encryption_handler = EncryptionHandler(simulator_gid)
encrypted_dir = os.path.join(encryption_handler.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
def setup_method(self):
self.files_helper = FilesHelper()
self.encryption_handler = EncryptionHandler('123')
self.clean_database()
self.test_user = TestFactory.create_user()
self.test_project = TestFactory.create_project(self.test_user)
class TestHPCSchedulerClient(BaseTestCase):
def setup_method(self):
self.files_helper = FilesHelper()
self.encryption_handler = EncryptionHandler('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.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.encryption_handler = EncryptionHandler(sim_gid)
job_encrypted_inputs = HPCSchedulerClient()._prepare_input(op, sim_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.files_helper.get_project_folder(
self.test_project, '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,
HPCSimulatorAdapter.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 = FilesHelper().get_project_folder(op.project, 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.encryption_handler.get_encrypted_dir()
if os.path.isdir(encrypted_dir):
shutil.rmtree(encrypted_dir)
passfile = self.encryption_handler.get_password_file()
if os.path.exists(passfile):
os.remove(passfile)
self.files_helper.remove_project_structure(self.test_project.name)
self.clean_database()
def test_encrypt_decrypt(self, dir_name, file_name):
# Generate a private key and public key
private_key = rsa.generate_private_key(
public_exponent=65537,
key_size=2048,
backend=default_backend()
)
public_key = private_key.public_key()
# Convert private key to bytes and save it so ABCUploader can use it
pem = private_key.private_bytes(
encoding=serialization.Encoding.PEM,
format=serialization.PrivateFormat.PKCS8,
encryption_algorithm=serialization.NoEncryption()
)
private_key_path = os.path.join(TvbProfile.current.TVB_TEMP_FOLDER, 'private_key.pem')
with open(private_key_path, 'wb') as f:
f.write(pem)
path_to_file = os.path.join(os.path.dirname(tvb_data.__file__), dir_name, file_name)
# Create model for ABCUploader
connectivity_model = ZIPConnectivityImporterModel()
# Generate password
pass_size = TvbProfile.current.hpc.CRYPT_PASS_SIZE
password = EncryptionHandler.generate_random_password(pass_size)
# Encrypt files using an AES symmetric key
encrypted_file_path = ABCUploader.get_path_to_encrypt(path_to_file)
buffer_size = TvbProfile.current.hpc.CRYPT_BUFFER_SIZE
pyAesCrypt.encryptFile(path_to_file, encrypted_file_path, password, buffer_size)
# Asynchronously encrypt the password used at the previous step for the symmetric encryption
password = str.encode(password)
encrypted_password = ABCUploader.encrypt_password(public_key, password)
# Save encrypted password
ABCUploader.save_encrypted_password(encrypted_password, TvbProfile.current.TVB_TEMP_FOLDER)
path_to_encrypted_password = os.path.join(TvbProfile.current.TVB_TEMP_FOLDER, ENCRYPTED_PASSWORD_NAME)
# Prepare model for decrypting
connectivity_model.uploaded = encrypted_file_path
connectivity_model.encrypted_aes_key = path_to_encrypted_password
TvbProfile.current.UPLOAD_KEY_PATH = os.path.join(TvbProfile.current.TVB_TEMP_FOLDER, 'private_key.pem')
# Decrypting
ABCUploader._decrypt_content(connectivity_model, 'uploaded')
decrypted_file_path = connectivity_model.uploaded.replace(ENCRYPTED_DATA_SUFFIX, DECRYPTED_DATA_SUFFIX)
with open(path_to_file, 'rb') as f_original:
with open(decrypted_file_path, 'rb') as f_decrypted:
while True:
original_content_chunk = f_original.read(buffer_size)
decrypted_content_chunk = f_decrypted.read(buffer_size)
assert original_content_chunk == decrypted_content_chunk,\
"Original and Decrypted chunks are not equal, so decryption hasn't been done correctly!"
# check if EOF was reached
if len(original_content_chunk) < buffer_size:
break
# Clean-up
os.remove(encrypted_file_path)
os.remove(decrypted_file_path)
os.remove(private_key_path)
os.remove(path_to_encrypted_password)