Execute certain jobs in different Conda environments

[JaGeo]

Hi @gpetretto ,

this is probably also something that we might need to solve on the jobflow level as well: do you think it might be feasible to execute certain jobs in different conda environments?

[gpetretto]

Hi @JaGeo, the first thing that comes to my mind would be to define different exec_config, each of which activating a different conda environment, and assign them to the different jobs. Would this be a solution for you?
Of course each environment would need to have the same version of jobflow-remote installed.

[JaGeo]

@gpetretto To me, this sounds like the best solution.

[JaGeo]

I am wondering if it makes sense to also add a conda environment keyword to the job submission. Similar to what you are proposing for #317

[gpetretto]

I think that the jobstore differs from the selection of the conda environment in this context. Mainly because the jobstore is a primary object in jobflow. But also because there would not be another option to set it, while the conda environment can be already handled in the exec_config, which is already an argument of submit_flow. Actually, I would not see any other way of implementing it, aside from updating the exec_config itself, which leaves margin for ambiguities.

However, to ease the selection of the conda environment I can propose the following changes:

• Add a conda attribute to the ExecutionConfig. If this is defined it will result in a line added to the submission script after the pre_run that looks like: f"conda activate {exec_config.conda}".
• aside from a str and an ExecutionConfig, allow also a dict as an argument of submit_flow, assuming that such a dictionary can be used to create an ExecutionConfig (e.g. `ExecutionConfig(**exec_config)).

This will allow to use something like this:

submit_flow(flow, exec_config={"conda": "my_other_environment"})

I suppose this is not as straightforward as submit_flow(flow, conda="my_other_environment"), but seems close enough.
Although it is not much better than

submit_flow(flow, exec_config=ExecutionConfig(pre_run="conda activate my_other_environment")

Which is already possible now...

What do you think?

[JaGeo]

Dear @gpetretto ,

as usual: thank you so much for the answer.

Do you see any other larger limitations of this approach?

My use case: incompatible force field installations as part of a workflow building upon atomate2 jobs.
I could likely only be able to run jobs in an environment different from the submitting environment, if I generate them in a Response, right?

[gpetretto]

Do you see any other larger limitations of this approach?

You mean using the exec_config to set a conda environment? One potential limitation that I could see, but likely not for the FF case, is that one might want to use multiple ExecutionConfig at the same time to combine the effect of different ones, in case one uses an exec_config to set up some environmental variable and avoiding many repetetions. I have been considering allowing to pass a list of exec_config instead of a single one, but I am afraid that this will lead to ambiguities to the order in which operations are applied.

My use case: incompatible force field installations as part of a workflow building upon atomate2 jobs. I could likely only be able to run jobs in an environment different from the submitting environment, if I generate them in a Response, right?

We are facing the same issue: needing to use different MLFF that do not fit in the same conda environement. We have created a "base" environment (e.g. test) , cloned it (e.g. test_uma, test_mace) and installed some force fields in each of those and using the exec_config to activate the different environment for the jobs requiring a specific MLFF.
I think that this is a rather involved solution, but given the limitation in compatibilities of the MLFFs, I cannot think of an easier one.