Introduction

The safest way to keep SageMaker training and inference consistent is to put preprocessing logic in one shared code path and reuse it in both places. If training normalizes one way and inference normalizes another, the model is effectively seeing different features in production than it saw during fitting.

Separate Stateless Logic From Fitted State

Preprocessing usually has two parts:

• stateless transformations, such as trimming strings, selecting columns, or parsing timestamps
• fitted state, such as category mappings, means and standard deviations, or token vocabularies

The stateless logic should live in shared code that both training and inference import. The fitted state should be learned during training, saved as an artifact, and loaded again during inference.

That design avoids a very common mistake: copying the same transformation code into two files and then letting them drift apart over time.

Put Shared Logic in a Common Module

A simple project layout might look like this:

1src/
2  preprocess.py
3  train.py
4  inference.py

The shared preprocessing module can expose functions for fitting and transforming:

1import joblib
2import pandas as pd
3from sklearn.preprocessing import StandardScaler
4
5
6def clean_frame(df: pd.DataFrame) -> pd.DataFrame:
7    df = df.copy()
8    df["age"] = df["age"].fillna(df["age"].median())
9    df["income"] = df["income"].fillna(0)
10    return df[["age", "income"]]
11
12
13def fit_preprocessor(df: pd.DataFrame):
14    features = clean_frame(df)
15    scaler = StandardScaler()
16    scaler.fit(features)
17    return scaler
18
19
20def transform_frame(df: pd.DataFrame, scaler):
21    features = clean_frame(df)
22    return scaler.transform(features)
23
24
25def save_preprocessor(scaler, path: str) -> None:
26    joblib.dump(scaler, path)
27
28
29def load_preprocessor(path: str):
30    return joblib.load(path)

Now there is one definition of the cleaning and feature ordering logic.

Reuse It During Training

Your training entry point should call the shared functions, fit the preprocessor, train the model, and save both artifacts into the SageMaker model directory.

1import os
2import joblib
3import pandas as pd
4from sklearn.linear_model import LogisticRegression
5
6from preprocess import fit_preprocessor, transform_frame, save_preprocessor
7
8
9def train():
10    train_path = "/opt/ml/input/data/training/train.csv"
11    model_dir = "/opt/ml/model"
12
13    df = pd.read_csv(train_path)
14    y = df["label"]
15
16    scaler = fit_preprocessor(df)
17    X = transform_frame(df, scaler)
18
19    model = LogisticRegression()
20    model.fit(X, y)
21
22    joblib.dump(model, os.path.join(model_dir, "model.joblib"))
23    save_preprocessor(scaler, os.path.join(model_dir, "preprocessor.joblib"))
24
25
26if __name__ == "__main__":
27    train()

The important detail is that the training script saves the fitted preprocessor, not just the model weights.

Reuse the Same Logic During Inference

In SageMaker inference, load both the model and the preprocessor from the model artifact. Then call the same shared transform function before prediction.

1import io
2import json
3import os
4
5import joblib
6import pandas as pd
7
8from preprocess import load_preprocessor, transform_frame
9
10
11def model_fn(model_dir):
12    model = joblib.load(os.path.join(model_dir, "model.joblib"))
13    preprocessor = load_preprocessor(os.path.join(model_dir, "preprocessor.joblib"))
14    return {
15        "model": model,
16        "preprocessor": preprocessor,
17    }
18
19
20def input_fn(request_body, content_type):
21    if content_type != "text/csv":
22        raise ValueError("unsupported content type")
23
24    return pd.read_csv(io.StringIO(request_body))
25
26
27def predict_fn(input_data, artifacts):
28    X = transform_frame(input_data, artifacts["preprocessor"])
29    predictions = artifacts["model"].predict(X)
30    return predictions.tolist()
31
32
33def output_fn(prediction, accept):
34    return json.dumps(prediction), "application/json"

That is the core pattern: one preprocessing module, one saved fitted preprocessor, one consistent transform path.

Why This Matters in SageMaker

SageMaker does not automatically guarantee that your training-time feature engineering and inference-time feature engineering match. It gives you the infrastructure hooks, but you still need to structure the code correctly.

Common mismatches include:

• columns reordered between training and inference
• missing values filled differently
• category encoding learned on training but rebuilt incorrectly at serving time
• normalization applied in one path and forgotten in the other

The shared-module approach removes most of those risks.

Consider Using a Pipeline Object

If you are using scikit-learn, a Pipeline or ColumnTransformer can make this even cleaner because the fitted preprocessing and model travel together as one artifact.

1from sklearn.compose import ColumnTransformer
2from sklearn.pipeline import Pipeline
3from sklearn.preprocessing import StandardScaler
4from sklearn.linear_model import LogisticRegression
5
6pipeline = Pipeline([
7    ("features", ColumnTransformer([
8        ("num", StandardScaler(), ["age", "income"]),
9    ])),
10    ("model", LogisticRegression()),
11])

You can fit and serialize that pipeline directly. In many SageMaker setups, that is the simplest way to ensure training and inference do exactly the same thing.

Common Pitfalls

The biggest pitfall is duplicating preprocessing logic in train.py and inference.py. The code starts identical, then one side changes and the model begins receiving different features in production.

Another common mistake is saving only the model and forgetting the fitted preprocessing state. If your encoder or scaler is rebuilt from scratch during inference, predictions can become meaningless.

Hardcoding feature order is another source of subtle bugs. Keep the feature-selection logic in one place so the same columns and order are used everywhere.

Finally, be clear about where preprocessing should run. For lightweight tabular transforms, putting the logic inside the model artifact is usually fine. For heavy image, NLP, or multi-stage ETL workflows, a separate SageMaker Processing or serial inference step may be cleaner.

Summary

• Put shared preprocessing logic in a reusable module.
• Learn preprocessing state during training and save it with the model artifact.
• Load the same preprocessor during SageMaker inference.
• Avoid duplicating feature engineering code across training and serving paths.
• If possible, serialize a single pipeline object so preprocessing and prediction stay coupled.