code: def run_mlflow_callback_test(
      mlflow_client,
      config,
      training_data,
      val_data,
      test_data,
      tmpdir,
      exp_name=None,
  ):
      ludwig_exp_name = "mlflow_test"
      callback = MlflowCallback()
      wrapped_callback = mock.Mock(wraps=callback)
  
      model = LudwigModel(
          config,
          callbacks=[wrapped_callback],
          backend=FakeRemoteBackend(),
      )
      model.train(
          training_set=training_data,
          validation_set=val_data,
          test_set=test_data,
          experiment_name=ludwig_exp_name,
      )
      expected_df, _ = model.predict(test_data)
  
      # Check mlflow artifacts
      assert callback.experiment_id is not None
      assert callback.run is not None
  
      mlflow_exp_name = exp_name or ludwig_exp_name
      experiment = mlflow.get_experiment_by_name(mlflow_exp_name)
      assert experiment.experiment_id == callback.experiment_id
  
      df = mlflow.search_runs([experiment.experiment_id])
      assert len(df) == 1
  
      run_id = df.run_id[0]
      assert run_id == callback.run.info.run_id
  
      run = mlflow.get_run(run_id)
      expected_status = "FINISHED" if exp_name is None else "RUNNING"
      assert run.info.status == expected_status
      assert wrapped_callback.on_trainer_train_setup.call_count == 1
      assert wrapped_callback.on_trainer_train_teardown.call_count == 1
  
      artifacts = [
          f.path
          for f in mlflow_client.list_artifacts(
              callback.run.info.run_id, ""
          )
      ]
      local_dir = f"{tmpdir}/local_artifacts"
      os.makedirs(local_dir)
  
      assert "config.yaml" in artifacts
      local_config_path = mlflow_client.download_artifacts(
          callback.run.info.run_id, "config.yaml", local_dir
      )
  
      with open(local_config_path) as f:
          config_artifact = yaml.safe_load(f)
      assert config_artifact == upgrade_config_dict_to_latest_version(
          config
      )
  
      model_path = f"runs:/{callback.run.info.run_id}/model"
      loaded_model = mlflow.pyfunc.load_model(model_path)
  
      assert "ludwig" in loaded_model.metadata.flavors
      flavor = loaded_model.metadata.flavors["ludwig"]
      config = model.config
  
      def compare_features(key):
          assert len(config[key]) == len(flavor["ludwig_schema"][key])
          for feature, schema_feature in zip(
              config[key], flavor["ludwig_schema"][key]
          ):
              assert feature["name"] == schema_feature["name"]
              assert feature["type"] == schema_feature["type"]
  
      compare_features("input_features")
      compare_features("output_features")
  
      test_df = pd.read_csv(test_data)
      pred_df = loaded_model.predict(test_df)
      assert pred_df.equals(expected_df)
      return run
  
imports: import os
  import shutil
  import uuid
  from unittest import mock
  import mlflow
  import pandas
  import pytest
  import yaml
  from mlflow.tracking import MlflowClient
  from ludwig.api import LudwigModel
  from ludwig.constants import TRAINER
  from ludwig.contribs.mlflow import MlflowCallback
  from ludwig.export import export_mlflow
  from ludwig.utils.backward_compatibility import upgrade_config_dict_to_latest_version
  from tests.integration_tests.utils import category_feature, FakeRemoteBackend, generate_data, sequence_feature
label_desc: Use image node_id 1 for calling a global function with example usage: run_mlflow_callback_test(mlflow_client, config, training_data, val_data, test_data, tmpdir, exp_name) and returns: run