How can I build the wheels necessary for a quick installation of PySCFad on a "compute node"?

Question

The following is the procedure I have been following:

Install PySCFad

module load python/3.10 flexiblas libxc/5.1.3 qcint # modify accordingly based on the modules available on your supercomputer
virtualenv --no-download --clear ~/ve_pyscfad_wheel && source ~/ve_pyscfad_wheel/bin/activate
pip install --no-index --upgrade pip
git clone https://github.com/fishjojo/pyscf.git --branch ad pyscf_ad
cd pyscf_ad
sed -i '38s/pyscf/pyscf-custom-ad/' setup.py
export VERBOSE=1; export CMAKE_CONFIGURE_ARGS="-DBLA_VENDOR=Flexiblas -DBLAS_LIBRARIES=-lflexiblas -DBUILD_LIBCINT=0 -DBUILD_LIBXC=0";
python setup.py bdist_wheel


# Install "ad" branch of PySCF in one line, then pyscfad and pyscf-properties next:
pip install --no-index dist/pyscf_custom_ad-2.1.1-cp310-cp310-linux_x86_64.whl
pip install --no-index pyscfad pyscf-properties
pip freeze > /home/nike/pyscfad-requirements.txt

# Make the ~/.pyscfad_conf.py file contain only the following text:
pyscfad = True
pyscf_numpy_backend = 'jax'
pyscf_scipy_linalg_backend = 'pyscfad'
pyscf_scipy_backend = 'jax'
# The followings are optional
pyscfad_scf_implicit_diff = True
pyscfad_ccsd_implicit_diff = True

Submit job

My Python script is the following executable file called 00-simple.py:

#!/usr/bin/env python

import pyscf
from pyscfad import gto, scf

"""
Analytic nuclear gradient for RHF computed by auto-differentiation
"""

mol = gto.Mole()
mol.atom    = 'H 0 0 0; H 0 0 0.74'  # in Angstrom
mol.basis   = '631g'
mol.verbose = 5
mol.build()

mf = scf.RHF(mol)
mf.kernel()
jac = mf.energy_grad()
print(f'Nuclaer gradient:\n{jac.coords}')
print(f'Gradient wrt basis exponents:\n{jac.exp}')
print(f'Gradient wrt basis contraction coefficients:\n{jac.ctr_coeff}')

My SLURM submission script contains the following:

module load python/3.10
virtualenv --no-download --clear $SLURM_TMPDIR/ENV && source $SLURM_TMPDIR/ENV/bin/activate
pip install --no-index --upgrade pip
pip install --no-index -r /home/nike/pyscfad-requirements.txt

export PYSCF_MAX_MEMORY=6000 # in MB
PYSCF_CONFIG_FILE=~/.pyscfad_conf.py

./00-simple.py

Output

The std output file contained:

Processing /home/nike/pyscfad/dist/pyscf_custom_ad-2.1.1-cp310-cp310-linux_x86_64.whl
...
Processing /cvmfs/soft.computecanada.ca/custom/python/wheelhouse/gentoo/avx2/pyscf-2.2.0+computecanada-cp310-cp310-linux_x86_64.whl
...
Installing collected packages: typing_extensions, six, pyparsing, Pillow, packaging, numpy, kiwisolver, fonttools, cycler, absl_py, scipy, python-dateutil, opt-einsum, h5py, contourpy, pyscf-custom-ad, pyscf, matplotlib, jaxlib, pyscf_properties, jax, jaxopt, pyscfad
Successfully installed Pillow-9.4.0+computecanada absl_py-1.4.0+computecanada contourpy-1.0.7+computecanada cycler-0.11.0+computecanada fonttools-4.39.0+computecanada h5py-3.8.0+computecanada jax-0.4.2+computecanada jaxlib-0.4.2+cuda11.cudnn82.computecanada jaxopt-0.6+computecanada kiwisolver-1.4.4+computecanada matplotlib-3.7.0+computecanada numpy-1.24.2+computecanada opt-einsum-3.3.0+computecanada packaging-23.0+computecanada pyparsing-3.0.9+computecanada pyscf-2.2.0+computecanada pyscf-custom-ad-2.1.1 pyscf_properties-0.1.0+computecanada pyscfad-0.1.2+computecanada python-dateutil-2.8.2+computecanada scipy-1.10.1+computecanada six-1.16.0+computecanada typing_extensions-4.5.0+computecanada

The error file contained the following:

Traceback (most recent call last):
  File "/project/6004934/nike/basis-set-optimization/./00-simple.py", line 4, in <module>
    from pyscfad import gto, scf
  File "/localscratch/nike.62605116.0/ENV/lib/python3.10/site-packages/pyscfad/scf/__init__.py", line 3, in <module>
    from pyscfad.scf import rohf
  File "/localscratch/nike.62605116.0/ENV/lib/python3.10/site-packages/pyscfad/scf/rohf.py", line 4, in <module>
    from pyscf.lib import stop_grad
ImportError: cannot import name 'stop_grad' from 'pyscf.lib' (/localscratch/nike.62605116.0/ENV/lib/python3.10/site-packages/pyscf/lib/__init__.py)

How can I run this code without getting an error?

The input file is taken precisely from the examples/scf/ folder of the main branch of the PySCFad GitHub repository, and it works on a "personal" laptop that uses the same version of PySCFad (the version that was released two days ago).

Answer 1

I finally got it working!

Explanation of the error

The above procedure installed the "main" branch of "pyscf_properties" rather than the "ad" branch. It seems that due to a requirement not being met, the original PySCF wheel on the cluster was being pulled, rather than installing the "ad" branch of PySCF. The following build procedure, constructed with the help of Charles Coulombe from Compute Canada.

Build

Set up the virtual environment, clone the "ad" branch of PySCF, add a custom tag to the version and build wheel:

module load python/3.10 flexiblas libxc/5.1.3 
virtualenv --no-download --clear ~/ve_pyscfad_wheel && source ~/ve_pyscfad_wheel/bin/activate
pip install --no-index --upgrade pip
git clone https://github.com/fishjojo/pyscf.git --branch ad pyscf_ad
cd pyscf_ad

sed -i '38s/2.1.1/2.1.1+ad/' pyscf/__init__.py 
export VERBOSE=1; export CMAKE_CONFIGURE_ARGS="-DBLA_VENDOR=Flexiblas -DBLAS_LIBRARIES=-lflexiblas -DBUILD_LIBCINT=1 -DBUILD_LIBXC=0";
time python setup.py bdist_wheel
cd

Clone the "ad" branch of "pyscf_properties", add a custom tag to the version and build wheel:

git clone https://github.com/fishjojo/properties.git --branch ad pyscf_properties_ad
cd pyscf_properties_ad/

sed -i '15s/0.1.0/0.1.0+ad/' pyscf/prop/__init__.py 
python setup.py bdist_wheel
cd

Install the "ad" branch of PySCF, the "ad" branch of "pyscf_properties" and the main branch of PySCFad, then freeze requirements:

pip install --no-index $HOME/pyscf_ad/dist/pyscf-2.1.1+ad-cp310-cp310-linux_x86_64.whl $HOME/pyscf_properties_ad/dist/pyscf_properties-0.1.0+ad-py3-none-any.whl pyscfad # Somehow this was working with $HOME/pyscf_ad/ and $HOME/properties/ even though these folders didn't exist?
pip freeze > pyscfad-requirements.txt

Test to see if the error mentioned in the original question happens, then deactivate the virtual environment:

python -c 'from pyscfad import gto, scf' # Will only work if we've correctly installed the "ad" branch of pyscf_properties
deactivate

Set the ~/.pyscfad_conf.py file exactly as described in the question:

pyscfad = True
pyscf_numpy_backend = 'jax'
pyscf_scipy_linalg_backend = 'pyscfad'
pyscf_scipy_backend = 'jax'
# The followings are optional
pyscfad_scf_implicit_diff = True
pyscfad_ccsd_implicit_diff = True

Run

The lines described for the SLURM submission script, which were mentioned in the "Submit job" section of the question, can seemingly remain unchanged. I used export PYSCF_CONFIG_FILE=~/.pyscfad_conf.py instead of just PYSCF_CONFIG_FILE=~/.pyscfad_conf.py but the output when the example script is run, isn't affected by whether or not export is included here. Add the following lines into your SLURM script:

module load python/3.10
virtualenv --no-download --clear $SLURM_TMPDIR/ENV && source $SLURM_TMPDIR/ENV/bin/activate
pip install --no-index --upgrade pip
pip install --no-index -r /home/nike/pyscfad-requirements.txt

export PYSCF_MAX_MEMORY=6000 # in MB
export PYSCF_CONFIG_FILE=~/.pyscfad_conf.py

./00-simple.py

Just for completeness I'll include the 00-simple.py file from the original question here again:

#!/usr/bin/env python

import pyscf
from pyscfad import gto, scf

"""
Analytic nuclear gradient for RHF computed by auto-differentiation
"""

mol = gto.Mole()
mol.atom    = 'H 0 0 0; H 0 0 0.74'  # in Angstrom
mol.basis   = '631g'
mol.verbose = 5
mol.build()

mf = scf.RHF(mol)
mf.kernel()
jac = mf.energy_grad()
print(f'Nuclaer gradient:\n{jac.coords}')
print(f'Gradient wrt basis exponents:\n{jac.exp}')
print(f'Gradient wrt basis contraction coefficients:\n{jac.ctr_coeff}')