Feature request: atomate2 worker config example

[Youjin1985]

Hi!
Could you please provide an example of config how to make the following code working?

In [1]: from atomate2.vasp.flows.core import RelaxMaker
   ...: from atomate2.vasp.powerups import add_metadata_to_flow
   ...: from jobflow_remote import submit_flow
   ...: from pymatgen.core import Structure
   ...:
   ...: structure = Structure.from_file('RbB4O6F.cif')
   ...: relax_flow = RelaxMaker().make(structure)
   ...: relax_flow = add_metadata_to_flow(
   ...:     flow=relax_flow,
   ...:     additional_fields={"camd_id": "camd-0"},
   ...: )
   ...:
   ...: submit_flow(
   ...:     relax_flow, worker="example_worker"
   ...: )

First, I am not sure where I should put atomate2 config.
I use (user/runner) + (worker) setup.
On both user/runner and worker I have export ATOMATE2_CONFIG_FILE pointing to atomate2 yaml config file which contains
VASP_CMD: srun /public/software/apps/vasp/intelmpi/5.4.4/bin/vasp_std

It does not work, automatically produced submit.sh looks like

#!/bin/bash

#SBATCH --job-name=relax
#SBATCH --output=/public/home/XXX/.jobflow_workdir/19/fc/3b/19fc3b0f-0958-4220-80df-470728f52250_1/queue.out
#SBATCH --error=/public/home/XXX/.jobflow_workdir/19/fc/3b/19fc3b0f-0958-4220-80df-470728f52250_1/queue.err
cd /public/home/XXX/.jobflow_workdir/19/fc/3b/19fc3b0f-0958-4220-80df-470728f52250_1
source /public/XXX/xtalgo/.bashrc
 conda deactivate
conda activate jobflow
module purge
source /public/software/profile.d/compiler_intel-2021.3.0.sh
source /public/software/profile.d/mpi_intelmpi-2021.3.0.sh
export MKL_DEBUG_CPU_TYPE=5
export I_MPI_PMI_LIBRARY=/opt/sothisai/slurm/lib/libpmi.so
export PATH=/public/software/apps/vasp/intelmpi/5.4.4/hpcx-2.4.1-intel2017:${{PATH}}
jf -fe execution run /public/home/XXX/.jobflow_workdir/19/fc/3b/19fc3b0f-0958-4220-80df-470728f52250_1
``` - vasp command is not inserted.
Thank you!

[JaGeo]

Hi @Youjin1985 ,

Thanks for raising this question.

Could you clarify how you export your atomate.yaml files? It might be that you have to add the export directly to your jobflow-remote config or that you need to source the bashrc.

[Youjin1985]

Hi @JaGeo !
I am sourcing ~/.bashrc of HPC in pre_run of my jobflow-remote yaml configuration at runner/worker machine.
I've found that it works and vasp relaxation is running, but I am not sure should I setup vasp command on runner/worker or HPC is enough. When I make sure, I'll be glad to provide working example right there.

[JaGeo]

Ah. I think this should be enough. Please let us know if it continues to work.

We can then surely add more info to the jobflow-remote documentation of atomate2 to clarify this.

Or should we add more here, @gpetretto ? (Btw, i did not realize until now that this issue is part the jobflow-remote issues and thought it was posted in the atomate2 issues)

[gpetretto]

Hi @Youjin1985, what you did should be fine. And if you run VASP only on the HPC worker it should be enough to create the configuration file just there, it should not be used on the runner machine.
However, to have a better control on the which values are used for each job, I prefer to set the atomate2 options as environmental variables and use the exec_config. I have started writing more detailed information for jobflow-remote in the atomate2 documentation, I will try to finish it and open a PR @JaGeo.

I paste below the section I have written so far about the atomate2 configuration options. Let me know if you think this is clear enough.

Atomate2 configurations

All atomate2 settings can be controlled defining values in an atomate2.yaml file or defining the same value as environment variables, prepended by atomate2. You can see the full list of available settings in the .Atomate2Settings docs. For now, we will just configure the commands used to run VASP.

In jobflow remote one way of setting the VASP options for a worker is to define an execution configuration, in the project configuration, including the necessary element to run VASP and setting the environmental variables for the atomate2 configuration.
An example of such configuration can be:

exec_config:
  vasp_6.4.3:
    modules:
    - gcc/11.3.0
    - openmpi/4.1.3
    - openblas/0.3.20
    - fftw/3.3.10
    export:
      PATH: /path/to/vasp/bin:$PATH
      atomate2_VASP_CMD: '"srun vasp_std"'
      atomate2_VASP_GAMMA_CMD: '"srun vasp_gam"'
      atomate2_VASP_NCL_CMD: '"srun vasp_ncl"'
      atomate2_VASP_STORE_ADDITIONAL_JSON: 'False'
    pre_run:
    post_run:

Note that several exec_config can be defined in the configuration, allowing to refer to different VASP versions.

Alternatively, it is possible to define an atomate2_vasp6.4.3.yaml file with the following content on the worker

VASP_CMD: <<VASP_CMD>>
VASP_GAMMA_CMD: <VASP_GAMMA_CMD>>
VASP_NCL_CMD: <VASP_NCL_CMD>>

and activate it in the execution configuration

exec_config:
  vasp_6.4.3:
    modules:
    - gcc/11.3.0
    - openmpi/4.1.3
    - openblas/0.3.20
    - fftw/3.3.10
    export:
      PATH: /path/to/vasp/bin:$PATH
      ATOMATE2_CONFIG_FILE: /path/to/atomate2_vasp6.4.3.yaml
    pre_run:
    post_run:

[Youjin1985]

I had to switch to another cluster, where there is different environment
In auto-produced submit.sh there is
source /public/home/xd_jwq/.bashrc; conda deactivate; conda activate jobflow_env
Seems it doesnt work, because I have a lot of import errors in resulting jfremote_out.json like
File \"/public/home/xd_jwq/miniconda3/envs/jobflow_env/lib/python3.12/site-packages/emmet/core/symmetry.py\", line 8, in <module>\n from pymatgen.symmetry.analyzer import (\nImportError: cannot import name 'SymmetryUndeterminedError' from 'pymatgen.symmetry.analyzer'
which is default environment - so my jobflow environment is somehow not activated despite activation command explicitly being in submit.sh, what can be done?

[Youjin1985]

There is complete example of working config:
USER+RUNNER yaml config:

exec_config:
  vasp:
    export:
      I_MPI_ADJUST_REDUCE: 3
      MKL_CBWR: AVX2
      MKL_DEBUG_CPU_TYPE: 5
      MKL_ENABLE_INSTRUCTIONS: AVX2
      OMP_NUM_THREADS: 1
      PATH: /public/home/XXX/miniconda3/envs/jobflow_env/bin:$PATH
      atomate2_VASP_CMD: '"mpirun -np 192 /public/software/apps/vasp/intelmpi/5.4.4/bin/vasp_std"'
    modules:
    - compiler/intel/2021.3.0
    - mpi/intelmpi/2021.3.0
    pre_run: source /public/home/XXX/.bashrc; conda deactivate; conda activate jobflow_env
jobstore:
  additional_stores:
    data:
      collection_name: outputs_blobs
      database: jf
      host: 172.27.160.1
      port: 27017
      type: GridFSStore
  docs_store:
    collection_name: outputs
    database: jf
    host: 172.27.160.1
    port: 27017
    type: MongoStore
name: CAMD_TEST
queue:
  store:
    collection_name: jobs
    database: jf
    host: localhost
    type: MongoStore
workers:
  example_worker:
    host: XXX
    pre_run: 'source /public/home/XXX/.bashrc; conda deactivate; conda activate jobflow_env'
    scheduler_type: slurm
    timeout_execute: 60
    type: remote
    user: XXX
    work_dir: /public/home/XXX/yj/jobflow_workdir

[Youjin1985]

I would like to ask the following: is there convenient way to retrieve final VASP results (relaxed structure and properties) or I am totally responsible for keeping track of ids which are returned by submit_flow (possibly constructing another database to keep track on project) and writing code for retrieving results from scratch?

[davidwaroquiers]

I would like to ask the following: is there convenient way to retrieve final VASP results (relaxed structure and properties) or I am totally responsible for keeping track of ids which are returned by submit_flow (possibly constructing another database to keep track on project) and writing code for retrieving results from scratch?

Hi @Youjin1985

The best practice is to add metadata to your jobs/flows to allow querying easily. This is a feature of jobflow itself, not directly jobflow-remote although of course you can query by metadata within jobflow-remote.

See here in jobflow: https://materialsproject.github.io/jobflow/tutorials/4-creating-flows.html#Setting-Metadata

In the command line interface of jobflow-remote, you have possibilities to query by data, e.g. "jf job list --metadata ...". Just do "jf job list -h" to see the help and how the "..." should be set. (you have of course other commands that have such a --metadata" option).

To retrieve final results programmatically, you would you use the JobStore of your project and query from there. Use the "get_jobstore" helper function of jobflow-remote to get the store for your project, which returns you a jobflow JobStore that you can use to query for your results, e.g.:

jobstore = get_jobstore("myprojectname")
output = jobstore.query_one(criteria={"metadata.sometag": "somevalueforthistag"})

The metadata can contain any dict like thing with lists or anything and your criteria above can follow any pymongo/mongodb allowed type of queries (standard equal as above, but also $all, $in, and all the others ...).

Hope this helps!