Add problem info to the NL sampler result

dwave/system/samplers/leap_hybrid_sampler.py

@@ -918,7 +918,7 @@ def parameters(self) -> Dict[str, List[str]]:
 
     class SampleResult(NamedTuple):
         model: dwave.optimization.Model
-        timing: dict
+        info: dict
 
     def sample(self, model: dwave.optimization.Model,
                time_limit: Optional[float] = None, **kwargs
@@ -945,10 +945,16 @@ def sample(self, model: dwave.optimization.Model,
 
         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 +977,20 @@ def hook(model, future):
         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)
 

tests/test_leaphybridnlsampler.py

@@ -71,56 +71,64 @@ def test_state_resolve(self, decode_response, upload_nlm, base_sample_problem):
         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
             answer_data.write(states_file.read())
             answer_data.seek(0)
             return {
                 'problem_type': 'nl',
-                'timing': mock_timing_info,
+                'timing': mock_timing | dict(warnings=mock_warnings),
                 'shape': {},
                 'answer': answer_data
             }
         decode_response.side_effect = mock_decode_response
 
         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()