wehub-resource-sync
1043 lines
40 KiB
Python
1043 lines
40 KiB
Python
import re
|
|
import time
|
|
import warnings
|
|
from unittest import mock
|
|
|
|
import numpy
|
|
import pytest
|
|
import requests
|
|
|
|
from mlflow.deployments.databricks import DatabricksDeploymentClient
|
|
from mlflow.environment_variables import (
|
|
_MLFLOW_DATABRICKS_TRAFFIC_ID,
|
|
MLFLOW_HTTP_REQUEST_TIMEOUT,
|
|
)
|
|
from mlflow.exceptions import InvalidUrlException, MlflowException, RestException
|
|
from mlflow.protos.databricks_pb2 import ENDPOINT_NOT_FOUND, ErrorCode
|
|
from mlflow.protos.service_pb2 import GetRun
|
|
from mlflow.pyfunc.scoring_server import NumpyEncoder
|
|
from mlflow.tracking.request_header.default_request_header_provider import (
|
|
_CLIENT_VERSION,
|
|
_USER_AGENT,
|
|
DefaultRequestHeaderProvider,
|
|
)
|
|
from mlflow.utils.rest_utils import (
|
|
_DATABRICKS_SDK_RETRY_AFTER_SECS_DEPRECATION_WARNING,
|
|
MlflowHostCreds,
|
|
_can_parse_as_json_object,
|
|
augmented_raise_for_status,
|
|
call_endpoint,
|
|
call_endpoints,
|
|
get_workspace_client,
|
|
http_request,
|
|
http_request_safe,
|
|
)
|
|
from mlflow.utils.workspace_context import WorkspaceContext
|
|
from mlflow.utils.workspace_utils import WORKSPACE_HEADER_NAME
|
|
|
|
from tests import helper_functions
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"response_mock",
|
|
[
|
|
helper_functions.create_mock_response(400, "Error message but not a JSON string"),
|
|
helper_functions.create_mock_response(400, ""),
|
|
helper_functions.create_mock_response(400, None),
|
|
],
|
|
)
|
|
def test_malformed_json_error_response(response_mock):
|
|
with mock.patch("requests.Session.request", return_value=response_mock):
|
|
host_only = MlflowHostCreds("http://my-host")
|
|
|
|
response_proto = GetRun.Response()
|
|
with pytest.raises(
|
|
MlflowException, match="API request to endpoint /my/endpoint failed with error code 400"
|
|
):
|
|
call_endpoint(host_only, "/my/endpoint", "GET", None, response_proto)
|
|
|
|
|
|
def test_call_endpoints():
|
|
with mock.patch("mlflow.utils.rest_utils.call_endpoint") as mock_call_endpoint:
|
|
response_proto = GetRun.Response()
|
|
mock_call_endpoint.side_effect = [
|
|
RestException({"error_code": ErrorCode.Name(ENDPOINT_NOT_FOUND)}),
|
|
None,
|
|
]
|
|
host_only = MlflowHostCreds("http://my-host")
|
|
endpoints = [("/my/endpoint", "POST"), ("/my/endpoint", "GET")]
|
|
resp = call_endpoints(host_only, endpoints, "", response_proto)
|
|
mock_call_endpoint.assert_has_calls([
|
|
mock.call(host_only, endpoint, method, "", response_proto, None)
|
|
for endpoint, method in endpoints
|
|
])
|
|
assert resp is None
|
|
|
|
|
|
def test_call_endpoints_raises_exceptions():
|
|
with mock.patch("mlflow.utils.rest_utils.call_endpoint") as mock_call_endpoint:
|
|
response_proto = GetRun.Response()
|
|
mock_call_endpoint.side_effect = [
|
|
RestException({"error_code": ErrorCode.Name(ENDPOINT_NOT_FOUND)}),
|
|
RestException({"error_code": ErrorCode.Name(ENDPOINT_NOT_FOUND)}),
|
|
]
|
|
host_only = MlflowHostCreds("http://my-host")
|
|
endpoints = [("/my/endpoint", "POST"), ("/my/endpoint", "GET")]
|
|
with pytest.raises(RestException, match="ENDPOINT_NOT_FOUND"):
|
|
call_endpoints(host_only, endpoints, "", response_proto)
|
|
mock_call_endpoint.side_effect = [RestException({}), None]
|
|
with pytest.raises(RestException, match="INTERNAL_ERROR"):
|
|
call_endpoints(host_only, endpoints, "", response_proto)
|
|
|
|
|
|
def test_http_request_hostonly():
|
|
host_only = MlflowHostCreds("http://my-host")
|
|
response = mock.MagicMock()
|
|
response.status_code = 200
|
|
with mock.patch("requests.Session.request", return_value=response) as mock_request:
|
|
http_request(host_only, "/my/endpoint", "GET")
|
|
mock_request.assert_called_with(
|
|
"GET",
|
|
"http://my-host/my/endpoint",
|
|
allow_redirects=True,
|
|
verify=True,
|
|
headers=DefaultRequestHeaderProvider().request_headers(),
|
|
timeout=120,
|
|
)
|
|
|
|
|
|
def test_http_request_includes_workspace_header_for_mlflow_routes():
|
|
host_only = MlflowHostCreds("http://my-host")
|
|
response = mock.MagicMock()
|
|
response.status_code = 200
|
|
with WorkspaceContext("team-a"):
|
|
with mock.patch("requests.Session.request", return_value=response) as mock_request:
|
|
http_request(host_only, "/api/2.0/mlflow/runs/search", "GET")
|
|
headers = mock_request.call_args.kwargs["headers"]
|
|
assert headers[WORKSPACE_HEADER_NAME] == "team-a"
|
|
|
|
|
|
def test_http_request_omits_workspace_header_for_workspace_admin_routes():
|
|
host_only = MlflowHostCreds("http://my-host")
|
|
response = mock.MagicMock()
|
|
response.status_code = 200
|
|
with WorkspaceContext("team-a"):
|
|
with mock.patch("requests.Session.request", return_value=response) as mock_request:
|
|
http_request(host_only, "/api/3.0/mlflow/workspaces/team-a", "GET")
|
|
headers = mock_request.call_args.kwargs["headers"]
|
|
assert WORKSPACE_HEADER_NAME not in headers
|
|
|
|
|
|
def test_http_request_cleans_hostname():
|
|
# Add a trailing slash, should be removed.
|
|
host_only = MlflowHostCreds("http://my-host/")
|
|
response = mock.MagicMock()
|
|
response.status_code = 200
|
|
with mock.patch("requests.Session.request", return_value=response) as mock_request:
|
|
http_request(host_only, "/my/endpoint", "GET")
|
|
mock_request.assert_called_with(
|
|
"GET",
|
|
"http://my-host/my/endpoint",
|
|
allow_redirects=True,
|
|
verify=True,
|
|
headers=DefaultRequestHeaderProvider().request_headers(),
|
|
timeout=120,
|
|
)
|
|
|
|
|
|
def test_http_request_with_basic_auth():
|
|
host_only = MlflowHostCreds("http://my-host", username="user", password="pass")
|
|
response = mock.MagicMock()
|
|
response.status_code = 200
|
|
with mock.patch("requests.Session.request", return_value=response) as mock_request:
|
|
http_request(host_only, "/my/endpoint", "GET")
|
|
headers = DefaultRequestHeaderProvider().request_headers()
|
|
headers["Authorization"] = "Basic dXNlcjpwYXNz"
|
|
mock_request.assert_called_with(
|
|
"GET",
|
|
"http://my-host/my/endpoint",
|
|
allow_redirects=True,
|
|
verify=True,
|
|
headers=headers,
|
|
timeout=120,
|
|
)
|
|
|
|
|
|
def test_http_request_with_aws_sigv4(monkeypatch):
|
|
from requests_auth_aws_sigv4 import AWSSigV4
|
|
|
|
monkeypatch.setenv("AWS_ACCESS_KEY_ID", "access-key")
|
|
monkeypatch.setenv("AWS_SECRET_ACCESS_KEY", "secret-key")
|
|
monkeypatch.setenv("AWS_DEFAULT_REGION", "eu-west-1")
|
|
aws_sigv4 = MlflowHostCreds("http://my-host", aws_sigv4=True)
|
|
response = mock.MagicMock()
|
|
response.status_code = 200
|
|
|
|
class AuthMatcher:
|
|
def __eq__(self, other):
|
|
return isinstance(other, AWSSigV4)
|
|
|
|
with mock.patch("requests.Session.request", return_value=response) as mock_request:
|
|
http_request(aws_sigv4, "/my/endpoint", "GET")
|
|
mock_request.assert_called_once_with(
|
|
"GET",
|
|
"http://my-host/my/endpoint",
|
|
allow_redirects=True,
|
|
verify=mock.ANY,
|
|
headers=mock.ANY,
|
|
timeout=mock.ANY,
|
|
auth=AuthMatcher(),
|
|
)
|
|
|
|
|
|
def test_http_request_with_auth():
|
|
mock_fetch_auth = {"test_name": "test_auth_value"}
|
|
auth = "test_auth_name"
|
|
host_only = MlflowHostCreds("http://my-host", auth=auth)
|
|
response = mock.MagicMock()
|
|
response.status_code = 200
|
|
with (
|
|
mock.patch("requests.Session.request", return_value=response) as mock_request,
|
|
mock.patch(
|
|
"mlflow.tracking.request_auth.registry.fetch_auth", return_value=mock_fetch_auth
|
|
) as mock_fetch_auth_call,
|
|
):
|
|
http_request(host_only, "/my/endpoint", "GET")
|
|
|
|
mock_fetch_auth_call.assert_called_with(auth)
|
|
|
|
mock_request.assert_called_with(
|
|
"GET",
|
|
"http://my-host/my/endpoint",
|
|
allow_redirects=True,
|
|
verify=mock.ANY,
|
|
headers=mock.ANY,
|
|
timeout=mock.ANY,
|
|
auth=mock_fetch_auth,
|
|
)
|
|
|
|
|
|
def test_http_request_with_token():
|
|
host_only = MlflowHostCreds("http://my-host", token="my-token")
|
|
response = mock.MagicMock()
|
|
response.status_code = 200
|
|
with mock.patch("requests.Session.request", return_value=response) as mock_request:
|
|
http_request(host_only, "/my/endpoint", "GET")
|
|
headers = DefaultRequestHeaderProvider().request_headers()
|
|
headers["Authorization"] = "Bearer my-token"
|
|
mock_request.assert_called_with(
|
|
"GET",
|
|
"http://my-host/my/endpoint",
|
|
allow_redirects=True,
|
|
verify=True,
|
|
headers=headers,
|
|
timeout=120,
|
|
)
|
|
|
|
|
|
def test_http_request_with_insecure():
|
|
host_only = MlflowHostCreds("http://my-host", ignore_tls_verification=True)
|
|
response = mock.MagicMock()
|
|
response.status_code = 200
|
|
with mock.patch("requests.Session.request", return_value=response) as mock_request:
|
|
http_request(host_only, "/my/endpoint", "GET")
|
|
mock_request.assert_called_with(
|
|
"GET",
|
|
"http://my-host/my/endpoint",
|
|
allow_redirects=True,
|
|
verify=False,
|
|
headers=DefaultRequestHeaderProvider().request_headers(),
|
|
timeout=120,
|
|
)
|
|
|
|
|
|
def test_http_request_client_cert_path():
|
|
host_only = MlflowHostCreds("http://my-host", client_cert_path="/some/path")
|
|
response = mock.MagicMock()
|
|
response.status_code = 200
|
|
with mock.patch("requests.Session.request", return_value=response) as mock_request:
|
|
http_request(host_only, "/my/endpoint", "GET")
|
|
mock_request.assert_called_with(
|
|
"GET",
|
|
"http://my-host/my/endpoint",
|
|
allow_redirects=True,
|
|
verify=True,
|
|
cert="/some/path",
|
|
headers=DefaultRequestHeaderProvider().request_headers(),
|
|
timeout=120,
|
|
)
|
|
|
|
|
|
def test_http_request_server_cert_path():
|
|
host_only = MlflowHostCreds("http://my-host", server_cert_path="/some/path")
|
|
response = mock.MagicMock()
|
|
response.status_code = 200
|
|
with mock.patch("requests.Session.request", return_value=response) as mock_request:
|
|
http_request(host_only, "/my/endpoint", "GET")
|
|
mock_request.assert_called_with(
|
|
"GET",
|
|
"http://my-host/my/endpoint",
|
|
allow_redirects=True,
|
|
verify="/some/path",
|
|
headers=DefaultRequestHeaderProvider().request_headers(),
|
|
timeout=120,
|
|
)
|
|
|
|
|
|
def test_http_request_with_content_type_header():
|
|
host_only = MlflowHostCreds("http://my-host", token="my-token")
|
|
response = mock.MagicMock()
|
|
response.status_code = 200
|
|
extra_headers = {"Content-Type": "text/plain"}
|
|
with mock.patch("requests.Session.request", return_value=response) as mock_request:
|
|
http_request(host_only, "/my/endpoint", "GET", extra_headers=extra_headers)
|
|
headers = DefaultRequestHeaderProvider().request_headers()
|
|
headers["Authorization"] = "Bearer my-token"
|
|
headers["Content-Type"] = "text/plain"
|
|
mock_request.assert_called_with(
|
|
"GET",
|
|
"http://my-host/my/endpoint",
|
|
allow_redirects=True,
|
|
verify=True,
|
|
headers=headers,
|
|
timeout=120,
|
|
)
|
|
|
|
|
|
def test_http_request_request_headers():
|
|
from mlflow_test_plugin.request_header_provider import PluginRequestHeaderProvider
|
|
|
|
# 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.
|
|
response = mock.MagicMock()
|
|
response.status_code = 200
|
|
with (
|
|
mock.patch("requests.Session.request", return_value=response) as mock_request,
|
|
mock.patch.object(PluginRequestHeaderProvider, "in_context", return_value=True),
|
|
):
|
|
host_only = MlflowHostCreds("http://my-host", server_cert_path="/some/path")
|
|
http_request(host_only, "/my/endpoint", "GET")
|
|
mock_request.assert_called_with(
|
|
"GET",
|
|
"http://my-host/my/endpoint",
|
|
allow_redirects=True,
|
|
verify="/some/path",
|
|
headers={**DefaultRequestHeaderProvider().request_headers(), "test": "header"},
|
|
timeout=120,
|
|
)
|
|
|
|
|
|
def test_http_request_request_headers_default():
|
|
from mlflow_test_plugin.request_header_provider import PluginRequestHeaderProvider
|
|
|
|
# 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.
|
|
host_only = MlflowHostCreds("http://my-host", server_cert_path="/some/path")
|
|
default_headers = DefaultRequestHeaderProvider().request_headers()
|
|
expected_headers = {
|
|
_USER_AGENT: "{} {}".format(default_headers[_USER_AGENT], "test_user_agent"),
|
|
_CLIENT_VERSION: "{} {}".format(default_headers[_CLIENT_VERSION], "test_client_version"),
|
|
}
|
|
|
|
response = mock.MagicMock()
|
|
response.status_code = 200
|
|
with (
|
|
mock.patch("requests.Session.request", return_value=response) as mock_request,
|
|
mock.patch.object(PluginRequestHeaderProvider, "in_context", return_value=True),
|
|
mock.patch.object(
|
|
PluginRequestHeaderProvider,
|
|
"request_headers",
|
|
return_value={_USER_AGENT: "test_user_agent", _CLIENT_VERSION: "test_client_version"},
|
|
),
|
|
):
|
|
http_request(host_only, "/my/endpoint", "GET")
|
|
mock_request.assert_called_with(
|
|
"GET",
|
|
"http://my-host/my/endpoint",
|
|
allow_redirects=True,
|
|
verify="/some/path",
|
|
headers=expected_headers,
|
|
timeout=120,
|
|
)
|
|
|
|
|
|
def test_http_request_request_headers_default_and_extra_header():
|
|
from mlflow_test_plugin.request_header_provider import PluginRequestHeaderProvider
|
|
|
|
# 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.
|
|
host_only = MlflowHostCreds("http://my-host", server_cert_path="/some/path")
|
|
default_headers = DefaultRequestHeaderProvider().request_headers()
|
|
expected_headers = {
|
|
_USER_AGENT: "{} {}".format(default_headers[_USER_AGENT], "test_user_agent"),
|
|
_CLIENT_VERSION: "{} {}".format(default_headers[_CLIENT_VERSION], "test_client_version"),
|
|
"header": "value",
|
|
}
|
|
|
|
response = mock.MagicMock()
|
|
response.status_code = 200
|
|
with (
|
|
mock.patch("requests.Session.request", return_value=response) as mock_request,
|
|
mock.patch.object(PluginRequestHeaderProvider, "in_context", return_value=True),
|
|
mock.patch.object(
|
|
PluginRequestHeaderProvider,
|
|
"request_headers",
|
|
return_value={
|
|
_USER_AGENT: "test_user_agent",
|
|
_CLIENT_VERSION: "test_client_version",
|
|
"header": "value",
|
|
},
|
|
),
|
|
):
|
|
http_request(host_only, "/my/endpoint", "GET")
|
|
mock_request.assert_called_with(
|
|
"GET",
|
|
"http://my-host/my/endpoint",
|
|
allow_redirects=True,
|
|
verify="/some/path",
|
|
headers=expected_headers,
|
|
timeout=120,
|
|
)
|
|
|
|
|
|
def test_http_request_with_invalid_url_raise_invalid_url_exception():
|
|
host_only = MlflowHostCreds("http://my-host")
|
|
|
|
with pytest.raises(InvalidUrlException, match="Invalid url: http://my-host/invalid_url"):
|
|
with mock.patch("requests.Session.request", side_effect=requests.exceptions.InvalidURL):
|
|
http_request(host_only, "/invalid_url", "GET")
|
|
|
|
|
|
def test_http_request_with_invalid_url_raise_mlflow_exception():
|
|
host_only = MlflowHostCreds("http://my-host")
|
|
|
|
with pytest.raises(MlflowException, match="Invalid url: http://my-host/invalid_url"):
|
|
with mock.patch("requests.Session.request", side_effect=requests.exceptions.InvalidURL):
|
|
http_request(host_only, "/invalid_url", "GET")
|
|
|
|
|
|
def test_ignore_tls_verification_not_server_cert_path():
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match="When 'ignore_tls_verification' is true then 'server_cert_path' must not be set",
|
|
):
|
|
MlflowHostCreds(
|
|
"http://my-host",
|
|
ignore_tls_verification=True,
|
|
server_cert_path="/some/path",
|
|
)
|
|
|
|
|
|
def test_http_request_wrapper():
|
|
host_only = MlflowHostCreds("http://my-host", ignore_tls_verification=True)
|
|
response = mock.MagicMock()
|
|
response.status_code = 200
|
|
response.text = "{}"
|
|
with mock.patch("requests.Session.request", return_value=response) as mock_request:
|
|
http_request_safe(host_only, "/my/endpoint", "GET")
|
|
mock_request.assert_called_with(
|
|
"GET",
|
|
"http://my-host/my/endpoint",
|
|
allow_redirects=True,
|
|
verify=False,
|
|
headers=DefaultRequestHeaderProvider().request_headers(),
|
|
timeout=120,
|
|
)
|
|
response.text = "non json"
|
|
http_request_safe(host_only, "/my/endpoint", "GET")
|
|
mock_request.assert_called_with(
|
|
"GET",
|
|
"http://my-host/my/endpoint",
|
|
allow_redirects=True,
|
|
verify=False,
|
|
headers=DefaultRequestHeaderProvider().request_headers(),
|
|
timeout=120,
|
|
)
|
|
response.status_code = 400
|
|
response.text = ""
|
|
with pytest.raises(MlflowException, match="Response body"):
|
|
http_request_safe(host_only, "/my/endpoint", "GET")
|
|
response.text = (
|
|
'{"error_code": "RESOURCE_DOES_NOT_EXIST", "message": "Node type not supported"}'
|
|
)
|
|
with pytest.raises(RestException, match="RESOURCE_DOES_NOT_EXIST: Node type not supported"):
|
|
http_request_safe(host_only, "/my/endpoint", "GET")
|
|
|
|
|
|
def test_numpy_encoder():
|
|
test_number = numpy.int64(42)
|
|
ne = NumpyEncoder()
|
|
defaulted_val = ne.default(test_number)
|
|
assert defaulted_val == 42
|
|
|
|
|
|
def test_numpy_encoder_fail():
|
|
if not hasattr(numpy, "float128"):
|
|
pytest.skip("numpy on exit this platform has no float128")
|
|
test_number = numpy.float128
|
|
ne = NumpyEncoder()
|
|
with pytest.raises(TypeError, match="not JSON serializable"):
|
|
ne.default(test_number)
|
|
|
|
|
|
def test_can_parse_as_json_object():
|
|
assert _can_parse_as_json_object("{}")
|
|
assert _can_parse_as_json_object('{"a": "b"}')
|
|
assert _can_parse_as_json_object('{"a": {"b": "c"}}')
|
|
assert not _can_parse_as_json_object("[0, 1, 2]")
|
|
assert not _can_parse_as_json_object('"abc"')
|
|
assert not _can_parse_as_json_object("123")
|
|
|
|
|
|
def test_http_request_customize_config(monkeypatch):
|
|
with mock.patch(
|
|
"mlflow.utils.rest_utils._get_http_response_with_retries"
|
|
) as mock_get_http_response_with_retries:
|
|
host_only = MlflowHostCreds("http://my-host")
|
|
monkeypatch.delenv("MLFLOW_HTTP_REQUEST_MAX_RETRIES", raising=False)
|
|
monkeypatch.delenv("MLFLOW_HTTP_REQUEST_BACKOFF_FACTOR", raising=False)
|
|
monkeypatch.delenv("MLFLOW_HTTP_REQUEST_TIMEOUT", raising=False)
|
|
monkeypatch.delenv("MLFLOW_HTTP_RESPECT_RETRY_AFTER_HEADER", raising=False)
|
|
http_request(host_only, "/my/endpoint", "GET")
|
|
mock_get_http_response_with_retries.assert_called_with(
|
|
mock.ANY,
|
|
mock.ANY,
|
|
7,
|
|
2,
|
|
1.0,
|
|
mock.ANY,
|
|
True,
|
|
headers=mock.ANY,
|
|
verify=mock.ANY,
|
|
timeout=120,
|
|
respect_retry_after_header=True,
|
|
)
|
|
mock_get_http_response_with_retries.reset_mock()
|
|
monkeypatch.setenv("MLFLOW_HTTP_REQUEST_MAX_RETRIES", "8")
|
|
monkeypatch.setenv("MLFLOW_HTTP_REQUEST_BACKOFF_FACTOR", "3")
|
|
monkeypatch.setenv("MLFLOW_HTTP_REQUEST_BACKOFF_JITTER", "1.0")
|
|
monkeypatch.setenv("MLFLOW_HTTP_REQUEST_TIMEOUT", "300")
|
|
monkeypatch.setenv("MLFLOW_HTTP_RESPECT_RETRY_AFTER_HEADER", "false")
|
|
http_request(host_only, "/my/endpoint", "GET")
|
|
mock_get_http_response_with_retries.assert_called_with(
|
|
mock.ANY,
|
|
mock.ANY,
|
|
8,
|
|
3,
|
|
1.0,
|
|
mock.ANY,
|
|
True,
|
|
headers=mock.ANY,
|
|
verify=mock.ANY,
|
|
timeout=300,
|
|
respect_retry_after_header=False,
|
|
)
|
|
|
|
|
|
def test_http_request_explains_how_to_increase_timeout_in_error_message():
|
|
with mock.patch("requests.Session.request", side_effect=requests.exceptions.Timeout):
|
|
with pytest.raises(
|
|
MlflowException,
|
|
match=(
|
|
r"To increase the timeout, set the environment variable "
|
|
+ re.escape(str(MLFLOW_HTTP_REQUEST_TIMEOUT))
|
|
),
|
|
):
|
|
http_request(MlflowHostCreds("http://my-host"), "/my/endpoint", "GET")
|
|
|
|
|
|
def test_augmented_raise_for_status():
|
|
response = requests.Response()
|
|
response.status_code = 403
|
|
response._content = b"Token expired"
|
|
|
|
with mock.patch("requests.Session.request", return_value=response) as mock_request:
|
|
response = requests.get("https://github.com/mlflow/mlflow.git")
|
|
mock_request.assert_called_once()
|
|
|
|
with pytest.raises(requests.HTTPError, match="Token expired") as e:
|
|
augmented_raise_for_status(response)
|
|
|
|
assert e.value.response == response
|
|
assert e.value.request == response.request
|
|
assert response.text in str(e.value)
|
|
|
|
|
|
def test_provide_redirect_kwarg():
|
|
with mock.patch("requests.Session.request") as mock_request:
|
|
mock_request.return_value.status_code = 302
|
|
mock_request.return_value.text = "mock response"
|
|
|
|
response = http_request(
|
|
MlflowHostCreds("http://my-host"),
|
|
"/my/endpoint",
|
|
"GET",
|
|
allow_redirects=False,
|
|
)
|
|
|
|
assert response.text == "mock response"
|
|
mock_request.assert_called_with(
|
|
"GET",
|
|
"http://my-host/my/endpoint",
|
|
allow_redirects=False,
|
|
headers=mock.ANY,
|
|
verify=mock.ANY,
|
|
timeout=120,
|
|
)
|
|
|
|
|
|
def test_http_request_max_retries(monkeypatch):
|
|
monkeypatch.setenv("_MLFLOW_HTTP_REQUEST_MAX_RETRIES_LIMIT", "15")
|
|
host_creds = MlflowHostCreds("http://example.com")
|
|
|
|
with mock.patch("requests.Session.request") as mock_request:
|
|
# Value exceeding limit should raise
|
|
with pytest.raises(MlflowException, match="The configured max_retries"):
|
|
http_request(host_creds, "/endpoint", "GET", max_retries=16)
|
|
mock_request.assert_not_called()
|
|
|
|
# Value equal to limit should succeed (boundary case)
|
|
http_request(host_creds, "/endpoint", "GET", max_retries=15)
|
|
assert mock_request.call_count == 1
|
|
|
|
# Value below limit should succeed
|
|
http_request(host_creds, "/endpoint", "GET", max_retries=3)
|
|
assert mock_request.call_count == 2
|
|
|
|
|
|
def test_http_request_backoff_factor(monkeypatch):
|
|
monkeypatch.setenv("_MLFLOW_HTTP_REQUEST_MAX_BACKOFF_FACTOR_LIMIT", "200")
|
|
host_creds = MlflowHostCreds("http://example.com")
|
|
|
|
with mock.patch("requests.Session.request") as mock_request:
|
|
# Value exceeding limit should raise
|
|
with pytest.raises(MlflowException, match="The configured backoff_factor"):
|
|
http_request(host_creds, "/endpoint", "GET", backoff_factor=250)
|
|
mock_request.assert_not_called()
|
|
|
|
# Value equal to limit should succeed (boundary case)
|
|
http_request(host_creds, "/endpoint", "GET", backoff_factor=200)
|
|
assert mock_request.call_count == 1
|
|
|
|
# Value below limit should succeed
|
|
http_request(host_creds, "/endpoint", "GET", backoff_factor=10)
|
|
assert mock_request.call_count == 2
|
|
|
|
|
|
def test_http_request_negative_max_retries():
|
|
host_creds = MlflowHostCreds("http://example.com")
|
|
|
|
with mock.patch("requests.Session.request") as mock_request:
|
|
with pytest.raises(MlflowException, match="The max_retries value must be either"):
|
|
http_request(host_creds, "/endpoint", "GET", max_retries=-1)
|
|
mock_request.assert_not_called()
|
|
|
|
|
|
def test_http_request_negative_backoff_factor():
|
|
host_creds = MlflowHostCreds("http://example.com")
|
|
|
|
with mock.patch("requests.Session.request") as mock_request:
|
|
with pytest.raises(MlflowException, match="The backoff_factor value must be"):
|
|
http_request(host_creds, "/endpoint", "GET", backoff_factor=-1)
|
|
mock_request.assert_not_called()
|
|
|
|
|
|
def test_suppress_databricks_retry_after_secs_warnings():
|
|
host_creds = MlflowHostCreds("http://example.com", use_databricks_sdk=True)
|
|
|
|
def mock_do(*args, **kwargs):
|
|
warnings.warn(_DATABRICKS_SDK_RETRY_AFTER_SECS_DEPRECATION_WARNING)
|
|
return mock.MagicMock()
|
|
|
|
with (
|
|
warnings.catch_warnings(record=True) as recorded_warnings,
|
|
mock.patch("mlflow.utils.rest_utils.get_workspace_client") as mock_get_workspace_client,
|
|
):
|
|
warnings.simplefilter("always")
|
|
mock_workspace_client = mock.MagicMock()
|
|
mock_workspace_client.api_client.do = mock_do
|
|
mock_get_workspace_client.return_value = mock_workspace_client
|
|
http_request(host_creds, "/endpoint", "GET")
|
|
mock_get_workspace_client.assert_called_once()
|
|
assert not any(
|
|
_DATABRICKS_SDK_RETRY_AFTER_SECS_DEPRECATION_WARNING in str(w.message)
|
|
for w in recorded_warnings
|
|
)
|
|
|
|
|
|
def test_databricks_sdk_retry_on_transient_errors():
|
|
host_creds = MlflowHostCreds("http://example.com", use_databricks_sdk=True)
|
|
|
|
call_count = 0
|
|
|
|
def mock_do_failing_then_success(*args, **kwargs):
|
|
nonlocal call_count
|
|
call_count += 1
|
|
if call_count <= 2: # Fail first 2 attempts
|
|
from databricks.sdk.errors import DatabricksError
|
|
|
|
from mlflow.protos.databricks_pb2 import ErrorCode
|
|
|
|
raise DatabricksError(
|
|
error_code=ErrorCode.Name(ErrorCode.INTERNAL_ERROR), message="Transient error"
|
|
)
|
|
# Success on 3rd attempt
|
|
response_mock = mock.MagicMock()
|
|
response_mock._response = mock.MagicMock()
|
|
return {"contents": response_mock}
|
|
|
|
with mock.patch("mlflow.utils.rest_utils.get_workspace_client") as mock_get_workspace_client:
|
|
mock_workspace_client = mock.MagicMock()
|
|
mock_workspace_client.api_client.do = mock_do_failing_then_success
|
|
mock_get_workspace_client.return_value = mock_workspace_client
|
|
|
|
# Use smaller retry timeout to make test run faster
|
|
response = http_request(
|
|
host_creds,
|
|
"/endpoint",
|
|
"GET",
|
|
retry_timeout_seconds=10,
|
|
backoff_factor=0.1, # Very small backoff for faster test
|
|
)
|
|
|
|
assert call_count == 3 # Should retry 2 times, succeed on 3rd
|
|
assert response is not None
|
|
|
|
|
|
def test_databricks_sdk_retry_max_retries_exceeded():
|
|
host_creds = MlflowHostCreds("http://example.com", use_databricks_sdk=True)
|
|
|
|
call_count = 0
|
|
|
|
def mock_do_always_fail(*args, **kwargs):
|
|
nonlocal call_count
|
|
call_count += 1
|
|
from databricks.sdk.errors import DatabricksError
|
|
|
|
raise DatabricksError(error_code="INTERNAL_ERROR", message="Always fails")
|
|
|
|
with (
|
|
mock.patch("mlflow.utils.rest_utils.get_workspace_client") as mock_get_workspace_client,
|
|
mock.patch("mlflow.utils.rest_utils._logger") as mock_logger,
|
|
):
|
|
mock_workspace_client = mock.MagicMock()
|
|
mock_workspace_client.api_client.do = mock_do_always_fail
|
|
mock_get_workspace_client.return_value = mock_workspace_client
|
|
|
|
response = http_request(host_creds, "/endpoint", "GET", max_retries=3)
|
|
|
|
assert call_count == 4 # Initial call + 3 retries
|
|
assert response.status_code == 500 # Should return error response
|
|
|
|
# Check that max retries warning was logged
|
|
mock_logger.warning.assert_called()
|
|
warning_call = mock_logger.warning.call_args[0][0]
|
|
assert "Max retries (3) exceeded" in warning_call
|
|
|
|
|
|
def test_databricks_sdk_retry_timeout_exceeded():
|
|
host_creds = MlflowHostCreds("http://example.com", use_databricks_sdk=True)
|
|
|
|
call_count = 0
|
|
|
|
def mock_do_always_fail(*args, **kwargs):
|
|
nonlocal call_count
|
|
call_count += 1
|
|
|
|
time.sleep(0.1) # Small delay to ensure timeout
|
|
from databricks.sdk.errors import DatabricksError
|
|
|
|
raise DatabricksError(error_code="INTERNAL_ERROR", message="Always fails")
|
|
|
|
with (
|
|
mock.patch("mlflow.utils.rest_utils.get_workspace_client") as mock_get_workspace_client,
|
|
mock.patch("mlflow.utils.rest_utils._logger") as mock_logger,
|
|
):
|
|
mock_workspace_client = mock.MagicMock()
|
|
mock_workspace_client.api_client.do = mock_do_always_fail
|
|
mock_get_workspace_client.return_value = mock_workspace_client
|
|
|
|
response = http_request(
|
|
host_creds,
|
|
"/endpoint",
|
|
"GET",
|
|
retry_timeout_seconds=0.2, # Very short timeout
|
|
max_retries=10, # High retry limit that shouldn't be reached
|
|
)
|
|
|
|
assert call_count >= 1 # At least initial call
|
|
assert response.status_code == 500 # Should return error response
|
|
|
|
# Check that timeout warning was logged
|
|
mock_logger.warning.assert_called()
|
|
warning_call = mock_logger.warning.call_args[0][0]
|
|
assert "Retry timeout (0.2s) exceeded" in warning_call
|
|
|
|
|
|
def test_databricks_sdk_retry_non_retryable_error():
|
|
host_creds = MlflowHostCreds("http://example.com", use_databricks_sdk=True)
|
|
|
|
call_count = 0
|
|
|
|
def mock_do_non_retryable_error(*args, **kwargs):
|
|
nonlocal call_count
|
|
call_count += 1
|
|
from databricks.sdk.errors import InvalidParameterValue
|
|
|
|
# Use an error code that maps to 400 (non-retryable)
|
|
raise InvalidParameterValue(error_code="INVALID_PARAMETER_VALUE", message="Bad request")
|
|
|
|
with mock.patch("mlflow.utils.rest_utils.get_workspace_client") as mock_get_workspace_client:
|
|
mock_workspace_client = mock.MagicMock()
|
|
mock_workspace_client.api_client.do = mock_do_non_retryable_error
|
|
mock_get_workspace_client.return_value = mock_workspace_client
|
|
|
|
response = http_request(host_creds, "/endpoint", "GET", max_retries=5)
|
|
|
|
assert call_count == 1 # Should not retry on non-retryable error
|
|
assert response.status_code == 400 # Should return 400 for INVALID_PARAMETER_VALUE
|
|
|
|
|
|
def test_databricks_sdk_retry_backoff_calculation():
|
|
from databricks.sdk.errors import DatabricksError
|
|
|
|
from mlflow.utils.request_utils import _TRANSIENT_FAILURE_RESPONSE_CODES
|
|
from mlflow.utils.rest_utils import _retry_databricks_sdk_call_with_exponential_backoff
|
|
|
|
call_count = 0
|
|
|
|
def mock_failing_call():
|
|
nonlocal call_count
|
|
call_count += 1
|
|
|
|
raise DatabricksError(error_code="INTERNAL_ERROR", message="Mock error")
|
|
|
|
with mock.patch("mlflow.utils.rest_utils._time_sleep") as mock_sleep:
|
|
with pytest.raises(DatabricksError, match="Mock error"):
|
|
_retry_databricks_sdk_call_with_exponential_backoff(
|
|
call_func=mock_failing_call,
|
|
retry_codes=_TRANSIENT_FAILURE_RESPONSE_CODES,
|
|
retry_timeout_seconds=10,
|
|
backoff_factor=1, # Use 1 for predictable calculation
|
|
backoff_jitter=0, # No jitter for predictable calculation
|
|
max_retries=3,
|
|
)
|
|
|
|
# Verify sleep was called with correct intervals
|
|
# attempt 0 (1st retry): 0 seconds (immediate)
|
|
# attempt 1 (2nd retry): 1 * (2^1) = 2 seconds
|
|
# attempt 2 (3rd retry): 1 * (2^2) = 4 seconds
|
|
expected_sleep_times = [0, 2, 4]
|
|
actual_sleep_times = [call.args[0] for call in mock_sleep.call_args_list]
|
|
assert actual_sleep_times == expected_sleep_times
|
|
assert call_count == 4 # Initial + 3 retries
|
|
|
|
|
|
@pytest.mark.skip
|
|
def test_timeout_parameter_propagation_with_timeout():
|
|
with (
|
|
mock.patch("databricks.sdk.WorkspaceClient") as mock_workspace_client,
|
|
mock.patch("databricks.sdk.config.Config") as mock_config,
|
|
):
|
|
# Test http_request with timeout via get_workspace_client directly
|
|
mock_workspace_client_instance = mock.MagicMock()
|
|
mock_workspace_client_instance.api_client.do.return_value = {"contents": mock.MagicMock()}
|
|
mock_workspace_client.return_value = mock_workspace_client_instance
|
|
|
|
get_workspace_client(
|
|
use_secret_scope_token=False,
|
|
host="http://my-host",
|
|
token=None,
|
|
databricks_auth_profile="my-profile",
|
|
retry_timeout_seconds=None,
|
|
timeout=180,
|
|
)
|
|
|
|
mock_config.assert_called_once_with(
|
|
profile="my-profile",
|
|
http_timeout_seconds=180,
|
|
retry_timeout_seconds=mock.ANY,
|
|
)
|
|
|
|
|
|
@pytest.mark.skip
|
|
def test_timeout_parameter_propagation_without_timeout():
|
|
with (
|
|
mock.patch("databricks.sdk.WorkspaceClient") as mock_workspace_client,
|
|
mock.patch("databricks.sdk.config.Config") as mock_config,
|
|
):
|
|
# Test http_request without timeout via get_workspace_client directly
|
|
mock_workspace_client_instance = mock.MagicMock()
|
|
mock_workspace_client_instance.api_client.do.return_value = {"contents": mock.MagicMock()}
|
|
mock_workspace_client.return_value = mock_workspace_client_instance
|
|
|
|
get_workspace_client(
|
|
use_secret_scope_token=False,
|
|
host="http://my-host",
|
|
token=None,
|
|
databricks_auth_profile="my-profile",
|
|
retry_timeout_seconds=None,
|
|
timeout=None,
|
|
)
|
|
|
|
mock_config.assert_called_once_with(
|
|
profile="my-profile",
|
|
retry_timeout_seconds=mock.ANY,
|
|
)
|
|
|
|
|
|
def test_deployment_client_timeout_propagation(monkeypatch):
|
|
with (
|
|
mock.patch("mlflow.utils.rest_utils.get_workspace_client") as mock_get_workspace_client,
|
|
mock.patch(
|
|
"mlflow.utils.databricks_utils.get_databricks_host_creds"
|
|
) as mock_get_databricks_host_creds,
|
|
mock.patch(
|
|
"mlflow.deployments.databricks.get_databricks_host_creds"
|
|
) as mock_deployment_host_creds,
|
|
):
|
|
# Mock the host creds to use Databricks SDK
|
|
mock_host_creds = MlflowHostCreds("http://my-host", use_databricks_sdk=True)
|
|
mock_get_databricks_host_creds.return_value = mock_host_creds
|
|
mock_deployment_host_creds.return_value = mock_host_creds
|
|
|
|
# Mock workspace client and its response
|
|
mock_workspace_client_instance = mock.MagicMock()
|
|
mock_workspace_client_instance.api_client.do.return_value = {"contents": mock.MagicMock()}
|
|
mock_get_workspace_client.return_value = mock_workspace_client_instance
|
|
|
|
# Set the environment variable to a custom value using monkeypatch
|
|
monkeypatch.setenv("MLFLOW_DEPLOYMENT_PREDICT_TIMEOUT", "300")
|
|
|
|
# Create deployment client and call predict
|
|
client = DatabricksDeploymentClient("databricks")
|
|
client.predict(endpoint="test-endpoint", inputs={"test": "data"})
|
|
|
|
# Verify get_workspace_client was called with the deployment predict timeout
|
|
mock_get_workspace_client.assert_called_once_with(
|
|
False, # use_secret_scope_token
|
|
"http://my-host", # host
|
|
None, # token
|
|
None, # databricks_auth_profile
|
|
retry_timeout_seconds=600,
|
|
timeout=300, # MLFLOW_DEPLOYMENT_PREDICT_TIMEOUT value
|
|
)
|
|
|
|
|
|
def test_http_request_with_databricks_traffic_id(monkeypatch: pytest.MonkeyPatch):
|
|
response = mock.MagicMock()
|
|
response.status_code = 200
|
|
|
|
# Test with env var set
|
|
monkeypatch.setenv(_MLFLOW_DATABRICKS_TRAFFIC_ID.name, "test-traffic-id-12345")
|
|
with mock.patch("requests.Session.request", return_value=response) as mock_request:
|
|
http_request(MlflowHostCreds("http://my-host"), "/my/endpoint", "GET")
|
|
headers = mock_request.call_args.kwargs["headers"]
|
|
assert headers["x-databricks-traffic-id"] == "test-traffic-id-12345"
|
|
|
|
# Test without env var set
|
|
monkeypatch.delenv(_MLFLOW_DATABRICKS_TRAFFIC_ID.name)
|
|
with mock.patch("requests.Session.request", return_value=response) as mock_request:
|
|
http_request(MlflowHostCreds("http://my-host"), "/my/endpoint", "GET")
|
|
headers = mock_request.call_args.kwargs["headers"]
|
|
assert "x-databricks-traffic-id" not in headers
|
|
|
|
|
|
def test_http_request_with_workspace_id():
|
|
response = mock.MagicMock()
|
|
response.status_code = 200
|
|
|
|
# With workspace_id set, header should be included
|
|
creds = MlflowHostCreds("http://my-host", workspace_id="6051921418418893")
|
|
with mock.patch("requests.Session.request", return_value=response) as mock_request:
|
|
http_request(creds, "/my/endpoint", "GET")
|
|
headers = mock_request.call_args.kwargs["headers"]
|
|
assert headers["x-databricks-org-id"] == "6051921418418893"
|
|
|
|
# Without workspace_id, header should not be present
|
|
creds = MlflowHostCreds("http://my-host")
|
|
with mock.patch("requests.Session.request", return_value=response) as mock_request:
|
|
http_request(creds, "/my/endpoint", "GET")
|
|
headers = mock_request.call_args.kwargs["headers"]
|
|
assert "x-databricks-org-id" not in headers
|
|
|
|
|
|
def test_mlflow_host_creds_workspace_id_equality():
|
|
creds1 = MlflowHostCreds("http://my-host", workspace_id="123")
|
|
creds2 = MlflowHostCreds("http://my-host", workspace_id="123")
|
|
creds3 = MlflowHostCreds("http://my-host", workspace_id="456")
|
|
creds4 = MlflowHostCreds("http://my-host")
|
|
|
|
assert creds1 == creds2
|
|
assert creds1 != creds3
|
|
assert creds1 != creds4
|
|
assert hash(creds1) == hash(creds2)
|
|
assert hash(creds1) != hash(creds3)
|
|
|
|
|
|
def test_oauth_error_when_sdk_enabled(monkeypatch):
|
|
# When MLFLOW_ENABLE_DB_SDK=true and the SDK failed to initialize, the OAuth error
|
|
# must acknowledge the env var is set and point at the preceding SDK warning, instead
|
|
# of telling the user to set an env var they have already set.
|
|
monkeypatch.setenv("MLFLOW_ENABLE_DB_SDK", "true")
|
|
creds = MlflowHostCreds(
|
|
"http://my-host",
|
|
client_id="some-id",
|
|
client_secret="some-secret",
|
|
)
|
|
with pytest.raises(MlflowException, match="SDK is enabled but failed to initialize") as exc:
|
|
http_request(creds, "/my/endpoint", "GET")
|
|
msg = str(exc.value)
|
|
assert "MLFLOW_ENABLE_DB_SDK is currently set to 'True'" in msg
|
|
assert "Failed to create databricks SDK workspace client" in msg
|
|
|
|
|
|
def test_oauth_error_when_sdk_disabled(monkeypatch):
|
|
monkeypatch.setenv("MLFLOW_ENABLE_DB_SDK", "false")
|
|
creds = MlflowHostCreds(
|
|
"http://my-host",
|
|
client_id="some-id",
|
|
client_secret="some-secret",
|
|
)
|
|
with pytest.raises(MlflowException, match="Set 'MLFLOW_ENABLE_DB_SDK' to true") as exc:
|
|
http_request(creds, "/my/endpoint", "GET")
|
|
msg = str(exc.value)
|
|
assert "MLFLOW_ENABLE_DB_SDK is currently set to 'False'" in msg
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("timeout", "retry_timeout_seconds", "should_warn"),
|
|
[
|
|
(300, 120, True),
|
|
(120, 600, False),
|
|
(300, 300, False),
|
|
(None, 120, False),
|
|
(300, None, False),
|
|
(None, None, False),
|
|
],
|
|
)
|
|
def test_validate_deployment_timeout_config(timeout, retry_timeout_seconds, should_warn):
|
|
from mlflow.utils.rest_utils import validate_deployment_timeout_config
|
|
|
|
if should_warn:
|
|
with warnings.catch_warnings(record=True) as w:
|
|
warnings.simplefilter("always")
|
|
validate_deployment_timeout_config(
|
|
timeout=timeout, retry_timeout_seconds=retry_timeout_seconds
|
|
)
|
|
assert len(w) == 1
|
|
warning_msg = str(w[0].message)
|
|
assert "MLFLOW_DEPLOYMENT_PREDICT_TOTAL_TIMEOUT" in warning_msg
|
|
assert f"({retry_timeout_seconds}s)" in warning_msg
|
|
assert f"({timeout}s)" in warning_msg
|
|
else:
|
|
with warnings.catch_warnings(record=True) as w:
|
|
warnings.simplefilter("always")
|
|
validate_deployment_timeout_config(
|
|
timeout=timeout, retry_timeout_seconds=retry_timeout_seconds
|
|
)
|
|
assert len(w) == 0
|