MLflow

Captures training runs with metrics, parameters, and artifacts through a fluent API that auto-logs framework-specific data. Uses a dual-layer storage architecture with a REST API server (mlflow/server) backed by pluggable storage backends (FileStore, SQLAlchemy, Databricks) that persist run metadata in structured tables and artifacts in cloud or local storage. The tracking system maintains parent-child run relationships and supports nested experiments for hierarchical organization.

Provides a centralized catalog for registered models with version control, stage management (Staging, Production, Archived), and metadata tracking. Implements a model registry store (mlflow/store/model_registry) with abstract interfaces backed by SQL or Databricks, allowing teams to promote models through lifecycle stages with approval workflows. Each model version maintains lineage to its source run, model signature, and custom tags for governance.

Provides a query language and API for searching experiments and runs by metrics, parameters, tags, and metadata. Implements a search backend (mlflow/store/tracking/search) that indexes run data for fast filtering and sorting. Supports complex queries (e.g., 'accuracy > 0.95 AND learning_rate < 0.01') via a SQL-like syntax or programmatic API.

Exposes all MLflow functionality via a REST API (mlflow/server) that enables remote clients to track experiments, manage models, and query runs. Implements a Flask-based server with request handlers for tracking, model registry, and artifact operations. Supports authentication via API tokens and integrates with Databricks for enterprise SSO.

Provides tight integration with Databricks workspace infrastructure, using Databricks volumes for artifact storage, Unity Catalog for model governance, and workspace authentication for access control. Enables seamless MLflow usage within Databricks notebooks and jobs without external server setup. Supports Databricks-native features like workspace secrets, cluster management, and audit logging.

Abstracts model serving across frameworks through a standardized PyFunc interface (mlflow/pyfunc) that wraps sklearn, TensorFlow, PyTorch, ONNX, and custom models. Enables deployment to MLflow Model Server, Spark UDFs, cloud platforms (SageMaker, AzureML), and serverless functions via a single model.yaml specification. The PyFunc loader handles environment reconstruction, dependency injection, and input/output schema validation at inference time.

Captures execution traces of LLM applications (chains, agents, function calls) with automatic instrumentation via MlflowLangchainTracer and OpenTelemetry integration. Records spans for each LLM call, tool invocation, and retrieval operation with latency, tokens, and error information. Stores traces in a dedicated backend (mlflow/store/trace) and provides a UI for visualization, latency analysis, and issue detection (e.g., high token usage, failed calls).

Manages prompts as first-class artifacts with versioning, metadata, and evaluation tracking. Stores prompts in the model registry (mlflow/entities/model_registry/prompt.py) with support for templating, variable substitution, and prompt chaining. Integrates with evaluation framework to track prompt performance metrics and enable A/B testing of prompt variants.

Provides a pluggable evaluation system (mlflow/entities/evaluation) that runs models against datasets and computes metrics using built-in evaluators (accuracy, F1, RMSE) or custom functions. Supports LLM-as-judge evaluation for generative tasks via integration with OpenAI, Anthropic, and other LLM providers. Stores evaluation results linked to model versions and runs, enabling comparison across model variants.

Abstracts artifact storage through a pluggable repository architecture (mlflow/store/artifact) supporting local filesystem, S3, Azure Blob Storage, GCS, and Databricks volumes. Handles artifact upload/download with automatic compression, deduplication, and URI resolution. Provides a unified artifact API regardless of backend, enabling seamless migration between storage tiers without code changes.

Provides a unified REST API gateway (mlflow/gateway) that abstracts multiple LLM providers (OpenAI, Anthropic, Azure OpenAI, Cohere, etc.) behind a single endpoint. Handles provider-specific request/response formatting, authentication, rate limiting, and cost tracking. Enables switching LLM providers without application code changes and supports request routing based on model availability or cost optimization.

Automatically captures training metrics, hyperparameters, and artifacts from ML frameworks without explicit logging code. Implements framework-specific autologgers (mlflow/ml_framework_integration) that hook into training loops via callbacks or decorators, extracting metrics from framework-native logging systems. Supports TensorFlow, PyTorch, scikit-learn, XGBoost, LightGBM, Keras, and others with minimal configuration.

Packages ML projects with code, dependencies, and configuration via MLflow Projects (mlflow/projects), enabling reproducible execution across environments. Captures environment specifications (conda.yaml, requirements.txt) and project metadata (entry points, parameters) in a declarative format. Reconstructs exact training environments on different machines or cloud platforms, ensuring reproducibility without manual dependency management.