def main():
batch_dir = os.path.basename(__file__).replace(".py", "").replace(
"run_", "flow_")
# intialize the BatchLauncher.
from pymatgen.io.abinit.launcher import BatchLauncher
batch = BatchLauncher(workdir=batch_dir)
# Build multiple flows and add them to the BatchLauncher.
for paral_kgb in range(2):
#flow_dir = os.path.join(batch.workdir, "flow_paral_kgb_%d" % paral_kgb)
# Get the SCF and the NSCF input and build the flow.
scf_input, nscf_input = make_scf_nscf_inputs(paral_kgb)
# Each flow will have a unique wordir inside batch.workdir.
# Note that we have to pass workdir=None and use set_name to specify the dir basename
flow = abilab.bandstructure_flow(None,
scf_input,
nscf_input,
allocate=False)
flow.set_name("flow_paral_kgb_%d" % paral_kgb)
batch.add_flow(flow)
# Submit to the queue in batch mode.
# Use abibatch.py to inspect the status or resubmit.
job = batch.submit()
print("batch.submit() returned: ", job.retcode)
return job.retcode
def test_batchlauncher(self):
"""Testing BatchLauncher methods."""
# Create the TaskManager.
manager = TaskManager.from_string(self.MANAGER)
print("batch_adapter", manager.batch_adapter)
assert manager.batch_adapter is not None
def build_flow_with_name(name):
"""Build a flow with workdir None and the given name."""
flow = Flow(workdir=None, manager=self.manager)
flow.set_name(name)
flow.register_task(self.fake_input)
work = Work()
work.register_scf_task(self.fake_input)
flow.register_work(work)
return flow
from pymatgen.io.abinit.launcher import BatchLauncher
tmpdir = tempfile.mkdtemp()
batch = BatchLauncher(workdir=tmpdir, manager=manager)
print(batch)
flow0 = build_flow_with_name("flow0")
flow1 = build_flow_with_name("flow1")
flow2_same_name = build_flow_with_name("flow1")
batch.add_flow(flow0)
# Cannot add the same flow twice.
with self.assertRaises(batch.Error):
batch.add_flow(flow0)
batch.add_flow(flow1)
# Cannot add two flows with the same name.
with self.assertRaises(batch.Error):
batch.add_flow(flow2_same_name)
batch.submit(dry_run=True)
for i, flow in enumerate([flow0, flow1]):
assert flow.workdir == os.path.join(batch.workdir, "flow%d" % i)
batch.pickle_dump()
batch_from_pickle = BatchLauncher.pickle_load(batch.workdir)
assert all(f1 == f2
for f1, f2 in zip(batch.flows, batch_from_pickle.flows))
