Skip to content

feat: upgrade to MLflow 3.x and simplify validation#223

Open
veenaramesh wants to merge 4 commits into
databricks:testfrom
veenaramesh:feat/upgrade-mlflow-v3-clean
Open

feat: upgrade to MLflow 3.x and simplify validation#223
veenaramesh wants to merge 4 commits into
databricks:testfrom
veenaramesh:feat/upgrade-mlflow-v3-clean

Conversation

@veenaramesh

Copy link
Copy Markdown
Collaborator

Summary

  • Bump mlflow>=3.1, unpin pandas and numpy
  • Update imports: from mlflow import MlflowClient (was mlflow.tracking)
  • Remove registry_uri="databricks-uc" and set_registry_uri calls (not needed in MLflow 3.x)
  • Fix empty run_id bug in ModelValidation (MLflow 3.x + UC model versions may have empty run_id)
  • Simplify validation: direct from validation import instead of importlib/getattr
  • Remove test_mlp.py (mlflow.recipes removed in MLflow 3.x)
  • Clean up CI workflow secrets

Depends on #219 (simplify-project-structure).

Test plan

  • Model training pipeline (Train, ModelValidation, ModelDeployment) passes on e2-demo-field-eng
  • Batch inference passes
  • ModelValidation handles empty run_id gracefully

Move all Databricks notebooks out of nested notebooks/ subdirectories
to sit alongside their sibling helper modules. This eliminates the
%cd .. hacks that were needed to import from parent directories and
simplifies the project layout for a data science workflow.

Moves:
- training/notebooks/*.py → training/*.py
- validation/notebooks/ModelValidation.py → validation/ModelValidation.py
- deployment/*/notebooks/*.py → deployment/*/*.py
- feature_engineering/notebooks/*.py → feature_engineering/*.py
- monitoring/notebooks/*.py → monitoring/*.py

Updates all resource YAMLs, update_layout.tmpl, and README references.

Co-authored-by: Isaac
Databricks bundle deployments set the notebook CWD correctly, so
the notebook_path/%cd $notebook_path pattern is unnecessary. Sibling
imports work naturally after the flatten-notebooks change.

Also inlines the trivial get_deployed_model_stage_for_env helper
(a dict lookup) to eliminate cross-directory sys.path.append in
deploy.py and BatchInference.py.

Co-authored-by: Isaac
No longer needed with databricks bundle init.

Co-authored-by: Isaac
- Bump mlflow to >=3.1, unpin pandas and numpy
- Update MlflowClient imports (from mlflow import MlflowClient)
- Remove registry_uri="databricks-uc" and set_registry_uri calls
  (no longer needed in MLflow 3.x with Unity Catalog)
- Fix empty run_id bug in ModelValidation (MLflow 3.x UC model
  versions may have empty run_id)
- Replace importlib/getattr dynamic loading of validation config
  with direct imports from validation.py
- Remove test_mlp.py (mlflow.recipes removed in MLflow 3.x)
- Clean up CI workflow secrets

Co-authored-by: Isaac
@veenaramesh veenaramesh force-pushed the feat/upgrade-mlflow-v3-clean branch from dab6937 to 19c7456 Compare June 12, 2026 17:08
@veenaramesh veenaramesh linked an issue Jun 12, 2026 that may be closed by this pull request
5 tasks
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Update MLOps Stacks to support MLflow 3.0

1 participant