Model Serving Endpoints and Traffic Splitting

What it is

Mosaic AI Model Serving puts a registered MLflow model behind a REST API for real time inference. You create an endpoint, attach one or more served entities (model versions), and callers send JSON requests that return predictions in milliseconds.

The exam asks how a custom model gets onto an endpoint, how the endpoint is queried, and how requests are divided between two served models, for example a champion and a challenger.

How it works

Deploying a custom model

Model Serving accepts any model logged in MLflow format, either a built in flavor (scikit-learn, XGBoost, PyTorch) or pyfunc for arbitrary Python code. A custom model subclasses mlflow.pyfunc.PythonModel and implements predict.

import mlflow
class ScoreTo100(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input):
        return model_input["score"] * 100
mlflow.set_registry_uri("databricks-uc")
mlflow.pyfunc.log_model("model", python_model=ScoreTo100(),
    input_example=sample_pdf, registered_model_name="main.ml.score_to_100")

Logging to Unity Catalog requires a model signature; passing an input example lets MLflow infer it. Custom logic such as pre processing, post processing or calling several models can all live inside predict.

Creating the endpoint

You can create an endpoint from the Serving page in the UI, the REST API or the MLflow Deployments client.

from mlflow.deployments import get_deploy_client
client = get_deploy_client("databricks")
client.create_endpoint(name="score-endpoint", config={"served_entities": [{
    "entity_name": "main.ml.score_to_100", "entity_version": "1",
    "workload_size": "Small", "scale_to_zero_enabled": True}]})

Scale to zero saves cost on development endpoints. Databricks advises against it for production endpoints that need consistent response times.

Querying for real time inference

Every endpoint exposes one URL: POST https://<workspace>/serving-endpoints/<endpoint-name>/invocations, authenticated with a Databricks token. For tabular models, the body holds a pandas style DataFrame.

resp = client.predict(endpoint="score-endpoint",
    inputs={"dataframe_records": [{"score": 0.42}, {"score": 0.91}]})

The same request can be sent with any HTTP client, and SQL users can call an endpoint with the ai_query function.

Splitting traffic between served models

One endpoint can hold several served entities. Its traffic_config lists routes, each naming a served entity and a traffic_percentage; the percentages sum to 100 and each request goes to exactly one entity. Each served entity has a name, and the route refers to it by that name.

client.update_endpoint_config(endpoint="score-endpoint", config={
    "served_entities": [champion_entity, challenger_entity],
    "traffic_config": {"routes": [
        {"served_model_name": "champion", "traffic_percentage": 90},
        {"served_model_name": "challenger", "traffic_percentage": 10}]}})

To test one entity directly, call /serving-endpoints/<endpoint-name>/served-models/<served-model-name>/invocations; traffic settings are ignored on that path. All served entities on one endpoint must be the same model type.

Worked example

An endpoint routes 90 percent to the champion and 10 percent to the challenger and receives 50,000 requests a day. About 45,000 go to the champion and 5,000 to the challenger, enough to compare live metrics before moving the challenger to 50 percent and then 100 percent.

Changing the split is a config update on the same endpoint, so callers keep using the same URL throughout.

How the exam asks about this

• "A model needs custom pre processing before predicting and must be served." Wrap it in an mlflow.pyfunc.PythonModel with predict, log and register it, then serve it.
• "A web app needs a prediction per user request in milliseconds." Deploy to a Model Serving endpoint and POST to its invocations URL.
• "Send 20 percent of live requests to a new version without changing the client." Add the version as a second served entity and set traffic_config routes to 80 and 20.