[MLflow Demo] Base implementation for demo framework (#19994)

Signed-off-by: Ben Wilson <benjamin.wilson@databricks.com>

Author: BenWilson2

mlflow/demo/README.md

@@ -0,0 +1,75 @@
+# MLflow Demo Data Framework
+
+This module generates demo data for MLflow's GenAI features. For user-facing documentation, see the MLflow docs.
+
+## User Entry Points
+
+1. **CLI:** `mlflow demo` - Launches a temporary server with demo data
+2. **Home Page:** "Launch Demo" button - Adds demo data to an existing server
+3. **Settings:** "Clear Demo Data" - Removes all demo data
+
+## Architecture
+
+```
+mlflow/demo/
+├── base.py                  # BaseDemoGenerator, DemoFeature enum, DemoResult
+├── registry.py              # DemoRegistry for managing generators
+└── generators/
+    ├── __init__.py          # Registers generators (order matters!)
+    ├── prompts.py           # Prompt versions and aliases
+    ├── traces.py            # Sample traces with various patterns
+    ├── evaluation.py        # Evaluation runs and datasets
+    └── scorers.py           # Registered LLM judges
+```
+
+**Generator order matters** - some generators depend on others (e.g., traces depend on prompts, evaluation depends on traces).
+
+## API Endpoints
+
+| Endpoint                             | Method | Description                     |
+| ------------------------------------ | ------ | ------------------------------- |
+| `/ajax-api/3.0/mlflow/demo/generate` | POST   | Generate demo data (idempotent) |
+| `/ajax-api/3.0/mlflow/demo/delete`   | POST   | Hard delete all demo data       |
+
+## Adding a New Generator
+
+1. Add feature to `DemoFeature` enum in `base.py`
+2. Create generator class extending `BaseDemoGenerator`
+3. Register in `generators/__init__.py` (respect dependency order)
+
+```python
+class MyFeatureDemoGenerator(BaseDemoGenerator):
+    name = DemoFeature.MY_FEATURE
+    version = 1
+
+    def generate(self) -> DemoResult:
+        # Create demo data, return DemoResult with navigation_url
+        ...
+
+    def _data_exists(self) -> bool:
+        # Return True if demo data exists
+        ...
+
+    def delete_demo(self) -> None:
+        # Clean up demo data (optional, has default no-op)
+        ...
+```
+
+## Naming Conventions
+
+| Entity Type | Convention                       | Example                   |
+| ----------- | -------------------------------- | ------------------------- |
+| Experiment  | `DEMO_EXPERIMENT_NAME` constant  | `"MLflow Demo"`           |
+| Prompts     | `{DEMO_PROMPT_PREFIX}.<name>`    | `"mlflow-demo.prompts.*"` |
+| Scorers     | `{DEMO_PROMPT_PREFIX}.scorers.*` | `"mlflow-demo.scorers.*"` |
+| Metadata    | `mlflow.demo.*`                  | `mlflow.demo.version`     |
+
+## Versioning
+
+Each generator has a `version` attribute. When the version changes, old demo data is automatically deleted and regenerated. Bump versions when changes make old demo data incompatible with the current UI.
+
+## Testing
+
+```bash
+uv run pytest tests/demo/ -v
+```

mlflow/demo/__init__.py

@@ -0,0 +1,30 @@
+import logging
+
+from mlflow.demo.base import DEMO_EXPERIMENT_NAME, DEMO_PROMPT_PREFIX, BaseDemoGenerator, DemoResult
+from mlflow.demo.registry import demo_registry
+
+_logger = logging.getLogger(__name__)
+
+__all__ = [
+    "DEMO_EXPERIMENT_NAME",
+    "DEMO_PROMPT_PREFIX",
+    "BaseDemoGenerator",
+    "DemoResult",
+    "demo_registry",
+    "generate_all_demos",
+]
+
+
+def generate_all_demos() -> list[DemoResult]:
+    results = []
+    for name in demo_registry.list_generators():
+        generator_cls = demo_registry.get(name)
+        generator = generator_cls()
+        if generator.is_generated():
+            _logger.debug(f"Demo '{name}' already exists, skipping")
+            continue
+        _logger.info(f"Generating demo data for '{name}'")
+        result = generator.generate()
+        generator.store_version()
+        results.append(result)
+    return results

mlflow/demo/base.py

@@ -0,0 +1,127 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from enum import Enum
+
+DEMO_EXPERIMENT_NAME = "MLflow Demo"
+DEMO_PROMPT_PREFIX = "mlflow-demo"
+
+
+class DemoFeature(str, Enum):
+    """Enumeration of demo features that can be generated."""
+
+    TRACES = "traces"
+    EVALUATION = "evaluation"
+
+
+@dataclass
+class DemoResult:
+    """Result returned by a demo generator after creating demo data.
+
+    Attributes:
+        feature: The demo feature that was generated. Use DemoFeature enum values.
+        entity_ids: List of identifiers for created entities (e.g., trace IDs, dataset names).
+        navigation_url: URL path to navigate to view the demo data in the UI.
+    """
+
+    feature: DemoFeature
+    entity_ids: list[str]
+    navigation_url: str
+
+
+class BaseDemoGenerator(ABC):
+    """Abstract base class for demo data generators.
+
+    Subclasses must define a `name` class attribute and implement the `generate()`
+    and `_data_exists()` methods. Generators are registered with the `demo_registry`
+    and invoked during server startup to populate demo data.
+
+    Versioning:
+        Each generator has a `version` class attribute (default: 1). When demo data
+        is generated, the version is stored as a tag on the MLflow Demo experiment.
+        On subsequent startups, if the stored version doesn't match the generator's
+        current version, stale data is cleaned up and regenerated.
+
+        Bump the version when making breaking changes to demo data format.
+
+    Example:
+        class MyDemoGenerator(BaseDemoGenerator):
+            name = DemoFeature.TRACES
+            version = 1  # Bump when demo format changes
+
+            def generate(self) -> DemoResult:
+                # Create demo data using MLflow APIs
+                return DemoResult(...)
+
+            def _data_exists(self) -> bool:
+                # Check if demo data exists (version handled by base class)
+                return True/False
+
+            def delete_demo(self) -> None:
+                # Optional: delete demo data (called on version mismatch or via UI)
+                pass
+    """
+
+    name: DemoFeature | None = None
+    version: int = 1
+
+    def __init__(self):
+        if self.name is None:
+            raise ValueError(f"{self.__class__.__name__} must define 'name' class attribute")
+
+    @abstractmethod
+    def generate(self) -> DemoResult:
+        """Generate demo data for this feature. Returns a DemoResult with details."""
+
+    @abstractmethod
+    def _data_exists(self) -> bool:
+        """Check if demo data exists (regardless of version)."""
+
+    def delete_demo(self) -> None:
+        """Delete demo data created by this generator.
+
+        Called automatically when version mismatches on startup, or can be called
+        directly via API for user-initiated deletion. Override to implement cleanup.
+        """
+
+    def is_generated(self) -> bool:
+        """Check if demo data exists with a matching version.
+
+        Returns True only if data exists AND the stored version matches the current
+        generator version. If version mismatches, calls delete_demo() and
+        returns False to trigger regeneration.
+        """
+        if not self._data_exists():
+            return False
+
+        stored_version = self._get_stored_version()
+        if stored_version is None or stored_version != self.version:
+            self.delete_demo()
+            return False
+
+        return True
+
+    def _get_stored_version(self) -> int | None:
+        """Get the stored version for this generator from experiment tags."""
+        from mlflow.tracking._tracking_service.utils import _get_store
+
+        store = _get_store()
+        try:
+            experiment = store.get_experiment_by_name(DEMO_EXPERIMENT_NAME)
+            if experiment is None:
+                return None
+            version_tag = experiment.tags.get(f"mlflow.demo.version.{self.name}")
+            return int(version_tag) if version_tag else None
+        except Exception:
+            return None
+
+    def store_version(self) -> None:
+        """Store the current version in experiment tags. Called after successful generation."""
+        from mlflow.tracking._tracking_service.utils import _get_store
+
+        store = _get_store()
+        if experiment := store.get_experiment_by_name(DEMO_EXPERIMENT_NAME):
+            store.set_experiment_tag(
+                experiment.experiment_id,
+                f"mlflow.demo.version.{self.name}",
+                str(self.version),
+            )

mlflow/demo/generators/__init__.py

@@ -0,0 +1,43 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+from mlflow.demo.base import DemoFeature
+
+if TYPE_CHECKING:
+    from mlflow.demo.base import BaseDemoGenerator
+
+
+class DemoRegistry:
+    """Registry for demo data generators.
+
+    Provides registration and lookup of BaseDemoGenerator subclasses by name.
+    The global `demo_registry` instance is used by `generate_all_demos()` to
+    discover and run all registered generators.
+    """
+
+    def __init__(self):
+        self._generators: dict[DemoFeature, type[BaseDemoGenerator]] = {}
+
+    def register(self, generator_cls: type[BaseDemoGenerator]) -> None:
+        name = generator_cls.name
+        if not name:
+            raise ValueError(f"{generator_cls.__name__} must define 'name' class attribute")
+        if name in self._generators:
+            raise ValueError(f"Generator '{name}' is already registered")
+        self._generators[name] = generator_cls
+
+    def get(self, name: DemoFeature) -> type[BaseDemoGenerator]:
+        if name not in self._generators:
+            available = list(self._generators.keys())
+            raise ValueError(f"Generator '{name}' not found. Available: {available}")
+        return self._generators[name]
+
+    def list_generators(self) -> list[DemoFeature]:
+        return list(self._generators.keys())
+
+    def __contains__(self, name: DemoFeature) -> bool:
+        return name in self._generators
+
+
+demo_registry = DemoRegistry()

mlflow/demo/registry.py

@@ -0,0 +1,73 @@
+import pytest
+
+from mlflow.demo.base import BaseDemoGenerator, DemoFeature, DemoResult
+from mlflow.demo.registry import DemoRegistry
+
+
+class StubGenerator(BaseDemoGenerator):
+    name = DemoFeature.TRACES
+
+    def __init__(self, version: int = 1):
+        self._version = version
+        self.generate_called = False
+        self.data_exists_value = False
+        self.stored_version_value = None
+        self.delete_demo_called = False
+        super().__init__()
+
+    @property
+    def version(self) -> int:
+        return self._version
+
+    @version.setter
+    def version(self, value: int) -> None:
+        self._version = value
+
+    def generate(self) -> DemoResult:
+        self.generate_called = True
+        return DemoResult(
+            feature=self.name,
+            entity_ids=["entity-1"],
+            navigation_url="/stub",
+        )
+
+    def _data_exists(self) -> bool:
+        return self.data_exists_value
+
+    def _get_stored_version(self) -> int | None:
+        return self.stored_version_value
+
+    def store_version(self) -> None:
+        self.stored_version_value = self.version
+
+    def delete_demo(self) -> None:
+        self.delete_demo_called = True
+
+
+class AnotherStubGenerator(BaseDemoGenerator):
+    name = DemoFeature.EVALUATION
+
+    def generate(self) -> DemoResult:
+        return DemoResult(
+            feature=self.name,
+            entity_ids=["entity-2"],
+            navigation_url="/evaluation",
+        )
+
+    def _data_exists(self) -> bool:
+        return False
+
+
+@pytest.fixture
+def stub_generator():
+    return StubGenerator()
+
+
+@pytest.fixture
+def another_stub_generator():
+    return AnotherStubGenerator()
+
+
+@pytest.fixture
+def fresh_registry():
+    return DemoRegistry()