From 9a6ca521483e83bdfb053c84297cae3e9849ce47 Mon Sep 17 00:00:00 2001 From: luandalmazo Date: Fri, 26 Jun 2026 14:52:11 -0300 Subject: [PATCH] add validate_model --- datamint/__init__.py | 5 +- datamint/mlflow/flavors/__init__.py | 5 + datamint/mlflow/flavors/validation.py | 228 ++++++++++++++++++ .../05_deployment/03_validate_model.ipynb | 207 ++++++++++++++++ .../slice_based/02_busi_segmentation.ipynb | 2 +- 5 files changed, 445 insertions(+), 2 deletions(-) create mode 100644 datamint/mlflow/flavors/validation.py create mode 100644 notebooks/05_deployment/03_validate_model.ipynb diff --git a/datamint/__init__.py b/datamint/__init__.py index d1cb4940..9986cdf7 100644 --- a/datamint/__init__.py +++ b/datamint/__init__.py @@ -9,7 +9,8 @@ # New modular datasets from .dataset.image_dataset import ImageDataset from .dataset.volume_dataset import VolumeDataset - + from .mlflow.flavors.validation import validate_model, ValidationReport, ValidationIssue, ModelValidationError + else: import lazy_loader as lazy @@ -21,6 +22,8 @@ # New modular dataset classes "dataset.image_dataset": ["ImageDataset"], "dataset.volume_dataset": ["VolumeDataset"], + "mlflow.flavors.validation": ["validate_model", "ValidationReport", + "ValidationIssue", "ModelValidationError"], }, ) diff --git a/datamint/mlflow/flavors/__init__.py b/datamint/mlflow/flavors/__init__.py index 95181169..dce6ea42 100644 --- a/datamint/mlflow/flavors/__init__.py +++ b/datamint/mlflow/flavors/__init__.py @@ -9,6 +9,7 @@ _load_pyfunc, ) from .task_type import TaskType +from .validation import validate_model, ValidationReport, ValidationIssue, ModelValidationError __all__ = [ "save_model", @@ -16,4 +17,8 @@ "load_model", "_load_pyfunc", "TaskType", + "validate_model", + "ValidationReport", + "ValidationIssue", + "ModelValidationError", ] diff --git a/datamint/mlflow/flavors/validation.py b/datamint/mlflow/flavors/validation.py new file mode 100644 index 00000000..fc17de63 --- /dev/null +++ b/datamint/mlflow/flavors/validation.py @@ -0,0 +1,228 @@ +"""Local validation of a DatamintModel before deployment.""" +from __future__ import annotations + +import logging +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Literal + +if TYPE_CHECKING: + from mlflow.models import ModelSignature + from datamint.mlflow.flavors.model import BaseDatamintModel + from datamint.dataset.base import DatamintBaseDataset + from datamint.entities.resource import BaseResource + +_LOGGER = logging.getLogger(__name__) + +_TASK_ANNOTATION_TYPE: dict[str, str] = { + 'image_segmentation': 'segmentation', + 'instance_segmentation': 'segmentation', + 'video_segmentation': 'segmentation', + 'volume_segmentation': 'segmentation', + 'image_classification': 'category', + 'multilabel_image_classification': 'category', + 'volume_classification': 'category', + 'video_frame_classification': 'category', + 'object_detection': 'square', +} + + +@dataclass +class ValidationIssue: + name: str + passed: bool + message: str + severity: Literal['error', 'warning', 'info'] = 'error' + + +@dataclass +class ValidationReport: + passed: bool + issues: list[ValidationIssue] = field(default_factory=list) + signature: ModelSignature | None = None + + def __str__(self) -> str: + lines = [] + for issue in self.issues: + if issue.passed: + marker = '[v]' + elif issue.severity == 'error': + marker = '[x]' + else: + marker = '[!]' + lines.append(f'{marker} {issue.message}') + + warnings = sum(1 for i in self.issues if not i.passed and i.severity == 'warning') + errors = sum(1 for i in self.issues if not i.passed and i.severity == 'error') + + if errors: + summary = f'Failed with {errors} error{"s" if errors > 1 else ""}' + if warnings: + summary += f', {warnings} warning{"s" if warnings > 1 else ""}' + elif warnings: + summary = f'Passed with {warnings} warning{"s" if warnings > 1 else ""}' + else: + summary = 'Passed' + + lines.append(summary + '.') + return '\n'.join(lines) + + +class ModelValidationError(Exception): + def __init__(self, report: ValidationReport) -> None: + self.report = report + super().__init__(str(report)) + + +# --- checks --- + +def _check_metadata(model: BaseDatamintModel) -> list[ValidationIssue]: + issues: list[ValidationIssue] = [] + + task_type = getattr(model, 'task_type', None) + if task_type is None: + issues.append(ValidationIssue('task_type', False, 'task_type is not set', 'warning')) + else: + issues.append(ValidationIssue('task_type', True, f'task_type: {task_type}')) + + specs = getattr(model, 'annotation_specs', None) + if not specs: + issues.append(ValidationIssue('annotation_specs', False, 'annotation_specs is not set', 'warning')) + else: + issues.append(ValidationIssue('annotation_specs', True, + f'annotation_specs: {", ".join(s.identifier for s in specs)}')) + + try: + modes = model.get_supported_modes() + if not modes: + issues.append(ValidationIssue('supported_modes', False, 'no prediction modes registered', 'error')) + else: + issues.append(ValidationIssue('supported_modes', True, f'supported modes: {", ".join(modes)}')) + except Exception as e: + issues.append(ValidationIssue('supported_modes', False, f'get_supported_modes() raised: {e}', 'error')) + + return issues + + +def _check_inference( + model: BaseDatamintModel, + sample: list[BaseResource], +) -> tuple[list[ValidationIssue], ModelSignature | None]: + issues: list[ValidationIssue] = [] + signature = None + + try: + output = model.predict(sample) + issues.append(ValidationIssue('predict', True, f'predict() OK on {len(sample)} sample(s)')) + except Exception as e: + issues.append(ValidationIssue('predict', False, f'predict() raised: {e}', 'error')) + return issues, None + + if not isinstance(output, list): + issues.append(ValidationIssue('output_type', False, + f'output must be a list, got {type(output).__name__}', 'error')) + return issues, None + + if len(output) != len(sample): + issues.append(ValidationIssue('output_length', False, + f'output length {len(output)} != input length {len(sample)}', 'error')) + else: + issues.append(ValidationIssue('output_length', True, f'output has {len(output)} element(s)')) + empty = sum(1 for inner in output if not inner) + if empty: + issues.append(ValidationIssue('output_nonempty', False, + f'{empty}/{len(output)} sample(s) returned no annotations', 'warning')) + else: + issues.append(ValidationIssue('output_nonempty', True, 'all predictions non-empty')) + + try: + from datamint.mlflow.flavors.datamint_flavor import _process_signature + signature = _process_signature(None, model) + issues.append(ValidationIssue('signature', True, 'MLflow signature inferred')) + except Exception as e: + issues.append(ValidationIssue('signature', False, f'signature inference failed: {e}', 'warning')) + + if len(output) != len(sample): + return issues, signature + + # annotation_specs consistency + specs = getattr(model, 'annotation_specs', None) + if specs: + declared = {s.identifier for s in specs} + required = {s.identifier for s in specs if s.required} + found: set[str] = { + ann.identifier + for inner in output + for ann in inner + if getattr(ann, 'identifier', None) + } + unknown = found - declared + if unknown: + issues.append(ValidationIssue('spec_identifiers', False, + f'identifiers not in annotation_specs: {", ".join(sorted(unknown))}', + 'error')) + else: + issues.append(ValidationIssue('spec_identifiers', True, 'output identifiers match annotation_specs')) + + missing = required - found + if missing: + issues.append(ValidationIssue('spec_required', False, + f'required annotations not produced: {", ".join(sorted(missing))}', + 'warning')) + else: + issues.append(ValidationIssue('spec_required', True, 'all required annotation_specs satisfied')) + + # task_type consistency + task_type = getattr(model, 'task_type', None) + if task_type is not None: + expected = _TASK_ANNOTATION_TYPE.get(str(task_type)) + if expected is not None: + wrong = { + str(getattr(ann, 'annotation_type', '')) + for inner in output + for ann in inner + if str(getattr(ann, 'annotation_type', '')) not in ('', expected) + } + if wrong: + issues.append(ValidationIssue('task_type_consistency', False, + f'task_type {str(task_type)!r} expects {expected!r}, ' + f'got: {", ".join(sorted(wrong))}', 'error')) + else: + issues.append(ValidationIssue('task_type_consistency', True, + f'annotation types consistent with task_type')) + + return issues, signature + + +# --- public API --- + +def validate_model( + model: BaseDatamintModel, + dataset: DatamintBaseDataset | None = None, + sample_input: list[BaseResource] | None = None, + *, + n_samples: int = 2, +) -> ValidationReport: + """Validate a DatamintModel locally before deployment. + + Runs metadata checks and, when sample data is available, an inference + smoke test with annotation_specs/task_type consistency checks. + The inferred MLflow signature is attached to the report. + """ + resolved: list[BaseResource] | None = sample_input + if resolved is None and dataset is not None: + resolved = list(dataset.resources[:n_samples]) + if not resolved: + _LOGGER.warning('Dataset has no resources; inference checks will be skipped.') + + issues = _check_metadata(model) + + signature = None + if resolved: + tier2, signature = _check_inference(model, resolved) + issues.extend(tier2) + else: + issues.append(ValidationIssue('inference', False, + 'no sample provided; inference checks skipped', 'warning')) + + passed = all(i.passed or i.severity != 'error' for i in issues) + return ValidationReport(passed=passed, issues=issues, signature=signature) diff --git a/notebooks/05_deployment/03_validate_model.ipynb b/notebooks/05_deployment/03_validate_model.ipynb new file mode 100644 index 00000000..c0e0ed95 --- /dev/null +++ b/notebooks/05_deployment/03_validate_model.ipynb @@ -0,0 +1,207 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "a1b2c3d4", + "metadata": {}, + "source": [ + "# Validating a Model Before Deployment\n", + "\n", + "Deploying a model just to discover it produces wrong outputs is slow and wasteful.\n", + "`validate_model` lets you catch problems locally — before anything is pushed to the registry.\n", + "\n", + "It runs two tiers of checks:\n", + "\n", + "| Tier | What it checks | Needs data? |\n", + "|------|---------------|-------------|\n", + "| Metadata | `task_type`, `annotation_specs`, `supported_modes` | No |\n", + "| Inference | `predict()` runs, output structure, label consistency | Yes |\n", + "\n", + "This notebook uses the `bccd_detection` project (blood-cell detection with YOLOX) as a concrete example." + ] + }, + { + "cell_type": "markdown", + "id": "b2c3d4e5", + "metadata": {}, + "source": [ + "## Setup" + ] + }, + { + "cell_type": "code", + "execution_count": 12, + "id": "c3d4e5f6", + "metadata": {}, + "outputs": [], + "source": [ + "import datamint.mlflow.flavors.datamint_flavor as datamint_flavor\n", + "from datamint import validate_model\n", + "from datamint.dataset import build_dataset\n", + "from datamint import Api\n", + "\n", + "\n", + "PROJECT_NAME = \"bccd_detection\"\n", + "MODEL_NAME = PROJECT_NAME\n", + "\n", + "api = Api()\n", + "proj = api.projects.create(\n", + " name=PROJECT_NAME,\n", + " description=\"Blood cell detection tutorial using YOLOX\",\n", + " exists_ok=True,\n", + ")" + ] + }, + { + "cell_type": "markdown", + "id": "d4e5f6a7", + "metadata": {}, + "source": [ + "## 1. Load the Dataset\n", + "\n", + "`validate_model` pulls sample resources directly from a `DatamintBaseDataset`.\n", + "No extra setup is needed — just pass `dataset=` and it handles the rest." + ] + }, + { + "cell_type": "code", + "execution_count": 13, + "id": "e5f6a7b8", + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Resources : 364\n", + "Box labels: ['Platelets', 'RBC', 'WBC']\n" + ] + } + ], + "source": [ + "dataset = build_dataset(project_name=PROJECT_NAME, allow_external_annotations=True)\n", + "\n", + "print(f\"Resources : {len(dataset.resources)}\")\n", + "print(f\"Box labels: {dataset.box_labels_set}\")" + ] + }, + { + "cell_type": "markdown", + "id": "f6a7b8c9", + "metadata": {}, + "source": [ + "## 2. Load the Registered Model\n", + "\n", + "Load any registered model from MLflow using its URI." + ] + }, + { + "cell_type": "code", + "execution_count": 14, + "id": "a7b8c9d0", + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "Downloading artifacts: 100%|██████████| 6/6 [00:27<00:00, 4.66s/it] " + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Model : YOLOXModule\n", + "task_type : TaskType.OBJECT_DETECTION\n", + "annotation_specs: [AnnotationSpec(type=, scope='image', required=False, identifier='Platelets'), AnnotationSpec(type=, scope='image', required=False, identifier='RBC'), AnnotationSpec(type=, scope='image', required=False, identifier='WBC')]\n" + ] + }, + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\n" + ] + } + ], + "source": [ + "model = datamint_flavor.load_model(f\"models:/{MODEL_NAME}/latest\")\n", + "\n", + "print(f\"Model : {model.__class__.__name__}\")\n", + "print(f\"task_type : {model.task_type}\")\n", + "print(f\"annotation_specs: {model.annotation_specs}\")" + ] + }, + { + "cell_type": "markdown", + "id": "b8c9d0e1", + "metadata": {}, + "source": "## 3. Run Validation\n\n`validate_model` accepts either a dataset or a list of resources directly.\n\n- `dataset` — pulls `n_samples` resources automatically\n- `sample_input` — pass a `list[Resource]` directly" + }, + { + "cell_type": "code", + "execution_count": null, + "id": "c9d0e1f2", + "metadata": {}, + "outputs": [], + "source": "report = validate_model(model, dataset=dataset, n_samples=2)\nprint(report)" + }, + { + "cell_type": "markdown", + "id": "d0e1f2a3", + "metadata": {}, + "source": [ + "## 4. Reading the Report\n", + "\n", + "Each line maps to one check:\n", + "\n", + "| Marker | Meaning |\n", + "|--------|---------|\n", + "| `[v]` | Check passed, ready to deploy. |\n", + "| `[!]` | Warning — model still deployable, but worth investigating |\n", + "| `[x]` | Error — model should not be deployed as-is |" + ] + }, + { + "cell_type": "markdown", + "id": "a3b4c5d6", + "metadata": {}, + "source": [ + "## 5. Passing Resources Directly\n", + "\n", + "If you already have specific resources you want to test against — for example,\n", + "a known edge case — pass them via `sample_input` instead of a dataset." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "b4c5d6e7", + "metadata": {}, + "outputs": [], + "source": "samples = list(dataset.resources[:3])\n\nreport = validate_model(model, sample_input=samples)\nprint(report)" + } + ], + "metadata": { + "kernelspec": { + "display_name": "env", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.3" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} \ No newline at end of file diff --git a/notebooks/06_end_to_end/slice_based/02_busi_segmentation.ipynb b/notebooks/06_end_to_end/slice_based/02_busi_segmentation.ipynb index 86943891..f052e975 100644 --- a/notebooks/06_end_to_end/slice_based/02_busi_segmentation.ipynb +++ b/notebooks/06_end_to_end/slice_based/02_busi_segmentation.ipynb @@ -52,7 +52,7 @@ "source": [ "from datamint import Api\n", "\n", - "PROJECT_NAME = \"Segmentation_Tutorial\"\n", + "PROJECT_NAME = \"UNetPP_Segmentation_Tutorial\"\n", "api = Api()" ] },