Skip to content

Commit

Permalink
Squashed commit of the following: (#465)
Browse files Browse the repository at this point in the history 
commit 7f494e4
Author: Matthew W. Thompson <mattwthompson@protonmail.com>
Date:   Fri May 13 12:01:44 2022 -0500

    Update unit formatting in test, pin Pint >=0.19

commit 220db6a
Author: Matthew W. Thompson <mattwthompson@protonmail.com>
Date:   Fri May 13 12:01:24 2022 -0500

    Run `pre-commit autoupdate`

commit 7ed2d52
Merge: 33408fa 9d17bc1
Author: Matthew W. Thompson <mattwthompson@protonmail.com>
Date:   Fri May 13 11:36:54 2022 -0500

    Merge branch 'v0-3-x' into add_opencl_option

commit 9d17bc1
Author: Matthew W. Thompson <mattwthompson@protonmail.com>
Date:   Fri May 13 11:36:28 2022 -0500

    Run CI on v0-3-x branch

commit 33408fa
Merge: c080868 d42e4f8
Author: Matthew W. Thompson <mattwthompson@protonmail.com>
Date:   Fri May 13 11:34:43 2022 -0500

    Merge remote-tracking branch 'upstream/v0-3-x' into add_opencl_option

commit c080868
Author: jeff231li <pea231@gmail.com>
Date:   Fri Mar 25 13:51:30 2022 -0700

    update docs

commit 863709d
Author: jeff231li <pea231@gmail.com>
Date:   Fri Mar 25 13:02:42 2022 -0700

    simplify GPU selection in Yank

commit 390d115
Author: jeff231li <pea231@gmail.com>
Date:   Fri Mar 25 12:59:11 2022 -0700

    update yank protocols with OpenCL

commit ea98260
Author: jeff231li <pea231@gmail.com>
Date:   Thu Mar 24 11:00:15 2022 -0700

    update utils/openmm.py

commit 6511b71
Merge: 3e89b01 89e9043
Author: jeff231li <pea231@gmail.com>
Date:   Thu Mar 24 10:20:37 2022 -0700

    Merge branch 'master' into add_opencl_option

commit 3e89b01
Author: jeff231li <pea231@gmail.com>
Date:   Mon Feb 21 15:16:28 2022 -0800

    remove testing print statements

commit 7b2e5c3
Author: jeff231li <pea231@gmail.com>
Date:   Mon Feb 21 13:35:52 2022 -0800

    initial commit
  • Loading branch information
@mattwthompson
mattwthompson authored Aug 1, 2022
1 parent 2fa7453 commit 7230c0f
Show file tree
Hide file tree
Showing 5 changed files with 106 additions and 43 deletions.
42 changes: 34 additions & 8 deletions docs/backends/daskbackends.rst
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
.. |base_dask_job_queue_backend| replace:: :py:class:`~openff.evaluator.backends.dask.BaseDaskJobQueueBackend`
.. |dask_lsf_backend| replace:: :py:class:`~openff.evaluator.backends.dask.DaskLSFBackend`
.. |dask_pbs_backend| replace:: :py:class:`~openff.evaluator.backends.dask.DaskPBSBackend`
.. |dask_slurm_backend| replace:: :py:class:`~openff.evaluator.backends.dask.DaskSLURMBackend`
.. |queue_worker_resources| replace:: :py:class:`~openff.evaluator.backends.dask.QueueWorkerResources`

Dask Backends
Expand All @@ -19,7 +20,7 @@ cluster.html>`_ class to distribute tasks on a single machine::
worker_resources = ComputeResources(
number_of_threads=1,
number_of_gpus=1,
preferred_gpu_toolkit=GPUToolkit.CUDA,
preferred_gpu_toolkit=ComputeResources.GPUToolkit.CUDA,
)

with DaskLocalCluster(number_of_workers=1, resources_per_worker=worker_resources) as local_backend:
Expand All @@ -32,12 +33,14 @@ numbers of CPUs or GPUs.
Dask HPC Cluster
----------------

The |dask_lsf_backend| and |dask_pbs_backend| backends wrap around the dask `LSFCluster <https://jobqueue.dask.org/en/
latest/generated/dask_jobqueue.LSFCluster.html#dask_jobqueue.LSFCluster>`_ and `PBSCluster <https://jobqueue.dask.org/
en/latest/generated/dask_jobqueue.PBSCluster.html#dask_jobqueue.PBSCluster>`_ classes respectively, and both inherit
the |base_dask_job_queue_backend| class which implements the core of their functionality. They predominantly run in an
adaptive mode, whereby the backend will automatically scale up or down the number of workers based on the current number
of tasks that the backend is trying to execute.
The |dask_lsf_backend|, |dask_pbs_backend|, and |dask_slurm_backend| backends wrap around the dask `LSFCluster
<https://jobqueue.dask.org/en/
latest/generated/dask_jobqueue.LSFCluster.html#dask_jobqueue.LSFCluster>`_, `PBSCluster <https://jobqueue.dask.org/
en/latest/generated/dask_jobqueue.PBSCluster.html#dask_jobqueue.PBSCluster>`_, and `SLURMCluster
<https://jobqueue.dask.org/en/latest/generated/dask_jobqueue.SLURMCluster.html#dask_jobqueue.SLURMCluster>`_ classes
respectively, and both inherit the |base_dask_job_queue_backend| class which implements the core of their
functionality. They predominantly run in an adaptive mode, whereby the backend will automatically scale up or down
the number of workers based on the current number of tasks that the backend is trying to execute.

These backends integrate with the queueing systems which most HPC cluster use to manage task execution. They work
by submitting jobs into the queueing system which themselves spawn `dask workers <https://distributed.dask.org/en/
Expand Down Expand Up @@ -96,4 +99,27 @@ configuration file (this can be found at ``~/.config/dask/distributed.yaml``)::


See the `dask documentation <https://docs.dask.org/en/latest/configuration.html>`_ for more information about changing
``dask`` settings.
``dask`` settings.

Selecting GPU Platform
----------------------
The calculation backends alos allows the user to specify the GPU platform and precision level. Users can specify
either ``auto``, ``CUDA`` or ``OpenCL`` as the `preferred_gpu_toolkit` using the ``GPUToolkit`` enum class. The
default precision level is set to ``mixed`` but can be overridden by specifying `preferred_gpu_precision` using the
``GPUPrecision`` enum class::

worker_resources = ComputeResources(
number_of_threads=1,
number_of_gpus=1,
preferred_gpu_toolkit=ComputeResources.GPUToolkit.OpenCL,
preferred_gpu_precision=ComputeResources.GPUPrecision.mixed,
)

With ``GPUToolkit.auto``, the framework will determine the fastest available platform based on the precision level::

worker_resources = ComputeResources(
number_of_threads=1,
number_of_gpus=1,
preferred_gpu_toolkit=ComputeResources.GPUToolkit.auto,
preferred_gpu_precision=ComputeResources.GPUPrecision.mixed,
)
23 changes: 22 additions & 1 deletion openff/evaluator/backends/backends.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,14 @@ class GPUToolkit(Enum):

CUDA = "CUDA"
OpenCL = "OpenCL"
auto = "auto"

class GPUPrecision(Enum):
"""An enumeration of the different precision for GPU calculations."""

single = "single"
mixed = "mixed"
double = "double"

@property
def number_of_threads(self):
Expand All @@ -36,6 +44,11 @@ def preferred_gpu_toolkit(self):
"""ComputeResources.GPUToolkit: The preferred toolkit to use when running on GPUs."""
return self._preferred_gpu_toolkit

@property
def preferred_gpu_precision(self):
"""ComputeResources.GPUPrecision: The preferred precision level to use when running on GPUs."""
return self._preferred_gpu_precision

@property
def gpu_device_indices(self):
"""str: The indices of the GPUs to run on. This is purely an internal
Expand All @@ -46,7 +59,8 @@ def __init__(
self,
number_of_threads=1,
number_of_gpus=0,
preferred_gpu_toolkit=GPUToolkit.CUDA,
preferred_gpu_toolkit=GPUToolkit.auto,
preferred_gpu_precision=GPUPrecision.mixed,
):
"""Constructs a new ComputeResources object.
Expand All @@ -58,12 +72,15 @@ def __init__(
The number of GPUs available to a calculation worker.
preferred_gpu_toolkit: ComputeResources.GPUToolkit, optional
The preferred toolkit to use when running on GPUs.
preferred_gpu_precision: ComputeResources.GPUPrecision, optional
The preferred precision level to use when runnin on GPUs.
"""

self._number_of_threads = number_of_threads
self._number_of_gpus = number_of_gpus

self._preferred_gpu_toolkit = preferred_gpu_toolkit
self._preferred_gpu_precision = preferred_gpu_precision
# A workaround for when using a local cluster backend which is
# strictly for internal purposes only for now.
self._gpu_device_indices = None
Expand All @@ -81,6 +98,7 @@ def __getstate__(self):
"number_of_threads": self.number_of_threads,
"number_of_gpus": self.number_of_gpus,
"preferred_gpu_toolkit": self.preferred_gpu_toolkit,
"preferred_gpu_precision": self.preferred_gpu_precision,
"_gpu_device_indices": self._gpu_device_indices,
}

Expand All @@ -89,13 +107,15 @@ def __setstate__(self, state):
self._number_of_threads = state["number_of_threads"]
self._number_of_gpus = state["number_of_gpus"]
self._preferred_gpu_toolkit = state["preferred_gpu_toolkit"]
self._preferred_gpu_precision = state["preferred_gpu_precision"]
self._gpu_device_indices = state["_gpu_device_indices"]

def __eq__(self, other):
return (
self.number_of_threads == other.number_of_threads
and self.number_of_gpus == other.number_of_gpus
and self.preferred_gpu_toolkit == other.preferred_gpu_toolkit
and self.preferred_gpu_precision == other.preferred_gpu_precision
)

def __ne__(self, other):
Expand Down Expand Up @@ -125,6 +145,7 @@ def __init__(
number_of_threads=1,
number_of_gpus=0,
preferred_gpu_toolkit=None,
preferred_gpu_precision=None,
per_thread_memory_limit=1 * unit.gigabytes,
wallclock_time_limit="01:00",
):
Expand Down
17 changes: 4 additions & 13 deletions openff/evaluator/backends/dask.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -298,12 +298,6 @@ def __init__(

if resources_per_worker.number_of_gpus > 0:

if (
resources_per_worker.preferred_gpu_toolkit
== ComputeResources.GPUToolkit.OpenCL
):
raise ValueError("The OpenCL gpu backend is not currently supported.")

if resources_per_worker.number_of_gpus > 1:
raise ValueError("Only one GPU per worker is currently supported.")

Expand Down Expand Up @@ -541,6 +535,7 @@ def __init__(
>>> resources = QueueWorkerResources(number_of_threads=1,
>>> number_of_gpus=1,
>>> preferred_gpu_toolkit=QueueWorkerResources.GPUToolkit.CUDA,
>>> preferred_gpu_precision=QueueWorkerResources.GPUPrecision.mixed,
>>> wallclock_time_limit='05:00')
>>>
>>> # Define the set of commands which will set up the correct environment
Expand Down Expand Up @@ -653,6 +648,7 @@ def __init__(
>>> resources = QueueWorkerResources(number_of_threads=1,
>>> number_of_gpus=1,
>>> preferred_gpu_toolkit=QueueWorkerResources.GPUToolkit.CUDA,
>>> preferred_gpu_precision=QueueWorkerResources.GPUPrecision.mixed,
>>> wallclock_time_limit='05:00')
>>>
>>> # Define the set of commands which will set up the correct environment
Expand Down Expand Up @@ -734,6 +730,7 @@ def __init__(
>>> resources = QueueWorkerResources(number_of_threads=1,
>>> number_of_gpus=1,
>>> preferred_gpu_toolkit=QueueWorkerResources.GPUToolkit.CUDA,
>>> preferred_gpu_precision=QueueWorkerResources.GPUPrecision.mixed,
>>> wallclock_time_limit='05:00')
>>>
>>> # Define the set of commands which will set up the correct environment
Expand Down Expand Up @@ -800,12 +797,6 @@ def __init__(self, number_of_workers=1, resources_per_worker=ComputeResources())

if resources_per_worker.number_of_gpus > 0:

if (
resources_per_worker.preferred_gpu_toolkit
== ComputeResources.GPUToolkit.OpenCL
):
raise ValueError("The OpenCL gpu backend is not currently supported.")

if resources_per_worker.number_of_gpus > 1:
raise ValueError("Only one GPU per worker is currently supported.")

Expand All @@ -827,7 +818,7 @@ def start(self):
self._cluster = distributed.LocalCluster(
name=None,
n_workers=self._number_of_workers,
threads_per_worker=self._resources_per_worker.number_of_threads,
threads_per_worker=1,
processes=False,
)

Expand Down
18 changes: 13 additions & 5 deletions openff/evaluator/protocols/yank.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -263,11 +263,19 @@ def _get_options_dictionary(self, available_resources):
)

# A platform which runs on GPUs has been requested.
platform_name = (
"CUDA"
if toolkit_enum == ComputeResources.GPUToolkit.CUDA
else ComputeResources.GPUToolkit.OpenCL
)
if toolkit_enum == ComputeResources.GPUToolkit.auto:
platform_name = "fastest"

elif toolkit_enum == ComputeResources.GPUToolkit.CUDA:
platform_name = "CUDA"

elif toolkit_enum == ComputeResources.GPUToolkit.OpenCL:
platform_name = "OpenCL"

else:
raise KeyError(
f"Specified GPU toolkit {toolkit_enum} is not supported in Yank."
)

return {
"verbose": self.verbose,
Expand Down
49 changes: 33 additions & 16 deletions openff/evaluator/utils/openmm.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,9 @@ def setup_platform_with_resources(compute_resources, high_precision=False):
high_precision: bool
If true, a platform with the highest possible precision (double
for CUDA and OpenCL, Reference for CPU only) will be returned.
For GPU platforms, this overrides the precision level in
`compute_resources`.
Returns
-------
Platform
Expand All @@ -44,43 +47,56 @@ def setup_platform_with_resources(compute_resources, high_precision=False):
# Setup the requested platform:
if compute_resources.number_of_gpus > 0:

# TODO: Make sure use mixing precision - CUDA, OpenCL.
# TODO: Deterministic forces = True

from openff.evaluator.backends import ComputeResources

toolkit_enum = ComputeResources.GPUToolkit(
compute_resources.preferred_gpu_toolkit
)

# A platform which runs on GPUs has been requested.
platform_name = (
"CUDA"
if toolkit_enum == ComputeResources.GPUToolkit.CUDA
else ComputeResources.GPUToolkit.OpenCL
precision_level = ComputeResources.GPUPrecision(
compute_resources.preferred_gpu_precision
)

# noinspection PyCallByClass,PyTypeChecker
platform = Platform.getPlatformByName(platform_name)
# Get platform for running on GPUs.
if toolkit_enum == ComputeResources.GPUToolkit.auto:

if compute_resources.gpu_device_indices is not None:
from openmmtools.utils import get_fastest_platform

# noinspection PyCallByClass,PyTypeChecker
platform = get_fastest_platform(minimum_precision=precision_level)

elif toolkit_enum == ComputeResources.GPUToolkit.CUDA:
# noinspection PyCallByClass,PyTypeChecker
platform = Platform.getPlatformByName("CUDA")

elif toolkit_enum == ComputeResources.GPUToolkit.OpenCL:
# noinspection PyCallByClass,PyTypeChecker
platform = Platform.getPlatformByName("OpenCL")

property_platform_name = platform_name
else:
raise KeyError(f"Specified GPU toolkit {toolkit_enum} is not supported.")

if toolkit_enum == ComputeResources.GPUToolkit.CUDA:
property_platform_name = platform_name.lower().capitalize()
# Set GPU device index
if compute_resources.gpu_device_indices is not None:

# `DeviceIndex` is used by both CUDA and OpenCL
platform.setPropertyDefaultValue(
property_platform_name + "DeviceIndex",
"DeviceIndex",
compute_resources.gpu_device_indices,
)

# Set GPU precision level
platform.setPropertyDefaultValue("Precision", precision_level.value)
if high_precision:
platform.setPropertyDefaultValue("Precision", "double")

# Print platform information
logger.info(
"Setting up an openmm platform on GPU {}".format(
compute_resources.gpu_device_indices or 0
"Setting up an openmm platform on GPU {} with {} kernel and {} precision".format(
compute_resources.gpu_device_indices or 0,
platform.getName(),
platform.getPropertyDefaultValue("Precision"),
)
)

Expand All @@ -92,6 +108,7 @@ def setup_platform_with_resources(compute_resources, high_precision=False):
platform.setPropertyDefaultValue(
"Threads", str(compute_resources.number_of_threads)
)

else:
# noinspection PyCallByClass,PyTypeChecker
platform = Platform.getPlatformByName("Reference")
Expand Down

0 comments on commit 7230c0f

Please sign in to comment.