AWS SageMakerをローカルで使うときに便利なスクリプト

get_predictor.py

from sagemaker.estimator import Estimator
from sagemaker.predictor import RealTimePredictor

def get_predictor(training_job_name: str) -> RealTimePredictor:
    """
    ジョブ名から RealTimePredictor を取得して返す

    すでに訓練ジョブ名と同名のエンドポイントがある場合には、エンドポイントから RealTimePredictor を返す

    Args:
        training_job_name (str): ジョブ名（あるいはエンドポイント名）

    Returns:
        RealtimePredictor: 予測エンドポイント
    """
    try:
        estimator = Estimator.attach(training_job_name=training_job_name)
        predictor = estimator.deploy(
            initial_instance_count=1, instance_type="ml.c4.xlarge"
        )
    except botocore.exceptions.ClientError as e:
        error = e.response["Error"]
        if error["Code"] == "ValidationException" and error["Message"].startswith(
            "Cannot create already existing endpoint"
        ):
            predictor = RealTimePredictor(training_job_name)
        else:
            raise e

    return predictor