Skip to content

Commit

Permalink
Implement method for JobCalculation to create a restart builder
Browse files Browse the repository at this point in the history
The `get_builder` method will return a "clean" builder for the given
`JobProcess`, without any inputs defined. Here we implement `get_builder_restart`
for the `JobCalculation` that will also return a `JobProcessBuilder` but with
the inputs and options pre-set with those that were used for the original node.
The builder can then immediately be resubmitted to run another instance of that
calculation, while of course giving the chance to change one or more inputs.

For example:

    calculation = load_node(<pk>)  # Some completed JobCalculation node
    builder = calculation.get_builder_restart()
    results = run(builder)
  • Loading branch information
sphuber committed Sep 18, 2018
1 parent b42decb commit 8a1e95d
Show file tree
Hide file tree
Showing 3 changed files with 82 additions and 3 deletions.
18 changes: 15 additions & 3 deletions aiida/backends/tests/work/job_processes.py
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
from aiida.orm.calculation.job.simpleplugins.templatereplacer import TemplatereplacerCalculation
from aiida.work.persistence import ObjectLoader
from aiida.work.job_processes import JobProcess
from aiida.work.process_builder import JobProcessBuilder
from aiida.work import Process

from . import utils
Expand Down Expand Up @@ -72,9 +73,20 @@ def tearDown(self):
self.runner = None
work.set_runner(None)

def test_class_loader(self):
cl = ObjectLoader()
TemplatereplacerProcess = JobProcess.build(TemplatereplacerCalculation)
def test_job_calculation_process(self):
"""Verify that JobCalculation.process returns a sub class of JobProcess with correct calculation class."""
process = TemplatereplacerCalculation.process()
self.assertTrue(issubclass(process, JobProcess))
self.assertEqual(process._calc_class, TemplatereplacerCalculation)

def test_job_calculation_get_builder(self):
"""Verify that JobCalculation.get_builder() returns an instance of JobProcessBuilder."""
process = TemplatereplacerCalculation.process()
builder = TemplatereplacerCalculation.get_builder()
self.assertTrue(isinstance(builder, JobProcessBuilder))

# Class objects are actually different memory instances so can't use assertEqual on simply instances
self.assertEqual(builder.process_class.__name__, process.__name__)

def test_job_process_set_label_and_description(self):
"""
Expand Down
17 changes: 17 additions & 0 deletions aiida/backends/tests/work/test_process_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -110,6 +110,23 @@ def test_dynamic_getters_doc_string(self):
self.assertEquals(builder.__class__.c.__doc__, str(TestWorkChain.spec().inputs['c']))
self.assertEquals(builder.c.__class__.d.__doc__, str(TestWorkChain.spec().inputs['c']['d']))

def test_job_calculation_get_builder_restart(self):
"""
Test the get_builder_restart method of JobCalculation class
"""
from aiida.orm.calculation.job import JobCalculation

original = JobCalculation()
original.set_option('resources', {'num_machines': 1, 'num_mpiprocs_per_machine': 1})
original.set_option('max_wallclock_seconds', 1800)
original.set_computer(self.computer)
original.label = 'original'
original.store()

builder = original.get_builder_restart()

self.assertDictEqual(builder.options, original.get_options())

def test_code_get_builder(self):
"""
Test that the get_builder method of Code returns a builder
Expand Down
50 changes: 50 additions & 0 deletions aiida/orm/implementation/general/calculation/job/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -112,13 +112,63 @@ def get_hash(self, ignore_errors=True, ignored_folder_content=('raw_input',), **

@classmethod
def process(cls):
"""
Return the JobProcess class constructed based on this JobCalculatin class
:return: JobProcess class
"""
from aiida.work.job_processes import JobProcess
return JobProcess.build(cls)

@classmethod
def get_builder(cls):
"""
Return a JobProcessBuilder instance, tailored for this calculation class
This builder is a mapping of the inputs of the JobCalculation class, supports tab-completion, automatic
validation when settings values as well as automated docstrings for each input
:return: JobProcessBuilder instance
"""
return cls.process().get_builder()

def get_builder_restart(self):
"""
Return a JobProcessBuilder instance, tailored for this calculation instance
This builder is a mapping of the inputs of the JobCalculation class, supports tab-completion, automatic
validation when settings values as well as automated docstrings for each input.
The fields of the builder will be pre-populated with all the inputs recorded for this instance as well as
settings all the options that were explicitly set for this calculation instance.
This builder can then directly be launched again to effectively run a duplicate calculation. But more useful
is that it serves as a starting point to, after changing one or more inputs, launch a similar calculation by
using this already completed calculation as a starting point.
:return: JobProcessBuilder instance
"""
from aiida.work.job_processes import JobProcess
from aiida.work.ports import PortNamespace

inputs = self.get_inputs_dict()
options = self.get_options()
builder = self.get_builder()

for port_name, port in self.process().spec().inputs.items():
if port_name == JobProcess.OPTIONS_INPUT_LABEL:
setattr(builder, port_name, options)
elif isinstance(port, PortNamespace):
namespace = port_name + '_'
sub = {link[len(namespace):]: node for link, node in inputs.items() if link.startswith(namespace)}
if sub:
setattr(builder, port_name, sub)
else:
if port_name in inputs:
setattr(builder, port_name, inputs[port_name])

return builder

def _init_internal_params(self):
"""
Define here internal parameters that should be defined right after the __init__
Expand Down

0 comments on commit 8a1e95d

Please sign in to comment.