Installing on midway2

[ap-uchi]

Module not found is thrown on import pysages if installation procedure is performed with mamba or micromamba on HPC clusters, particularly if the underlying python is loaded as a module.

Steps to reproduce (on midway2):

# Load necessary modules
module load python/intel-2020.up1
module load gcc/10.2.0 cuda/11.2 cmake/3.19

# Create environment
mamba create -n pysages
mamba activate pysages
mamba install openmm cupy
pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

# Install OpenMM dlext
git clone https://github.com/SSAGESLabs/openmm-dlext.git
cmake -S . -B build
cmake --build build --target install

# Install pysages
pip install git+https://github.com/SSAGESLabs/PySAGES.git

# Attempt to import
python
import pysages

[ap-uchi]

Proceeding with conda instead of mamba, I receive the following error when trying to run the SpectralABF tutorial:

        RuntimeError                              Traceback (most recent call last)
        Input In [4], in <cell line: 1>()
        ----> 1 run_result = pysages.run(method, generate_simulation, timesteps)
            2 with open('results.pkl', 'wb') as f:
            3     pkl.dump(run_result, f) 

        File /project2/depablo/apeng/.local/conda/envs/pys-om/lib/python3.9/site-packages/plum/function.py:399, in Function.__call__(self, *args, **kw_args)
            397 def __call__(self, *args, **kw_args):
            398     method, return_type = self._resolve_method_with_cache(args=args)
        --> 399     return _convert(method(*args, **kw_args), return_type)

        File /project2/depablo/apeng/.local/conda/envs/pys-om/lib/python3.9/site-packages/pysages/methods/core.py:222, in run(method, context_generator, timesteps, callback, context_args, post_run_action, config, **kwargs)
            210     return executor.submit(
            211         _run_replica,
            212         method,
        (...)
            218         **kwargs,
            219     )
            221 with config.executor as ex:
        --> 222     futures = [submit_work(ex, method, callback) for _ in range(config.copies)]
            223 results = [future.result() for future in futures]
            224 states = [r.states for r in results]

        File /project2/depablo/apeng/.local/conda/envs/pys-om/lib/python3.9/site-packages/pysages/methods/core.py:222, in <listcomp>(.0)
            210     return executor.submit(
            211         _run_replica,
            212         method,
        (...)
            218         **kwargs,
            219     )
            221 with config.executor as ex:
        --> 222     futures = [submit_work(ex, method, callback) for _ in range(config.copies)]
            223 results = [future.result() for future in futures]
            224 states = [r.states for r in results]

        File /project2/depablo/apeng/.local/conda/envs/pys-om/lib/python3.9/site-packages/pysages/methods/core.py:210, in run.<locals>.submit_work(executor, method, callback)
            209 def submit_work(executor, method, callback):
        --> 210     return executor.submit(
            211         _run_replica,
            212         method,
            213         context_generator,
            214         timesteps,
            215         context_args,
            216         callback,
            217         post_run_action,
            218         **kwargs,
            219     )

        File /project2/depablo/apeng/.local/conda/envs/pys-om/lib/python3.9/site-packages/pysages/methods/utils.py:33, in SerialExecutor.submit(self, fn, *args, **kwargs)
            29 """
            30 Executes `fn(*args, **kwargs)` and returns a `Future` object wrapping the result.
            31 """
            32 future = Future()
        ---> 33 future.set_result(fn(*args, **kwargs))
            34 return future

        File /project2/depablo/apeng/.local/conda/envs/pys-om/lib/python3.9/site-packages/pysages/methods/core.py:316, in _run_replica(method, *args, **kwargs)
            313 def _run_replica(method, *args, **kwargs):
            314     # Trampoline method to enable multiple replicas to be run with mpi4py.
            315     run = dispatch_table(dispatch)["_run"]
        --> 316     return run(method, *args, **kwargs)

        File /project2/depablo/apeng/.local/conda/envs/pys-om/lib/python3.9/site-packages/plum/function.py:399, in Function.__call__(self, *args, **kw_args)
            397 def __call__(self, *args, **kw_args):
            398     method, return_type = self._resolve_method_with_cache(args=args)
        --> 399     return _convert(method(*args, **kw_args), return_type)

        File /project2/depablo/apeng/.local/conda/envs/pys-om/lib/python3.9/site-packages/pysages/methods/core.py:363, in _run(method, context_generator, timesteps, context_args, callback, post_run_action, **kwargs)
            329 """
            330 Base implementation for running a single simulation with the specified `SamplingMethod`.
            331 
        (...)
            359 **Note**: All arguments must be pickable.
            360 """
            361 timesteps = int(timesteps)
        --> 363 sampling_context = SamplingContext(method, context_generator, callback, context_args)
            364 context_args["context"] = sampling_context.context
            365 sampler = sampling_context.sampler

        File /project2/depablo/apeng/.local/conda/envs/pys-om/lib/python3.9/site-packages/pysages/backends/core.py:102, in SamplingContext.__init__(self, sampling_method, context_generator, callback, context_args, **kwargs)
            99 self.run = None
            101 backend = import_module("." + self._backend_name, package="pysages.backends")
        --> 102 self.sampler = backend.bind(self, callback, **kwargs)
            104 # `self.view` and `self.run` *must* be set by the backend bind function.
            105 assert self.view is not None

        File /project2/depablo/apeng/.local/conda/envs/pys-om/lib/python3.9/site-packages/pysages/backends/openmm.py:195, in bind(sampling_context, callback, **kwargs)
            193 sampling_method = sampling_context.method
            194 force = Force()
        --> 195 force.add_to(context)  # OpenMM will handle the lifetime of the force
            196 sampling_context.view = force.view(context)
            197 sampling_context.run = simulation.step

        File /project2/depablo/apeng/.local/conda/envs/pys-om/lib/python3.9/site-packages/openmm_dlext/__init__.py:29, in Force.add_to(self, context)
            27 if not self.thisown and not self.is_present_in(system):
            28     raise RuntimeError("Force already added to another system")
        ---> 29 self.__alt__.add_to(_to_capsule(context), _to_capsule(system))
            30 self.thisown = False

        RuntimeError: Unsupported platform

Attached is my environment and the modules loaded are below:

Currently Loaded Modulefiles:
 1) serf/1.3.9         5) slurm/current            
 2) subversion/1.9.4   6) python/anaconda-2020.02  
 3) git/2.10           7) cuda/11.5                
 4) env/rcc

[pabloferz]

Hey @ap-uchi!

On midway2 the installed nvidia drivers are only compatible with CUDA 11.2. I found the following to work (tested only with python 3.9.x, where x>1, but 3.10 or newer might work as well):

# Create environment
mamba create -n pysages -c conda-forge python=3.9
mamba activate pysages
mamba install -c conda-forge cuda-version=11.2 cudatoolkit-dev=11.2 cudnn=8.2 cutensor nccl cupy  # CUDA deps
mamba install -c conda-forge cython numba  # compiler tools
mamba install -c conda-forge typing-extensions plum-dispatch  # typing tools
mamba install -c conda-forge importlib-metadata  # python extras

# Install jax
mamba install -c conda-forge ml_dtypes opt-einsum scipy zipp  # jax deps
# There are wheels available for the following jax version for the corresponding python, cuda, and cudnn versions,
JAX_VERSION=0.4.6
pip install "jax==$JAX_VERSION" "https://storage.googleapis.com/jax-releases/cuda11/jaxlib-$JAX_VERSION+cuda11.cudnn82-cp39-cp39-manylinux2014_x86_64.whl"

# Install openmm-dlext, installs openmm as well
mamba install -c conda-forge openmm-dlext

# Install pysages
pip install git+https://github.com/SSAGESLabs/PySAGES.git

Notice that there's no need to load any modules or to build openmm-dlext (it's already available in conda-forge)

Finally, you might need to do the following

export LD_LIBRARY_PATH="$CONDA_PREFIX/lib:$LD_LIBRARY_PATH"

before running a simulation, though.

[ricalessandri]

OK, we should add this to: https://hackmd.io/Jbpc1E2kRbKPLUnmLKwINA

[pabloferz]

Better yet, we should turn that document into Wiki sections within the repo.

[ricalessandri]

I agree. And, perhaps, also write up some of the tips that we deem "transferable" beyond midway2/midway3 (CUDA compatibilities? what else?) and include them as part of the documentation?

[pabloferz]

Yes, that's a good idea. That should probably be a subsection within the installation guide in the documentation.

[ricalessandri]

New home for PySAGES installation notes: https://github.com/SSAGESLabs/PySAGES/wiki/PySAGES-installation-notes