Skip to main content

MLflow Plugins

MLflow's plugin architecture enables seamless integration with third-party tools and custom infrastructure. As a framework-agnostic platform, MLflow provides developer APIs for extending functionality across storage, authentication, execution backends, and model evaluation.

Quick Start​

Installing and Using a Plugin​

Try the built-in test plugin to see how plugins work:

# Clone MLflow and install example plugin
git clone https://github.com/mlflow/mlflow
cd mlflow
pip install -e tests/resources/mlflow-test-plugin
# Use the plugin with custom tracking URI scheme
MLFLOW_TRACKING_URI=file-plugin:$(PWD)/mlruns python examples/quickstart/mlflow_tracking.py

# Launch MLflow UI to view results
mlflow server --backend-store-uri ./mlruns

Open http://localhost:5000 to see your tracked experiment:

Quickstart UI

Plugin Benefits

Plugins let you integrate MLflow with your existing infrastructure without modifying core MLflow code, ensuring smooth upgrades and maintenance.

Plugin Types & Use Cases​

MLflow supports eight types of plugins, each addressing different integration needs:

Storage & Persistence​

Plugin TypePurposeExample Use Cases
Tracking StoreCustom experiment data storageEnterprise databases, cloud data warehouses
Artifact RepositoryCustom artifact storageIn-house blob storage, specialized file systems
Model Registry StoreCustom model registry backendEnterprise model catalogs, version control systems

Authentication & Headers​

Plugin TypePurposeExample Use Cases
Request Auth ProviderCustom authenticationOAuth, API keys, certificate-based auth
Request Header ProviderCustom HTTP headersEnvironment identification, compliance headers
Run Context ProviderAutomatic run metadataGit info, environment details, custom tags

Execution & Evaluation​

Plugin TypePurposeExample Use Cases
Project BackendCustom execution environmentsInternal clusters, job schedulers, cloud platforms
Model EvaluatorCustom evaluation metricsDomain-specific validation, custom test suites
DeploymentCustom serving platformsInternal serving infrastructure, edge deployment

Developing Custom Plugins​

Plugin Structure​

Create a plugin as a standalone Python package:

# setup.py
from setuptools import setup

setup(
    name="my-mlflow-plugin",
    version="0.1.0",
    install_requires=["mlflow>=2.0.0"],
    entry_points={
        # Define plugin entry points
        "mlflow.tracking_store": "my-scheme=my_plugin.store:MyTrackingStore",
        "mlflow.artifact_repository": "my-scheme=my_plugin.artifacts:MyArtifactRepo",
        "mlflow.run_context_provider": "unused=my_plugin.context:MyContextProvider",
        "mlflow.request_auth_provider": "unused=my_plugin.auth:MyAuthProvider",
        "mlflow.model_evaluator": "my-evaluator=my_plugin.evaluator:MyEvaluator",
        "mlflow.project_backend": "my-backend=my_plugin.backend:MyBackend",
        "mlflow.deployments": "my-target=my_plugin.deployment",
        "mlflow.app": "my-app=my_plugin.app:create_app",
    },
)

Storage Plugins​

# my_plugin/store.py
from mlflow.store.tracking.abstract_store import AbstractStore


class MyTrackingStore(AbstractStore):
    """Custom tracking store for scheme 'my-scheme://'"""

    def __init__(self, store_uri):
        super().__init__()
        self.store_uri = store_uri
        # Initialize your custom storage backend

    def create_experiment(self, name, artifact_location=None, tags=None):
        # Implement experiment creation logic
        pass

    def log_metric(self, run_id, metric):
        # Implement metric logging logic
        pass

    def log_param(self, run_id, param):
        # Implement parameter logging logic
        pass

    # Implement other required AbstractStore methods...

Authentication Plugins​

# my_plugin/auth.py
from mlflow.tracking.request_auth.abstract_request_auth_provider import (
    RequestAuthProvider,
)


class MyAuthProvider(RequestAuthProvider):
    """Custom authentication provider"""

    def get_name(self):
        return "my_auth_provider"

    def get_auth(self):
        # Return authentication object for HTTP requests
        # Can be anything that requests.auth accepts
        return MyCustomAuth()


class MyCustomAuth:
    """Custom authentication class"""

    def __call__(self, request):
        # Add authentication headers to request
        token = self._get_token()
        request.headers["Authorization"] = f"Bearer {token}"
        return request

    def _get_token(self):
        # Implement token retrieval logic
        # E.g., read from file, environment, or API call
        pass

Usage:

export MLFLOW_TRACKING_AUTH=my_auth_provider
python your_mlflow_script.py

Execution Plugins​

Project Backend Plugin​

# my_plugin/backend.py
from mlflow.projects.backend import AbstractBackend
from mlflow.projects.submitted_run import SubmittedRun


class MyBackend(AbstractBackend):
    """Custom execution backend for MLflow Projects"""

    def run(
        self,
        project_uri,
        entry_point,
        parameters,
        version,
        backend_config,
        tracking_uri,
        experiment_id,
    ):
        """Execute project on custom infrastructure"""

        # Parse backend configuration
        cluster_config = backend_config.get("cluster_config", {})

        # Submit job to your execution system
        job_id = self._submit_job(
            project_uri=project_uri,
            entry_point=entry_point,
            parameters=parameters,
            cluster_config=cluster_config,
        )

        # Return SubmittedRun for monitoring
        return MySubmittedRun(job_id, tracking_uri)

    def _submit_job(self, project_uri, entry_point, parameters, cluster_config):
        # Implement job submission to your infrastructure
        # Return job ID for monitoring
        pass


class MySubmittedRun(SubmittedRun):
    """Handle for submitted run"""

    def __init__(self, job_id, tracking_uri):
        self.job_id = job_id
        self.tracking_uri = tracking_uri
        super().__init__()

    def wait(self):
        # Wait for job completion and return success status
        return self._poll_job_status()

    def cancel(self):
        # Cancel the running job
        self._cancel_job()

    def get_status(self):
        # Get current job status
        return self._get_job_status()

Model Evaluation Plugin​

# my_plugin/evaluator.py
from mlflow.models.evaluation import ModelEvaluator
from mlflow.models import EvaluationResult


class MyEvaluator(ModelEvaluator):
    """Custom model evaluator"""

    def can_evaluate(self, *, model_type, evaluator_config, **kwargs):
        """Check if this evaluator can handle the model type"""
        supported_types = ["classifier", "regressor"]
        return model_type in supported_types

    def evaluate(
        self, *, model, model_type, dataset, run_id, evaluator_config, **kwargs
    ):
        """Perform custom evaluation"""

        # Get predictions
        predictions = model.predict(dataset.features_data)

        # Compute custom metrics
        metrics = self._compute_custom_metrics(
            predictions, dataset.labels_data, evaluator_config
        )

        # Generate custom artifacts
        artifacts = self._generate_artifacts(predictions, dataset, evaluator_config)

        return EvaluationResult(metrics=metrics, artifacts=artifacts)

    def _compute_custom_metrics(self, predictions, labels, config):
        # Implement domain-specific metrics
        return {
            "custom_score": self._calculate_custom_score(predictions, labels),
            "business_metric": self._calculate_business_metric(predictions, labels),
        }

    def _generate_artifacts(self, predictions, dataset, config):
        # Generate custom plots, reports, etc.
        return {}

SQL Server Plugin​

Store artifacts directly in SQL Server databases:

pip install mlflow[sqlserver]
import mlflow

# Use SQL Server as artifact store
db_uri = "mssql+pyodbc://user:pass@host:port/db?driver=ODBC+Driver+17+for+SQL+Server"
mlflow.create_experiment("sql_experiment", artifact_location=db_uri)

with mlflow.start_run():
    mlflow.onnx.log_model(model, name="model")  # Stored as BLOB in SQL Server

Alibaba Cloud OSS Plugin​

Integrate with Aliyun Object Storage Service:

pip install mlflow[aliyun-oss]
import os
import mlflow

# Configure OSS credentials
os.environ["MLFLOW_OSS_ENDPOINT_URL"] = "https://oss-region.aliyuncs.com"
os.environ["MLFLOW_OSS_KEY_ID"] = "your_access_key"
os.environ["MLFLOW_OSS_KEY_SECRET"] = "your_secret_key"

# Use OSS as artifact store
mlflow.create_experiment("oss_experiment", artifact_location="oss://bucket/path")

XetHub Plugin​

Use XetHub for versioned artifact storage:

pip install mlflow[xethub]
import mlflow

# Authenticate with XetHub (via CLI or environment variables)
mlflow.create_experiment(
    "xet_experiment", artifact_location="xet://username/repo/branch"
)

Elasticsearch Plugin​

Use Elasticsearch for experiment tracking:

pip install mlflow-elasticsearchstore

Testing Your Plugin​

# tests/test_my_plugin.py
import pytest
import mlflow
from my_plugin.store import MyTrackingStore


class TestMyTrackingStore:
    def setup_method(self):
        self.store = MyTrackingStore("my-scheme://test")

    def test_create_experiment(self):
        experiment_id = self.store.create_experiment("test_exp")
        assert experiment_id is not None

    def test_log_metric(self):
        experiment_id = self.store.create_experiment("test_exp")
        run = self.store.create_run(experiment_id, "user", "test_run")

        metric = mlflow.entities.Metric("accuracy", 0.95, 12345, 0)
        self.store.log_metric(run.info.run_id, metric)

        # Verify metric was logged correctly
        stored_run = self.store.get_run(run.info.run_id)
        assert "accuracy" in stored_run.data.metrics
        assert stored_run.data.metrics["accuracy"] == 0.95

    def test_log_artifact(self):
        # Test artifact logging functionality
        pass

Distribution & Publishing​

Package Structure​

my-mlflow-plugin/
├── setup.py                    # Package configuration
├── README.md                   # Plugin documentation
├── my_plugin/
│   ├── __init__.py
│   ├── store.py               # Tracking store implementation
│   ├── artifacts.py           # Artifact repository implementation
│   ├── auth.py                # Authentication provider
│   └── evaluator.py           # Model evaluator
├── tests/
│   ├── test_store.py
│   ├── test_artifacts.py
│   └── test_integration.py
└── examples/
    └── example_usage.py

Publishing to PyPI​

# Build distribution packages
python setup.py sdist bdist_wheel

# Upload to PyPI
pip install twine
twine upload dist/*

Documentation Template​

# My MLflow Plugin

## Installation
```bash
pip install my-mlflow-plugin

Configuration​

export MY_PLUGIN_CONFIG="value"

Usage​

import mlflow

mlflow.set_tracking_uri("my-scheme://config")

Features​

  • Feature 1
  • Feature 2

Examples​

See examples/ directory for complete usage examples.

Best Practices​

Plugin Development​

  • Follow MLflow interfaces - Implement all required abstract methods
  • Handle errors gracefully - Provide clear error messages for configuration issues
  • Support authentication - Integrate with existing credential systems
  • Add comprehensive logging - Help users debug configuration problems
  • Version compatibility - Test against multiple MLflow versions

Performance Optimization​

  • Batch operations - Implement efficient bulk logging when possible
  • Connection pooling - Reuse connections to external systems
  • Async operations - Use async I/O for storage operations when beneficial
  • Caching - Cache frequently accessed metadata

Security Considerations​

  • Credential management - Never log or expose sensitive credentials
  • Input validation - Validate all user inputs and URIs
  • Access controls - Respect existing authentication and authorization
  • Secure communication - Use TLS/SSL for network communications

Testing Strategy​

  • Unit tests - Test individual plugin components
  • Integration tests - Test full workflows with MLflow
  • Performance tests - Verify acceptable performance characteristics
  • Compatibility tests - Test with different MLflow versions

Ready to extend MLflow? Start with the example test plugin to see all plugin types in action, then build your custom integration!