Mlflow Python

# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "mlflow>=2.9.0",
# ]
# ///
"""Create a new MLflow experiment with metadata.

ADR: 2025-12-12-mlflow-python-skill

This script creates MLflow experiments with optional description and tags.

Usage:
    uv run scripts/create_experiment.py --name "crypto-backtests-2025"
    uv run scripts/create_experiment.py --name "crypto-backtests-2025" --description "Q1 2025 cryptocurrency trading strategy backtests"
    uv run scripts/create_experiment.py --name "crypto-backtests-2025" --tags '{"team": "quant", "asset_class": "crypto"}'
"""

from __future__ import annotations

import argparse
import json
import sys

import mlflow


def create_experiment(
    name: str,
    description: str | None = None,
    tags: dict[str, str] | None = None,
    artifact_location: str | None = None,
) -> str:
    """Create a new MLflow experiment.

    Args:
        name: Experiment name (must be unique)
        description: Optional description
        tags: Optional tags dict
        artifact_location: Optional artifact storage location

    Returns:
        Experiment ID.
    """
    # Check if experiment already exists
    existing = mlflow.get_experiment_by_name(name)
    if existing is not None:
        print(f"Experiment '{name}' already exists with ID: {existing.experiment_id}")
        return existing.experiment_id

    # Create experiment
    experiment_id = mlflow.create_experiment(
        name=name,
        artifact_location=artifact_location,
        tags=tags,
    )

    # Set description as a tag if provided (MLflow stores description in tags)
    if description:
        client = mlflow.tracking.MlflowClient()
        client.set_experiment_tag(experiment_id, "mlflow.note.content", description)

    return experiment_id


def main() -> int:
    """Main entry point."""
    parser = argparse.ArgumentParser(description="Create a new MLflow experiment")
    parser.add_argument(
        "--name",
        "-n",
        required=True,
        help="Experiment name (must be unique)",
    )
    parser.add_argument(
        "--description",
        "-d",
        help="Experiment description",
    )
    parser.add_argument(
        "--tags",
        "-t",
        help="Tags as JSON string (e.g., '{\"team\": \"quant\"}')",
    )
    parser.add_argument(
        "--artifact-location",
        help="Artifact storage location (default: MLflow server default)",
    )

    args = parser.parse_args()

    # Parse tags if provided
    tags = None
    if args.tags:
        try:
            tags = json.loads(args.tags)
        except json.JSONDecodeError as e:
            print(f"Error parsing --tags JSON: {e}", file=sys.stderr)
            return 1

    # Create experiment
    experiment_id = create_experiment(
        name=args.name,
        description=args.description,
        tags=tags,
        artifact_location=args.artifact_location,
    )

    print("Experiment created:")
    print(f"  Name: {args.name}")
    print(f"  ID: {experiment_id}")
    if args.description:
        print(f"  Description: {args.description}")
    if tags:
        print(f"  Tags: {tags}")

    return 0


if __name__ == "__main__":
    sys.exit(main())

# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "mlflow>=2.9.0",
#     "pandas>=2.0",
#     "tabulate>=0.9.0",
# ]
# ///
"""Retrieve metric history (time-series data) for a run.

ADR: 2025-12-12-mlflow-python-skill

This script retrieves the full history of logged metrics for a run,
which is a Python API-only feature (not available via CLI).

Usage:
    uv run scripts/get_metric_history.py --run-id abc123 --metrics sharpe_ratio
    uv run scripts/get_metric_history.py --run-id abc123 --metrics sharpe_ratio,cumulative_return
    uv run scripts/get_metric_history.py --run-id abc123 --all
"""

from __future__ import annotations

import argparse
import sys

import mlflow
import pandas as pd
from tabulate import tabulate


def get_metric_history(run_id: str, metric_name: str) -> pd.DataFrame:
    """Get the full history of a metric for a run.

    Args:
        run_id: MLflow run ID
        metric_name: Name of the metric

    Returns:
        DataFrame with timestamp, step, and value columns.
    """
    client = mlflow.tracking.MlflowClient()

    try:
        history = client.get_metric_history(run_id, metric_name)
    except Exception as e:
        print(f"Error getting metric '{metric_name}': {e}", file=sys.stderr)
        return pd.DataFrame()

    if not history:
        return pd.DataFrame()

    data = [
        {
            "metric": metric_name,
            "timestamp": pd.Timestamp(m.timestamp, unit="ms"),
            "step": m.step,
            "value": m.value,
        }
        for m in history
    ]

    return pd.DataFrame(data)


def get_all_metrics(run_id: str) -> list[str]:
    """Get all metric names for a run.

    Args:
        run_id: MLflow run ID

    Returns:
        List of metric names.
    """
    client = mlflow.tracking.MlflowClient()
    run = client.get_run(run_id)
    return list(run.data.metrics.keys())


def main() -> int:
    """Main entry point."""
    parser = argparse.ArgumentParser(description="Get metric history for an MLflow run")
    parser.add_argument(
        "--run-id",
        "-r",
        required=True,
        help="MLflow run ID",
    )
    parser.add_argument(
        "--metrics",
        "-m",
        help="Comma-separated list of metrics to retrieve",
    )
    parser.add_argument(
        "--all",
        "-a",
        action="store_true",
        help="Retrieve history for all metrics",
    )
    parser.add_argument(
        "--format",
        choices=["table", "csv", "json"],
        default="table",
        help="Output format",
    )

    args = parser.parse_args()

    if not args.metrics and not args.all:
        print("Error: Must specify --metrics or --all", file=sys.stderr)
        return 1

    # Get metric names
    if args.all:
        metric_names = get_all_metrics(args.run_id)
        if not metric_names:
            print(f"No metrics found for run {args.run_id}")
            return 0
        print(f"Found {len(metric_names)} metrics: {', '.join(metric_names)}")
    else:
        metric_names = [m.strip() for m in args.metrics.split(",")]

    # Collect all metric histories
    all_data = []
    for metric_name in metric_names:
        df = get_metric_history(args.run_id, metric_name)
        if not df.empty:
            all_data.append(df)

    if not all_data:
        print("No metric history found")
        return 0

    combined = pd.concat(all_data, ignore_index=True)
    combined = combined.sort_values(["metric", "step"])

    # Output
    if args.format == "table":
        print(tabulate(combined, headers="keys", tablefmt="simple", showindex=False))
    elif args.format == "csv":
        print(combined.to_csv(index=False))
    elif args.format == "json":
        print(combined.to_json(orient="records", indent=2))

    print(f"\nTotal: {len(combined)} data points across {len(metric_names)} metrics")

    return 0


if __name__ == "__main__":
    sys.exit(main())

# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "mlflow>=2.9.0",
#     "quantstats>=0.0.77",
#     "pandas>=2.0",
#     "pydantic>=2.0",
# ]
# ///
"""Log backtest metrics to MLflow using QuantStats for comprehensive calculations.

ADR: 2025-12-12-mlflow-python-skill

This script calculates 70+ trading metrics from a returns series using QuantStats,
then logs them to MLflow. Supports both daily returns CSV files and inline data.

Usage:
    uv run scripts/log_backtest.py --experiment "crypto-backtests" --run-name "btc_momentum_v2" --returns path/to/returns.csv
    uv run scripts/log_backtest.py --experiment "crypto-backtests" --run-name "eth_mean_rev" --returns data.csv --params '{"strategy": "mean_reversion", "lookback": 20}'
"""

from __future__ import annotations

import argparse
import json
import os
import sys
from pathlib import Path

import mlflow
import pandas as pd
import quantstats as qs
from pydantic import BaseModel, Field


class BacktestConfig(BaseModel):
    """Configuration for backtest logging.

    ADR: 2025-12-12-mlflow-python-skill
    """

    experiment_name: str = Field(description="MLflow experiment name")
    run_name: str = Field(description="MLflow run name")
    returns_path: Path = Field(description="Path to returns CSV file")
    params: dict = Field(default_factory=dict, description="Strategy parameters to log")
    benchmark: str | None = Field(default=None, description="Benchmark symbol for comparison")
    risk_free_rate: float = Field(default=0.0, description="Risk-free rate for Sharpe calculation")


def load_returns(path: Path) -> pd.Series:
    """Load returns from CSV file.

    Expected format: CSV with 'date' and 'returns' columns, or single column of returns.
    """
    df = pd.read_csv(path, parse_dates=True, index_col=0)

    # Handle different CSV formats
    if isinstance(df, pd.DataFrame):
        if "returns" in df.columns:
            returns = df["returns"]
        elif len(df.columns) == 1:
            returns = df.iloc[:, 0]
        else:
            raise ValueError(f"Cannot determine returns column from: {df.columns.tolist()}")
    else:
        returns = df

    returns.index = pd.to_datetime(returns.index)
    returns = returns.sort_index()

    return returns


def calculate_quantstats_metrics(returns: pd.Series, rf: float = 0.0) -> dict[str, float]:
    """Calculate comprehensive trading metrics using QuantStats.

    Returns 70+ metrics grouped by category.
    """
    metrics = {}

    # Core ratios
    metrics["sharpe_ratio"] = float(qs.stats.sharpe(returns, rf=rf) or 0)
    metrics["sortino_ratio"] = float(qs.stats.sortino(returns, rf=rf) or 0)
    metrics["calmar_ratio"] = float(qs.stats.calmar(returns) or 0)
    metrics["omega_ratio"] = float(qs.stats.omega(returns, rf=rf) or 0)

    # Returns metrics
    metrics["cagr"] = float(qs.stats.cagr(returns) or 0)
    metrics["total_return"] = float(qs.stats.comp(returns) or 0)
    metrics["avg_return"] = float(qs.stats.avg_return(returns) or 0)
    metrics["avg_win"] = float(qs.stats.avg_win(returns) or 0)
    metrics["avg_loss"] = float(qs.stats.avg_loss(returns) or 0)
    metrics["best_day"] = float(qs.stats.best(returns) or 0)
    metrics["worst_day"] = float(qs.stats.worst(returns) or 0)

    # Drawdown metrics
    metrics["max_drawdown"] = float(qs.stats.max_drawdown(returns) or 0)
    metrics["avg_drawdown"] = float(qs.stats.avg_drawdown(returns) or 0)
    metrics["avg_drawdown_days"] = float(qs.stats.avg_drawdown_days(returns) or 0)

    # Trade metrics
    metrics["win_rate"] = float(qs.stats.win_rate(returns) or 0)
    metrics["profit_factor"] = float(qs.stats.profit_factor(returns) or 0)
    metrics["payoff_ratio"] = float(qs.stats.payoff_ratio(returns) or 0)
    metrics["consecutive_wins"] = float(qs.stats.consecutive_wins(returns) or 0)
    metrics["consecutive_losses"] = float(qs.stats.consecutive_losses(returns) or 0)

    # Risk metrics
    metrics["volatility"] = float(qs.stats.volatility(returns) or 0)
    metrics["var"] = float(qs.stats.var(returns) or 0)
    metrics["cvar"] = float(qs.stats.cvar(returns) or 0)
    metrics["ulcer_index"] = float(qs.stats.ulcer_index(returns) or 0)

    # Advanced metrics
    metrics["kelly_criterion"] = float(qs.stats.kelly_criterion(returns) or 0)
    metrics["recovery_factor"] = float(qs.stats.recovery_factor(returns) or 0)
    metrics["risk_of_ruin"] = float(qs.stats.risk_of_ruin(returns) or 0)
    metrics["tail_ratio"] = float(qs.stats.tail_ratio(returns) or 0)
    metrics["common_sense_ratio"] = float(qs.stats.common_sense_ratio(returns) or 0)
    metrics["cpc_index"] = float(qs.stats.cpc_index(returns) or 0)
    metrics["outlier_win_ratio"] = float(qs.stats.outlier_win_ratio(returns) or 0)
    metrics["outlier_loss_ratio"] = float(qs.stats.outlier_loss_ratio(returns) or 0)

    # Skew and kurtosis
    metrics["skew"] = float(qs.stats.skew(returns) or 0)
    metrics["kurtosis"] = float(qs.stats.kurtosis(returns) or 0)

    # Clean up any NaN/inf values
    metrics = {k: v if pd.notna(v) and abs(v) != float("inf") else 0.0 for k, v in metrics.items()}

    return metrics


def log_to_mlflow(config: BacktestConfig, metrics: dict[str, float], returns: pd.Series) -> str:
    """Log metrics and parameters to MLflow.

    Returns the run ID.
    """
    # Set experiment
    mlflow.set_experiment(config.experiment_name)

    with mlflow.start_run(run_name=config.run_name) as run:
        # Log parameters
        if config.params:
            mlflow.log_params(config.params)

        # Log metadata
        mlflow.log_param("returns_file", str(config.returns_path))
        mlflow.log_param("returns_start", str(returns.index.min()))
        mlflow.log_param("returns_end", str(returns.index.max()))
        mlflow.log_param("returns_count", len(returns))
        mlflow.log_param("risk_free_rate", config.risk_free_rate)

        # Log all metrics
        mlflow.log_metrics(metrics)

        return run.info.run_id


def main() -> int:
    """Main entry point."""
    parser = argparse.ArgumentParser(description="Log backtest metrics to MLflow using QuantStats")
    parser.add_argument(
        "--experiment",
        "-e",
        required=True,
        help="MLflow experiment name",
    )
    parser.add_argument(
        "--run-name",
        "-r",
        required=True,
        help="MLflow run name",
    )
    parser.add_argument(
        "--returns",
        type=Path,
        required=True,
        help="Path to returns CSV file",
    )
    parser.add_argument(
        "--params",
        type=str,
        default="{}",
        help="Strategy parameters as JSON string",
    )
    parser.add_argument(
        "--risk-free-rate",
        type=float,
        default=0.0,
        help="Risk-free rate for Sharpe calculation (default: 0.0)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Calculate metrics without logging to MLflow",
    )

    args = parser.parse_args()

    # Parse parameters
    try:
        params = json.loads(args.params)
    except json.JSONDecodeError as e:
        print(f"Error parsing --params JSON: {e}", file=sys.stderr)
        return 1

    # Validate returns file exists
    if not args.returns.exists():
        print(f"Error: Returns file not found: {args.returns}", file=sys.stderr)
        return 1

    # Load returns
    print(f"Loading returns from: {args.returns}")
    returns = load_returns(args.returns)
    print(f"  Loaded {len(returns)} data points")
    print(f"  Date range: {returns.index.min()} to {returns.index.max()}")

    # Calculate metrics
    print("Calculating QuantStats metrics...")
    metrics = calculate_quantstats_metrics(returns, rf=args.risk_free_rate)
    print(f"  Calculated {len(metrics)} metrics")

    # Print key metrics
    print("\nKey Metrics:")
    print(f"  Sharpe Ratio: {metrics['sharpe_ratio']:.4f}")
    print(f"  Sortino Ratio: {metrics['sortino_ratio']:.4f}")
    print(f"  Max Drawdown: {metrics['max_drawdown']:.2%}")
    print(f"  CAGR: {metrics['cagr']:.2%}")
    print(f"  Win Rate: {metrics['win_rate']:.2%}")
    print(f"  Profit Factor: {metrics['profit_factor']:.4f}")

    if args.dry_run:
        print("\n[Dry run] Metrics calculated but not logged to MLflow")
        print(f"\nAll metrics: {json.dumps(metrics, indent=2)}")
        return 0

    # Verify MLflow connection
    tracking_uri = os.environ.get("MLFLOW_TRACKING_URI", "http://localhost:5000")
    print(f"\nConnecting to MLflow: {tracking_uri}")

    # Create config and log
    config = BacktestConfig(
        experiment_name=args.experiment,
        run_name=args.run_name,
        returns_path=args.returns,
        params=params,
        risk_free_rate=args.risk_free_rate,
    )

    run_id = log_to_mlflow(config, metrics, returns)
    print("\nLogged to MLflow:")
    print(f"  Experiment: {config.experiment_name}")
    print(f"  Run ID: {run_id}")
    print(f"  Metrics logged: {len(metrics)}")

    return 0


if __name__ == "__main__":
    sys.exit(main())

# /// script
# requires-python = ">=3.11"
# dependencies = [
#     "mlflow>=2.9.0",
#     "pandas>=2.0",
#     "tabulate>=0.9.0",
# ]
# ///
"""Query MLflow experiments and runs with DataFrame output.

ADR: 2025-12-12-mlflow-python-skill

This script replaces CLI-based `mlflow experiments search` and `mlflow runs list`
with Python API calls that return DataFrames for easier filtering and analysis.

Usage:
    uv run scripts/query_experiments.py experiments
    uv run scripts/query_experiments.py runs --experiment "crypto-backtests"
    uv run scripts/query_experiments.py runs --experiment "crypto-backtests" --filter "metrics.sharpe_ratio > 1.5"
    uv run scripts/query_experiments.py runs --experiment "crypto-backtests" --order-by "metrics.sharpe_ratio DESC"
"""

from __future__ import annotations

import argparse
import sys

import mlflow
import pandas as pd
from tabulate import tabulate


def list_experiments(show_deleted: bool = False) -> pd.DataFrame:
    """List all MLflow experiments.

    Returns DataFrame with experiment_id, name, artifact_location, lifecycle_stage.
    """
    client = mlflow.tracking.MlflowClient()
    experiments = client.search_experiments()

    if not show_deleted:
        experiments = [e for e in experiments if e.lifecycle_stage == "active"]

    data = [
        {
            "experiment_id": e.experiment_id,
            "name": e.name,
            "artifact_location": e.artifact_location,
            "lifecycle_stage": e.lifecycle_stage,
        }
        for e in experiments
    ]

    return pd.DataFrame(data)


def search_runs(
    experiment_name: str,
    filter_string: str | None = None,
    order_by: list[str] | None = None,
    max_results: int = 100,
) -> pd.DataFrame:
    """Search runs in an experiment with optional filtering.

    Args:
        experiment_name: Name of the experiment to search
        filter_string: SQL-like filter (e.g., "metrics.sharpe_ratio > 1.5")
        order_by: List of columns to order by (e.g., ["metrics.sharpe_ratio DESC"])
        max_results: Maximum number of runs to return

    Returns:
        DataFrame with run info, params, and metrics.
    """
    # Get experiment by name
    experiment = mlflow.get_experiment_by_name(experiment_name)
    if experiment is None:
        print(f"Experiment '{experiment_name}' not found", file=sys.stderr)
        return pd.DataFrame()

    # Search runs
    runs = mlflow.search_runs(
        experiment_ids=[experiment.experiment_id],
        filter_string=filter_string or "",
        order_by=order_by,
        max_results=max_results,
    )

    return runs


def format_runs_output(df: pd.DataFrame, columns: list[str] | None = None) -> str:
    """Format runs DataFrame for display.

    Args:
        df: Runs DataFrame from search_runs
        columns: Specific columns to display (default: key columns)

    Returns:
        Formatted table string.
    """
    if df.empty:
        return "No runs found"

    # Default columns to show
    if columns is None:
        # Find metric and param columns
        metric_cols = [c for c in df.columns if c.startswith("metrics.")]
        param_cols = [c for c in df.columns if c.startswith("params.")]

        # Show key columns first, then metrics, then params
        base_cols = ["run_id", "run_name", "status", "start_time"]
        columns = [c for c in base_cols if c in df.columns]
        columns.extend(sorted(metric_cols)[:5])  # Top 5 metrics
        columns.extend(sorted(param_cols)[:3])  # Top 3 params

    # Filter to existing columns
    columns = [c for c in columns if c in df.columns]

    return tabulate(df[columns], headers="keys", tablefmt="simple", showindex=False)


def main() -> int:
    """Main entry point."""
    parser = argparse.ArgumentParser(description="Query MLflow experiments and runs")
    subparsers = parser.add_subparsers(dest="command", required=True)

    # experiments subcommand
    exp_parser = subparsers.add_parser("experiments", help="List all experiments")
    exp_parser.add_argument(
        "--show-deleted",
        action="store_true",
        help="Include deleted experiments",
    )
    exp_parser.add_argument(
        "--format",
        choices=["table", "csv", "json"],
        default="table",
        help="Output format",
    )

    # runs subcommand
    runs_parser = subparsers.add_parser("runs", help="Search runs in an experiment")
    runs_parser.add_argument(
        "--experiment",
        "-e",
        required=True,
        help="Experiment name to search",
    )
    runs_parser.add_argument(
        "--filter",
        "-f",
        help="SQL-like filter (e.g., 'metrics.sharpe_ratio > 1.5')",
    )
    runs_parser.add_argument(
        "--order-by",
        "-o",
        help="Order by clause (e.g., 'metrics.sharpe_ratio DESC')",
    )
    runs_parser.add_argument(
        "--max-results",
        type=int,
        default=100,
        help="Maximum number of runs (default: 100)",
    )
    runs_parser.add_argument(
        "--columns",
        "-c",
        help="Comma-separated list of columns to display",
    )
    runs_parser.add_argument(
        "--format",
        choices=["table", "csv", "json"],
        default="table",
        help="Output format",
    )

    args = parser.parse_args()

    if args.command == "experiments":
        df = list_experiments(show_deleted=args.show_deleted)
        if args.format == "table":
            print(tabulate(df, headers="keys", tablefmt="simple", showindex=False))
        elif args.format == "csv":
            print(df.to_csv(index=False))
        elif args.format == "json":
            print(df.to_json(orient="records", indent=2))

    elif args.command == "runs":
        order_by = [args.order_by] if args.order_by else None
        df = search_runs(
            experiment_name=args.experiment,
            filter_string=args.filter,
            order_by=order_by,
            max_results=args.max_results,
        )

        if df.empty:
            print("No runs found")
            return 0

        columns = args.columns.split(",") if args.columns else None

        if args.format == "table":
            print(format_runs_output(df, columns=columns))
        elif args.format == "csv":
            if columns:
                df = df[[c for c in columns if c in df.columns]]
            print(df.to_csv(index=False))
        elif args.format == "json":
            if columns:
                df = df[[c for c in columns if c in df.columns]]
            print(df.to_json(orient="records", indent=2))

        print(f"\nTotal: {len(df)} runs")

    return 0


if __name__ == "__main__":
    sys.exit(main())