Skip to content

Integrate mlflow tracking for sklearn based tasks#151

Open
lorenzomammana wants to merge 7 commits intomainfrom
feature/sklearn-mlflow-tracking
Open

Integrate mlflow tracking for sklearn based tasks#151
lorenzomammana wants to merge 7 commits intomainfrom
feature/sklearn-mlflow-tracking

Conversation

@lorenzomammana
Copy link
Contributor

@lorenzomammana lorenzomammana commented Feb 3, 2026

Summary

### [2.7.2]

#### Added

- MLflow tracking for sklearn classification and patch classification tasks (metrics, artifacts, and model uploads)

#### Updated

- Minor UI improvements to segmentation datamodule, reduce OOM risk in segmentation test when running with thousands of images.

Type of Change

  • New feature (non-breaking change that adds functionality)

Checklist

Please confirm that the following tasks have been completed:

  • I have tested my changes locally and they work as expected. (Please describe the tests you performed.)
  • I have added unit tests for my changes, or updated existing tests if necessary.
  • I have updated the documentation, if applicable.
  • I have installed pre-commit and run locally for my code changes.

Copy link

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull request overview

This pull request integrates MLflow tracking capabilities for sklearn-based classification tasks, extending existing MLflow support beyond PyTorch Lightning workflows.

Changes:

  • Refactored MLflow artifact uploading logic from utils.py into a dedicated mlflow.py module
  • Implemented SklearnMLflowClient and mixin classes for sklearn task tracking
  • Added MLflow logging hooks to sklearn classification and patch classification tasks

Reviewed changes

Copilot reviewed 9 out of 9 changed files in this pull request and generated 1 comment.

Show a summary per file
File Description
quadra/utils/utils.py Removed deprecated finish() function and related imports after moving MLflow logic to dedicated module
quadra/utils/mlflow.py Added sklearn MLflow client, mixins, and helper functions; refactored Lightning artifact upload logic
quadra/tasks/base.py Updated to call new upload_lightning_artifacts() instead of deprecated finish()
quadra/tasks/classification.py Integrated SklearnClassificationMLflowMixin and added MLflow logging lifecycle hooks
quadra/tasks/patch.py Integrated SklearnPatchMLflowMixin and added MLflow logging lifecycle hooks
quadra/tasks/segmentation.py Added .cpu() calls to prevent GPU memory accumulation during test data collection
quadra/datamodules/segmentation.py Added progress bar for split file reading; corrected docstring default value
quadra/configs/experiment/base/classification/sklearn_classification.yaml Added MLflow logger configuration and artifact upload settings
quadra/configs/experiment/base/classification/sklearn_classification_patch.yaml Added MLflow logger configuration and artifact upload settings

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

"""Upload config files to MLFlow server.

Args:
config: Configuration composed by Hydra..
Copy link

Copilot AI Feb 3, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Corrected double period '..' to single period '.' in docstring

Suggested change
config: Configuration composed by Hydra..
config: Configuration composed by Hydra.

Copilot uses AI. Check for mistakes.
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.

1 participant