{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ee0bbaae",
   "metadata": {},
   "outputs": [],
   "source": [
    "%load_ext nb_black\n",
    "\n",
    "from problems.tsp.problem_tsp import TSP\n",
    "from utils import load_model, move_to\n",
    "from torch.utils.data import DataLoader\n",
    "import boto3\n",
    "import sagemaker"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc98bfee",
   "metadata": {},
   "outputs": [],
   "source": [
    "from inference import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f8471da",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d8a96f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.serializers import JSONLinesSerializer\n",
    "from sagemaker.deserializers import JSONLinesDeserializer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42192ca1",
   "metadata": {},
   "outputs": [],
   "source": [
    "session = sagemaker.Session()\n",
    "BUCKET = session.default_bucket()  # Set a default S3 bucket"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0df89443",
   "metadata": {},
   "outputs": [],
   "source": [
    "# set USE_PRETRAINED_MODEL to False if you have trained a model using pytorch_training.ipynb\n",
    "USE_PRETRAINED_MODEL = True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99ad81fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "PRETRAINED_MODEL_PATH = \"../learning-tsp/pretrained/tsp_20-50/rl-ar-var-20pnn-gnn-max_20200313T002243/model.tar.gz\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "102dd028",
   "metadata": {},
   "source": [
    "# 1. Test inference code locally"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f746e077",
   "metadata": {},
   "source": [
    "## Prepare data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca2b081e",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_path = None\n",
    "batch_size = 1\n",
    "accumulation_steps = 80\n",
    "num_samples = 2  # 1280 samples per TSP size\n",
    "\n",
    "neighbors = 0.20\n",
    "knn_strat = \"percentage\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bf153a2",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = TSP.make_dataset(\n",
    "    filename=dataset_path,\n",
    "    batch_size=batch_size,\n",
    "    num_samples=num_samples,\n",
    "    min_size=10,\n",
    "    max_size=10,\n",
    "    neighbors=neighbors,\n",
    "    knn_strat=knn_strat,\n",
    "    supervised=False,\n",
    ")\n",
    "dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=False, num_workers=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66982dfa",
   "metadata": {},
   "outputs": [],
   "source": [
    "# transform data\n",
    "data = []\n",
    "for bat_idx, bat in enumerate(dataloader):\n",
    "    input = {}\n",
    "    input[\"nodes\"] = bat[\"nodes\"].tolist()\n",
    "    data.append(input)\n",
    "for record in data:\n",
    "    record[\"neighbors\"] = neighbors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7e7fc961",
   "metadata": {},
   "outputs": [],
   "source": [
    "data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27cfc2ab",
   "metadata": {},
   "source": [
    "## Prepare the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a19edbb0",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Getting the latest model data from the training jobs\n",
    "def get_latest_model():\n",
    "    client = boto3.client(\"sagemaker\")\n",
    "\n",
    "    # Get the trained sklearn model\n",
    "    response = client.list_training_jobs(\n",
    "        NameContains=\"pytorch-smdataparallel-tsp\",\n",
    "        StatusEquals=\"Completed\",\n",
    "        SortBy=\"CreationTime\",\n",
    "        SortOrder=\"Descending\",\n",
    "    )\n",
    "    training_job_name = response[\"TrainingJobSummaries\"][0][\"TrainingJobName\"]\n",
    "    model_s3 = client.describe_training_job(TrainingJobName=training_job_name)[\n",
    "        \"ModelArtifacts\"\n",
    "    ][\"S3ModelArtifacts\"]\n",
    "    return model_s3\n",
    "\n",
    "\n",
    "# Upload a pretrained model to s3\n",
    "def upload_pretrained_model():\n",
    "    s3 = boto3.resource(\"s3\")\n",
    "    S3_PATH = PRETRAINED_MODEL_PATH.lstrip(\"../\")\n",
    "    s3.meta.client.upload_file(PRETRAINED_MODEL_PATH, BUCKET, S3_PATH)\n",
    "    return f\"s3://{BUCKET}/{S3_PATH}\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d88c8bf",
   "metadata": {},
   "outputs": [],
   "source": [
    "if USE_PRETRAINED_MODEL == True:\n",
    "    model_data = upload_pretrained_model()\n",
    "else:\n",
    "    model_data = get_latest_model()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9c161a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6dbb706",
   "metadata": {},
   "source": [
    "## Download the model locally for testing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd663c92",
   "metadata": {},
   "outputs": [],
   "source": [
    "!aws s3 cp $model_data ./"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "88f6b417",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfa50bdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "!tar -xvzf ./model.tar.gz -C ./model/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "873306c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "562477ed",
   "metadata": {},
   "source": [
    "## Load model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74709ee6",
   "metadata": {},
   "outputs": [],
   "source": [
    "model_dir = \"./model\"\n",
    "model = model_fn(model_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08a3748f",
   "metadata": {},
   "outputs": [],
   "source": [
    "model"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f19e66f8",
   "metadata": {},
   "source": [
    "## Define input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7c8e1cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "serializer = JSONLinesSerializer()\n",
    "\n",
    "data_jsonlines = serializer.serialize(data)\n",
    "\n",
    "request_body = data_jsonlines.encode(\"utf-8\")\n",
    "\n",
    "input_data = input_fn(request_body)\n",
    "\n",
    "with open(\"inference_input\", \"w\") as file:\n",
    "    file.write(data_jsonlines)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "667dfd95",
   "metadata": {},
   "outputs": [],
   "source": [
    "# upload to S3 for batch transform\n",
    "!aws s3 cp inference_input s3://$BUCKET/data/inference/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebcd49b1",
   "metadata": {},
   "source": [
    "## Prediction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea804a50",
   "metadata": {},
   "outputs": [],
   "source": [
    "input_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c013a50c",
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction = predict_fn(input_data, model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b3c9ee3",
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68fe93f4",
   "metadata": {},
   "source": [
    "## Define output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57c70d8a",
   "metadata": {},
   "outputs": [],
   "source": [
    "output = output_fn(prediction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "40b58920",
   "metadata": {},
   "outputs": [],
   "source": [
    "output"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec96bbb4",
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(\"prediction\", \"w\") as file:\n",
    "    file.write(output[0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0af288e",
   "metadata": {},
   "source": [
    "# 2. Test inference code via endpoints"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d7d2f53",
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "\n",
    "role = sagemaker.get_execution_role()\n",
    "\n",
    "from sagemaker.pytorch import PyTorchModel\n",
    "\n",
    "model_sm = PyTorchModel(\n",
    "    model_data=model_data,\n",
    "    source_dir=\"../src\",\n",
    "    entry_point=\"inference.py\",\n",
    "    role=role,\n",
    "    framework_version=\"1.8.1\",\n",
    "    py_version=\"py36\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2b892e7a",
   "metadata": {},
   "outputs": [],
   "source": [
    "predictor = model_sm.deploy(\n",
    "    initial_instance_count=1,\n",
    "    instance_type=\"ml.m4.xlarge\",\n",
    "    serializer=JSONLinesSerializer(),\n",
    "    deserializer=JSONLinesDeserializer(),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9377b689",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Send the sampled images to endpoint for inference\n",
    "prediction = predictor.predict(data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "894d55d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a54ac57c",
   "metadata": {},
   "source": [
    "# 3. Batch Transform"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d0fffba",
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer = model_sm.transformer(\n",
    "    instance_count=1,\n",
    "    instance_type=\"ml.m5.2xlarge\",\n",
    "    strategy=\"MultiRecord\",\n",
    "    assemble_with=\"Line\",\n",
    "    accept=\"application/jsonlines\",\n",
    "    max_concurrent_transforms=4,\n",
    "    env={\"SAGEMAKER_MODEL_SERVER_TIMEOUT\": \"3600\"},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a8d4830",
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer.transform(\n",
    "    f\"s3://{BUCKET}/data/inference/inference_input\",\n",
    "    content_type=\"application/jsonlines\",\n",
    "    split_type=\"Line\",\n",
    "    wait=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f92bb4a0",
   "metadata": {},
   "source": [
    "# 4. Clean up (Optional)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ba0c8e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Delete the SageMaker endpoint\n",
    "predictor.delete_endpoint()\n",
    "\n",
    "# Delete the SageMaker model\n",
    "model_sm.delete_model()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_tsp_deep_rl",
   "language": "python",
   "name": "conda_tsp_deep_rl"
  },
  "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": 5
}