wehub-resource-sync
850 lines
30 KiB
Python
850 lines
30 KiB
Python
import json
|
|
import posixpath
|
|
import warnings
|
|
from typing import Any, Iterator
|
|
|
|
from mlflow.deployments import BaseDeploymentClient
|
|
from mlflow.deployments.constants import (
|
|
MLFLOW_DEPLOYMENT_CLIENT_REQUEST_RETRY_CODES,
|
|
)
|
|
from mlflow.environment_variables import (
|
|
MLFLOW_DEPLOYMENT_PREDICT_TIMEOUT,
|
|
MLFLOW_DEPLOYMENT_PREDICT_TOTAL_TIMEOUT,
|
|
MLFLOW_HTTP_REQUEST_TIMEOUT,
|
|
)
|
|
from mlflow.exceptions import MlflowException
|
|
from mlflow.utils import AttrDict
|
|
from mlflow.utils.annotations import deprecated
|
|
from mlflow.utils.databricks_utils import get_databricks_host_creds
|
|
from mlflow.utils.rest_utils import (
|
|
augmented_raise_for_status,
|
|
http_request,
|
|
validate_deployment_timeout_config,
|
|
)
|
|
|
|
|
|
class DatabricksEndpoint(AttrDict):
|
|
"""
|
|
A dictionary-like object representing a Databricks serving endpoint.
|
|
|
|
.. code-block:: python
|
|
|
|
endpoint = DatabricksEndpoint({
|
|
"name": "chat",
|
|
"creator": "alice@company.com",
|
|
"creation_timestamp": 0,
|
|
"last_updated_timestamp": 0,
|
|
"state": {...},
|
|
"config": {...},
|
|
"tags": [...],
|
|
"id": "88fd3f75a0d24b0380ddc40484d7a31b",
|
|
})
|
|
assert endpoint.name == "chat"
|
|
"""
|
|
|
|
|
|
class DatabricksDeploymentClient(BaseDeploymentClient):
|
|
"""
|
|
Client for interacting with Databricks serving endpoints.
|
|
|
|
Example:
|
|
|
|
First, set up credentials for authentication:
|
|
|
|
.. code-block:: bash
|
|
|
|
export DATABRICKS_HOST=...
|
|
export DATABRICKS_TOKEN=...
|
|
|
|
.. seealso::
|
|
|
|
See https://docs.databricks.com/en/dev-tools/auth.html for other authentication methods.
|
|
|
|
Then, create a deployment client and use it to interact with Databricks serving endpoints:
|
|
|
|
.. code-block:: python
|
|
|
|
from mlflow.deployments import get_deploy_client
|
|
|
|
client = get_deploy_client("databricks")
|
|
endpoints = client.list_endpoints()
|
|
assert endpoints == [
|
|
{
|
|
"name": "chat",
|
|
"creator": "alice@company.com",
|
|
"creation_timestamp": 0,
|
|
"last_updated_timestamp": 0,
|
|
"state": {...},
|
|
"config": {...},
|
|
"tags": [...],
|
|
"id": "88fd3f75a0d24b0380ddc40484d7a31b",
|
|
},
|
|
]
|
|
"""
|
|
|
|
def create_deployment(self, name, model_uri, flavor=None, config=None, endpoint=None):
|
|
"""
|
|
.. warning::
|
|
|
|
This method is not implemented for `DatabricksDeploymentClient`.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def update_deployment(self, name, model_uri=None, flavor=None, config=None, endpoint=None):
|
|
"""
|
|
.. warning::
|
|
|
|
This method is not implemented for `DatabricksDeploymentClient`.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def delete_deployment(self, name, config=None, endpoint=None):
|
|
"""
|
|
.. warning::
|
|
|
|
This method is not implemented for `DatabricksDeploymentClient`.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def list_deployments(self, endpoint=None):
|
|
"""
|
|
.. warning::
|
|
|
|
This method is not implemented for `DatabricksDeploymentClient`.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def get_deployment(self, name, endpoint=None):
|
|
"""
|
|
.. warning::
|
|
|
|
This method is not implemented for `DatabricksDeploymentClient`.
|
|
"""
|
|
raise NotImplementedError
|
|
|
|
def _call_endpoint(
|
|
self,
|
|
*,
|
|
method: str,
|
|
prefix: str = "/api/2.0",
|
|
route: str | None = None,
|
|
json_body: dict[str, Any] | None = None,
|
|
timeout: int | None = None,
|
|
retry_timeout_seconds: int | None = None,
|
|
):
|
|
"""
|
|
Args:
|
|
method: HTTP method (GET, POST, etc.).
|
|
prefix: API prefix path.
|
|
route: Endpoint route.
|
|
json_body: Request payload.
|
|
timeout: Maximum time (in seconds) for a single HTTP request.
|
|
retry_timeout_seconds: Maximum time (in seconds) for all retry attempts combined.
|
|
"""
|
|
validate_deployment_timeout_config(timeout, retry_timeout_seconds)
|
|
|
|
call_kwargs = {}
|
|
if method.lower() == "get":
|
|
call_kwargs["params"] = json_body
|
|
else:
|
|
call_kwargs["json"] = json_body
|
|
|
|
response = http_request(
|
|
host_creds=get_databricks_host_creds(self.target_uri),
|
|
endpoint=posixpath.join(prefix, "serving-endpoints", route or ""),
|
|
method=method,
|
|
timeout=MLFLOW_HTTP_REQUEST_TIMEOUT.get() if timeout is None else timeout,
|
|
retry_timeout_seconds=retry_timeout_seconds,
|
|
raise_on_status=False,
|
|
retry_codes=MLFLOW_DEPLOYMENT_CLIENT_REQUEST_RETRY_CODES,
|
|
extra_headers={"X-Databricks-Endpoints-API-Client": "Databricks Deployment Client"},
|
|
**call_kwargs,
|
|
)
|
|
augmented_raise_for_status(response)
|
|
return DatabricksEndpoint(response.json())
|
|
|
|
def _call_endpoint_stream(
|
|
self,
|
|
*,
|
|
method: str,
|
|
prefix: str = "/api/2.0",
|
|
route: str | None = None,
|
|
json_body: dict[str, Any] | None = None,
|
|
timeout: int | None = None,
|
|
retry_timeout_seconds: int | None = None,
|
|
) -> Iterator[str]:
|
|
validate_deployment_timeout_config(timeout, retry_timeout_seconds)
|
|
|
|
call_kwargs = {}
|
|
if method.lower() == "get":
|
|
call_kwargs["params"] = json_body
|
|
else:
|
|
call_kwargs["json"] = json_body
|
|
|
|
response = http_request(
|
|
host_creds=get_databricks_host_creds(self.target_uri),
|
|
endpoint=posixpath.join(prefix, "serving-endpoints", route or ""),
|
|
method=method,
|
|
timeout=MLFLOW_HTTP_REQUEST_TIMEOUT.get() if timeout is None else timeout,
|
|
retry_timeout_seconds=retry_timeout_seconds,
|
|
raise_on_status=False,
|
|
retry_codes=MLFLOW_DEPLOYMENT_CLIENT_REQUEST_RETRY_CODES,
|
|
extra_headers={"X-Databricks-Endpoints-API-Client": "Databricks Deployment Client"},
|
|
stream=True, # Receive response content in streaming way.
|
|
**call_kwargs,
|
|
)
|
|
augmented_raise_for_status(response)
|
|
|
|
# Streaming response content are composed of multiple lines.
|
|
# Each line format depends on specific endpoint
|
|
# Explicitly set the encoding to `utf-8` so the `decode_unicode` in the next line
|
|
# will decode correctly
|
|
response.encoding = "utf-8"
|
|
return (
|
|
line.strip()
|
|
for line in response.iter_lines(decode_unicode=True)
|
|
if line.strip() # filter out keep-alive new lines
|
|
)
|
|
|
|
def predict(self, deployment_name=None, inputs=None, endpoint=None):
|
|
"""
|
|
Query a serving endpoint with the provided model inputs.
|
|
See https://docs.databricks.com/api/workspace/servingendpoints/query for request/response
|
|
schema.
|
|
|
|
Args:
|
|
deployment_name: Unused.
|
|
inputs: A dictionary containing the model inputs to query.
|
|
endpoint: The name of the serving endpoint to query.
|
|
|
|
Returns:
|
|
A :py:class:`DatabricksEndpoint` object containing the query response.
|
|
|
|
Example:
|
|
|
|
.. code-block:: python
|
|
|
|
from mlflow.deployments import get_deploy_client
|
|
|
|
client = get_deploy_client("databricks")
|
|
response = client.predict(
|
|
endpoint="chat",
|
|
inputs={
|
|
"messages": [
|
|
{"role": "user", "content": "Hello!"},
|
|
],
|
|
},
|
|
)
|
|
assert response == {
|
|
"id": "chatcmpl-8OLm5kfqBAJD8CpsMANESWKpLSLXY",
|
|
"object": "chat.completion",
|
|
"created": 1700814265,
|
|
"model": "gpt-4-0613",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": "Hello! How can I assist you today?",
|
|
},
|
|
"finish_reason": "stop",
|
|
}
|
|
],
|
|
"usage": {
|
|
"prompt_tokens": 9,
|
|
"completion_tokens": 9,
|
|
"total_tokens": 18,
|
|
},
|
|
}
|
|
"""
|
|
return self._call_endpoint(
|
|
method="POST",
|
|
prefix="/",
|
|
route=posixpath.join(endpoint, "invocations"),
|
|
json_body=inputs,
|
|
timeout=MLFLOW_DEPLOYMENT_PREDICT_TIMEOUT.get(),
|
|
retry_timeout_seconds=MLFLOW_DEPLOYMENT_PREDICT_TOTAL_TIMEOUT.get(),
|
|
)
|
|
|
|
def predict_stream(
|
|
self, deployment_name=None, inputs=None, endpoint=None
|
|
) -> Iterator[dict[str, Any]]:
|
|
"""
|
|
Submit a query to a configured provider endpoint, and get streaming response
|
|
|
|
Args:
|
|
deployment_name: Unused.
|
|
inputs: The inputs to the query, as a dictionary.
|
|
endpoint: The name of the endpoint to query.
|
|
|
|
Returns:
|
|
An iterator of dictionary containing the response from the endpoint.
|
|
|
|
Example:
|
|
|
|
.. code-block:: python
|
|
|
|
from mlflow.deployments import get_deploy_client
|
|
|
|
client = get_deploy_client("databricks")
|
|
chunk_iter = client.predict_stream(
|
|
endpoint="databricks-llama-2-70b-chat",
|
|
inputs={
|
|
"messages": [{"role": "user", "content": "Hello!"}],
|
|
"temperature": 0.0,
|
|
"n": 1,
|
|
"max_tokens": 500,
|
|
},
|
|
)
|
|
for chunk in chunk_iter:
|
|
print(chunk)
|
|
# Example:
|
|
# {
|
|
# "id": "82a834f5-089d-4fc0-ad6c-db5c7d6a6129",
|
|
# "object": "chat.completion.chunk",
|
|
# "created": 1712133837,
|
|
# "model": "llama-2-70b-chat-030424",
|
|
# "choices": [
|
|
# {
|
|
# "index": 0, "delta": {"role": "assistant", "content": "Hello"},
|
|
# "finish_reason": None,
|
|
# }
|
|
# ],
|
|
# "usage": {"prompt_tokens": 11, "completion_tokens": 1, "total_tokens": 12},
|
|
# }
|
|
"""
|
|
inputs = inputs or {}
|
|
|
|
# Add stream=True param in request body to get streaming response
|
|
# See https://docs.databricks.com/api/workspace/servingendpoints/query#stream
|
|
chunk_line_iter = self._call_endpoint_stream(
|
|
method="POST",
|
|
prefix="/",
|
|
route=posixpath.join(endpoint, "invocations"),
|
|
json_body={**inputs, "stream": True},
|
|
timeout=MLFLOW_DEPLOYMENT_PREDICT_TIMEOUT.get(),
|
|
retry_timeout_seconds=MLFLOW_DEPLOYMENT_PREDICT_TOTAL_TIMEOUT.get(),
|
|
)
|
|
|
|
for line in chunk_line_iter:
|
|
splits = line.split(":", 1)
|
|
if len(splits) < 2:
|
|
raise MlflowException(
|
|
f"Unknown response format: '{line}', "
|
|
"expected 'data: <value>' for streaming response."
|
|
)
|
|
key, value = splits
|
|
if key != "data":
|
|
raise MlflowException(
|
|
f"Unknown response format with key '{key}'. "
|
|
f"Expected 'data: <value>' for streaming response, got '{line}'."
|
|
)
|
|
|
|
value = value.strip()
|
|
if value == "[DONE]":
|
|
# Databricks endpoint streaming response ends with
|
|
# a line of "data: [DONE]"
|
|
return
|
|
|
|
yield json.loads(value)
|
|
|
|
def create_endpoint(self, name=None, config=None, route_optimized=False):
|
|
"""
|
|
Create a new serving endpoint with the provided name and configuration.
|
|
|
|
See https://docs.databricks.com/api/workspace/servingendpoints/create for request/response
|
|
schema.
|
|
|
|
Args:
|
|
name: The name of the serving endpoint to create.
|
|
|
|
.. warning::
|
|
Deprecated. Include `name` in `config` instead.
|
|
|
|
config: A dictionary containing either the full API request payload
|
|
or the configuration of the serving endpoint to create.
|
|
route_optimized: A boolean which defines whether databricks serving endpoint
|
|
is optimized for routing traffic. Only used in the deprecated approach.
|
|
|
|
.. warning::
|
|
Deprecated. Include `route_optimized` in `config` instead.
|
|
|
|
Returns:
|
|
A :py:class:`DatabricksEndpoint` object containing the request response.
|
|
|
|
Example:
|
|
|
|
.. code-block:: python
|
|
|
|
from mlflow.deployments import get_deploy_client
|
|
|
|
client = get_deploy_client("databricks")
|
|
endpoint = client.create_endpoint(
|
|
config={
|
|
"name": "test",
|
|
"config": {
|
|
"served_entities": [
|
|
{
|
|
"external_model": {
|
|
"name": "gpt-4",
|
|
"provider": "openai",
|
|
"task": "llm/v1/chat",
|
|
"openai_config": {
|
|
"openai_api_key": "{{secrets/scope/key}}",
|
|
},
|
|
},
|
|
}
|
|
],
|
|
"route_optimized": True,
|
|
},
|
|
},
|
|
)
|
|
assert endpoint == {
|
|
"name": "test",
|
|
"creator": "alice@company.com",
|
|
"creation_timestamp": 0,
|
|
"last_updated_timestamp": 0,
|
|
"state": {...},
|
|
"config": {...},
|
|
"tags": [...],
|
|
"id": "88fd3f75a0d24b0380ddc40484d7a31b",
|
|
"permission_level": "CAN_MANAGE",
|
|
"route_optimized": False,
|
|
"task": "llm/v1/chat",
|
|
"endpoint_type": "EXTERNAL_MODEL",
|
|
"creator_display_name": "Alice",
|
|
"creator_kind": "User",
|
|
}
|
|
|
|
"""
|
|
warnings_list = []
|
|
|
|
if config and "config" in config:
|
|
# Using new style: full API request payload
|
|
payload = config.copy()
|
|
|
|
# Validate name conflicts
|
|
if "name" in payload:
|
|
if name is not None:
|
|
if payload["name"] == name:
|
|
warnings_list.append(
|
|
"Passing 'name' as a parameter is deprecated. "
|
|
"Please specify 'name' only within the config dictionary."
|
|
)
|
|
else:
|
|
raise MlflowException(
|
|
f"Name mismatch. Found '{name}' as parameter and '{payload['name']}' "
|
|
"in config. Please specify 'name' only within the config dictionary "
|
|
"as this parameter is deprecated."
|
|
)
|
|
else:
|
|
if name is None:
|
|
raise MlflowException(
|
|
"The 'name' field is required. Please specify it within the config "
|
|
"dictionary."
|
|
)
|
|
payload["name"] = name
|
|
warnings_list.append(
|
|
"Passing 'name' as a parameter is deprecated. "
|
|
"Please specify 'name' within the config dictionary."
|
|
)
|
|
|
|
# Validate route_optimized conflicts
|
|
if "route_optimized" in payload:
|
|
if route_optimized is not None:
|
|
if payload["route_optimized"] != route_optimized:
|
|
raise MlflowException(
|
|
"Conflicting 'route_optimized' values found. "
|
|
"Please specify 'route_optimized' only within the config dictionary "
|
|
"as this parameter is deprecated."
|
|
)
|
|
warnings_list.append(
|
|
"Passing 'route_optimized' as a parameter is deprecated. "
|
|
"Please specify 'route_optimized' only within the config dictionary."
|
|
)
|
|
else:
|
|
if route_optimized:
|
|
payload["route_optimized"] = route_optimized
|
|
warnings_list.append(
|
|
"Passing 'route_optimized' as a parameter is deprecated. "
|
|
"Please specify 'route_optimized' within the config dictionary."
|
|
)
|
|
else:
|
|
# Handle legacy format (backwards compatibility)
|
|
warnings_list.append(
|
|
"Passing 'name', 'config', and 'route_optimized' as separate parameters is "
|
|
"deprecated. Please pass the full API request payload as a single dictionary "
|
|
"in the 'config' parameter."
|
|
)
|
|
config = config.copy() if config else {} # avoid mutating config
|
|
extras = {}
|
|
for key in ("tags", "rate_limits"):
|
|
if tags := config.pop(key, None):
|
|
extras[key] = tags
|
|
payload = {"name": name, "config": config, "route_optimized": route_optimized, **extras}
|
|
|
|
if warnings_list:
|
|
warnings.warn("\n".join(warnings_list), UserWarning)
|
|
|
|
return self._call_endpoint(method="POST", json_body=payload)
|
|
|
|
@deprecated(
|
|
alternative=(
|
|
"update_endpoint_config, update_endpoint_tags, update_endpoint_rate_limits, "
|
|
"or update_endpoint_ai_gateway"
|
|
)
|
|
)
|
|
def update_endpoint(self, endpoint, config=None):
|
|
"""
|
|
Update a specified serving endpoint with the provided configuration.
|
|
See https://docs.databricks.com/api/workspace/servingendpoints/updateconfig for
|
|
request/response schema.
|
|
|
|
Args:
|
|
endpoint: The name of the serving endpoint to update.
|
|
config: A dictionary containing the configuration of the serving endpoint to update.
|
|
|
|
Returns:
|
|
A :py:class:`DatabricksEndpoint` object containing the request response.
|
|
|
|
Example:
|
|
|
|
.. code-block:: python
|
|
|
|
from mlflow.deployments import get_deploy_client
|
|
|
|
client = get_deploy_client("databricks")
|
|
endpoint = client.update_endpoint(
|
|
endpoint="chat",
|
|
config={
|
|
"served_entities": [
|
|
{
|
|
"name": "test",
|
|
"external_model": {
|
|
"name": "gpt-4",
|
|
"provider": "openai",
|
|
"task": "llm/v1/chat",
|
|
"openai_config": {
|
|
"openai_api_key": "{{secrets/scope/key}}",
|
|
},
|
|
},
|
|
}
|
|
],
|
|
},
|
|
)
|
|
assert endpoint == {
|
|
"name": "chat",
|
|
"creator": "alice@company.com",
|
|
"creation_timestamp": 0,
|
|
"last_updated_timestamp": 0,
|
|
"state": {...},
|
|
"config": {...},
|
|
"tags": [...],
|
|
"id": "88fd3f75a0d24b0380ddc40484d7a31b",
|
|
}
|
|
|
|
rate_limits = client.update_endpoint(
|
|
endpoint="chat",
|
|
config={
|
|
"rate_limits": [
|
|
{
|
|
"key": "user",
|
|
"renewal_period": "minute",
|
|
"calls": 10,
|
|
}
|
|
],
|
|
},
|
|
)
|
|
assert rate_limits == {
|
|
"rate_limits": [
|
|
{
|
|
"key": "user",
|
|
"renewal_period": "minute",
|
|
"calls": 10,
|
|
}
|
|
],
|
|
}
|
|
"""
|
|
warnings.warn(
|
|
"The `update_endpoint` method is deprecated. Use the specific update methods—"
|
|
"`update_endpoint_config`, `update_endpoint_tags`, `update_endpoint_rate_limits`, "
|
|
"`update_endpoint_ai_gateway`—instead.",
|
|
UserWarning,
|
|
)
|
|
|
|
if list(config) == ["rate_limits"]:
|
|
return self._call_endpoint(
|
|
method="PUT", route=posixpath.join(endpoint, "rate-limits"), json_body=config
|
|
)
|
|
else:
|
|
return self._call_endpoint(
|
|
method="PUT", route=posixpath.join(endpoint, "config"), json_body=config
|
|
)
|
|
|
|
def update_endpoint_config(self, endpoint, config):
|
|
"""
|
|
Update the configuration of a specified serving endpoint. See
|
|
https://docs.databricks.com/api/workspace/servingendpoints/updateconfig for request/response
|
|
request/response schema.
|
|
|
|
Args:
|
|
endpoint: The name of the serving endpoint to update.
|
|
config: A dictionary containing the configuration of the serving endpoint to update.
|
|
|
|
Returns:
|
|
A :py:class:`DatabricksEndpoint` object containing the request response.
|
|
|
|
Example:
|
|
|
|
.. code-block:: python
|
|
|
|
from mlflow.deployments import get_deploy_client
|
|
|
|
client = get_deploy_client("databricks")
|
|
updated_endpoint = client.update_endpoint_config(
|
|
endpoint="test",
|
|
config={
|
|
"served_entities": [
|
|
{
|
|
"name": "gpt-4o-mini",
|
|
"external_model": {
|
|
"name": "gpt-4o-mini",
|
|
"provider": "openai",
|
|
"task": "llm/v1/chat",
|
|
"openai_config": {
|
|
"openai_api_key": "{{secrets/scope/key}}",
|
|
},
|
|
},
|
|
}
|
|
]
|
|
},
|
|
)
|
|
assert updated_endpoint == {
|
|
"name": "test",
|
|
"creator": "alice@company.com",
|
|
"creation_timestamp": 1729527763000,
|
|
"last_updated_timestamp": 1729530896000,
|
|
"state": {"ready": "READY", "config_update": "NOT_UPDATING"},
|
|
"config": {...},
|
|
"id": "44b258fb39804564b37603d8d14b853e",
|
|
"permission_level": "CAN_MANAGE",
|
|
"route_optimized": False,
|
|
"task": "llm/v1/chat",
|
|
"endpoint_type": "EXTERNAL_MODEL",
|
|
"creator_display_name": "Alice",
|
|
"creator_kind": "User",
|
|
}
|
|
"""
|
|
|
|
return self._call_endpoint(
|
|
method="PUT", route=posixpath.join(endpoint, "config"), json_body=config
|
|
)
|
|
|
|
def update_endpoint_tags(self, endpoint, config):
|
|
"""
|
|
Update the tags of a specified serving endpoint. See
|
|
https://docs.databricks.com/api/workspace/servingendpoints/patch for request/response
|
|
schema.
|
|
|
|
Args:
|
|
endpoint: The name of the serving endpoint to update.
|
|
config: A dictionary containing tags to add and/or remove.
|
|
|
|
Returns:
|
|
A :py:class:`DatabricksEndpoint` object containing the request response.
|
|
|
|
Example:
|
|
|
|
.. code-block:: python
|
|
|
|
from mlflow.deployments import get_deploy_client
|
|
|
|
client = get_deploy_client("databricks")
|
|
updated_tags = client.update_endpoint_tags(
|
|
endpoint="test", config={"add_tags": [{"key": "project", "value": "test"}]}
|
|
)
|
|
assert updated_tags == {"tags": [{"key": "project", "value": "test"}]}
|
|
"""
|
|
return self._call_endpoint(
|
|
method="PATCH", route=posixpath.join(endpoint, "tags"), json_body=config
|
|
)
|
|
|
|
def update_endpoint_rate_limits(self, endpoint, config):
|
|
"""
|
|
Update the rate limits of a specified serving endpoint.
|
|
See https://docs.databricks.com/api/workspace/servingendpoints/put for request/response
|
|
schema.
|
|
|
|
Args:
|
|
endpoint: The name of the serving endpoint to update.
|
|
config: A dictionary containing the updated rate limit configuration.
|
|
|
|
Returns:
|
|
A :py:class:`DatabricksEndpoint` object containing the updated rate limits.
|
|
|
|
Example:
|
|
|
|
.. code-block:: python
|
|
|
|
from mlflow.deployments import get_deploy_client
|
|
|
|
client = get_deploy_client("databricks")
|
|
name = "databricks-dbrx-instruct"
|
|
rate_limits = {
|
|
"rate_limits": [{"calls": 10, "key": "endpoint", "renewal_period": "minute"}]
|
|
}
|
|
updated_rate_limits = client.update_endpoint_rate_limits(
|
|
endpoint=name, config=rate_limits
|
|
)
|
|
assert updated_rate_limits == {
|
|
"rate_limits": [{"calls": 10, "key": "endpoint", "renewal_period": "minute"}]
|
|
}
|
|
"""
|
|
return self._call_endpoint(
|
|
method="PUT", route=posixpath.join(endpoint, "rate-limits"), json_body=config
|
|
)
|
|
|
|
def update_endpoint_ai_gateway(self, endpoint, config):
|
|
"""
|
|
Update the AI Gateway configuration of a specified serving endpoint.
|
|
|
|
Args:
|
|
endpoint (str): The name of the serving endpoint to update.
|
|
config (dict): A dictionary containing the AI Gateway configuration to update.
|
|
|
|
Returns:
|
|
dict: A dictionary containing the updated AI Gateway configuration.
|
|
|
|
Example:
|
|
|
|
.. code-block:: python
|
|
|
|
from mlflow.deployments import get_deploy_client
|
|
|
|
client = get_deploy_client("databricks")
|
|
name = "test"
|
|
|
|
gateway_config = {
|
|
"usage_tracking_config": {"enabled": True},
|
|
"inference_table_config": {
|
|
"enabled": True,
|
|
"catalog_name": "my_catalog",
|
|
"schema_name": "my_schema",
|
|
},
|
|
}
|
|
|
|
updated_gateway = client.update_endpoint_ai_gateway(
|
|
endpoint=name, config=gateway_config
|
|
)
|
|
assert updated_gateway == {
|
|
"usage_tracking_config": {"enabled": True},
|
|
"inference_table_config": {
|
|
"catalog_name": "my_catalog",
|
|
"schema_name": "my_schema",
|
|
"table_name_prefix": "test",
|
|
"enabled": True,
|
|
},
|
|
}
|
|
"""
|
|
return self._call_endpoint(
|
|
method="PUT", route=posixpath.join(endpoint, "ai-gateway"), json_body=config
|
|
)
|
|
|
|
def delete_endpoint(self, endpoint):
|
|
"""
|
|
Delete a specified serving endpoint.
|
|
See https://docs.databricks.com/api/workspace/servingendpoints/delete for request/response
|
|
schema.
|
|
|
|
Args:
|
|
endpoint: The name of the serving endpoint to delete.
|
|
|
|
Returns:
|
|
A DatabricksEndpoint object containing the request response.
|
|
|
|
Example:
|
|
|
|
.. code-block:: python
|
|
|
|
from mlflow.deployments import get_deploy_client
|
|
|
|
client = get_deploy_client("databricks")
|
|
client.delete_endpoint(endpoint="chat")
|
|
"""
|
|
return self._call_endpoint(method="DELETE", route=endpoint)
|
|
|
|
def list_endpoints(self):
|
|
"""
|
|
Retrieve all serving endpoints.
|
|
|
|
See https://docs.databricks.com/api/workspace/servingendpoints/list for request/response
|
|
schema.
|
|
|
|
Returns:
|
|
A list of :py:class:`DatabricksEndpoint` objects containing the request response.
|
|
|
|
Example:
|
|
|
|
.. code-block:: python
|
|
|
|
from mlflow.deployments import get_deploy_client
|
|
|
|
client = get_deploy_client("databricks")
|
|
endpoints = client.list_endpoints()
|
|
assert endpoints == [
|
|
{
|
|
"name": "chat",
|
|
"creator": "alice@company.com",
|
|
"creation_timestamp": 0,
|
|
"last_updated_timestamp": 0,
|
|
"state": {...},
|
|
"config": {...},
|
|
"tags": [...],
|
|
"id": "88fd3f75a0d24b0380ddc40484d7a31b",
|
|
},
|
|
]
|
|
|
|
"""
|
|
return self._call_endpoint(method="GET").endpoints
|
|
|
|
def get_endpoint(self, endpoint):
|
|
"""
|
|
Get a specified serving endpoint.
|
|
See https://docs.databricks.com/api/workspace/servingendpoints/get for request/response
|
|
schema.
|
|
|
|
Args:
|
|
endpoint: The name of the serving endpoint to get.
|
|
|
|
Returns:
|
|
A DatabricksEndpoint object containing the request response.
|
|
|
|
Example:
|
|
|
|
.. code-block:: python
|
|
|
|
from mlflow.deployments import get_deploy_client
|
|
|
|
client = get_deploy_client("databricks")
|
|
endpoint = client.get_endpoint(endpoint="chat")
|
|
assert endpoint == {
|
|
"name": "chat",
|
|
"creator": "alice@company.com",
|
|
"creation_timestamp": 0,
|
|
"last_updated_timestamp": 0,
|
|
"state": {...},
|
|
"config": {...},
|
|
"tags": [...],
|
|
"id": "88fd3f75a0d24b0380ddc40484d7a31b",
|
|
}
|
|
"""
|
|
return self._call_endpoint(method="GET", route=endpoint)
|
|
|
|
|
|
def run_local(name, model_uri, flavor=None, config=None):
|
|
pass
|
|
|
|
|
|
def target_help():
|
|
pass
|