【夏休みの自由研究リレー】 Chronos-2をSageMaker AIのリアルタイム推論で使う方法を調査してみた

# AutoGluon の推論ペイロードで pyarrow が必要だったが、
# PyExtensionType エラーが起きたためバージョンを指定した。
!pip uninstall -y pyarrow
!pip install -U "autogluon.cloud>=0.5.0" "pandas[pyarrow]" "pyarrow>=17,<20"

# 共通設定
from autogluon.cloud import bootstrap
from autogluon.cloud import TimeSeriesFoundationModel, TimeSeriesEndpoint
import warnings
from pathlib import Path
import matplotlib.pyplot as plt
import pandas as pd
warnings.filterwarnings("ignore", category=FutureWarning)

INSTANCE_TYPE = "ml.c5.xlarge" # 今回は検証なのでサポートされている中で最も安価なインスタンスタイプ
REGION = "ap-northeast-1"
HISTORY_YEARS = 3   # 使用する履歴データの期間
PERIODS = 365        # 予測期間 = テスト期間（日数）
FREQ = "D"
SEASONAL_PERIOD = 7  # MASE 計算用（日次データの週次季節性）

# 必要リソースの作成、初回のみ
bootstrap()

# 検証データの読み込み
df = pd.read_csv(
    "https://raw.githubusercontent.com/facebook/prophet/main/examples/example_wp_log_peyton_manning.csv"
)
df["ds"] = pd.to_datetime(df["ds"])
df = df.sort_values("ds")

# 直近 HISTORY_YEARS 年に絞る（末尾 PERIODS 日をテストにすると学習は約 2 年）
history_start = df["ds"].max() - pd.DateOffset(years=HISTORY_YEARS)
df = df[df["ds"] >= history_start].reset_index(drop=True)

# 末尾 PERIODS 日をテストデータ（評価用）に分割
df_train = df.iloc[:-PERIODS].reset_index(drop=True)
df_test = df.iloc[-PERIODS:].reset_index(drop=True)

# Chronos-2 をリアルタイムエンドポイントに新規デプロイ
model = TimeSeriesFoundationModel(model_id="chronos-2")
endpoint = model.deploy(instance_type=INSTANCE_TYPE)

# AutoGluon 形式に整形（item_id / timestamp / target）
df_chronos = df_train.rename(columns={"ds": "timestamp", "y": "target"}).copy()
df_chronos["item_id"] = "peyton_manning"

# 日付欠損があると freq を推論できないため、日次に揃える
full_idx = pd.date_range(
    df_chronos["timestamp"].min(),
    df_chronos["timestamp"].max(),
    freq=FREQ,
)
df_chronos = (
    df_chronos.set_index("timestamp")
    .reindex(full_idx)
    .assign(item_id=lambda d: d["item_id"].fillna("peyton_manning"))
)
df_chronos["target"] = df_chronos["target"].interpolate(limit_direction="both")
df_chronos = df_chronos.reset_index(names="timestamp")

# 推論実行
predictions = endpoint.predict(
    data=df_chronos,
    target="target",
    id_column="item_id",
    timestamp_column="timestamp",
    prediction_length=PERIODS,
)

from autogluon.cloud import register

SAGEMAKER_ROLE = "" # あらかじめ作成していたIAMロールARN
S3_BUCKET = "" # あらかじめ作成していたS3バケット名

register(
    role=SAGEMAKER_ROLE,
    bucket=S3_BUCKET,
    region=REGION,
)

import json
import warnings
from pathlib import Path

import boto3
import matplotlib.pyplot as plt
import pandas as pd
# 続く検証で使うためimport
from prophet import Prophet

warnings.filterwarnings("ignore", category=FutureWarning)

ENDPOINT_NAME = "" # デプロイしたエンドポイント
REGION = "ap-northeast-1"
HISTORY_YEARS = 3    # 使用する履歴データの期間
PERIODS = 365        # 予測期間 = テスト期間（日数）
FREQ = "D"
SEASONAL_PERIOD = 7  # MASE 計算用（日次データの週次季節性）

# 学習・テストデータの用意
df = pd.read_csv(
    "https://raw.githubusercontent.com/facebook/prophet/main/examples/example_wp_log_peyton_manning.csv"
)
df["ds"] = pd.to_datetime(df["ds"])
df = df.sort_values("ds")

# 直近 HISTORY_YEARS 年に絞る（末尾 PERIODS 日をテストにすると学習は約 2 年）
history_start = df["ds"].max() - pd.DateOffset(years=HISTORY_YEARS)
df = df[df["ds"] >= history_start].reset_index(drop=True)

# 末尾 PERIODS 日をテストデータ（評価用）に分割
df_train = df.iloc[:-PERIODS].reset_index(drop=True)
df_test = df.iloc[-PERIODS:].reset_index(drop=True)

print(f"全体:  {df['ds'].min().date()} 〜 {df['ds'].max().date()} ({len(df)} 日)")
print(f"Train: {df_train['ds'].min().date()} 〜 {df_train['ds'].max().date()} ({len(df_train)} 日)")
print(f"Test:  {df_test['ds'].min().date()} 〜 {df_test['ds'].max().date()} ({len(df_test)} 日)")
df_train.head()

# Chronos-2での推論実行
client = boto3.client("sagemaker-runtime", region_name=REGION)

payload = {
    "inputs": [
        {
            "target": df_train["y"].tolist(),
            "item_id": "peyton_manning",
            "start": df_train["ds"].iloc[0].strftime("%Y-%m-%d"),
        }
    ],
    "parameters": {
        "prediction_length": PERIODS,
        "freq": FREQ,
        "quantile_levels": [0.1, 0.5, 0.9],
    },
}

response = client.invoke_endpoint(
    EndpointName=ENDPOINT_NAME,
    ContentType="application/json",
    Body=json.dumps(payload),
)
result = json.loads(response["Body"].read().decode("utf-8"))