From e3d2944c4578fdbb1b64bc997f42cfbb5e751206 Mon Sep 17 00:00:00 2001 From: Timo Reents Date: Tue, 21 May 2024 18:22:40 +0200 Subject: [PATCH 1/2] First draft to enable multiple submissions per parent node Currently, one needs to specify the `unique_extras` to uniquely identify the parent nodes in the `parent_group`. This logic can cause some limitations, e.g. in the context of convergence studies. In such scenarios, one wants to submit several calculations (with different parameters) for a certain parent node. Here, the concept of `dynamic_extra`s is introduced. These extras are not used to identify a parent node but rather to identify a certain result, e.g. a `WorkChain` with certain parameters. Moreover, the `uuid` is used as the attribute that defines a parent node in the `FromGroupSubmissionController`. The `unique_extra_kyes` can still be used to take advantage of passing the extras between the parent node and the result node, but they become optional in this version. --- aiida_submission_controller/from_group.py | 49 ++++++++++++++--------- 1 file changed, 31 insertions(+), 18 deletions(-) diff --git a/aiida_submission_controller/from_group.py b/aiida_submission_controller/from_group.py index 1eea6a2..7145a59 100644 --- a/aiida_submission_controller/from_group.py +++ b/aiida_submission_controller/from_group.py @@ -3,7 +3,7 @@ from typing import Optional from aiida import orm -from pydantic import validator +from pydantic import Field, PrivateAttr, validator from .base import BaseSubmissionController, validate_group_exists @@ -15,6 +15,10 @@ class FromGroupSubmissionController(BaseSubmissionController): # pylint: disabl and define the abstract methods. """ + dynamic_extra: dict = Field(default_factory=dict) + """A dictionary of dynamic extras to be added to the extras of the process.""" + unique_extra_keys: tuple = Field(default_factory=tuple) + """List of keys defined in the extras that uniquely define each process to be run.""" parent_group_label: str """Label of the parent group from which to construct the process inputs.""" filters: Optional[dict] = None @@ -22,29 +26,32 @@ class FromGroupSubmissionController(BaseSubmissionController): # pylint: disabl order_by: Optional[dict] = None """Ordering applied to the query of the nodes in the parent group.""" + _dynamic_extra_keys: tuple = PrivateAttr(default_factory=tuple) + _dynamic_extra_values: tuple = PrivateAttr(default_factory=tuple) + _validate_group_exists = validator("parent_group_label", allow_reuse=True)(validate_group_exists) + def __init__(self, **kwargs): + """Initialize the instance.""" + super().__init__(**kwargs) + + if self.dynamic_extra: + self._dynamic_extra_keys, self._dynamic_extra_values = zip(*self.dynamic_extra.items()) + @property def parent_group(self): """Return the AiiDA ORM Group instance of the parent group.""" return orm.Group.objects.get(label=self.parent_group_label) - def get_parent_node_from_extras(self, extras_values): - """Return the Node instance (in the parent group) from the (unique) extras identifying it.""" - extras_projections = self.get_process_extra_projections() - assert len(extras_values) == len(extras_projections), f"The extras must be of length {len(extras_projections)}" - filters = dict(zip(extras_projections, extras_values)) + def get_extra_unique_keys(self): + """Return a tuple of the keys of the unique extras that will be used to uniquely identify your workchains.""" + # `_parent_uuid` will be replaced by the `uuid` attribute in the queries + combined_extras = ["_parent_uuid"] + list(self.unique_extra_keys) + list(self._dynamic_extra_keys) + return tuple(combined_extras) - qbuild = orm.QueryBuilder() - qbuild.append(orm.Group, filters={"id": self.parent_group.pk}, tag="group") - qbuild.append(orm.Node, project="*", filters=filters, tag="process", with_group="group") - qbuild.limit(2) - results = qbuild.all(flat=True) - if len(results) != 1: - raise ValueError( - "I would have expected only 1 result for extras={extras}, I found {'>1' if len(qbuild) else '0'}" - ) - return results[0] + def get_parent_node_from_extras(self, extras_values): + """Return the Node instance (in the parent group) from the `uuid` identifying it.""" + return orm.load_node(extras_values[0]) def get_all_extras_to_submit(self): """Return a *set* of the values of all extras uniquely identifying all simulations that you want to submit. @@ -57,11 +64,15 @@ def get_all_extras_to_submit(self): """ extras_projections = self.get_process_extra_projections() + # Use only the unique extras (and the parent uuid) to identify the processes to be submitted + if self._dynamic_extra_keys: + extras_projections = extras_projections[: -len(self._dynamic_extra_keys)] + qbuild = orm.QueryBuilder() qbuild.append(orm.Group, filters={"id": self.parent_group.pk}, tag="group") qbuild.append( orm.Node, - project=extras_projections, + project=["uuid"] + extras_projections[1:], # Replace `_parent_uuid` with `uuid` filters=self.filters, tag="process", with_group="group", @@ -76,10 +87,12 @@ def get_all_extras_to_submit(self): # First, however, convert to a list of tuples otherwise # the inner lists are not hashable results = [tuple(_) for _ in results] - for res in results: + for i, res in enumerate(results): assert all( extra is not None for extra in res ), "There is at least one of the nodes in the parent group that does not define one of the required extras." + results[i] = (*res, *self._dynamic_extra_values) # Add the dynamic extras to the results + results_set = set(results) assert len(results) == len(results_set), "There are duplicate extras in the parent group" From c70b418e275a03b7b1ce36e1cd7fa85a097d3007 Mon Sep 17 00:00:00 2001 From: Timo Reents Date: Fri, 24 May 2024 10:57:48 +0200 Subject: [PATCH 2/2] Make `unique_extra_keys` optional --- aiida_submission_controller/from_group.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/aiida_submission_controller/from_group.py b/aiida_submission_controller/from_group.py index 7145a59..36a53e7 100644 --- a/aiida_submission_controller/from_group.py +++ b/aiida_submission_controller/from_group.py @@ -17,8 +17,6 @@ class FromGroupSubmissionController(BaseSubmissionController): # pylint: disabl dynamic_extra: dict = Field(default_factory=dict) """A dictionary of dynamic extras to be added to the extras of the process.""" - unique_extra_keys: tuple = Field(default_factory=tuple) - """List of keys defined in the extras that uniquely define each process to be run.""" parent_group_label: str """Label of the parent group from which to construct the process inputs.""" filters: Optional[dict] = None @@ -46,7 +44,8 @@ def parent_group(self): def get_extra_unique_keys(self): """Return a tuple of the keys of the unique extras that will be used to uniquely identify your workchains.""" # `_parent_uuid` will be replaced by the `uuid` attribute in the queries - combined_extras = ["_parent_uuid"] + list(self.unique_extra_keys) + list(self._dynamic_extra_keys) + unique_extra_keys = self.unique_extra_keys or [] + combined_extras = ["_parent_uuid"] + list(unique_extra_keys) + list(self._dynamic_extra_keys) return tuple(combined_extras) def get_parent_node_from_extras(self, extras_values):