Add SampleSet.wait_id()

Author: randomir

dimod/sampleset.py

@@ -1489,9 +1489,28 @@ def resolve(self):
         if hasattr(self, '_future'):
             samples = self._result_hook(self._future)
             self.__init__(samples.record, samples.variables, samples.info, samples.vartype)
+
+            # cache problem id if we just resolved from a cloud-client's future
+            if hasattr(self._future, 'wait_id'):
+                self._future_wait_id_result = self._future.wait_id(timeout=0)
+
             del self._future
             del self._result_hook
 
+    def wait_id(self, timeout: Optional[float] = None) -> Optional[str]:
+        """Return a QPU problem ID when the sample set is constructed from a QPU
+        computation answer (future) using :meth:`.from_future`.
+
+        For details, see :meth:`~dwave.cloud.client.computation.Future.wait_id`.
+        """
+        # note: after the sample set is resolved and the underlying future deleted,
+        # we cache the ``wait_id`` result so that the future is not referenced and
+        # it can be gc'ed.
+        if hasattr(self, '_future_wait_id_result'):
+            return self._future_wait_id_result
+        if hasattr(self, '_future') and hasattr(self._future, 'wait_id'):
+            return self._future.wait_id(timeout=timeout)
+
     def aggregate(self):
         """Create a new SampleSet with repeated samples aggregated.
 

releasenotes/notes/add-sampleset.wait_id-as-proxy-to-qpu-future-9de9fbbc3dbc2285.yaml

@@ -0,0 +1,6 @@
+---
+features:
+  - |
+    Add ``SampleSet.wait_id()`` as a proxy to the underlying QPU result future,
+    when sample set is constructed from such a future.
+    See `#1392 <https://github.com/dwavesystems/dimod/issues/1392>`_.

tests/test_sampleset.py

@@ -786,6 +786,29 @@ def result_to_response(future):
 
         np.testing.assert_equal(matrix, result['samples'])
 
+    def test_cloud_client_future(self):
+        answer = dimod.SampleSet.from_samples(samples_like=[[1]], vartype='SPIN', energy=[1])
+        pid = 'problem-id'
+
+        future = concurrent.futures.Future()
+        future.wait_id = lambda timeout=None: pid
+
+        ss = dimod.SampleSet.from_future(future)
+
+        with self.subTest('unresolved sampleset from qpu computation has a functional wait_id'):
+            self.assertFalse(ss.done())
+            self.assertTrue(hasattr(ss, 'wait_id'))
+            self.assertEqual(ss.wait_id(), pid)
+
+        future.set_result(answer)
+        ss.resolve()
+
+        with self.subTest('resolved sampleset caches wait_id result'):
+            self.assertTrue(ss.done())
+            self.assertTrue(hasattr(ss, 'wait_id'))
+            self.assertEqual(ss.wait_id(), pid)
+            self.assertFalse(hasattr(ss, '_future'))
+
 
 class TestKeepVariables(unittest.TestCase):
     def test_empty(self):