{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install sagemaker xgboost==1.5.1 scikit-learn install sm-serverless-benchmarking -Uqq"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train and deploy an XGBooost Model\n",
    "In this example, we'll train an XGBoost model on an synthetic dataset and then deploy it for benchmarking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tarfile\n",
    "\n",
    "\n",
    "import sagemaker\n",
    "import xgboost as xgb\n",
    "import numpy as np\n",
    "from sagemaker.model import Model\n",
    "from sagemaker.serializers import CSVSerializer\n",
    "from sklearn.datasets import make_classification\n",
    "\n",
    "role = (\n",
    "    sagemaker.get_execution_role()\n",
    ")  # manually provide role if using non role based identity\n",
    "sess = sagemaker.Session()\n",
    "region = sess.boto_region_name\n",
    "bucket = sess.default_bucket()\n",
    "image_uri = sagemaker.image_uris.retrieve(\n",
    "    framework=\"xgboost\", region=region, version=\"1.5-1\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# train a model on some synthetic data\n",
    "X, y = make_classification(n_samples=1000, n_features=20)\n",
    "data = xgb.DMatrix(data=X, label=y)\n",
    "bst = xgb.train(params={}, dtrain=data)\n",
    "bst.save_model(\"./xgboost-model\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# package the model for deployment and upload to S3\n",
    "with tarfile.open(\"model.tar.gz\", \"w:gz\") as f:\n",
    "    f.add(\"xgboost-model\")\n",
    "model_uri = sess.upload_data(\n",
    "    \"model.tar.gz\", bucket=bucket, key_prefix=\"sm-sl-bench-xgboost\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm_model = Model(image_uri=image_uri, model_data=model_uri, role=role)\n",
    "sm_model.create(instance_type=\"ml.m5.xlarge\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Validate Endpoint\n",
    "Before launching a full benchmarking job, it is a good idea to first deploy the model on a test endpoint to ensure everything is functioning as it should. Here we will deploy a temporary endpoint and test it with an example payload. Afterwards, the endpoint is deleted. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create a temporary endpoint\n",
    "from sm_serverless_benchmarking.endpoint import ServerlessEndpoint\n",
    "\n",
    "endpoint = ServerlessEndpoint(model_name=sm_model.name, memory_size=6144)\n",
    "endpoint.create_endpoint()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# convert first 10 records into csv format\n",
    "ser = CSVSerializer()\n",
    "payload = ser.serialize(X[:10])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# invoke endpoint and print predictions\n",
    "response = endpoint.invoke_endpoint({\"Body\": payload, \"ContentType\": \"text/csv\"})\n",
    "print(response[\"Body\"].read().decode(\"utf-8\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint.clean_up()  # delete the endpoint"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Launch Benchmarking SageMaker Job"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sm_serverless_benchmarking.utils import convert_invoke_args_to_jsonl\n",
    "from sm_serverless_benchmarking.sagemaker_runner import run_as_sagemaker_job\n",
    "\n",
    "# we'll use 20 random inputs\n",
    "sample_inputs = X[np.random.choice(X.shape[0], size=20, replace=False)]\n",
    "\n",
    "example_invoke_args = [\n",
    "    {\"Body\": ser.serialize(inp), \"ContentType\": \"text/csv\"} for inp in sample_inputs\n",
    "]\n",
    "\n",
    "example_invoke_file = convert_invoke_args_to_jsonl(example_invoke_args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "processor = run_as_sagemaker_job(\n",
    "    role=role,\n",
    "    model_name=sm_model.name,\n",
    "    invoke_args_examples_file=example_invoke_file,\n",
    "    stability_benchmark_invocations=2500,\n",
    "    concurrency_benchmark_invocations=2500,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\n",
    "    f\"Once the job is finished, the outputs will be uploaded to {processor.latest_job.outputs[0].destination}\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can optionally run the command below to copy all of the benchmark output artifacts into the current directory. The primary report output will be under the `benchmarking_report/` directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!aws s3 cp --recursive {processor.latest_job.outputs[0].destination} ."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run a Local Benchmarking Job [OPTIONAL]\n",
    "You can also run the same benchmark locally "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sm_serverless_benchmarking.benchmark import run_serverless_benchmarks\n",
    "\n",
    "report = run_serverless_benchmarks(\n",
    "    model_name=sm_model.name, invoke_args_examples_file=example_invoke_file\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "37058495916f7ab2e7db9963171426deb73c0dc04073ed3a56b3427789bc2f48"
  },
  "kernelspec": {
   "display_name": "Python 3.8.13 ('serverless-benchmarking')",
   "language": "python",
   "name": "python3"
  },
  "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.8.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}