Support model evaluation with an endpoint URL

[B-Step62]

🛠 DevTools 🛠

Install mlflow from this PR

pip install git+https://github.com/mlflow/mlflow.git@refs/pull/11262/merge

Checkout with GitHub CLI

gh pr checkout 11262

Related Issues/PRs

#xxx

What changes are proposed in this pull request?

This change enhances mlflow.evaluate() function to accepts an endpoint URL to model argument. We start to get frequent requests for this from users who want to evaluate the proprietary model endpoint or Databricks served models. Without this feature, they need to implement a custom predict function, which is not super difficult but less convenient.

Important notes

• We only support the OpenAI-compatible chat/completion models, to avoid asking too many configurations to users.
• The input dataset should contain a single text input and MLflow constructs chat/completion request from it. This limits some flexibility for input e.g. system prompt for chat model, but is consistent with other flavors like OpenAI and Transformers. The more complex use case should be handled by a custom function.
• Inference params is not supported in current API, but they are indispensable for chat/completion models. This PR adds a new argument inference_params to the API, but it is only used for endpoint evaluation, not for other type of models. We can do follow-up to support inference params for all model types.

Example

eval_data = pd.DataFrame(
    {
        # Input data must be a string column and named "inputs".
        "inputs": [
            "Write 3 reasons why you should use MLflow?",
            "Can you explain the difference between classification and regression?",
        ],
    }
)


with mlflow.start_run() as run:
    results = mlflow.evaluate(
        # URL of the OpenAI-compatible chat/completion endpoint
        model="https://example.com/serving-endpoints/databricks-mpt-30b-instruct/invocations",
        data=eval_data,
        # New argument for support inference params
        inference_parameters={"max_tokens": 100, "temperature": 0.0},
        # New argument for passing auth info
        headers={"Authorization": "Bearer YOUR API KEY"},
        # By specifying the "completion" model type, MOLflow automatically format the input text to the completion request format,
        # and extract the generated text from the completion response.
        model_type="completion",
    )

How is this PR tested?

• Existing unit/integration tests
• New unit/integration tests
• Manual tests (WIP)

Tested in Databricks. Sample code for one test case:

import mlflow
import pandas as pd

eval_data = pd.DataFrame(
    {
        "inputs": [
            "What is MLflow?",
            "What is Spark?",
        ],
        "ground_truth": [
            "MLflow is an open-source platform ....",
            "Apache Spark is an open-source, ...",
        ],
    }
)


with mlflow.start_run(run_name="mlflow-eval-chat") as run:
    results = mlflow.evaluate(
        # Pass the URL of an OpenAI-compatible chat endpoint
        model="https://${{DB_HOST}}/serving-endpoints/databricks-mixtral-8x7b-instruct/invocations",
        data=eval_data,
        targets="ground_truth",
        model_type="chat",
        inference_params={"max_tokens": 100, "temperature": 0.0},
        headers={"Authentication": "Bearer ${{DB_TOKEN}}"},
        extra_metrics=[mlflow.metrics.exact_match()],
    )

Does this PR require documentation update?

• No. You can skip the rest of this section.
• Yes. I've updated:
  • Examples
  • API references
  • Instructions

Release Notes

Is this a user-facing change?

• No. You can skip the rest of this section.
• Yes. Give a description of this change to be included in the release notes for MLflow users.

The mlflow.evaluate() API supports a model endpoint URL as a evaluation subject. This is particularly useful when you want to evaluate deployed models or proprietary APIs such as Databricks Foundation Model APIs.

What component(s), interfaces, languages, and integrations does this PR affect?

Components

• area/artifacts: Artifact stores and artifact logging
• area/build: Build and test infrastructure for MLflow
• area/deployments: MLflow Deployments client APIs, server, and third-party Deployments integrations
• area/docs: MLflow documentation pages
• area/examples: Example code
• area/model-registry: Model Registry service, APIs, and the fluent client calls for Model Registry
• area/models: MLmodel format, model serialization/deserialization, flavors
• area/recipes: Recipes, Recipe APIs, Recipe configs, Recipe Templates
• area/projects: MLproject format, project running backends
• area/scoring: MLflow Model server, model deployment tools, Spark UDFs
• area/server-infra: MLflow Tracking server backend
• area/tracking: Tracking Service, tracking client APIs, autologging

Interface

• area/uiux: Front-end, user experience, plotting, JavaScript, JavaScript dev server
• area/docker: Docker use across MLflow's components, such as MLflow Projects and MLflow Models
• area/sqlalchemy: Use of SQLAlchemy in the Tracking Service or Model Registry
• area/windows: Windows support

Language

• language/r: R APIs and clients
• language/java: Java APIs and clients
• language/new: Proposals for new client languages

Integrations

• integrations/azure: Azure and Azure ML integrations
• integrations/sagemaker: SageMaker integrations
• integrations/databricks: Databricks integrations

How should the PR be classified in the release notes? Choose one:

• rn/none - No description will be included. The PR will be mentioned only by the PR number in the "Small Bugfixes and Documentation Updates" section
• rn/breaking-change - The PR will be mentioned in the "Breaking Changes" section
• rn/feature - A new user-facing feature worth mentioning in the release notes
• rn/bug-fix - A user-facing bug fix worth mentioning in the release notes
• rn/documentation - A user-facing documentation change worth mentioning in the release notes

[github-actions]

Documentation preview for 20a9741 will be available when this CircleCI job
completes successfully.

• Top page
• Changed pages

More info

• Ignore this comment if this PR does not change the documentation.
• It takes a few minutes for the preview to be available.
• The preview is updated when a new commit is pushed to this PR.
• This comment was created by https://github.com/mlflow/mlflow/actions/runs/8106426949.

[B-Step62]

note: This hashing support is required for passing a list of dictionary (chat/completion payload) or similar pandas DataFrame as a input dataset. The hashing is required for generating dataset digest, but it doesn't handle a list or array contains dictionary. The implementation is not super clean, just band-aiding without breaking existing use cases.

[B-Step62]

note: This is small but important change. In the original logic, targets param is always wrapped by numpy array when the input data is list, then if targets is None doesn't work in downstream. Concretely, when we pass a list as an input data and leave targets None, it will raise a weird error.

This is not specific to the endpoint evaluation, however, implicitly requiring targets only for list input is confusing, so I think we should fix it in this PR.

[dbczumar]

Can we move all of these hashing methods into a separate module, ideally located in mlflow.data? since these are all about hashing for use with MLflow Datasets

[B-Step62]

Agree this should be in a separate module, but for the location, I would still put this under evaluation module. These methods are used for hashing various input like list, numpy, dataframe, and only used for class EvaluationDataset digest.

Now that the digesting method for each flavor e.g. TensorflowDataset is moved to respective module (ref), so I think we can have evaluation/dataset.py and put all these methods. WDYT?

[dbczumar]

I would prefer to put all of EvaluationDataset in mlflow/data to maintain consistency with other datasets - e.g. we don't put TensorflowDataset in mlflow/tensorflow

[dbczumar]

Feel free to tackle this in a follow-up ticket as long as we track it :)

[B-Step62]

I would prefer to put all of EvaluationDataset in mlflow/data

Makes sense.

Feel free to tackle this in a follow-up ticket as long as we track it :)

Yeah let me do that, created a JIRA in next sprint (ML-39039), but may be able to pick up if no fun surprise happens in other tasks:)

[harupy]

LGTM!

[dbczumar]

LGTM with docs nits! Thanks @B-Step62