sagemaker-pipelines/tabular/tuning-step/sagemaker-pipelines-tuning-step.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "#### SageMaker Pipelines Tuning Step\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook. \n",
    "\n",
    "![This us-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-2/sagemaker-pipelines|tabular|tuning-step|sagemaker-pipelines-tuning-step.ipynb)\n",
    "\n",
    "---"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "\n",
    "This notebook illustrates how a Hyperparameter Tuning Job can be run as a step in a SageMaker Pipeline.\n",
    "\n",
    "The steps in this pipeline include -\n",
    "* Preprocessing the abalone dataset\n",
    "* Running a Hyperparameter Tuning job\n",
    "* Creating the 2 best models\n",
    "* Evaluating the performance of the top performing model of the HPO step\n",
    "* Registering the top model in the model registry using a conditional step based on evaluation metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "!pip install -U sagemaker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import boto3\n",
    "import sagemaker\n",
    "\n",
    "from sagemaker.estimator import Estimator\n",
    "from sagemaker.inputs import TrainingInput\n",
    "\n",
    "from sagemaker.processing import (\n",
    "    ProcessingInput,\n",
    "    ProcessingOutput,\n",
    "    Processor,\n",
    "    ScriptProcessor,\n",
    ")\n",
    "\n",
    "from sagemaker import Model\n",
    "from sagemaker.xgboost import XGBoostPredictor\n",
    "from sagemaker.sklearn.processing import SKLearnProcessor\n",
    "from sagemaker.model_metrics import (\n",
    "    MetricsSource,\n",
    "    ModelMetrics,\n",
    ")\n",
    "from sagemaker.workflow.parameters import (\n",
    "    ParameterInteger,\n",
    "    ParameterString,\n",
    ")\n",
    "from sagemaker.workflow.pipeline import Pipeline\n",
    "from sagemaker.workflow.properties import PropertyFile\n",
    "from sagemaker.workflow.steps import ProcessingStep, CacheConfig, TuningStep\n",
    "from sagemaker.workflow.model_step import ModelStep\n",
    "from sagemaker.workflow.conditions import ConditionLessThanOrEqualTo\n",
    "from sagemaker.workflow.condition_step import ConditionStep\n",
    "\n",
    "from sagemaker.workflow.functions import Join, JsonGet\n",
    "from sagemaker.workflow.execution_variables import ExecutionVariables\n",
    "from sagemaker.workflow.pipeline_context import PipelineSession\n",
    "from sagemaker.tuner import (\n",
    "    ContinuousParameter,\n",
    "    HyperparameterTuner,\n",
    "    WarmStartConfig,\n",
    "    WarmStartTypes,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Create the SageMaker Session\n",
    "\n",
    "region = sagemaker.Session().boto_region_name\n",
    "sm_client = boto3.client(\"sagemaker\")\n",
    "boto_session = boto3.Session(region_name=region)\n",
    "sagemaker_session = sagemaker.session.Session(boto_session=boto_session, sagemaker_client=sm_client)\n",
    "\n",
    "# Create a Pipeline Session\n",
    "pipeline_session = PipelineSession()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Define variables and parameters needed for the Pipeline steps\n",
    "\n",
    "role = sagemaker.get_execution_role()\n",
    "default_bucket = sagemaker_session.default_bucket()\n",
    "base_job_prefix = \"tuning-step-example\"\n",
    "model_package_group_name = \"tuning-job-model-packages\"\n",
    "\n",
    "processing_instance_count = ParameterInteger(name=\"ProcessingInstanceCount\", default_value=1)\n",
    "training_instance_type = ParameterString(name=\"TrainingInstanceType\", default_value=\"ml.m5.xlarge\")\n",
    "model_approval_status = ParameterString(\n",
    "    name=\"ModelApprovalStatus\", default_value=\"PendingManualApproval\"\n",
    ")\n",
    "input_data = ParameterString(\n",
    "    name=\"InputDataUrl\",\n",
    "    default_value=f\"s3://sagemaker-example-files-prod-{region}/datasets/tabular/uci_abalone/abalone.csv\",\n",
    ")\n",
    "model_approval_status = ParameterString(\n",
    "    name=\"ModelApprovalStatus\", default_value=\"PendingManualApproval\"\n",
    ")\n",
    "\n",
    "# Cache Pipeline steps to reduce execution time on subsequent executions\n",
    "cache_config = CacheConfig(enable_caching=True, expire_after=\"30d\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "#### Data Preparation\n",
    "\n",
    "An SKLearn processor is used to prepare the dataset for the Hyperparameter Tuning job. Using the script `preprocess.py`, the dataset is featurized and split into train, test, and validation datasets.\n",
    "\n",
    "The output of this step is used as the input to the TuningStep"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "!mkdir -p code"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%%writefile code/preprocess.py\n",
    "\n",
    "\"\"\"Feature engineers the abalone dataset.\"\"\"\n",
    "import argparse\n",
    "import logging\n",
    "import os\n",
    "import pathlib\n",
    "import requests\n",
    "import tempfile\n",
    "\n",
    "import boto3\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "from sklearn.compose import ColumnTransformer\n",
    "from sklearn.impute import SimpleImputer\n",
    "from sklearn.pipeline import Pipeline\n",
    "from sklearn.preprocessing import StandardScaler, OneHotEncoder\n",
    "\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.INFO)\n",
    "logger.addHandler(logging.StreamHandler())\n",
    "\n",
    "\n",
    "# Since we get a headerless CSV file we specify the column names here.\n",
    "feature_columns_names = [\n",
    "    \"sex\",\n",
    "    \"length\",\n",
    "    \"diameter\",\n",
    "    \"height\",\n",
    "    \"whole_weight\",\n",
    "    \"shucked_weight\",\n",
    "    \"viscera_weight\",\n",
    "    \"shell_weight\",\n",
    "]\n",
    "label_column = \"rings\"\n",
    "\n",
    "feature_columns_dtype = {\n",
    "    \"sex\": str,\n",
    "    \"length\": np.float64,\n",
    "    \"diameter\": np.float64,\n",
    "    \"height\": np.float64,\n",
    "    \"whole_weight\": np.float64,\n",
    "    \"shucked_weight\": np.float64,\n",
    "    \"viscera_weight\": np.float64,\n",
    "    \"shell_weight\": np.float64,\n",
    "}\n",
    "label_column_dtype = {\"rings\": np.float64}\n",
    "\n",
    "\n",
    "def merge_two_dicts(x, y):\n",
    "    \"\"\"Merges two dicts, returning a new copy.\"\"\"\n",
    "    z = x.copy()\n",
    "    z.update(y)\n",
    "    return z\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    logger.debug(\"Starting preprocessing.\")\n",
    "    parser = argparse.ArgumentParser()\n",
    "    parser.add_argument(\"--input-data\", type=str, required=True)\n",
    "    args = parser.parse_args()\n",
    "\n",
    "    base_dir = \"/opt/ml/processing\"\n",
    "    pathlib.Path(f\"{base_dir}/data\").mkdir(parents=True, exist_ok=True)\n",
    "    input_data = args.input_data\n",
    "    bucket = input_data.split(\"/\")[2]\n",
    "    key = \"/\".join(input_data.split(\"/\")[3:])\n",
    "\n",
    "    logger.info(\"Downloading data from bucket: %s, key: %s\", bucket, key)\n",
    "    fn = f\"{base_dir}/data/abalone-dataset.csv\"\n",
    "    s3 = boto3.resource(\"s3\")\n",
    "    s3.Bucket(bucket).download_file(key, fn)\n",
    "\n",
    "    logger.debug(\"Reading downloaded data.\")\n",
    "    df = pd.read_csv(\n",
    "        fn,\n",
    "        header=None,\n",
    "        names=feature_columns_names + [label_column],\n",
    "        dtype=merge_two_dicts(feature_columns_dtype, label_column_dtype),\n",
    "    )\n",
    "    os.unlink(fn)\n",
    "\n",
    "    logger.debug(\"Defining transformers.\")\n",
    "    numeric_features = list(feature_columns_names)\n",
    "    numeric_features.remove(\"sex\")\n",
    "    numeric_transformer = Pipeline(\n",
    "        steps=[\n",
    "            (\"imputer\", SimpleImputer(strategy=\"median\")),\n",
    "            (\"scaler\", StandardScaler()),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    categorical_features = [\"sex\"]\n",
    "    categorical_transformer = Pipeline(\n",
    "        steps=[\n",
    "            (\"imputer\", SimpleImputer(strategy=\"constant\", fill_value=\"missing\")),\n",
    "            (\"onehot\", OneHotEncoder(handle_unknown=\"ignore\")),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    preprocess = ColumnTransformer(\n",
    "        transformers=[\n",
    "            (\"num\", numeric_transformer, numeric_features),\n",
    "            (\"cat\", categorical_transformer, categorical_features),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "    logger.info(\"Applying transforms.\")\n",
    "    y = df.pop(\"rings\")\n",
    "    X_pre = preprocess.fit_transform(df)\n",
    "    y_pre = y.to_numpy().reshape(len(y), 1)\n",
    "\n",
    "    X = np.concatenate((y_pre, X_pre), axis=1)\n",
    "\n",
    "    logger.info(\"Splitting %d rows of data into train, validation, test datasets.\", len(X))\n",
    "    np.random.shuffle(X)\n",
    "    train, validation, test = np.split(X, [int(0.7 * len(X)), int(0.85 * len(X))])\n",
    "\n",
    "    logger.info(\"Writing out datasets to %s.\", base_dir)\n",
    "    pd.DataFrame(train).to_csv(f\"{base_dir}/train/train.csv\", header=False, index=False)\n",
    "    pd.DataFrame(validation).to_csv(\n",
    "        f\"{base_dir}/validation/validation.csv\", header=False, index=False\n",
    "    )\n",
    "    pd.DataFrame(test).to_csv(f\"{base_dir}/test/test.csv\", header=False, index=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Process the training data step using a python script.\n",
    "# Split the training data set into train, test, and validation datasets\n",
    "# When defining the ProcessingOutput destination as a dynamic value using the\n",
    "# Pipeline Execution ID, caching will not be in effect as each time the step runs,\n",
    "# the step definition changes resulting in new execution. If caching is required,\n",
    "# the ProcessingOutput definition should be status\n",
    "\n",
    "sklearn_processor = SKLearnProcessor(\n",
    "    framework_version=\"1.2-1\",\n",
    "    instance_type=\"ml.m5.xlarge\",\n",
    "    instance_count=processing_instance_count,\n",
    "    base_job_name=f\"{base_job_prefix}/sklearn-abalone-preprocess\",\n",
    "    sagemaker_session=pipeline_session,\n",
    "    role=role,\n",
    ")\n",
    "\n",
    "processor_run_args = sklearn_processor.run(\n",
    "    outputs=[\n",
    "        ProcessingOutput(\n",
    "            output_name=\"train\",\n",
    "            source=\"/opt/ml/processing/train\",\n",
    "            destination=Join(\n",
    "                on=\"/\",\n",
    "                values=[\n",
    "                    \"s3:/\",\n",
    "                    default_bucket,\n",
    "                    base_job_prefix,\n",
    "                    ExecutionVariables.PIPELINE_EXECUTION_ID,\n",
    "                    \"PreprocessAbaloneDataForHPO\",\n",
    "                ],\n",
    "            ),\n",
    "        ),\n",
    "        ProcessingOutput(\n",
    "            output_name=\"validation\",\n",
    "            source=\"/opt/ml/processing/validation\",\n",
    "            destination=Join(\n",
    "                on=\"/\",\n",
    "                values=[\n",
    "                    \"s3:/\",\n",
    "                    default_bucket,\n",
    "                    base_job_prefix,\n",
    "                    ExecutionVariables.PIPELINE_EXECUTION_ID,\n",
    "                    \"PreprocessAbaloneDataForHPO\",\n",
    "                ],\n",
    "            ),\n",
    "        ),\n",
    "        ProcessingOutput(\n",
    "            output_name=\"test\",\n",
    "            source=\"/opt/ml/processing/test\",\n",
    "            destination=Join(\n",
    "                on=\"/\",\n",
    "                values=[\n",
    "                    \"s3:/\",\n",
    "                    default_bucket,\n",
    "                    base_job_prefix,\n",
    "                    ExecutionVariables.PIPELINE_EXECUTION_ID,\n",
    "                    \"PreprocessAbaloneDataForHPO\",\n",
    "                ],\n",
    "            ),\n",
    "        ),\n",
    "    ],\n",
    "    code=\"code/preprocess.py\",\n",
    "    arguments=[\"--input-data\", input_data],\n",
    ")\n",
    "\n",
    "step_process = ProcessingStep(\n",
    "    name=\"PreprocessAbaloneDataForHPO\",\n",
    "    step_args=processor_run_args,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "#### Hyperparameter Tuning\n",
    "\n",
    "Amazon SageMaker automatic model tuning, also known as hyperparameter tuning, finds the best version of a model by running many training jobs on your dataset using the algorithm and ranges of hyperparameters that you specify. It then chooses the hyperparameter values that result in a model that performs the best, as measured by a metric that you choose.\n",
    "\n",
    "[Valid metrics](https://github.com/dmlc/xgboost/blob/master/doc/parameter.rst#learning-task-parameters) for XGBoost Tuning Job\n",
    "\n",
    "You can learn more about [Hyperparameter Tuning](https://docs.aws.amazon.com/sagemaker/latest/dg/automatic-model-tuning-how-it-works.html) in the SageMaker docs."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Define the output path for the model artifacts from the Hyperparameter Tuning Job\n",
    "model_path = f\"s3://{default_bucket}/{base_job_prefix}/AbaloneTrain\"\n",
    "\n",
    "image_uri = sagemaker.image_uris.retrieve(\n",
    "    framework=\"xgboost\",\n",
    "    region=region,\n",
    "    version=\"1.0-1\",\n",
    "    py_version=\"py3\",\n",
    "    instance_type=\"ml.m5.xlarge\",\n",
    ")\n",
    "\n",
    "xgb_train = Estimator(\n",
    "    image_uri=image_uri,\n",
    "    instance_type=training_instance_type,\n",
    "    instance_count=1,\n",
    "    output_path=model_path,\n",
    "    base_job_name=f\"{base_job_prefix}/abalone-train\",\n",
    "    sagemaker_session=pipeline_session,\n",
    "    role=role,\n",
    ")\n",
    "\n",
    "xgb_train.set_hyperparameters(\n",
    "    eval_metric=\"rmse\",\n",
    "    objective=\"reg:squarederror\",  # Define the object metric for the training job\n",
    "    num_round=50,\n",
    "    max_depth=5,\n",
    "    eta=0.2,\n",
    "    gamma=4,\n",
    "    min_child_weight=6,\n",
    "    subsample=0.7,\n",
    "    silent=0,\n",
    ")\n",
    "\n",
    "objective_metric_name = \"validation:rmse\"\n",
    "\n",
    "hyperparameter_ranges = {\n",
    "    \"alpha\": ContinuousParameter(0.01, 10, scaling_type=\"Logarithmic\"),\n",
    "    \"lambda\": ContinuousParameter(0.01, 10, scaling_type=\"Logarithmic\"),\n",
    "}\n",
    "\n",
    "tuner_log = HyperparameterTuner(\n",
    "    xgb_train,\n",
    "    objective_metric_name,\n",
    "    hyperparameter_ranges,\n",
    "    max_jobs=3,\n",
    "    max_parallel_jobs=3,\n",
    "    strategy=\"Random\",\n",
    "    objective_type=\"Minimize\",\n",
    ")\n",
    "\n",
    "hpo_args = tuner_log.fit(\n",
    "    inputs={\n",
    "        \"train\": TrainingInput(\n",
    "            s3_data=step_process.properties.ProcessingOutputConfig.Outputs[\"train\"].S3Output.S3Uri,\n",
    "            content_type=\"text/csv\",\n",
    "        ),\n",
    "        \"validation\": TrainingInput(\n",
    "            s3_data=step_process.properties.ProcessingOutputConfig.Outputs[\n",
    "                \"validation\"\n",
    "            ].S3Output.S3Uri,\n",
    "            content_type=\"text/csv\",\n",
    "        ),\n",
    "    }\n",
    ")\n",
    "\n",
    "step_tuning = TuningStep(\n",
    "    name=\"HPTuning\",\n",
    "    step_args=hpo_args,\n",
    "    cache_config=cache_config,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "#### Warm start for Hyperparameter Tuning Job\n",
    "\n",
    "Use warm start to start a hyperparameter tuning job using one or more previous tuning jobs as a starting point. The results of previous tuning jobs are used to inform which combinations of hyperparameters to search over in the new tuning job. Hyperparameter tuning uses either Bayesian or random search to choose combinations of hyperparameter values from ranges that you specify.\n",
    "\n",
    "Find more information on [Warm Starts](https://docs.aws.amazon.com/sagemaker/latest/dg/automatic-model-tuning-warm-start.html) in the SageMaker docs.\n",
    "\n",
    "In a training pipeline, the parent tuning job name can be provided as a pipeline parameter if there is an already complete Hyperparameter tuning job that should be used as the basis for the warm start.\n",
    "\n",
    "This step is left out of the pipeline steps in this notebook. It can be added into the steps while defining the pipeline and the appropriate parent tuning job should be specified."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# This is an example to illustrate how a the name of the tuning job from the previous step can be used as the parent tuning job, in practice,\n",
    "# it is unlikely to have the parent job run before the warm start job on each run. Typically the first tuning job would run and the pipeline\n",
    "# would be altered to use tuning jobs with a warm start using the first job as the parent job.\n",
    "\n",
    "parent_tuning_job_name = (\n",
    "    step_tuning.properties.HyperParameterTuningJobName\n",
    ")  # Use the parent tuning job specific to the use case\n",
    "\n",
    "warm_start_config = WarmStartConfig(\n",
    "    WarmStartTypes.IDENTICAL_DATA_AND_ALGORITHM, parents={parent_tuning_job_name}\n",
    ")\n",
    "\n",
    "tuner_log_warm_start = HyperparameterTuner(\n",
    "    xgb_train,\n",
    "    objective_metric_name,\n",
    "    hyperparameter_ranges,\n",
    "    max_jobs=3,\n",
    "    max_parallel_jobs=3,\n",
    "    strategy=\"Random\",\n",
    "    objective_type=\"Minimize\",\n",
    "    warm_start_config=warm_start_config,\n",
    ")\n",
    "\n",
    "tuner_run_args = tuner_log_warm_start.fit(\n",
    "    inputs={\n",
    "        \"train\": TrainingInput(\n",
    "            s3_data=step_process.properties.ProcessingOutputConfig.Outputs[\"train\"].S3Output.S3Uri,\n",
    "            content_type=\"text/csv\",\n",
    "        ),\n",
    "        \"validation\": TrainingInput(\n",
    "            s3_data=step_process.properties.ProcessingOutputConfig.Outputs[\n",
    "                \"validation\"\n",
    "            ].S3Output.S3Uri,\n",
    "            content_type=\"text/csv\",\n",
    "        ),\n",
    "    }\n",
    ")\n",
    "\n",
    "step_tuning_warm_start = TuningStep(\n",
    "    name=\"HPTuningWarmStart\",\n",
    "    step_args=tuner_run_args,\n",
    "    cache_config=cache_config,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "#### Creating and Registering the best models\n",
    "\n",
    "After successfully completing the Hyperparameter Tuning job. You can either create SageMaker models from the model artifacts created by the training jobs from the TuningStep or register the models into the Model Registry.\n",
    "\n",
    "When using the Model Registry, if you register multiple models from the TuningStep, they will be registered as versions within the same model package group unless unique model package groups are specified for each `ModelStep` that is part of the pipeline.\n",
    "\n",
    "In this example, the two best models from the TuningStep are added to the same model package group in the Model Registry as v0 and v1.\n",
    "\n",
    "You use the `get_top_model_s3_uri` method of the TuningStep class to get the model artifact from one of the top performing model versions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Creating 2 SageMaker Models\n",
    "model_prefix = f\"{base_job_prefix}/AbaloneTrain\"\n",
    "\n",
    "best_model = Model(\n",
    "    image_uri=image_uri,\n",
    "    model_data=step_tuning.get_top_model_s3_uri(\n",
    "        top_k=0, s3_bucket=default_bucket, prefix=model_prefix\n",
    "    ),\n",
    "    predictor_cls=XGBoostPredictor,\n",
    "    sagemaker_session=pipeline_session,\n",
    "    role=role,\n",
    ")\n",
    "\n",
    "step_create_first = ModelStep(\n",
    "    name=\"CreateBestModel\",\n",
    "    step_args=best_model.create(instance_type=\"ml.m5.xlarge\"),\n",
    ")\n",
    "\n",
    "second_best_model = Model(\n",
    "    image_uri=image_uri,\n",
    "    model_data=step_tuning.get_top_model_s3_uri(\n",
    "        top_k=1, s3_bucket=default_bucket, prefix=model_prefix\n",
    "    ),\n",
    "    predictor_cls=XGBoostPredictor,\n",
    "    sagemaker_session=pipeline_session,\n",
    "    role=role,\n",
    ")\n",
    "\n",
    "step_create_second = ModelStep(\n",
    "    name=\"CreateSecondBestModel\",\n",
    "    step_args=second_best_model.create(instance_type=\"ml.m5.xlarge\"),\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "#### Evaluate the top model\n",
    "\n",
    "Use a processing job to evaluate the top model from the tuning step"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "%%writefile code/evaluate.py\n",
    "\n",
    "\"\"\"Evaluation script for measuring mean squared error.\"\"\"\n",
    "import json\n",
    "import logging\n",
    "import pathlib\n",
    "import pickle\n",
    "import tarfile\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import xgboost\n",
    "\n",
    "from sklearn.metrics import mean_squared_error\n",
    "\n",
    "logger = logging.getLogger()\n",
    "logger.setLevel(logging.INFO)\n",
    "logger.addHandler(logging.StreamHandler())\n",
    "\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "    logger.debug(\"Starting evaluation.\")\n",
    "    model_path = \"/opt/ml/processing/model/model.tar.gz\"\n",
    "    with tarfile.open(model_path) as tar:\n",
    "        tar.extractall(path=\".\")\n",
    "\n",
    "    logger.debug(\"Loading xgboost model.\")\n",
    "    model = pickle.load(open(\"xgboost-model\", \"rb\"))\n",
    "\n",
    "    logger.debug(\"Reading test data.\")\n",
    "    test_path = \"/opt/ml/processing/test/test.csv\"\n",
    "    df = pd.read_csv(test_path, header=None)\n",
    "\n",
    "    logger.debug(\"Reading test data.\")\n",
    "    y_test = df.iloc[:, 0].to_numpy()\n",
    "    df.drop(df.columns[0], axis=1, inplace=True)\n",
    "    X_test = xgboost.DMatrix(df.values)\n",
    "\n",
    "    logger.info(\"Performing predictions against test data.\")\n",
    "    predictions = model.predict(X_test)\n",
    "\n",
    "    logger.debug(\"Calculating mean squared error.\")\n",
    "    mse = mean_squared_error(y_test, predictions)\n",
    "    std = np.std(y_test - predictions)\n",
    "    report_dict = {\n",
    "        \"regression_metrics\": {\n",
    "            \"mse\": {\"value\": mse, \"standard_deviation\": std},\n",
    "        },\n",
    "    }\n",
    "\n",
    "    output_dir = \"/opt/ml/processing/evaluation\"\n",
    "    pathlib.Path(output_dir).mkdir(parents=True, exist_ok=True)\n",
    "\n",
    "    logger.info(\"Writing out evaluation report with mse: %f\", mse)\n",
    "    evaluation_path = f\"{output_dir}/evaluation.json\"\n",
    "    with open(evaluation_path, \"w\") as f:\n",
    "        f.write(json.dumps(report_dict))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# A ProcessingStep is used to evaluate the performance of a selected model from the HPO step. In this case, the top performing model\n",
    "# is evaluated. Based on the results of the evaluation, the model is registered into the Model Registry using a ConditionStep.\n",
    "\n",
    "script_eval = ScriptProcessor(\n",
    "    image_uri=image_uri,\n",
    "    command=[\"python3\"],\n",
    "    instance_type=\"ml.m5.xlarge\",\n",
    "    instance_count=1,\n",
    "    base_job_name=f\"{base_job_prefix}/script-tuning-step-eval\",\n",
    "    sagemaker_session=pipeline_session,\n",
    "    role=role,\n",
    ")\n",
    "\n",
    "evaluation_report = PropertyFile(\n",
    "    name=\"BestTuningModelEvaluationReport\",\n",
    "    output_name=\"evaluation\",\n",
    "    path=\"evaluation.json\",\n",
    ")\n",
    "\n",
    "processor_args = script_eval.run(\n",
    "    inputs=[\n",
    "        ProcessingInput(\n",
    "            source=step_tuning.get_top_model_s3_uri(\n",
    "                top_k=0, s3_bucket=default_bucket, prefix=model_prefix\n",
    "            ),\n",
    "            destination=\"/opt/ml/processing/model\",\n",
    "        ),\n",
    "        ProcessingInput(\n",
    "            source=step_process.properties.ProcessingOutputConfig.Outputs[\"test\"].S3Output.S3Uri,\n",
    "            destination=\"/opt/ml/processing/test\",\n",
    "        ),\n",
    "    ],\n",
    "    outputs=[\n",
    "        ProcessingOutput(output_name=\"evaluation\", source=\"/opt/ml/processing/evaluation\"),\n",
    "    ],\n",
    "    code=\"code/evaluate.py\",\n",
    ")\n",
    "\n",
    "# This can be extended to evaluate multiple models from the HPO step\n",
    "step_eval = ProcessingStep(\n",
    "    name=\"EvaluateTopModel\",\n",
    "    step_args=processor_args,\n",
    "    property_files=[evaluation_report],\n",
    "    cache_config=cache_config,\n",
    ")\n",
    "\n",
    "model_metrics = ModelMetrics(\n",
    "    model_statistics=MetricsSource(\n",
    "        s3_uri=\"{}/evaluation.json\".format(\n",
    "            step_eval.arguments[\"ProcessingOutputConfig\"][\"Outputs\"][0][\"S3Output\"][\"S3Uri\"]\n",
    "        ),\n",
    "        content_type=\"application/json\",\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# Register the model in the Model Registry\n",
    "# Multiple models can be registered into the Model Registry using multiple ModelSteps. These models can either be added to the\n",
    "# same model package group as different versions within the group or the models can be added to different model package groups.\n",
    "\n",
    "register_args = best_model.register(\n",
    "    content_types=[\"text/csv\"],\n",
    "    response_types=[\"text/csv\"],\n",
    "    inference_instances=[\"ml.t2.medium\", \"ml.m5.large\"],\n",
    "    transform_instances=[\"ml.m5.large\"],\n",
    "    model_package_group_name=model_package_group_name,\n",
    "    approval_status=model_approval_status,\n",
    ")\n",
    "\n",
    "\n",
    "step_register_best = ModelStep(name=\"RegisterBestAbaloneModel\", step_args=register_args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# condition step for evaluating model quality and branching execution\n",
    "\n",
    "cond_lte = ConditionLessThanOrEqualTo(\n",
    "    left=JsonGet(\n",
    "        step_name=step_eval.name,\n",
    "        property_file=evaluation_report,\n",
    "        json_path=\"regression_metrics.mse.value\",\n",
    "    ),\n",
    "    right=6.0,\n",
    ")\n",
    "step_cond = ConditionStep(\n",
    "    name=\"CheckMSEAbaloneEvaluation\",\n",
    "    conditions=[cond_lte],\n",
    "    if_steps=[step_register_best],\n",
    "    else_steps=[],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "pipeline = Pipeline(\n",
    "    name=\"tuning-step-pipeline\",\n",
    "    parameters=[\n",
    "        processing_instance_count,\n",
    "        training_instance_type,\n",
    "        input_data,\n",
    "        model_approval_status,\n",
    "    ],\n",
    "    steps=[\n",
    "        step_process,\n",
    "        step_tuning,\n",
    "        step_create_first,\n",
    "        step_create_second,\n",
    "        step_eval,\n",
    "        step_cond,\n",
    "    ],\n",
    "    sagemaker_session=pipeline_session,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "#### Execute the Pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "definition = json.loads(pipeline.definition())\n",
    "definition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "pipeline.upsert(role_arn=role)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "pipeline.start()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {
    "pycharm": {
     "name": "#%% md\n"
    }
   },
   "source": [
    "#### Cleaning up resources\n",
    "\n",
    "Users are responsible for cleaning up resources created when running this notebook. Specify the ModelName, ModelPackageName, and ModelPackageGroupName that need to be deleted. The model names are generated by the CreateModel step of the Pipeline and the property values are available only in the Pipeline context. To delete the models created by this pipeline, navigate to the Model Registry and Console to find the models to delete."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "# # Create a SageMaker client\n",
    "# sm_client = boto3.client(\"sagemaker\")\n",
    "\n",
    "# # Delete SageMaker Models\n",
    "# sm_client.delete_model(ModelName=\"...\")\n",
    "\n",
    "# # Delete Model Packages\n",
    "# sm_client.delete_model_package(ModelPackageName=\"...\")\n",
    "\n",
    "# # Delete the Model Package Group\n",
    "# sm_client.delete_model_package_group(ModelPackageGroupName=\"...\")\n",
    "\n",
    "# # Delete the Pipeline\n",
    "# sm_client.delete_pipeline(PipelineName=\"tuning-step-pipeline\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook CI Test Results\n",
    "\n",
    "This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.\n",
    "\n",
    "![This us-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-1/sagemaker-pipelines|tabular|tuning-step|sagemaker-pipelines-tuning-step.ipynb)\n",
    "\n",
    "![This us-east-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-2/sagemaker-pipelines|tabular|tuning-step|sagemaker-pipelines-tuning-step.ipynb)\n",
    "\n",
    "![This us-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-1/sagemaker-pipelines|tabular|tuning-step|sagemaker-pipelines-tuning-step.ipynb)\n",
    "\n",
    "![This ca-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ca-central-1/sagemaker-pipelines|tabular|tuning-step|sagemaker-pipelines-tuning-step.ipynb)\n",
    "\n",
    "![This sa-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/sa-east-1/sagemaker-pipelines|tabular|tuning-step|sagemaker-pipelines-tuning-step.ipynb)\n",
    "\n",
    "![This eu-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-1/sagemaker-pipelines|tabular|tuning-step|sagemaker-pipelines-tuning-step.ipynb)\n",
    "\n",
    "![This eu-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-2/sagemaker-pipelines|tabular|tuning-step|sagemaker-pipelines-tuning-step.ipynb)\n",
    "\n",
    "![This eu-west-3 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-3/sagemaker-pipelines|tabular|tuning-step|sagemaker-pipelines-tuning-step.ipynb)\n",
    "\n",
    "![This eu-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-central-1/sagemaker-pipelines|tabular|tuning-step|sagemaker-pipelines-tuning-step.ipynb)\n",
    "\n",
    "![This eu-north-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-north-1/sagemaker-pipelines|tabular|tuning-step|sagemaker-pipelines-tuning-step.ipynb)\n",
    "\n",
    "![This ap-southeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-1/sagemaker-pipelines|tabular|tuning-step|sagemaker-pipelines-tuning-step.ipynb)\n",
    "\n",
    "![This ap-southeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-2/sagemaker-pipelines|tabular|tuning-step|sagemaker-pipelines-tuning-step.ipynb)\n",
    "\n",
    "![This ap-northeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-1/sagemaker-pipelines|tabular|tuning-step|sagemaker-pipelines-tuning-step.ipynb)\n",
    "\n",
    "![This ap-northeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-2/sagemaker-pipelines|tabular|tuning-step|sagemaker-pipelines-tuning-step.ipynb)\n",
    "\n",
    "![This ap-south-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-south-1/sagemaker-pipelines|tabular|tuning-step|sagemaker-pipelines-tuning-step.ipynb)\n"
   ]
  }
 ],
 "metadata": {
  "instance_type": "ml.t3.medium",
  "kernelspec": {
   "display_name": "Python 3 (Data Science 3.0)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-east-1:081325390199:image/sagemaker-data-science-310-v1"
  },
  "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.10.6"
  },
  "metadata": {
   "interpreter": {
    "hash": "ac2eaa0ea0ebeafcc7822e65e46aa9d4f966f30b695406963e145ea4a91cd4fc"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}