How to run on Midway3?

[ricalessandri]

On midway3, I've been trying to run the ABF alanine dipeptide on GPU but I always get:

2022-01-12 11:12:21.762567: E external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_driver.cc:271] failed call to cuInit: CUDA_ERROR_SYSTEM_DRIVER_MISMATCH: system has unsupported display driver / cuda driver combination
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)

What I load (as recommended by @pabloferz some time ago):

source activate pysages
module load python/anaconda-2020.11
module load mvapich-gdr
module load nvhpc

which shows:

Currently Loaded Modulefiles:
 1) slurm/current   2) rcc/default   3) python/anaconda-2020.11(default)   4) cuda/11.2   5) mvapich-gdr/2.3.5(default)   6) nvhpc/20.11(default)

Then run my sbatch script:

#!/bin/bash
#SBATCH --nodes=1
#SBATCH --account=pi-depablo
#SBATCH --job-name=test
#SBATCH --time=0-01:00:00
#SBATCH --ntasks=1
#SBATCH --gres=gpu:1
#SBATCH --partition=depablo-gpu
python alanine-dipeptide_openmm.py

(or --partition=gpu).

What am I doing wrong?

[johnaparker]

I created a convenience module that loads the correct modules, "module load pysages". The current install is compiled against openmpi/3.1.4 instead of mvapich-gdr (until mvapich-gdr supports cuda-11.5)

You should also be able to create your own conda environment with your dev version of PySAGES by cloning the existing module.

conda create --prefix=/path/to/env --clone pysages   # clone env
pip install -e .  # run inside pysages directory to get latest dev in editable mode

[ricalessandri]

This works, thanks a lot!
Namely, I can now load:

source activate pysages
module load pysages

which shows:

Currently Loaded Modulefiles:
 1) slurm/current   3) python/anaconda-2021.05   5) cuda/11.2               
 2) rcc/default     4) openmpi/3.1.4             6) pysages/0.1.0(default) 

and then run the same script.

Thanks also for the cloning tip. It works. I did something like:

conda create --clone pysages --name mypysages   # clone env

[ricalessandri]

So, I've been struggling this morning to get PySAGES on midway3 going again - it seems that the procedure here doesn't work anymore. I have tried quite a few things but I'm now converging to the following error:

WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
Traceback (most recent call last):
  File "alanine-dipeptide.py", line 126, in <module>
    main(sys.argv[1:])
  File "alanine-dipeptide.py", line 117, in main
    pysages.run(method, generate_simulation, timesteps, callback)
  File "plum/function.py", line 537, in plum.function.Function.__call__
  File "/project/depablo/riccardo/PySG/PySAGES/pysages/methods/core.py", line 115, in run
    wrapped_context = ContextWrapper(context, method, callback)
  File "/project/depablo/riccardo/PySG/PySAGES/pysages/backends/core.py", line 44, in __init__
    self.sampler = self._backend.bind(self, sampling_method, callback, **kwargs)
  File "/project/depablo/riccardo/PySG/PySAGES/pysages/backends/openmm.py", line 194, in bind
    snapshot = take_snapshot(wrapped_context)
  File "/project/depablo/riccardo/PySG/PySAGES/pysages/backends/openmm.py", line 59, in take_snapshot
    positions = asarray(dlext.positions(context_view))
  File "/home/alessandri/.conda/envs/mypysages/lib/python3.8/site-packages/jax/_src/dlpack.py", line 60, in from_dlpack
    buf = xla_client._xla.dlpack_managed_tensor_to_buffer(
jaxlib.xla_extension.XlaRuntimeError: INVALID_ARGUMENT: DLPack tensor is on GPU, but no GPU backend was provided.

which gives the GPU/TPU not found + there's some new DLPack error.

This is basically what I do:

module load python/anaconda-2021.05
conda create --clone pysages --name mypysages   # clone the env
conda activate mypysages
cd PySAGES
pip install -e .  # run inside pysages directory to get latest dev in editable mode
pip install --upgrade jaxlib # necessary because by default some previous version of jaxlib gets installed

maybe you guys @sivadasetty @gustavor101 @ndtrung81 have more recent experience on running PySAGES on midway3.

P.S. And the jobscript:

#SBATCH --job-name=metad
#SBATCH --output=metad.o%j.out
#SBATCH --time=0-01:00:00
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --gres=gpu:1
#SBATCH --partition=gpu
python alanine-dipeptide.py

Edit: Problem must be in openmm-dlext:

>>> import openmm_dlext
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/software/python-anaconda-2021.05-el8-x86_64/envs/pysages/lib/python3.8/site-packages/openmm_dlext/__init__.py", line 1, in <module>
    from .dlpack_extension import (
ImportError: libcuda.so.1: cannot open shared object file: No such file or directory

So I've been trying to re-build openmm-dlext following the latest install instructions but that doesn't quite work on midway3. I'm getting stuck at some permission issues:

...
[ 90%] Building CXX object wrappers/python/CMakeFiles/dlpack_extension.dir/PyDLExt.cpp.o
[100%] Linking CXX shared module dlpack_extension.cpython-38-x86_64-linux-gnu.so
[100%] Built target dlpack_extension
Install the project...
-- Install configuration: "Release"
-- Up-to-date: /software/python-anaconda-2021.05-el8-x86_64/envs/pysages/include
CMake Error at extern/dlpack/cmake_install.cmake:46 (file):
  file INSTALL cannot set permissions on
  "/software/python-anaconda-2021.05-el8-x86_64/envs/pysages/include":
  Operation not permitted.
Call Stack (most recent call first):
  cmake_install.cmake:47 (include)

gmake: *** [Makefile:137: install] Error 1

[ndtrung81]

@ricalessandri It seems that you were trying to installing into the pysages env. Can you try to install into your own mypysages to avoid the permission issue?

[ricalessandri]

Thanks @ndtrung81 . Let me try a few more things later today and get back to you

[ndtrung81]

@ricalessandri I will take a closer look at this issue.

Btw, in the job script you posted I didn't see you have something like:

module load python/anaconda-2021.05
conda activate mypysages

[ndtrung81]

@ricalessandri I can reproduce the errors with no TPU/GPU backend you encountered from my end. What I did to resolve this issue is to upgrade both jax and jaxlib with the following:

pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html

With this one, without reinstalling openmm_dlext nor hoomd_dlext I can run the examples under openmm/abf and hoomd-blue/harmonic_bias successfully.

Can you give it another try to see if the issue persists on your end? You can simply try module load pysages, or clone pysages to your mypysages and activate it for use as you did earlier.

[ndtrung81]

A little update to the command:
pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

[ricalessandri]

Just an update on this. Recently, I was getting the error:

2023-01-31 17:00:26.244195: E external/org_tensorflow/tensorflow/compiler/xla/stream_executor/cuda/cuda_dnn.cc:418] Loaded runtime CuDNN library: 8.2.1 but source was compiled with: 8.6.0.  CuDNN library needs to have matching major version and equal or higher minor version. If using a binary install, upgrade your CuDNN library.  If building from sources, make sure the library loaded at runtime is compatible with the version specified during compile configuration.

...

openmm.OpenMMException: XlaRuntimeError: UNKNOWN: Failed to determine best cudnn convolution algorithm for:
%cudnn-conv.1 = (f64[2,1,50,50]{3,2,1,0}, u8[0]{0}) custom-call(f64[2,1,54,54]{3,2,1,0} %bitcast.962, f64[1,1,5,5]{3,2,1,0} %bitcast.963), window={size=5x5}, dim_labels=bf01_oi01->bf01, custom_call_target="__cudnn$convForward", metadata={op_name="jit(update)/jit(main)/jit(update)/cond/branch_1_fun/jit(<lambda>)/conv_general_dilated[window_strides=(1, 1) padding=((0, 0), (0, 0)) lhs_dilation=(1, 1) rhs_dilation=(1, 1) dimension_numbers=ConvDimensionNumbers(lhs_spec=(0, 1, 2, 3), rhs_spec=(0, 1, 2, 3), out_spec=(0, 1, 2, 3)) feature_group_count=1 batch_group_count=1 lhs_shape=(2, 1, 54, 54) rhs_shape=(1, 1, 5, 5) precision=None preferred_element_type=None]" source_file="/project/depablo/riccardo/PySG/PySAGES/pysages/ml/training.py" source_line=41}, backend_config="{\"conv_result_scale\":1,\"activation_mode\":\"0\",\"side_input_scale\":0}"

Original error: UNIMPLEMENTED: DNN library is not found.
...

This can be fixed by slightly modifiying the pip install --upgrade "jax[cuda]" command following the jax repo instructions:

pip install --upgrade pip
# Installs the wheel compatible with Cuda >= 11.8 and cudnn >= 8.6
pip install "jax[cuda11_cudnn86]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

I've update the hackmd.io doc accordingly.

[ricalessandri]

Actually, it now (April 2023) seems that we need to use the following cuda & cudnn versions on midway3:

pip install --upgrade pip
# Installs the wheel compatible with Cuda 11.4+ and cudnn 8.2+ (deprecated).
pip install "jax[cuda11_cudnn82]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

[ricalessandri]

As agreed, I summarized the discussion on "how to use/install PySAGES on midway3" in the hackmd.io doc. See here: https://hackmd.io/Jbpc1E2kRbKPLUnmLKwINA?view#Installation-of-PySAGES-on-midway3
Thanks again @ndtrung81 (and previously @johnaparker ) for the help!