mlflow/mlflow · error · NotImplementedError

`version` is only supported for Databricks datasets.

Error message

`version` is only supported for Databricks datasets.

What it means

The `version` parameter of mlflow.genai.datasets.get_dataset is only meaningful for Databricks-hosted evaluation datasets. On a non-Databricks (OSS) tracking URI, passing version != None raises NotImplementedError because OSS MLflowClient datasets have no version concept.

Source

Thrown at mlflow/genai/datasets/__init__.py:395

            ]
            dataset.merge_records(new_test_cases)
    """

    if is_databricks_uri(get_tracking_uri()):
        _validate_databricks_params(name, dataset_id)
        resolved_version = _resolve_dataset_version_arg(version)
        try:
            from databricks.agents.datasets import get_dataset as db_get

            with _databricks_profile_env():
                if version is not None:
                    return EvaluationDataset(db_get(name, version=resolved_version))
                return EvaluationDataset(db_get(name))
        except ImportError as e:
            raise ImportError(_ERROR_MSG) from e
    else:
        if version is not None:
            raise NotImplementedError("`version` is only supported for Databricks datasets.")
        _validate_non_databricks_get_params(name, dataset_id)

        if name is not None:
            return EvaluationDataset(_get_dataset_by_name(name))

        return EvaluationDataset(MlflowClient().get_dataset(dataset_id))


def search_datasets(
    experiment_ids: str | list[str] | None = None,
    filter_string: str | None = None,
    max_results: int | None = None,
    order_by: list[str] | None = None,
) -> list[EvaluationDataset]:
    """
    Search for datasets.

    .. warning::

View on GitHub (pinned to 6a27f2decc)

Solutions

Drop the `version` argument when targeting an OSS MLflow backend
Only pass version when is_databricks_uri(mlflow.get_tracking_uri()) is True
Fetch by dataset_id (which implies a specific dataset) instead of name+version

Example fix

// before
get_dataset(name="eval", version=2)  # on OSS server
// after
if is_databricks_uri(mlflow.get_tracking_uri()):
    get_dataset(name="eval", version=2)
else:
    get_dataset(name="eval")

Defensive patterns

Strategy: validation

Validate before calling

from mlflow.tracking._tracking_service.utils import is_databricks_uri
from mlflow.tracking.fluent import get_tracking_uri
assert version is None or is_databricks_uri(get_tracking_uri()), "version only valid on Databricks"

Type guard

def supports_version(version) -> bool:
    return version is None or is_databricks_uri(get_tracking_uri())

Try / catch

try:
    ds = mlflow.genai.datasets.get_dataset(name="eval", version=ver)
except NotImplementedError:
    ds = mlflow.genai.datasets.get_dataset(name="eval")

Prevention

Gate Databricks-only kwargs behind backend checks
Avoid copying version= from Databricks notebooks to OSS code

When it happens

Trigger: Calling get_dataset(name="x", version=2) while the tracking URI is a local file store, SQLite, or HTTP server (not databricks://).

Common situations: Porting Databricks notebook code to a self-hosted MLflow server; shared helper that always passes version.

Related errors

AI-assisted analysis of mlflow/mlflow@6a27f2decc (2026-08-29). Data as JSON: /api/errors/5a7ac6264cb3c394. Report an issue: GitHub.