Merge pull request #9985 from mvdbeek/flush_once_per_mapped_over_execution

Flush once per mapped over execution

lib/galaxy/managers/collections.py

@@ -91,7 +91,8 @@ class DatasetCollectionManager(object):
 
     def create(self, trans, parent, name, collection_type, element_identifiers=None,
                elements=None, implicit_collection_info=None, trusted_identifiers=None,
-               hide_source_items=False, tags=None, copy_elements=False, history=None):
+               hide_source_items=False, tags=None, copy_elements=False, history=None,
+               set_hid=True, flush=True):
         """
         PRECONDITION: security checks on ability to add to parent
         occurred during load.
@@ -122,10 +123,10 @@ class DatasetCollectionManager(object):
             implicit_output_name = implicit_collection_info["implicit_output_name"]
 
         return self._create_instance_for_collection(
-            trans, parent, name, dataset_collection, implicit_inputs=implicit_inputs, implicit_output_name=implicit_output_name, tags=tags
+            trans, parent, name, dataset_collection, implicit_inputs=implicit_inputs, implicit_output_name=implicit_output_name, tags=tags, set_hid=set_hid, flush=flush,
         )
 
-    def _create_instance_for_collection(self, trans, parent, name, dataset_collection, implicit_output_name=None, implicit_inputs=None, tags=None, flush=True):
+    def _create_instance_for_collection(self, trans, parent, name, dataset_collection, implicit_output_name=None, implicit_inputs=None, tags=None, set_hid=True, flush=True):
         if isinstance(parent, model.History):
             dataset_collection_instance = self.model.HistoryDatasetCollectionAssociation(
                 collection=dataset_collection,
@@ -139,8 +140,9 @@ class DatasetCollectionManager(object):
                 dataset_collection_instance.implicit_output_name = implicit_output_name
 
             log.debug("Created collection with %d elements" % (len(dataset_collection_instance.collection.elements)))
-            # Handle setting hid
-            parent.add_dataset_collection(dataset_collection_instance)
+
+            if set_hid:
+                parent.add_dataset_collection(dataset_collection_instance)
 
         elif isinstance(parent, model.LibraryFolder):
             dataset_collection_instance = self.model.LibraryDatasetCollectionAssociation(

lib/galaxy/model/__init__.py

@@ -1765,9 +1765,10 @@ class History(HasTags, Dictifiable, UsesAnnotations, HasName, RepresentById):
             dataset = HistoryDatasetAssociation(dataset=dataset)
             object_session(self).add(dataset)
             object_session(self).flush()
-        elif not isinstance(dataset, HistoryDatasetAssociation):
+        elif not isinstance(dataset, (HistoryDatasetAssociation, HistoryDatasetCollectionAssociation)):
             raise TypeError("You can only add Dataset and HistoryDatasetAssociation instances to a history"
                             + " ( you tried to add %s )." % str(dataset))
+        is_dataset = is_hda(dataset)
         if parent_id:
             for data in self.datasets:
                 if data.id == parent_id:
@@ -1779,24 +1780,23 @@ class History(HasTags, Dictifiable, UsesAnnotations, HasName, RepresentById):
         else:
             if set_hid:
                 dataset.hid = self._next_hid()
-        if quota and self.user:
+        if quota and is_dataset and self.user:
             self.user.adjust_total_disk_usage(dataset.quota_amount(self.user))
         dataset.history = self
-        if genome_build not in [None, '?']:
+        if is_dataset and genome_build not in [None, '?']:
             self.genome_build = genome_build
         dataset.history_id = self.id
         return dataset
 
     def add_datasets(self, sa_session, datasets, parent_id=None, genome_build=None, set_hid=True, quota=True, flush=False):
         """ Optimized version of add_dataset above that minimizes database
-        interactions when adding many datasets to history at once.
+        interactions when adding many datasets and collections to history at once.
         """
-        all_hdas = all(is_hda(_) for _ in datasets)
-        optimize = len(datasets) > 1 and parent_id is None and all_hdas and set_hid
+        optimize = len(datasets) > 1 and parent_id is None and set_hid
         if optimize:
             self.__add_datasets_optimized(datasets, genome_build=genome_build)
             if quota and self.user:
-                disk_usage = sum([d.get_total_size() for d in datasets])
+                disk_usage = sum([d.get_total_size() for d in datasets if is_hda(d)])
                 self.user.adjust_total_disk_usage(disk_usage)
             sa_session.add_all(datasets)
             if flush:
@@ -1821,7 +1821,7 @@ class History(HasTags, Dictifiable, UsesAnnotations, HasName, RepresentById):
             dataset.hid = base_hid + i
             dataset.history = self
             dataset.history_id = cached_id(self)
-            if set_genome:
+            if set_genome and is_hda(dataset):
                 self.genome_build = genome_build
         return datasets
 
@@ -5300,6 +5300,8 @@ class WorkflowInvocation(UsesCreateAndUpdateTime, Dictifiable, RepresentById):
             for step in self.steps:
                 if step.workflow_step.type == 'tool':
                     for job in step.jobs:
+                        if job is None:
+                            continue
                         for step_input in step.workflow_step.input_connections:
                             output_step_type = step_input.output_step.type
                             if output_step_type in ['data_input', 'data_collection_input']:

lib/galaxy/model/tags.py

@@ -132,7 +132,9 @@ class TagHandler(object):
         """Delete tags from an item."""
         # Delete item-tag associations.
         for tag in item.tags:
-            self.sa_session.delete(tag)
+            if tag.id:
+                # Only can and need to delete tag if tag is persisted
+                self.sa_session.delete(tag)
         # Delete tags from item.
         del item.tags[:]
 

lib/galaxy/tools/__init__.py

@@ -1584,13 +1584,13 @@ class Tool(Dictifiable):
                         output_collections=execution_tracker.output_collections,
                         implicit_collections=execution_tracker.implicit_collections)
 
-    def handle_single_execution(self, trans, rerun_remap_job_id, execution_slice, history, execution_cache=None, completed_job=None, collection_info=None):
+    def handle_single_execution(self, trans, rerun_remap_job_id, execution_slice, history, execution_cache=None, completed_job=None, collection_info=None, flush_job=True):
         """
         Return a pair with whether execution is successful as well as either
         resulting output data or an error message indicating the problem.
         """
         try:
-            job, out_data = self.execute(
+            rval = self.execute(
                 trans,
                 incoming=execution_slice.param_combination,
                 history=history,
@@ -1599,7 +1599,13 @@ class Tool(Dictifiable):
                 dataset_collection_elements=execution_slice.dataset_collection_elements,
                 completed_job=completed_job,
                 collection_info=collection_info,
+                flush_job=flush_job,
             )
+            job = rval[0]
+            out_data = rval[1]
+            if len(rval) == 4:
+                execution_slice.datasets_to_persist = rval[2]
+                execution_slice.history = rval[3]
         except (webob.exc.HTTPFound, exceptions.MessageException) as e:
             # if it's a webob redirect exception, pass it up the stack
             raise e

lib/galaxy/tools/actions/__init__.py

@@ -33,6 +33,7 @@ class ToolExecutionCache(object):
         self.trans = trans
         self.current_user_roles = trans.get_current_user_roles()
         self.chrom_info = {}
+        self.cached_collection_elements = {}
 
     def get_chrom_info(self, tool_id, input_dbkey):
         genome_builds = self.trans.app.genome_builds
@@ -276,7 +277,7 @@ class DefaultToolAction(object):
                         preserved_tags[tag.value] = tag
         return history, inp_data, inp_dataset_collections, preserved_tags, all_permissions
 
-    def execute(self, tool, trans, incoming=None, return_job=False, set_output_hid=True, history=None, job_params=None, rerun_remap_job_id=None, execution_cache=None, dataset_collection_elements=None, completed_job=None, collection_info=None):
+    def execute(self, tool, trans, incoming=None, return_job=False, set_output_hid=True, history=None, job_params=None, rerun_remap_job_id=None, execution_cache=None, dataset_collection_elements=None, completed_job=None, collection_info=None, flush_job=True):
         """
         Executes a tool, creating job and tool outputs, associating them, and
         submitting the job to the job queue. If history is not specified, use
@@ -381,6 +382,7 @@ class DefaultToolAction(object):
                     inp_dataset_collections,
                     input_ext,
                     python_template_version=tool.python_template_version,
+                    execution_cache=execution_cache,
                 )
                 create_datasets = True
                 dataset = None
@@ -408,7 +410,7 @@ class DefaultToolAction(object):
                     dataset_collection_elements[name].hda = data
                 trans.sa_session.add(data)
                 if not completed_job:
-                    trans.app.security_agent.set_all_dataset_permissions(data.dataset, output_permissions, new=True)
+                    trans.app.security_agent.set_all_dataset_permissions(data.dataset, output_permissions, new=True, flush=False)
             data.copy_tags_to(preserved_tags)
 
             if not completed_job and trans.app.config.legacy_eager_objectstore_initialization:
@@ -454,6 +456,7 @@ class DefaultToolAction(object):
             # Flush all datasets at once.
             return data
 
+        datasets_to_persist = []
         for name, output in tool.outputs.items():
             if not filter_output(output, incoming):
                 handle_output_timer = ExecutionTimer()
@@ -461,7 +464,6 @@ class DefaultToolAction(object):
                     collections_manager = app.dataset_collections_service
                     element_identifiers = []
                     known_outputs = output.known_outputs(input_collections, collections_manager.type_registry)
-                    created_element_datasets = []
                     # Just to echo TODO elsewhere - this should be restructured to allow
                     # nested collections.
                     for output_part_def in known_outputs:
@@ -488,7 +490,7 @@ class DefaultToolAction(object):
 
                         effective_output_name = output_part_def.effective_output_name
                         element = handle_output(effective_output_name, output_part_def.output_def, hidden=True)
-                        created_element_datasets.append(element)
+                        datasets_to_persist.append(element)
                         # TODO: this shouldn't exist in the top-level of the history at all
                         # but for now we are still working around that by hiding the contents
                         # there.
@@ -500,17 +502,20 @@ class DefaultToolAction(object):
                             "name": output_part_def.element_identifier,
                         })
 
-                    history.add_datasets(trans.sa_session, created_element_datasets, set_hid=set_output_hid, quota=False, flush=True)
                     if output.dynamic_structure:
                         assert not element_identifiers  # known_outputs must have been empty
                         element_kwds = dict(elements=collections_manager.ELEMENTS_UNINITIALIZED)
                     else:
                         element_kwds = dict(element_identifiers=element_identifiers)
-                    output_collections.create_collection(
+                    hdca = output_collections.create_collection(
                         output=output,
                         name=name,
+                        set_hid=True if flush_job else False,
+                        flush=flush_job,
                         **element_kwds
                     )
+                    if hdca:
+                        datasets_to_persist.append(hdca)
                     log.info("Handled collection output named %s for tool %s %s" % (name, tool.id, handle_output_timer))
                 else:
                     handle_output(name, output)
@@ -521,13 +526,9 @@ class DefaultToolAction(object):
             'Added output datasets to history',
         )
         # Add all the top-level (non-child) datasets to the history unless otherwise specified
-        datasets_to_persist = []
         for name, data in out_data.items():
             if name not in child_dataset_names and name not in incoming:  # don't add children; or already existing datasets, i.e. async created
                 datasets_to_persist.append(data)
-        # Set HID and add to history.
-        # This is brand new and certainly empty so don't worry about quota.
-        history.add_datasets(trans.sa_session, datasets_to_persist, set_hid=set_output_hid, quota=False, flush=False)
 
         # Add all the children to their parents
         for parent_name, child_name in parent_to_child_pairs:
@@ -550,11 +551,16 @@ class DefaultToolAction(object):
         # Now that we have a job id, we can remap any outputs if this is a rerun and the user chose to continue dependent jobs
         # This functionality requires tracking jobs in the database.
         if app.config.track_jobs_in_database and rerun_remap_job_id is not None:
+            # We need a flush here and get hids in order to rewrite jobs parameter,
+            # but remapping jobs should only affect single jobs anyway, so this is not too costly.
+            trans.sa_session.flush()
+            history.add_datasets(trans.sa_session, datasets_to_persist, set_hid=set_output_hid, quota=False, flush=False)
             self._remap_job_on_rerun(trans=trans,
                                      galaxy_session=galaxy_session,
                                      rerun_remap_job_id=rerun_remap_job_id,
                                      current_job=job,
                                      out_data=out_data)
+            datasets_to_persist = []
         log.info("Setup for job %s complete, ready to be enqueued %s" % (job.log_str(), job_setup_timer))
 
         # Some tools are not really executable, but jobs are still created for them ( for record keeping ).
@@ -578,10 +584,18 @@ class DefaultToolAction(object):
             trans.sa_session.flush()
             trans.response.send_redirect(url_for(controller='tool_runner', action='redirect', redirect_url=redirect_url))
         else:
-            # Dispatch to a job handler. enqueue() is responsible for flushing the job
-            app.job_manager.enqueue(job, tool=tool)
-            trans.log_event("Added job to the job queue, id: %s" % str(job.id), tool_id=job.tool_id)
-            return job, out_data
+            if flush_job:
+                # Set HID and add to history.
+                # This is brand new and certainly empty so don't worry about quota.
+                history.add_datasets(trans.sa_session, datasets_to_persist, set_hid=set_output_hid, quota=False, flush=False)
+                job_flush_timer = ExecutionTimer()
+                trans.sa_session.flush()
+                log.info("Flushed transaction for job %s %s" % (job.log_str(), job_flush_timer))
+
+                # Dispatch to a job handler. enqueue() is responsible for flushing the job
+                app.job_manager.enqueue(job, tool=tool)
+                trans.log_event("Added job to the job queue, id: %s" % str(job.id), tool_id=job.tool_id)
+            return job, out_data, datasets_to_persist, history
 
     def _remap_job_on_rerun(self, trans, galaxy_session, rerun_remap_job_id, current_job, out_data):
         """
@@ -804,7 +818,7 @@ class OutputCollections(object):
         self.out_collection_instances = {}
         self.tags = tags
 
-    def create_collection(self, output, name, collection_type=None, **element_kwds):
+    def create_collection(self, output, name, collection_type=None, set_hid=True, flush=True, **element_kwds):
         input_collections = self.input_collections
         collections_manager = self.trans.app.dataset_collections_service
         collection_type = collection_type or output.structure.collection_type
@@ -883,11 +897,14 @@ class OutputCollections(object):
                 collection_type=collection_type,
                 trusted_identifiers=True,
                 tags=self.tags,
+                set_hid=set_hid,
+                flush=flush,
                 **element_kwds
             )
             # name here is name of the output element - not name
             # of the hdca.
             self.out_collection_instances[name] = hdca
+            return hdca
 
 
 def on_text_for_names(input_names):
@@ -934,7 +951,7 @@ def get_ext_or_implicit_ext(hda):
     return hda.ext
 
 
-def determine_output_format(output, parameter_context, input_datasets, input_dataset_collections, random_input_ext, python_template_version='3'):
+def determine_output_format(output, parameter_context, input_datasets, input_dataset_collections, random_input_ext, python_template_version='3', execution_cache=None):
     """ Determines the output format for a dataset based on an abstract
     description of the output (galaxy.tool_util.parser.ToolOutput), the parameter
     wrappers, a map of the input datasets (name => HDA), and the last input
@@ -975,7 +992,13 @@ def determine_output_format(output, parameter_context, input_datasets, input_dat
                     try:
                         input_element = input_collection_collection[element_index]
                     except KeyError:
-                        for element in input_collection_collection.dataset_elements:
+                        if execution_cache:
+                            dataset_elements = execution_cache.cached_collection_elements.get(input_collection_collection.id)
+                            if dataset_elements is None:
+                                dataset_elements = execution_cache.cached_collection_elements[input_collection_collection.id] = input_collection_collection.dataset_elements
+                        else:
+                            dataset_elements = input_collection_collection.dataset_elements
+                        for element in dataset_elements:
                             if element.element_identifier == element_index:
                                 input_element = element
                                 break

lib/galaxy/tools/actions/data_manager.py

@@ -10,7 +10,7 @@ class DataManagerToolAction(DefaultToolAction):
 
     def execute(self, tool, trans, **kwds):
         rval = super(DataManagerToolAction, self).execute(tool, trans, **kwds)
-        if isinstance(rval, tuple) and len(rval) == 2 and isinstance(rval[0], trans.app.model.Job):
+        if isinstance(rval, tuple) and len(rval) >= 2 and isinstance(rval[0], trans.app.model.Job):
             assoc = trans.app.model.DataManagerJobAssociation(job=rval[0], data_manager_id=tool.data_manager_id)
             trans.sa_session.add(assoc)
             trans.sa_session.flush()

lib/galaxy/tools/execute.py

@@ -68,7 +68,7 @@ def execute(trans, tool, mapping_params, history, rerun_remap_job_id=None, colle
             del params['__workflow_resource_params__']
         if validate_outputs:
             params['__validate_outputs__'] = True
-        job, result = tool.handle_single_execution(trans, rerun_remap_job_id, execution_slice, history, execution_cache, completed_job, collection_info)
+        job, result = tool.handle_single_execution(trans, rerun_remap_job_id, execution_slice, history, execution_cache, completed_job, collection_info, flush_job=False)
         if job:
             log.debug(job_timer.to_str(tool_id=tool.id, job_id=job.id))
             execution_tracker.record_success(execution_slice, job, result)
@@ -93,6 +93,7 @@ def execute(trans, tool, mapping_params, history, rerun_remap_job_id=None, colle
 
     jobs_executed = 0
     has_remaining_jobs = False
+    datasets_to_persist = []
 
     for i, execution_slice in enumerate(execution_tracker.new_execution_slices()):
         if max_num_jobs and jobs_executed >= max_num_jobs:
@@ -100,6 +101,19 @@ def execute(trans, tool, mapping_params, history, rerun_remap_job_id=None, colle
             break
         else:
             execute_single_job(execution_slice, completed_jobs[i])
+            if execution_slice.datasets_to_persist:
+                datasets_to_persist.extend(execution_slice.datasets_to_persist)
+
+    if datasets_to_persist:
+        execution_slice.history.add_datasets(trans.sa_session, datasets_to_persist, set_hid=True, quota=False, flush=False)
+        # a side effect of history.add_datasets is a commit within db_next_hid (even with flush=False).
+    else:
+        # Make sure collections, implicit jobs etc are flushed even if there are no precreated output datasets
+        trans.sa_session.flush()
+    for job in execution_tracker.successful_jobs:
+        # Put the job in the queue if tracking in memory
+        tool.app.job_manager.enqueue(job, tool=tool)
+        trans.log_event("Added job to the job queue, id: %s" % str(job.id), tool_id=job.tool_id)
 
     if has_remaining_jobs:
         raise PartialJobExecution(execution_tracker)
@@ -116,6 +130,8 @@ class ExecutionSlice(object):
         self.job_index = job_index
         self.param_combination = param_combination
         self.dataset_collection_elements = dataset_collection_elements
+        self.datasets_to_persist = None
+        self.history = None
 
 
 class ExecutionTracker(object):
@@ -361,7 +377,7 @@ class ExecutionTracker(object):
             job_assoc = model.ImplicitCollectionJobsJobAssociation()
             job_assoc.order_index = execution_slice.job_index
             job_assoc.implicit_collection_jobs = implicit_collection_jobs
-            job_assoc.job_id = job.id
+            job_assoc.job = job
             self.trans.sa_session.add(job_assoc)
 
 

lib/galaxy/workflow/modules.py

@@ -1802,17 +1802,16 @@ class ToolModule(WorkflowModule):
 
         # Combine workflow and runtime post job actions into the effective post
         # job actions for this execution.
-        flush_required = False
         effective_post_job_actions = self._effective_post_job_actions(step)
         for pja in effective_post_job_actions:
             if pja.action_type in ActionBox.immediate_actions or isinstance(self.tool, DatabaseOperationTool):
                 ActionBox.execute(self.trans.app, self.trans.sa_session, pja, job, replacement_dict)
             else:
-                pjaa = model.PostJobActionAssociation(pja, job_id=job.id)
+                if job.id:
+                    pjaa = model.PostJobActionAssociation(pja, job_id=job.id)
+                else:
+                    pjaa = model.PostJobActionAssociation(pja, job=job)
                 self.trans.sa_session.add(pjaa)
-                flush_required = True
-        if flush_required:
-            self.trans.sa_session.flush()
 
     def __restore_step_meta_runtime_state(self, step_runtime_state):
         if RUNTIME_POST_JOB_ACTIONS_KEY in step_runtime_state:

test/unit/tools/test_actions.py

@@ -134,12 +134,13 @@ class DefaultToolActionTestCase(unittest.TestCase, tools_support.UsesApp, tools_
         if incoming is None:
             incoming = dict(param1="moo")
         self._init_tool(contents)
-        return self.action.execute(
+        job, out_data, _, _ = self.action.execute(
             tool=self.tool,
             trans=self.trans,
             history=self.history,
             incoming=incoming,
         )
+        return job, out_data
 
 
 def test_determine_output_format():

test/unit/tools/test_execution.py

@@ -211,6 +211,9 @@ class MockTrans(object):
     def get_current_user_roles(self):
         return []
 
+    def log_event(self, *args, **kwds):
+        pass
+
 
 class MockCollectionService(object):