import base64 import json from functools import lru_cache import requests from mlflow.environment_variables import ( _MLFLOW_HTTP_REQUEST_MAX_BACKOFF_FACTOR_LIMIT, _MLFLOW_HTTP_REQUEST_MAX_RETRIES_LIMIT, MLFLOW_DATABRICKS_ENDPOINT_HTTP_RETRY_TIMEOUT, MLFLOW_ENABLE_DB_SDK, MLFLOW_HTTP_REQUEST_BACKOFF_FACTOR, MLFLOW_HTTP_REQUEST_BACKOFF_JITTER, MLFLOW_HTTP_REQUEST_MAX_RETRIES, MLFLOW_HTTP_REQUEST_TIMEOUT, MLFLOW_HTTP_RESPECT_RETRY_AFTER_HEADER, ) from mlflow.exceptions import ( CUSTOMER_UNAUTHORIZED, ERROR_CODE_TO_HTTP_STATUS, INVALID_PARAMETER_VALUE, InvalidUrlException, MlflowException, RestException, get_error_code, ) from mlflow.protos import databricks_pb2 from mlflow.protos.databricks_pb2 import ENDPOINT_NOT_FOUND, ErrorCode from mlflow.utils.proto_json_utils import parse_dict from mlflow.utils.request_utils import ( _TRANSIENT_FAILURE_RESPONSE_CODES, _get_http_response_with_retries, augmented_raise_for_status, # noqa: F401 cloud_storage_http_request, # noqa: F401 ) from mlflow.utils.string_utils import strip_suffix RESOURCE_NON_EXISTENT = "RESOURCE_DOES_NOT_EXIST" _REST_API_PATH_PREFIX = "/api/2.0" _UC_OSS_REST_API_PATH_PREFIX = "/api/2.1" _TRACE_REST_API_PATH_PREFIX = f"{_REST_API_PATH_PREFIX}/mlflow/traces" _ARMERIA_OK = "200 OK" def http_request( host_creds, endpoint, method, max_retries=None, backoff_factor=None, backoff_jitter=None, extra_headers=None, retry_codes=_TRANSIENT_FAILURE_RESPONSE_CODES, timeout=None, raise_on_status=True, respect_retry_after_header=None, retry_timeout_seconds=None, **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. Args: host_creds: A :py:class:`mlflow.rest_utils.MlflowHostCreds` object containing hostname and optional authentication. endpoint: A string for service endpoint, e.g. "/path/to/object". method: A string indicating the method to use, e.g. "GET", "POST", "PUT". max_retries: Maximum number of retries before throwing an exception. 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. backoff_jitter: A random jitter to add to the backoff interval. extra_headers: A dict of HTTP header name-value pairs to be included in the request. retry_codes: A list of HTTP response error codes that qualifies for retry. timeout: Wait for timeout seconds for response from remote server for connect and read request. raise_on_status: Whether to raise an exception, or return a response, if status falls in retry_codes range and retries have been exhausted. respect_retry_after_header: Whether to respect Retry-After header on status codes defined as Retry.RETRY_AFTER_STATUS_CODES or not. retry_timeout_seconds: Timeout for retires. Only effective when using Databricks SDK. kwargs: Additional keyword arguments to pass to `requests.Session.request()` Returns: requests.Response object. """ cleaned_hostname = strip_suffix(host_creds.host, "/") url = f"{cleaned_hostname}{endpoint}" if host_creds.use_databricks_sdk: from databricks.sdk.errors import DatabricksError ws_client = get_workspace_client( host_creds.use_secret_scope_token, host_creds.host, host_creds.token, host_creds.databricks_auth_profile, retry_timeout_seconds=retry_timeout_seconds, ) try: # Databricks SDK `APIClient.do` API is for making request using # HTTP # https://github.com/databricks/databricks-sdk-py/blob/a714146d9c155dd1e3567475be78623f72028ee0/databricks/sdk/core.py#L134 raw_response = ws_client.api_client.do( method=method, path=endpoint, headers=extra_headers, raw=True, query=kwargs.get("params"), body=kwargs.get("json"), files=kwargs.get("files"), data=kwargs.get("data"), ) return raw_response["contents"]._response except DatabricksError as e: response = requests.Response() response.url = url response.status_code = ERROR_CODE_TO_HTTP_STATUS.get(e.error_code, 500) response.reason = str(e) response.encoding = "UTF-8" response._content = json.dumps( { "error_code": e.error_code, "message": str(e), } ).encode("UTF-8") return response max_retries = MLFLOW_HTTP_REQUEST_MAX_RETRIES.get() if max_retries is None else max_retries backoff_factor = ( MLFLOW_HTTP_REQUEST_BACKOFF_FACTOR.get() if backoff_factor is None else backoff_factor ) _validate_max_retries(max_retries) _validate_backoff_factor(backoff_factor) respect_retry_after_header = ( MLFLOW_HTTP_RESPECT_RETRY_AFTER_HEADER.get() if respect_retry_after_header is None else respect_retry_after_header ) backoff_jitter = ( MLFLOW_HTTP_REQUEST_BACKOFF_JITTER.get() if backoff_jitter is None else backoff_jitter ) timeout = MLFLOW_HTTP_REQUEST_TIMEOUT.get() if timeout is None else timeout auth_str = None if host_creds.username and host_creds.password: basic_auth_str = f"{host_creds.username}:{host_creds.password}".encode() auth_str = "Basic " + base64.standard_b64encode(basic_auth_str).decode("utf-8") elif host_creds.token: auth_str = f"Bearer {host_creds.token}" elif host_creds.client_secret: raise MlflowException( "To use OAuth authentication, set environmental variable " f"'{MLFLOW_ENABLE_DB_SDK.name}' to true", error_code=CUSTOMER_UNAUTHORIZED, ) from mlflow.tracking.request_header.registry import resolve_request_headers headers = dict(**resolve_request_headers()) if extra_headers: headers = dict(**headers, **extra_headers) if auth_str: headers["Authorization"] = auth_str if host_creds.client_cert_path is not None: kwargs["cert"] = host_creds.client_cert_path if host_creds.aws_sigv4: # will overwrite the Authorization header from requests_auth_aws_sigv4 import AWSSigV4 kwargs["auth"] = AWSSigV4("execute-api") elif host_creds.auth: from mlflow.tracking.request_auth.registry import fetch_auth kwargs["auth"] = fetch_auth(host_creds.auth) try: return _get_http_response_with_retries( method, url, max_retries, backoff_factor, backoff_jitter, retry_codes, raise_on_status, headers=headers, verify=host_creds.verify, timeout=timeout, respect_retry_after_header=respect_retry_after_header, **kwargs, ) except requests.exceptions.Timeout as to: raise MlflowException( f"API request to {url} failed with timeout exception {to}." " To increase the timeout, set the environment variable " f"{MLFLOW_HTTP_REQUEST_TIMEOUT!s} to a larger value." ) from to except requests.exceptions.InvalidURL as iu: raise InvalidUrlException(f"Invalid url: {url}") from iu except Exception as e: raise MlflowException(f"API request to {url} failed with exception {e}") @lru_cache(maxsize=1) def get_workspace_client( use_secret_scope_token, host, token, databricks_auth_profile, retry_timeout_seconds=None, ): from databricks.sdk import WorkspaceClient from databricks.sdk.config import Config if use_secret_scope_token: kwargs = {"host": host, "token": token} else: kwargs = {"profile": databricks_auth_profile} config = Config( **kwargs, retry_timeout_seconds=retry_timeout_seconds or MLFLOW_DATABRICKS_ENDPOINT_HTTP_RETRY_TIMEOUT.get(), ) # Note: If we use `config` param, all SDK configurations must be set in `config` object. return WorkspaceClient(config=config) def _can_parse_as_json_object(string): try: return isinstance(json.loads(string), dict) except Exception: return False def http_request_safe(host_creds, endpoint, method, **kwargs): """ Wrapper around ``http_request`` that also verifies that the request succeeds with code 200. """ response = http_request(host_creds=host_creds, endpoint=endpoint, method=method, **kwargs) return verify_rest_response(response, endpoint) def verify_rest_response(response, endpoint): """Verify the return code and format, raise exception if the request was not successful.""" # Handle Armeria-specific response case where response text is "200 OK" if response.status_code == 200 and response.text.strip() == _ARMERIA_OK: response._content = b"{}" # Update response content to be an empty JSON dictionary return response # Handle non-200 status codes if response.status_code != 200: if _can_parse_as_json_object(response.text): raise RestException(json.loads(response.text)) else: base_msg = ( f"API request to endpoint {endpoint} " f"failed with error code {response.status_code} != 200" ) raise MlflowException( f"{base_msg}. Response body: '{response.text}'", error_code=get_error_code(response.status_code), ) # Skip validation for endpoints (e.g. DBFS file-download API) which may return a non-JSON # response if endpoint.startswith(_REST_API_PATH_PREFIX) and not _can_parse_as_json_object(response.text): base_msg = ( "API request to endpoint was successful but the response body was not " "in a valid JSON format" ) raise MlflowException(f"{base_msg}. Response body: '{response.text}'") return response def _validate_max_retries(max_retries): max_retry_limit = _MLFLOW_HTTP_REQUEST_MAX_RETRIES_LIMIT.get() if max_retry_limit < 0: raise MlflowException( message=f"The current maximum retry limit is invalid ({max_retry_limit}). " "Cannot be negative.", error_code=INVALID_PARAMETER_VALUE, ) if max_retries >= max_retry_limit: raise MlflowException( message=f"The configured max_retries value ({max_retries}) is " f"in excess of the maximum allowable retries ({max_retry_limit})", error_code=INVALID_PARAMETER_VALUE, ) if max_retries < 0: raise MlflowException( message=f"The max_retries value must be either 0 a positive integer. Got {max_retries}", error_code=INVALID_PARAMETER_VALUE, ) def _validate_backoff_factor(backoff_factor): max_backoff_factor_limit = _MLFLOW_HTTP_REQUEST_MAX_BACKOFF_FACTOR_LIMIT.get() if max_backoff_factor_limit < 0: raise MlflowException( message="The current maximum backoff factor limit is invalid " f"({max_backoff_factor_limit}). Cannot be negative.", error_code=INVALID_PARAMETER_VALUE, ) if backoff_factor >= max_backoff_factor_limit: raise MlflowException( message=f"The configured backoff_factor value ({backoff_factor}) is in excess " "of the maximum allowable backoff_factor limit " f"({max_backoff_factor_limit})", error_code=INVALID_PARAMETER_VALUE, ) if backoff_factor < 0: raise MlflowException( message="The backoff_factor value must be either 0 a positive integer. " f"Got {backoff_factor}", error_code=INVALID_PARAMETER_VALUE, ) def _get_path(path_prefix, endpoint_path): return f"{path_prefix}{endpoint_path}" def extract_api_info_for_service(service, path_prefix): """Return a dictionary mapping each API method to a tuple (path, HTTP method)""" service_methods = service.DESCRIPTOR.methods res = {} for service_method in service_methods: endpoints = service_method.GetOptions().Extensions[databricks_pb2.rpc].endpoints endpoint = endpoints[0] endpoint_path = _get_path(path_prefix, endpoint.path) res[service().GetRequestClass(service_method)] = (endpoint_path, endpoint.method) return res def extract_all_api_info_for_service(service, path_prefix): """Return a dictionary mapping each API method to a list of tuples [(path, HTTP method)]""" service_methods = service.DESCRIPTOR.methods res = {} for service_method in service_methods: endpoints = service_method.GetOptions().Extensions[databricks_pb2.rpc].endpoints res[service().GetRequestClass(service_method)] = [ (_get_path(path_prefix, endpoint.path), endpoint.method) for endpoint in endpoints ] return res def get_single_trace_endpoint(request_id): return f"{_TRACE_REST_API_PATH_PREFIX}/{request_id}" def get_trace_info_endpoint(request_id): return f"{get_single_trace_endpoint(request_id)}/info" def get_trace_assessment_endpoint(request_id): # TEMPORARY ENDPOINT: this is currently hosted at /api/2.0/... but will be moved to /api/3.0/... return f"{get_single_trace_endpoint(request_id)}" def get_set_trace_tag_endpoint(request_id): return f"{get_single_trace_endpoint(request_id)}/tags" def get_create_assessment_endpoint(trace_id: str): return f"{_TRACE_REST_API_PATH_PREFIX}/{trace_id}/assessments" def get_single_assessment_endpoint(trace_id: str, assessment_id: str): return f"{_TRACE_REST_API_PATH_PREFIX}/{trace_id}/assessments/{assessment_id}" def call_endpoint(host_creds, endpoint, method, json_body, response_proto, extra_headers=None): # Convert json string to json dictionary, to pass to requests if json_body is not None: json_body = json.loads(json_body) call_kwargs = { "host_creds": host_creds, "endpoint": endpoint, "method": method, } if extra_headers is not None: call_kwargs["extra_headers"] = extra_headers if method == "GET": call_kwargs["params"] = json_body response = http_request(**call_kwargs) else: call_kwargs["json"] = json_body response = http_request(**call_kwargs) response = verify_rest_response(response, endpoint) response_to_parse = response.text js_dict = json.loads(response_to_parse) parse_dict(js_dict=js_dict, message=response_proto) return response_proto def call_endpoints(host_creds, endpoints, json_body, response_proto, extra_headers=None): # The order that the endpoints are called in is defined by the order # specified in ModelRegistryService in model_registry.proto for i, (endpoint, method) in enumerate(endpoints): try: return call_endpoint( host_creds, endpoint, method, json_body, response_proto, extra_headers ) except RestException as e: if e.error_code != ErrorCode.Name(ENDPOINT_NOT_FOUND) or i == len(endpoints) - 1: raise e class MlflowHostCreds: """ Provides a hostname and optional authentication for talking to an MLflow tracking server. Args: host: Hostname (e.g., http://localhost:5000) to MLflow server. Required. username: Username to use with Basic authentication when talking to server. If this is specified, password must also be specified. password: Password to use with Basic authentication when talking to server. If this is specified, username must also be specified. token: Token to use with Bearer authentication when talking to server. If provided, user/password authentication will be ignored. aws_sigv4: If true, we will create a signature V4 to be added for any outgoing request. Keys for signing the request can be passed via ENV variables, or will be fetched via boto3 session. auth: If set, the auth will be added for any outgoing request. Keys for signing the request can be passed via ENV variables, ignore_tls_verification: If true, we will not verify the server's hostname or TLS certificate. This is useful for certain testing situations, but should never be true in production. If this is set to true ``server_cert_path`` must not be set. client_cert_path: Path to ssl client cert file (.pem). Sets the cert param of the ``requests.request`` function (see https://requests.readthedocs.io/en/master/api/). server_cert_path: Path to a CA bundle to use. Sets the verify param of the ``requests.request`` function (see https://requests.readthedocs.io/en/master/api/). If this is set ``ignore_tls_verification`` must be false. use_databricks_sdk: A boolean value represent whether using Databricks SDK for authentication. databricks_auth_profile: The name of the profile used by Databricks SDK for authentication. client_id: The client ID used by Databricks OAuth client_secret: The client secret used by Databricks OAuth """ def __init__( self, host, username=None, password=None, token=None, aws_sigv4=False, auth=None, ignore_tls_verification=False, client_cert_path=None, server_cert_path=None, use_databricks_sdk=False, databricks_auth_profile=None, client_id=None, client_secret=None, use_secret_scope_token=False, ): if not host: raise MlflowException( message="host is a required parameter for MlflowHostCreds", error_code=INVALID_PARAMETER_VALUE, ) if ignore_tls_verification and (server_cert_path is not None): raise MlflowException( message=( "When 'ignore_tls_verification' is true then 'server_cert_path' " "must not be set! This error may have occurred because the " "'MLFLOW_TRACKING_INSECURE_TLS' and 'MLFLOW_TRACKING_SERVER_CERT_PATH' " "environment variables are both set - only one of these environment " "variables may be set." ), error_code=INVALID_PARAMETER_VALUE, ) self.host = host self.username = username self.password = password self.token = token self.aws_sigv4 = aws_sigv4 self.auth = auth self.ignore_tls_verification = ignore_tls_verification self.client_cert_path = client_cert_path self.server_cert_path = server_cert_path self.use_databricks_sdk = use_databricks_sdk self.databricks_auth_profile = databricks_auth_profile self.client_id = client_id self.client_secret = client_secret self.use_secret_scope_token = use_secret_scope_token def __eq__(self, other): if isinstance(other, self.__class__): return self.__dict__ == other.__dict__ return NotImplemented @property def verify(self): if self.use_databricks_sdk: # Let databricks-sdk set HTTP request `verify` param. return None if self.server_cert_path is None: return not self.ignore_tls_verification else: return self.server_cert_path