Categories: AI, Databricks

by Artjoms Rebrovs

by Artjoms Rebrovs

A practical guide to packaging and serving a tabular foundation model as a real-time endpoint on Databricks.

What Is TabFM?

TabFM (Tabular Foundation Model) is an open-source foundation model from Google Research designed specifically for structured, tabular data — the kind of data that lives in database tables, CSV exports, and feature stores.

Unlike traditional ML models that must be trained from scratch on every new dataset, TabFM uses in-context learning: you provide a handful of labeled examples ("context rows") alongside the rows you want predicted ("query rows") in a single inference call. The model generalizes from those few-shot examples on the fly — no gradient updates, no training pipeline, no retraining when your schema changes.

This makes TabFM particularly interesting for scenarios where:

  • Labeled data is scarce. A few dozen examples are enough for the model to produce meaningful predictions. Classical approaches need hundreds or thousands of rows before they converge.
  • Schemas change frequently. Because TabFM consumes feature key-value pairs at inference time rather than relying on a fixed input shape learned during training, you can add, remove, or rename columns without retraining.
  • Fast iteration matters more than peak accuracy. TabFM lets you go from raw data to a live prediction endpoint in minutes rather than days. It won't beat a carefully tuned XGBoost on a well-understood dataset, but it eliminates the entire train-evaluate-register-deploy cycle when you need a working baseline quickly.
  • You need a single model for many tasks. One TabFM checkpoint handles both classification and regression across arbitrary feature sets, reducing the number of model artifacts you manage.

Under the hood, TabFM is a transformer that treats each row as a sequence of (column-name, value) tokens. Context rows and their labels teach the model the structure of the current task; query rows are then processed through the same attention mechanism to produce predictions. The pretrained checkpoint — available via HuggingFace Hub — has been trained on a large corpus of public tabular datasets, giving it broad prior knowledge about feature-target relationships across domains.

Using TabFM: fit() and predict()

TabFM follows the scikit-learn estimator pattern, so if you've used RandomForestClassifier or XGBRegressor, the API will feel familiar — but the semantics are fundamentally different.

Loading the pretrained model

Before anything else, you load the pretrained checkpoint. TabFM ships a single model that handles both classification and regression; you specify the task when loading:

from tabfm import TabFMClassifier, TabFMRegressor
from tabfm.src.pytorch.tabfm_v1_0_0 import load as load_tabfm

# Load the pretrained transformer (downloads weights from HuggingFace Hub)
pretrained_model = load_tabfm("classification")  # or "regression"

# Wrap in the sklearn-compatible interface
model = TabFMClassifier(model=pretrained_model)

fit() — preparing context, not training

The fit() call accepts your labeled examples — the rows that will teach the model about the current task. In traditional ML this is where gradient descent happens; here it only preprocesses the data and stages it for inference:

import pandas as pd

# Your labeled examples (the "context")
context_features = pd.DataFrame()
context_labels = pd.Series([0, 1, 0, 1])  # e.g., churn / no-churn

# "Fit" — encodes features, builds ensemble views, optionally caches KV tensors
model.fit(context_features, context_labels)

After this call, the pretrained transformer weights are untouched. What changed is the preprocessing state (encoders, scalers) and — if caching is enabled — a precomputed KV cache of the context rows' attention representations.

predict() — few-shot inference in one forward pass

Once context is staged, you pass in the rows you want classified (or regressed). The transformer processes them alongside the context in a single forward pass:

# New, unlabeled rows you want predictions for
query_features = pd.DataFrame()

predictions = model.predict(query_features)
# array([0, 1])  — predicted labels based on the context examples

The model looks at the context rows and their labels, learns the pattern (via attention, not gradient updates), and applies it to the query rows — all in one call, with no training loop.

The full cycle is lightweight and repeatable. You can fit() with a different dataset and predict() again immediately — each cycle is independent. This is what makes TabFM viable as a real-time serving endpoint: context and query data arrive together in every request, and the model produces predictions without any persistent state from previous requests.

How In-Context Learning Actually Works (It's Not Training)

A common first question is: "If I fit() the model on one dataset and then fit() it on a completely different one, does it remember the first?" The answer is no — and understanding why requires a closer look at what fit() actually does.

TabFM's TabFMClassifier and TabFMRegressor follow the scikit-learn estimator interface, so they expose familiar fit() and predict() methods. But despite the familiar API, fit() does not perform gradient descent or update the model's pretrained weights in any way. The source code states this explicitly:

"The model itself is not trained on the data; it uses in-context learning at inference time. This method only prepares the data transformations."

— TabFMClassifier.fit() docstring (line 2525)

When you call fit() a second time with different data, every piece of state from the first call — the label encoder, the feature transformer, the ensemble generator, and the KV cache — is fully replaced (see line 2699 where self.context_caches_ is reassigned). There is no accumulation, no memory of the previous dataset, and no risk of cross-contamination between tasks.

This statelessness is what makes TabFM suitable for multi-tenant serving: a single endpoint can handle fraud detection in one request and churn prediction in the next, with complete isolation guaranteed by the architecture.

Deploying TabFM to Databricks Model Serving

Why MLflow?

Databricks Model Serving doesn't run arbitrary Python scripts — it serves MLflow models. Every endpoint expects a model artifact that conforms to MLflow's packaging contract: a predict() method, a declared input/output signature, a pinned set of dependencies, and a registered artifact in Unity Catalog.

This means you can't just pip install tabfm on a VM and open a port. To get TabFM behind a Databricks endpoint, you need to:

  1. Wrap the model in an mlflow.pyfunc.PythonModel subclass that implements load_context() (one-time model loading) and predict() (per-request inference).
  2. Log the wrapped model with mlflow.pyfunc.log_model(), which snapshots the code, dependencies, and a declared input/output signature into a versioned artifact.
  3. Register that artifact in Unity Catalog (models:/catalog.schema.model_name) so it can be referenced by the serving infrastructure.
  4. Deploy the registered model version to a Model Serving endpoint, which spins up a container, calls your load_context() once, and routes every incoming HTTP POST to your predict() method.

MLflow is not optional here — it's the contract between your model code and the serving infrastructure. The good news is that once the wrapper is right, registration and deployment are a few API calls.

Challenge 1: Only predict() Is Reachable

Here's the core tension. TabFM's scikit-learn interface exposes two methods:

  • fit(X, y) — stage context rows and their labels
  • predict(X) — generate predictions for query rows using the staged context

But Databricks Model Serving maps every incoming HTTP request to a single method: predict(). There is no /fit route, no way to call fit() separately, and no session state between requests. The endpoint is stateless — each request arrives, predict() runs, and everything is discarded.

This means the classic two-step workflow (model.fit(context); model.predict(query)) simply doesn't work as-is behind a serving endpoint. If fit() isn't callable from the outside, how do you get context data into the model?

The solution: collapse fit + predict into a single predict() call. The MLflow wrapper must accept both context and query rows in one request, partition them internally, call fit() on the context portion, and then call predict() on the query portion — all within a single predict() invocation.

This challenge can be handled by adding metadata columns to the input DataFrame:

Column Purpose
_data JSON string containing the feature key-value pairs for the row
_role "context" for labeled examples, "query" for rows to predict
_label Target value for context rows; ignored for query rows

Inside the wrapper's predict() method, the logic is straightforward:

def predict(self, context, model_input: pd.DataFrame) -> pd.DataFrame:
    # Partition by role
    context_mask = model_input["_role"] == "context"
    query_mask = model_input["_role"] == "query"

    # Parse features from JSON
    context_X = parse_features(model_input.loc[context_mask, "_data"])
    query_X = parse_features(model_input.loc[query_mask, "_data"])
    context_y = model_input.loc[context_mask, "_label"]

    # fit + predict in one call
    self.model.fit(context_X, context_y)
    predictions = self.model.predict(query_X)

    return pd.DataFrame()

Every request carries its own context. No state leaks between requests. The serving endpoint stays stateless, and TabFM's in-context learning works exactly as designed.

Challenge 2: A Fixed Signature for Unfixed Data

MLflow requires every logged model to declare an input signature — the exact column names and types the predict() method expects. This signature is enforced at serving time: if an incoming request doesn't match, it's rejected before your code even runs.

For a traditional model, this is fine. A churn model always expects age, income, region — the schema is known at training time and never changes.

TabFM is different. The entire point of a foundation model for tabular data is that it works across arbitrary feature sets. One request might send customer demographics; the next might send sensor readings with completely different column names and types. If the signature locks in a fixed set of columns at logging time, you've defeated the purpose of using a foundation model.

The solution: encode features as JSON strings. Instead of declaring one column per feature, the wrapper uses a generic three-column schema:

_data:  string   — JSON-encoded feature key-value pairs
_role:  string   — "context" or "query"
_label: string   — target value (for context rows)

This signature is fixed (MLflow is happy), but the actual feature space is dynamic in the _data column (TabFM is happy). The _data column carries an arbitrary JSON object per row — in one request, in the next — and the wrapper deserializes it into a DataFrame inside predict().

The tradeoff is clear: you lose schema validation at the serving layer (the endpoint can't reject a malformed feature payload until your code parses it) and callers need to know the JSON contract. But for a model that's designed to handle any tabular schema, this is the right compromise — a rigid per-column signature would be more limiting than helpful.

Here's what the MLflow signature looks like at logging time:

from mlflow.models import infer_signature

signature = infer_signature(
    model_input=sample_input_df[["_data", "_role", "_label"]],
    model_output=pd.DataFrame()
)

mlflow.pyfunc.log_model(
    artifact_path="tabfm",
    python_model=MyCustomWrapper(task="classification"),
    signature=signature,
    ...
)

With these two design choices — collapsing fit/predict into one call, and encoding features as JSON — TabFM fits neatly into the Model Serving contract while preserving the flexibility that makes it useful in the first place.

Challenge 3: Classification and Regression on One Endpoint

TabFM's pretrained checkpoint is task-agnostic — the same transformer weights handle both classification and regression. But the wrapper we saw above is not. The task type is fixed at instantiation:

tabfm_wrapper = MyCustomWrapper(task="classification")

Inside load_context(), this parameter determines which scikit-learn interface wraps the pretrained model:

if self.task == "classification":
    self.model = TabFMClassifier(model=pretrained_model)
else:
    self.model = TabFMRegressor(model=pretrained_model)

TabFMClassifier and TabFMRegressor have different preprocessing pipelines (label encoding vs. continuous target scaling), different output formats (class labels/probabilities vs. continuous values), and different calibration logic. They are not interchangeable at inference time. A model version logged with task="classification" will always run classification, regardless of what the caller sends. Sending regression-style continuous labels to a classification endpoint will treat each unique value as a discrete class — likely producing nonsensical predictions or outright errors.

The solution: we have options here.

The simplest option is two separate endpoints — one per task. This is operationally clean and avoids any cross-task interference. The downside is double the GPU cost, though the scheduled warm-up/cool-down pattern described later can mitigate that for dev environments.

If you choose this option make sure you name and register models using different names in Unity Catalog, otherwise the deployment will overwrite one model with the other in the Serving Endpoint.

The single-endpoint option is a dual-task wrapper that loads both models at startup and dispatches per-request based on a _task field in the payload. The signature gains a fourth column:

_data:  string   — JSON-encoded feature key-value pairs
_role:  string   — "context" or "query"
_label: string   — target value (for context rows)
_task:  string   — "classification" or "regression"

So how do you serve both tasks without deploying two endpoints? The wrapper loads both model types in load_context() and routes each request to the right one:

def load_context(self, context):
    from tabfm import TabFMClassifier, TabFMRegressor
    from tabfm.src.pytorch.tabfm_v1_0_0 import load as load_tabfm

    # Load both models — doubles GPU memory for weights
    clf_pretrained = load_tabfm("classification")
    reg_pretrained = load_tabfm("regression")

    if torch.cuda.is_available():
        clf_pretrained = clf_pretrained.to("cuda")
        reg_pretrained = reg_pretrained.to("cuda")

    self.models = 

def predict(self, context, model_input: pd.DataFrame) -> pd.DataFrame:
    task = model_input["_task"].iloc[0]  # "classification" or "regression"
    model = self.models[task]

    # ... same _role / _data / _label partitioning as before ...
    # For regression, cast _label to float before fit()
    model.fit(context_X, context_y)
    predictions = model.predict(query_X)
    return pd.DataFrame()

A classification request looks exactly like before, with _task added:

", "context", "low_risk", "classification"],
      ["", "context", "high_risk", "classification"],
      ["",  "query",   "",          "classification"]
    ]
  }
}

A regression request hits the same endpoint with different _task and _label values:

", "context", "250000", "regression"],
      ["", "context", "520000", "regression"],
      ["", "query",   "",       "regression"]
    ]
  }
}

The tradeoff is real. Loading two pretrained checkpoints roughly doubles the GPU memory baseline, so a GPU_SMALL (16 GB) that comfortably fits one model may need to step up to GPU_MEDIUM (24 GB) to fit both. You're saving on endpoint count but potentially paying more per endpoint. All rows in a single request must share the same _task — the wrapper reads it from the first row and dispatches the entire batch to that model.

A note on Databricks' native traffic routing: Model Serving does support multiple served entities behind a single endpoint, but traffic is split by percentage (e.g., 70/30 for A/B testing), not by request content. There is no built-in "route classification requests to model A, regression requests to model B" mechanism — which is why the dual-task wrapper handles dispatching in application code instead.

Approach GPU cost Complexity When to use
Two endpoints 2× GPU (mitigated by scheduled warm-up/cool-down) Low — each endpoint is isolated Most teams; operational simplicity wins
Dual-task wrapper 1× GPU (but likely a larger tier) Medium — shared wrapper, _task routing Cost-sensitive setups or when you want a single URL for all tabular predictions

Request Payload Examples

Once the endpoint is live, every request is an HTTP POST with a JSON body containing a dataframe_split key — the format Databricks Model Serving uses for DataFrame-oriented payloads. The four columns (_data, _role, _label, _task) map directly to the dual-task signature we defined above.

Classification request — customer churn

Four context rows teach the model a churn pattern; two query rows get predictions. Note _task: "classification" on every row:

",  "context", "0", "classification"],
      ["",  "context", "1", "classification"],
      ["",  "context", "0", "classification"],
      ["", "context", "1", "classification"],
      ["",  "query",   "", "classification"],
      ["",  "query",   "", "classification"]
    ]
  }
}

The response returns predictions only for the query rows:

,
    
  ]
}

Regression request — house price estimation

Same endpoint, completely different task and features. Note _task: "regression" and numeric _label values (passed as strings):

", "context", "250000",  "regression"],
      ["", "context", "520000",  "regression"],
      ["", "context", "380000",  "regression"],
      ["", "context", "650000",  "regression"],
      ["", "query",   "",        "regression"],
      ["", "query",   "",        "regression"]
    ]
  }
}

The response returns continuous numeric predictions:

,
    
  ]
}

Classification request — different domain, same endpoint

This is where TabFM's flexibility truly shines. The same endpoint handles sensor-anomaly classification with entirely different features — no redeployment, no new model version:

", "context", "normal",  "classification"],
      ["", "context", "anomaly", "classification"],
      ["", "context", "normal",  "classification"],
      ["", "context", "anomaly", "classification"],
      ["", "query",   "",        "classification"]
    ]
  }
}

Three examples, three completely different feature sets and label vocabularies, two different task types — all served by a single endpoint with a single model artifact.

Calling the endpoint from Python

In practice, you'd call the endpoint using the Databricks SDK or a simple HTTP request. Use _task to switch between classification and regression on the same endpoint:

import json
import requests

endpoint_url = "https:///serving-endpoints/tabfm-dual-serving-endpoint/invocations"
token = dbutils.notebook.entry_point.getDbutils().notebook().getContext().apiToken().get()
headers = ", "Content-Type": "application/json"}

# Classification request
clf_payload = {
    "dataframe_split": {
        "columns": ["_data", "_role", "_label", "_task"],
        "data": [
            [json.dumps(), "context", "0", "classification"],
            [json.dumps(), "context", "1", "classification"],
            [json.dumps(), "query", "", "classification"],
        ]
    }
}

# Regression request (same endpoint!)
reg_payload = ), "context", "250000", "regression"],
            [json.dumps(), "context", "520000", "regression"],
            [json.dumps(), "query", "", "regression"],
        ]
    }
}

response = requests.post(endpoint_url, headers=headers, json=clf_payload)
print("Classification:", response.json())

response = requests.post(endpoint_url, headers=headers, json=reg_payload)
print("Regression:", response.json())

Key things to note about the payload format:

  • The examples above use dataframe_split (column names listed once, data as a 2D array), but Databricks Model Serving also accepts dataframe_records — an array of row dictionaries that's more readable for small payloads:
    ", "_role": "context", "_label": "0", "_task": "classification"},
        {"_data": "", "_role": "context", "_label": "1", "_task": "classification"},
        {"_data": "",  "_role": "query",   "_label": "",  "_task": "classification"}
      ]
    }

    Both formats produce the same DataFrame inside predict(). Use dataframe_split when payload size matters (column names aren't repeated per row); use dataframe_records when readability matters (e.g., the Send request box in the Serving UI).

  • _task must be "classification" or "regression". All rows in a single request must share the same _task — the wrapper reads it from the first row and dispatches the entire batch to that model.
  • _data values must be valid JSON strings — double-escaped when embedded inside the outer JSON body. In Python, use json.dumps() to avoid manual escaping.
  • _label for query rows can be an empty string or any placeholder — the wrapper ignores it when _role is "query". For regression context rows, _label must be a numeric value passed as a string (e.g., "250000").
  • Context and query rows can be interleaved in any order; the wrapper partitions by _role, not by position.
  • There is no hard limit on the number of context rows, but more context means a larger forward pass and higher latency. In practice, 20–100 context rows per class is a reasonable range for TabFM.

Serving Endpoint GPU Requirements

TabFM is a transformer. Every inference call runs a full forward pass over all context and query rows through multi-head attention — and attention scales quadratically with sequence length. On CPU, this is slow enough to make real-time serving impractical for anything beyond toy payloads. GPU acceleration isn't a nice-to-have; it's effectively required for production use.

Workload Types in Databricks Model Serving

Databricks Model Serving offers three GPU tiers, each mapping to different hardware:

Workload type GPU GPU memory Best for
GPU_SMALL NVIDIA T4 16 GB Small-to-medium context sizes (up to ~200 rows), light concurrent traffic
GPU_MEDIUM NVIDIA A10G 24 GB Larger context windows, moderate concurrency, or multiple ensemble members in memory
GPU_LARGE NVIDIA A100 40/80 GB Very large contexts (500+ rows), high throughput, or latency-critical production

The dev or staging environments usually go well with GPU_SMALL and workload_size="Small":

ServedEntityInput(
    entity_name=MODEL_NAME,
    entity_version=str(latest_version),
    scale_to_zero_enabled=True,
    workload_size="Small",
    workload_type=ServingModelWorkloadType.GPU_SMALL,
)

This is a reasonable starting point for dev and experimentation. The T4's 16 GB of GPU memory comfortably holds the TabFM checkpoint (~1–2 GB depending on precision) plus the attention KV cache for typical context sizes.

Context size Ensemble members Recommended minimum
20–100 rows Default GPU_SMALL (T4, 16 GB)
100–300 rows Default GPU_MEDIUM (A10G, 24 GB)
300+ rows or high concurrency Default or increased GPU_LARGE (A100, 40+ GB)

These are rough guidelines — actual memory usage depends on the number of features per row, label cardinality, and whether KV cache quantisation is enabled.

If you deploy to a non-GPU endpoint (omitting workload_type), the model will run on CPU. This works for testing the serving contract — verifying payloads, debugging the wrapper, checking JSON parsing — but expect latency in the range of seconds per request rather than milliseconds. Don't use CPU for any workload where response time matters.

Provisioning Time

GPU endpoints take longer to provision than CPU ones. Set a 45-minute timeout on create_and_wait(), which accounts for:

  • Container image build (first deployment or dependency changes)
  • GPU instance allocation (subject to cloud provider capacity)
  • Model download and load_context() execution (loading the TabFM checkpoint from HuggingFace or the MLflow artifact store)

Subsequent deployments of new model versions to the same endpoint are faster since the container image is cached. Scale-from-zero events (when the endpoint wakes up after being idle) typically take 5–15 minutes, mostly spent on GPU provisioning — this is the cold-start cost discussed in the next section.

Right-Sizing Recommendations

Start with GPU_SMALL / workload_size="Small" for dev and staging. Monitor two signals to know when to scale up:

  1. GPU memory utilisation — if you see OOM errors or the endpoint becomes unresponsive with larger context payloads, step up to GPU_MEDIUM.
  2. Latency under load — if P95 latency exceeds your SLA with concurrent requests, either increase workload_size (adds replicas behind a load balancer) or move to a larger GPU tier.

For production, GPU_MEDIUM with workload_size="Medium" is a solid default that balances cost, memory headroom, and throughput. Reserve GPU_LARGE for use cases with unusually large context windows or stricter latency requirements.

Cost Optimisation for Dev Environments

For most production environments, a TabFM serving endpoint should be always-on with scale-to-zero disabled — where latency matters, and traffic is continuous. But in a dev or staging environment, the usage pattern is very different: the endpoint sits idle every evening, night, and weekend while you pay for a GPU that nobody is using.

Databricks Model Serving gives you two knobs:

Configuration Behaviour Cost Developer experience
Scale-to-zero disabled Endpoint stays warm 24/7 Full cost even during idle hours Instant responses, no cold starts
Scale-to-zero enabled Endpoint shuts down after idle timeout Pay only for active hours Cold start on first request after idle — typically 5–15 minutes to re-provision the GPU and reload the model

Neither option is great for dev. You either burn budget overnight or frustrate developers with cold starts every morning and through the day when they sit down and fire the first request after the endpoint has been scaled down.

The Middle Ground: Scheduled Warm-Up and Cool-Down

The idea is simple: disable scale-to-zero before business hours (so the endpoint is warm when developers arrive) and re-enable it after business hours (so the GPU shuts down overnight). Two short notebooks and two scheduled jobs are all it takes.

Notebook 1: Morning warm-up (disable scale-to-zero)

This notebook runs before the workday starts — say, 07:30. It patches the endpoint config to set scale_to_zero_enabled=False, which wakes the endpoint and keeps it alive:

import requests
import json

# Configuration
WORKSPACE_URL = "https://.cloud.databricks.com"
ENDPOINT_NAME = "tabfm-serving-endpoint"
TOKEN = dbutils.secrets.get(scope="your-scope", key="serving-token")

headers = ",
    "Content-Type": "application/json"
}

# Get current endpoint config
get_response = requests.get(
    f"/api/2.0/serving-endpoints/",
    headers=headers
)
get_response.raise_for_status()
endpoint_config = get_response.json()

# Extract the served model config
served_entities = endpoint_config["config"]["served_entities"]

# Disable scale-to-zero — keep the endpoint warm during the workday
payload = 
        for entity in served_entities
    ]
}

update_response = requests.put(
    f"/api/2.0/serving-endpoints//config",
    headers=headers,
    json=payload
)
update_response.raise_for_status()
print(f"Scale-to-zero DISABLED for  — endpoint will stay warm.")

Notebook 2: Evening cool-down (enable scale-to-zero)

This notebook runs after the workday ends — say, 19:00. It patches the same config to re-enable scale-to-zero, so the GPU shuts down once idle:

import requests
import json

# Configuration (same as above)
WORKSPACE_URL = "https://.cloud.databricks.com"
ENDPOINT_NAME = "tabfm-serving-endpoint"
TOKEN = dbutils.secrets.get(scope="your-scope", key="serving-token")

headers = ",
    "Content-Type": "application/json"
}

# Get current endpoint config
get_response = requests.get(
    f"/api/2.0/serving-endpoints/",
    headers=headers
)
get_response.raise_for_status()
endpoint_config = get_response.json()

served_entities = endpoint_config["config"]["served_entities"]

# Enable scale-to-zero — let the endpoint shut down overnight
payload = 
        for entity in served_entities
    ]
}

update_response = requests.put(
    f"/api/2.0/serving-endpoints//config",
    headers=headers,
    json=payload
)
update_response.raise_for_status()
print(f"Scale-to-zero ENABLED for  — endpoint will shut down after idle timeout.")

Scheduling the notebooks

Create two Lakeflow Jobs, each running one notebook on a cron schedule:

Job Notebook Schedule Purpose
tabfm-warmup-weekday Morning warm-up 0 30 7 ? * MON-FRI (07:30 weekdays) Wake endpoint before devs arrive
tabfm-cooldown-weekday Evening cool-down 0 0 19 ? * MON-FRI (19:00 weekdays) Let endpoint scale down overnight

Both jobs can run on serverless compute (the notebooks just make REST calls — no GPU needed). The endpoint itself handles the heavy lifting.

What this buys you:

  • During business hours (07:30–19:00 weekdays): Endpoint is warm. Zero cold starts. Developers get instant predictions.
  • Evenings, nights, weekends: Endpoint scales to zero after the idle timeout. No GPU cost.
  • Effective cost: Roughly 50–60% of always-on, depending on your timezone and work patterns. For a GPU_SMALL workload type, that's a meaningful saving on a dev endpoint that would otherwise run 168 hours/week but only gets used for 50–60.

Wrapping Up

This guide covered the end-to-end journey of taking Google Research's TabFM — a pretrained transformer for tabular data — and turning it into a live, GPU-backed prediction endpoint on Databricks.

We started with what TabFM is and why its in-context learning approach is fundamentally different from traditional ML: no training loop, no persistent state, no redeployment when your data changes. We walked through the fit() and predict() API, then went deeper into the source code to confirm that fit() truly doesn't update model weights — it stages preprocessing and optionally builds a KV cache, all of which is fully replaced on the next call.

From there, we tackled the practical challenges of deployment: why MLflow is the mandatory packaging layer, how to collapse TabFM's two-method interface into a single predict() call that the serving endpoint can route to, and how to use a JSON-encoded schema to preserve TabFM's ability to handle arbitrary feature sets within MLflow's fixed-signature contract.

Finally, we covered the operational side — choosing the right GPU tier for your workload, understanding what drives memory usage, and a scheduling pattern that keeps dev endpoints warm during business hours while scaling to zero overnight to cut costs.

The result is an endpoint that can serve predictions across completely different datasets and feature schemas, with no retraining and no redeployment — just context rows and query rows in every request. And with the dual-task wrapper from Challenge 3, a single endpoint can handle both classification and regression, dispatched per-request via a _task field in the payload.