def test_resolve_request_headers_no_arg(mock_request_header_providers):
assert resolve_request_headers() == {
"one": "override",
"two": "two-val",
"three": "three-val",
"new": "new-val",
}
def test_resolve_request_headers(mock_request_header_providers):
request_headers_arg = {"two": "arg-override", "arg": "arg-val"}
assert resolve_request_headers(request_headers_arg) == {
"one": "override",
"two": "arg-override",
"three": "three-val",
"new": "new-val",
"arg": "arg-val",
}
def test_run_context_provider_registry_with_installed_plugin():
"""This test requires the package in tests/resources/mlflow-test-plugin to be installed"""
reload(mlflow.tracking.request_header.registry)
from mlflow_test_plugin.request_header_provider import PluginRequestHeaderProvider
assert PluginRequestHeaderProvider in _currently_registered_request_header_provider_classes(
)
# The test plugin's request header provider always returns False from in_context to avoid
# polluting request headers in developers' environments. The following mock overrides this to
# perform the integration test.
with mock.patch.object(PluginRequestHeaderProvider,
"in_context",
return_value=True):
assert resolve_request_headers()["test"] == "header"
def http_request(
host_creds,
endpoint,
method,
max_retries=5,
backoff_factor=2,
retry_codes=_TRANSIENT_FAILURE_RESPONSE_CODES,
timeout=120,
**kwargs,
):
"""
Makes an HTTP request with the specified method to the specified hostname/endpoint. Transient
errors such as Rate-limited (429), service unavailable (503) and internal error (500) are
retried with an exponential back off with backoff_factor * (1, 2, 4, ... seconds).
The function parses the API response (assumed to be JSON) into a Python object and returns it.
:param host_creds: A :py:class:`mlflow.rest_utils.MlflowHostCreds` object containing
hostname and optional authentication.
:param endpoint: a string for service endpoint, e.g. "/path/to/object".
:param method: a string indicating the method to use, e.g. "GET", "POST", "PUT".
:param max_retries: maximum number of retries before throwing an exception.
:param backoff_factor: a time factor for exponential backoff. e.g. value 5 means the HTTP
request will be retried with interval 5, 10, 20... seconds. A value of 0 turns off the
exponential backoff.
:param retry_codes: a list of HTTP response error codes that qualifies for retry.
:param timeout: wait for timeout seconds for response from remote server for connect and
read request.
:param kwargs: Additional keyword arguments to pass to `requests.Session.request()`
:return: requests.Response object.
"""
hostname = host_creds.host
auth_str = None
if host_creds.username and host_creds.password:
basic_auth_str = (
"%s:%s" %
(host_creds.username, host_creds.password)).encode("utf-8")
auth_str = "Basic " + base64.standard_b64encode(basic_auth_str).decode(
"utf-8")
elif host_creds.token:
auth_str = "Bearer %s" % host_creds.token
from mlflow.tracking.request_header.registry import resolve_request_headers
headers = dict({**_DEFAULT_HEADERS, **resolve_request_headers()})
if auth_str:
headers["Authorization"] = auth_str
if host_creds.server_cert_path is None:
verify = not host_creds.ignore_tls_verification
else:
verify = host_creds.server_cert_path
if host_creds.client_cert_path is not None:
kwargs["cert"] = host_creds.client_cert_path
cleaned_hostname = strip_suffix(hostname, "/")
url = "%s%s" % (cleaned_hostname, endpoint)
try:
return _get_http_response_with_retries(
method,
url,
max_retries,
backoff_factor,
retry_codes,
headers=headers,
verify=verify,
timeout=timeout,
**kwargs,
)
except Exception as e:
raise MlflowException("API request to %s failed with exception %s" %
(url, e))
def http_request(host_creds,
endpoint,
retries=3,
retry_interval=3,
max_rate_limit_interval=60,
**kwargs):
"""
Makes an HTTP request with the specified method to the specified hostname/endpoint. Ratelimit
error code (429) will be retried with an exponential back off (1, 2, 4, ... seconds) for at most
`max_rate_limit_interval` seconds. Internal errors (500s) will be retried up to `retries` times
, waiting `retry_interval` seconds between successive retries. Parses the API response
(assumed to be JSON) into a Python object and returns it.
:param host_creds: A :py:class:`mlflow.rest_utils.MlflowHostCreds` object containing
hostname and optional authentication.
:return: Parsed API response
"""
hostname = host_creds.host
auth_str = None
if host_creds.username and host_creds.password:
basic_auth_str = (
"%s:%s" %
(host_creds.username, host_creds.password)).encode("utf-8")
auth_str = "Basic " + base64.standard_b64encode(basic_auth_str).decode(
"utf-8")
elif host_creds.token:
auth_str = "Bearer %s" % host_creds.token
from mlflow.tracking.request_header.registry import resolve_request_headers
headers = dict({**_DEFAULT_HEADERS, **resolve_request_headers()})
if auth_str:
headers["Authorization"] = auth_str
if host_creds.server_cert_path is None:
verify = not host_creds.ignore_tls_verification
else:
verify = host_creds.server_cert_path
if host_creds.client_cert_path is not None:
kwargs["cert"] = host_creds.client_cert_path
def request_with_ratelimit_retries(max_rate_limit_interval, **kwargs):
response = requests.request(**kwargs)
time_left = max_rate_limit_interval
sleep = 1
while response.status_code == 429 and time_left > 0:
_logger.warning(
"API request to {path} returned status code 429 (Rate limit exceeded). "
"Retrying in %d seconds. "
"Will continue to retry 429s for up to %d seconds.",
sleep,
time_left,
)
time.sleep(sleep)
time_left -= sleep
response = requests.request(**kwargs)
sleep = min(time_left,
sleep * 2) # sleep for 1, 2, 4, ... seconds;
return response
cleaned_hostname = strip_suffix(hostname, "/")
url = "%s%s" % (cleaned_hostname, endpoint)
for i in range(retries):
response = request_with_ratelimit_retries(max_rate_limit_interval,
url=url,
headers=headers,
verify=verify,
**kwargs)
if response.status_code >= 200 and response.status_code < 500:
return response
else:
_logger.error(
"API request to %s failed with code %s != 200, retrying up to %s more times. "
"API response body: %s",
url,
response.status_code,
retries - i - 1,
response.text,
)
time.sleep(retry_interval)
raise MlflowException(
"API request to %s failed to return code 200 after %s tries" %
(url, retries))
