Offline SageMaker local mode — CatBoost batch transform

Batch inference against moto through sagemaker-local (no AWS). Same training flow as training.ipynb, scored offline instead of through an endpoint:

  • Train: image_uri=sagemaker-catboost:train (generic Estimator) on diabetes.
  • Serve: the inference-only image sagemaker-catboost:inference runs the same serving entry point as an online endpoint, but the local batch-transform job calls it.
  • Input: a local CSV (first 20 rows of the training dataset), no header.
  • Output: one prediction per input row.

The default MultiRecord batch strategy groups all rows into one payload, so the default text/csv input handler parses them as a 2-D matrix. SingleRecord would send one 1-D row per request, which that handler rejects.

Cell 02Code
import json
import os
import shutil
from dataclasses import replace

import numpy as np
from sagemaker.estimator import Estimator
from sagemaker.model import Model
from sagemaker_local.config import config_from_env
from sagemaker_local.patches import cleanup_stale_serving_containers
from sagemaker_local.session import make_local_session
from sklearn.datasets import load_diabetes

# A killed process can leave a serving container bound to the shared
# 8080 host port; drop any such leftovers before we start (idempotent).
cleanup_stale_serving_containers()

PROJECT_DIR = os.path.join(
    os.environ.get("SAGEMAKER_LOCAL_REPO_PATH", "/workspace"),
    "projects",
    "sagemaker_catboost",
)
INPUT_DIR = os.path.join(PROJECT_DIR, "data", "input")
OUTPUT_DIR = os.path.join(PROJECT_DIR, "data", "output")

cfg = replace(
    config_from_env(),
    bucket="sagemaker-catboost",
    image_tag="sagemaker-catboost:train",
)
boto_session, sm_session = make_local_session(cfg)
TRAIN_IMAGE = "sagemaker-catboost:train"
INFERENCE_IMAGE = "sagemaker-catboost:inference"

1. Train

train.py fits a model on the scikit-learn dataset named by the dataset hyperparameter and writes model.joblib to /opt/ml/model, which the SDK uploads to moto S3.

Cell 04Code
est = Estimator(
    entry_point="train.py",
    source_dir=os.path.join(PROJECT_DIR, "src", "sagemaker_catboost"),
    image_uri=TRAIN_IMAGE,
    role=cfg.role_arn,
    instance_type="local",
    instance_count=1,
    sagemaker_session=sm_session,
    output_path=f"s3://{cfg.bucket}/models",
    hyperparameters={"dataset": "diabetes"},
)
est.fit()

2. Write the input CSV

Take the first 20 feature rows from diabetes and write them as a headerless CSV. Each line is one record.

Cell 06Code
x_test = load_diabetes().data[:20]
os.makedirs(INPUT_DIR, exist_ok=True)
input_csv = os.path.join(INPUT_DIR, "x_test.csv")
np.savetxt(input_csv, x_test, delimiter=",")
assert x_test.shape[1] == 10, (
    f"expected 10 features per row, got {x_test.shape[1]}"
)
print(f"wrote {len(x_test)} rows to {input_csv}")

3. Transform

A Model pointed at the inference-only image and the trained artifacts is wrapped in a transformer. The transform job starts a serving container, asks it to publish its batch contract via /execution-parameters, then writes x_test.csv.out next to x_test.csv. The serving container produces the output file, one row per input record.

Cell 08Code
shutil.rmtree(OUTPUT_DIR, ignore_errors=True)  # fresh output each run

transformer = Model(
    image_uri=INFERENCE_IMAGE,
    model_data=est.model_data,
    role=cfg.role_arn,
    sagemaker_session=sm_session,
    name="sagemaker-catboost-batch-model",
).transformer(
    instance_count=1,
    instance_type="local",
    output_path=f"file:{OUTPUT_DIR}",
)
transformer.transform(
    f"file://{INPUT_DIR}",
    split_type="Line",
    content_type="text/csv",
)

4. Read the predictions

The transform output is a JSON array with one prediction per input row, in the same order.

Cell 10Code
predictions_path = os.path.join(OUTPUT_DIR, "x_test.csv.out")
predictions = json.loads(open(predictions_path, encoding="utf-8").read())

assert len(predictions) == len(x_test), (
    f"expected one prediction per input row: "
    f"{len(x_test)} rows -> {len(predictions)} predictions"
)
assert all(np.isfinite(float(p)) for p in predictions)
print(f"{len(predictions)} predictions from {predictions_path}")
for i, prediction in enumerate(predictions[:5]):
    print(f"row {i}: {prediction:.4f}")