{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "057c8631",
   "metadata": {},
   "source": [
    "# Module 2. Deploy Multi-container Endpoint\n",
    "---\n",
    "\n",
    "## Overview\n",
    "\n",
    "SageMaker 멀티 컨테이너 엔드포인트를 사용하면 서로 다른 serving 스택(예: 모델 서버, 머신 러닝 프레임워크, 프레임워크 버전, 알고리즘 등)에 구축된 여러 추론 컨테이너를 하나의 엔드포인트에서 실행하고 독립적으로 각 추론 컨테이너를 호출할 수 있습니다. \n",
    "\n",
    "- 인스턴스의 전체 수용량을 포화시킬 정도의 트래픽이 없는 경우에 여러 모델(예: Object Detection, Named Entity Recognition)을 서빙\n",
    "- A/B 테스트와 같은 시나리오에서 서로 다른 프레임워크 버전(예: TensorFlow 1.x vs. TensorFlow 2.x)에서 실행되는 유사한 아키텍처의 비교"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2d86db0",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 1. Upload Model Artifacts\n",
    "---\n",
    "\n",
    "모델을 아카이빙하여 S3로 업로드합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cfcb20f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision\n",
    "import torchvision.models as models\n",
    "import sagemaker\n",
    "from sagemaker import get_execution_role\n",
    "from sagemaker.utils import name_from_base\n",
    "from sagemaker.pytorch import PyTorchModel\n",
    "import boto3\n",
    "import datetime\n",
    "import time\n",
    "from time import strftime,gmtime\n",
    "import json\n",
    "import os\n",
    "import io\n",
    "import torchvision.transforms as transforms\n",
    "from src.utils import print_outputs, upload_model_artifact_to_s3, NLPPredictor \n",
    "\n",
    "role = get_execution_role()\n",
    "boto_session = boto3.session.Session()\n",
    "sm_session = sagemaker.session.Session()\n",
    "sm_client = boto_session.client(\"sagemaker\")\n",
    "sm_runtime = boto_session.client(\"sagemaker-runtime\")\n",
    "region = boto_session.region_name\n",
    "bucket = sm_session.default_bucket()\n",
    "prefix = 'multi-container-nlp'\n",
    "\n",
    "print(f'region = {region}')\n",
    "print(f'role = {role}')\n",
    "print(f'bucket = {bucket}')\n",
    "print(f'prefix = {prefix}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca4c93a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "modelA_variant = 'modelA'\n",
    "modelA_path = 'model-nsmc'\n",
    "modelA_s3_uri = upload_model_artifact_to_s3(modelA_variant, modelA_path, bucket, prefix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "461fea3a",
   "metadata": {},
   "outputs": [],
   "source": [
    "modelB_variant = 'modelB'\n",
    "modelB_path = 'model-korsts'\n",
    "modelB_s3_uri = upload_model_artifact_to_s3(modelB_variant, modelB_path, bucket, prefix)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c26f015b",
   "metadata": {},
   "outputs": [],
   "source": [
    "modelC_variant = 'modelC'\n",
    "modelC_path = 'model-kobart'\n",
    "modelC_s3_uri = upload_model_artifact_to_s3(modelC_variant, modelC_path, bucket, prefix)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0abac512",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## 2. Create Multi-container endpoint\n",
    "---\n",
    "\n",
    "SageMaker 멀티 컨테이너 엔드포인트를 사용하면 여러 컨테이너들을 동일한 엔드포인트에 배포할 수 있으며, 각 컨테이너에 개별적으로 액세스하여 비용을 최적화할 수 있습니다. 다중 컨테이너 엔드포인트 설정은 기존 엔드포인트와 유사하지만, SageMaker 모델 생성 시 여러 컨테이너들을 명시해 줘야 합니다.\n",
    "\n",
    "- 배포에 필요한 각 컨테이너에 대한 추론 컨테이너 정의 생성\n",
    "- `create_model` API를 사용하여 SageMaker 모델 생성; PrimaryContainer 대신 Containers 매개변수를 사용하고 Containers 매개변수에 두 개 이상의 컨테이너를 포함합니다. (최대 15개까지 지원)\n",
    "- `create_endpoint_config` API를 사용하여 SageMaker 엔드포인트 설정 생성\n",
    "- `create_endpoint` API를 사용하여 SageMaker 엔드포인트 생성; 생성 시 반드시 Direct 모드를 사용해야 합니다.\n",
    "\n",
    "자세한 내용은 아래 링크를 참조해 주세요.\n",
    "- Host Multiple Models with Multi-Model Endpoints: https://docs.aws.amazon.com/sagemaker/latest/dg/multi-model-endpoints.html"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6b7de177",
   "metadata": {},
   "source": [
    "### Create Inference containter definition for Model A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b3ae708",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.image_uris import retrieve\n",
    "\n",
    "deploy_instance_type = 'ml.c5.xlarge'\n",
    "pt_ecr_image_uriA = retrieve(\n",
    "    framework='pytorch',\n",
    "    region=region,\n",
    "    version='1.7.1',\n",
    "    py_version='py3',\n",
    "    instance_type = deploy_instance_type,\n",
    "    accelerator_type=None,\n",
    "    image_scope='inference'\n",
    ")\n",
    "\n",
    "pt_containerA = {\n",
    "    \"ContainerHostname\": \"pytorch-kornlp-nsmc\",\n",
    "    \"Image\": pt_ecr_image_uriA,\n",
    "    \"ModelDataUrl\": modelA_s3_uri,\n",
    "    \"Environment\": {\n",
    "        \"SAGEMAKER_PROGRAM\": \"inference_nsmc.py\",\n",
    "        \"SAGEMAKER_SUBMIT_DIRECTORY\": modelA_s3_uri,\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3446e7f8",
   "metadata": {},
   "source": [
    "### Create Inference containter definition for Model B"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1aa16444",
   "metadata": {},
   "outputs": [],
   "source": [
    "pt_ecr_image_uriB = retrieve(\n",
    "    framework='pytorch',\n",
    "    region=region,\n",
    "    version='1.8.1',\n",
    "    py_version='py3',\n",
    "    instance_type = deploy_instance_type,\n",
    "    accelerator_type=None,\n",
    "    image_scope='inference'\n",
    ")\n",
    "\n",
    "pt_containerB = {\n",
    "    \"ContainerHostname\": \"pytorch-kornlp-korsts\",\n",
    "    \"Image\": pt_ecr_image_uriB,\n",
    "    \"ModelDataUrl\": modelB_s3_uri,\n",
    "    \"Environment\": {\n",
    "        \"SAGEMAKER_PROGRAM\": \"inference_korsts.py\",\n",
    "        \"SAGEMAKER_SUBMIT_DIRECTORY\": modelB_s3_uri,\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d9a9483",
   "metadata": {},
   "source": [
    "### Create Inference containter definition for Model C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f4485b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "pt_ecr_image_uriC = retrieve(\n",
    "    framework='pytorch',\n",
    "    region=region,\n",
    "    version='1.8.1',\n",
    "    py_version='py3',\n",
    "    instance_type = deploy_instance_type,\n",
    "    accelerator_type=None,\n",
    "    image_scope='inference'\n",
    ")\n",
    "\n",
    "pt_containerC = {\n",
    "    \"ContainerHostname\": \"pytorch-kornlp-kobart\",\n",
    "    \"Image\": pt_ecr_image_uriC,\n",
    "    \"ModelDataUrl\": modelC_s3_uri,\n",
    "    \"Environment\": {\n",
    "        \"SAGEMAKER_PROGRAM\": \"inference_kobart.py\",\n",
    "        \"SAGEMAKER_SUBMIT_DIRECTORY\": modelC_s3_uri,\n",
    "    },\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a125c4cf",
   "metadata": {},
   "source": [
    "### Create a SageMaker Model\n",
    "\n",
    "`create_model` API를 호출하여 위 코드 셀에서 생성한 다중 컨테이너의 정의를 포함하는 모델을 생성합니다. 기존 엔드포인트와의 차이점은 `Containers` 매개 변수를 설정하고 `InferenceExecutionConfig` 매개변수의 Mode를 `Direct`나 `Serial`로 설정하는 것입니다. 기본 모드는 `Serial`이지만, 각 컨테이너를 직접 호출하려면 `Direct`로 설정해야 합니다. 자세한 내용은 [멀티 컨테이너 엔드포인트 배포](https://docs.aws.amazon.com/sagemaker/latest/dg/multi-container-endpoints.html)를 확인하세요."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee2c950e",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_name = f\"KorNLPMultiContainer-{strftime('%Y-%m-%d-%H-%M-%S', gmtime())}\"\n",
    "\n",
    "create_model_response = sm_client.create_model(\n",
    "    ModelName=model_name,\n",
    "    Containers=[pt_containerA, pt_containerB, pt_containerC],\n",
    "    InferenceExecutionConfig={\"Mode\": \"Direct\"},\n",
    "    ExecutionRoleArn=role,\n",
    ")\n",
    "print(f\"Created Model: {create_model_response['ModelArn']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c7c558a",
   "metadata": {},
   "source": [
    "### Create Endpoint Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37db052e",
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint_config_name = f\"KorNLPMultiContainerEndpointConfig-{strftime('%Y-%m-%d-%H-%M-%S', gmtime())}\"\n",
    "endpoint_config_response = sm_client.create_endpoint_config(\n",
    "    EndpointConfigName=endpoint_config_name,\n",
    "    ProductionVariants=[\n",
    "        {\n",
    "            \"VariantName\": \"prod\",\n",
    "            \"ModelName\": model_name,\n",
    "            \"InitialInstanceCount\": 1,\n",
    "            \"InstanceType\": deploy_instance_type,\n",
    "        },\n",
    "    ],\n",
    ")\n",
    "print(f\"Created EndpointConfig: {endpoint_config_response['EndpointConfigArn']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e82aae26",
   "metadata": {},
   "source": [
    "### Create a SageMaker Multi-container endpoint\n",
    "\n",
    "create_endpoint API로 멀티 컨테이너 엔드포인트를 생성합니다. 기존의 엔드포인트 생성 방법과 동일합니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40a8246a",
   "metadata": {},
   "outputs": [],
   "source": [
    "endpoint_name = f\"KorNLPMultiContainerEndpoint-{strftime('%Y-%m-%d-%H-%M-%S', gmtime())}\"\n",
    "endpoint_response = sm_client.create_endpoint(\n",
    "    EndpointName=endpoint_name, EndpointConfigName=endpoint_config_name\n",
    ")\n",
    "print(f\"Creating Endpoint: {endpoint_response['EndpointArn']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d687e54e",
   "metadata": {},
   "source": [
    "`describe_endpoint` API를 사용하여 엔드포인트 생성 상태를 확인할 수 있습니다. 엔드포인트 생성은 약 5분에서 10분이 소요됩니다."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eabb724b",
   "metadata": {},
   "outputs": [],
   "source": [
    "waiter = boto3.client('sagemaker').get_waiter('endpoint_in_service')\n",
    "print(\"Waiting for endpoint to create...\")\n",
    "waiter.wait(EndpointName=endpoint_name)\n",
    "resp = sm_client.describe_endpoint(EndpointName=endpoint_name)\n",
    "print(f\"Endpoint Status: {resp['EndpointStatus']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "30e5e5b7",
   "metadata": {},
   "source": [
    "### Direct Invocation for Model A\n",
    "\n",
    "두 문장간의 유사도를 정량화하는 예시입니다.\n",
    "- KorNLI and KorSTS: https://github.com/kakaobrain/KorNLUDatasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cafb3e21",
   "metadata": {},
   "outputs": [],
   "source": [
    "modelA_sample_path = 'samples/nsmc.txt'\n",
    "!cat $modelA_sample_path\n",
    "with open(modelA_sample_path, mode='rb') as file:\n",
    "    modelA_input_data = file.read()  \n",
    "\n",
    "modelA_response = sm_runtime.invoke_endpoint(\n",
    "    EndpointName=endpoint_name,\n",
    "    ContentType=\"application/jsonlines\",\n",
    "    Accept=\"application/jsonlines\",\n",
    "    TargetContainerHostname=\"pytorch-kornlp-nsmc\",\n",
    "    Body=modelA_input_data\n",
    ")\n",
    "\n",
    "modelA_outputs = modelA_response['Body'].read().decode()\n",
    "print()\n",
    "print_outputs(modelA_outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "024c97ca",
   "metadata": {},
   "source": [
    "### Direct Invocation for Model B\n",
    "\n",
    "두 문장간의 유사도를 정량화하는 예시입니다.\n",
    "- KorNLI and KorSTS: https://github.com/kakaobrain/KorNLUDatasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e0afbce",
   "metadata": {},
   "outputs": [],
   "source": [
    "modelB_sample_path = 'samples/korsts.txt'\n",
    "!cat $modelB_sample_path\n",
    "with open(modelB_sample_path, mode='rb') as file:\n",
    "    modelB_input_data = file.read()    \n",
    "    \n",
    "modelB_response = sm_runtime.invoke_endpoint(\n",
    "    EndpointName=endpoint_name,\n",
    "    ContentType=\"application/jsonlines\",\n",
    "    Accept=\"application/jsonlines\",\n",
    "    TargetContainerHostname=\"pytorch-kornlp-korsts\",\n",
    "    Body=modelB_input_data\n",
    ")\n",
    "\n",
    "modelB_outputs = modelB_response['Body'].read().decode()\n",
    "print()\n",
    "print_outputs(modelB_outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "972c22c3",
   "metadata": {},
   "source": [
    "### Direct Invocation for Model C\n",
    "\n",
    "문서 내용(예: 뉴스 기사)을 요약하는 예시입니다.\n",
    "\n",
    "- KoBART: https://github.com/SKT-AI/KoBART\n",
    "- KoBART Summarization: https://github.com/seujung/KoBART-summarization"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b8d6f44",
   "metadata": {},
   "outputs": [],
   "source": [
    "modelC_sample_path = 'samples/kobart.txt'\n",
    "!cat $modelC_sample_path\n",
    "with open(modelC_sample_path, mode='rb') as file:\n",
    "    modelC_input_data = file.read()    \n",
    "    \n",
    "modelC_response = sm_runtime.invoke_endpoint(\n",
    "    EndpointName=endpoint_name,\n",
    "    ContentType=\"application/jsonlines\",\n",
    "    Accept=\"application/jsonlines\",\n",
    "    TargetContainerHostname=\"pytorch-kornlp-kobart\",\n",
    "    Body=modelC_input_data\n",
    ")\n",
    "\n",
    "modelC_outputs = modelC_response['Body'].read().decode()\n",
    "print()\n",
    "print_outputs(modelC_outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8b5fef3e",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "## Clean Up\n",
    "---"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a752ae5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "sm_client.delete_endpoint(EndpointName=endpoint_name)\n",
    "sm_client.delete_endpoint_config(EndpointConfigName=endpoint_config_name)\n",
    "sm_client.delete_model(ModelName=model_name)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_latest_p37",
   "language": "python",
   "name": "conda_pytorch_latest_p37"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}