One endpoint, many models — cutting serving cost with an MLflow PyFunc router

Every model you give its own serving endpoint is another always-on box on the bill. At some point the math stops making sense. Here's how I put a whole fleet of models behind a single endpoint without changing how anything calls them.

There's a moment in every growing ML system where the serving bill stops matching the value the models produce.

It creeps up on you. You train one model, deploy it to a serving endpoint, and it works. Then you train a second model for a slightly different slice of the problem, and it gets its own endpoint. Then a third. Then someone splits one of them into two variants. Six months later you're staring at a list of a dozen serving endpoints, most of them handling a trickle of traffic, every one of them either kept warm (and billed for it) or scaling to zero (and cold-starting on the first request after every quiet period).

The models were small. The traffic was modest. But the endpoints — the always-on compute behind them — were the expensive part, and they multiplied one-for-one with the models.

This is the moment I reached for a single multi-model endpoint. The idea is almost embarrassingly simple: instead of one serving endpoint per model, bundle every model into a single MLflow PyFunc that routes each request to the right sub-model, and deploy that one PyFunc to one endpoint. This post is about what that actually looks like — the mechanism, the parts that bite, and where consolidating is the wrong call.

The pain, made concrete

Here's the shape of the problem. Say you've got a handful of tabular regressors — a price estimator, an ETA predictor, a demand forecaster, plus a couple of variants of each. The "normal" deployment looks like this:

price-model           → endpoint: price-model
eta-model             → endpoint: eta-model
demand-model          → endpoint: demand-model
price-model-region-b  → endpoint: price-model-region-b
eta-model-heavy       → endpoint: eta-model-heavy
...

Each endpoint is a piece of provisioned compute. Each one has to be created, sized, permissioned, monitored, warmed, and paid for — independently. And the cost doesn't scale with how useful a model is; it scales with how many models you have. Eight low-traffic models cost roughly eight times one low-traffic model, even if all eight together wouldn't saturate a single box.

You're left choosing between two bad options for each endpoint:

  1. Keep it warm. Fast responses, but you're paying for idle compute around the clock on a model that gets a few hundred requests a day.
  2. Let it scale to zero. Cheaper at rest, but now every model has its own cold start, and the first caller after a quiet stretch eats tens of seconds of latency. Multiply that annoyance by the number of endpoints.

There's also the operational sameness. Every endpoint is the same create → set permissions → point at a model version → monitor dance. Add a model, repeat the dance. The boilerplate isn't hard, it's just per-model, forever.

The thing all of this has in common is that the unit of cost and operation is the endpoint, not the model. So the fix is to make those two things stop being one-to-one.

The idea in one sentence

Wrap every model in a single PyFunc model that holds all of them, looks at an incoming model_name column, and dispatches each row to the sub-model it names. Deploy that one PyFunc to one endpoint. Callers pick a model by setting a column, not by choosing a URL.

That's the whole trick. The rest is making it robust.

What an MLflow PyFunc model actually is

If you've only ever logged scikit-learn or XGBoost models with their flavor-specific helpers, the generic PyFunc interface is the thing that unlocks this. A PyFunc model is just a Python class with two methods:

import mlflow.pyfunc

class MyModel(mlflow.pyfunc.PythonModel):
    def load_context(self, context):
        # Called once, when the model is loaded into the serving container.
        # context.artifacts is a dict: {artifact_name: local_path_on_disk}
        ...

    def predict(self, context, model_input, params=None):
        # Called per request. model_input is a pandas DataFrame.
        ...

Two ideas do all the work here:

artifacts let you bundle arbitrary files into the model. When you log a PyFunc, you hand MLflow a dict of {name: path}. MLflow copies those paths into the model package. At serving time, context.artifacts[name] gives you the local path to that file inside the container. Those artifacts can be anything — including other MLflow models. That's the seam we're going to exploit: a PyFunc whose artifacts are a pile of already-trained models.

load_context runs once; predict runs per request. So we load every sub-model a single time at startup, keep them in memory, and just route on each call. No per-request loading, no per-request cold start beyond the one container.

Building the router

The router model itself

The core is maybe twenty lines. At load time it reads a small metadata file listing the bundled models, then loads each one from its artifact path. At predict time it groups the input by model_name and hands each group to the matching sub-model.

import json
from pathlib import Path

import mlflow.pyfunc
import pandas as pd


class ModelRouter(mlflow.pyfunc.PythonModel):
    """Routes each row to one of several bundled sub-models.

    The input DataFrame must carry a 'model_name' column naming the target
    model. Every other column is a feature for that model.
    """

    def load_context(self, context):
        # A tiny JSON file we generate at bundle time: {"models": ["price", "eta", ...]}
        with Path(context.artifacts["metadata"]).open() as f:
            self.metadata = json.load(f)

        self.models = {}
        for key in self.metadata["models"]:
            self.models[key] = mlflow.pyfunc.load_model(context.artifacts[key])

    def predict(self, context, model_input, params=None):
        if "model_name" not in model_input.columns:
            raise ValueError("Input must contain a 'model_name' column")

        results = []
        for raw_name, group in model_input.groupby("model_name"):
            name = str(raw_name)
            if name not in self.models:
                raise ValueError(f"Unknown model: {name}")
            features = group.drop(columns=["model_name"])
            preds = self.models[name].predict(features)
            results.append(pd.DataFrame({"prediction": preds}, index=group.index))

        # Reindex back to the caller's original row order.
        return pd.concat(results).loc[model_input.index]

Grouping by model_name instead of looping row-by-row matters: a single request can carry rows for several different models, and each sub-model still gets one vectorized predict call over its slice. The final .loc[model_input.index] restitches the per-model results back into the order the caller sent.

Gathering the models to bundle

Before you can bundle models, you have to find them. In a registry-backed setup, each model has a "production" version you want to serve. So you resolve every model you intend to include down to a concrete version URI:

from mlflow.tracking import MlflowClient

def resolve_production_models(names: list[str], alias: str = "Production") -> dict[str, str]:
    """Map a friendly key -> a concrete model URI for each model's production alias."""
    client = MlflowClient()
    uris = {}
    for registered_name in names:
        try:
            version = client.get_model_version_by_alias(registered_name, alias)
            key = registered_name.split(".")[-1]  # short, route-friendly key
            uris[key] = f"models:/{registered_name}/{version.version}"
        except Exception:
            # No production version yet — skip it rather than failing the whole bundle.
            logging.warning("No %s alias for %s — skipping", alias, registered_name)
    return uris

The "skip if there's no production version" behavior is deliberate. A half-trained model that hasn't been promoted shouldn't be able to break the bundling job for every other model.

Downloading them into one place

The router's artifacts have to be actual files on disk at log time. So you download each model's artifacts locally and write the little metadata file the router reads in load_context:

import tempfile
import mlflow.artifacts

def download_artifacts(model_uris: dict[str, str]) -> dict[str, str]:
    tmp = Path(tempfile.mkdtemp())
    artifacts = {}
    for key, uri in model_uris.items():
        local = mlflow.artifacts.download_artifacts(artifact_uri=uri, dst_path=str(tmp / key))
        # download can hand back a nested dir; walk to the folder holding MLmodel.
        model_dir = Path(local)
        if not (model_dir / "MLmodel").exists():
            found = list(model_dir.rglob("MLmodel"))
            if found:
                model_dir = found[0].parent
        artifacts[key] = str(model_dir)

    metadata = tmp / "metadata.json"
    metadata.write_text(json.dumps({"models": list(model_uris.keys())}))
    artifacts["metadata"] = str(metadata)
    return artifacts

That MLmodel walk is one of those small things you only learn by hitting it: depending on how a model was logged, the download can land you one directory above the actual model folder. Better to find the MLmodel file than to assume the path.

Merging the dependencies, automatically

Here's a problem that consolidation creates: each sub-model was logged with its own requirements.txt. The router's serving environment needs the union of all of them. You could maintain that list by hand, but it'll drift the first time someone bumps a dependency. So merge it programmatically — every downloaded artifact carries its requirements, and you dedupe by package name:

def merge_pip_requirements(artifacts: dict[str, str]) -> list[str]:
    seen: dict[str, str] = {}
    for key, path in artifacts.items():
        if key == "metadata":
            continue
        req_file = Path(path) / "requirements.txt"
        if not req_file.exists():
            continue
        for line in req_file.read_text().splitlines():
            line = line.strip()
            if not line or line.startswith("#"):
                continue
            pkg = line.split("==")[0].split(">=")[0].split("<=")[0].strip().lower()
            seen.setdefault(pkg, line)
    return list(seen.values())

Now adding a model contributes its dependencies to the router automatically. No manual environment maintenance, no "works in training, missing a package in serving" surprise.

The signature trick that makes the whole thing work

This is the part worth slowing down on, because it's the least obvious and the most important.

When you serve a model, the platform enforces the model's signature — it validates and coerces the incoming columns against the declared input schema. That's great for a single model. It's a disaster for a router, because every sub-model has a different, often conflicting set of feature columns. There is no one input schema that fits all of them.

The fix is counterintuitive: give the router a signature that declares almost nothing.

from mlflow.models.signature import ModelSignature
from mlflow.types.schema import ColSpec, Schema

def build_router_signature() -> ModelSignature:
    # Declare only the routing column. Everything else passes through untouched,
    # so the platform won't try to enforce feature schemas that vary per model.
    return ModelSignature(
        inputs=Schema([ColSpec("string", "model_name")]),
        outputs=Schema([ColSpec("double", "prediction")]),
    )

By declaring only model_name, you tell the serving layer "don't enforce anything about the feature columns — just pass them through." The router then takes responsibility for getting each sub-model the exact columns it expects. We'll get to that responsibility in a moment, because it's where the real work moved.

Logging and pointing one endpoint at it

With artifacts gathered, requirements merged, and a signature built, you log the router as a normal PyFunc and register it:

import cloudpickle, sys

with mlflow.start_run(run_name="bundle-multi-model-router"):
    cloudpickle.register_pickle_by_value(sys.modules[__name__])
    mlflow.pyfunc.log_model(
        artifact_path="router_model",
        python_model=ModelRouter(),
        artifacts=artifacts,
        signature=build_router_signature(),
        pip_requirements=merge_pip_requirements(artifacts),
        registered_model_name="catalog.schema.model_router",
    )
    cloudpickle.unregister_pickle_by_value(sys.modules[__name__])

That register_pickle_by_value line is a production detail you'll be glad you knew up front. By default, cloudpickle serializes your router class by reference — it stores "this class lives in module model_bundler" and expects that module to be importable when the model is loaded. In a serving container, it usually isn't. Registering the module by value serializes the class definition itself into the model, so the serving environment can reconstruct it without your source package installed. Skip this, and your endpoint fails to load with a confusing ModuleNotFoundError that has nothing to do with your actual dependencies. (I learned this one the fun way.)

After that, it's the same create-or-update logic any single-model deployment uses, except you run it exactly once, for one endpoint:

from databricks.sdk.service.serving import (
    EndpointCoreConfigInput,
    ServedModelInput,
)

served_model = ServedModelInput(
    model_name="catalog.schema.model_router",   # the router, not any one model
    model_version=new_version,
    scale_to_zero_enabled=False,   # one warm box for everything
    workload_size="Small",
)

if endpoint_exists(name):
    client.serving_endpoints.update_config(name=name, served_models=[served_model])
else:
    client.serving_endpoints.create(
        name=name, config=EndpointCoreConfigInput(served_models=[served_model])
    )

The multi-model nature of the system is invisible at this layer: as far as the infrastructure is concerned, the endpoint serves exactly one model. That it happens to be a router holding a dozen others is the router's private business. One endpoint. One thing to permission, monitor, and pay for.

Calling it

From the caller's side, picking a model is now picking a column value. You build your features the same as always, add a model_name column, and send the records to the single endpoint:

features["model_name"] = "eta"          # choose the sub-model
records = features.to_dict(orient="records")
response = client.serving_endpoints.query(name="model_router", dataframe_records=records)

One endpoint URL reaches every model. Want a different model? Change the string. You can even mix models in a single request — rows naming "price" and rows naming "eta" come back correctly routed, because the router groups before it dispatches.

One real-world wrinkle worth mentioning: JSON has no representation for NaN or Infinity, and tabular feature frames are full of them. Clean those to null before you serialize, or the endpoint will reject the payload:

for record in records:
    for k, v in record.items():
        if isinstance(v, float) and (math.isnan(v) or math.isinf(v)):
            record[k] = None

There's one piece of coupling hiding in that "eta" string. The value a caller sends has to match a key the router stored in its self.models dict at load time — get it wrong and you get a runtime Unknown model error, not a friendly failure at import. So don't hand-write that string at the call site. Derive it from the same naming convention the bundler uses (the short key off the registered model name) and centralize it in one helper that both sides import. The router and its callers should agree on names in exactly one place, not at every call.

The part nobody warns you about: you now own schema enforcement

Remember how the router's signature deliberately declares nothing about features? That convenience has a bill attached, and it comes due inside predict.

When each model had its own endpoint, the serving layer's signature enforcement quietly did two jobs for you: it selected the right columns and coerced them to the dtypes the model trained on. The moment you collapse everything behind a pass-through router signature, both of those jobs become your code's problem.

The symptom, if you skip this, is ugly and model-specific. A tree model like LightGBM will reject input whose categorical columns don't match what it saw in training — "categorical_feature do not match" — or blow up comparing a float against a string category. The incoming payload has columns the model never saw (shared context that's relevant to other sub-models), and it's missing columns this one expects, and the dtypes are whatever JSON-to-pandas inference guessed.

So the router has to realign each group to its sub-model before predicting. Two steps: fix the columns, then fix the types.

def align_features(self, name, features):
    expected = self.expected_columns.get(name)   # captured from the sub-model's signature at load time
    if not expected:
        return features

    aligned = features.copy()
    for col in expected:
        if col not in aligned.columns:
            aligned[col] = 0.0                   # add what's missing
    aligned = aligned.reindex(columns=expected)  # drop extras + fix order
    return self.coerce_to_schema(name, aligned)  # then fix dtypes

Capturing expected is just reading each sub-model's own input signature in load_context and stashing its column names and types. The model already told you what it wants; you're just honoring it manually now.

There's a subtlety here that bites if you miss it. Each sub-model still carries its own signature, so the moment the router loads it, that inner signature is live too — and it would happily try to enforce itself a second time, on the very columns you just carefully aligned. That's the per-model enforcement you escaped at the router boundary, sneaking back in one level down. The fix is to read the schema, remember it in the router's own lookup, and then strip it off the loaded sub-model so the router is the only thing doing enforcement. Read it, stash it, clear it.

The dtype coercion is the half that's easy to get subtly wrong. The instinct is to reach for MLflow's built-in _enforce_schema, but it's all-or-nothing: it aborts on the first incompatible column and leaves every column after it untouched. One int64-where-double-expected column early in the frame, and your carefully aligned categorical columns downstream never get coerced at all. Per-column is the resilient choice:

from mlflow.types.schema import DataType

def coerce_to_schema(self, name, features):
    schema = self.expected_schema.get(name)
    if schema is None:
        return features
    result = features.copy()
    for spec in schema.inputs:
        col = spec.name
        if col not in result.columns:
            continue
        try:
            if spec.type == DataType.string:
                # Critically: this turns columns we 0.0-filled above into "0.0"
                # strings, so a tree model never compares a float to a real category.
                result[col] = result[col].astype("string").astype(object)
            elif spec.type in {DataType.double, DataType.float}:
                result[col] = pd.to_numeric(result[col], errors="coerce").astype(float)
            elif spec.type in {DataType.integer, DataType.long}:
                result[col] = pd.to_numeric(result[col], errors="coerce")
            elif spec.type == DataType.boolean:
                result[col] = result[col].astype(bool)
        except Exception:
            # A failure on one column is logged and skipped — it doesn't poison the rest.
            logging.warning("Could not coerce %s on model %s", col, name, exc_info=True)
    return result

This is the conceptual mirror of the SQL-templating lesson I keep relearning in different forms: the framework was doing invisible work for you, and the moment you route around it for flexibility, you have to reproduce that work explicitly and own it. It's not a lot of code. But it's code you didn't have to write when every model had its own endpoint, and it's the honest cost of consolidation.

What it bought us

The bill stopped scaling with the model count. This was the entire point, and it delivered. A dozen endpoints became one. The compute that used to sit warm under twelve models — or cold-start twelve separate times — is now a single box. Adding the thirteenth model adds zero endpoints.

Cold starts amortize. With one shared container, the first request after a quiet period pays one cold start, and then every model is warm. Per-model endpoints each had their own cold start to suffer through. Consolidation turns N intermittent cold starts into one.

Adding a model became a no-op for ops. Because the bundler discovers production models and packages whatever it finds, shipping a new model is "promote it to production and re-run the bundle." No new endpoint to create, size, permission, or monitor. The per-model operational dance is gone.

There's one place to look. One endpoint to watch, one set of logs, one permission grant, one thing to warm. The cognitive overhead of "which of these fourteen endpoints is misbehaving" disappears.

Batched, mixed-model requests come for free. Because predict groups by model_name, a single request can carry rows bound for several different models and they all get routed in one round trip. That fell out of the design rather than being designed for, but it turned out to matter for batch scoring, where a chunk of rows naturally spans more than one model type — one call, one network hop, every row routed correctly.

When one endpoint is the wrong call

I'd be selling you something if I pretended this is free. Consolidation trades a cost problem for a set of coupling problems, and they're real.

Sub-models share one container's resources. They split the same CPU and memory. A memory-hungry model can starve its neighbors, and a sudden burst of traffic to one model contends with every other. If one model is meaningfully heavier or hotter than the rest, it probably deserves its own endpoint where it can scale on its own terms.

They share one dependency environment. Merging requirements assumes the union is compatible. The day two models genuinely need conflicting versions of the same library, the merge can't save you — that's a signal those two don't belong in the same container.

The blast radius is bigger. One sub-model that fails to load can take down startup for the whole endpoint, which means every model behind it. Per-model endpoints fail independently; a router fails together. You're trading many small blast radii for one large one, and you should make that trade with your eyes open.

Release cadences get coupled. Re-bundling to ship a new version of one model also re-ships the currently-pinned versions of every other model in the bundle. When all of them are owned by one team on one cadence, that's fine. The day they're owned by different teams on different schedules — or one needs an urgent hotfix while another is mid-experiment — that coupling becomes coordination overhead you simply didn't have when each model owned its own endpoint.

Per-model autoscaling and metrics get muddier. You can't independently scale the one hot model, and request metrics are aggregated across everything unless you add your own per-model instrumentation. If you need to reason about each model's latency and throughput in isolation, separate endpoints give you that for free.

The break-even, roughly, is: many models, each with modest and similar traffic, similar runtimes, and compatible dependencies, where cost matters more than isolation. That's exactly the "fleet of small models each paying for its own box" situation that started this post. If instead you have a few heavyweight models with very different scaling profiles, the isolation of separate endpoints is worth paying for.

One last honest cost, independent of when you choose the pattern: debugging gains a layer. A wrong prediction could now be the sub-model itself, the routing, or the feature alignment in between — and the stack trace surfaces from inside the router rather than from a clean single-model endpoint. It's rarely a dealbreaker, but the first baffling trace out of the alignment code is a rite of passage worth bracing for.

The bottom line

A single multi-model endpoint didn't make any individual model faster or smarter. It changed what you pay per model — from "one endpoint each" to "one column value each."

The mechanism is just MLflow's generic PyFunc interface used a little unusually: a model whose artifacts are other models, a signature that deliberately declares almost nothing, and a predict that routes. The work that used to be spread across a dozen endpoints — and a dozen line items — collapses into one artifact you build once and deploy once.

The first time you ship a new model and realize there's no new endpoint to stand up — you just promote it and re-bundle — is when the approach earns its keep. And the first time you compare the new serving bill to the old one, you stop missing the dozen endpoints you used to maintain.

If you're running more than a handful of models with traffic that doesn't justify a dedicated box each, it's worth the afternoon it takes to build the router.