Source code for tvb.core.services.burst_service

# -*- coding: utf-8 -*-
#
#
# TheVirtualBrain-Framework Package. This package holds all Data Management, and
# Web-UI helpful to run brain-simulations. To use it, you also need to download
# TheVirtualBrain-Scientific Package (for simulators). See content of the
# documentation-folder for more details. See also http://www.thevirtualbrain.org
#
# (c) 2012-2023, Baycrest Centre for Geriatric Care ("Baycrest") and others
#
# This program is free software: you can redistribute it and/or modify it under the
# terms of the GNU General Public License as published by the Free Software Foundation,
# either version 3 of the License, or (at your option) any later version.
# This program is distributed in the hope that it will be useful, but WITHOUT ANY
# WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR A
# PARTICULAR PURPOSE.  See the GNU General Public License for more details.
# You should have received a copy of the GNU General Public License along with this
# program.  If not, see <http://www.gnu.org/licenses/>.
#
#
#   CITATION:
# When using The Virtual Brain for scientific publications, please cite it as explained here:
# https://www.thevirtualbrain.org/tvb/zwei/neuroscience-publications
#
#

import json
import os
from datetime import datetime

from tvb.basic.logger.builder import get_logger
from tvb.config import MEASURE_METRICS_MODULE, MEASURE_METRICS_CLASS
from tvb.core.entities.file.simulator.burst_configuration_h5 import BurstConfigurationH5
from tvb.core.entities.file.simulator.datatype_measure_h5 import DatatypeMeasureH5
from tvb.core.entities.file.simulator.view_model import SimulatorAdapterModel
from tvb.core.entities.generic_attributes import GenericAttributes
from tvb.core.entities.model.model_burst import BurstConfiguration
from tvb.core.entities.model.model_datatype import DataTypeGroup
from tvb.core.entities.model.model_operation import Operation, STATUS_FINISHED, STATUS_PENDING, STATUS_CANCELED
from tvb.core.entities.model.model_operation import OperationGroup, STATUS_ERROR, STATUS_STARTED, has_finished
from tvb.core.entities.storage import dao
from tvb.core.entities.transient.range_parameter import RangeParameter
from tvb.core.neocom import h5
from tvb.core.neocom.h5 import DirLoader
from tvb.core.services.import_service import ImportService
from tvb.core.utils import format_bytes_human, format_timedelta
from tvb.storage.storage_interface import StorageInterface

MAX_BURSTS_DISPLAYED = 50
STATUS_FOR_OPERATION = {
    STATUS_PENDING: BurstConfiguration.BURST_RUNNING,
    STATUS_STARTED: BurstConfiguration.BURST_RUNNING,
    STATUS_CANCELED: BurstConfiguration.BURST_CANCELED,
    STATUS_ERROR: BurstConfiguration.BURST_ERROR,
    STATUS_FINISHED: BurstConfiguration.BURST_FINISHED
}


[docs] class BurstService(object): LAUNCH_NEW = 'new' LAUNCH_BRANCH = 'branch' def __init__(self): self.logger = get_logger(self.__class__.__module__) self.storage_interface = StorageInterface()
[docs] def mark_burst_finished(self, burst_entity, burst_status=None, error_message=None, store_h5_file=True): """ Mark Burst status field. Also compute 'weight' for current burst: no of operations inside, estimate time on disk... :param burst_entity: BurstConfiguration to be updated, at finish time. :param burst_status: BurstConfiguration status. By default BURST_FINISHED :param error_message: If given, set the status to error and perpetuate the message. """ if burst_status is None: burst_status = BurstConfiguration.BURST_FINISHED if error_message is not None: burst_status = BurstConfiguration.BURST_ERROR try: # If there are any DataType Groups in current Burst, update their counter. burst_dt_groups = dao.get_generic_entity(DataTypeGroup, burst_entity.gid, "fk_parent_burst") for dt_group in burst_dt_groups: dt_group.count_results = dao.count_datatypes_in_group(dt_group.id) dt_group.disk_size, dt_group.subject = dao.get_summary_for_group(dt_group.id) dao.store_entity(dt_group) # Update actual Burst entity fields burst_entity.datatypes_number = dao.count_datatypes_in_burst(burst_entity.gid) burst_entity.status = burst_status burst_entity.error_message = error_message burst_entity.finish_time = datetime.now() dao.store_entity(burst_entity) if store_h5_file: self.store_burst_configuration(burst_entity) except Exception: self.logger.exception("Could not correctly update Burst status and meta-data!") burst_entity.status = burst_status burst_entity.error_message = "Error when updating Burst Status" burst_entity.finish_time = datetime.now() dao.store_entity(burst_entity) if store_h5_file: self.store_burst_configuration(burst_entity)
[docs] def persist_operation_state(self, operation, operation_status, message=None): """ Update Operation instance state. Store it in DB and on HDD/ :param operation: Operation instance :param operation_status: new status :param message: message in case of error :return: operation instance changed """ operation.mark_complete(operation_status, message) operation.queue_full = False operation = dao.store_entity(operation) # update burst also burst_config = self.get_burst_for_operation_id(operation.id) if burst_config is not None: burst_status = STATUS_FOR_OPERATION.get(operation_status) self.mark_burst_finished(burst_config, burst_status, message) return operation
[docs] @staticmethod def get_burst_for_operation_id(operation_id, is_group=False): return dao.get_burst_for_operation_id(operation_id, is_group)
[docs] def rename_burst(self, burst_id, new_name): """ Rename the burst given by burst_id, setting it's new name to burst_name. """ burst = dao.get_burst_by_id(burst_id) burst.name = new_name dao.store_entity(burst) self.store_burst_configuration(burst)
[docs] @staticmethod def get_available_bursts(project_id): """ Return all the burst for the current project. """ bursts = dao.get_bursts_for_project(project_id, page_size=MAX_BURSTS_DISPLAYED) or [] return bursts
[docs] @staticmethod def populate_burst_disk_usage(bursts): """ Adds a disk_usage field to each burst object. The disk usage is computed as the sum of the datatypes generated by a burst """ sizes = dao.compute_bursts_disk_size([b.gid for b in bursts]) for b in bursts: b.disk_size = format_bytes_human(sizes[b.gid])
[docs] def update_history_status(self, id_list): """ For each burst_id received in the id_list read new status from DB and return a list [id, new_status, is_group, message, running_time] tuple. """ result = [] for b_id in id_list: burst = dao.get_burst_by_id(b_id) if burst is not None: if burst.status == burst.BURST_RUNNING: running_time = datetime.now() - burst.start_time else: running_time = burst.finish_time - burst.start_time running_time = format_timedelta(running_time, most_significant2=False) if burst.status == burst.BURST_ERROR: msg = 'Check Operations page for error Message' else: msg = '' result.append([burst.id, burst.status, burst.is_group, msg, running_time]) else: self.logger.debug("Could not find burst with id=" + str(b_id) + ". Might have been deleted by user!!") return result
[docs] @staticmethod def update_simulation_fields(burst, op_simulation_id, simulation_gid): burst.fk_simulation = op_simulation_id burst.simulator_gid = simulation_gid.hex burst = dao.store_entity(burst) return burst
[docs] @staticmethod def load_burst_configuration(burst_config_id): # type: (int) -> BurstConfiguration burst_config = dao.get_burst_by_id(burst_config_id) return burst_config
[docs] @staticmethod def remove_burst_configuration(burst_config_id): # type: (int) -> None dao.remove_entity(BurstConfiguration, burst_config_id)
[docs] @staticmethod def prepare_burst_for_pse(burst_config): # type: (BurstConfiguration) -> (BurstConfiguration) operation_group = OperationGroup(burst_config.fk_project, ranges=burst_config.ranges) operation_group = dao.store_entity(operation_group) metric_operation_group = OperationGroup(burst_config.fk_project, ranges=burst_config.ranges) metric_operation_group = dao.store_entity(metric_operation_group) burst_config.operation_group = operation_group burst_config.fk_operation_group = operation_group.id burst_config.metric_operation_group = metric_operation_group burst_config.fk_metric_operation_group = metric_operation_group.id return dao.store_entity(burst_config)
[docs] @staticmethod def store_burst_configuration(burst_config): project = dao.get_project_by_id(burst_config.fk_project) bc_path = h5.path_for(burst_config.fk_simulation, BurstConfigurationH5, burst_config.gid, project.name) with BurstConfigurationH5(bc_path) as bc_h5: bc_h5.store(burst_config)
[docs] @staticmethod def load_burst_configuration_from_folder(simulator_folder, project): bc_h5_filename = DirLoader(simulator_folder, None).find_file_for_has_traits_type(BurstConfiguration) burst_config = BurstConfiguration(project.id) with BurstConfigurationH5(os.path.join(simulator_folder, bc_h5_filename)) as bc_h5: bc_h5.load_into(burst_config) return burst_config
[docs] @staticmethod def prepare_simulation_name(burst, project_id): simulation_number = dao.get_number_of_bursts(project_id) + 1 if burst.name is None: simulation_name = 'simulation_' + str(simulation_number) else: simulation_name = burst.name return simulation_name, simulation_number
[docs] def prepare_indexes_for_simulation_results(self, operation, result_filenames, burst): indexes = list() self.logger.debug("Preparing indexes for simulation results in operation {}...".format(operation.id)) for filename in result_filenames: try: self.logger.debug("Preparing index for filename: {}".format(filename)) index = h5.index_for_h5_file(filename)() h5_class = h5.REGISTRY.get_h5file_for_index(type(index)) with h5_class(filename) as index_h5: index.fill_from_h5(index_h5) index.fill_from_generic_attributes(index_h5.load_generic_attributes()) index.fk_parent_burst = burst.gid index.fk_from_operation = operation.id if operation.fk_operation_group: datatype_group = dao.get_datatypegroup_by_op_group_id(operation.fk_operation_group) self.logger.debug( "Found DatatypeGroup with id {} for operation {}".format(datatype_group.id, operation.id)) index.fk_datatype_group = datatype_group.id # Update the operation group name operation_group = dao.get_operationgroup_by_id(operation.fk_operation_group) operation_group.fill_operationgroup_name("TimeSeriesRegionIndex") dao.store_entity(operation_group) self.logger.debug( "Prepared index {} for file {} in operation {}".format(index.summary_info, filename, operation.id)) indexes.append(index) except Exception as e: self.logger.debug("Skip preparing index {} because there was an error.".format(filename)) self.logger.error(e) self.logger.debug("Prepared {} indexes for results in operation {}...".format(len(indexes), operation.id)) return indexes
[docs] def prepare_index_for_metric_result(self, operation, result_filename, burst): self.logger.debug("Preparing index for metric result in operation {}...".format(operation.id)) index = h5.index_for_h5_file(result_filename)() with DatatypeMeasureH5(result_filename) as dti_h5: index.gid = dti_h5.gid.load().hex index.metrics = json.dumps(dti_h5.metrics.load()) index.fk_source_gid = dti_h5.analyzed_datatype.load().hex index.fk_from_operation = operation.id index.fk_parent_burst = burst.gid datatype_group = dao.get_datatypegroup_by_op_group_id(operation.fk_operation_group) self.logger.debug("Found DatatypeGroup with id {} for operation {}".format(datatype_group.id, operation.id)) index.fk_datatype_group = datatype_group.id self.logger.debug("Prepared index {} for results in operation {}...".format(index.summary_info, operation.id)) return index
def _update_pse_burst_status(self, burst_config): operations_in_group = dao.get_operations_in_group(burst_config.fk_operation_group) if burst_config.fk_metric_operation_group: operations_in_group.extend(dao.get_operations_in_group(burst_config.fk_metric_operation_group)) operation_statuses = list() for operation in operations_in_group: if not has_finished(operation.status): self.logger.debug( 'Operation {} in group {} is not finished, burst status will not be updated'.format( operation.id, operation.fk_operation_group)) return operation_statuses.append(operation.status) self.logger.debug( 'All operations in burst {} have finished. Will update burst status'.format(burst_config.id)) if STATUS_ERROR in operation_statuses: self.mark_burst_finished(burst_config, BurstConfiguration.BURST_ERROR, 'Some operations in PSE have finished with errors') elif STATUS_CANCELED in operation_statuses: self.mark_burst_finished(burst_config, BurstConfiguration.BURST_CANCELED) else: self.mark_burst_finished(burst_config)
[docs] def update_burst_status(self, burst_config): if burst_config.fk_operation_group: self._update_pse_burst_status(burst_config) else: operation = dao.get_operation_by_id(burst_config.fk_simulation) message = operation.additional_info if len(message) == 0: message = None self.mark_burst_finished(burst_config, STATUS_FOR_OPERATION[operation.status], message)
[docs] @staticmethod def prepare_metrics_operation(operation): # TODO reuse from OperationService and do not duplicate logic here parent_burst = dao.get_generic_entity(BurstConfiguration, operation.fk_operation_group, 'fk_operation_group')[0] metric_operation_group_id = parent_burst.fk_metric_operation_group range_values = operation.range_values metric_algo = dao.get_algorithm_by_module(MEASURE_METRICS_MODULE, MEASURE_METRICS_CLASS) metric_operation = Operation(None, operation.fk_launched_by, operation.fk_launched_in, metric_algo.id, status=STATUS_FINISHED, op_group_id=metric_operation_group_id, range_values=range_values) metric_operation.visible = False metric_operation = dao.store_entity(metric_operation) op_dir = StorageInterface().get_project_folder(operation.project.name, str(metric_operation.id)) return op_dir, metric_operation
[docs] @staticmethod def get_range_param_by_name(param_name, all_range_parameters): for range_param in all_range_parameters: if param_name == range_param.name: return range_param return None
[docs] @staticmethod def handle_range_params_at_loading(burst_config, all_range_parameters): param1, param2 = None, None if burst_config.range1: param1 = RangeParameter.from_json(burst_config.range1) param1.fill_from_default(BurstService.get_range_param_by_name(param1.name, all_range_parameters)) if burst_config.range2 is not None: param2 = RangeParameter.from_json(burst_config.range2) param2.fill_from_default(BurstService.get_range_param_by_name(param2.name, all_range_parameters)) return param1, param2
[docs] def prepare_data_for_burst_copy(self, burst_config_id, burst_name_format, project): burst_config = self.load_burst_configuration(burst_config_id) burst_config_copy = burst_config.clone() count = dao.count_bursts_with_name(burst_config.name, burst_config.fk_project) burst_config_copy.name = burst_name_format.format(burst_config.name, count + 1) storage_path = self.storage_interface.get_project_folder(project.name, str(burst_config.fk_simulation)) simulator = h5.load_view_model(burst_config.simulator_gid, storage_path) simulator.generic_attributes = GenericAttributes() return simulator, burst_config_copy
[docs] @staticmethod def store_burst(burst_config): return dao.store_entity(burst_config)
[docs] def load_simulation_from_zip(self, zip_file, project): import_service = ImportService() simulator_folder = import_service.import_simulator_configuration_zip(zip_file) simulator_h5_filename = DirLoader(simulator_folder, None).find_file_for_has_traits_type(SimulatorAdapterModel) simulator_h5_filepath = os.path.join(simulator_folder, simulator_h5_filename) simulator = h5.load_view_model_from_file(simulator_h5_filepath) burst_config = self.load_burst_configuration_from_folder(simulator_folder, project) burst_config_copy = burst_config.clone() simulator.generic_attributes.parent_burst = burst_config_copy.gid return simulator, burst_config_copy, simulator_folder