tessl/pypi-mlflow

MLflow is an open source platform for the complete machine learning lifecycle

—

Pending

Overview

Eval results

Files

Tracking and Experiment Management

Name: tessl/pypi-mlflow
Author: tessl

MLflow's tracking system provides comprehensive experiment and run management with parameter, metric, and artifact logging. It supports both interactive (fluent API) and programmatic workflows with automatic logging capabilities for popular ML frameworks.

Capabilities

Run Management

Core functionality for starting, managing, and ending MLflow runs with support for nested runs and comprehensive metadata tracking.

def start_run(run_id=None, experiment_id=None, run_name=None, nested=False, parent_run_id=None, tags=None, description=None, log_system_metrics=None):
    """
    Start a new MLflow run.
    
    Parameters:
    - run_id: str, optional - Existing run ID to resume
    - experiment_id: str, optional - Experiment ID for the run
    - run_name: str, optional - Human-readable name for the run
    - nested: bool - Whether this is a nested run (default False)
    - parent_run_id: str, optional - Parent run ID for nested runs
    - tags: dict, optional - Dictionary of string key-value pairs
    - description: str, optional - Description of the run
    - log_system_metrics: bool, optional - Whether to log system metrics
    
    Returns:
    ActiveRun object representing the started run
    """

def end_run(status=None):
    """
    End the currently active run.
    
    Parameters:
    - status: str, optional - Status of the run ('FINISHED', 'FAILED', 'KILLED')
    """

def active_run():
    """
    Get the currently active run.
    
    Returns:
    ActiveRun object or None if no run is active
    """

def get_run(run_id):
    """
    Get run by ID.
    
    Parameters:
    - run_id: str - The run ID
    
    Returns:
    Run object containing run information and data
    """

def delete_run(run_id):
    """
    Delete a run.
    
    Parameters:
    - run_id: str - The run ID to delete
    """

def search_runs(experiment_ids=None, filter_string="", run_view_type=ViewType.ACTIVE_ONLY, max_results=SEARCH_MAX_RESULTS_DEFAULT, order_by=None, page_token=None, search_all_experiments=False, experiment_names=None, output_format=""):
    """
    Search for runs matching criteria.
    
    Parameters:
    - experiment_ids: list, optional - List of experiment IDs to search
    - filter_string: str - Filter expression for runs
    - run_view_type: ViewType - Type of runs to return (ACTIVE_ONLY, DELETED_ONLY, ALL)
    - max_results: int - Maximum number of runs to return
    - order_by: list, optional - List of columns to order by
    - page_token: str, optional - Token for pagination
    - search_all_experiments: bool - Whether to search all experiments
    - experiment_names: list, optional - List of experiment names to search
    - output_format: str - Output format ('list', 'pandas')
    
    Returns:
    List of Run objects or pandas DataFrame
    """

Experiment Management

Functions for creating, managing, and organizing experiments to group related runs and provide hierarchical organization of ML projects.

def create_experiment(name, artifact_location=None, tags=None):
    """
    Create a new experiment.
    
    Parameters:
    - name: str - Experiment name (must be unique)
    - artifact_location: str, optional - Location to store artifacts
    - tags: dict, optional - Dictionary of string key-value pairs
    
    Returns:
    str - Experiment ID of the created experiment
    """

def set_experiment(experiment_name=None, experiment_id=None):
    """
    Set the active experiment by name or ID.
    
    Parameters:
    - experiment_name: str, optional - Name of the experiment
    - experiment_id: str, optional - ID of the experiment
    
    Returns:
    Experiment object
    """

def get_experiment(experiment_id):
    """
    Get experiment by ID.
    
    Parameters:
    - experiment_id: str - The experiment ID
    
    Returns:
    Experiment object
    """

def get_experiment_by_name(name):
    """
    Get experiment by name.
    
    Parameters:
    - name: str - The experiment name
    
    Returns:
    Experiment object or None if not found
    """

def delete_experiment(experiment_id):
    """
    Delete an experiment.
    
    Parameters:
    - experiment_id: str - The experiment ID to delete
    """

def search_experiments(view_type=ViewType.ACTIVE_ONLY, max_results=None, filter_string=None, order_by=None, page_token=None):
    """
    Search experiments matching criteria.
    
    Parameters:
    - view_type: ViewType - Type of experiments to return
    - max_results: int, optional - Maximum number of experiments
    - filter_string: str, optional - Filter expression
    - order_by: list, optional - List of columns to order by
    - page_token: str, optional - Token for pagination
    
    Returns:
    List of Experiment objects
    """

Parameter and Metric Logging

Core logging functions for tracking experiment parameters, metrics, and performance data with support for batch operations and time-series metrics.

def log_param(key, value):
    """
    Log a parameter for the current run.
    
    Parameters:
    - key: str - Parameter name
    - value: str/int/float/bool - Parameter value
    """

def log_params(params):
    """
    Log multiple parameters for the current run.
    
    Parameters:
    - params: dict - Dictionary of parameter name-value pairs
    """

def log_metric(key, value, step=None, timestamp=None, synchronous=None):
    """
    Log a metric for the current run.
    
    Parameters:
    - key: str - Metric name
    - value: float - Metric value
    - step: int, optional - Training step/epoch number
    - timestamp: int, optional - Unix timestamp in milliseconds
    - synchronous: bool, optional - Whether to log synchronously
    """

def log_metrics(metrics, step=None, timestamp=None, synchronous=None):
    """
    Log multiple metrics for the current run.
    
    Parameters:
    - metrics: dict - Dictionary of metric name-value pairs
    - step: int, optional - Training step/epoch number
    - timestamp: int, optional - Unix timestamp in milliseconds
    - synchronous: bool, optional - Whether to log synchronously
    """

Artifact Logging

Functions for logging files, directories, and various data formats as artifacts with support for different storage backends and artifact organization.

def log_artifact(local_path, artifact_path=None, synchronous=None):
    """
    Log a local file as an artifact.
    
    Parameters:
    - local_path: str - Local file path
    - artifact_path: str, optional - Relative artifact path within run
    - synchronous: bool, optional - Whether to log synchronously
    """

def log_artifacts(local_dir, artifact_path=None, synchronous=None):
    """
    Log all files in a local directory as artifacts.
    
    Parameters:
    - local_dir: str - Local directory path
    - artifact_path: str, optional - Relative artifact path within run
    - synchronous: bool, optional - Whether to log synchronously
    """

def log_dict(dictionary, artifact_file):
    """
    Log a dictionary as a JSON artifact.
    
    Parameters:
    - dictionary: dict - Dictionary to log
    - artifact_file: str - Name of the artifact file
    """

def log_figure(figure, artifact_file, save_kwargs=None):
    """
    Log a matplotlib figure as an artifact.
    
    Parameters:
    - figure: matplotlib.figure.Figure - Figure to log
    - artifact_file: str - Name of the artifact file
    - save_kwargs: dict, optional - Keyword arguments for savefig
    """

def log_image(image, artifact_file=None, key=None, step=None, timestamp=None, synchronous=None):
    """
    Log an image as an artifact.
    
    Parameters:
    - image: PIL Image, numpy array, or file path - Image to log
    - artifact_file: str, optional - Name of the artifact file
    - key: str, optional - Metric key for image logging
    - step: int, optional - Training step number
    - timestamp: int, optional - Unix timestamp in milliseconds
    - synchronous: bool, optional - Whether to log synchronously
    """

def log_table(data, artifact_file, extra_tags=None):
    """
    Log tabular data as an artifact.
    
    Parameters:
    - data: pandas.DataFrame, numpy.ndarray, or list - Tabular data
    - artifact_file: str - Name of the artifact file
    - extra_tags: dict, optional - Additional tags for the artifact
    """

def log_text(text, artifact_file):
    """
    Log text content as an artifact.
    
    Parameters:
    - text: str - Text content to log
    - artifact_file: str - Name of the artifact file
    """

Tagging

Functions for adding metadata tags to runs and experiments for organization, filtering, and categorization purposes.

def set_tag(key, value):
    """
    Set a tag on the current run.
    
    Parameters:
    - key: str - Tag key
    - value: str - Tag value
    """

def set_tags(tags):
    """
    Set multiple tags on the current run.
    
    Parameters:
    - tags: dict - Dictionary of tag key-value pairs
    """

def delete_tag(key):
    """
    Delete a tag from the current run.
    
    Parameters:
    - key: str - Tag key to delete
    """

def set_experiment_tag(key, value):
    """
    Set a tag on the current experiment.
    
    Parameters:
    - key: str - Tag key
    - value: str - Tag value
    """

def set_experiment_tags(tags):
    """
    Set multiple tags on the current experiment.
    
    Parameters:
    - tags: dict - Dictionary of tag key-value pairs
    """

def delete_experiment_tag(key):
    """
    Delete a tag from the current experiment.
    
    Parameters:
    - key: str - Tag key to delete
    """

Automatic Logging

Configuration and management of automatic logging for supported ML frameworks with customizable logging behavior and exclusion options.

def autolog(log_input_examples=False, log_model_signatures=True, log_models=True, log_datasets=True, disable=False, exclusive=False, disable_for_unsupported_versions=False, silent=False, extra_tags=None, registered_model_name=None):
    """
    Enable automatic logging for supported ML frameworks.
    
    Parameters:
    - log_input_examples: bool - Whether to log input examples
    - log_model_signatures: bool - Whether to log model signatures
    - log_models: bool - Whether to log models
    - log_datasets: bool - Whether to log datasets
    - disable: bool - Whether to disable autologging
    - exclusive: bool - Whether to disable other autologging integrations
    - disable_for_unsupported_versions: bool - Disable for unsupported versions
    - silent: bool - Whether to suppress autolog warnings
    - extra_tags: dict, optional - Additional tags for autolog runs
    - registered_model_name: str, optional - Name for model registration
    """

Utility Functions

Helper functions for accessing run information, artifact URIs, and managing run context within MLflow workflows.

def get_artifact_uri(artifact_path=None):
    """
    Get the artifact URI for the current run.
    
    Parameters:
    - artifact_path: str, optional - Relative path within artifacts
    
    Returns:
    str - URI to the artifact location
    """

def get_parent_run():
    """
    Get the parent run of the current nested run.
    
    Returns:
    Run object or None if not nested
    """

def last_active_run():
    """
    Get the last active run.
    
    Returns:
    Run object or None
    """

Usage Examples

Basic Experiment Tracking

import mlflow
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error

# Set experiment
mlflow.set_experiment("housing-price-prediction")

# Generate sample data
X = np.random.rand(1000, 10)
y = np.random.rand(1000)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)

# Start run
with mlflow.start_run(run_name="random-forest-experiment"):
    # Log parameters
    n_estimators = 100
    max_depth = 10
    mlflow.log_param("n_estimators", n_estimators)
    mlflow.log_param("max_depth", max_depth)
    
    # Train model
    model = RandomForestRegressor(n_estimators=n_estimators, max_depth=max_depth)
    model.fit(X_train, y_train)
    
    # Make predictions and calculate metrics
    y_pred = model.predict(X_test)
    mse = mean_squared_error(y_test, y_pred)
    
    # Log metrics
    mlflow.log_metric("mse", mse)
    mlflow.log_metric("rmse", np.sqrt(mse))
    
    # Log model
    mlflow.sklearn.log_model(model, "model")
    
    # Log additional artifacts
    import matplotlib.pyplot as plt
    
    plt.figure(figsize=(8, 6))
    plt.scatter(y_test, y_pred, alpha=0.6)
    plt.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'r--', lw=2)
    plt.xlabel('Actual')
    plt.ylabel('Predicted')
    plt.title('Actual vs Predicted Values')
    
    mlflow.log_figure(plt.gcf(), "predictions_scatter.png")
    plt.close()
    
    print(f"Run ID: {mlflow.active_run().info.run_id}")

Nested Runs for Hyperparameter Tuning

import mlflow
from sklearn.model_selection import ParameterGrid
from sklearn.ensemble import RandomForestRegressor
from sklearn.metrics import mean_squared_error
import numpy as np

# Set experiment
mlflow.set_experiment("hyperparameter-tuning")

# Parameter grid
param_grid = {
    'n_estimators': [50, 100, 200],
    'max_depth': [5, 10, 15]
}

# Sample data
X = np.random.rand(1000, 10)
y = np.random.rand(1000)

best_mse = float('inf')
best_params = None

# Parent run for the entire tuning process
with mlflow.start_run(run_name="hyperparameter-tuning-parent"):
    for params in ParameterGrid(param_grid):
        # Nested run for each parameter combination
        with mlflow.start_run(nested=True, run_name=f"rf-{params['n_estimators']}-{params['max_depth']}"):
            # Log parameters
            mlflow.log_params(params)
            
            # Train model
            model = RandomForestRegressor(**params)
            model.fit(X, y)
            
            # Calculate and log metrics
            y_pred = model.predict(X)
            mse = mean_squared_error(y, y_pred)
            mlflow.log_metric("mse", mse)
            
            # Track best parameters
            if mse < best_mse:
                best_mse = mse
                best_params = params
    
    # Log best results to parent run
    mlflow.log_params(best_params)
    mlflow.log_metric("best_mse", best_mse)

Types

from mlflow.entities import Experiment, Run, RunInfo, RunData, Metric, Param, RunTag
from mlflow.tracking.fluent import ActiveRun
from mlflow.utils.mlflow_tags import MLFLOW_RUN_NAME, MLFLOW_USER, MLFLOW_SOURCE_NAME

class Experiment:
    experiment_id: str
    name: str
    artifact_location: str
    lifecycle_stage: str
    tags: Dict[str, str]
    creation_time: int
    last_update_time: int

class Run:
    info: RunInfo
    data: RunData

class RunInfo:
    run_id: str
    run_uuid: str  # Deprecated, use run_id
    run_name: str
    experiment_id: str
    user_id: str
    status: str
    start_time: int
    end_time: int
    artifact_uri: str
    lifecycle_stage: str

class RunData:
    metrics: List[Metric]
    params: List[Param]
    tags: List[RunTag]

class ActiveRun:
    info: RunInfo
    data: RunData

class Metric:
    key: str
    value: float
    timestamp: int
    step: int

class Param:
    key: str
    value: str

class RunTag:
    key: str
    value: str

class ViewType:
    ACTIVE_ONLY: str = "1"
    DELETED_ONLY: str = "2"
    ALL: str = "3"

Install with Tessl CLI