How to properly handle (multiple) restarts?

[tomdemeyere]

Very often jobs will have to be restarted because of timeout mechanisms on HPC environments. This causes two issues when using Quacc:

1. Files from the timed-out job will often need to be copied over to the restarting jobs, this piece of code would not work because there would be no way to know where the previous timed-out directory is:

@flow
def my_flow(...):
   try:
       results = static_job(...)
   except SomeTimedOutError:
       results = static_job(restart_parameters = True, prev_dir = ???)

2. Changing parameters when restarting like in the code above. Some workflow engines seem to be able to do it (Covalent) although this requires input from the user.

Currently this is blocking us to go forward on a project where we decided to switch to Quacc/Parsl to perform the calculations (before asking on the parsl-help I preferred to open something here). Is this something Quacc should do or this is solely the responsibility of the workflow engines?

[Andrew-S-Rosen]

Hi @tomdemeyere. Thanks for your question.

This is definitely something I would like to be able to have in quacc directly, but I have not had time to sort out a mechanism that would achieve this. You're right that various workflow engines (e.g. Covalent, Prefect) do have options for restarts, including with parameter updates. Although even in those cases, as you alluded to, you'll still have file shuttling issues if you want to do a "continuation" rather than a full restart, which is normally the case.

I would be more than happy to consider any ideas you have in terms of a mechanism for this, ideally one that is generic (i.e. not tied to a specific recipe).

[Andrew-S-Rosen]

For timeout error, the ASE check_call can also take a time_out parameter, python will then throw a specific error if the program runs more than the specified time.

One challenge with this approach is that I believe the check_call is only used by the GenericFileIOCalculator and not the older calculator classes, which are still the most dominant in quacc.

At minimum, I don't think there would be any issue with catching the error and returning the original error plus key details, most importantly the temporary directory where the calculation is run. The current working directory of the calculation is normally reported by ASE when a calculation fails but of course won't be reported if there is a timeout.

With that in mind, one limitation: The actual restart will have to be taken care of by the user, by doing a custom flow.

I do think it will probably be default to have a clean way to do it all automatically in quacc. At least in the short-term, this might be the best approach.

One other challenge with trying to do everything restart-related in quacc is that it won't be effective for workflow engines. If Job 4 fails in a flow and you rerun it as a new job, the original task graph will never see that Job 4 successfully completes. Seems very messy to get right if dealing with it outside of the given workflow management solution, and of course they all behave differently. Not sure how to do it right...

[tomdemeyere]

I do think it will probably be default to have a clean way to do it all automatically in quacc. At least in the short-term, this might be the best approach.

Just thought about it: wouldn't it be possible to tweak the job decorator? The last example on this page (the retry_attempt code) is pretty much what I am thinking about:

https://rinaarts.com/declutter-python-code-with-error-handling-decorators/#:~:text=The%20simplest%20example%20is%20a,error%20makes%20absolutely%20no%20sense.

The user would send "restart_attempts: int" and "restart_param: dict", "restart_copy_files" to the decorator?

[Andrew-S-Rosen]

Preamble

It's an interesting idea. Now that I think about this some more, I'm starting to think that it might make sense to have most of this logic be handled by the workflow engine. That said, I have not been able to come up with a clean solution. Here are my scattered thoughts:

Challenges

Most, if not all, of the workflow managers offer support for retries where the job itself is simply reran. If a user wants just vanilla retry behavior, it seems like suggesting them to use any workflow engine would be the wiser option. There's no real benefit to adding that feature to quacc itself.

Where the workflow engine retry stuff is not sufficient is when some action needs to take place when the retry occurs. For instance, if you trigger a retry in Parsl/Prefect/Covalent/whatever, it simply reruns everything. This is fine if you genuinely want a clean restart, but it may not be fine if you actually want to continue where you left off. For instance, in some codes (e.g. VASP), you'll need to do something like mv CONTCAR POSCAR during the retry. Of course, there are also files you likely want to copy to the new job's working directory, like converged wavefunctions.

In principle, the user could make a new @job that does the file shuttling operation and then runs the job again with the given parameters. However, this is not appropriate because the @job in the prior @flow's task graph would never be updated, as this @job is entirely separate now. It's important that the failed @job is the one that is reran so any downstream tasks can then continue and so the provenance is retained.

So, that suggests restart_attempts may not be wise to add. The restart_copy_files also has problems because you would want to be able to specify everything when you launch the job the first time. The user can't reinstantiate the @job because that will be a new @job in an unrelated flow, as mentioned above. But the absolute paths of the files to copy on a restart are not known a priori, meaning the user doesn't have enough info to specify restart_copy_files when the @job is defined and invoked.

Proposal

With all this being said, I think the only way forward might be:

1. Ensure that the user can easily get the temporary working directory of any job that has failed. This is already done in the traceback when it is due to a calculator failure but is not always accessible when there is a timeout or hardware issue. Some workflow engines, like Covalent and FireWorks, do report the working directory of the job in the UI, so this is not a universal problem.
2. Have the workflow engine update the job's function signature on retry to include a value for copy_files that is a list of paths for files in the aforementioned working directory.

There are some minor but solvable problems with Point 2: a) not all workflow engines have that option (not necessarily a problem; they are open-source and PRs exist for this reason); b) some actions may need to take place besides simple copying (like the file renaming I mentioned with VASP). This latter point seems quite manageable to address though.

Sidenote

As an aside, from a cursory glance, it seems like several of the workflow engines support retrying with a custom function. For instance, in Parsl there is the concept of retry handlers. They are meant to introduce custom logic for whether a retry should be run, but (as suggested on Slack just now) it is possible that one could introduce logic in the retry_handler function to do some custom action like copy files or rename things. Prefect has basically the same thing, which they call the retry_condition_fn. However, even if you can inject custom logic into these functions, it's impossible to copy the files to the new job's working directory since it is not known until runtime. The only mechanism that could be used with quacc as-is right now is to update copy_files in the function signature on retry.

[tomdemeyere]

Interesting! I am speaking with Yadu Babuji from Parsl, funnily enough they pretty much have the opposite philosophy (not to be taken too seriously, maybe with more details they would have answered differently):

"Hi Tom, the retry_handler is designed to change the priority of tasks being retried, you can use it as a way to do additional things, but I have a feeling that if you have complex actions you need done, retry_handler isn't the way to go. Based on what you just said you want to be change params, that itself says you are no longer truly retrying a task but launching a new task."

The advise was then to use try, except in my code. I think this just show that this problem is very much on the line of what should be taken care of by the engine vs the actual code. I will trust your judgment on it!

[Andrew-S-Rosen]

😅 Guess the answer is ambiguous!

Yeah, I remain unsure. It's not that I wouldn't want to have logic in quacc per se but rather I'm not sure it's possible to add the logic in quacc itself since so much of this depends on the workflow engine. For instance, we don't actually want to launch a new task even if we're updating the function arguments because we don't want the DAG to change and the new task would then be disjointed from prior tasks in the workflow. For something like the grid phonon flow, if you ran one of those tasks separately since it failed, then the main phonon workflow would never complete.

[tomdemeyere]

Alright, whether Quacc is taking care of it or not we will probably need to do it for our project anyway. We will do that on a separate fork.

We will probably add error handling on the read_output as well with other parts of the code, if at some point you are interested in one / multiple specific parts ping me and I can open a PR on master.

Also: do you think the ASE people would be interested in seeing an optional timeout arg for Profile or should we just create our custom Profile?

[Andrew-S-Rosen]

We will probably add error handling on the read_output as well with other parts of the code, if at some point you are interested in one / multiple specific parts ping me and I can open a PR on master.

Please do share a link when you're done, if it ends up being public!

Also: do you think the ASE people would be interested in seeing an optional timeout arg for Profile or should we just create our custom Profile?

Since timeout is already in check_call as a kwarg, I think this would certainly be welcome. I see no issues, personally.