Repository: dwavesystems/dwave-system (python)

Issue:
Turn DWaveSampler into context manager
Blocked by https://github.com/dwavesystems/dwave-system/issues/77.

Pull request description:
Add support for context manager protocol to all samplers
Now all of QPU (`DWaveSampler`, `DWaveCliqueSampler`), and hybrid (`LeapHybrid*Sampler`) samplers support `.close()` and runtime context use.

Close #91.

Composites will be addressed in a follow-up as per #555.

Patch:
diff --git a/dwave/system/samplers/clique.py b/dwave/system/samplers/clique.py
index bc38e7d..28ca3b4 100644
--- a/dwave/system/samplers/clique.py
+++ b/dwave/system/samplers/clique.py
@@ -12,6 +12,7 @@
 #    See the License for the specific language governing permissions and
 #    limitations under the License.
 
+from contextlib import AbstractContextManager
 from numbers import Number
 from typing import Tuple
 
@@ -115,7 +116,7 @@ class _QubitCouplingComposite(dimod.ComposedSampler):
 
         yield sampleset 
 
-class DWaveCliqueSampler(dimod.Sampler):
+class DWaveCliqueSampler(dimod.Sampler, AbstractContextManager):
     r"""A sampler for solving clique binary quadratic models on the D-Wave system.
 
     This sampler wraps
@@ -159,6 +160,23 @@ class DWaveCliqueSampler(dimod.Sampler):
         **config:
             Keyword arguments, as accepted by :class:`.DWaveSampler`
 
+    .. versionadded:: 1.29.0
+        Support for context manager protocol.
+
+    Note:
+        The recommended way to use :class:`DWaveCliqueSampler` is from a
+        `runtime context <https://docs.python.org/3/reference/datamodel.html#with-statement-context-managers>`_:
+
+        >>> with DWaveCliqueSampler() as sampler:   # doctest: +SKIP
+        ...     sampler.sample(...)
+
+        Alternatively, call the :meth:`~DWaveCliqueSampler.close` method to
+        terminate the sampler resources:
+
+        >>> sampler = DWaveCliqueSampler()      # doctest: +SKIP
+        ...
+        >>> sampler.close()     # doctest: +SKIP
+
     Examples:
         This example creates a BQM based on a 6-node clique (complete graph),
         with random :math:`\pm 1` values assigned to nodes, and submits it to
@@ -172,11 +190,10 @@ class DWaveCliqueSampler(dimod.Sampler):
         ...
         >>> bqm = dimod.generators.ran_r(1, 6)
         ...
-        >>> sampler = DWaveCliqueSampler()   # doctest: +SKIP
-        >>> sampler.largest_clique_size > 5  # doctest: +SKIP
+        >>> with DWaveCliqueSampler() as sampler:   # doctest: +SKIP
+        ...     print(sampler.largest_clique_size > 5)
+        ...     sampleset = sampler.sample(bqm, num_reads=100)
         True
-        >>> sampleset = sampler.sample(bqm, num_reads=100)   # doctest: +SKIP
-        >>> sampler.close()     # doctest: +SKIP
 
     """
     def __init__(self, *,
@@ -197,6 +214,13 @@ class DWaveCliqueSampler(dimod.Sampler):
         """
         self.child.close()
 
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Release system resources allocated and raise any exception triggered
+        within the runtime context.
+        """
+        self.close()
+        return None
+
     @property
     def parameters(self) -> dict:
         try:
diff --git a/dwave/system/samplers/dwave_sampler.py b/dwave/system/samplers/dwave_sampler.py
index f39b60d..a4735fb 100644
--- a/dwave/system/samplers/dwave_sampler.py
+++ b/dwave/system/samplers/dwave_sampler.py
@@ -23,6 +23,7 @@ for explanations of technical terms in descriptions of Ocean tools.
 import copy
 import collections.abc as abc
 from collections import defaultdict
+from contextlib import AbstractContextManager
 from typing import Optional, Dict
 
 import dimod
@@ -86,7 +87,7 @@ def qpu_graph(topology_type, topology_shape, nodelist, edgelist):
     return G
 
 
-class DWaveSampler(dimod.Sampler, dimod.Structured):
+class DWaveSampler(dimod.Sampler, dimod.Structured, AbstractContextManager):
     """A class for using D-Wave quantum computers as samplers for binary quadratic models.
 
     You can configure your :term:`solver` selection and usage by setting parameters,
@@ -127,12 +128,29 @@ class DWaveSampler(dimod.Sampler, dimod.Structured):
         **config:
             Keyword arguments passed to :meth:`~dwave.cloud.client.Client.from_config`.
 
+    .. versionadded:: 1.29.0
+        Support for context manager protocol.
+
     Note:
         Prior to version 1.0.0, :class:`.DWaveSampler` used the ``base`` client,
         allowing non-QPU solvers to be selected.
         To reproduce the old behavior, instantiate :class:`.DWaveSampler` with
         ``client='base'``.
 
+    Note:
+        The recommended way to use :class:`DWaveSampler` is from a
+        `runtime context <https://docs.python.org/3/reference/datamodel.html#with-statement-context-managers>`_:
+
+        >>> with DWaveSampler() as sampler:
+        ...     sampler.sample_ising(...)       # doctest: +SKIP
+
+        Alternatively, call the :meth:`~DWaveSampler.close` method to
+        terminate the sampler resources:
+
+        >>> sampler = DWaveSampler()
+        ...
+        >>> sampler.close()
+
     Examples:
         This example submits a two-variable Ising problem mapped directly to two
         adjacent qubits on a D-Wave system. ``qubit_a`` is the first qubit in
@@ -148,16 +166,14 @@ class DWaveSampler(dimod.Sampler, dimod.Structured):
 
         >>> from dwave.system import DWaveSampler
         ...
-        >>> sampler = DWaveSampler()
-        ...
-        >>> qubit_a = sampler.nodelist[0]
-        >>> qubit_b = next(iter(sampler.adjacency[qubit_a]))
-        >>> sampleset = sampler.sample_ising({qubit_a: -1, qubit_b: 1},
-        ...                                  {},
-        ...                                  num_reads=100)
-        >>> print(sampleset.first.sample[qubit_a] == 1 and sampleset.first.sample[qubit_b] == -1)
+        >>> with DWaveSampler() as sampler:
+        ...     qubit_a = sampler.nodelist[0]
+        ...     qubit_b = next(iter(sampler.adjacency[qubit_a]))
+        ...     sampleset = sampler.sample_ising({qubit_a: -1, qubit_b: 1},
+        ...                                      {},
+        ...                                      num_reads=100)
+        ...     print(sampleset.first.sample[qubit_a] == 1 and sampleset.first.sample[qubit_b] == -1)
         True
-        >>> sampler.close()
 
     See `Ocean Glossary <https://docs.ocean.dwavesys.com/en/stable/concepts/index.html>`_
     for explanations of technical terms in descriptions of Ocean tools.
@@ -194,6 +210,13 @@ class DWaveSampler(dimod.Sampler, dimod.Structured):
         """
         self.client.close()
 
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Release system resources allocated and raise any exception triggered
+        within the runtime context.
+        """
+        self.close()
+        return None
+
     def _get_solver(self, *, refresh: bool = False, penalty: Optional[Dict[str, int]] = None):
         """Get the least penalized solver from the list of solvers filtered and
         ordered according to user config.
@@ -234,8 +257,8 @@ class DWaveSampler(dimod.Sampler, dimod.Structured):
         Examples:
 
             >>> from dwave.system import DWaveSampler
-            >>> sampler = DWaveSampler()
-            >>> sampler.properties    # doctest: +SKIP
+            >>> with DWaveSampler() as sampler:     # doctest: +SKIP
+            ...     sampler.properties
             {'anneal_offset_ranges': [[-0.2197463755538704, 0.03821687759418928],
               [-0.2242514597680286, 0.01718456460967399],
               [-0.20860153999435985, 0.05511969218508182],
@@ -265,8 +288,8 @@ class DWaveSampler(dimod.Sampler, dimod.Structured):
         Examples:
 
             >>> from dwave.system import DWaveSampler
-            >>> sampler = DWaveSampler()
-            >>> sampler.parameters    # doctest: +SKIP
+            >>> with DWaveSampler() as sampler:     # doctest: +SKIP
+            ...     sampler.parameters
             {'anneal_offsets': ['parameters'],
              'anneal_schedule': ['parameters'],
              'annealing_time': ['parameters'],
@@ -296,8 +319,8 @@ class DWaveSampler(dimod.Sampler, dimod.Structured):
             First 5 entries of the coupler list for one Advantage system.
 
             >>> from dwave.system import DWaveSampler
-            >>> sampler = DWaveSampler()
-            >>> sampler.edgelist[:5]    # doctest: +SKIP
+            >>> with DWaveSampler() as sampler:     # doctest: +SKIP
+            ...     sampler.edgelist[:5]
             [(30, 31), (30, 45), (30, 2940), (30, 2955), (30, 2970)]
 
         See `Ocean Glossary <https://docs.ocean.dwavesys.com/en/stable/concepts/index.html>`_
@@ -320,8 +343,8 @@ class DWaveSampler(dimod.Sampler, dimod.Structured):
             First 5 entries of the node list for one Advantage system.
 
             >>> from dwave.system import DWaveSampler
-            >>> sampler = DWaveSampler()
-            >>> sampler.nodelist[:5]    # doctest: +SKIP
+            >>> with DWaveSampler() as sampler:     # doctest: +SKIP
+            ...     sampler.nodelist[:5]
             [30, 31, 32, 33, 34]
 
         See `Ocean Glossary <https://docs.ocean.dwavesys.com/en/stable/concepts/index.html>`_
@@ -403,14 +426,13 @@ class DWaveSampler(dimod.Sampler, dimod.Structured):
 
             >>> from dwave.system import DWaveSampler
             ...
-            >>> sampler = DWaveSampler()
-            ...
-            >>> qubit_a = sampler.nodelist[0]
-            >>> qubit_b = next(iter(sampler.adjacency[qubit_a]))
-            >>> sampleset = sampler.sample_ising({qubit_a: -1, qubit_b: 1},
-            ...                                  {},
-            ...                                  num_reads=100)
-            >>> print(sampleset.first.sample[qubit_a] == 1 and sampleset.first.sample[qubit_b] == -1)
+            >>> with DWaveSampler() as sampler:
+            ...     qubit_a = sampler.nodelist[0]
+            ...     qubit_b = next(iter(sampler.adjacency[qubit_a]))
+            ...     sampleset = sampler.sample_ising({qubit_a: -1, qubit_b: 1},
+            ...                                      {},
+            ...                                      num_reads=100)
+            ...     print(sampleset.first.sample[qubit_a] == 1 and sampleset.first.sample[qubit_b] == -1)
             True
 
         See `Ocean Glossary <https://docs.ocean.dwavesys.com/en/stable/concepts/index.html>`_
@@ -523,10 +545,9 @@ class DWaveSampler(dimod.Sampler, dimod.Structured):
             This example sets a quench schedule on a D-Wave system.
 
             >>> from dwave.system import DWaveSampler
-            >>> sampler = DWaveSampler()
-            >>> quench_schedule=[[0.0, 0.0], [12.0, 0.6], [12.8, 1.0]]
-            >>> DWaveSampler().validate_anneal_schedule(quench_schedule)    # doctest: +SKIP
-            >>>
+            >>> with DWaveSampler() as sampler:     # doctest: +SKIP
+            ...     quench_schedule=[[0.0, 0.0], [12.0, 0.6], [12.8, 1.0]]
+            ...     DWaveSampler().validate_anneal_schedule(quench_schedule)
 
         """
         if 'anneal_schedule' not in self.parameters:
@@ -597,9 +618,9 @@ class DWaveSampler(dimod.Sampler, dimod.Structured):
 
             >>> from dwave.system import DWaveSampler
             ...
-            >>> sampler = DWaveSampler()
-            >>> g = sampler.to_networkx_graph()      # doctest: +SKIP
-            >>> len(g.nodes) > 5000                  # doctest: +SKIP
+            >>> with DWaveSampler() as sampler:         # doctest: +SKIP
+            ...     g = sampler.to_networkx_graph()
+            ...     len(g.nodes) > 5000
             True
         """
         return qpu_graph(self.properties['topology']['type'],
diff --git a/dwave/system/samplers/leap_hybrid_sampler.py b/dwave/system/samplers/leap_hybrid_sampler.py
index bb917cc..f2caf34 100644
--- a/dwave/system/samplers/leap_hybrid_sampler.py
+++ b/dwave/system/samplers/leap_hybrid_sampler.py
@@ -19,6 +19,7 @@ A :std:doc:`dimod sampler <oceandocs:docs_dimod/reference/samplers>` for Leap's
 import concurrent.futures
 import warnings
 from collections import abc
+from contextlib import AbstractContextManager
 from numbers import Number
 from typing import Any, Dict, List, NamedTuple, Optional
 
@@ -38,7 +39,34 @@ __all__ = ['LeapHybridSampler',
            ]
 
 
-class LeapHybridSampler(dimod.Sampler):
+class _ClosableClientBaseMixin(AbstractContextManager):
+    """A mixin that implements ``close`` method to close the underlying cloud
+    client. It also implements a default context manager that closes resources
+    on exit.
+    """
+
+    def close(self):
+        """Close the underlying cloud client to release system resources such as
+        threads.
+
+        .. note::
+
+            The method blocks for all the currently scheduled work (sampling
+            requests) to finish.
+
+        See: :meth:`~dwave.cloud.client.Client.close`.
+        """
+        self.client.close()
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Release system resources allocated and raise any exception triggered
+        within the runtime context.
+        """
+        self.close()
+        return None
+
+
+class LeapHybridSampler(dimod.Sampler, _ClosableClientBaseMixin):
     """A class for using Leap's cloud-based hybrid BQM solvers.
 
     Leap's quantum-classical hybrid binary quadratic models (BQM) solvers are
@@ -82,23 +110,8 @@ class LeapHybridSampler(dimod.Sampler):
         >>> bqm = dimod.BQM.from_qubo(qubo)
         ...
         >>> # Find a good solution
-        >>> sampler = LeapHybridSampler()       # doctest: +SKIP
-        >>> sampleset = sampler.sample(bqm)     # doctest: +SKIP
-
-        This example specializes the default solver selection by filtering out
-        bulk BQM solvers. (Bulk solvers are throughput-optimal for heavy/batch
-        workloads, have a higher start-up latency, and are not well suited for
-        live workloads. Not all Leap accounts have access to bulk solvers.)
-
-        >>> from dwave.system import LeapHybridSampler
-        ...
-        >>> solver = LeapHybridSampler.default_solver
-        >>> solver.update(name__regex=".*(?<!bulk)$")       # name shouldn't end with "bulk"
-        >>> sampler = LeapHybridSampler(solver=solver)      # doctest: +SKIP
-        >>> sampler.solver        # doctest: +SKIP
-        BQMSolver(id='hybrid_binary_quadratic_model_version2')
-        ...
-        >>> sampler.close()       # doctest: +SKIP
+        >>> with LeapHybridSampler() as sampler:    # doctest: +SKIP
+        ...     sampleset = sampler.sample(bqm)
     """
 
     _INTEGER_BQM_SIZE_THRESHOLD = 10000
@@ -139,19 +152,6 @@ class LeapHybridSampler(dimod.Sampler):
         if 'bqm' not in self.solver.supported_problem_types:
             raise ValueError("selected solver does not support the 'bqm' problem type.")
 
-    def close(self):
-        """Close the underlying cloud client to release system resources such as
-        threads.
-
-        .. note::
-
-            The method blocks for all the currently scheduled work (sampling
-            requests) to finish.
-
-        See: :meth:`~dwave.cloud.client.Client.close`.
-        """
-        self.client.close()
-
     @property
     def properties(self) -> Dict[str, Any]:
         """Solver properties as returned by a SAPI query.
@@ -222,10 +222,8 @@ class LeapHybridSampler(dimod.Sampler):
             >>> bqm = dimod.BQM.from_qubo(qubo)
             ...
             >>> # Find a good solution
-            >>> sampler = LeapHybridSampler()       # doctest: +SKIP
-            >>> sampleset = sampler.sample(bqm)     # doctest: +SKIP
-            ...
-            >>> sampler.close()                     # doctest: +SKIP
+            >>> with LeapHybridSampler() as sampler:    # doctest: +SKIP
+            ...     sampleset = sampler.sample(bqm)
         """
 
         if not isinstance(bqm, dimod.BQM):
@@ -297,7 +295,7 @@ class LeapHybridSampler(dimod.Sampler):
 LeapHybridBQMSampler = LeapHybridSampler
 
 
-class LeapHybridDQMSampler:
+class LeapHybridDQMSampler(_ClosableClientBaseMixin):
     """A class for using Leap's cloud-based hybrid DQM solvers.
 
     Leap's quantum-classical hybrid DQM solvers are intended to solve arbitrary
@@ -353,14 +351,11 @@ class LeapHybridDQMSampler:
         ...          dqm.set_quadratic('my_hand', 'their_hand',
         ...                            {(my_idx, their_idx): 1})
         ...
-        >>> dqm_sampler = LeapHybridDQMSampler()      # doctest: +SKIP
-        ...
-        >>> sampleset = dqm_sampler.sample_dqm(dqm)   # doctest: +SKIP
-        >>> print("{} beats {}".format(cases[sampleset.first.sample['my_hand']],
-        ...                            cases[sampleset.first.sample['their_hand']]))   # doctest: +SKIP
+        >>> with LeapHybridDQMSampler() as dqm_sampler:     # doctest: +SKIP
+        ...     sampleset = dqm_sampler.sample_dqm(dqm)
+        ...     print(f"{} beats {}".format(cases[sampleset.first.sample['my_hand']],
+        ...                                 cases[sampleset.first.sample['their_hand']]))
         rock beats scissors
-        ...
-        >>> dqm_sampler.close()                       # doctest: +SKIP
     """
 
     @classproperty
@@ -398,19 +393,6 @@ class LeapHybridDQMSampler:
         if 'dqm' not in self.solver.supported_problem_types:
             raise ValueError("selected solver does not support the 'dqm' problem type.")
 
-    def close(self):
-        """Close the underlying cloud client to release system resources such as
-        threads.
-
-        .. note::
-
-            The method blocks for all the currently scheduled work (sampling
-            requests) to finish.
-
-        See: :meth:`~dwave.cloud.client.Client.close`.
-        """
-        self.client.close()
-
     @property
     def properties(self) -> Dict[str, Any]:
         """Solver properties as returned by a SAPI query.
@@ -561,7 +543,7 @@ class LeapHybridDQMSampler:
         return max([5, t])
 
 
-class LeapHybridCQMSampler:
+class LeapHybridCQMSampler(_ClosableClientBaseMixin):
     """A class for using Leap's cloud-based hybrid CQM solvers.
 
     Leap's quantum-classical hybrid CQM solvers are intended to solve
@@ -616,17 +598,15 @@ class LeapHybridCQMSampler:
         a remote solver provided by the Leap quantum cloud service:
 
         >>> from dwave.system import LeapHybridCQMSampler   # doctest: +SKIP
-        >>> sampler = LeapHybridCQMSampler()                # doctest: +SKIP
-        >>> sampleset = sampler.sample_cqm(cqm)             # doctest: +SKIP
-        >>> print(sampleset.first)                          # doctest: +SKIP
+        >>> with LeapHybridCQMSampler() as sampler:         # doctest: +SKIP
+        ...     sampleset = sampler.sample_cqm(cqm)
+        ...     print(sampleset.first)
         Sample(sample={'i': 2.0, 'j': 2.0}, energy=-4.0, num_occurrences=1,
         ...            is_feasible=True, is_satisfied=array([ True]))
 
         The best (lowest-energy) solution found has :math:`i=j=2` as expected,
         a solution that is feasible because all the constraints (one in this
         example) are satisfied.
-        ...
-        >>> sampler.close()                                 # doctest: +SKIP
     """
 
     def __init__(self, **config):
@@ -658,19 +638,6 @@ class LeapHybridCQMSampler:
         if 'cqm' not in self.solver.supported_problem_types:
             raise ValueError("selected solver does not support the 'cqm' problem type.")
 
-    def close(self):
-        """Close the underlying cloud client to release system resources such as
-        threads.
-
-        .. note::
-
-            The method blocks for all the currently scheduled work (sampling
-            requests) to finish.
-
-        See: :meth:`~dwave.cloud.client.Client.close`.
-        """
-        self.client.close()
-
     @classproperty
     def default_solver(cls) -> Dict[str, str]:
         """Features used to select the latest accessible hybrid CQM solver."""
@@ -832,7 +799,7 @@ class LeapHybridCQMSampler:
             )
 
 
-class LeapHybridNLSampler:
+class LeapHybridNLSampler(_ClosableClientBaseMixin):
     r"""A class for using Leap's cloud-based hybrid nonlinear-model solvers.
 
     Leap's quantum-classical hybrid nonlinear-model solvers are intended to
@@ -863,18 +830,15 @@ class LeapHybridNLSampler:
         >>> from dwave.optimization.generators import flow_shop_scheduling
         >>> from dwave.system import LeapHybridNLSampler
         ...
-        >>> sampler = LeapHybridNLSampler()     # doctest: +SKIP
-        ...
-        >>> processing_times = [[10, 5, 7], [20, 10, 15]]
-        >>> model = flow_shop_scheduling(processing_times=processing_times)
-        >>> results = sampler.sample(model, label="Small FSS problem")    # doctest: +SKIP
-        >>> job_order = next(model.iter_decisions())  # doctest: +SKIP
-        >>> print(f"State 0 of {model.objective.state_size()} has an "\   # doctest: +SKIP
-        ... f"objective value {model.objective.state(0)} for order " \    # doctest: +SKIP
-        ... f"{job_order.state(0)}.")     # doctest: +SKIP
+        >>> with LeapHybridNLSampler() as sampler:      # doctest: +SKIP
+        ...     processing_times = [[10, 5, 7], [20, 10, 15]]
+        ...     model = flow_shop_scheduling(processing_times=processing_times)
+        ...     results = sampler.sample(model, label="Small FSS problem")
+        ...     job_order = next(model.iter_decisions())
+        ...     print(f"State 0 of {model.objective.state_size()} has an "
+        ...           f"objective value {model.objective.state(0)} for order "
+        ...           f"{job_order.state(0)}.")
         State 0 of 8 has an objective value 50.0 for order [1. 2. 0.].
-        ...
-        >>> sampler.close()       # doctest: +SKIP
     """
 
     def __init__(self, **config):
@@ -919,7 +883,7 @@ class LeapHybridNLSampler:
 
         See: :meth:`~dwave.cloud.client.Client.close`.
         """
-        self.client.close()
+        super().close()
         self._executor.shutdown()
 
     @classproperty


Test patch:
diff --git a/dwave/system/testing.py b/dwave/system/testing.py
index 801ae4b..cb563f2 100644
--- a/dwave/system/testing.py
+++ b/dwave/system/testing.py
@@ -15,6 +15,7 @@
 import concurrent.futures
 import warnings
 import weakref
+from contextlib import AbstractContextManager
 from uuid import uuid4
 
 import numpy as np
@@ -25,7 +26,7 @@ from dwave.samplers import SteepestDescentSampler
 from dwave.system import qpu_graph
 
 
-class MockDWaveSampler(dimod.Sampler, dimod.Structured):
+class MockDWaveSampler(dimod.Sampler, dimod.Structured, AbstractContextManager):
     """Mock sampler modeled after DWaveSampler that can be used for tests.
 
     Properties and topology parameters are populated to qualitatively match
@@ -356,6 +357,10 @@ class MockDWaveSampler(dimod.Sampler, dimod.Structured):
     def close(self):
         pass
 
+    def __exit__(self, *_, **__):
+        self.close()
+        return None
+
     @classmethod
     def from_qpu_sampler(cls, sampler):
         return cls(properties=sampler.properties,
@@ -444,7 +449,7 @@ class MockDWaveSampler(dimod.Sampler, dimod.Structured):
                          self.nodelist, self.edgelist)
     
     
-class MockLeapHybridDQMSampler:
+class MockLeapHybridDQMSampler(AbstractContextManager):
     """Mock sampler modeled after LeapHybridDQMSampler that can be used for tests."""
     def __init__(self, **config):
         self.parameters = {'time_limit': ['parameters'],
@@ -466,6 +471,13 @@ class MockLeapHybridDQMSampler:
                            'maximum_number_of_variables': 3000,
                            'maximum_number_of_biases': 3000000000}
 
+    def close(self):
+        pass
+
+    def __exit__(self, *_, **__):
+        self.close()
+        return None
+
     def sample_dqm(self, dqm, **kwargs):
         num_samples = 12    # min num of samples from dqm solver
         samples = np.empty((num_samples, dqm.num_variables()), dtype=int)
@@ -482,7 +494,7 @@ class MockLeapHybridDQMSampler:
         return self.properties['minimum_time_limit'][0][1]
 
 
-class MockLeapHybridSolver:
+class MockLeapHybridSolver(AbstractContextManager):
 
     properties = {'supported_problem_types': ['bqm'],
                   'minimum_time_limit': [[1, 1.0], [1024, 1.0],
@@ -493,6 +505,13 @@ class MockLeapHybridSolver:
 
     supported_problem_types = ['bqm']
 
+    def close(self):
+        pass
+
+    def __exit__(self, *_, **__):
+        self.close()
+        return None
+
     def upload_bqm(self, bqm, **parameters):
         bqm_adjarray = dimod.serialization.fileview.load(bqm)
         future = concurrent.futures.Future()
diff --git a/tests/qpu/test_dwavesampler.py b/tests/qpu/test_dwavesampler.py
index b15dbd9..8836456 100644
--- a/tests/qpu/test_dwavesampler.py
+++ b/tests/qpu/test_dwavesampler.py
@@ -62,6 +62,15 @@ class TestDWaveSampler(unittest.TestCase):
                     h, J = self.nonzero_ising_problem(sampler)
                     sampler.sample_ising(h, J)
 
+        with self.subTest('verify context manager calls close'):
+            n_initial = threading.active_count()
+            with DWaveSampler():
+                n_active = threading.active_count()
+            n_closed = threading.active_count()
+
+            self.assertGreater(n_active, n_initial)
+            self.assertEqual(n_closed, n_initial)
+
     @staticmethod
     def nonzero_ising_problem(sampler):
         max_h = max(sampler.properties.get('h_range', [-1, 1]))
diff --git a/tests/qpu/test_leaphybrid_common.py b/tests/qpu/test_leaphybrid_common.py
index b85c5ee..092b982 100644
--- a/tests/qpu/test_leaphybrid_common.py
+++ b/tests/qpu/test_leaphybrid_common.py
@@ -143,3 +143,12 @@ class TestSamplerInterface(unittest.TestCase):
                 with self.assertRaises(UseAfterCloseError):
                     problem = self.problem_gen()
                     getattr(sampler, self.sample_meth)(problem)
+
+        with self.subTest('verify context manager calls close'):
+            n_initial = threading.active_count()
+            with self.sampler_cls():
+                n_active = threading.active_count()
+            n_closed = threading.active_count()
+
+            self.assertGreater(n_active, n_initial)
+            self.assertEqual(n_closed, n_initial)
diff --git a/tests/test_dwavecliquesampler.py b/tests/test_dwavecliquesampler.py
index bfaaea8..160f5e4 100644
--- a/tests/test_dwavecliquesampler.py
+++ b/tests/test_dwavecliquesampler.py
@@ -147,7 +147,6 @@ class TestDWaveCliqueSampler(unittest.TestCase):
         pegasus_sampler.sample(bqm, chain_strength=-0.5).resolve()
 
 
-
 class TestFailover(unittest.TestCase):
 
     @staticmethod
@@ -221,3 +220,10 @@ class TestFailover(unittest.TestCase):
         sampler.close()
 
         mock_child_sampler.return_value.close.assert_called_once()
+
+    @unittest.mock.patch('dwave.system.samplers.clique.DWaveSampler')
+    def test_context_manager(self, mock_child_sampler):
+        with DWaveCliqueSampler() as sampler:
+            ...
+
+        mock_child_sampler.return_value.close.assert_called_once()
diff --git a/tests/test_dwavesampler.py b/tests/test_dwavesampler.py
index f10ac7f..ca51a00 100644
--- a/tests/test_dwavesampler.py
+++ b/tests/test_dwavesampler.py
@@ -144,6 +144,17 @@ class TestDWaveSampler(unittest.TestCase):
 
         MockClient.from_config.return_value.close.assert_called_once()
 
+    @mock.patch('dwave.system.samplers.dwave_sampler.Client')
+    def test_context_manager(self, MockClient):
+        """DWaveSampler() can be used as a context manager and the underlying
+        client is closed on exit.
+        """
+
+        with DWaveSampler() as sampler:
+            ...
+
+        MockClient.from_config.return_value.close.assert_called_once()
+
     def test_sample_ising_variables(self):
 
         sampler = self.sampler
diff --git a/tests/test_leaphybridcqmsampler.py b/tests/test_leaphybridcqmsampler.py
index 36d1260..fa4c566 100644
--- a/tests/test_leaphybridcqmsampler.py
+++ b/tests/test_leaphybridcqmsampler.py
@@ -78,7 +78,14 @@ class TestTimeLimit(unittest.TestCase):
         mock_solver.properties = {'category': 'hybrid'}
         mock_solver.supported_problem_types = ['cqm']
 
-        sampler = LeapHybridCQMSampler()
-        sampler.close()
+        with self.subTest('manual close'):
+            sampler = LeapHybridCQMSampler()
+            sampler.close()
+            mock_client.from_config.return_value.close.assert_called_once()
+
+        mock_client.reset_mock()
 
-        mock_client.from_config.return_value.close.assert_called_once()
+        with self.subTest('context manager'):
+            with LeapHybridCQMSampler():
+                ...
+            mock_client.from_config.return_value.close.assert_called_once()
diff --git a/tests/test_leaphybriddqmsolver.py b/tests/test_leaphybriddqmsolver.py
index cd4d7ff..eac14ed 100644
--- a/tests/test_leaphybriddqmsolver.py
+++ b/tests/test_leaphybriddqmsolver.py
@@ -67,7 +67,14 @@ class TestLeapHybridDQMSampler(unittest.TestCase):
         mock_solver.properties = {'category': 'hybrid'}
         mock_solver.supported_problem_types = ['dqm']
 
-        sampler = LeapHybridDQMSampler()
-        sampler.close()
+        with self.subTest('manual close'):
+            sampler = LeapHybridDQMSampler()
+            sampler.close()
+            mock_client.from_config.return_value.close.assert_called_once()
+
+        mock_client.reset_mock()
 
-        mock_client.from_config.return_value.close.assert_called_once()
+        with self.subTest('context manager'):
+            with LeapHybridDQMSampler():
+                ...
+            mock_client.from_config.return_value.close.assert_called_once()
diff --git a/tests/test_leaphybridnlsampler.py b/tests/test_leaphybridnlsampler.py
index 1600120..4520da1 100644
--- a/tests/test_leaphybridnlsampler.py
+++ b/tests/test_leaphybridnlsampler.py
@@ -130,7 +130,14 @@ class TestNLSampler(unittest.TestCase):
         mock_solver.properties = {'category': 'hybrid'}
         mock_solver.supported_problem_types = ['nl']
 
-        sampler = LeapHybridNLSampler()
-        sampler.close()
+        with self.subTest('manual close'):
+            sampler = LeapHybridNLSampler()
+            sampler.close()
+            mock_client.from_config.return_value.close.assert_called_once()
+
+        mock_client.reset_mock()
 
-        mock_client.from_config.return_value.close.assert_called_once()
+        with self.subTest('context manager'):
+            with LeapHybridNLSampler():
+                ...
+            mock_client.from_config.return_value.close.assert_called_once()
diff --git a/tests/test_leaphybridsolver.py b/tests/test_leaphybridsolver.py
index 1fa44dc..7ea494a 100644
--- a/tests/test_leaphybridsolver.py
+++ b/tests/test_leaphybridsolver.py
@@ -115,10 +115,17 @@ class TestLeapHybridSampler(unittest.TestCase):
         mock_solver.properties = {'category': 'hybrid'}
         mock_solver.supported_problem_types = ['bqm']
 
-        sampler = LeapHybridSampler()
-        sampler.close()
+        with self.subTest('manual close'):
+            sampler = LeapHybridSampler()
+            sampler.close()
+            mock_client.from_config.return_value.close.assert_called_once()
+
+        mock_client.reset_mock()
 
-        mock_client.from_config.return_value.close.assert_called_once()
+        with self.subTest('context manager'):
+            with LeapHybridSampler():
+                ...
+            mock_client.from_config.return_value.close.assert_called_once()
 
     @mock.patch('dwave.system.samplers.leap_hybrid_sampler.Client')
     def test_sample_bqm(self, mock_client):