Repository: dwavesystems/dimod
Pull request #1399 (2025-03-11 15:32:19)

Add scoped resource support
Addresses https://github.com/dwavesystems/dwave-system/issues/555 on a lower-level.


Patch:
diff --git a/dimod/core/__init__.py b/dimod/core/__init__.py
index ada7a61a..bb8280a0 100644
--- a/dimod/core/__init__.py
+++ b/dimod/core/__init__.py
@@ -16,4 +16,5 @@ from dimod.core.composite import Composite, ComposedSampler
 from dimod.core.initialized import Initialized
 from dimod.core.polysampler import PolySampler, ComposedPolySampler
 from dimod.core.sampler import Sampler
+from dimod.core.scoped import Scoped
 from dimod.core.structured import Structured
diff --git a/dimod/core/composite.py b/dimod/core/composite.py
index af5532e7..753d6a47 100644
--- a/dimod/core/composite.py
+++ b/dimod/core/composite.py
@@ -48,17 +48,26 @@ Examples:
 import abc
 
 from dimod.core.sampler import Sampler
+from dimod.core.scoped import Scoped
 
 __all__ = ['Composite', 'ComposedSampler']
 
 
-class Composite(abc.ABC):
+class Composite(Scoped):
     """Abstract base class for dimod composites.
 
     Provides the :attr:`Composite.child` mixin property and defines the :attr:`Composite.children`
     abstract property to be implemented. These define the supported samplers for the composed sampler.
 
+    :class:`.Composite` implements the :class:`~dimod.core.scoped.Scoped` interface, and
+    provides :meth:`.close` method that closes all child samplers or composites.
+
+    .. versionchanged:: 0.12.19
+        Implemented the :class:`~dimod.core.scoped.Scoped` interface. Now all
+        composites support context manager protocol and release scope-based
+        resources of sub-samplers/composites by default.
     """
+
     @abc.abstractproperty
     def children(self):
         """list[ :obj:`.Sampler`]: List of child samplers that that are used by
@@ -74,6 +83,29 @@ class Composite(abc.ABC):
         except IndexError:
             raise RuntimeError("A Composite must have at least one child Sampler")
 
+    def close(self):
+        """Release any scope-bound resources of child samplers or composites.
+
+        .. note::
+            If a :class:`.Composite` subclass doesn't allocate resources that have
+            to be explicitly released, there's no need to override the default
+            :meth:`~.Composite.close` implementation.
+
+            However, if you do implement :meth:`~.Composite.close` on a subclass,
+            make sure to either call ``super().close()``, or to explicitly close
+            all child samplers/composites.
+
+        .. versionadded:: 0.12.19
+            :class:`.Composite` now implements the :class:`~dimod.core.scoped.Scoped`
+            interface. The default :meth:`~.Composite.close` method recursively closes
+            composite's children.
+        """
+
+        for child in self.children:
+            if hasattr(child, 'close') and callable(child.close):
+                child.close()
+        super().close()
+
 
 class ComposedSampler(Sampler, Composite):
     """Abstract base class for dimod composed samplers.
diff --git a/dimod/core/initialized.py b/dimod/core/initialized.py
index 20ea8571..32c0ed49 100644
--- a/dimod/core/initialized.py
+++ b/dimod/core/initialized.py
@@ -21,7 +21,6 @@ this abstract base class in a sampler.
 """
 import abc
 
-from collections import namedtuple
 from numbers import Integral
 
 import typing
diff --git a/dimod/core/polysampler.py b/dimod/core/polysampler.py
index d225b766..5802c1a8 100644
--- a/dimod/core/polysampler.py
+++ b/dimod/core/polysampler.py
@@ -20,9 +20,9 @@ not constrained to quadratic interactions.
 
 import abc
 import collections.abc
-import warnings
 
 from dimod.core.composite import Composite
+from dimod.core.scoped import Scoped
 from dimod.higherorder.polynomial import BinaryPolynomial
 from dimod.sampleset import SampleSet
 
@@ -31,12 +31,16 @@ from dimod.typing import Bias, Variable
 __all__ = 'PolySampler', 'ComposedPolySampler'
 
 
-class PolySampler(abc.ABC):
+class PolySampler(Scoped):
     """Sampler that supports binary polynomials.
 
     Binary polynomials are an extension of binary quadratic models that allow
     higher-order interactions.
 
+    .. versionchanged:: 0.12.19
+        :class:`.PolySampler` now implements the :class:`~dimod.core.scoped.Scoped`
+        interface, so it supports context manager protocol by default.
+
     """
     @abc.abstractproperty  # for python2 compatibility
     def parameters(self):
@@ -105,6 +109,13 @@ class PolySampler(abc.ABC):
         """
         return self.sample_poly(BinaryPolynomial.from_hubo(H), **kwargs)
 
+    def close(self):
+        """Release allocated resources.
+
+        Override to release sampler-allocated resources.
+        """
+        pass
+
 
 class ComposedPolySampler(PolySampler, Composite):
     """Abstract base class for dimod composed polynomial samplers.
diff --git a/dimod/core/sampler.py b/dimod/core/sampler.py
index 7eb8c3c9..fd95d676 100644
--- a/dimod/core/sampler.py
+++ b/dimod/core/sampler.py
@@ -112,6 +112,7 @@ import typing
 import warnings
 
 from dimod.binary_quadratic_model import BinaryQuadraticModel
+from dimod.core.scoped import Scoped
 from dimod.exceptions import InvalidSampler, SamplerUnknownArgWarning
 from dimod.sampleset import SampleSet
 from dimod.typing import Bias, Variable
@@ -164,7 +165,7 @@ def samplemixinmethod(method):
     return method
 
 
-class Sampler(metaclass=SamplerABCMeta):
+class Sampler(Scoped, metaclass=SamplerABCMeta):
     """Abstract base class for dimod samplers.
 
     Provides all methods :meth:`~.Sampler.sample`, :meth:`~.Sampler.sample_ising`,
@@ -172,6 +173,15 @@ class Sampler(metaclass=SamplerABCMeta):
 
     Also includes utility method :meth:`~.Sampler.remove_unknown_kwargs`,
     which may be used in sample methods to handle unknown kwargs.
+
+    :class:`.Sampler` implements the :class:`~dimod.core.scoped.Scoped` interface,
+    but the default implementation of :meth:`~.Sampler.close` does nothing. It
+    should be overridden to release any scope-based sampler resources.
+
+    .. versionchanged:: 0.12.19
+        :class:`.Sampler` now implements the :class:`~dimod.core.scoped.Scoped`
+        interface, meaning all samplers support context manager protocol by
+        default.
     """
 
     @abc.abstractproperty  # for python2 compatibility
@@ -317,3 +327,10 @@ class Sampler(metaclass=SamplerABCMeta):
             kwargs.pop(kw)
 
         return kwargs
+
+    def close(self):
+        """Release allocated resources.
+
+        Override to release sampler-allocated resources.
+        """
+        pass
diff --git a/dimod/core/scoped.py b/dimod/core/scoped.py
new file mode 100644
index 00000000..52e71bf4
--- /dev/null
+++ b/dimod/core/scoped.py
@@ -0,0 +1,60 @@
+# Copyright 2025 D-Wave
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import abc
+
+__all__ = ['Scoped']
+
+
+class Scoped(abc.ABC):
+    """Abstract base class for components that allocate scope-bound resources.
+
+    :class:`.Scoped` requires the concrete implementation to provide resources
+    clean-up via :meth:`~.Scoped.close` method. With that, context manager
+    protocol is automatically supported (resources are released on context exit).
+
+    The scoped resource should either be used from a `runtime context`_, e.g.:
+
+    >>> with SomeScopedSampler() as sampler:       # doctest: +SKIP
+    ...     sampler.sample(...)
+
+    or explicitly disposed with a call to the :meth:`.close` method.
+
+    .. _runtime context: https://docs.python.org/3/reference/datamodel.html#with-statement-context-managers
+
+    .. versionadded:: 0.12.19
+        :class:`.Scoped` abstract base class.
+
+    Both :class:`~dimod.code.sampler.Sampler` and :class:`~dimod.core.composite.Composite`
+    implement :class:`.Scoped` interface.
+    """
+
+    @abc.abstractmethod
+    def close(self):
+        """Release allocated system resources that are hard or impossible to
+        garbage-collect.
+        """
+        pass
+
+    def __enter__(self):
+        """Return `self` upon entering the runtime context."""
+        return self
+
+    def __exit__(self, exc_type, exc_value, traceback):
+        """Release system resources allocated and raise any exception triggered
+        within the runtime context.
+        """
+        self.close()
+        # raise exceptions from the runtime context as they're not handled here
+        return None
diff --git a/dimod/core/structured.py b/dimod/core/structured.py
index 3a2b059e..65e733a3 100644
--- a/dimod/core/structured.py
+++ b/dimod/core/structured.py
@@ -72,8 +72,6 @@ Examples:
 """
 from __future__ import annotations
 
-from collections import namedtuple
-
 import abc
 import typing
 
diff --git a/dimod/discrete/discrete_quadratic_model.py b/dimod/discrete/discrete_quadratic_model.py
index cd1d660e..b5cdf9a9 100644
--- a/dimod/discrete/discrete_quadratic_model.py
+++ b/dimod/discrete/discrete_quadratic_model.py
@@ -14,14 +14,12 @@
 
 import collections.abc as abc
 import io
-import json
 import warnings
 
 from collections import defaultdict, namedtuple
 from typing import List, Tuple, Union, Generator, Iterator
 
 import numpy as np
-from numpy.core.shape_base import stack
 
 from dimod.discrete.cydiscrete_quadratic_model import cyDiscreteQuadraticModel
 from dimod.sampleset import as_samples
diff --git a/dimod/quadratic/quadratic_model.py b/dimod/quadratic/quadratic_model.py
index 8cb2c6a0..1ca6cbbe 100644
--- a/dimod/quadratic/quadratic_model.py
+++ b/dimod/quadratic/quadratic_model.py
@@ -30,7 +30,6 @@ variables and :math:`a_{i}, b_{ij}, c` are real values.
 from __future__ import annotations
 
 import collections.abc
-import struct
 import tempfile
 import typing
 
diff --git a/dimod/reference/composites/higherordercomposites.py b/dimod/reference/composites/higherordercomposites.py
index faeb53ea..e47566ce 100644
--- a/dimod/reference/composites/higherordercomposites.py
+++ b/dimod/reference/composites/higherordercomposites.py
@@ -24,7 +24,6 @@ Higher-order composites implement three sampling methods (similar to
 * :meth:`.PolySampler.sample_hubo`
 
 """
-import warnings
 
 from collections import defaultdict
 
@@ -32,7 +31,7 @@ import numpy as np
 
 from dimod.core.polysampler import ComposedPolySampler, PolySampler
 from dimod.higherorder.polynomial import BinaryPolynomial
-from dimod.higherorder.utils import make_quadratic, poly_energies
+from dimod.higherorder.utils import make_quadratic
 from dimod.sampleset import SampleSet, append_variables
 
 
diff --git a/dimod/reference/composites/truncatecomposite.py b/dimod/reference/composites/truncatecomposite.py
index 744a0e12..da987397 100644
--- a/dimod/reference/composites/truncatecomposite.py
+++ b/dimod/reference/composites/truncatecomposite.py
@@ -17,8 +17,6 @@ A composite that truncates the returned :obj:`dimod.SampleSet` based on options
 specified by the user.
 """
 
-import numpy as np
-
 from dimod.core.composite import ComposedSampler
 
 __all__ = 'TruncateComposite',
diff --git a/docs/reference/generators.rst b/docs/reference/generators.rst
index 408a4b22..e89092de 100644
--- a/docs/reference/generators.rst
+++ b/docs/reference/generators.rst
@@ -51,9 +51,9 @@ Optimization
    maximum_weight_independent_set
    mimo
    multi_knapsack
+   multi_quadratic_knapsack
    quadratic_assignment
    quadratic_knapsack
-   quadratic_multi_knapsack
    random_bin_packing
    random_knapsack
    random_multi_knapsack
diff --git a/docs/reference/sampler_composites/api.rst b/docs/reference/sampler_composites/api.rst
index 0324414b..37f166bd 100644
--- a/docs/reference/sampler_composites/api.rst
+++ b/docs/reference/sampler_composites/api.rst
@@ -22,44 +22,51 @@ ABCs.
         - Abstract Properties
         - Abstract Methods
         - Mixin Methods
-    *   - :class:`.Sampler`
+    *   - :class:`.Scoped`
+        -
         -
+        - :meth:`~.Scoped.close`
+        - :meth:`~.Scoped.__enter__`, :meth:`~.Scoped.__exit__`
+    *   - :class:`.Sampler`
+        - :class:`.Scoped`
         - :attr:`~.Sampler.parameters`, :attr:`~.Sampler.properties`
         - at least one of
           :meth:`~.Sampler.sample`, :meth:`~.Sampler.sample_ising`, :meth:`~.Sampler.sample_qubo`
-        - :meth:`~.Sampler.sample`, :meth:`~.Sampler.sample_ising`, :meth:`~.Sampler.sample_qubo`
+        - :meth:`~.Sampler.sample`, :meth:`~.Sampler.sample_ising`, :meth:`~.Sampler.sample_qubo`,
+          :meth:`~.Sampler.close`
     *   - :class:`.Structured`
         -
         - :attr:`~.Structured.nodelist`, :attr:`~.Structured.edgelist`
         -
         - :attr:`~.Structured.structure`, :attr:`~.Structured.adjacency`, :meth:`~.Structured.to_networkx_graph`
     *   - :class:`.Composite`
-        -
+        - :class:`.Scoped`
         - :attr:`~.Composite.children`
         -
-        - :attr:`~.Composite.child`
+        - :attr:`~.Composite.child`, :meth:`~.Composite.close`
     *   - :class:`.ComposedSampler`
         - :class:`.Sampler`, :class:`.Composite`
         - :attr:`~.Sampler.parameters`, :attr:`~.Sampler.properties`, :attr:`~.Composite.children`
         - at least one of
           :meth:`~.Sampler.sample`, :meth:`~.Sampler.sample_ising`, :meth:`~.Sampler.sample_qubo`
         - :meth:`~.Sampler.sample`, :meth:`~.Sampler.sample_ising`, :meth:`~.Sampler.sample_qubo`,
-          :attr:`~.Composite.child`
+          :attr:`~.Composite.child`, :meth:`~.Composite.close`
     *   - :class:`.Initialized`
         -
         -
         -
         - :attr:`~.Initialized.parse_initial_states`
     *   - :class:`.PolySampler`
-        -
+        - :class:`.Scoped`
         - :attr:`~.PolySampler.parameters`, :attr:`~.PolySampler.properties`
         - :meth:`~.PolySampler.sample_poly`
-        - :meth:`~.PolySampler.sample_hising`, :meth:`~.PolySampler.sample_hubo`
+        - :meth:`~.PolySampler.sample_hising`, :meth:`~.PolySampler.sample_hubo`, :meth:`~.PolySampler.close`
     *   - :class:`.ComposedPolySampler`
         - :class:`.PolySampler`, :class:`.Composite`
         - :attr:`~.PolySampler.parameters`, :attr:`~.PolySampler.properties`, :attr:`~.Composite.children`
         - :meth:`~.Sampler.sample_poly`
-        - :meth:`~.Sampler.sample_hising`, :meth:`~.Sampler.sample_hubo`, :attr:`~.Composite.child`
+        - :meth:`~.Sampler.sample_hising`, :meth:`~.Sampler.sample_hubo`, :attr:`~.Composite.child`,
+          :meth:`~.PolySampler.close`
 
 The table shows, for example, that the :class:`.Sampler` class requires that you implement
 the :attr:`~.Sampler.parameters` and :attr:`~.Sampler.properties` properties and at least
@@ -99,6 +106,7 @@ Methods
    :toctree: generated/
 
    Sampler.remove_unknown_kwargs
+   Sampler.close
 
 Creating a Composed Sampler
 ===========================
@@ -127,7 +135,6 @@ Abstract Properties
 
    Composite.children
 
-
 Mixin Properties
 ----------------
 
@@ -136,6 +143,13 @@ Mixin Properties
 
    Composite.child
 
+Methods
+-------
+
+.. autosummary::
+   :toctree: generated/
+
+   Composite.close
 
 
 Creating a Structured Sampler
@@ -173,8 +187,8 @@ Mixin Methods
    Structured.to_networkx_graph
 
 
-Creating a Initialized Sampler
-==============================
+Creating an Initialized Sampler
+===============================
 
 .. automodule:: dimod.core.initialized
 
@@ -190,6 +204,32 @@ Mixin Methods
    Initialized.parse_initial_states
 
 
+Creating a Scoped Sampler
+=========================
+
+.. automodule:: dimod.core.scoped
+
+.. currentmodule:: dimod
+.. autoclass:: Scoped
+
+Abstract Methods
+----------------
+
+.. autosummary::
+   :toctree: generated/
+
+   Scoped.close
+
+Mixin Methods
+-------------
+
+.. autosummary::
+   :toctree: generated/
+
+   Scoped.__enter__
+   Scoped.__exit__
+
+
 Creating a Binary Polynomial Sampler
 ====================================
 
@@ -224,6 +264,15 @@ Mixin Methods
    PolySampler.sample_hising
    PolySampler.sample_hubo
 
+Methods
+-------
+
+.. autosummary::
+   :toctree: generated/
+
+   PolySampler.close
+
+
 Creating a Composed Binary Polynomial Sampler
 =============================================
 
diff --git a/releasenotes/notes/add-scoped-abc-53b64b69ddd5c1ac.yaml b/releasenotes/notes/add-scoped-abc-53b64b69ddd5c1ac.yaml
new file mode 100644
index 00000000..cdbeab61
--- /dev/null
+++ b/releasenotes/notes/add-scoped-abc-53b64b69ddd5c1ac.yaml
@@ -0,0 +1,9 @@
+---
+features:
+  - |
+    Add ``dimod.Scoped`` abstract base class for components that allocate
+    scope-bound resources.
+  - |
+    Implement ``dimod.Scoped`` interface on ``dimod.Sampler`` and ``dimod.Composite``.
+
+    See `dwave-system#555 <https://github.com/dwavesystems/dwave-system/issues/555>`_.


Test patch:
diff --git a/tests/test_composite.py b/tests/test_composite.py
index cdf8b65e..20b0738c 100644
--- a/tests/test_composite.py
+++ b/tests/test_composite.py
@@ -13,6 +13,7 @@
 #    limitations under the License.
 
 import unittest
+from unittest import mock
 
 import dimod
 
@@ -39,3 +40,63 @@ class TestCompositeClass(unittest.TestCase):
         sampler = Dummy()
         with self.assertRaises(RuntimeError):
             sampler.child
+
+    def test_close(self):
+        sampler = mock.MagicMock()
+
+        class Dummy(dimod.Composite):
+            @property
+            def children(self):
+                return [sampler]
+
+        composite = Dummy()
+        composite.close()
+
+        composite.child.close.assert_called_once()
+
+    def test_context_cleanup(self):
+        class Inner(dimod.Sampler):
+            parameters = None
+            properties = None
+
+            def __init__(self):
+                self.resource = mock.Mock()
+
+            def close(self):
+                self.resource.close()
+
+            def sample(self, bqm):
+                pass
+
+        class Mid(dimod.Composite):
+            def __init__(self, *children):
+                self._children = children
+
+            @property
+            def children(self):
+                return self._children
+
+        class Outer(dimod.Composite):
+            def __init__(self, child):
+                self._child = child
+                self.resource = mock.Mock()
+
+            @property
+            def children(self):
+                return [self._child]
+
+            def close(self):
+                super().close()
+                self.resource.close()
+
+        inner_1 = Inner()
+        inner_2 = Inner()
+        mid = Mid(inner_1, inner_2)
+        with Outer(mid) as outer:
+            pass
+
+        # closed by Composite base close impl
+        inner_1.resource.close.assert_called_once()
+        inner_2.resource.close.assert_called_once()
+        # closed by subclass
+        outer.resource.close.assert_called_once()
diff --git a/tests/test_sampler.py b/tests/test_sampler.py
index 3906bf80..0216df55 100644
--- a/tests/test_sampler.py
+++ b/tests/test_sampler.py
@@ -14,6 +14,7 @@
 
 import concurrent.futures
 import unittest
+from unittest import mock
 
 import numpy as np
 
@@ -362,3 +363,20 @@ class TestSamplerClass(unittest.TestCase):
 
         # Check that known kwargs are kept and unknown kwargs are removed
         self.assertDictEqual(kwargs, {'a': 1})
+
+    def test_context_manager_support(self):
+        class Sampler(dimod.Sampler):
+            parameters = None
+            properties = None
+
+            def sample(self, bqm):
+                pass
+
+            def close(self):
+                pass
+
+        with mock.patch.object(Sampler, 'close'):
+            with Sampler() as sampler:
+                pass
+
+            sampler.close.assert_called_once()
diff --git a/tests/test_scoped.py b/tests/test_scoped.py
new file mode 100644
index 00000000..d2585fac
--- /dev/null
+++ b/tests/test_scoped.py
@@ -0,0 +1,35 @@
+# Copyright 2025 D-Wave
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import unittest
+from unittest import mock
+
+import dimod
+
+
+class TestScopedClass(unittest.TestCase):
+    def test_instantiation_base_class(self):
+        with self.assertRaises(TypeError):
+            dimod.Scoped()
+
+    def test_context(self):
+        class ScopedResource(dimod.Scoped):
+            def close(self):
+                pass
+
+        with mock.patch.object(ScopedResource, 'close'):
+            with ScopedResource() as scoped:
+                pass
+
+            scoped.close.assert_called_once()