mlflow-backend-azure.py

#!/usr/bin/env -S uv run
# /// script
# requires-python = ">=3.13"
# dependencies = [
#   "mlflow",
#   "azureml-mlflow",
#   "azure-ai-ml",
#   "azure-identity",
# ]
# [tool.uv]
# exclude-newer = "2025-06-12T00:00:00Z"
# ///
import mlflow
import os
import random
from azure.ai.ml import MLClient
from azure.identity import DefaultAzureCredential


def run_experiment():
    """
    Runs a fake MLflow experiment, attempting to track to Azure ML.
    Falls back to local MLflow server if Azure ML connection fails.
    """
    credential = DefaultAzureCredential()
    ml_client = MLClient.from_config(credential=credential)

    azureml_tracking_uri = ml_client.workspaces.get(
        ml_client.workspace_name
    ).mlflow_tracking_uri
    mlflow.set_tracking_uri(azureml_tracking_uri)
    print(
        f"Successfully configured MLflow to track to Azure ML Workspace: {ml_client.workspace_name}"
    )
    print(f"Azure ML MLflow Tracking URI: {azureml_tracking_uri}")

    # Set an experiment name
    experiment_name = "My Fake Experiment"
    mlflow.set_experiment(experiment_name)

    with mlflow.start_run() as run:
        print(f"Starting run: {run.info.run_id}")

        # Log parameters
        params = {
            "learning_rate": random.uniform(0.001, 0.1),
            "epochs": random.randint(5, 50),
            "optimizer": random.choice(["adam", "sgd", "rmsprop"]),
        }
        mlflow.log_params(params)
        print(f"Logged parameters: {params}")

        # Log metrics (simulating training progress)
        for epoch in range(params["epochs"]):
            accuracy = (
                0.6 + (epoch / params["epochs"]) * 0.35 + random.uniform(-0.05, 0.05)
            )
            loss = 0.5 - (epoch / params["epochs"]) * 0.4 + random.uniform(-0.05, 0.05)
            mlflow.log_metric("accuracy", accuracy, step=epoch)
            mlflow.log_metric("loss", loss, step=epoch)

        print(f"Logged metrics for {params['epochs']} epochs.")

        # Create a dummy artifact
        artifact_dir = "artifacts"
        if not os.path.exists(artifact_dir):
            os.makedirs(artifact_dir)

        dummy_file_path = os.path.join(artifact_dir, "dummy_output.txt")
        with open(dummy_file_path, "w") as f:
            f.write(f"This is a dummy output for run {run.info.run_id}.\n")
            f.write(f"Parameters used: {params}\n")

        mlflow.log_artifact(dummy_file_path, artifact_path="outputs")
        print(f"Logged artifact: {dummy_file_path}")

        print(f"Run {run.info.run_id} finished.")
        print(f"MLflow Run URI: mlflow:///{run.info.experiment_id}/{run.info.run_id}")


if __name__ == "__main__":
    run_experiment()