mlflow/mlflow · error · NotImplementedError

Dataset association operations are not available in Databric

Error message

Dataset association operations are not available in Databricks yet. Associations are managed through Unity Catalog.

What it means

_validate_association_operation rejects dataset<->experiment association operations (add_dataset_to_experiments / remove_dataset_from_experiments) on Databricks tracking URIs. Associations there are managed through Unity Catalog, so the MLflow-level API intentionally raises NotImplementedError.

Source

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

    """
    if is_databricks_uri(get_tracking_uri()):
        raise NotImplementedError(
            "Dataset tag operations are not available in Databricks yet. "
            "Tags are managed through Unity Catalog."
        )

    from mlflow.tracking.client import MlflowClient

    MlflowClient().delete_dataset_tag(dataset_id, key)


def _validate_association_operation():
    """Validate that dataset association operations can be performed."""
    from mlflow.store.tracking.file_store import FileStore
    from mlflow.tracking._tracking_service.utils import _get_store

    if is_databricks_uri(get_tracking_uri()):
        raise NotImplementedError(
            "Dataset association operations are not available in Databricks yet. "
            "Associations are managed through Unity Catalog."
        )

    store = _get_store()
    if isinstance(store, FileStore):
        raise NotImplementedError(
            "Dataset association operations are not supported with FileStore backend. "
            "Please use a database-backed tracking store."
        )


def add_dataset_to_experiments(dataset_id: str, experiment_ids: list[str]) -> "EvaluationDataset":
    """
    Add a dataset to additional experiments.

    This allows reusing datasets across multiple experiments for evaluation purposes.

View on GitHub (pinned to 6a27f2decc)

Solutions

  1. Associate datasets to experiments via Unity Catalog-managed mechanisms
  2. Run association operations against an OSS database-backed tracking server
  3. Branch on is_databricks_uri() to choose the right path

Example fix

// before
add_dataset_to_experiments(dataset_id=ds_id, experiment_ids=[exp_id])  # on Databricks
// after
if not is_databricks_uri(mlflow.get_tracking_uri()):
    add_dataset_to_experiments(dataset_id=ds_id, experiment_ids=[exp_id])
else:
    # manage associations via Unity Catalog
Defensive patterns

Strategy: validation

Validate before calling

assert not is_databricks_uri(get_tracking_uri()), "dataset associations must be managed via Unity Catalog on Databricks"

Try / catch

try:
    mlflow.genai.datasets.add_dataset_to_experiments(ds_id, exp_ids)
except NotImplementedError:
    associate_via_uc(ds_id, exp_ids)

Prevention

When it happens

Trigger: Calling mlflow.genai.datasets.add_dataset_to_experiments(...) or remove_dataset_from_experiments(...) while the tracking URI is databricks://.

Common situations: Wiring evaluation datasets to experiments inside a Databricks workspace; migrating OSS association scripts to Databricks.

Related errors


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