{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Tensorflow V2로 학습한 모델을 SageMaker로 배포하기\n",
    "\n",
    "본 노트북에서는 학습된 모델을 SageMaker endpoint로 배포하는 프로세스를 살펴봅니다. [첫번째 노트북](1.mnist_train.ipynb)에서 매직명령어 %store% 로 저장했던 `model_data`의 모델 아티팩트를 로드하여 사용합니다. (만약 이전에 생성한 모델 아티팩트가 없다면 공개 S3 버킷에서 해당 파일을 다운로드하게 됩니다.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker \n",
    "sagemaker.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# setups\n",
    "\n",
    "import os\n",
    "import json\n",
    "\n",
    "import sagemaker\n",
    "from sagemaker.tensorflow import TensorFlowModel\n",
    "from sagemaker import get_execution_role, Session\n",
    "import boto3\n",
    "\n",
    "# Get global config\n",
    "with open('code/config.json', 'r') as f:\n",
    "    CONFIG=json.load(f)\n",
    "\n",
    "sess = Session()\n",
    "role = get_execution_role()\n",
    "\n",
    "%store -r tf_mnist_model_data\n",
    "\n",
    "\n",
    "# store -r 시도 후 모델이 없는 경우 publc s3 bucket에서 다운로드\n",
    "try: \n",
    "    tf_mnist_model_data\n",
    "except NameError:\n",
    "    import json\n",
    "    # copy a pretrained model from a public public to your default bucket\n",
    "    s3 = boto3.client('s3')\n",
    "    bucket = CONFIG['public_bucket']\n",
    "    key = 'datasets/image/MNIST/model/tensorflow-training-2020-11-20-23-57-13-077/model.tar.gz'\n",
    "    s3.download_file(bucket, key, 'model.tar.gz')\n",
    "    tf_mnist_model_data = sess.upload_data(\n",
    "        path='model.tar.gz', bucket=sess.default_bucket(), key_prefix='model/tensorflow')\n",
    "    os.remove('model.tar.gz')\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(tf_mnist_model_data)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## TensorFlow Model Object\n",
    "\n",
    "SageMaker에서 제공하는 `TensorFlowModel` 클래스는 여러분의 모델 아티팩트를 이용하여 추론을 실행하는 환경을 정의하도록 해 줍니다. 이는 [첫번째 노트북](1.mnist_train.ipynb)에서 `TensorFlow` estimator를 정의했던 것과 유사한 방식으로, 학습된 모델을 SageMaker에서 호스팅하도록 도커 이미지를 정의하는 하이레벨 API입니다. \n",
    "\n",
    "해당 API를 통해 모델을 추론할 환경을 설정하고 나면 SageMaker에서 관리하는 EC2 인스턴스에서 SageMaker Endpoint 형태로 실행할 수 있습니다. SageMaker Endpoint는 학습된 모델을 RESTful API를 통해 추론하도록 하는 컨테이너기반 환경입니다. \n",
    "\n",
    "`TensorFlowModel` 클래스를 초기화할 때 사용되는 파라미터들은 다음과 같습니다.\n",
    "- role: AWS 리소스 사용을 위한 An IAM 역할(role) \n",
    "- model_data: 압축된 모델 아티팩트가 있는 S3 bucket URI. local mode로 실행시에는 로컬 파일경로 사용가능함\n",
    "- framework_version: 사용하 프레임워크의 버전\n",
    "- py_version: 파이썬 버전"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "model = TensorFlowModel(\n",
    "    role=role,\n",
    "    model_data=tf_mnist_model_data,\n",
    "    framework_version='2.3.1'\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 추론 컨테이너 실행\n",
    "\n",
    "`TensorFlowModel` 클래스가 초기화되고 나면 `deploy`메소드를 이용하여 호스팅용 컨테이너를 실행할 수 있습니다.  \n",
    "\n",
    "`deploy`메소드 실행시 사용되는 파라미터들은 다음과 같습니다.\n",
    "- initial_instance_count: 호스팅 서비스에 사용할 SageMaker 인스턴스의 숫자 \n",
    "- instance_type: 호스팅 서비스를 실행할 SageMaker 인스턴스 타입. 이 값을 `local` 로 선택하면 로컬 인스턴스(SageMaker Jupyter notebook)에 호스팅 컨테이너가 실행됩니다. local mode는 주로 디버깅 단계에서 사용하게 됩니다. \n",
    "\n",
    "<span style=\"color:red\"> 주의 : SageMaker Studio 환경에서는 local mode 가 지원되지 않습니다. </span>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# from sagemaker.serializers import JSONSerializer\n",
    "# from sagemaker.deserializers import JSONDeserializer\n",
    "\n",
    "instance_type='ml.c4.xlarge'\n",
    "\n",
    "predictor = model.deploy(\n",
    "    initial_instance_count=1,\n",
    "    instance_type=instance_type,\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SageMaker endpoint를 이용한 예측 실행\n",
    "\n",
    "`model.deploy(...)`에 의해 리턴된 `Predictor` 인스턴스를 이용하여 예측 요청을 endpoint에 보낼 수 있습니다. 이 경우 모델은 정규화 된 배치 이미지를 받습니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# use some dummy inputs\n",
    "import numpy as np\n",
    "\n",
    "dummy_inputs = {\n",
    "    'instances': np.random.rand(4, 28, 28, 1).tolist()\n",
    "}\n",
    "\n",
    "res = predictor.predict(dummy_inputs)\n",
    "print(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "입출력 데이터 포맷이 [TensorFlow Serving REST API](https://www.tensorflow.org/tfx/serving/api_rest)의 `Predict`에서 정의된 request, respoinst 포맷과 일치하는 지 확인합니다. \n",
    "\n",
    "예를 들어 본 코드에서 `dummy_inputs`은 `instances`를 키로 하여 배열의 형태로 전달하고 있습니다. 또한 입력데이터는 batch dimension을 포함한 4차원 배열로 구성되어 있습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # Uncomment the following lines to see an example that cannot be processed by the endpoint\n",
    "\n",
    "# dummy_data = {\n",
    "#    'instances': np.random.rand(28, 28, 1).tolist()\n",
    "# }\n",
    "# print(predictor.predict(dummy_data))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "이제 실제 MNIST 테스트 데이터로 엔드포인트를 호출해 봅니다. 여기서는 MNIST 데이터를 다운로드하고 normalize하기 위해 `code.utils` 의 헬퍼함수를 사용하였습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from utils.mnist import mnist_to_numpy, normalize\n",
    "import random\n",
    "import matplotlib.pyplot as plt\n",
    "%matplotlib inline\n",
    "\n",
    "data_dir = '/tmp/data'\n",
    "X, _ = mnist_to_numpy(data_dir, train=False)\n",
    "\n",
    "# randomly sample 16 images to inspect\n",
    "mask = random.sample(range(X.shape[0]), 16)\n",
    "samples = X[mask]\n",
    "\n",
    "# plot the images \n",
    "fig, axs = plt.subplots(nrows=1, ncols=16, figsize=(16, 1))\n",
    "\n",
    "for i, splt in enumerate(axs):\n",
    "    splt.imshow(samples[i])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "모델이 nomalized 된 입력을 받게 되어있으므로 normalize 처리 후 엔디포인트를 호출합니다. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "samples = normalize(samples, axis=(1, 2))\n",
    "predictions = predictor.predict(\n",
    "    np.expand_dims(samples, 3) # add channel dim\n",
    ")['predictions'] \n",
    "\n",
    "# softmax to logit\n",
    "predictions = np.array(predictions, dtype=np.float32)\n",
    "predictions = np.argmax(predictions, axis=1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Predictions: \", predictions.tolist())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## (Optional) 새로운 환경에서 추론 endpoint 호출\n",
    "\n",
    "SageMaker는 배포된 endpoint를 호출하는 `ReatTimePredictor` 클래스를 제공합니다. 이는 별도의 새로운 환경에서 endpoint를 호출하는 방식을 예제로 보여줍니다.\n",
    "\n",
    "먼저 생성된 endpoint의 이름을 기억합니다. 여기서는 앞서 생성한 predictor객체로부터 가져오겠습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "my_endpoint = predictor.endpoint_name\n",
    "print(my_endpoint)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "`endpoint_name`을 이용하여 `ReatTimePredictor` 오브젝트를 생성합니다. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.predictor import  Predictor \n",
    "from sagemaker.serializers import JSONSerializer\n",
    "from sagemaker.deserializers import JSONDeserializer\n",
    "\n",
    "my_predictor = Predictor(endpoint_name=my_endpoint, \n",
    "                         sagemaker_session=sess, \n",
    "                         serializer=JSONSerializer(), \n",
    "                         deserializer=JSONDeserializer())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "이제 predict 함수를 이용하여 추론을 요청할 수 있습니다. 이전 코드에서 사용했던 dummy_inputs을 테스트용 데이터로 이용하겠습니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# dummy_inputs = {\n",
    "#     'instances': np.random.rand(4, 28, 28, 1).tolist()\n",
    "# }\n",
    "\n",
    "my_predictor.predict(dummy_inputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "추론 실행을 위해 boto3 SDK를 이용할 수도 있습니다. 아래 코드를 참조합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "sm_runtime = boto3.Session().client(service_name='sagemaker-runtime',region_name=sess.boto_region_name)\n",
    "\n",
    "response = sm_runtime.invoke_endpoint(EndpointName=my_endpoint, \n",
    "                                      ContentType='application/json', \n",
    "                                      Accept='application/json',\n",
    "                                      Body=json.dumps(dummy_inputs))\n",
    "\n",
    "json.loads(response.get('Body').read().decode())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 리소스 삭제\n",
    "\n",
    "endpoint의 사용이 끝나면 추가과금을 막기 위해 endpoint를 삭제합니다. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor.delete_endpoint()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "## (Optional) Hander debugging in local mode\n",
    "\n",
    "본 섹션에서는 Tensorflow Serving container에 hander 코드를 추가하고 테스트하는 방법을 살펴보겠습니다. \n",
    "\n",
    "먼저 로컬모드에서 handler를 테스트하기 위해 S3의 model.tar.gz파일을 로컬로 복사합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# !rm -Rf model_with_handler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p model_with_handler/model\n",
    "!mkdir -p model_with_handler/code\n",
    "!aws s3 cp {tf_mnist_model_data} model_with_handler/model.tar.gz\n",
    "!tar -zvxf model_with_handler/model.tar.gz -C model_with_handler/model\n",
    "!rm model_with_handler/model.tar.gz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### hanlder 코드 생성\n",
    "\n",
    "다음은 이 model.tar.gz에 request handler 코드를 추가하겠습니다. Tensorflow serving container v1.11부터 사용되고 있는 model.tar.gz 구조는 다음과 같습니다. (프레임워크의 종류와 버전별로 여러가지 방식이 가능합니다. Tensorflow구성과 관련된 내용은 https://github.com/aws/sagemaker-tensorflow-serving-container 를 참조합니다. 해당경로의 가이드를 통해 inference.py 파일의 생성방법 또한 참고할 수 있습니다.)\n",
    "\n",
    "```\n",
    "    model.tar.gz/\n",
    "    |- model\n",
    "    |   |--[model_version_number]\n",
    "    |       |--variables\n",
    "    |       |--saved_model.pb\n",
    "    |- code\n",
    "        |--inference.py\n",
    "        |--requirements.txt\n",
    "```\n",
    "\n",
    "다음 셀의 코드를 통해 커스텀 inference.py 파일을 생성합니다. Serving container에서 Request를 받아 디코딩하고 추론 Response에서 content 부분만 추출하는 간단한 handler입니다. (디버깅 확인 목적으로 위해 간단한 print문을 넣었습니다.)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile model_with_handler/code/inference.py\n",
    "\n",
    "import json\n",
    "\n",
    "def input_handler(data, context):\n",
    "    # decode request payload\n",
    "    payload = data.read().decode('utf-8')\n",
    "    print(\"========== my debugging message ===============\")\n",
    "    print(payload[:30])\n",
    "    return  payload\n",
    "    \n",
    "\n",
    "def output_handler(response, context):\n",
    "    # Serialize the prediction result\n",
    "    response_content_type = context.accept_header\n",
    "    prediction = response.content\n",
    "\n",
    "    return prediction, response_content_type\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "방금 생성한 handler를 포함하여 `model.tar.gz` 파일을 다시 생성합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%sh\n",
    "cd model_with_handler\n",
    "tar -czvf model.tar.gz model code"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 로컬모드 테스트\n",
    "\n",
    "로컬모드로 컨테이너를 실행하기에 앞서 기존 실행중인 로컬 컨테이너가 있다면 중지합니다. (기존 로컬모드로 8080 포트를 사용하는 Docker 컨테이너를 삭제합니다.)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.system(\"docker container ls | grep 8080 | awk '{print $1}' | xargs docker container rm -f\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "전 단계에서 생성한 model.tar.gz 로컬파일을 이용하여 로컬모드로 Tensorflow 추론 컨테이너를 실행합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "local_model = TensorFlowModel(role=role,\n",
    "                              model_data='file://model_with_handler/model.tar.gz',\n",
    "                              framework_version='2.3.1'\n",
    "                              )\n",
    "\n",
    "instance_type='local'\n",
    "\n",
    "local_predictor = local_model.deploy(initial_instance_count=1,\n",
    "                                     instance_type=instance_type,\n",
    "                                     )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "로컬 추론 컨테이너로 추론 요청을 보냅니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "res = local_predictor.predict(dummy_inputs)\n",
    "print(res)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "사용자 메시지가 출력되었습니다. 유사한 방식으로 다른 로그를 출력해 봅니다.\n",
    "\n",
    "### 로컬 리소스 삭제\n",
    "\n",
    "테스트가 완료되면 로컬의 컨테이너를 삭제합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.system(\"docker container ls | grep 8080 | awk '{print $1}' | xargs docker container rm -f\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "instance_type": "ml.t3.medium",
  "kernelspec": {
   "display_name": "conda_tensorflow2_p36",
   "language": "python",
   "name": "conda_tensorflow2_p36"
  },
  "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.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}