Merge pull request #556 from randomir/feature/add-ctx-mgr-support-to-dwavesamplers/issue-91

Add support for context manager protocol to all samplers

Author: randomir

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 @@ def sample(self, bqm, **parameters):
 
         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 @@ def close(self):
         """
         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:

dwave/system/samplers/dwave_sampler.py

@@ -23,6 +23,7 @@
 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 @@ def close(self):
         """
         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 @@ def properties(self):
         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 @@ def parameters(self):
         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 @@ def edgelist(self):
             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 @@ def nodelist(self):
             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 @@ def sample(self, bqm, warnings=None, **kwargs):
 
             >>> 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 @@ def validate_anneal_schedule(self, anneal_schedule):
             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 @@ def to_networkx_graph(self):
 
             >>> 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'],