Allow configuring object stores on a per-job-destination basis.

Setup a hierarchical objectstore with ids - these ids can be zero weight for objectstores that should only be used on a per-job-destination basis. Job destinations can then set object_store_id as a param to force a particular object store from the hierarchical configuration to be used. This can be set in job_conf.xml directly for static mappings or combined with dynamic job destinations for dynamic mapping to object stores. Dynamic job destinations can dispatch on cluster conditions, user, user preferences, etc.. Additionally, job resource parameters can be used with dynamic job destinations to allow user input on the object store to use.

An integration test case is included that tests and demonstrates all these different pieces. It sets up a Galaxy with a job and object store configuration that routes one tool to a specific disk store (called "static"), and routes another tool dynamically using job resource parameters that allow the user to pick between "slow, cheap" and "fast, expensive" storage options (for the test it is all just disk stores used - but you can imagine choosing between EBS and S3 or something this way).

I'd be happy to see something more formal than job resource parameters or user preferences developed to help users route job outputs to particular object stores - this is just a demonstration of what can be done today after these changes using extension points already being used by admins and users.

lib/galaxy/jobs/__init__.py

@@ -913,9 +913,9 @@ class JobWrapper(HasResourceParameters):
                             job.id)
 
     def _create_working_directory(self, job):
-        self.app.object_store.create(
+        self.object_store.create(
             job, base_dir='job_work', dir_only=True, obj_dir=True)
-        working_directory = self.app.object_store.get_filename(
+        working_directory = self.object_store.get_filename(
             job, base_dir='job_work', dir_only=True, obj_dir=True)
         return working_directory
 
@@ -928,10 +928,10 @@ class JobWrapper(HasResourceParameters):
                         self.working_directory)
             return
 
-        self.app.object_store.create(
+        self.object_store.create(
             job, base_dir='job_work', dir_only=True, obj_dir=True,
             extra_dir='_cleared_contents', extra_dir_at_root=True)
-        base = self.app.object_store.get_filename(
+        base = self.object_store.get_filename(
             job, base_dir='job_work', dir_only=True, obj_dir=True,
             extra_dir='_cleared_contents', extra_dir_at_root=True)
         date_str = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
@@ -1169,11 +1169,12 @@ class JobWrapper(HasResourceParameters):
 
     def enqueue(self):
         job = self.get_job()
-        self._set_object_store_ids(job)
         # Change to queued state before handing to worker thread so the runner won't pick it up again
         self.change_state(model.Job.states.QUEUED, flush=False, job=job)
         # Persist the destination so that the job will be included in counts if using concurrency limits
         self.set_job_destination(self.job_destination, None, flush=False, job=job)
+        # Set object store after job destination so can leverage parameters...
+        self._set_object_store_ids(job)
         self.sa_session.flush()
 
     def _set_object_store_ids(self, job):
@@ -1186,6 +1187,9 @@ class JobWrapper(HasResourceParameters):
             return
 
         object_store_populator = ObjectStorePopulator(self.app)
+        object_store_id = self.get_destination_configuration("object_store_id", None)
+        if object_store_id:
+            object_store_populator.object_store_id = object_store_id
 
         # Ideally we would do this without loading the actual job association
         # objects but change_state isn't yet optimized to do that so we need to
@@ -1316,7 +1320,7 @@ class JobWrapper(HasResourceParameters):
                     dataset.dataset.uuid = context['uuid']
                 # Update (non-library) job output datasets through the object store
                 if dataset not in job.output_library_datasets:
-                    self.app.object_store.update_from_file(dataset.dataset, create=True)
+                    self.object_store.update_from_file(dataset.dataset, create=True)
                 self.__update_output(job, dataset)
                 if not purged:
                     self._collect_extra_files(dataset.dataset, self.working_directory)
@@ -1326,7 +1330,7 @@ class JobWrapper(HasResourceParameters):
                         with NamedTemporaryFile() as temp_fh:
                             temp_fh.write(dataset.datatype.generate_primary_file(dataset))
                             temp_fh.flush()
-                            self.app.object_store.update_from_file(dataset.dataset, file_name=temp_fh.name, create=True)
+                            self.object_store.update_from_file(dataset.dataset, file_name=temp_fh.name, create=True)
                             dataset.set_size()
                     except Exception as e:
                         log.warning('Unable to generate primary composite file automatically for %s: %s', dataset.dataset.id, e)
@@ -1522,7 +1526,7 @@ class JobWrapper(HasResourceParameters):
             galaxy.tools.imp_exp.JobExportHistoryArchiveWrapper(self.job_id).cleanup_after_job(self.sa_session)
             galaxy.tools.imp_exp.JobImportHistoryArchiveWrapper(self.app, self.job_id).cleanup_after_job()
             if delete_files:
-                self.app.object_store.delete(self.get_job(), base_dir='job_work', entire_dir=True, dir_only=True, obj_dir=True)
+                self.object_store.delete(self.get_job(), base_dir='job_work', entire_dir=True, dir_only=True, obj_dir=True)
         except Exception:
             log.exception("Unable to cleanup job %d", self.job_id)
 
@@ -1537,7 +1541,7 @@ class JobWrapper(HasResourceParameters):
             for root, dirs, files in os.walk(temp_file_path):
                 extra_dir = root.replace(job_working_directory, '', 1).lstrip(os.path.sep)
                 for f in files:
-                    self.app.object_store.update_from_file(
+                    self.object_store.update_from_file(
                         dataset,
                         extra_dir=extra_dir,
                         alt_name=f,
@@ -1682,6 +1686,10 @@ class JobWrapper(HasResourceParameters):
                 return dp.dataset_id
         return None
 
+    @property
+    def object_store(self):
+        return self.app.object_store
+
     @property
     def tmp_dir_creation_statement(self):
         tmp_dir = self.get_destination_configuration("tmp_dir", None)
@@ -1815,7 +1823,7 @@ class JobWrapper(HasResourceParameters):
         if dataset not in job.output_library_datasets:
             purged = dataset.purged
             if not purged and not clean_only:
-                self.app.object_store.update_from_file(dataset, create=True)
+                self.object_store.update_from_file(dataset, create=True)
             else:
                 # If the dataset is purged and Galaxy is configured to write directly
                 # to the object store from jobs - be sure that file is cleaned up. This

lib/galaxy/objectstore/__init__.py

@@ -632,7 +632,7 @@ class DistributedObjectStore(NestedObjectStore):
     def create(self, obj, **kwargs):
         """The only method in which obj.object_store_id may be None."""
         if obj.object_store_id is None or not self.exists(obj, **kwargs):
-            if obj.object_store_id is None or obj.object_store_id not in self.weighted_backend_ids:
+            if obj.object_store_id is None or obj.object_store_id not in self.backends:
                 try:
                     obj.object_store_id = random.choice(self.weighted_backend_ids)
                 except IndexError:

test/integration/objectstore_selection_job_conf.xml

@@ -0,0 +1,47 @@
+<?xml version="1.0"?>
+<!-- 
+    - Add ready_for_resubmission to job_wrapper.
+    - Use in the handler.
+-->
+<job_conf>
+    <plugins>
+        <plugin id="local" type="runner" load="galaxy.jobs.runners.local:LocalJobRunner" workers="2"/>
+        <plugin id="dynamic" type="runner">
+            <param id="rules_module">integration.objectstore_selection_rules</param>
+        </plugin>
+    </plugins>
+
+    <handlers>
+        <handler id="main"/>
+    </handlers>
+
+    <destinations default="local">
+        <!-- Upload destination. -->
+        <destination id="local" runner="local">
+        </destination>
+
+        <!-- Static object store configuration per destination. -->
+        <destination id="static_object_store" runner="local">
+            <param id="object_store_id">static</param>
+        </destination>
+
+        <!-- Dynamic object store configuration - can dispatch on user, tool, destination, etc.. -->
+        <destination id="dynamic_object_store_destination" runner="dynamic">
+            <param id="type">python</param>
+            <param id="function">the_destination</param>
+        </destination>
+
+    </destinations>
+
+    <resources default="test">
+      <group id="upload"></group>
+      <group id="storage_params">how_store</group>
+    </resources>
+
+    <tools>
+        <tool id="upload1" destination="local" resources="upload" />
+        <tool id="multi_data_param" destination="static_object_store" />
+        <tool id="create_10" destination="dynamic_object_store_destination" resources="storage_params" />
+    </tools>
+
+</job_conf>

test/integration/objectstore_selection_job_resource_parameters_conf.xml

@@ -0,0 +1,6 @@
+<parameters>
+  <param label="How should this job's outputs be stored?" name="how_store" type="select">
+    <option selected="true" value="fast">Fast, expensive disk</option>
+    <option value="slow">Slower, cheaper disk</option>
+  </param>
+</parameters>

test/integration/objectstore_selection_rules/rules.py

@@ -0,0 +1,14 @@
+from galaxy.jobs import JobDestination
+
+
+def the_destination(resource_params):
+    job_destination = JobDestination()
+    job_destination.runner = "local"
+    how_store = resource_params.get("how_store", None)
+    # Retrieve answer from user about whether to store on fast or slow disk,
+    # translate to object store ID to give to job handling code.
+    object_store_id = "dynamic_ebs"
+    if how_store == "slow":
+        object_store_id = "dynamic_s3"
+    job_destination.params['object_store_id'] = object_store_id
+    return job_destination

test/integration/test_objectstore_selection.py

@@ -0,0 +1,128 @@
+"""Integration tests for object stores."""
+
+import os
+import string
+
+from base import integration_util  # noqa: I202
+from base.populators import (
+    DatasetPopulator,
+)
+
+SCRIPT_DIRECTORY = os.path.abspath(os.path.dirname(__file__))
+JOB_CONFIG_FILE = os.path.join(SCRIPT_DIRECTORY, "objectstore_selection_job_conf.xml")
+JOB_RESOURCE_PARAMETERS_CONFIG_FILE = os.path.join(SCRIPT_DIRECTORY, "objectstore_selection_job_resource_parameters_conf.xml")
+
+DISTRIBUTED_OBJECT_STORE_CONFIG_TEMPLATE = string.Template("""<?xml version="1.0"?>
+<object_store type="distributed" id="primary" order="0">
+    <backends>
+        <backend id="default" type="disk" weight="1">
+            <files_dir path="${temp_directory}/files_default"/>
+            <extra_dir type="temp" path="${temp_directory}/tmp_default"/>
+            <extra_dir type="job_work" path="${temp_directory}/job_working_directory_default"/>
+        </backend>
+        <backend id="static" type="disk" weight="0">
+            <files_dir path="${temp_directory}/files_static"/>
+            <extra_dir type="temp" path="${temp_directory}/tmp_static"/>
+            <extra_dir type="job_work" path="${temp_directory}/job_working_directory_static"/>
+        </backend>
+        <backend id="dynamic_ebs" type="disk" weight="0">
+            <files_dir path="${temp_directory}/files_dynamic_ebs"/>
+            <extra_dir type="temp" path="${temp_directory}/tmp_dynamic_ebs"/>
+            <extra_dir type="job_work" path="${temp_directory}/job_working_directory_dynamic_ebs"/>
+        </backend>
+        <backend id="dynamic_s3" type="disk" weight="0">
+            <files_dir path="${temp_directory}/files_dynamic_s3"/>
+            <extra_dir type="temp" path="${temp_directory}/tmp_dynamic_s3"/>
+            <extra_dir type="job_work" path="${temp_directory}/job_working_directory_dynamic_s3"/>
+        </backend>
+    </backends>
+</object_store>
+""")
+
+
+class ObjectStoreJobsIntegrationTestCase(integration_util.IntegrationTestCase):
+
+    framework_tool_and_types = True
+
+    @classmethod
+    def handle_galaxy_config_kwds(cls, config):
+        temp_directory = cls._test_driver.mkdtemp()
+        cls.object_stores_parent = temp_directory
+        for disk_store_file_name in ["files_default", "files_static", "files_dynamic_ebs", "files_dynamic_s3"]:
+            disk_store_path = os.path.join(temp_directory, disk_store_file_name)
+            os.makedirs(disk_store_path)
+            setattr(cls, "%s_path" % disk_store_file_name, disk_store_path)
+        config_path = os.path.join(temp_directory, "object_store_conf.xml")
+        with open(config_path, "w") as f:
+            f.write(DISTRIBUTED_OBJECT_STORE_CONFIG_TEMPLATE.safe_substitute({"temp_directory": temp_directory}))
+        config["object_store_config_file"] = config_path
+        config["job_config_file"] = JOB_CONFIG_FILE
+        config["job_resource_params_file"] = JOB_RESOURCE_PARAMETERS_CONFIG_FILE
+
+    def setUp(self):
+        super(ObjectStoreJobsIntegrationTestCase, self).setUp()
+        self.dataset_populator = DatasetPopulator(self.galaxy_interactor)
+
+    def _object_store_counts(self):
+        files_default_count = _files_count(self.files_default_path)
+        files_static_count = _files_count(self.files_static_path)
+        files_dynamic_count = _files_count(self.files_dynamic_path)
+        return files_default_count, files_static_count, files_dynamic_count
+
+    def _assert_file_counts(self, default, static, dynamic_ebs, dynamic_s3):
+        files_default_count = _files_count(self.files_default_path)
+        files_static_count = _files_count(self.files_static_path)
+        files_dynamic_ebs_count = _files_count(self.files_dynamic_ebs_path)
+        files_dynamic_s3_count = _files_count(self.files_dynamic_s3_path)
+        assert default == files_default_count
+        assert static == files_static_count
+        assert dynamic_ebs == files_dynamic_ebs_count
+        assert dynamic_s3 == files_dynamic_s3_count
+
+    def test_tool_simple_constructs(self):
+
+        with self.dataset_populator.test_history() as history_id:
+
+            def _run_tool(tool_id, inputs):
+                self.dataset_populator.run_tool(
+                    tool_id,
+                    inputs,
+                    history_id,
+                    assert_ok=True,
+                )
+                self.dataset_populator.wait_for_history(history_id)
+
+            self._assert_file_counts(0, 0, 0, 0)
+
+            hda1 = self.dataset_populator.new_dataset(history_id, content="1 2 3")
+            self.dataset_populator.wait_for_history(history_id)
+            hda1_input = {"src": "hda", "id": hda1["id"]}
+
+            # One file uploaded, added to default object store ID.
+            self._assert_file_counts(1, 0, 0, 0)
+
+            # should create two files in static object store.
+            _run_tool("multi_data_param", {"f1": hda1_input, "f2": hda1_input})
+            self._assert_file_counts(1, 2, 0, 0)
+
+            # should create two files in ebs object store.
+            create_10_inputs = {
+                "input1": hda1_input,
+                "input2": hda1_input,
+            }
+            _run_tool("create_10", create_10_inputs)
+            self._assert_file_counts(1, 2, 10, 0)
+
+            # should create 10 files in S3 object store.
+            create_10_inputs = {
+                "__job_resource|__job_resource__select": "yes",
+                "__job_resource|how_store": "slow",
+                "input1": hda1_input,
+                "input2": hda1_input,
+            }
+            _run_tool("create_10", create_10_inputs)
+            self._assert_file_counts(1, 2, 10, 10)
+
+
+def _files_count(directory):
+    return sum(len(files) for _, _, files in os.walk(directory))