Use websocket stream to retrieve runtime job results

qiskit_ibm/runtime/runtime_job.py

@@ -13,14 +13,12 @@
 """Qiskit runtime job."""
 
 from typing import Any, Optional, Callable, Dict, Type
-import time
 import logging
 from concurrent import futures
 import traceback
 import queue
 from datetime import datetime
 
-from qiskit.providers.exceptions import JobTimeoutError
 from qiskit.providers.backend import Backend
 from qiskit.providers.jobstatus import JobStatus, JOB_FINAL_STATES
 
@@ -128,14 +126,12 @@ def __init__(
     def result(
             self,
             timeout: Optional[float] = None,
-            wait: float = 5,
             decoder: Optional[Type[ResultDecoder]] = None
     ) -> Any:
         """Return the results of the job.
 
         Args:
             timeout: Number of seconds to wait for job.
-            wait: Seconds between queries.
             decoder: A :class:`ResultDecoder` subclass used to decode job results.
 
         Returns:
@@ -146,7 +142,7 @@ def result(
         """
         _decoder = decoder or self._result_decoder
         if self._results is None or (_decoder != self._result_decoder):
-            self.wait_for_final_state(timeout=timeout, wait=wait)
+            self.wait_for_final_state(timeout=timeout)
             if self._status == JobStatus.ERROR:
                 raise RuntimeJobFailureError(f"Unable to retrieve job result. "
                                              f"{self.error_message()}")
@@ -188,30 +184,20 @@ def error_message(self) -> Optional[str]:
         self._set_status_and_error_message()
         return self._error_message
 
-    def wait_for_final_state(
-            self,
-            timeout: Optional[float] = None,
-            wait: float = 5
-    ) -> None:
-        """Poll the job status until it progresses to a final state such as ``DONE`` or ``ERROR``.
+    def wait_for_final_state(self, timeout: Optional[float] = None) -> None:
+        """If the websocket server stream endpoint is open, wait for it to close.
 
         Args:
             timeout: Seconds to wait for the job. If ``None``, wait indefinitely.
-            wait: Seconds between queries.
 
         Raises:
             JobTimeoutError: If the job does not reach a final state before the
                 specified timeout.
         """
-        start_time = time.time()
-        status = self.status()
-        while status not in JOB_FINAL_STATES:
-            elapsed_time = time.time() - start_time
-            if timeout is not None and elapsed_time >= timeout:
-                raise JobTimeoutError(
-                    'Timeout while waiting for job {}.'.format(self.job_id))
-            time.sleep(wait)
-            status = self.status()
+        if self._status not in JOB_FINAL_STATES:
+            self._ws_client_future = self._executor.submit(self._start_websocket_client)
+        self._ws_client_future.result(timeout)
+        self.status()
 
     def stream_results(
             self,
@@ -233,16 +219,18 @@ def stream_results(
             RuntimeInvalidStateError: If a callback function is already streaming results or
                 if the job already finished.
         """
-        if self._is_streaming():
-            raise RuntimeInvalidStateError("A callback function is already streaming results.")
-
         if self._status in JOB_FINAL_STATES:
             raise RuntimeInvalidStateError("Job already finished.")
-
-        self._ws_client_future = self._executor.submit(self._start_websocket_client)
-        self._executor.submit(self._stream_results,
-                              result_queue=self._result_queue, user_callback=callback,
-                              decoder=decoder)
+        if self._is_streaming():
+            try:
+                self._executor.submit(self._stream_results, result_queue=self._result_queue,
+                                      user_callback=callback, decoder=decoder)
+            except Exception:
+                raise RuntimeInvalidStateError("A callback function is already streaming results.")
+        else:
+            self._ws_client_future = self._executor.submit(self._start_websocket_client)
+            self._executor.submit(self._stream_results, result_queue=self._result_queue,
+                                  user_callback=callback, decoder=decoder)
 
     def cancel_result_streaming(self) -> None:
         """Cancel result streaming."""

test/ibm/runtime/test_runtime.py

@@ -424,9 +424,11 @@ def test_run_program(self):
         self.assertIsInstance(job, RuntimeJob)
         self.assertIsInstance(job.status(), JobStatus)
         self.assertEqual(job.inputs, params)
-        job.wait_for_final_state()
+        with mock.patch.object(RuntimeJob, 'wait_for_final_state',
+                               side_effect=time.sleep(3)):
+            job.wait_for_final_state()
+            self.assertTrue(job.result())
         self.assertEqual(job.status(), JobStatus.DONE)
-        self.assertTrue(job.result())
 
     def test_run_program_with_custom_runtime_image(self):
         """Test running program."""
@@ -437,9 +439,11 @@ def test_run_program_with_custom_runtime_image(self):
         self.assertIsInstance(job, RuntimeJob)
         self.assertIsInstance(job.status(), JobStatus)
         self.assertEqual(job.inputs, params)
-        job.wait_for_final_state()
+        with mock.patch.object(RuntimeJob, 'wait_for_final_state',
+                               side_effect=time.sleep(3)):
+            job.wait_for_final_state()
+            self.assertTrue(job.result())
         self.assertEqual(job.status(), JobStatus.DONE)
-        self.assertTrue(job.result())
         self.assertEqual(job.image, image)
 
     def test_retrieve_program_data(self):
@@ -479,24 +483,32 @@ def test_program_params_namespace(self):
     def test_run_program_failed(self):
         """Test a failed program execution."""
         job = self._run_program(job_classes=FailedRuntimeJob)
-        job.wait_for_final_state()
+        with mock.patch.object(RuntimeJob, 'wait_for_final_state',
+                               side_effect=time.sleep(5)):
+            job.wait_for_final_state()
         job_result_raw = self.runtime._api_client.job_results(job.job_id)
         self.assertEqual(JobStatus.ERROR, job.status())
         self.assertEqual(API_TO_JOB_ERROR_MESSAGE['FAILED'].format(
             job.job_id, job_result_raw), job.error_message())
         with self.assertRaises(RuntimeJobFailureError):
-            job.result()
+            with mock.patch.object(RuntimeJob, 'wait_for_final_state',
+                                   side_effect=time.sleep(5)):
+                job.result()
 
     def test_run_program_failed_ran_too_long(self):
         """Test a program that failed since it ran longer than maxiumum execution time."""
         job = self._run_program(job_classes=FailedRanTooLongRuntimeJob)
-        job.wait_for_final_state()
+        with mock.patch.object(RuntimeJob, 'wait_for_final_state',
+                               side_effect=time.sleep(3)):
+            job.wait_for_final_state()
         job_result_raw = self.runtime._api_client.job_results(job.job_id)
         self.assertEqual(JobStatus.ERROR, job.status())
         self.assertEqual(API_TO_JOB_ERROR_MESSAGE['CANCELLED - RAN TOO LONG'].format(
             job.job_id, job_result_raw), job.error_message())
         with self.assertRaises(RuntimeJobFailureError):
-            job.result()
+            with mock.patch.object(RuntimeJob, 'wait_for_final_state',
+                                   side_effect=time.sleep(3)):
+                job.result()
 
     def test_retrieve_job(self):
         """Test retrieving a job."""
@@ -630,8 +642,10 @@ def test_jobs_filter_by_program_id(self):
         program_id_1 = self._upload_program()
         job = self._run_program(program_id=program_id)
         job_1 = self._run_program(program_id=program_id_1)
-        job.wait_for_final_state()
-        job_1.wait_for_final_state()
+        with mock.patch.object(RuntimeJob, 'wait_for_final_state',
+                               side_effect=time.sleep(1)):
+            job.wait_for_final_state()
+            job_1.wait_for_final_state()
         rjobs = self.runtime.jobs(program_id=program_id)
         self.assertEqual(program_id, rjobs[0].program_id)
         self.assertEqual(1, len(rjobs))
@@ -641,7 +655,9 @@ def test_jobs_filter_by_provider(self):
         program_id = self._upload_program()
         job = self._run_program(program_id=program_id,
                                 hub="defaultHub", group="defaultGroup", project="defaultProject")
-        job.wait_for_final_state()
+        with mock.patch.object(RuntimeJob, 'wait_for_final_state',
+                               side_effect=time.sleep(1)):
+            job.wait_for_final_state()
         rjobs = self.runtime.jobs(program_id=program_id,
                                   hub="defaultHub", group="defaultGroup", project="defaultProject")
         self.assertEqual(program_id, rjobs[0].program_id)
@@ -664,7 +680,9 @@ def test_cancel_job(self):
     def test_final_result(self):
         """Test getting final result."""
         job = self._run_program()
-        result = job.result()
+        with mock.patch.object(RuntimeJob, 'wait_for_final_state',
+                               side_effect=time.sleep(5)):
+            result = job.result()
         self.assertTrue(result)
 
     def test_job_status(self):
@@ -688,7 +706,9 @@ def test_job_program_id(self):
     def test_wait_for_final_state(self):
         """Test wait for final state."""
         job = self._run_program()
-        job.wait_for_final_state()
+        with mock.patch.object(RuntimeJob, 'wait_for_final_state',
+                               side_effect=time.sleep(3)):
+            job.wait_for_final_state()
         self.assertEqual(JobStatus.DONE, job.status())
 
     def test_result_decoder(self):
@@ -701,7 +721,9 @@ def test_result_decoder(self):
         for result_decoder, decoder in sub_tests:
             with self.subTest(decoder=decoder):
                 job = self._run_program(job_classes=job_cls, decoder=result_decoder)
-                result = job.result(decoder=decoder)
+                with mock.patch.object(RuntimeJob, 'wait_for_final_state',
+                                       side_effect=time.sleep(10)):
+                    result = job.result(decoder=decoder)
                 self.assertIsInstance(result['serializable_class'], SerializableClass)
 
     def test_get_result_twice(self):
@@ -711,8 +733,10 @@ def test_get_result_twice(self):
         job_cls.custom_result = custom_result
 
         job = self._run_program(job_classes=job_cls)
-        _ = job.result()
-        _ = job.result()
+        with mock.patch.object(RuntimeJob, 'wait_for_final_state',
+                               side_effect=time.sleep(5)):
+            _ = job.result()
+            _ = job.result()
 
     def test_program_metadata(self):
         """Test program metadata."""