Merge pull request #1399 from randomir/feature/closeable-sampler-composite

Add scoped resource support

Author: randomir

dimod/core/__init__.py

@@ -16,4 +16,5 @@
 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

dimod/core/composite.py

@@ -48,17 +48,26 @@
 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 @@ def child(self):
         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.

dimod/core/initialized.py

@@ -21,7 +21,6 @@
 """
 import abc
 
-from collections import namedtuple
 from numbers import Integral
 
 import typing

dimod/core/polysampler.py

@@ -20,9 +20,9 @@
 
 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 @@
 __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 @@ def sample_hubo(self, H: collections.abc.Mapping[tuple[Variable, Variable], Bias
         """
         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.

dimod/core/sampler.py

@@ -112,6 +112,7 @@ def parameters(self):
 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,14 +165,23 @@ 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`,
     :meth:`~.Sampler.sample_qubo` assuming at least one is implemented.
 
     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 @@ def remove_unknown_kwargs(self, **kwargs) -> dict[str, typing.Any]:
             kwargs.pop(kw)
 
         return kwargs
+
+    def close(self):
+        """Release allocated resources.
+
+        Override to release sampler-allocated resources.
+        """
+        pass

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

dimod/core/structured.py

@@ -72,8 +72,6 @@ def sample(self, bqm):
 """
 from __future__ import annotations
 
-from collections import namedtuple
-
 import abc
 import typing
 

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

dimod/quadratic/quadratic_model.py

@@ -30,7 +30,6 @@
 from __future__ import annotations
 
 import collections.abc
-import struct
 import tempfile
 import typing
 

dimod/reference/composites/higherordercomposites.py

@@ -24,15 +24,14 @@
 * :meth:`.PolySampler.sample_hubo`
 
 """
-import warnings
 
 from collections import defaultdict
 
 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