Modify ASE NEB to make it parallel via workflow engines

[tomdemeyere]

Not all calculators are compatible with the parallel version of ASE's NEBs. Similarly, it might be complicated to cleanly use mpi4py in some HPC environment. It might be possible to make it inherently parallel using workflow engine instead.

At some point in the NEB code there is something which looks like:

def get_forces(...)
...
    for image in images[1:-1]:
        image.get_forces()
...

If we could call a function with a @job decorator in the for loop without having to change the whole thing that would do the trick. However I understand that this is far from being trivial since that's not how workflows work. In my understanding this would require to manually resolve the futures after the for loop since the rest of the code depends on these results.

Writing the idea here, which unfortunately might not even be possible.

[Andrew-S-Rosen]

@tomdemeyere: In principle, the pattern you're asking about is definitely doable and something that would be added to quacc/recipes/common. In fact, you will find very similar patterns there.

What you are describing is a dynamic workflow where the DAG is not known until runtime. This can be achieved using the @subflow pattern. For a minimal example of what a @subflow is, see the callout box here. In essence, a @subflow as the name suggests is a sub-workflow. The @subflow generally involves some looping pattern where @jobs are called, but the length of the loop and/or items in the loop may or may not be known in advance. The return of a @subflow is a list of @job outputs (i.e. list of futures). In this case, I will note that the images (both the entries and its length) may not be known a priori if a @job comes before the proposed NEB workflow.

Most importantly, the @subflow pattern will automatically resolve the list[future] into a list[dict]. So, for instance, instead of having to do [future.result() for future in futures], which would screw things up downstream, you just do futures.result(). Check out the example in the docs to see what I mean.

This is very similar to the slab recipes. Check out quacc.recipes.emt.slabs.bulk_to_slabs_flow, which calls quacc.recipes.common.slabs.bulk_to_slabs_subflow. Like in the example you asked about, with the bulk-to-slabs workflow, we are: 1) Taking in an Atoms object (potentially from some prior @job); 2) Generating a bunch of slabs (images); 3) Running a relax job on each (instead of static job in the case of NEB); 4) Returning a list of results, potentially to be used elsewhere downstream.

[Andrew-S-Rosen]

The tricky thing in this scenario is not the iterative call to get_forces() but rather that NEB requires many calls to this subflow pattern for the flow to complete. If each @job is dispatched as a different Slurm job that may not all run at the same time, the workflow is going to go nowhere. In practice, what this means is that the "conventional" approach of launching each @job as a Slurm job is going to be highly impractical. You'd end up with hundreds or thousands of queued jobs over the course of the run. In practice, this kind of workflow would only make sense within the pilot job model.

[tomdemeyere]

[...] something that would be added to quacc/recipes/common

That's what I was thinking about!

Just to make things clear: the get_forces I am describing in the original post is well hidden behind many calls in ase.mep.NEB. Additionally, the function interact and modifies many attributes of the NEB class but also of the calculators and the Atoms objects such as self.results, self.optimizables etc... My doubts here were really about how to adapt that to Quacc without the need of re-writing the whole NEB class.

If you think that this is doable then I am really happy to work on it later on (when I have more time, and fixed the already existing issues with Espresso)

You'd end up with hundreds or thousands of queued jobs over the course of the run.

Mmmh we are probably not thinking about the same thing, the way I am thinking about it there can be only len(images) jobs running at the same time. Which can still be an issue if one job stays in queue much longer than the others of course, although there are other solutions that I see without the pilot job concept (interactive sessions, job arrays)

[Andrew-S-Rosen]

Just to make things clear: the get_forces I am describing in the original post is well hidden behind many calls in ase.mep.NEB. Additionally, the function interact and modifies many attributes of the NEB class but also of the calculators and the Atoms objects such as self.results, self.optimizables etc... My doubts here were really about how to adapt that to Quacc without the need of re-writing the whole NEB class.

Yeah, that was my main fear. It's not immediately obvious to me how one could go about this without rewriting the NEB class itself. I haven't given it much thought, and indeed I would be hesitant to simply copy/paste most of the code. Certainly open to suggestions if you have any. I have not really looked at the inner workings much.

Mmmh we are probably not thinking about the same thing, the way I am thinking about it there can be only len(images) jobs running at the same time. Which can still be an issue if one job stays in queue much longer than the others of course, although there are other solutions that I see without the pilot job concept (interactive sessions, job arrays)

If one wraps a @job decorator around a single-point calculation, then a new job will be launched for every image every time a force needs to be called. That's what I was thinking about, but I imagine it would be trivial to restructure it so it's one @job per image and each force call is really just the underlying function and not a @job. Maybe I'm misunderstanding that part though.

[tomdemeyere]

Giving this more thoughts (this would be very nice):

For basic NEB and DynNEB I was able to make sure that all Atoms.get_forces() and Atoms.get_potential_energy() calls are coming from this function: ase.mep.neb.BaseNEB.get_forces() (I placed assert statements in Atoms.get_...() using the inspect module to make sure of the caller)

Didn't try AutoNEB for now, as I don't know how to tame the beast.

The interesting bit in the NEB.get_forces() function:

        if not self.parallel:
            # Do all images - one at a time:
            for i in range(1, self.nimages - 1):
                forces[i - 1] = images[i].get_forces()
                energies[i] = images[i].get_potential_energy()

Can this be changed into calling a subflow that would dispatch all calculations in parallel?

But then later these are set to attributes:

        self.energies = energies
        self.real_forces = np.zeros((self.nimages, self.natoms, 3))
        self.real_forces[1:-1] = forces

Which means that we need to resolve the futures before this happens? I am not sure what would be the concurrent friendly way of doing this? I can imagine having a function somewhere else in Quacc that forces all workflow engines to stop and resolve, although that is probably a terrible concept.

If we can do this, it's pretty much done, we would just need to subclass this particular class where get_forces() belongs.

[Andrew-S-Rosen]

Can this be changed into calling a subflow that would dispatch all calculations in parallel?

If we were writing it ourselves, then yes, but it's not immediately obvious to me how one can do that without rewriting chunks of ASE. Maybe by monkeypatching the BaseNEB class? Not ideal, but if we have ample tests, it is not the end of the world.

Which means that we need to resolve the futures before this happens?

Yes, probably so. The challenge is that you'll now have futures floating around everywhere inside the ASE class. Messy.

I can imagine having a function somewhere else in Quacc that forces all workflow engines to stop and resolve, although that is probably a terrible concept.

Yeah, it's generally considered bad design to include blocking calls inside a function. I'd like to not promote that kind of behavior because it is tempting to do but often will break concurrency in unexpected ways.

If we can do this, it's pretty much done, we would just need to subclass this particular class where get_forces() belongs.

I don't doubt it can be done in some way, although because I have not put nearly as much time into thinking about it as you have, the path forward isn't clear to me yet.

[tomdemeyere]

Alright, I got it to work. Only the NEB.get_forces() function had to be copy-pasted from ASE, the function now call a "concurrent_calculate" which is a subflow that iterate through the images and launch static_jobs (which can be user provided.)

The catch is... I resolve the subflow manually not far after because the results are needed in many other places. I have tried it with Parsl and so far it seems to be quite robust. I will progressively run bigger calculations but I fail to see edge cases where this can fail.

This is very much a "the end justifies the means" solution and you will probably not like it at all. In any case I will let that on a custom branch "neb" on my fork.

[Andrew-S-Rosen]

Looks very cool! I'm glad to see you were able to get something to work. After all, that is the main goal of quacc in the end --- to provide the building blocks needed to write efficient workflows in a consistent manner. I'll hold off on considering this for merge at this time (since it's a fair bit complex), but let's keep this discussion open. I think you have a nice approach, and others may be interested.