mlflow/mlflow · error · ImportError

The `databricks-agents` package is required to use `mlflow.g

Error message

The `databricks-agents` package is required to use `mlflow.genai.datasets`. Please install it with `pip install databricks-agents`.

What it means

create_dataset delegates to the databricks.agents.datasets package on Databricks. If that optional dependency is not installed, the ImportError is re-raised as an ImportError telling you to install databricks-agents. The genai.datasets module is unusable on Databricks without it.

Source

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

    """
    if name is None:
        raise ValueError("Parameter 'name' is required.")

    experiment_ids = [experiment_id] if isinstance(experiment_id, str) else experiment_id

    if is_databricks_uri(get_tracking_uri()):
        if tags is not None:
            raise NotImplementedError(
                "Tags are not supported in Databricks environments. "
                "Tags are managed through Unity Catalog."
            )
        try:
            from databricks.agents.datasets import create_dataset as db_create

            with _databricks_profile_env():
                return EvaluationDataset(db_create(name, experiment_ids))
        except ImportError as e:
            raise ImportError(_ERROR_MSG) from e
    else:
        from mlflow.tracking.client import MlflowClient

        if experiment_ids is None:
            from mlflow.tracking.fluent import _get_experiment_id

            current_exp_id = _get_experiment_id()
            if current_exp_id:
                experiment_ids = [current_exp_id]

        mlflow_dataset = MlflowClient().create_dataset(
            name=name,
            experiment_id=experiment_ids,
            tags=tags,
        )
        return EvaluationDataset(mlflow_dataset)

View on GitHub (pinned to 6a27f2decc)

Solutions

Run `pip install databricks-agents` in the active environment
Pin databricks-agents in requirements and rebuild the deployment image
Verify import works: `from databricks.agents.datasets import create_dataset`

Example fix

// before
mlflow.genai.datasets.create_dataset(name="eval")  # ImportError
// after
# shell: pip install databricks-agents
mlflow.genai.datasets.create_dataset(name="eval")

Defensive patterns

Strategy: validation

Validate before calling

try:
    import databricks.agents  # noqa
except ImportError:
    raise RuntimeError("pip install databricks-agents required for mlflow.genai.datasets on Databricks")

Try / catch

try:
    ds = mlflow.genai.datasets.create_dataset(name="eval")
except ImportError as e:
    logger.error("install databricks-agents: %s", e)
    raise

Prevention

Include databricks-agents in requirements for Databricks-targeting code
Smoke-test imports in CI before running genai dataset jobs
Use uv/pip extras so the dependency travels with the package

When it happens

Trigger: Calling mlflow.genai.datasets.create_dataset() with a databricks:// tracking URI in an environment where `pip install databricks-agents` was never run.

Common situations: Fresh CI runners or local environments; deploying code to a container image that only includes mlflow core; sklearn-style skinny installs without extras.

Related errors

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