Skip to content

Commit

Permalink
👌 Make unique_extra_keys a pydantic field
Browse files Browse the repository at this point in the history
Add the `unique_extra_keys` field to the `BaseSubmissionContoller`. This makes
it possible to set the keys that used to keep track of which processes are
already run upon instantiation of the controller, instead of having to write
a new class for each set of extras.
  • Loading branch information
mbercx committed Dec 15, 2023
1 parent 98771d9 commit 6662ac8
Showing 1 changed file with 5 additions and 3 deletions.
8 changes: 5 additions & 3 deletions aiida_submission_controller/base.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
"""A prototype class to submit processes in batches, avoiding to submit too many."""
import abc
import logging
from typing import Optional

from aiida import engine, orm
from aiida.common import NotExistent
Expand Down Expand Up @@ -30,6 +31,8 @@ class BaseSubmissionController(BaseModel):
"""Label of the group to store the process nodes in."""
max_concurrent: int
"""Maximum concurrent active processes."""
unique_extra_keys: Optional[tuple]
"""List of keys defined in the extras that uniquely define each process to be run."""

_validate_group_exists = validator("group_label", allow_reuse=True)(validate_group_exists)

Expand Down Expand Up @@ -177,7 +180,7 @@ def submit_new_batch(self, dry_run=False, sort=True, verbose=False):
# Actually submit
wc_node = engine.submit(builder)

except (ValueError, TypeError) as exc:
except (ValueError, TypeError, AttributeError) as exc:
CMDLINE_LOGGER.error(f"Failed to submit work chain for extras <{workchain_extras}>: {exc}")
else:
CMDLINE_LOGGER.report(f"Submitted work chain <{wc_node}> for extras <{workchain_extras}>.")
Expand All @@ -188,10 +191,9 @@ def submit_new_batch(self, dry_run=False, sort=True, verbose=False):

return submitted

@abc.abstractmethod
def get_extra_unique_keys(self):
"""Return a tuple of the kes of the unique extras that will be used to uniquely identify your workchains."""
return
return self.unique_extra_keys

@abc.abstractmethod
def get_all_extras_to_submit(self):
Expand Down

0 comments on commit 6662ac8

Please sign in to comment.