train.py

# Start an MLflow experiment run
with mlflow.start_run():
  
  lr = ElasticNet(alpha=alpha, l1_ratio=l1_ratio, random_state=42)
  lr.fit(train_x, train_y)

  predicted_qualities = lr.predict(test_x)

  rmse = np.sqrt(mean_squared_error(actual, pred))
  mae = mean_absolute_error(actual, pred)
  r2 = r2_score(actual, pred)

  mlflow.log_param("alpha", alpha)
  mlflow.log_param("l1_ratio", l1_ratio)
  
  mlflow.log_metric("rmse", rmse)
  mlflow.log_metric("r2", r2)
  mlflow.log_metric("mae", mae)

  # Create a model version from this run
  mlflow.sklearn.log_model(lr, "model", registered_model_name=MODEL_NAME)