Repository: dwavesystems/dwave-cloud-client
Pull request #222 (2018-07-27 23:17:35)

Add timeout


Patch:
diff --git a/dwave/cloud/client.py b/dwave/cloud/client.py
index d64c188..4748e6b 100644
--- a/dwave/cloud/client.py
+++ b/dwave/cloud/client.py
@@ -44,7 +44,7 @@ from dwave.cloud.package_info import __packagename__, __version__
 from dwave.cloud.exceptions import *
 from dwave.cloud.config import load_config, legacy_load_config
 from dwave.cloud.solver import Solver
-from dwave.cloud.utils import datetime_to_timestamp
+from dwave.cloud.utils import datetime_to_timestamp, TimeoutingHTTPAdapter
 
 __all__ = ['Client']
 
@@ -383,7 +383,7 @@ class Client(object):
         return _clients[_client](**config)
 
     def __init__(self, endpoint=None, token=None, solver=None, proxy=None,
-                 permissive_ssl=False, **kwargs):
+                 permissive_ssl=False, timeout=60, **kwargs):
         """To setup the connection a pipeline of queues/workers is constructed.
 
         There are five interactions with the server the connection manages:
@@ -405,9 +405,12 @@ class Client(object):
         self.endpoint = endpoint
         self.token = token
         self.default_solver = solver
+        self.timeout = float(timeout)
 
         # Create a :mod:`requests` session. `requests` will manage our url parsing, https, etc.
         self.session = requests.Session()
+        self.session.mount('http://', TimeoutingHTTPAdapter(timeout=self.timeout))
+        self.session.mount('https://', TimeoutingHTTPAdapter(timeout=self.timeout))
         self.session.headers.update({'X-Auth-Token': self.token,
                                      'User-Agent': self.USER_AGENT})
         self.session.proxies = {'http': proxy, 'https': proxy}
@@ -580,8 +583,10 @@ class Client(object):
                 return self._solvers
 
             _LOGGER.debug("Requesting list of all solver data.")
-            response = self.session.get(
-                posixpath.join(self.endpoint, 'solvers/remote/'))
+            try:
+                response = self.session.get(posixpath.join(self.endpoint, 'solvers/remote/'))
+            except requests.exceptions.Timeout:
+                raise ConnectionTimeout
 
             if response.status_code == 401:
                 raise SolverAuthenticationError
@@ -638,7 +643,7 @@ class Client(object):
         Examples:
             Get a solver not based on VFYC, with 2048 qubits:
 
-                solver = client.solvers(vfyc=True, num_qubits=2048)[0]
+                solver = client.solvers(vfyc=False, num_qubits=2048)[0]
 
             Get all solvers that have between 1024 and 2048 qubits:
 
@@ -647,6 +652,10 @@ class Client(object):
             Get all solvers that have at least 2000 qubits:
 
                 solvers = client.solvers(num_qubits=[2000, None])
+
+            Get the first QPU solver that accepts flux biases:
+
+                solver = client.solvers(qpu=True, flux_biases=True)[0]
         """
 
         def predicate(solver):
@@ -727,8 +736,11 @@ class Client(object):
 
         with self._solvers_lock:
             if refresh or name not in self._solvers:
-                response = self.session.get(
-                    posixpath.join(self.endpoint, 'solvers/remote/{}/'.format(name)))
+                try:
+                    response = self.session.get(
+                        posixpath.join(self.endpoint, 'solvers/remote/{}/'.format(name)))
+                except requests.exceptions.Timeout:
+                    raise ConnectionTimeout
 
                 if response.status_code == 401:
                     raise SolverAuthenticationError
@@ -782,7 +794,10 @@ class Client(object):
                 _LOGGER.debug("Submitting %d problems", len(ready_problems))
                 body = '[' + ','.join(mess.body for mess in ready_problems) + ']'
                 try:
-                    response = self.session.post(posixpath.join(self.endpoint, 'problems/'), body)
+                    try:
+                        response = self.session.post(posixpath.join(self.endpoint, 'problems/'), body)
+                    except requests.exceptions.Timeout:
+                        raise ConnectionTimeout
 
                     if response.status_code == 401:
                         raise SolverAuthenticationError()
@@ -917,7 +932,12 @@ class Client(object):
                 # body of the delete query.
                 try:
                     body = [item[0] for item in item_list]
-                    self.session.delete(posixpath.join(self.endpoint, 'problems/'), json=body)
+
+                    try:
+                        self.session.delete(posixpath.join(self.endpoint, 'problems/'), json=body)
+                    except requests.exceptions.Timeout:
+                        raise ConnectionTimeout
+
                 except Exception as err:
                     for _, future in item_list:
                         if future is not None:
@@ -1031,7 +1051,11 @@ class Client(object):
 
                 try:
                     _LOGGER.trace("Executing poll API request")
-                    response = self.session.get(posixpath.join(self.endpoint, query_string))
+
+                    try:
+                        response = self.session.get(posixpath.join(self.endpoint, query_string))
+                    except requests.exceptions.Timeout:
+                        raise ConnectionTimeout
 
                     if response.status_code == 401:
                         raise SolverAuthenticationError()
@@ -1086,7 +1110,10 @@ class Client(object):
                 # Submit the query
                 query_string = 'problems/{}/'.format(future.id)
                 try:
-                    response = self.session.get(posixpath.join(self.endpoint, query_string))
+                    try:
+                        response = self.session.get(posixpath.join(self.endpoint, query_string))
+                    except requests.exceptions.Timeout:
+                        raise ConnectionTimeout
 
                     if response.status_code == 401:
                         raise SolverAuthenticationError()
diff --git a/dwave/cloud/exceptions.py b/dwave/cloud/exceptions.py
index ae49435..5d4bd6f 100644
--- a/dwave/cloud/exceptions.py
+++ b/dwave/cloud/exceptions.py
@@ -24,6 +24,16 @@ class UnsupportedSolverError(SolverError):
     """The solver we received from the API is not supported by the client."""
 
 
+class Timeout(Exception):
+    """General timeout error."""
+
+class ConnectionTimeout(Timeout):
+    """REST API request timed out."""
+
+class PollingTimeout(Timeout):
+    """Problem polling timed out."""
+
+
 class CanceledFutureError(Exception):
     """An exception raised when code tries to read from a canceled future."""
 
diff --git a/dwave/cloud/utils.py b/dwave/cloud/utils.py
index 4769301..7bd8568 100644
--- a/dwave/cloud/utils.py
+++ b/dwave/cloud/utils.py
@@ -8,6 +8,7 @@ import random
 
 import six
 import click
+import requests
 
 # Use numpy if available for fast decoding
 try:
@@ -174,3 +175,31 @@ def datetime_to_timestamp(dt):
 def strtrunc(s, maxlen=60):
     s = str(s)
     return s[:(maxlen-3)]+'...' if len(s) > maxlen else s
+
+
+class TimeoutingHTTPAdapter(requests.adapters.HTTPAdapter):
+    """Sets a default timeout for all adapter (think session) requests. It is
+    overridden with per-request timeout. But it can not be reset back to
+    infinite wait (``None``).
+
+    Usage:
+
+        s = requests.Session()
+        s.mount("http://", TimeoutingHTTPAdapter(timeout=5))
+        s.mount("https://", TimeoutingHTTPAdapter(timeout=5))
+
+        s.get('http://httpbin.org/delay/6')                 # -> timeouts after 5sec
+        s.get('http://httpbin.org/delay/6', timeout=10)     # -> completes after 6sec
+
+    The alternative is to set ``timeout`` on each request manually/explicitly,
+    subclass ``Session``, or monkeypatch ``Session.request()``.
+    """
+
+    def __init__(self, timeout=None, *args, **kwargs):
+        self.timeout = timeout
+        super(TimeoutingHTTPAdapter, self).__init__(*args, **kwargs)
+
+    def send(self, *args, **kwargs):
+        # can't use setdefault because caller always sets timeout kwarg
+        kwargs['timeout'] = self.timeout
+        return super(TimeoutingHTTPAdapter, self).send(*args, **kwargs)


Test patch:
diff --git a/tests/test_client.py b/tests/test_client.py
index 1c91639..f49bfae 100644
--- a/tests/test_client.py
+++ b/tests/test_client.py
@@ -6,6 +6,7 @@ test_mock_solver_loading.py duplicates some of these tests against a mock server
 from __future__ import absolute_import
 
 import unittest
+import requests.exceptions
 
 from dwave.cloud.config import load_config
 from dwave.cloud.client import Client
@@ -43,6 +44,16 @@ class ConnectivityTests(unittest.TestCase):
             self.assertTrue(len(client.get_solvers()) > 0)
 
 
+@unittest.skipUnless(config, "No live server configuration available.")
+class TestTimeout(unittest.TestCase):
+    """Test timeout works for all Client connections."""
+
+    def test_timeout(self):
+        with self.assertRaises(dwave.cloud.exceptions.ConnectionTimeout):
+            with Client(timeout=0.00001, **config) as client:
+                client.solvers()
+
+
 @unittest.skipUnless(config, "No live server configuration available.")
 class SolverLoading(unittest.TestCase):
     """Test loading solvers in a few different configurations."""
diff --git a/tests/test_mock_solver_loading.py b/tests/test_mock_solver_loading.py
index 04af5be..953e78e 100644
--- a/tests/test_mock_solver_loading.py
+++ b/tests/test_mock_solver_loading.py
@@ -4,6 +4,8 @@ from __future__ import division, absolute_import, print_function, unicode_litera
 import os
 import json
 import unittest
+
+import requests
 import requests_mock
 
 from dwave.cloud.qpu import Client, Solver
@@ -228,6 +230,12 @@ token = file-prod-token
 endpoint = file-alpha-url
 token = file-alpha-token
 solver = alpha-solver
+
+[custom]
+endpoint = http://httpbin.org/delay/10
+token = 123
+permissive_ssl = True
+timeout = 15
 """
 
 
@@ -328,3 +336,24 @@ class MockLegacyConfiguration(unittest.TestCase):
         with mock.patch("dwave.cloud.config.open", side_effect=OSError, create=True):
             with mock.patch("os.path.exists", lambda fn: True):
                 self.assertRaises(ConfigFileReadError, legacy_load_config)
+
+    def test_custom_options(self):
+        """Test custom options (timeout, permissive_ssl) are propagated to Client."""
+        timeout = 15
+
+        with mock.patch("dwave.cloud.config.open", iterable_mock_open(config_body), create=True):
+            with Client.from_config('config_file', profile='custom') as client:
+                # check permissive_ssl and timeout custom params passed-thru
+                self.assertFalse(client.session.verify)
+                self.assertEqual(client.timeout, timeout)
+
+                # verify client uses those properly
+                def mock_send(*args, **kwargs):
+                    self.assertEqual(kwargs.get('timeout'), timeout)
+                    response = requests.Response()
+                    response.status_code = 200
+                    response._content = b'{}'
+                    return response
+
+                with mock.patch("requests.adapters.HTTPAdapter.send", mock_send):
+                    client.solvers()