Repository: dwavesystems/dwave-system

Issue:
Expose submitted problem ID from `LeapHybridNLSampler`
Problem ID is available in `SampleSet.info` for other solver/samplers, but not in the `LeapHybridNLSampler.SampleResult` named tuple.

One option is to expand the `SampleResult` tuple with a third field, `info`, but that breaks backward compatibility.

Another option is to generalize the `timing` field, but that again breaks compatibility if the field is renamed.

So the only "safe" option is to tuck `problem_id` in the `timing` dict, but that's too ugly.

Hints from issue comments:
@rahmanshamit: Is the Problem ID referred in this issue, the problem_data_id from here in line 945? https://github.com/dwavesystems/dwave-system/blob/1bb84c6e73513c66f31dc4bee06aeab6b9e6743d/dwave/system/samplers/leap_hybrid_sampler.py#L945

If that's the case, one still slightly ugly but backwards compatible way might be to have a simple wrapper for the SampleResult class: https://github.com/dwavesystems/dwave-system/blob/1bb84c6e73513c66f31dc4bee06aeab6b9e6743d/dwave/system/samplers/leap_hybrid_sampler.py#L901

The updated code would be:
```
class SampleResult(NamedTuple):
        model: "dwave.optimization.Model"
        timing: dict

"""  Wrapper for SampleResult to include additional metadata like Problem ID """
class ExtendedSampleResult:
      def __init__(self, sample_result, problem_id=None):
          self._sample_result = sample_result
          self.info = {"problem_id": problem_data_id} if problem_data_id else { }

      def __getattr__(self, item):
          """ Delegate access to the original SampleResult attributes """
          return getattr(self._sample_result, item)
```

and then slightly update the collect function: https://github.com/dwavesystems/dwave-system/blob/1bb84c6e73513c66f31dc4bee06aeab6b9e6743d/dwave/system/samplers/leap_hybrid_sampler.py#L961
Right now it returns LeapHybridNLSampler.SampleResult(model, timing).

The update would need to be:
```
sample_result = LeapHybridNLSampler.SampleResult(model, timing)
return LeapHybridNLSampler.WrappedSampleResult(sample_result, problem_id=problem_data_id)
```

This method doesn't effect any past uses and introduces no changes to existing fields or variables. @randomir: > Is the Problem ID referred in this issue, the problem_data_id from here in line 945?

No, it's the problem ID as returned by SAPI on sampling job submit. We refer to `problem_data_id` (problem data) when submitting the sampling job.@rahmanshamit: I see, hope I'm not hogging up the comments of this issue - would we be able to do sampler.properties['problem_id'] and then pass that to the SampleResult wrapper from the previous comment to still maintain all current uses of SampleResult and its fields?@randomir: No, `problem_id` is not a sampler/solver property. it's only available on the future returned by `sample_nlm()` on line 947:

https://github.com/dwavesystems/dwave-system/blob/1bb84c6e73513c66f31dc4bee06aeab6b9e6743d/dwave/system/samplers/leap_hybrid_sampler.py#L947

To retrieve it, you can use [`.id`](https://github.com/dwavesystems/dwave-cloud-client/blob/90de738edddded5a48b56a324d75573901e76c03/dwave/cloud/computation.py#L528), or safer, [`.wait_id()`](https://github.com/dwavesystems/dwave-cloud-client/blob/90de738edddded5a48b56a324d75573901e76c03/dwave/cloud/computation.py#L502). So, you would want to modify `collect()`:

https://github.com/dwavesystems/dwave-system/blob/1bb84c6e73513c66f31dc4bee06aeab6b9e6743d/dwave/system/samplers/leap_hybrid_sampler.py#L955-L961

to look like:
```python
    def collect():
        timing = future.timing
        problem_id = future.wait_id()
        ...
        # store problem_id in the result
```

@rahmanshamit: I see, thank you!

In that case, could we get the problem_id in collect as shown above and then rather than returning 'LeapHybridNLSampler.SampleResult(model, timing)', return using the SampleResult wrapper? 
This way we still don't change any existing use of SampleResult, don't generalize timing or tuck it in the timing dict.

All code changes would look like this:
```
class SampleResult(NamedTuple):
        model: "dwave.optimization.Model"
        timing: dict

class ExtendedSampleResult:
    """A wrapper for SampleResult to include problem_id."""
    def __init__(self, sample_result, problem_id):
        self.model = sample_result.model
        self.timing = sample_result.timing
        self.problem_id = problem_id
```

```
def collect():
      timing = future.timing
      problem_id = future.wait_id()
      for msg in timing.get('warnings', []):
          # note: no point using stacklevel, as this is a different thread
          warnings.warn(msg, category=UserWarning)

      sample_result = LeapHybridNLSampler.SampleResult(model, timing)
      return ExtendedSampleResult(sample_result, problem_id)
```@randomir: Unfortunately, a `namedtuple` (what `SampleResult` is) can't be extended that easily. Its `__slots__` are set to empty, and we can't just add any attribute to it. Note that
```
model, timing = sample_result
```
is a valid assignment. That's why adding a third attribute would break backwards compatibility.

But looks like we'll have to break it in the end if we want to return problem metadata like the id, label, etc.@rahmanshamit: I see, that makes sense.
Thank you for the feedback! Unfortunately at the moment that might be our best option.

Solution patch:
diff --git a/dwave/system/samplers/__init__.py b/dwave/system/samplers/__init__.py
index a228235..bbc8ef5 100644
--- a/dwave/system/samplers/__init__.py
+++ b/dwave/system/samplers/__init__.py
@@ -12,6 +12,22 @@
 #    See the License for the specific language governing permissions and
 #    limitations under the License.
 
+import typing
+
+
+class ResultInfoDict(typing.TypedDict, total=False):
+    """Returned in ``SampleSet.info`` and ``LeapHybridNLSampler.SampleResult.info``.
+
+    Not all fields defined below are always set, and additional might be set.
+    """
+
+    timing: dict[str, float]
+    warnings: list[str]
+    problem_id: str
+    problem_label: str
+    problem_data_id: str
+
+
 from dwave.system.samplers.clique import *
 from dwave.system.samplers.dwave_sampler import *
 from dwave.system.samplers.leap_hybrid_sampler import *
diff --git a/dwave/system/samplers/leap_hybrid_sampler.py b/dwave/system/samplers/leap_hybrid_sampler.py
index 8bf3fb5..8812482 100644
--- a/dwave/system/samplers/leap_hybrid_sampler.py
+++ b/dwave/system/samplers/leap_hybrid_sampler.py
@@ -27,6 +27,7 @@ import dwave.optimization
 import numpy
 from dwave.cloud.client import Client
 
+from dwave.system.samplers import ResultInfoDict
 from dwave.system.utilities import classproperty, FeatureFlags
 
 
@@ -918,7 +919,7 @@ class LeapHybridNLSampler(_ScopedSamplerMixin):
 
     class SampleResult(NamedTuple):
         model: dwave.optimization.Model
-        timing: dict
+        info: ResultInfoDict
 
     def sample(self, model: dwave.optimization.Model,
                time_limit: Optional[float] = None, **kwargs
@@ -945,10 +946,16 @@ class LeapHybridNLSampler(_ScopedSamplerMixin):
 
         Returns:
             :class:`~concurrent.futures.Future` [SampleResult]:
-                Named tuple containing nonlinear model and timing info, in a Future.
+                Named tuple, in a Future, containing the nonlinear model and general
+                result information such as timing and the identity of the problem data.
+
+        .. versionchanged:: 1.31.0
+            The return value includes timing information as part of the ``info``
+            field dictionary, which now replaces the previous ``timing`` field.
         """
 
         if not isinstance(model, dwave.optimization.Model):
+
             raise TypeError("first argument 'model' must be a dwave.optimization.Model, "
                             f"received {type(model).__name__}")
 
@@ -971,12 +978,20 @@ class LeapHybridNLSampler(_ScopedSamplerMixin):
         model.states.from_future(future, hook)
 
         def collect():
-            timing = future.timing
-            for msg in timing.get('warnings', []):
+            timing = future.timing.copy()
+            info = dict(
+                timing=timing,
+                warnings=timing.pop('warnings', []),
+                # match SampleSet.info fields (see :meth:`~dwave.cloud.computation.Future._get_problem_info`)
+                problem_id=future.id,
+                problem_label=future.label,
+                problem_data_id=problem_data_id,
+            )
+            for msg in info['warnings']:
                 # note: no point using stacklevel, as this is a different thread
                 warnings.warn(msg, category=UserWarning)
 
-            return LeapHybridNLSampler.SampleResult(model, timing)
+            return LeapHybridNLSampler.SampleResult(model, info)
 
         result = self._executor.submit(collect)
 


Test patch:
diff --git a/tests/test_leaphybridnlsampler.py b/tests/test_leaphybridnlsampler.py
index 4520da1..e80137a 100644
--- a/tests/test_leaphybridnlsampler.py
+++ b/tests/test_leaphybridnlsampler.py
@@ -71,16 +71,20 @@ class TestNLSampler(unittest.TestCase):
         self.assertEqual(model.states.size(), 2)
 
         # upload is tested in dwave-cloud-client, we can just mock it here
+        mock_problem_id = '321'
         mock_problem_data_id = '123'
-        mock_timing_info = {'qpu_access_time': 1, 'run_time': 2}
+        mock_problem_label = 'mock-label'
+        mock_timing = {'qpu_access_time': 1, 'run_time': 2}
+        mock_warnings = ['solved by classical presolve techniques']
 
         # note: instead of simply mocking `sampler.solver`, we mock a set of
         # solver methods minimally required to fully test `NLSolver.sample_problem`
 
         upload_nlm.return_value.result.return_value = mock_problem_data_id
 
-        base_sample_problem.return_value = Future(solver=sampler.solver, id_="x")
+        base_sample_problem.return_value = Future(solver=sampler.solver, id_=mock_problem_id)
         base_sample_problem.return_value._set_message({"answer": {}})
+        base_sample_problem.return_value.label = mock_problem_label
 
         def mock_decode_response(msg, answer_data: io.IOBase):
             # copy model states to the "received" answer_data
@@ -88,7 +92,7 @@ class TestNLSampler(unittest.TestCase):
             answer_data.seek(0)
             return {
                 'problem_type': 'nl',
-                'timing': mock_timing_info,
+                'timing': mock_timing | dict(warnings=mock_warnings),
                 'shape': {},
                 'answer': answer_data
             }
@@ -96,31 +100,35 @@ class TestNLSampler(unittest.TestCase):
 
         time_limit = 5
 
-        result = sampler.sample(model, time_limit=time_limit)
+        result = sampler.sample(model, time_limit=time_limit, label=mock_problem_label)
 
         with self.subTest('low-level sample_nlm called'):
             base_sample_problem.assert_called_with(
-                mock_problem_data_id, label=None, upload_params=None, time_limit=time_limit)
+                mock_problem_data_id, label=mock_problem_label,
+                upload_params=None, time_limit=time_limit)
 
         with self.subTest('max_num_states is respected on upload'):
             upload_nlm.assert_called_with(
                 model, max_num_states=sampler.properties['maximum_number_of_states'])
 
-        with self.subTest('timing returned in sample future'):
+        with self.subTest('timing returned in sample result'):
             self.assertIsInstance(result, concurrent.futures.Future)
             self.assertIsInstance(result.result(), LeapHybridNLSampler.SampleResult)
-            self.assertEqual(result.result().timing, mock_timing_info)
+            self.assertEqual(result.result().info['timing'], mock_timing)
+            # warnings separate from timing
+            self.assertEqual(result.result().info['warnings'], mock_warnings)
+
+        with self.subTest('problem info returned in sample result'):
+            self.assertEqual(result.result().info['problem_id'], mock_problem_id)
+            # self.assertEqual(result.result().info['problem_label'], mock_problem_label)
+            self.assertEqual(result.result().info['problem_data_id'], mock_problem_data_id)
 
         with self.subTest('model states updated'):
             self.assertEqual(result.result().model.states.size(), num_states)
             self.assertEqual(model.states.size(), num_states)
 
         with self.subTest('warnings raised'):
-            # add a warning to timing info (inplace)
-            msg = 'solved by classical presolve techniques'
-            mock_timing_info['warnings'] = [msg]
-
-            with self.assertWarns(UserWarning, msg=msg):
+            with self.assertWarns(UserWarning, msg=mock_warnings[0]):
                 result = sampler.sample(model, time_limit=time_limit)
                 result.result()