Add HTTP header with MLflow version in Python client (mlflow#1131)

brianhendee · Apr 17, 2019 · 44a501b · 44a501b
1 parent 9199dfc
commit 44a501b
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 17 deletions.
diff --git a/mlflow/utils/rest_utils.py b/mlflow/utils/rest_utils.py
@@ -7,6 +7,7 @@
 import numpy
 import requests
 
+from mlflow import __version__
 from mlflow.utils.string_utils import strip_suffix
 from mlflow.exceptions import MlflowException, RestException
 
@@ -17,6 +18,11 @@
 _logger = logging.getLogger(__name__)
 
 
+_DEFAULT_HEADERS = {
+    'User-Agent': 'mlflow-python-client/%s' % __version__
+}
+
+
 def http_request(host_creds, endpoint, retries=3, retry_interval=3, **kwargs):
     """
     Makes an HTTP request with the specified method to the specified hostname/endpoint. Retries
@@ -36,7 +42,7 @@ def http_request(host_creds, endpoint, retries=3, retry_interval=3, **kwargs):
     elif host_creds.token:
         auth_str = "Bearer %s" % host_creds.token
 
-    headers = {}
+    headers = dict(_DEFAULT_HEADERS)
     if auth_str:
         headers['Authorization'] = auth_str
 

diff --git a/tests/projects/test_databricks.py b/tests/projects/test_databricks.py
@@ -18,6 +18,8 @@
 from mlflow.utils.mlflow_tags import MLFLOW_DATABRICKS_RUN_URL, \
     MLFLOW_DATABRICKS_SHELL_JOB_RUN_ID, \
     MLFLOW_DATABRICKS_WEBAPP_URL
+from mlflow.utils.rest_utils import _DEFAULT_HEADERS
+
 
 from tests.projects.utils import validate_exit_status, TEST_PROJECT_DIR
 from tests.projects.utils import tracking_uri_mock  # pylint: disable=unused-import
@@ -273,12 +275,12 @@ def get_config(self):
 def test_databricks_http_request_integration(get_config, request):
     """Confirms that the databricks http request params can in fact be used as an HTTP request"""
     def confirm_request_params(**kwargs):
+        headers = dict(_DEFAULT_HEADERS)
+        headers['Authorization'] = 'Basic dXNlcjpwYXNz'
         assert kwargs == {
             'method': 'PUT',
             'url': 'host/clusters/list',
-            'headers': {
-                'Authorization': 'Basic dXNlcjpwYXNz'
-            },
+            'headers': headers,
             'verify': True,
             'json': {'a': 'b'}
         }

diff --git a/tests/store/test_rest_store.py b/tests/store/test_rest_store.py
@@ -12,7 +12,7 @@
 from mlflow.store.rest_store import RestStore
 from mlflow.utils.proto_json_utils import message_to_json
 
-from mlflow.utils.rest_utils import MlflowHostCreds
+from mlflow.utils.rest_utils import MlflowHostCreds, _DEFAULT_HEADERS
 
 
 class TestRestStore(unittest.TestCase):
@@ -25,7 +25,7 @@ def mock_request(**kwargs):
                 'method': 'GET',
                 'params': {'view_type': 'ACTIVE_ONLY'},
                 'url': 'https://hello/api/2.0/preview/mlflow/experiments/list',
-                'headers': {},
+                'headers': _DEFAULT_HEADERS,
                 'verify': True,
             }
             response = mock.MagicMock

diff --git a/tests/utils/test_rest_utils.py b/tests/utils/test_rest_utils.py
@@ -5,7 +5,7 @@
 import pytest
 
 from mlflow.utils.rest_utils import NumpyEncoder, http_request, http_request_safe,\
-    MlflowHostCreds
+    MlflowHostCreds, _DEFAULT_HEADERS
 from mlflow.exceptions import MlflowException, RestException
 
 
@@ -19,7 +19,7 @@ def test_http_request_hostonly(request):
     request.assert_called_with(
         url='http://my-host/my/endpoint',
         verify=True,
-        headers={},
+        headers=_DEFAULT_HEADERS,
     )
 
 
@@ -34,7 +34,7 @@ def test_http_request_cleans_hostname(request):
     request.assert_called_with(
         url='http://my-host/my/endpoint',
         verify=True,
-        headers={},
+        headers=_DEFAULT_HEADERS,
     )
 
 
@@ -45,12 +45,12 @@ def test_http_request_with_basic_auth(request):
     response.status_code = 200
     request.return_value = response
     http_request(host_only, '/my/endpoint')
+    headers = dict(_DEFAULT_HEADERS)
+    headers['Authorization'] = 'Basic dXNlcjpwYXNz'
     request.assert_called_with(
         url='http://my-host/my/endpoint',
         verify=True,
-        headers={
-            'Authorization': 'Basic dXNlcjpwYXNz'
-        },
+        headers=headers,
     )
 
 
@@ -61,12 +61,12 @@ def test_http_request_with_token(request):
     response.status_code = 200
     request.return_value = response
     http_request(host_only, '/my/endpoint')
+    headers = dict(_DEFAULT_HEADERS)
+    headers['Authorization'] = 'Bearer my-token'
     request.assert_called_with(
         url='http://my-host/my/endpoint',
         verify=True,
-        headers={
-            'Authorization': 'Bearer my-token'
-        },
+        headers=headers,
     )
 
 
@@ -80,7 +80,7 @@ def test_http_request_with_insecure(request):
     request.assert_called_with(
         url='http://my-host/my/endpoint',
         verify=False,
-        headers={},
+        headers=_DEFAULT_HEADERS,
     )
 
 
@@ -94,7 +94,7 @@ def test_http_request_wrapper(request):
     request.assert_called_with(
         url='http://my-host/my/endpoint',
         verify=False,
-        headers={},
+        headers=_DEFAULT_HEADERS,
     )
     response.status_code = 400
     response.text = ""