Update xgboost churn neo example for sagemaker v2 (#1591)

* Update xgboost churn neo example for sagemaker v2

* Remove use of latest version

* Add sagemaker installation command and remove duplicate import

* Use sagemaker pysdk v2

* Add setup and cleanup steps

* clear output

* Revert kernel metadata

Co-authored-by: Nikhil Kulkarni <[email protected]>

Author: nikhil-sk

sagemaker_neo_compilation_jobs/xgboost_customer_churn/xgboost_customer_churn_neo.ipynb

@@ -36,7 +36,42 @@
     "\n",
     "## Setup\n",
     "\n",
-    "_This notebook was created and tested on an ml.m4.xlarge notebook instance._\n",
+    "_This notebook was created and tested on an ml.m4.xlarge notebook instance. It assumes that the 'sagemaker' package is already installed in the environment._ \n",
+    "\n",
+    "_It has been tested on the following versions of important python packages:_\n",
+    "\n",
+    "_sagemaker>=2.14_,\n",
+    "_pandas==1.1.2_"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import sys\n",
+    "import sagemaker\n",
+    "import importlib\n",
+    "orig_sm_version = sagemaker.__version__\n",
+    "\n",
+    "with open('orig_sm_version.txt', \"w\") as f:\n",
+    "    f.write(orig_sm_version)\n",
+    "\n",
+    "%pip install --no-cache-dir -qU \"sagemaker>=2.14.0\""
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Please restart the kernel to use the updated sagemaker package and continue below"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
     "\n",
     "Let's start by specifying:\n",
     "\n",
@@ -86,8 +121,7 @@
     "import json\n",
     "from IPython.display import display\n",
     "from time import strftime, gmtime\n",
-    "import sagemaker\n",
-    "from sagemaker.predictor import csv_serializer"
+    "from sagemaker.serializers import CSVSerializer"
    ]
   },
   {
@@ -325,8 +359,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "from sagemaker.amazon.amazon_estimator import get_image_uri\n",
-    "container = get_image_uri(boto3.Session().region_name, 'xgboost')"
+    "from sagemaker import image_uris\n",
+    "container = image_uris.retrieve(framework='xgboost', region=boto3.Session().region_name, version='1')"
    ]
   },
   {
@@ -342,8 +376,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "s3_input_train = sagemaker.s3_input(s3_data='s3://{}/{}/train'.format(bucket, prefix), content_type='csv')\n",
-    "s3_input_validation = sagemaker.s3_input(s3_data='s3://{}/{}/validation/'.format(bucket, prefix), content_type='csv')"
+    "s3_input_train = sagemaker.inputs.TrainingInput(s3_data='s3://{}/{}/train'.format(bucket, prefix), content_type='csv')\n",
+    "s3_input_validation = sagemaker.inputs.TrainingInput(s3_data='s3://{}/{}/validation/'.format(bucket, prefix), content_type='csv')"
    ]
   },
   {
@@ -370,8 +404,8 @@
     "\n",
     "xgb = sagemaker.estimator.Estimator(container,\n",
     "                                    role, \n",
-    "                                    train_instance_count=1, \n",
-    "                                    train_instance_type='ml.m4.xlarge',\n",
+    "                                    instance_count=1, \n",
+    "                                    instance_type='ml.m4.xlarge',\n",
     "                                    output_path='s3://{}/{}/output'.format(bucket, prefix),\n",
     "                                    sagemaker_session=sess)\n",
     "xgb.set_hyperparameters(max_depth=5,\n",
@@ -393,7 +427,7 @@
     "---\n",
     "## Host\n",
     "\n",
-    "Now that we've trained the algorithm, let's create a model and deploy it to a hosted endpoint."
+    "Now that we've trained the algorithm, let's create a model and deploy it to a hosted endpoint. We'll also need to setup serializers and deserializers for passing our `test_data` NumPy arrays to the model behind the endpoint."
    ]
   },
   {
@@ -403,7 +437,7 @@
    "outputs": [],
    "source": [
     "xgb_predictor = xgb.deploy(initial_instance_count=1,\n",
-    "                           instance_type='ml.m4.xlarge')"
+    "                           instance_type='ml.m4.xlarge', serializer=CSVSerializer())"
    ]
   },
   {
@@ -412,25 +446,14 @@
    "source": [
     "### Evaluate\n",
     "\n",
-    "Now that we have a hosted endpoint running, we can make real-time predictions from our model very easily, simply by making an http POST request.  But first, we'll need to setup serializers and deserializers for passing our `test_data` NumPy arrays to the model behind the endpoint."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "xgb_predictor.content_type = 'text/csv'\n",
-    "xgb_predictor.serializer = csv_serializer\n",
-    "xgb_predictor.deserializer = None"
+    "Now that we have a hosted endpoint running, we can make real-time predictions from our model very easily, simply by making an http POST request."
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Now, we'll use a simple function to:\n",
+    "We'll use a simple function to:\n",
     "1. Loop over our test dataset\n",
     "1. Split it into mini-batches of rows \n",
     "1. Convert those mini-batchs to CSV string payloads\n",
@@ -449,10 +472,10 @@
     "    predictions = ''\n",
     "    for array in split_array:\n",
     "        predictions = ','.join([predictions, xgb_predictor.predict(array).decode('utf-8')])\n",
-    "\n",
+    "    \n",
     "    return np.fromstring(predictions[1:], sep=',')\n",
     "\n",
-    "predictions = predict(test_data.as_matrix()[:, 1:])"
+    "predictions = predict(test_data.values[:, 1:])"
    ]
   },
   {
@@ -604,6 +627,15 @@
     "**Important. If the following command result in a permission error, scroll up and locate the value of execution role returned by `get_execution_role()`. The role must have access to the S3 bucket specified in ``output_path``.**"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "'/'.join(xgb.output_path.split('/')[:-1])"
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": null,
@@ -615,7 +647,7 @@
     "                                   input_shape={'data':[1, 69]},\n",
     "                                   role=role,\n",
     "                                   framework='xgboost',\n",
-    "                                   framework_version='0.7',\n",
+    "                                   framework_version='0.9',\n",
     "                                   output_path=output_path)"
    ]
   },
@@ -627,7 +659,9 @@
     "\n",
     "We can deploy this compiled model, note that we need to use the same instance that the target we used for compilation. This creates a SageMaker endpoint that we can use to perform inference. \n",
     "\n",
-    "The arguments to the ``deploy`` function allow us to set the number and type of instances that will be used for the Endpoint. Make sure to choose an instance for which you have compiled your model, so in our case  `ml_c5`. Neo API uses a special runtime (DLR runtime), in which our optimzed model will run."
+    "The arguments to the ``deploy`` function allow us to set the number and type of instances that will be used for the Endpoint. Make sure to choose an instance for which you have compiled your model, so in our case  `ml_c5`. Neo API uses a special runtime (DLR runtime), in which our optimzed model will run.\n",
+    "\n",
+    "The compiled model accepts CSV content type:"
    ]
   },
   {
@@ -638,18 +672,15 @@
    "source": [
     "# known issue: need to manually specify endpoint name\n",
     "compiled_model.name = 'deployed-xgboost-customer-churn'\n",
-    "# There is a known issue where SageMaker SDK locates the incorrect docker image URI for XGBoost\n",
-    "# For now, we manually set Image URI\n",
-    "compiled_model.image = get_image_uri(sess.boto_region_name, 'xgboost-neo', repo_version='latest')\n",
-    "compiled_predictor = compiled_model.deploy(initial_instance_count = 1, instance_type = 'ml.c5.4xlarge')"
+    "\n",
+    "compiled_predictor = compiled_model.deploy(initial_instance_count = 1, instance_type = 'ml.c5.4xlarge', serializer=CSVSerializer())"
    ]
   },
   {
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "### Making an inference request\n",
-    "The compiled model accepts CSV content type:"
+    "### Making an inference request\n"
    ]
   },
   {
@@ -658,9 +689,22 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "compiled_predictor.content_type = 'text/csv'\n",
-    "compiled_predictor.serializer = csv_serializer\n",
-    "compiled_predictor.deserializer = None"
+    "def optimized_predict(data, rows=500):\n",
+    "    split_array = np.array_split(data, int(data.shape[0] / float(rows) + 1))\n",
+    "    predictions = ''\n",
+    "    for array in split_array:\n",
+    "        predictions = ','.join([predictions, compiled_predictor.predict(array).decode('utf-8')])\n",
+    "\n",
+    "    return np.fromstring(predictions[1:], sep=',')\n",
+    "\n",
+    "predictions = optimized_predict(test_data.values[:, 1:])"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "### Clean up sagemaker version"
    ]
   },
   {
@@ -669,15 +713,16 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "def optimized_predict(data, rows=500):\n",
-    "    split_array = np.array_split(data, int(data.shape[0] / float(rows) + 1))\n",
-    "    predictions = ''\n",
-    "    for array in split_array:\n",
-    "        predictions = ','.join([predictions, compiled_predictor.predict(array).decode('utf-8')])\n",
-    "\n",
-    "    return np.fromstring(predictions[1:], sep=',')\n",
+    "# rollback the SageMaker Python SDK to the kernel's original version\n",
+    "with open('orig_sm_version.txt', \"r\") as f:\n",
+    "    orig_sm_version = f.read()\n",
     "\n",
-    "predictions = optimized_predict(test_data.as_matrix()[:, 1:])"
+    "print(\"Original version: {}\".format(orig_sm_version))\n",
+    "print(\"Current version: {}\".format(sagemaker.__version__))\n",
+    "s = 'sagemaker=={}'.format(orig_sm_version)\n",
+    "print(\"Rolling back to... {}\".format(s))\n",
+    "%pip install --no-cache-dir -qU {s}\n",
+    "%rm orig_sm_version.txt"
    ]
   },
   {
@@ -695,8 +740,8 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "sagemaker.Session().delete_endpoint(xgb_predictor.endpoint)\n",
-    "sagemaker.Session().delete_endpoint(compiled_predictor.endpoint)"
+    "xgb_predictor.delete_endpoint()\n",
+    "compiled_predictor.delete_endpoint()"
    ]
   }
  ],