From 31aeefd65f7fd311bf0ff62429ae9fd46a9d0be8 Mon Sep 17 00:00:00 2001 From: mvdbeek Date: Thu, 13 Mar 2025 16:22:03 +0100 Subject: [PATCH] Use materialized datasets in pulsar --- .../job_execution/compute_environment.py | 4 ++++ lib/galaxy/job_execution/datasets.py | 9 ++++++++ lib/galaxy/job_execution/setup.py | 23 +++++++++++++------ lib/galaxy/jobs/runners/pulsar.py | 4 +++- lib/galaxy/tools/evaluation.py | 8 ++----- test/integration/test_pulsar_embedded.py | 2 ++ 6 files changed, 36 insertions(+), 14 deletions(-) diff --git a/lib/galaxy/job_execution/compute_environment.py b/lib/galaxy/job_execution/compute_environment.py index de81b0d4e2f..bdd5b2ed96e 100644 --- a/lib/galaxy/job_execution/compute_environment.py +++ b/lib/galaxy/job_execution/compute_environment.py @@ -8,6 +8,7 @@ from typing import ( Dict, ) +from galaxy.job_execution.datasets import DeferrableObjectsT from galaxy.job_execution.setup import JobIO from galaxy.model import Job @@ -22,6 +23,9 @@ class ComputeEnvironment(metaclass=ABCMeta): compute server. """ + def __init__(self): + self.materialized_objects: Dict[str, DeferrableObjectsT] = {} + @abstractmethod def output_names(self): """Output unqualified filenames defined by job.""" diff --git a/lib/galaxy/job_execution/datasets.py b/lib/galaxy/job_execution/datasets.py index 69ed7a90224..529f54832a5 100644 --- a/lib/galaxy/job_execution/datasets.py +++ b/lib/galaxy/job_execution/datasets.py @@ -7,6 +7,15 @@ from abc import ( ABCMeta, abstractmethod, ) +from typing import Union + +from galaxy.model import ( + DatasetCollectionElement, + DatasetInstance, + HistoryDatasetCollectionAssociation, +) + +DeferrableObjectsT = Union[DatasetInstance, HistoryDatasetCollectionAssociation, DatasetCollectionElement] def dataset_path_rewrites(dataset_paths): diff --git a/lib/galaxy/job_execution/setup.py b/lib/galaxy/job_execution/setup.py index 433350ab3ae..135aab46c8d 100644 --- a/lib/galaxy/job_execution/setup.py +++ b/lib/galaxy/job_execution/setup.py @@ -22,6 +22,7 @@ from galaxy.files import ( from galaxy.job_execution.datasets import ( DatasetPath, DatasetPathRewriter, + DeferrableObjectsT, get_path_rewriter, ) from galaxy.model import ( @@ -148,7 +149,7 @@ class JobIO(UsesDictVisibleKeys): self._dataset_path_rewriter: Optional[DatasetPathRewriter] = None @property - def job(self): + def job(self) -> Job: return self.sa_session.get(Job, self.job_id) @classmethod @@ -217,11 +218,19 @@ class JobIO(UsesDictVisibleKeys): filenames.append(ds.dataset.extra_files_path) return filenames - def get_input_datasets(self) -> List[DatasetInstance]: + def get_input_datasets( + self, materialized_objects: Optional[Dict[str, DeferrableObjectsT]] = None + ) -> List[DatasetInstance]: job = self.job - return [ - da.dataset for da in job.input_datasets + job.input_library_datasets if da.dataset - ] # da is JobToInputDatasetAssociation object + datasets: List[DatasetInstance] = [] + for da in job.input_datasets + job.input_library_datasets: + if materialized_objects and da.name in materialized_objects: + materialized_object = materialized_objects[da.name] + if isinstance(materialized_object, DatasetInstance): + datasets.append(materialized_object) + elif da.dataset: + datasets.append(da.dataset) + return datasets def get_input_fnames(self) -> List[str]: filenames = [] @@ -229,9 +238,9 @@ class JobIO(UsesDictVisibleKeys): filenames.extend(self.get_input_dataset_fnames(ds)) return filenames - def get_input_paths(self) -> List[DatasetPath]: + def get_input_paths(self, materialized_objects: Optional[Dict[str, DeferrableObjectsT]]) -> List[DatasetPath]: paths = [] - for ds in self.get_input_datasets(): + for ds in self.get_input_datasets(materialized_objects): paths.append(self.get_input_path(ds)) return paths diff --git a/lib/galaxy/jobs/runners/pulsar.py b/lib/galaxy/jobs/runners/pulsar.py index a2bb8ff6c76..6a8bedec914 100644 --- a/lib/galaxy/jobs/runners/pulsar.py +++ b/lib/galaxy/jobs/runners/pulsar.py @@ -365,7 +365,9 @@ class PulsarJobRunner(AsynchronousJobRunner): output_names = compute_environment.output_names() client_inputs_list = [] - for input_dataset_wrapper in job_wrapper.job_io.get_input_paths(): + for input_dataset_wrapper in job_wrapper.job_io.get_input_paths( + compute_environment.materialized_objects + ): # str here to resolve false_path if set on a DatasetPath object. path = str(input_dataset_wrapper) object_store_ref = { diff --git a/lib/galaxy/tools/evaluation.py b/lib/galaxy/tools/evaluation.py index 8eb2129a501..e42a5f68f9b 100644 --- a/lib/galaxy/tools/evaluation.py +++ b/lib/galaxy/tools/evaluation.py @@ -13,7 +13,6 @@ from typing import ( List, Optional, TYPE_CHECKING, - Union, ) from packaging.version import Version @@ -21,6 +20,7 @@ from packaging.version import Version from galaxy import model from galaxy.authnz.util import provider_name_to_backend from galaxy.job_execution.compute_environment import ComputeEnvironment +from galaxy.job_execution.datasets import DeferrableObjectsT from galaxy.job_execution.setup import ensure_configs_directory from galaxy.model.deferred import ( materialize_collection_input, @@ -117,11 +117,6 @@ def global_tool_logs(func, config_file: str, action_str: str, tool: "Tool"): ) from e -DeferrableObjectsT = Union[ - model.DatasetInstance, model.HistoryDatasetCollectionAssociation, model.DatasetCollectionElement -] - - class ToolEvaluator: """An abstraction linking together a tool and a job runtime to evaluate tool inputs in an isolated, testable manner. @@ -168,6 +163,7 @@ class ToolEvaluator: # materialize deferred datasets materialized_objects = self._materialize_objects(deferred_objects, self.local_working_directory) + self.compute_environment.materialized_objects = materialized_objects # replace materialized objects back into tool input parameters self._replaced_deferred_objects(inp_data, incoming, materialized_objects) diff --git a/test/integration/test_pulsar_embedded.py b/test/integration/test_pulsar_embedded.py index c2d37c26527..b0962cdcc7f 100644 --- a/test/integration/test_pulsar_embedded.py +++ b/test/integration/test_pulsar_embedded.py @@ -23,6 +23,8 @@ instance = integration_util.integration_module_instance(EmbeddedPulsarIntegratio test_tools = integration_util.integration_tool_runner( [ + "cat_default", + "collection_nested_default", "collection_creates_dynamic_nested_from_json", "composite", "simple_constructs",