Custom `pyfunc` Models

Wrapping anything

Sometimes a model is more than one estimator — pre-processing + inference + post-processing, or a rule engine, or an LLM call. Subclass mlflow.pyfunc.PythonModel and provide your own predict:

import mlflow.pyfunc

class Recommender(mlflow.pyfunc.PythonModel):
    def load_context(self, ctx):
        import joblib
        self.embeddings = joblib.load(ctx.artifacts['emb'])

    def predict(self, ctx, df):
        return df['user_id'].map(self.embeddings)

mlflow.pyfunc.log_model(
    artifact_path='reco',
    python_model=Recommender(),
    artifacts={'emb': '/tmp/embeddings.joblib'},
    pip_requirements=['joblib', 'pandas'],
)

The result behaves like any other MLflow model: it can be registered, staged, served and audited — with no framework assumptions.

A custom pyfunc is the USB-C adapter of MLflow. The outside is the universal connector everyone expects (.predict(df)); the inside hides whatever proprietary circuitry your model actually needs.