[新機能] Amazon SageMaker AIがサーバーレスなMLflow機能をサポートしました

import sagemaker
from sagemaker import get_execution_role
from sagemaker.sklearn.estimator import SKLearn

import pandas as pd
import os

sagemaker_session = sagemaker.Session()
role = get_execution_role()
region = sagemaker_session.boto_region_name

tracking_server_arn = f"{アプリのARN}"

!mkdir -p training_code

%%writefile training_code/requirements.txt
mlflow==3.1.4
sagemaker-mlflow==0.2.0

%%writefile training_code/train.py

from __future__ import print_function

import argparse
import joblib
import mlflow
import os
import pandas as pd
from sklearn import datasets
from sklearn import tree
from sklearn.metrics import f1_score, accuracy_score
from sklearn.model_selection import train_test_split

if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument('--max_leaf_nodes', type=int, default=-1)
    parser.add_argument('--output-data-dir', type=str, default=os.environ['SM_OUTPUT_DATA_DIR'])
    parser.add_argument('--model-dir', type=str, default=os.environ['SM_MODEL_DIR'])

    args = parser.parse_args()
    iris = datasets.load_iris()
    train_data = pd.DataFrame(iris.data, columns=iris.feature_names)
    train_data.insert(0, 'target', iris.target)

    mlflow.set_tracking_uri(os.environ['MLFLOW_TRACKING_ARN'])
    mlflow.autolog()

    # labels are in the first column
    y = train_data.iloc[:, 0]
    X = train_data.iloc[:, 1:]

    X_train, X_test, y_train, y_test = train_test_split(
        X, y, test_size=0.2, random_state=46, stratify=y
    )

    max_leaf_nodes = args.max_leaf_nodes

    clf = tree.DecisionTreeClassifier(max_leaf_nodes=max_leaf_nodes)
    clf = clf.fit(X_train, y_train)

    test_pred = clf.predict(X_test)
    test_f1 = f1_score(y_test, test_pred, average='weighted')
    test_acc = accuracy_score(y_test, test_pred)

    mlflow.log_metric("test_f1_score", test_f1)
    mlflow.log_metric("test_accuracy", test_acc)

    mlflow.sklearn.log_model(clf, "model")

    joblib.dump(clf, os.path.join(args.model_dir, 'model.joblib'))

sklearn = SKLearn(
    entry_point="train.py",
    source_dir="training_code",
    framework_version="1.2-1",
    instance_type="ml.c4.xlarge",
    role=role,
    sagemaker_session=sagemaker_session,
    hyperparameters={"max_leaf_nodes": 30},
    keep_alive_period_in_seconds=3600,
    environment={"MLFLOW_TRACKING_ARN": tracking_server_arn},
)

sklearn.fit()