sagemaker/mlops/step-functions-data-science-sdk/studio_step_functions_mlworkflow_scikit_learn_data_processing_and_model_evaluation_with_experiments.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Amazon SageMaker Processing と AWS Step Functions Data Science SDK で機械学習ワークフローを構築する (SageMaker Studio)\n",
    "\n",
    "Amazon SageMaker Processing を使うと、データの前/後処理やモデル評価のワークロードを Amazon SageMaker platform 上で簡単に実行することができます。Processingジョブは Amazon Simple Storage Service (Amazon S3) から入力データをダウンロードし、処理結果を Amazon S3 にアップロードします。\n",
    "\n",
    "Step Functions SDK は AWS Step Function と Amazon SageMaker を使って、データサイエンティストが機械学習ワークフローを簡単に作成して実行するためのものです。詳しい情報は以下のドキュメントをご参照ください。\n",
    "\n",
    "* [AWS Step Functions](https://aws.amazon.com/step-functions/)\n",
    "* [AWS Step Functions Developer Guide](https://docs.aws.amazon.com/step-functions/latest/dg/welcome.html)\n",
    "* [AWS Step Functions Data Science SDK](https://aws-step-functions-data-science-sdk.readthedocs.io)\n",
    "\n",
    "AWS Step Functions Data Science SDK の SageMaker Processing Step [ProcessingStep](https://aws-step-functions-data-science-sdk.readthedocs.io/en/stable/sagemaker.html#stepfunctions.steps.sagemaker.ProcessingStep) によって、AWS Step Functions ワークフローで実装された Sageaker Processing を機械学習エンジニアが直接システムに統合することができます。\n",
    "\n",
    "このノートブックは、SageMaker Processing Job を使ってデータの前処理、モデルの学習、モデルの精度評価の機械学習ワークフローを AWS Step Functions Data Science SDK を使って作成する方法をご紹介します。大まかな流れは以下の通りです。\n",
    "\n",
    "1. AWS Step Functions Data Science SDK の `ProcessingStep` を使ってデータの前処理、特徴量エンジニアリング、学習用とテスト用への分割を行う scikit-learn スクリプトを実行する SageMaker Processing Job を実行\n",
    "1. AWS Step Functions Data Science SDK の `TrainingStep` を使って前処理された学習データを使ったモデルの学習を実行\n",
    "1. AWS Step Functions Data Science SDK の `ProcessingStep` を使って前処理したテスト用データを使った学習済モデルの評価を実行\n",
    "1. AWS Step Functions Data Science SDK の `LambdaStep`を使って最新のモデルと過去のモデルの評価指標の比較を実行\n",
    "\n",
    "\n",
    "このノートブックで使用するデータは [Census-Income KDD Dataset](https://archive.ics.uci.edu/ml/datasets/Census-Income+%28KDD%29) です。このデータセットから特徴量を選択し、データクレンジングを実施し、二値分類モデルの利用できる形にデータを変換し、最後にデータを学習用とテスト用に分割します。このノートブックではロジスティック回帰モデルを使って、国勢調査の回答者の収入が 5万ドル以上か 5万ドル未満かを予測します。このデータセットはクラスごとの不均衡が大きく、ほとんどのデータに 5万ドル以下というラベルが付加されています。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "このノートブックを実行するのに必要なライブラリをインストールします。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Import the latest sagemaker, stepfunctions and boto3 SDKs\n",
    "import sys\n",
    "\n",
    "!{sys.executable} -m pip install --upgrade pip\n",
    "!{sys.executable} -m pip install -qU awscli boto3 \"sagemaker==2.50.0\"\n",
    "!{sys.executable} -m pip install -qU \"stepfunctions==2.0.0\"\n",
    "!{sys.executable} -m pip install sagemaker-experiments\n",
    "!{sys.executable} -m pip show sagemaker stepfunctions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 必要なモジュールのインポート"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import io\n",
    "import logging\n",
    "import os\n",
    "import random\n",
    "import time\n",
    "\n",
    "import boto3\n",
    "import stepfunctions\n",
    "from stepfunctions import steps\n",
    "from stepfunctions.inputs import ExecutionInput\n",
    "from stepfunctions.steps import (\n",
    "    Chain,\n",
    "    ChoiceRule,\n",
    "    ModelStep,\n",
    "    ProcessingStep,\n",
    "    TrainingStep,\n",
    "    TransformStep,\n",
    ")\n",
    "from stepfunctions.template import TrainingPipeline\n",
    "from stepfunctions.template.utils import replace_parameters_with_jsonpath\n",
    "from stepfunctions.workflow import Workflow\n",
    "\n",
    "import sagemaker\n",
    "from sagemaker import get_execution_role\n",
    "from sagemaker.amazon.amazon_estimator import image_uris\n",
    "from sagemaker.processing import ProcessingInput, ProcessingOutput\n",
    "from sagemaker.s3 import S3Uploader\n",
    "from sagemaker.sklearn.processing import SKLearnProcessor\n",
    "\n",
    "from sagemaker.analytics import ExperimentAnalytics\n",
    "from smexperiments.experiment import Experiment\n",
    "from smexperiments.trial import Trial\n",
    "from smexperiments.trial_component import TrialComponent\n",
    "from smexperiments.tracker import Tracker\n",
    "import time\n",
    "\n",
    "# SageMaker Session\n",
    "sagemaker_session = sagemaker.Session()\n",
    "region = sagemaker_session.boto_region_name\n",
    "\n",
    "# SageMaker Execution Role\n",
    "# You can use sagemaker.get_execution_role() if running inside sagemaker's notebook instance\n",
    "role = get_execution_role()\n",
    "\n",
    "account_id = boto3.client('sts').get_caller_identity().get('Account')\n",
    "bucket = sagemaker_session.default_bucket()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "次に、ノートブックから Step Functions を実行するための IAM ロール設定を行います。\n",
    "\n",
    "### ノートブックインスタンスの IAM ロールに権限を追加\n",
    "\n",
    "https://sagemaker-immersionday.workshop.aws/ja/lab3/option2.html に従って、\n",
    "\n",
    "1. AWS CodeBuild の信頼関係を追加\n",
    "1. https://github.com/aws-samples/amazon-sagemaker-immersion-day/blob/master/iam-policy-sm-cb.txt のポリシーを IAM ロールに追加\n",
    "\n",
    "続けて、IAM ロールの画面で\n",
    "1.  `AWSStepFunctionsFullAccess` のポリシーも追加します。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "次に Step Functions で使用する実行ロールを作成します。\n",
    "\n",
    "### Step Functions の実行ロールの作成\n",
    "\n",
    " 作成した Step Functions ワークフローは、AWS の他のサービスと連携するための IAM ロールを必要とします。\n",
    "\n",
    "1. [IAM console](https://console.aws.amazon.com/iam/) にアクセス\n",
    "2. 左側のメニューの **ロール** を選択し **ロールの作成** をクリック\n",
    "3. **ユースケースの選択** で **Step Functions** をクリック\n",
    "4. **次のステップ：アクセス権限** **次のステップ：タグ** **次のステップ：確認**をクリック\n",
    "5. **ロール名** に `AmazonSageMaker-StepFunctionsWorkflowExecutionRole` と入力して **ロールの作成** をクリック\n",
    "\n",
    "次に、作成したロールに AWS マネージド IAM ポリシーをアタッチします。\n",
    "\n",
    "1. [IAM console](https://console.aws.amazon.com/iam/) にアクセス\n",
    "2. 左側のメニューの **ロール** を選択\n",
    "3. 先ほど作成した `AmazonSageMaker-StepFunctionsWorkflowExecutionRole`を検索\n",
    "4. **ポリシーをアタッチします** をクリックして `CloudWatchEventsFullAccess` を検索\n",
    "5. `CloudWatchEventsFullAccess` の横のチェックボックスをオンにして **ポリシーのアタッチ** をクリック\n",
    "\n",
    "次に、別の新しいポリシーをロールにアタッチします。ベストプラクティスとして、以下のステップで特定のリソースのみのアクセス権限とこのサンプルを実行するのに必要なアクションのみを有効にします。\n",
    "\n",
    "1. 左側のメニューの **ロール** を選択\n",
    "1. 先ほど作成した `AmazonSageMaker-StepFunctionsWorkflowExecutionRole`を検索\n",
    "1. **ポリシーをアタッチします** をクリックして **ポリシーの作成** をクリック\n",
    "1. **JSON** タブをクリックして以下の内容をペースト<br>\n",
    "NOTEBOOK_ROLE_ARN の部分をノートブックインスタンスで使用している IAM ロールの ARN に置き換えてください。\n",
    "\n",
    "```json\n",
    "{\n",
    "    \"Version\": \"2012-10-17\",\n",
    "    \"Statement\": [\n",
    "        {\n",
    "            \"Sid\": \"VisualEditor0\",\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": [\n",
    "                \"events:PutTargets\",\n",
    "                \"events:DescribeRule\",\n",
    "                \"events:PutRule\"\n",
    "            ],\n",
    "            \"Resource\": [\n",
    "                \"arn:aws:events:*:*:rule/StepFunctionsGetEventsForSageMakerTrainingJobsRule\",\n",
    "                \"arn:aws:events:*:*:rule/StepFunctionsGetEventsForSageMakerTransformJobsRule\",\n",
    "                \"arn:aws:events:*:*:rule/StepFunctionsGetEventsForSageMakerTuningJobsRule\",\n",
    "                \"arn:aws:events:*:*:rule/StepFunctionsGetEventsForECSTaskRule\",\n",
    "                \"arn:aws:events:*:*:rule/StepFunctionsGetEventsForBatchJobsRule\"\n",
    "            ]\n",
    "        },\n",
    "        {\n",
    "            \"Sid\": \"VisualEditor1\",\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": \"iam:PassRole\",\n",
    "            \"Resource\": \"NOTEBOOK_ROLE_ARN\",\n",
    "            \"Condition\": {\n",
    "                \"StringEquals\": {\n",
    "                    \"iam:PassedToService\": \"sagemaker.amazonaws.com\"\n",
    "                }\n",
    "            }\n",
    "        },\n",
    "        {\n",
    "            \"Sid\": \"VisualEditor2\",\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": [\n",
    "                \"batch:DescribeJobs\",\n",
    "                \"batch:SubmitJob\",\n",
    "                \"batch:TerminateJob\",\n",
    "                \"dynamodb:DeleteItem\",\n",
    "                \"dynamodb:GetItem\",\n",
    "                \"dynamodb:PutItem\",\n",
    "                \"dynamodb:UpdateItem\",\n",
    "                \"ecs:DescribeTasks\",\n",
    "                \"ecs:RunTask\",\n",
    "                \"ecs:StopTask\",\n",
    "                \"glue:BatchStopJobRun\",\n",
    "                \"glue:GetJobRun\",\n",
    "                \"glue:GetJobRuns\",\n",
    "                \"glue:StartJobRun\",\n",
    "                \"lambda:InvokeFunction\",\n",
    "                \"sagemaker:CreateEndpoint\",\n",
    "                \"sagemaker:CreateEndpointConfig\",\n",
    "                \"sagemaker:CreateHyperParameterTuningJob\",\n",
    "                \"sagemaker:CreateModel\",\n",
    "                \"sagemaker:CreateProcessingJob\",\n",
    "                \"sagemaker:CreateTrainingJob\",\n",
    "                \"sagemaker:CreateTransformJob\",\n",
    "                \"sagemaker:DeleteEndpoint\",\n",
    "                \"sagemaker:DeleteEndpointConfig\",\n",
    "                \"sagemaker:DescribeHyperParameterTuningJob\",\n",
    "                \"sagemaker:DescribeProcessingJob\",\n",
    "                \"sagemaker:DescribeTrainingJob\",\n",
    "                \"sagemaker:DescribeTransformJob\",\n",
    "                \"sagemaker:ListProcessingJobs\",\n",
    "                \"sagemaker:ListTags\",\n",
    "                \"sagemaker:StopHyperParameterTuningJob\",\n",
    "                \"sagemaker:StopProcessingJob\",\n",
    "                \"sagemaker:StopTrainingJob\",\n",
    "                \"sagemaker:StopTransformJob\",\n",
    "                \"sagemaker:UpdateEndpoint\",\n",
    "                \"sns:Publish\",\n",
    "                \"sqs:SendMessage\"\n",
    "            ],\n",
    "            \"Resource\": \"*\"\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "```\n",
    "\n",
    "5. **次のステップ：タグ** **次のステップ：確認**をクリック\n",
    "6. **名前** に `AmazonSageMaker-StepFunctionsWorkflowExecutionPolicy` と入力して **ポリシーの作成** をクリック\n",
    "7. 左側のメニューで **ロール** を選択して `AmazonSageMaker-StepFunctionsWorkflowExecutionRole` を検索\n",
    "8. **ポリシーをアタッチします** をクリック\n",
    "9. 前の手順で作成した `AmazonSageMaker-StepFunctionsWorkflowExecutionPolicy` ポリシーを検索してチェックボックスをオンにして **ポリシーのアタッチ** をクリック\n",
    "11. AmazonSageMaker-StepFunctionsWorkflowExecutionRole の **Role ARN** をコピーして以下のセルにペースト"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# paste the AmazonSageMaker-StepFunctionsWorkflowExecutionRole ARN from above\n",
    "workflow_execution_role = \"arn:aws:iam::\"+account_id+\":role/AmazonSageMaker-StepFunctionsWorkflowExecutionRole\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Step Functions ワークフロー実行時の入力スキーマ作成\n",
    "\n",
    "Step Functions ワークフローを実行する際に、パラメタなどを引数として渡すことができます。ここではそれらの引数のスキーマを作成します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SageMaker expects unique names for each job, model and endpoint.\n",
    "# If these names are not unique the execution will fail. Pass these\n",
    "# dynamically for each execution using placeholders.\n",
    "execution_input = ExecutionInput(\n",
    "    schema={\n",
    "        \"PreprocessingJobName\": str,\n",
    "        \"PreprocessingInputData\": str,\n",
    "        \"PreprocessingOutputDataTrain\": str,\n",
    "        \"PreprocessingOutputDataTest\": str,\n",
    "        \"TrainingJobName\": str,\n",
    "        \"TrainingParameters\": dict,\n",
    "        \"TrainingOutputModel\": str,\n",
    "        \"ExperimentName\": str,\n",
    "        \"EvaluationProcessingJobName\": str,\n",
    "        \"EvaluationProcessingOutput\": str,\n",
    "        \"EvaluationExperimentArgs\": list,\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Amazon SageMaker Experiments のセットアップ\n",
    "\n",
    "このノートブックでは、モデルの評価メトリクスを記録するために Amazon SageMaker Experiments を使用します。`experiment_name` に任意の Experiment 名をセットして以下のセルを実行してください。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "experiment_name = \"sfn-evaluate-model\"\n",
    "\n",
    "# create the experiment if it doesn't exist\n",
    "try:\n",
    "    experiment_evaluate = Experiment.load(experiment_name=experiment_name)\n",
    "except Exception as ex:\n",
    "    if \"ResourceNotFound\" in str(ex):\n",
    "        experiment_evaluate = Experiment.create(\n",
    "            experiment_name=experiment_name, \n",
    "            description=\"model evaluation\", \n",
    "            sagemaker_boto_client=boto3.client('sagemaker'))\n",
    "\n",
    "print(experiment_evaluate.experiment_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## データの前処理と特徴量エンジニアリング"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "データクレンジング 、前処理、特徴量エンジニアリングのスクリプトの前に、データセットの初めの 20行をのぞいてみましょう。ターゲット変数は `income` 列です。選択する特徴量は `age`, `education`, `major industry code`, `class of worker`, `num persons worked for employer`, `capital gains`, `capital losses`,  `dividends from stocks` です。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "\n",
    "input_data = \"s3://sagemaker-sample-data-{}/processing/census/census-income.csv\".format(region)\n",
    "df = pd.read_csv(input_data, nrows=10)\n",
    "df.head(n=10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "scikit-learn の前処理スクリプトを実行するために `SKLearnProcessor`を作成します。これは、SageMaker が用意している scikit-learn のコンテナイメージを使って Processing ジョブを実行するためのクラスです。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sklearn_processor = SKLearnProcessor(\n",
    "    framework_version=\"0.20.0\",\n",
    "    role=role,\n",
    "    instance_type=\"ml.m5.xlarge\",\n",
    "    instance_count=1,\n",
    "    max_runtime_in_seconds=1200,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "以下のセルを実行すると `preprocessing.py` が作成されます。これは前処理のためのスクリプトです。以下のセルを書き換えて実行すれば、`preprocessing.py` が上書き保存されます。このスクリプトでは、以下の処理が実行されます。\n",
    "\n",
    "* 重複データやコンフリクトしているデータの削除\n",
    "* ターゲット変数 `income` 列をカテゴリ変数から 2つのラベルを持つ列に変換\n",
    "* `age` と `num persons worked for employer` をビニングして数値からカテゴリ変数に変換\n",
    "* 連続値である`capital gains`, `capital losses`, `dividends from stocks` を学習しやすいようスケーリング\n",
    "* `education`, `major industry code`, `class of worker`を学習しやすいようエンコード\n",
    "* データを学習用とテスト用に分割し特徴量とラベルの値をそれぞれ保存\n",
    "\n",
    "学習スクリプトでは、前処理済みの学習用データとラベル情報を使用してモデルを学習します。また、モデル評価スクリプトでは学習済みモデルと前処理済みのテスト用データトラベル情報を使用してモデルを評価します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile preprocessing.py\n",
    "\n",
    "import argparse\n",
    "import os\n",
    "import warnings\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.model_selection import train_test_split\n",
    "from sklearn.preprocessing import StandardScaler, OneHotEncoder, LabelBinarizer, KBinsDiscretizer\n",
    "from sklearn.preprocessing import PolynomialFeatures\n",
    "from sklearn.compose import make_column_transformer\n",
    "\n",
    "from sklearn.exceptions import DataConversionWarning\n",
    "\n",
    "warnings.filterwarnings(action=\"ignore\", category=DataConversionWarning)\n",
    "\n",
    "\n",
    "columns = [\n",
    "    \"age\",\n",
    "    \"education\",\n",
    "    \"major industry code\",\n",
    "    \"class of worker\",\n",
    "    \"num persons worked for employer\",\n",
    "    \"capital gains\",\n",
    "    \"capital losses\",\n",
    "    \"dividends from stocks\",\n",
    "    \"income\",\n",
    "]\n",
    "class_labels = [\" - 50000.\", \" 50000+.\"]\n",
    "\n",
    "\n",
    "def print_shape(df):\n",
    "    negative_examples, positive_examples = np.bincount(df[\"income\"])\n",
    "    print(\n",
    "        \"Data shape: {}, {} positive examples, {} negative examples\".format(\n",
    "            df.shape, positive_examples, negative_examples\n",
    "        )\n",
    "    )\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument(\"--train-test-split-ratio\", type=float, default=0.3)\n",
    "    args, _ = parser.parse_known_args()\n",
    "\n",
    "    print(\"Received arguments {}\".format(args))\n",
    "\n",
    "    input_data_path = os.path.join(\"/opt/ml/processing/input\", \"census-income.csv\")\n",
    "\n",
    "    print(\"Reading input data from {}\".format(input_data_path))\n",
    "    df = pd.read_csv(input_data_path)\n",
    "    df = pd.DataFrame(data=df, columns=columns)\n",
    "    df.dropna(inplace=True)\n",
    "    df.drop_duplicates(inplace=True)\n",
    "    df.replace(class_labels, [0, 1], inplace=True)\n",
    "\n",
    "    negative_examples, positive_examples = np.bincount(df[\"income\"])\n",
    "    print(\n",
    "        \"Data after cleaning: {}, {} positive examples, {} negative examples\".format(\n",
    "            df.shape, positive_examples, negative_examples\n",
    "        )\n",
    "    )\n",
    "\n",
    "    split_ratio = args.train_test_split_ratio\n",
    "    print(\"Splitting data into train and test sets with ratio {}\".format(split_ratio))\n",
    "    X_train, X_test, y_train, y_test = train_test_split(\n",
    "        df.drop(\"income\", axis=1), df[\"income\"], test_size=split_ratio, random_state=0\n",
    "    )\n",
    "\n",
    "    preprocess = make_column_transformer(\n",
    "        (\n",
    "            [\"age\", \"num persons worked for employer\"],\n",
    "            KBinsDiscretizer(encode=\"onehot-dense\", n_bins=10),\n",
    "        ),\n",
    "        (\n",
    "            [\"capital gains\", \"capital losses\", \"dividends from stocks\"],\n",
    "            StandardScaler(),\n",
    "        ),\n",
    "        (\n",
    "            [\"education\", \"major industry code\", \"class of worker\"],\n",
    "            OneHotEncoder(sparse=False),\n",
    "        ),\n",
    "    )\n",
    "    print(\"Running preprocessing and feature engineering transformations\")\n",
    "    train_features = preprocess.fit_transform(X_train)\n",
    "    test_features = preprocess.transform(X_test)\n",
    "\n",
    "    print(\"Train data shape after preprocessing: {}\".format(train_features.shape))\n",
    "    print(\"Test data shape after preprocessing: {}\".format(test_features.shape))\n",
    "\n",
    "    train_features_output_path = os.path.join(\"/opt/ml/processing/train\", \"train_features.csv\")\n",
    "    train_labels_output_path = os.path.join(\"/opt/ml/processing/train\", \"train_labels.csv\")\n",
    "\n",
    "    test_features_output_path = os.path.join(\"/opt/ml/processing/test\", \"test_features.csv\")\n",
    "    test_labels_output_path = os.path.join(\"/opt/ml/processing/test\", \"test_labels.csv\")\n",
    "\n",
    "    print(\"Saving training features to {}\".format(train_features_output_path))\n",
    "    pd.DataFrame(train_features).to_csv(train_features_output_path, header=False, index=False)\n",
    "\n",
    "    print(\"Saving test features to {}\".format(test_features_output_path))\n",
    "    pd.DataFrame(test_features).to_csv(test_features_output_path, header=False, index=False)\n",
    "\n",
    "    print(\"Saving training labels to {}\".format(train_labels_output_path))\n",
    "    y_train.to_csv(train_labels_output_path, header=False, index=False)\n",
    "\n",
    "    print(\"Saving test labels to {}\".format(test_labels_output_path))\n",
    "    y_test.to_csv(test_labels_output_path, header=False, index=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "前処理用スクリプトを S3 にアップロードします。Processing ジョブの Input にアップロードした S3 パスを指定することで、ジョブ起動後にそのパスから前処理用スクリプトが Processing 用コンテナにダウンロードされます。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "PREPROCESSING_SCRIPT_LOCATION = \"preprocessing.py\"\n",
    "\n",
    "input_code = sagemaker_session.upload_data(\n",
    "    PREPROCESSING_SCRIPT_LOCATION,\n",
    "    bucket=bucket,\n",
    "    key_prefix=\"data/sklearn_processing/code\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `ProcessingStep` の作成\n",
    "\n",
    "それでは、SageMaker Processing ジョブを起動するための [ProcessingStep](https://aws-step-functions-data-science-sdk.readthedocs.io/en/stable/sagemaker.html#stepfunctions.steps.sagemaker.ProcessingStep) を作成しましょう。\n",
    "\n",
    "このステップは、前の手順で定義した SKLearnProcessor に入力と出力の情報を追加して使用します。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### [ProcessingInputs](https://sagemaker.readthedocs.io/en/stable/api/training/processing.html#sagemaker.processing.ProcessingInput) と [ProcessingOutputs](https://sagemaker.readthedocs.io/en/stable/api/training/processing.html#sagemaker.processing.ProcessingOutput)  オブジェクトを作成して SageMaker Processing ジョブに入力と出力の情報を追加"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "inputs = [\n",
    "    ProcessingInput(\n",
    "        source=execution_input[\"PreprocessingInputData\"], destination=\"/opt/ml/processing/input\", input_name=\"input-1\"\n",
    "    ),\n",
    "    ProcessingInput(\n",
    "        source=input_code,\n",
    "        destination=\"/opt/ml/processing/input/code\",\n",
    "        input_name=\"code\",\n",
    "    ),\n",
    "]\n",
    "\n",
    "outputs = [\n",
    "    ProcessingOutput(\n",
    "        source=\"/opt/ml/processing/train\",\n",
    "        destination=execution_input[\"PreprocessingOutputDataTrain\"],\n",
    "        output_name=\"train_data\",\n",
    "    ),\n",
    "    ProcessingOutput(\n",
    "        source=\"/opt/ml/processing/test\",\n",
    "        destination=execution_input[\"PreprocessingOutputDataTest\"],\n",
    "        output_name=\"test_data\",\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "####  `ProcessingStep` の作成"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "processing_step = ProcessingStep(\n",
    "    \"SageMaker pre-processing step\",\n",
    "    processor=sklearn_processor,\n",
    "    job_name=execution_input[\"PreprocessingJobName\"],\n",
    "    inputs=inputs,\n",
    "    outputs=outputs,\n",
    "    container_arguments=[\"--train-test-split-ratio\", \"0.2\"],\n",
    "    container_entrypoint=[\"python3\", \"/opt/ml/processing/input/code/preprocessing.py\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 前処理済みデータを使ったモデルの学習\n",
    "\n",
    "学習スクリプト `train.py` を使って学習ジョブを実行するための `SKLearn` インスタンスを作成します。これはあとで `TrainingStep` を作成する際に使用します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.sklearn.estimator import SKLearn\n",
    "\n",
    "sklearn = SKLearn(\n",
    "    entry_point=\"train.py\",\n",
    "    train_instance_type=\"ml.m5.xlarge\",\n",
    "    role=role,\n",
    "    framework_version=\"0.20.0\",\n",
    "    py_version=\"py3\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "学習スクリプト `train.py` は、ロジスティック回帰モデルを学習し、学習済みモデルを `/opt/ml/model` に保存します。Amazon SageMaker は、学習ジョブの最後に `/opt/ml/model` に保存されているモデルを `model.tar.gz` に圧縮して S3 にアップロードします。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile train.py\n",
    "\n",
    "import os\n",
    "import argparse\n",
    "\n",
    "import pandas as pd\n",
    "from sklearn.linear_model import LogisticRegression\n",
    "from sklearn.externals import joblib\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    \n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument(\"-c\", type=str, default='1.0')\n",
    "    args, _ = parser.parse_known_args()\n",
    "    \n",
    "    print(args.c)\n",
    "    \n",
    "    training_data_directory = \"/opt/ml/input/data/train\"\n",
    "    train_features_data = os.path.join(training_data_directory, \"train_features.csv\")\n",
    "    train_labels_data = os.path.join(training_data_directory, \"train_labels.csv\")\n",
    "    print(\"Reading input data\")\n",
    "    X_train = pd.read_csv(train_features_data, header=None)\n",
    "    y_train = pd.read_csv(train_labels_data, header=None)\n",
    "\n",
    "    model = LogisticRegression(class_weight=\"balanced\", solver=\"lbfgs\", C=float(args.c), verbose=1)\n",
    "    print(\"Training LR model\")\n",
    "    model.fit(X_train, y_train)\n",
    "    model_output_directory = os.path.join(\"/opt/ml/model\", \"model.joblib\")\n",
    "    print(\"Saving model to {}\".format(model_output_directory))\n",
    "    joblib.dump(model, model_output_directory)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "学習用スクリプトを source.tar.gz に固めて S3 にアップロードします。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "TRAINNING_SCRIPT_LOCATION = \"source.tar.gz\"\n",
    "!tar zcvf $TRAINNING_SCRIPT_LOCATION train.py\n",
    "\n",
    "train_code = sagemaker_session.upload_data(\n",
    "    TRAINNING_SCRIPT_LOCATION,\n",
    "    bucket=bucket,\n",
    "    key_prefix=\"data/sklearn_train/code\",\n",
    ")\n",
    "train_code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### `TrainingStep` の作成"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "training_step = steps.TrainingStep(\n",
    "    \"SageMaker Training Step\",\n",
    "    estimator=sklearn,\n",
    "    data={\"train\": sagemaker.TrainingInput(execution_input[\"PreprocessingOutputDataTrain\"], content_type=\"text/csv\")},\n",
    "    job_name=execution_input[\"TrainingJobName\"],\n",
    "    hyperparameters=execution_input[\"TrainingParameters\"],\n",
    "    wait_for_completion=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## モデルの評価\n",
    "\n",
    "`evaluation.py` はモデル評価用のスクリプトです。このスクリプトは scikit-learn と Amazon SageMake Experiments を用いるため、カスタムコンテナを利用できる`ScriptProcessor` を使用します。このスクリプトは学習済みモデルとテスト用データセットを入力として受け取り、各分類クラスの分類評価メトリクス、precision、リコール、F1スコア、accuracy と ROC AUC が記載された JSON ファイルを出力します。\n",
    "\n",
    "パラメタや学習データを変えて複数のモデルを学習させ、それらの精度を比較する際に評価メトリクスの情報を Amazon SageMaker Experiments で管理しておくと、作成した複数のモデルを比較しやすくなります。以下のモデル評価用のスクリプトでは、モデルの評価メトリクスを算出してそれらを Experiment に登録し、また、最新の学習済みモデルが既存のモデルと比べて良いかどうかを評価メトリクスを使って判定し、その結果も Experiment に記録しています。学習ジョブや Processing ジョブの中で Experiments の Tracker を使う際は、`Tracker.load() ` を使用することでジョブが使用している Trial を自動的に読み出して tracking を開始します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile evaluation.py\n",
    "\n",
    "import json\n",
    "import os\n",
    "import tarfile\n",
    "import argparse\n",
    "\n",
    "import pandas as pd\n",
    "\n",
    "from sklearn.externals import joblib\n",
    "from sklearn.metrics import classification_report, roc_auc_score, accuracy_score\n",
    "\n",
    "from smexperiments.tracker import Tracker\n",
    "from smexperiments.trial import Trial\n",
    "from sagemaker.analytics import ExperimentAnalytics\n",
    "import sagemaker\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    parser = argparse.ArgumentParser()\n",
    "    \n",
    "    parser.add_argument('--experiment-name', type=str, default=None,\n",
    "                        help='Experiment name')\n",
    "\n",
    "    args = parser.parse_args()\n",
    "    \n",
    "    model_path = os.path.join(\"/opt/ml/processing/model\", \"model.tar.gz\")\n",
    "    print(\"Extracting model from path: {}\".format(model_path))\n",
    "    with tarfile.open(model_path) as tar:\n",
    "        tar.extractall(path=\".\")\n",
    "    print(\"Loading model\")\n",
    "    model = joblib.load(\"model.joblib\")\n",
    "\n",
    "    print(\"Loading test input data\")\n",
    "    test_features_data = os.path.join(\"/opt/ml/processing/test\", \"test_features.csv\")\n",
    "    test_labels_data = os.path.join(\"/opt/ml/processing/test\", \"test_labels.csv\")\n",
    "\n",
    "    X_test = pd.read_csv(test_features_data, header=None)\n",
    "    y_test = pd.read_csv(test_labels_data, header=None)\n",
    "    predictions = model.predict(X_test)\n",
    "\n",
    "    print(\"Creating classification evaluation report\")\n",
    "    report_dict = classification_report(y_test, predictions, output_dict=True)\n",
    "    report_dict[\"accuracy\"] = accuracy_score(y_test, predictions)\n",
    "    report_dict[\"roc_auc\"] = roc_auc_score(y_test, predictions)\n",
    "\n",
    "    print(args.experiment_name)\n",
    "    trial_component_analytics = ExperimentAnalytics(\n",
    "        experiment_name=args.experiment_name,\n",
    "        sort_by=\"parameters.accuracy\",\n",
    "        sort_order=\"Descending\",# Ascending or Descending\n",
    "    )\n",
    "    \n",
    "    df = trial_component_analytics.dataframe()\n",
    "    is_best = 0\n",
    "    try:\n",
    "        best_acc = df.iloc[0]['accuracy'] \n",
    "        if best_acc < report_dict[\"accuracy\"]:\n",
    "            print('This model is the best ever!!')\n",
    "            is_best = 1\n",
    "        else:\n",
    "            print('This model is not so good.')\n",
    "    except:\n",
    "        is_best = 1\n",
    "        print('This model is the first one.')\n",
    "    \n",
    "    print('Recording metrics to Experiments...')\n",
    "    with Tracker.load() as processing_tracker: # Tracker requires with keyword\n",
    "        processing_tracker.log_parameters({ \"accuracy\": report_dict[\"accuracy\"], \n",
    "                                                                           \"roc_auc\": report_dict[\"roc_auc\"], \n",
    "                                                                           \"is_best\": is_best})\n",
    "\n",
    "    print(\"Classification report:\\n{}\".format(report_dict))\n",
    "\n",
    "    evaluation_output_path = os.path.join(\"/opt/ml/processing/evaluation\", \"evaluation.json\")\n",
    "    print(\"Saving classification report to {}\".format(evaluation_output_path))\n",
    "\n",
    "    with open(evaluation_output_path, \"w\") as f:\n",
    "        f.write(json.dumps(report_dict))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "評価用スクリプトを S3 にアップロードします。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "MODELEVALUATION_SCRIPT_LOCATION = \"evaluation.py\"\n",
    "\n",
    "input_evaluation_code = sagemaker_session.upload_data(\n",
    "    MODELEVALUATION_SCRIPT_LOCATION,\n",
    "    bucket=bucket,\n",
    "    key_prefix=\"data/sklearn_processing/code\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "モデル評価用の ProcessingStep の入力と出力オブジェクトを作成します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "inputs_evaluation = [\n",
    "    ProcessingInput(\n",
    "        source=execution_input[\"PreprocessingOutputDataTest\"],\n",
    "        destination=\"/opt/ml/processing/test\",\n",
    "        input_name=\"input-1\",\n",
    "    ),\n",
    "    ProcessingInput(\n",
    "        source=execution_input[\"TrainingOutputModel\"],\n",
    "        destination=\"/opt/ml/processing/model\",\n",
    "        input_name=\"input-2\",\n",
    "    ),\n",
    "    ProcessingInput(\n",
    "        source=input_evaluation_code,\n",
    "        destination=\"/opt/ml/processing/input/code\",\n",
    "        input_name=\"code\",\n",
    "    ),\n",
    "]\n",
    "\n",
    "outputs_evaluation = [\n",
    "    ProcessingOutput(\n",
    "        source=\"/opt/ml/processing/evaluation\",\n",
    "        destination=execution_input[\"EvaluationProcessingOutput\"],\n",
    "        output_name=\"evaluation\",\n",
    "    ),\n",
    "]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "モデルの評価結果を SageMaker Experiments で管理するため、Experiments のライブラリを含んだコンテナイメージを作成します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p docker/processing/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile docker/processing/Dockerfile\n",
    "FROM public.ecr.aws/docker/library/python:3.7-slim-buster\n",
    "    \n",
    "ENV AWS_DEFAULT_REGION us-east-1\n",
    "\n",
    "RUN pip3 install --upgrade pip\n",
    "RUN pip3 install -qU boto3 pandas==0.25.3 scikit-learn==0.20.0 sagemaker-experiments sagemaker\n",
    "\n",
    "ENTRYPOINT [\"python3\", \"/opt/ml/processing/input/code/evaluation.py\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ecr_repository = 'sagemaker-studio-sf-model-evaluate'\n",
    "tag = ':latest'\n",
    "uri_suffix = 'amazonaws.com'\n",
    "processing_repository_uri = '{}.dkr.ecr.{}.{}/{}'.format(account_id, region, uri_suffix, ecr_repository + tag)\n",
    "processing_image_name = ecr_repository + tag"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SageMaker Studio では Docker イメージのビルドを行うことができないため、Docker イメージを CodeBuild でビルドし、ECR にプッシュする sm-docker をインストールして利用します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install sagemaker-studio-image-build"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cd ./docker/processing/ && sm-docker build . --repository $processing_image_name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "データの前処理の際は SKProcessor を使用しましたが、モデルの評価にはカスタムコンテナを使用するため ScriptProcessor を使用します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.processing import ScriptProcessor\n",
    "\n",
    "model_evaluation_processor = ScriptProcessor(\n",
    "    command=['python3'],\n",
    "    image_uri=processing_repository_uri,\n",
    "    role=role,\n",
    "    sagemaker_session=sagemaker_session,\n",
    "    instance_count=1,\n",
    "    instance_type='ml.m5.xlarge',\n",
    "    max_runtime_in_seconds=1200\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "ProcessingStep の引数 experiment_config に Experiment 名をセットすることで、起動した Processing ジョブの中で自動的に Trial が作成されます。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "processing_evaluation_step = ProcessingStep(\n",
    "    \"SageMaker Processing Model Evaluation step\",\n",
    "    processor=model_evaluation_processor,\n",
    "    job_name=execution_input[\"EvaluationProcessingJobName\"],\n",
    "    inputs=inputs_evaluation,\n",
    "    outputs=outputs_evaluation,\n",
    "    experiment_config={\n",
    "             \"ExperimentName\":execution_input[\"ExperimentName\"]},\n",
    "    container_arguments=execution_input[\"EvaluationExperimentArgs\"],\n",
    "    container_entrypoint=[\"python3\", \"/opt/ml/processing/input/code/evaluation.py\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## モデルの評価結果に応じた処理\n",
    "\n",
    "モデルの評価ジョブでの評価結果をもとに、Lambda 関数で後処理を行います。後処理としては例えば、Slack にモデルの評価結果をポストしたり、別の Workflow を起動したりなどが考えられます。このノートブックでは、評価結果に応じて表示するテキストの内容を変えています。\n",
    "\n",
    "このノートブックでは、Experiment で管理されている情報を取得するために `sagemaker.analytics.ExperimentAnalytics` を使用するため、必要なライブラリが入ったコンテナを使って Lambda 関数を実行します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p docker/lambda/app"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile ./docker/lambda/app/app.py\n",
    "\n",
    "import sys\n",
    "import time\n",
    "import os\n",
    "import glob\n",
    "import numpy as np\n",
    "import json\n",
    "import boto3\n",
    "from sagemaker.analytics import ExperimentAnalytics\n",
    "\n",
    "def handler(event, context):\n",
    "\n",
    "    experiment_name = event['experiment-name']\n",
    "    job_name = event['evaluation-job-name']\n",
    "    \n",
    "    print('job_name: ', job_name)\n",
    "    \n",
    "    search_expression = {\n",
    "        \"Filters\":[\n",
    "            {\n",
    "                \"Name\": \"TrialComponentName\",\n",
    "                \"Operator\": \"Contains\",\n",
    "                \"Value\": job_name,\n",
    "            }\n",
    "        ],\n",
    "    }\n",
    "\n",
    "    trial_component_analytics = ExperimentAnalytics(\n",
    "        experiment_name=experiment_name,\n",
    "        search_expression=search_expression,\n",
    "    )\n",
    "    \n",
    "    df = trial_component_analytics.dataframe()\n",
    "    print('is_best: ', str(df['is_best']))\n",
    "\n",
    "    result = False\n",
    "    if int(df['is_best']) > 0:\n",
    "        print('This model is the best ever!')\n",
    "        result = True\n",
    "    else:\n",
    "        print('This model is not so good!')\n",
    "    \n",
    "    return {\n",
    "        'statusCode'        : 200,\n",
    "        'result':result\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile ./docker/lambda/Dockerfile\n",
    "\n",
    "# Define custom function directory\n",
    "ARG FUNCTION_DIR=\"/function\"\n",
    "\n",
    "FROM public.ecr.aws/docker/library/python:3.7-slim-buster\n",
    "    \n",
    "# Include global arg in this stage of the build\n",
    "ARG FUNCTION_DIR\n",
    "\n",
    "# Install aws-lambda-cpp build dependencies\n",
    "RUN apt-get update && \\\n",
    "  apt-get install -y \\\n",
    "  g++ \\\n",
    "  make \\\n",
    "  cmake \\\n",
    "  unzip \\\n",
    "  libcurl4-openssl-dev \\\n",
    "  libopencv-dev\n",
    "  \n",
    "RUN pip3 install --upgrade pip\n",
    "RUN pip3 install -qU boto3 pandas==0.25.3 sagemaker-experiments sagemaker\n",
    "\n",
    "# Copy function code\n",
    "RUN mkdir -p ${FUNCTION_DIR}\n",
    "COPY app/ ${FUNCTION_DIR}/\n",
    "\n",
    "# Install the function's dependencies\n",
    "RUN pip install \\\n",
    "    --target ${FUNCTION_DIR} \\\n",
    "        awslambdaric\n",
    "\n",
    "\n",
    "# Set working directory to function root directory\n",
    "WORKDIR ${FUNCTION_DIR}\n",
    "\n",
    "\n",
    "ENTRYPOINT [ \"/usr/local/bin/python\", \"-m\", \"awslambdaric\" ]\n",
    "CMD [ \"app.handler\" ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ecr_repository_lambda = 'sagemaker-studio-sf-lambda'\n",
    "tag = ':latest'\n",
    "uri_suffix = 'amazonaws.com'\n",
    "lambda_repository_uri = '{}.dkr.ecr.{}.{}/{}'.format(account_id, region, uri_suffix, ecr_repository_lambda + tag)\n",
    "lambda_image_name = ecr_repository_lambda + tag"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!cd ./docker/lambda/ && sm-docker build . --repository $lambda_image_name"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "###### コンテナイメージを使って AWS Lambda 関数を作成\n",
    "\n",
    "Lambda で使用するコンテナイメージができたら、コンテナイメージを使って Lambda 関数を作成します。ここからは、ノートブックインスタンスを離れて、AWS Lambda のコンソール操作になります。AWS コンソールから AWS Lambda のコンソールにアクセスしてください。その後、以下の手順を実施して先ほど作成したコンテナイメージが動作する Lambda 関数を作成してください。\n",
    "\n",
    "1. AWS Lambda コンソールで、「関数の作成」をクリック\n",
    "1. 「以下のいずれかのオプションを選択して、関数を作成します。」と書かれた部分で「コンテナイメージ」を選択\n",
    "1. 「関数名」に任意の関数名を入力\n",
    "1. 「イメージを参照」ボタンをクリックして先ほど作成したコンテナイメージを選択<br>\n",
    "「Amazon ECR イメージリポジトリ」のプルダウンメニューに作成したはずのリポジトリがなかったり、コンテナイメージがない場合は一つ上のセルの実行時に何らかのエラーが出ている可能性があるので確認してください。\n",
    "1. 「関数の作成」をクリック\n",
    "\n",
    "#### Lambda 関数の IAM ロールに Experiment 参照権限を追加\n",
    "\n",
    "1. Lambda コンソールで先ほど作成した関数の詳細画面を表示し、設定タブをクリックして左側のメニューの「アクセス権限」をクリック\n",
    "1. **ロール名** と書かれた部分にあるリンクをクリック（IAM のコンソールが表示される）\n",
    "1. **ポリシーをアタッチします** をクリックして **ポリシーの作成** をクリック\n",
    "1. **JSON** タブをクリックして以下の内容をペースト\n",
    "\n",
    "```json\n",
    "{\n",
    "    \"Version\": \"2012-10-17\",\n",
    "    \"Statement\": [\n",
    "        {\n",
    "            \"Sid\": \"VisualEditor0\",\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": \"sagemaker:Search\",\n",
    "            \"Resource\": \"*\"\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "```\n",
    "\n",
    "5. **次のステップ：タグ** **次のステップ：確認**をクリック\n",
    "6. **名前** に `SageMakerSearchPolicy` と入力して **ポリシーの作成** をクリック\n",
    "7. 作成したポリシーを検索して IAM ロールにアタッチ\n",
    "\n",
    "Lambda 関数が作成できたら、以下のセルの `FunctionName` に関数名を入力して実行します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from stepfunctions.steps.states import Retry\n",
    "lambda_step = stepfunctions.steps.compute.LambdaStep(\n",
    "    \"Query Evaluation Results\",\n",
    "    parameters={\n",
    "        \"FunctionName\": 'query_experiment_and_evaluate_container',\n",
    "        \"Payload\": {\n",
    "            \"experiment-name\": execution_input[\"ExperimentName\"],\n",
    "             \"evaluation-job-name\": execution_input[\"EvaluationProcessingJobName\"],\n",
    "            \n",
    "        },\n",
    "    },\n",
    ")\n",
    "lambda_step.add_retry(\n",
    "    Retry(error_equals=[\"States.TaskFailed\"], interval_seconds=15, max_attempts=2, backoff_rate=4.0)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Fail 状態の作成\n",
    "いずれかのステップが失敗したときにワークフローが失敗だとわかるように `Fail` 状態を作成します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "failed_state_sagemaker_processing_failure = stepfunctions.steps.states.Fail(\n",
    "    \"ML Workflow failed\", cause=\"SageMakerProcessingJobFailed\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### ワークフローの中のエラーハンドリングを追加\n",
    "\n",
    "エラーハンドリングのために [Catch Block](https://aws-step-functions-data-science-sdk.readthedocs.io/en/stable/states.html#stepfunctions.steps.states.Catch) を使用します。もし いずれかの Step が失敗したら、`Fail` 状態に遷移します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "catch_state_processing = stepfunctions.steps.states.Catch(\n",
    "    error_equals=[\"States.TaskFailed\"],\n",
    "    next_step=failed_state_sagemaker_processing_failure,\n",
    ")\n",
    "\n",
    "processing_step.add_catch(catch_state_processing)\n",
    "processing_evaluation_step.add_catch(catch_state_processing)\n",
    "training_step.add_catch(catch_state_processing)\n",
    "lambda_step.add_catch(catch_state_processing)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## `Workflow` の作成と実行\n",
    "\n",
    "ここまでで Step Functions のワークフローを作成する準備が完了しました。それでは、ワークフローを作成して実行してみましょう。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime, timezone, timedelta\n",
    "\n",
    "JST = timezone(timedelta(hours=+9), 'JST')\n",
    "\n",
    "timestamp = datetime.now(JST).strftime(\"%Y-%m-%d-%H-%M-%S\")\n",
    "\n",
    "# Generate unique names for Pre-Processing Job, Training Job, and Model Evaluation Job for the Step Functions Workflow\n",
    "training_job_name = \"scikit-learn-training-{}\".format(\n",
    "    timestamp\n",
    ")  # Each Training Job requires a unique name\n",
    "preprocessing_job_name = \"scikit-learn-sm-preprocessing-{}\".format(\n",
    "    timestamp\n",
    ")  # Each Preprocessing job requires a unique name,\n",
    "evaluation_job_name = \"scikit-learn-sm-evaluation-{}\".format(\n",
    "    timestamp\n",
    ")  # Each Evaluation Job requires a unique name\n",
    "timestamp"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Processing ジョブの出力を保存する S3 パスを作成します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 前処理ジョブ用\n",
    "s3_bucket_base_uri = \"{}{}\".format(\"s3://\", bucket)\n",
    "output_data = \"{}/{}-{}\".format(s3_bucket_base_uri, \"data/sklearn_processing/output\", timestamp)\n",
    "preprocessed_training_data = \"{}/{}\".format(output_data, \"train_data\")\n",
    "\n",
    "# モデル評価ジョブ用\n",
    "preprocessed_testing_data = \"{}/{}\".format(output_data, \"test_data\")\n",
    "model_data_s3_uri = \"{}/{}/{}\".format(s3_bucket_base_uri, training_job_name, \"output/model.tar.gz\")\n",
    "output_model_evaluation_s3_uri = \"{}/{}/{}\".format(\n",
    "    s3_bucket_base_uri, training_job_name, \"evaluation\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Chain を使って各 Step を連結してワークフローを作成します。既存のワークフローを変更する場合は、update() を実行します。ログに ERROR が表示された場合は、以下のセルを再度実行してください。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow_graph = Chain([processing_step, training_step, processing_evaluation_step, lambda_step])\n",
    "# workflow_graph = Chain([training_step])\n",
    "# workflow_graph = Chain([processing_evaluation_step,  lambda_step])\n",
    "# workflow_graph = Chain([ lambda_step])\n",
    "branching_workflow = Workflow(\n",
    "    name=\"SageMakerProcessingWorkflow\",\n",
    "    definition=workflow_graph,\n",
    "    role=workflow_execution_role,\n",
    ")\n",
    "\n",
    "branching_workflow.create()\n",
    "branching_workflow.update(workflow_graph)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "パラメタを指定して、ワークフローを実行します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Execute workflow\n",
    "execution = branching_workflow.execute(\n",
    "    inputs={\n",
    "        \"PreprocessingJobName\": preprocessing_job_name,  # Each pre processing job (SageMaker processing job) requires a unique name,\n",
    "        \"PreprocessingInputData\": input_data,\n",
    "        \"PreprocessingOutputDataTrain\": output_data+'/train_data',\n",
    "        \"PreprocessingOutputDataTest\": output_data+'/test_data',\n",
    "        \"TrainingJobName\": training_job_name,  # Each Sagemaker Training job requires a unique name,\n",
    "        \"TrainingParameters\": {\n",
    "                                     \"sagemaker_program\": \"train.py\",\n",
    "                                     \"sagemaker_submit_directory\":  train_code,\n",
    "                                      \"c\": '1.1'\n",
    "},\n",
    "        \"TrainingOutputModel\": model_data_s3_uri,\n",
    "        \"ExperimentName\": experiment_evaluate.experiment_name,\n",
    "        \"EvaluationProcessingJobName\": evaluation_job_name,  # Each SageMaker processing job requires a unique name,\n",
    "        \"EvaluationProcessingOutput\": output_model_evaluation_s3_uri,\n",
    "        \"EvaluationExperimentArgs\": ['--experiment-name', experiment_evaluate.experiment_name]\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "以下のセルを実行することで、Workflow の進行状況がわかります。実行開始から12分程度で完了します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "execution.render_progress()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Workflow の実行が完了したら、Experiment の中をのぞいてみましょう。まずは Experiment のデータを ExperimentAnalytics を使って読み出します。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "search_expression = {\n",
    "    \"Filters\":[\n",
    "        {\n",
    "            \"Name\": \"TrialComponentName\",\n",
    "            \"Operator\": \"Contains\",\n",
    "            \"Value\": evaluation_job_name,\n",
    "        }\n",
    "    ],\n",
    "}\n",
    "\n",
    "trial_component_analytics = ExperimentAnalytics(\n",
    "    experiment_name=experiment_evaluate.experiment_name,\n",
    "    sort_by=\"parameters.accuracy\",\n",
    "#     search_expression=search_expression,\n",
    "#     sort_by=\"metrics.acc.max\",\n",
    "#     sort_order=\"Ascending\",# Ascending or Descending\n",
    "#     metric_names=['metric1', 'metric2'],\n",
    "#     parameter_names=['accuracy', 'roc_auc'],\n",
    "    input_artifact_names=[]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "読み出したデータを DataFrame 形式に変換して表示します。accuracy や roc_auc が記録されていることがわかります。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "df = trial_component_analytics.dataframe()\n",
    "pd.set_option('display.max_columns', None)\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ワークフローの出力を確認\n",
    "\n",
    "Amazon S3 から `evaluation.json` を取得して確認します。ここにはモデルの評価レポートが書かれています。なお、以下のセルは Step Functions でワークフローの実行が完了してから（`evaluation.json` が出力されてから）実行してください。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow_execution_output_json = execution.get_output(wait=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.s3 import S3Downloader\n",
    "import json\n",
    "\n",
    "evaluation_s3_uri = output_model_evaluation_s3_uri + '/evaluation.json'\n",
    "evaluation_output = S3Downloader.read_file(evaluation_s3_uri)\n",
    "evaluation_output_dict = json.loads(evaluation_output)\n",
    "print(json.dumps(evaluation_output_dict, sort_keys=True, indent=4))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step Functions Workflow を実行のみ行う場合\n",
    "\n",
    "ここまでで、Step Functions Data Science SDK を使って Workflow を作成し、実行するところまでご紹介しました。実際に使用する際は、作成済みの Workflow をパラメタを指定して実行する部分を繰り返すことになります。ここでは、作成済みの Workflow を実行する部分のみを抜き出してご紹介します。\n",
    "\n",
    "Workflow の ARN を使って既存の Workflow を読み出します。ARN はわからないけれど Workflow 名はわかる、という場合は以下のセルを実行することで ARN を取得することができます。（`workflow_name` に Workflow 名を入れてください）"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "workflow_name = 'SageMakerProcessingWorkflow'\n",
    "workflow_list = Workflow.list_workflows()\n",
    "workflow_arn = [d['stateMachineArn'] for d in workflow_list  if d['name']==workflow_name][0]\n",
    "workflow_arn"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Workflow を読み出して `execute` で実行します。実行する際にパラメタを指定することができます。SageMaker ジョブはユニークな名前である必要があるため、このノートブックでは datetime を使ってジョブ名を生成しています。以下のセルの `train_code` には、学習ジョブで使用するスクリプトをアップロードしてある S3 パスを設定してください。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datetime import datetime, timezone, timedelta\n",
    "from stepfunctions.workflow import Workflow\n",
    "\n",
    "region = sagemaker_session.boto_region_name\n",
    "input_data = \"s3://sagemaker-sample-data-{}/processing/census/census-income.csv\".format(region)\n",
    "# train_code = 's3://xxx/code/source.tar.gz'\n",
    "\n",
    "bucket = sagemaker_session.default_bucket()\n",
    "\n",
    "JST = timezone(timedelta(hours=+9), 'JST')\n",
    "\n",
    "timestamp = datetime.now(JST).strftime(\"%Y-%m-%d-%H-%M-%S\")\n",
    "\n",
    "# Generate unique names for Pre-Processing Job, Training Job, and Model Evaluation Job for the Step Functions Workflow\n",
    "training_job_name = \"scikit-learn-training-{}\".format(\n",
    "    timestamp\n",
    ")  # Each Training Job requires a unique name\n",
    "preprocessing_job_name = \"scikit-learn-sm-preprocessing-{}\".format(\n",
    "    timestamp\n",
    ")  # Each Preprocessing job requires a unique name,\n",
    "evaluation_job_name = \"scikit-learn-sm-evaluation-{}\".format(\n",
    "    timestamp\n",
    ")  # Each Evaluation Job requires a unique name\n",
    "\n",
    "# 前処理ジョブ用\n",
    "s3_bucket_base_uri = \"{}{}\".format(\"s3://\", bucket)\n",
    "output_data = \"{}/{}-{}\".format(s3_bucket_base_uri, \"data/sklearn_processing/output\", timestamp)\n",
    "preprocessed_training_data = \"{}/{}\".format(output_data, \"train_data\")\n",
    "\n",
    "# モデル評価ジョブ用\n",
    "preprocessed_testing_data = \"{}/{}\".format(output_data, \"test_data\")\n",
    "model_data_s3_uri = \"{}/{}/{}\".format(s3_bucket_base_uri, training_job_name, \"output/model.tar.gz\")\n",
    "output_model_evaluation_s3_uri = \"{}/{}/{}\".format(\n",
    "    s3_bucket_base_uri, training_job_name, \"evaluation\"\n",
    ")\n",
    "\n",
    "existing_workflow = Workflow.attach(workflow_arn)\n",
    "execution = existing_workflow.execute(\n",
    "    inputs={\n",
    "        \"PreprocessingJobName\": preprocessing_job_name,  # Each pre processing job (SageMaker processing job) requires a unique name,\n",
    "        \"PreprocessingInputData\": input_data,\n",
    "        \"PreprocessingOutputDataTrain\": output_data+'/train_data',\n",
    "        \"PreprocessingOutputDataTest\": output_data+'/test_data',\n",
    "        \"TrainingJobName\": training_job_name,  # Each Sagemaker Training job requires a unique name,\n",
    "        \"TrainingParameters\": {\n",
    "                                     \"sagemaker_program\": \"train.py\",\n",
    "                                     \"sagemaker_submit_directory\":  train_code,\n",
    "                                      \"c\": '1.1'\n",
    "        },\n",
    "        \"TrainingOutputModel\": model_data_s3_uri,\n",
    "        \"ExperimentName\": experiment_evaluate.experiment_name,\n",
    "        \"EvaluationProcessingJobName\": evaluation_job_name,  # Each SageMaker processing job requires a unique name,\n",
    "        \"EvaluationProcessingOutput\": output_model_evaluation_s3_uri,\n",
    "        \"EvaluationExperimentArgs\": ['--experiment-name', experiment_evaluate.experiment_name]\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "以下のセルを実行することで、Workflow の進行状況がわかります。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "execution.render_progress()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## リソースの削除\n",
    "\n",
    "このノートブックの実行が終わったら、不要なリソースを削除することを忘れないでください。以下のコードのコメントアウトを外してから実行すると、このノートブックで作成した Step Functions のワークフローを削除することができます。ノートブックインスタンス、各種データを保存した S3 バケットも不要であれば削除してください。\n",
    "\n",
    "### Step Functions Workflow の削除"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# branching_workflow.delete()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Experiment の削除"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm = boto3.Session().client('sagemaker')\n",
    "def cleanup(experiment):\n",
    "    for trial_summary in experiment.list_trials():\n",
    "        trial = Trial.load(sagemaker_boto_client=sm, trial_name=trial_summary.trial_name)\n",
    "        for trial_component_summary in trial.list_trial_components():\n",
    "            tc = TrialComponent.load(\n",
    "                sagemaker_boto_client=sm,\n",
    "                trial_component_name=trial_component_summary.trial_component_name)\n",
    "            trial.remove_trial_component(tc)\n",
    "            try:\n",
    "                # comment out to keep trial components\n",
    "                tc.delete()\n",
    "            except:\n",
    "                # tc is associated with another trial\n",
    "                continue\n",
    "            # to prevent throttling\n",
    "            time.sleep(.5)\n",
    "        trial.delete()\n",
    "    experiment.delete()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "cleanup(experiment_evaluate)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "instance_type": "ml.t3.medium",
  "kernelspec": {
   "display_name": "Python 3 (Data Science)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-east-1:081325390199:image/datascience-1.0"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}