---
type: concept
x-llmwiki:
  schemaVersion: "0.1"
  contentHash: e2cb3b14fb5e2f78dae49217660f19f84c3c0ddec9f5dcb41aa8887e9b1efe9e
  pageDirectory: concepts
  sources:
    - distributed-training-of-two-tower-recommendation-model-using-lightning-databricks-on-aws.md
  confidence: 0.9
  provenanceState: extracted
  freshnessStatus: unverified
  aliases:
    - mlflow-pyfunc-wrapper-for-model-serving
    - MPWFMS
  citations:
    - file: distributed-training-of-two-tower-recommendation-model-using-lightning-databricks-on-aws.md
title: MLflow PyFunc Wrapper for Model Serving
description: Wrapping a TwoTower model in an MLflow PythonModel (PyFunc) to abstract away TorchRec internals and enable simple dictionary-based serving
tags:
  - mlflow
  - model-serving
  - mlops
timestamp: "2026-06-18T15:31:31.858Z"
---

## MLflow PyFunc Wrapper for Model Serving

The **MLflow PyFunc Wrapper** is a pattern for packaging custom machine learning models — such as [PyTorch Lightning](/concepts/pytorch-lightning-for-ddp.md) or TorchRec models — into a standardized Python function interface using the `mlflow.pyfunc.PythonModel` class. This wrapper abstracts away the model’s internal implementation details and exposes a simple `predict` method, which makes the model straightforward to deploy for serving through MLflow’s inference infrastructure.^[distributed-training-of-two-tower-recommendation-model-using-lightning-databricks-on-aws.md]

### Example: Wrapping a Two‑Tower Recommendation Model

In the context of a two‑tower recommendation model, a PyFunc wrapper is created to transform the model’s native batch input (a dictionary of tensors) into a simpler dictionary-of-lists format, and to return a list of float predictions. The wrapper inherits from `PythonModel` and implements two methods:

- `__init__(self, two_tower_model)` – stores the trained model instance.
- `predict(self, model_input)` – converts the input to tensors, runs inference inside `torch.no_grad()`, applies a sigmoid activation, and returns the result as a list of floats.^[distributed-training-of-two-tower-recommendation-model-using-lightning-databricks-on-aws.md]

```python
class TwoTowerWrapper(PythonModel):
    def __init__(self, two_tower_model):
        self.two_tower_model = two_tower_model

    def predict(self, model_input: Dict[str, List]) -> List[float]:
        batch = {key: torch.tensor(value) for key, value in model_input.items()}
        if "label" not in batch:
            batch["label"] = torch.zeros(len(next(iter(batch.values()))))
        with torch.no_grad():
            output = self.two_tower_model(batch).cpu()
        output = torch.sigmoid(output)
        return output.tolist()
```

The comment in the source code explains the motivation: “To make this easy to serve, create a PyFunc that wraps the two tower model to take in a simpler input: `(Dict[str, List] -> List[float])`.”^[distributed-training-of-two-tower-recommendation-model-using-lightning-databricks-on-aws.md]

### Logging and Registering a PyFunc Wrapper

After creating the wrapper instance, the model is logged using `mlflow.pyfunc.log_model()`. This call requires an `artifact_path`, the `python_model` object, an inferred [Model Signature](/concepts/model-signatures-in-unity-catalog.md) (generated by `mlflow.models.signature.infer_signature`), and an optional `input_example`.^[distributed-training-of-two-tower-recommendation-model-using-lightning-databricks-on-aws.md]

```python
pyfunc_two_tower_model = TwoTowerWrapper(two_tower_model)
current_output = pyfunc_two_tower_model.predict(current_batch)
signature = infer_signature(current_batch, current_output)
logged_model = mlflow.pyfunc.log_model(
    artifact_path="two_tower_pyfunc",
    python_model=pyfunc_two_tower_model,
    signature=signature,
    input_example=current_batch
)
```

The logged model can then be registered to the [Unity Catalog](/concepts/unity-catalog.md) model registry using `mlflow.register_model()`, which makes it available for deployment as a serving endpoint.^[distributed-training-of-two-tower-recommendation-model-using-lightning-databricks-on-aws.md]

```python
uc_model_version = mlflow.register_model(
    f"models:/{logged_model.model_id}",
    name=f"{uc_catalog}.{uc_schema}.{model_name}"
)
```

### Benefits

- **Simplified serving interface** – The wrapper hides complex model internals (embedding tables, sparse features, etc.) behind a uniform `predict` contract.^[distributed-training-of-two-tower-recommendation-model-using-lightning-databricks-on-aws.md]
- **Interoperability** – Any `PythonModel` subclass can be served by the same MLflow inference infrastructure, regardless of the underlying framework.^[distributed-training-of-two-tower-recommendation-model-using-lightning-databricks-on-aws.md]
- **Reproducibility** – Input/output signatures and examples are captured at log time, ensuring consistent serialization and deserialization.^[distributed-training-of-two-tower-recommendation-model-using-lightning-databricks-on-aws.md]

### Related Concepts

- [PythonModel](/concepts/custom-mlflow-pythonmodel.md) – The base class for custom PyFunc wrappers.
- MLflow Model Signature – Defines the expected input and output schema.
- [Model Serving](/concepts/model-serving.md) – Deploying models as REST endpoints.
- [Unity Catalog Model Registry](/concepts/unity-catalog-model-registry.md) – Central repository for managing model versions.

## Sources

- distributed-training-of-two-tower-recommendation-model-using-lightning-databricks-on-aws.md

# Citations

1. [distributed-training-of-two-tower-recommendation-model-using-lightning-databricks-on-aws.md](/references/distributed-training-of-two-tower-recommendation-model-using-lightning-databricks-on-aws-093d5979.md)
