Modo local offline do SageMaker — transformação em lote com CatBoost

Inferência em lote contra o moto através de sagemaker-local (sem AWS). Mesmo fluxo de treinamento de training.ipynb, mas a pontuação é feita offline em vez de através de um endpoint:

  • Treinamento: image_uri=sagemaker-catboost:train (Estimator genérico) em diabetes.
  • Serviço: a imagem somente de inferência sagemaker-catboost:inference executa o mesmo ponto de entrada de serviço que um endpoint online, mas o trabalho local de transformação em lote a invoca.
  • Entrada: um CSV local (as primeiras 20 linhas do conjunto de dados de treinamento), sem cabeçalho.
  • Saída: uma previsão por linha de entrada.

O manipulador de entrada padrão MultiRecord agrupa todas as linhas em uma única carga útil, de modo que o manipulador padrão text/csv analisa-os como uma matriz 2-D. SingleRecord enviaria uma linha 1-D por pedido, que esse manipulador rejeita.

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. Treine

O train.py ajusta um modelo no conjunto de dados do scikit-learn indicado pelo hiperparâmetro dataset e grava model.joblib em /opt/ml/model, que o SDK envia ao S3 do moto.

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. Escreva a entrada CSV

Pegue as primeiras 20 linhas de recurso do diabetes e escreva-as como um CSV sem cabeçalho. Cada linha é um registro.

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. Transforme

Um Model apontado para a imagem somente de inferência e os artefatos treinados é encapsulado em um transformer. O trabalho de transformação inicia um contêiner de serviço, pede que ele publique seu contrato de lote via /execution-parameters e depois grava x_test.csv.out ao lado de x_test.csv. O contêiner de serviço produz o arquivo de saída com uma linha por registro de entrada.

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. Leia as previsões

A saída de transformação é uma matriz JSON com uma previsão por linha de entrada, na mesma ordem.

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}")