Use mlflow to serve a custom python model for scoring

Question

I am using Python code generated from an ML software with mlflow to read a data frame, perform some table operations and output a data frame. I am able to run the code successfully and save the new data frame as an artifact. However, I am unable to log the model using log_model because it is not an LR or classifier model where we train and fit. I want to log a model for this so that it can be served with new data and deployed with a rest API

df = pd.read_csv(r"/home/xxxx.csv")


with mlflow.start_run():

    def getPrediction(row):
        perform_some_python_operations 
        return [Status_prediction, Status_0_probability, Status_1_probability]

    columnValues = []
    for column in columns:
        columnValues.append([])

    for index, row in df.iterrows():
        results = getPrediction(row)
        for n in range(len(results)):
            columnValues[n].append(results[n])

    for n in range(len(columns)):
        df[columns[n]] = columnValues[n]

    df.to_csv('dataset_statistics.csv')
    mlflow.log_artifact('dataset_statistics.csv')
   

Answer 1

MLflow supports custom models of mlflow.pyfunc flavor. You can create a custom class inherited from the mlflow.pyfunc.PythonModel, that needs to provide function predict for performing predictions, and optional load_context to load the necessary artifacts, like this (adopted from the docs):

class MyModel(mlflow.pyfunc.PythonModel):

    def load_context(self, context):
        # load your artifacts

    def predict(self, context, model_input):
        return my_predict(model_input.values)

You can log to MLflow whatever artifacts you need for your models, define Conda environment if necessary, etc.
Then you can use save_model with your class to save your implementation, that could be loaded with load_model and do the predict using your model:

mlflow.pyfunc.save_model(
        path=mlflow_pyfunc_model_path, 
        python_model=MyModel(), 
        artifacts=artifacts)

# Load the model in `python_function` format
loaded_model = mlflow.pyfunc.load_model(mlflow_pyfunc_model_path)