mlflow doesn't autolog artifacts while logging images

I am fairly new to mlflow. I stumbled across some strange behavior. When I am running a simple Keras model fit with MLFlow Autolog like this:

mlflow.set_tracking_uri("sqlite:///mlflow.db")
mlflow.set_experiment("keras_log")
mlflow.tensorflow.autolog()

# Define the parameters.
num_epochs = 10
batch_size = 256

# Train the model.
history = model.fit(X_train,
                      y_train,
                      epochs=num_epochs,
                      batch_size=batch_size,
                      validation_data=(X_test, y_test))
mlflow.end_run()

This produces the expected behavior. I can see the model in artifacts and metrics.

However, when I autolog along with custom figure containing the training accuracy and loss , the artifacts contain only the custom images. As if autolog didn't work at all.

mlflow.set_tracking_uri("sqlite:///mlflow.db")
mlflow.set_experiment("keras_log")
mlflow.tensorflow.autolog()

# Define the parameters.
num_epochs = 10
batch_size = 256
    
# Train the model.
history = model.fit(X_train,
                    y_train,
                    epochs=num_epochs,
                    batch_size=batch_size,
                    validation_data=(X_test, y_test))

##_________ Problematic Part
fig, ax = plt.subplots(1,2,figsize=(10,4))
ax[0].plot(history.history['accuracy'], label='Accuracy' )
ax[0].plot(history.history['val_accuracy'], label='Val Accuracy' )
ax[0].set_title('Accuracy')
ax[0].legend(loc='best')
ax[1].plot(history.history['loss'], label='Loss' )
ax[1].plot(history.history['val_loss'], label='Val Loss' )
ax[1].set_title('Loss')
ax[1].legend(loc='best')
mlflow.log_figure(fig,'training_history.png')
# _________

mlflow.end_run()

The model artifacts aren't there. The metrics also aren't recorded.

Am I missing something straightforward?

Please help.

You are using a library-specific autolog for tensorflow/keras. So, according to this part of the mlflow documentation:

MLflow will then automatically end the run once training ends via calls to tf.keras.fit().

In this regard, when you call the manual log mlflow.log_figure and mix Manual Logging and Auto Logging tracking APIs, a new run automatically starts and the last run will end, and that is why you only see the figure is logged. You may probably find the artifact and metrics in a previous run.