TopolearnReal predictors include feature engineering, thresholds and multiple models; pyfunc ships them as one deployable unit.
When a model is more than one estimator's predict, subclass mlflow.pyfunc.PythonModel:
load_context(self, context)runs once at load time;context.artifactsmaps names to local paths of files you attached, such as a tokenizer or a second model.predict(self, context, model_input, params=None)implements inference.paramsreceives per-request parameters declared in the signature.
Log it with mlflow.pyfunc.log_model(name=..., python_model=MyModel(), artifacts={...}, pip_requirements=[...]). Artifacts are copied into the model directory, so the result is self-contained.
There are two ways to provide the model:
- An instance (
python_model=MyModel()), which MLflow serializes with cloudpickle. Easy, but sensitive to Python and library versions. - Models from code: pass a path to a Python file that defines the model and calls
mlflow.models.set_model(MyModel()). Nothing is pickled; the code itself is stored, which is more robust and readable. This is the recommended approach, especially for GenAI apps and agents.
Type hints on predict (e.g. model_input: list[str]) let MLflow infer the signature and validate inputs automatically.
Keep predict free of training logic and network calls you cannot control, and pin every dependency it needs in pip_requirements; the serving environment is built from them.
Example
import mlflow
import pandas as pd
class ThresholdModel(mlflow.pyfunc.PythonModel):
def load_context(self, context):
cfg = mlflow.artifacts.load_dict(context.artifacts["config"])
self.threshold = cfg["threshold"]
def predict(self, context, model_input: pd.DataFrame, params=None):
thr = (params or {}).get("threshold", self.threshold)
score = model_input["a"] * 0.6 + model_input["b"] * 0.4
return (score > thr).astype(int)
with mlflow.start_run() as run:
mlflow.log_dict({"threshold": 0.5}, "config.json")
info = mlflow.pyfunc.log_model(
name="thresh", python_model=ThresholdModel(),
artifacts={"config": f"runs:/{run.info.run_id}/config.json"},
input_example=pd.DataFrame({"a": [0.9, 0.1], "b": [0.8, 0.2]}),
)
m = mlflow.pyfunc.load_model(info.model_uri)
print(list(m.predict(pd.DataFrame({"a": [0.9, 0.1], "b": [0.8, 0.2]})))) # [1, 0]Remember
- Subclass PythonModel; load resources in
load_context, infer inpredict - Attach files via
artifactsso the model is self-contained - Prefer models-from-code over pickled instances for robustness
Watch out for
- Loading large resources inside predict on every call instead of once in load_context.