domino-experiment-setup

/domino-experiment-setup Command

Set up MLflow experiment tracking for traditional machine learning projects in Domino.

Usage

/domino-experiment-setup [experiment_name]

What This Command Does

1. Creates experiment setup code with unique naming
2. Configures auto-logging for detected frameworks
3. Adds Domino context as MLflow tags
4. Creates example training script with best practices

Output

experiment_setup.py

"""
Domino Experiment Tracking Setup
Generated by /domino-experiment-setup
"""

import mlflow
import os

def setup_experiment(base_name: str = "experiment"):
    """
    Set up a Domino-compatible MLflow experiment.

    IMPORTANT: Experiment names must be unique across the entire
    Domino deployment. This function appends username and project
    to ensure uniqueness.
    """
    username = os.environ.get('DOMINO_STARTING_USERNAME', 'unknown')
    project = os.environ.get('DOMINO_PROJECT_NAME', 'unknown')

    # Create unique experiment name
    experiment_name = f"{base_name}-{project}-{username}"

    mlflow.set_experiment(experiment_name)
    print(f"Experiment set: {experiment_name}")

    return experiment_name

def log_domino_context():
    """Log Domino environment information as tags."""
    mlflow.set_tags({
        "domino.user": os.environ.get('DOMINO_STARTING_USERNAME', 'unknown'),
        "domino.project": os.environ.get('DOMINO_PROJECT_NAME', 'unknown'),
        "domino.run_id": os.environ.get('DOMINO_RUN_ID', 'unknown'),
        "domino.hardware_tier": os.environ.get('DOMINO_HARDWARE_TIER_NAME', 'unknown'),
    })

# Auto-detect and enable framework logging
def setup_autolog():
    """Enable auto-logging for detected ML frameworks."""
    try:
        import sklearn
        mlflow.sklearn.autolog()
        print("Enabled sklearn auto-logging")
    except ImportError:
        pass

    try:
        import tensorflow
        mlflow.tensorflow.autolog()
        print("Enabled TensorFlow auto-logging")
    except ImportError:
        pass

    try:
        import torch
        mlflow.pytorch.autolog()
        print("Enabled PyTorch auto-logging")
    except ImportError:
        pass

    try:
        import xgboost
        mlflow.xgboost.autolog()
        print("Enabled XGBoost auto-logging")
    except ImportError:
        pass

    try:
        import lightgbm
        mlflow.lightgbm.autolog()
        print("Enabled LightGBM auto-logging")
    except ImportError:
        pass