advanced_functionality/pytorch_extend_container_train_deploy_bertopic/BERTtopic_extending_container.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Fine-tuning and deploying a BERTopic model on SageMaker with your own scripts and dataset, by extending existing PyTorch containers\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "---\n",
    "\n",
    "This notebook's CI test result for us-west-2 is as follows. CI test results in other regions can be found at the end of the notebook. \n",
    "\n",
    "![This us-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-2/advanced_functionality|pytorch_extend_container_train_deploy_bertopic|BERTtopic_extending_container.ipynb)\n",
    "\n",
    "---"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "### INFO: this notebook follows and extends the structure of the [Extending our PyTorch containers example](https://github.com/aws/amazon-sagemaker-examples/blob/main/advanced_functionality/pytorch_extending_our_containers/pytorch_extending_our_containers.ipynb)\n",
    "\n",
    "With Amazon SageMaker, you can package your own algorithms that can then be trained and deployed in the SageMaker environment. This notebook guides you through an example on how to extend one of our existing and predefined SageMaker deep learning framework containers.\n",
    "\n",
    "By packaging an algorithm in a container, you can bring almost any code to the Amazon SageMaker environment, regardless of programming language, environment, framework, or dependencies. \n",
    "\n",
    "1. [Extending our PyTorch containers](#Extending-our-PyTorch-containers)\n",
    "    1. [When should I extend a SageMaker container?](#When-should-I-extend-a-SageMaker-container?)\n",
    "    1. [Permissions](#Permissions)\n",
    "    1. [This example](#This-example)\n",
    "    1. [This notebook's structure](#This-notebook's-structure)\n",
    "1. [Part 1: Docker containers and their use in Amazon SageMaker](#Part-1:-Docker-containers-and-their-use-in-Amazon-SageMaker)\n",
    "    1. [An overview of Docker](#An-overview-of-Docker)\n",
    "    1. [How Amazon SageMaker runs your Docker container](#How-Amazon-SageMaker-runs-your-Docker-container)\n",
    "      1. [Running your container during training](#Running-your-container-during-training)\n",
    "        1. [The input](#The-input)\n",
    "        1. [The output](#The-output)\n",
    "      1. [Running your container during hosting](#Running-your-container-during-hosting)\n",
    "    1. [Custom files available to build the container used in this example](#Custom-files-available-to-build-the-container-used-in-this-example)\n",
    "1. [Part 2: Packaging and Uploading your Algorithm for training](#Part-2:-Packaging-and-Uploading-your-Algorithm-for-training)\n",
    "    1. [The training Dockerfile](#The-training-Dockerfile)\n",
    "    1. [Writing the training script (bert-topic.py)](#Writing-the-training-script-(bert-topic.py))\n",
    "    1. [Building and registering the training container](#Building-and-registering-the-training-container)\n",
    "    1. [Testing your algorithm](#Testing-your-algorithm)\n",
    "        1. [Download the 20 Newsgroups training dataset](#Download-the-20-Newsgroups-training-dataset)\n",
    "        1. [SageMaker Python SDK Local Training](#SageMaker-Python-SDK-Local-Training)\n",
    "        1. [Check the local SageMaker setup](#Check-the-local-SageMaker-setup)\n",
    "    1. [Fit (+ Deploy, Predict)](#Fit-(+-Deploy,-Predict))\n",
    "        1. [Local training: Fit](#Local-training:-Fit)\n",
    "        1. [A sleek alternative to \"Bring your own container\": bring your own requirements.txt](#A-sleek-alternative-to-\"Bring-your-own-container\":-bring-your-own-requirements.txt)\n",
    "    1.[Training your Algorithm in Amazon SageMaker in batch mode](#Training-your-Algorithm-in-Amazon-SageMaker-in-batch-mode)\n",
    "        1. [Set up the environment](#Set-up-the-environment)\n",
    "        1. [Create the session](#Create-the-session)\n",
    "        1. [Upload the data for training](#Upload-the-data-for-training)\n",
    "        1. [Training on SageMaker in batch mode](#Training-on-SageMaker-in-batch-mode)\n",
    "1. [Part 3: Packaging and Uploading your Algorithm for inference](#Part-3:-Packaging-and-Uploading-your-Algorithm-for-inference)\n",
    "    1. [The inference Dockerfile](#The-inference-Dockerfile)\n",
    "    1. [Writing your own inference script (bert-topic-inference.py)](#Writing-your-own-inference-script-(bert-topic-inference.py))\n",
    "    1. [Inference from Containerized SageMaker Model)](#Inference-from-Containerized-SageMaker-Model)\n",
    "        1. [Prepare session to run container locally](#Prepare-session-to-run-container-locally)\n",
    "        1. [Check for locally running Docker containers (and stop them if needed)](#Check-for-locally-running-Docker-containers-(and-stop-them-if-needed))\n",
    "        1. [Deploy container locally to create a local endpoint](#Deploy-container-locally-to-create-a-local-endpoint)\n",
    "        1. [Invoke locally deployed endpoint](#Invoke-locally-deployed-endpoint)\n",
    "        1. [A sleek alternative to \"Bring your own container\": bring your own requirements.txt](#A-sleek-alternative-to-\"Bring-your-own-container\":-bring-your-own-requirements.txt)\n",
    "        1. [Invoke locally deployed endpoint](#A-sleek-alternative-to-\"Bring-your-own-container\":-bring-your-own-requirements.txt)\n",
    "    1. [Deploy container remotely to create a managed Amazon SageMaker endpoint for real-time inference](#Deploy-container-remotely-to-create-a-managed-Amazon-SageMaker-endpoint-for-real-time-inference)\n",
    "        1. [Invoke remotely deployed endpoint](#Invoke-remotely-deployed-endpoint)\n",
    "    1. [Optional cleanup](#Optional-cleanup)  \n",
    "1. [Reference](#Reference)\n",
    "\n",
    "\n",
    "\n",
    "## Extending our PyTorch containers\n",
    "\n",
    "### When should I extend a SageMaker container?\n",
    "\n",
    "You may not need to create a container to bring your own code to Amazon SageMaker. When you are using a framework such as [TensorFlow](https://github.com/aws/sagemaker-tensorflow-container), [MXNet](https://github.com/aws/sagemaker-mxnet-container), [PyTorch](https://github.com/aws/sagemaker-pytorch-container) or [Chainer](https://github.com/aws/sagemaker-chainer-container) that has direct support in SageMaker, you can simply supply the Python code that implements your algorithm using the SDK entry points for that framework.\n",
    "\n",
    "Even if there is direct SDK support for your environment or framework, you may want to add additional functionality or configure your container environment differently while utilizing our container to use on SageMaker.\n",
    "\n",
    "**Some of the reasons to extend a SageMaker deep learning framework container are:**\n",
    "1. Install additional dependencies. (E.g. I want to install a specific Python library, that the current SageMaker containers don't install.)\n",
    "2. Configure your environment. (E.g. I want to add an environment variable to my container.)\n",
    "\n",
    "**Although it is possible to extend any of our framework containers as a parent image, the example this notebook covers is currently only intended to work with our PyTorch container.**\n",
    "\n",
    "This walkthrough shows that it is quite straightforward to extend one of our containers to build your own custom container for PyTorch.\n",
    "\n",
    "### Permissions\n",
    "\n",
    "Running this notebook requires permissions in addition to the normal `SageMakerFullAccess` permissions. This is because it creates new repositories in Amazon ECR. The easiest way to add these permissions is simply to add the managed policy `AmazonEC2ContainerRegistryFullAccess` to the role that you used to start your notebook instance. There's no need to restart your notebook instance when you do this, the new permissions will be available immediately.\n",
    "\n",
    "### This example\n",
    "\n",
    "In this example we show how to package a PyTorch container, extending the SageMaker PyTorch container, with a Python example which works with the BERTopic model and the 20 Newsgroups dataset. BERTopic comes now as standalone library and requires GPU support, so the most effective way to use it is to extend the SageMaker PyTorch container, to utilize the existing training and hosting solution made to work on SageMaker. By comparison, if one were to build their own custom framework container from scratch, they would need to implement a training and hosting solution in order to use SageMaker. Here is an example showing [how to create a SageMaker TensorFlow container from scratch](https://github.com/awslabs/amazon-sagemaker-examples/blob/master/advanced_functionality/tensorflow_bring_your_own/tensorflow_bring_your_own.ipynb).\n",
    "\n",
    "In this example, we need to use separate base images to support training and hosting, as they are provided separately for [PyTorch](https://github.com/aws/deep-learning-containers/blob/master/available_images.md#sagemaker-framework-containers-sm-support-only). Sometimes you may want to use the same image for training and hosting because they have the same requirements. This simplifies the procedure because we only need to manage one image for both tasks. In that case, merge the parts discussed below into the same Dockerfiles and build one image. Choosing whether to use a single image or two images is a matter of what is most convenient for you to develop and manage.\n",
    "\n",
    "If you're only using Amazon SageMaker for training or hosting, but not both, only the functionality used needs to be built into your container.\n",
    "Finally, when using different containers, make sure that the libraries and their versions match, otherwise this could create problems. In fact, some libraries versions are different in the PyTorch training and inference containers.\n",
    "\n",
    "### This notebook's structure\n",
    "\n",
    "This notebook is divided into three parts: \n",
    "1. Docker containers and their use in Amazon SageMaker\n",
    "2. _building_ and _using_ the training container \n",
    "3. _building_ and _using_ the inference container "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 1: Docker containers and their use in Amazon SageMaker\n",
    "\n",
    "### An overview of Docker\n",
    "\n",
    "If you're familiar with Docker already, you can skip ahead to the next section.\n",
    "\n",
    "For many data scientists, Docker containers are a new technology. But they are not difficult and can significantly simplify the deployment of your software packages. \n",
    "\n",
    "Docker provides a simple way to package arbitrary code into an _image_ that is totally self-contained. Once you have an image, you can use Docker to run a _container_ based on that image. Running a container is just like running a program on the machine except that the container creates a fully self-contained environment for the program to run. Containers are isolated from each other and from the host environment, so the way your program is set up is the way it runs, no matter where you run it.\n",
    "\n",
    "Docker is more powerful than environment managers like Conda or virtualenv because (a) it is completely language independent and (b) it comprises your whole operating environment, including startup commands, and environment variable.\n",
    "\n",
    "A Docker container is like a virtual machine, but it is much lighter weight. For example, a program running in a container can start in less than a second and many containers can run simultaneously on the same physical or virtual machine instance.\n",
    "\n",
    "Docker uses a simple file called a `Dockerfile` to specify how the image is assembled. An example is provided below. You can build your Docker images based on Docker images built by yourself or by others, which can simplify things quite a bit.\n",
    "\n",
    "Docker has become very popular in programming and DevOps communities due to its flexibility and its well-defined specification of how code can be run in its containers. It is the underpinning of many services built in the past few years, such as [Amazon ECS].\n",
    "\n",
    "Amazon SageMaker uses Docker to allow users to train and deploy arbitrary algorithms.\n",
    "\n",
    "In Amazon SageMaker, Docker containers are invoked in a one way for training and another, slightly different, way for hosting. The following sections outline how to build containers for the SageMaker environment.\n",
    "\n",
    "Some helpful links:\n",
    "\n",
    "* [Docker home page](http://www.docker.com)\n",
    "* [Getting started with Docker](https://docs.docker.com/get-started/)\n",
    "* [Dockerfile reference](https://docs.docker.com/engine/reference/builder/)\n",
    "* [`docker run` reference](https://docs.docker.com/engine/reference/run/)\n",
    "\n",
    "[Amazon ECS]: https://aws.amazon.com/ecs/\n",
    "\n",
    "### How Amazon SageMaker runs your Docker container\n",
    "\n",
    "* Typically you specify a program (e.g. script) as an `ENTRYPOINT` in the Dockerfile, that program will be run at startup and decide what to do. The original `ENTRYPOINT` specified within the SageMaker PyTorch is [here](https://github.com/aws/deep-learning-containers/blob/master/pytorch/training/docker/1.5.1/py3/Dockerfile.cpu#L142).\n",
    "\n",
    "#### Running your container during training\n",
    "\n",
    "Currently, our SageMaker PyTorch container utilizes [console_scripts](http://python-packaging.readthedocs.io/en/latest/command-line-scripts.html#the-console-scripts-entry-point) to make use of the `train` command issued at training time. The line that gets invoked during `train` is defined within the setup.py file inside [SageMaker Containers](https://github.com/aws/sagemaker-containers/blob/master/setup.py#L48), our common SageMaker deep learning container framework. When this command is run, it will invoke the [trainer class](https://github.com/aws/sagemaker-containers/blob/master/src/sagemaker_containers/cli/train.py) to run, which will finally invoke our [PyTorch container code](https://github.com/aws/sagemaker-pytorch-container/blob/master/src/sagemaker_pytorch_container/training.py) to run your Python file.\n",
    "\n",
    "A number of files are laid out for your use, under the `/opt/ml` directory:\n",
    "\n",
    "    /opt/ml\n",
    "    |-- input\n",
    "    |   |-- config\n",
    "    |   |   |-- hyperparameters.json\n",
    "    |   |   `-- resourceConfig.json\n",
    "    |   `-- data\n",
    "    |       `-- <channel_name>\n",
    "    |           `-- <input data>\n",
    "    |-- model\n",
    "    |   `-- <model files>\n",
    "    `-- output\n",
    "        `-- failure\n",
    "\n",
    "##### The input\n",
    "\n",
    "* `/opt/ml/input/config` contains information to control how your program runs. `hyperparameters.json` is a JSON-formatted dictionary of hyperparameter names to values. These values are always strings, so you may need to convert them. `resourceConfig.json` is a JSON-formatted file that describes the network layout used for distributed training.\n",
    "* `/opt/ml/input/data/<channel_name>/` (for File mode) contains the input data for that channel. The channels are created based on the call to CreateTrainingJob, but it's generally important that channels match algorithm expectations. The files for each channel are copied from S3 to this directory, preserving the tree structure indicated by the S3 key structure. \n",
    "* `/opt/ml/input/data/<channel_name>_<epoch_number>` (for Pipe mode) is the pipe for a given epoch. Epochs start at zero and go up by one each time you read them. There is no limit to the number of epochs that you can run, but you must close each pipe before reading the next epoch.\n",
    "\n",
    "##### The output\n",
    "\n",
    "* `/opt/ml/model/` is the directory where you write the model that your algorithm generates. Your model can be in any format that you want. It can be a single file or a whole directory tree. SageMaker packages any files in this directory into a compressed tar archive file. This file is made available at the S3 location returned in the `DescribeTrainingJob` result.\n",
    "* `/opt/ml/output` is a directory where the algorithm can write a file `failure` that describes why the job failed. The contents of this file are returned in the `FailureReason` field of the `DescribeTrainingJob` result. For jobs that succeed, there is no reason to write this file as it is ignored.\n",
    "\n",
    "#### Running your container during hosting\n",
    "\n",
    "Hosting has a very different model than training because hosting is responding to inference requests that come in via HTTP. Currently, the SageMaker PyTorch containers [uses](https://github.com/aws/sagemaker-pytorch-container/blob/master/src/sagemaker_pytorch_container/serving.py#L103) our [recommended Python serving stack](https://github.com/aws/sagemaker-containers/blob/master/src/sagemaker_containers/_server.py#L44) to provide robust and scalable serving of inference requests:\n",
    "\n",
    "![Request serving stack](stack.png)\n",
    "\n",
    "Amazon SageMaker uses two URLs in the container:\n",
    "\n",
    "* `/ping` receives `GET` requests from the infrastructure. Your program returns 200 if the container is up and accepting requests.\n",
    "* `/invocations` is the endpoint that receives client inference `POST` requests. The format of the request and the response is up to the algorithm. If the client supplied `ContentType` and `Accept` headers, these are passed in as well. \n",
    "\n",
    "The container has the model files in the same place that they were written to during training:\n",
    "\n",
    "    /opt/ml\n",
    "    `-- model\n",
    "        `-- <model files>\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Custom files available to build the container used in this example\n",
    "\n",
    "The `container` directory has all the components you need to extend the SageMaker PyTorch container to use as a sample algorithm:\n",
    "\n",
    "    .\n",
    "    |-- Dockerfile\n",
    "    |-- Dockerfile-inference\n",
    "    `-- bert-topic\n",
    "        `-- bert-topic.py\n",
    "        `-- bert-topic-inference.py\n",
    "\n",
    "Let's discuss each of these in turn:\n",
    "\n",
    "* __`Dockerfile`__ describes how to build your Docker container image for *training*. More details are provided below.\n",
    "* __`Dockerfile-inference`__ describes how to build your Docker container image for *inference*. More details are provided below.\n",
    "* __`build_and_push.sh`__ is a script that uses the Dockerfile to build your container images and then pushes it to ECR. We invoke the commands directly later in this notebook, but you can just copy and run the script for your own algorithms.\n",
    "* __`bert-topic`__ is the directory which contains our user code to be invoked.\n",
    "\n",
    "In this application, we install and/or update a few libraries, and copy one script in the container, which will be used as `ENTRYPOINT`. You may only need that many, but if you have many supporting routines, you may wish to install more and use more files.\n",
    "\n",
    "The files that we put in the container are:\n",
    "\n",
    "* __`bert-topic.py`__ is the program that implements our training algorithm (used only for training container)\n",
    "* __`bert-topic-inference.py`__ is the program that handles loading our model for inferences (used only for inference container)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 2: Packaging and Uploading your Algorithm for training\n",
    "\n",
    "### The training Dockerfile\n",
    "\n",
    "The Dockerfile describes the image that we want to build. You can think of it as describing the complete operating system installation of the system that you want to run. A Docker container running is quite a bit lighter than a full operating system, however, because it takes advantage of Linux on the host machine for the basic operations. \n",
    "\n",
    "We start from the SageMaker PyTorch image as the base. The base image is an ECR image, so it will have the following pattern.\n",
    "* {account}.dkr.ecr.{region}.amazonaws.com/sagemaker-{framework}:{framework_version}-{processor_type}-{python_version}\n",
    "\n",
    "Here is an explanation of each field.\n",
    "1. account - AWS account ID the ECR image belongs to. Our public deep learning framework images are all under the 763104351884 account.\n",
    "2. region - The region the ECR image belongs to. [Available regions](https://aws.amazon.com/about-aws/global-infrastructure/regional-product-services/).\n",
    "3. framework - The deep learning framework.\n",
    "4. framework_version - The version of the deep learning framework.\n",
    "5. processor_type - CPU or GPU.\n",
    "6. python_version - The supported version of Python.\n",
    "\n",
    "So the SageMaker PyTorch ECR image would be:\n",
    "* 763104351884.dkr.ecr.{region}.amazonaws.com/pytorch-training:1.12.1-gpu-py38-cu113-ubuntu20.04-sagemaker\n",
    "\n",
    "Information on supported frameworks and versions can be found in this [README](https://github.com/aws/sagemaker-python-sdk).\n",
    "\n",
    "Next, we install the required additional libraries like BERTopic and add the code that implements our specific algorithm to the container, and set up the right environment for it to run under.\n",
    "\n",
    "Finally, we need to specify two environment variables.\n",
    "1. SAGEMAKER_SUBMIT_DIRECTORY - the directory within the container containing our Python script for training and inference.\n",
    "2. SAGEMAKER_PROGRAM - the Python script that should be invoked for training and inference.\n",
    "\n",
    "Let's look at the Dockerfile for this example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!cat container/Dockerfile"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Writing the training script (bert-topic.py)\n",
    "\n",
    "To fine-tune a BERTopic model with a custom dataset on Amazon SageMaker, we will write a training script to be used by the Amazon SageMaker Training Job (or locally).\n",
    "\n",
    "The training script will need to do the following steps:\n",
    "- Load a pretrained Model\n",
    "- Load the dataset\n",
    "- Define the Training Arguments\n",
    "- Define a Trainer\n",
    "- Train the model and save the checkpoint on the validation set\n",
    "\n",
    "These steps will be done in a `_train()` function.\n",
    "\n",
    "The script also uses some hyperparameters, which can be extended depending on what is needed for each model. Here we used `language` to choose the training language."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!pygmentize container/bert-topic/bert-topic.py"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Building and registering the training container\n",
    "\n",
    "The following shell code shows how to build the container image using `docker build` and push the container image to ECR using `docker push`. \n",
    "\n",
    "This code looks for an ECR repository in the account you're using and the current default region (if you're using a SageMaker notebook instance, this is the region where the notebook instance was created). If the repository doesn't exist, the script will create it. In addition, since we are using the SageMaker PyTorch image as the base, we will need to retrieve ECR credentials to pull this public image.\n",
    "\n",
    "Note that your notebook role needs to have the permission to push images to the ECR registry."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# The name of our algorithm -- i.e. the name of the produced training container\n",
    "training_algorithm_name = \"bert-topic-training-example\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# building and pushing the container\n",
    "! cd container && sh build_and_push.sh {training_algorithm_name}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# An alternative and simplified command which outsources the docker build process to codebuild.\n",
    "# This is especially useful when using Sagemaker studio notebooks where docker is not running.\n",
    "# Details in https://aws.amazon.com/blogs/machine-learning/using-the-amazon-sagemaker-studio-image-build-cli-to-build-container-images-from-your-studio-notebooks/\n",
    "#!cd container && sm-docker build {training_algorithm_name}"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Testing your algorithm\n",
    "\n",
    "When you're packaging your first algorithm to use with Amazon SageMaker, you probably want to test it yourself to make sure it's working correctly. We use the [SageMaker Python SDK](https://github.com/aws/sagemaker-python-sdk) to test both locally and on SageMaker. For more examples with the SageMaker Python SDK, see [Amazon SageMaker Examples](https://github.com/awslabs/amazon-sagemaker-examples/tree/master/sagemaker-python-sdk). In order to test our algorithm, we need our dataset."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Download the 20 Newsgroups training dataset\n",
    "We will be utilizing the 20 Newsgroups dataset loader provided within AWS to download and load our data for training."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "training_file_name = \"training_file.txt\"\n",
    "working_dir = os.getcwd()\n",
    "training_file_path = os.path.join(working_dir, training_file_name)\n",
    "\n",
    "print(f\"Working Dir: {working_dir}\")\n",
    "print(f\"Training File: {training_file_path}\")\n",
    "docs = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# data from AWS internal repository, external repo can be used as alternative, see below\n",
    "import boto3\n",
    "\n",
    "s3 = boto3.client(\"s3\")\n",
    "s3.download_file(\n",
    "    f\"sagemaker-example-files-prod-{boto3.session.Session().region_name}\",\n",
    "    \"datasets/text/20_newsgroups/20_newsgroups.tar.gz\",\n",
    "    training_file_path.replace(\".txt\", \"_temp.tar.gz\"),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# extract few files from the rec.autos topic, without untarring the whole file. No pre-processing is applied\n",
    "import tarfile\n",
    "import codecs\n",
    "import gzip\n",
    "\n",
    "tar = tarfile.open(training_file_path.replace(\".txt\", \"_temp.tar.gz\"), encoding=\"utf-8\")\n",
    "utf8reader = codecs.getreader(\"utf-8\")\n",
    "for i, name in enumerate(tar.getmembers()):\n",
    "    if \"rec.autos\" in name.name:\n",
    "        f = tar.extractfile(name)\n",
    "        if f is not None:\n",
    "            content = f.read()\n",
    "            docs = docs + [content.decode(\"utf-8\").split(\"Message-ID:\")[1]]\n",
    "        if len(docs) > 100:\n",
    "            break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# # data from external repository, pre-processing is already applied\n",
    "# from sklearn.datasets import fetch_20newsgroups\n",
    "\n",
    "# docs = fetch_20newsgroups(subset=\"all\", remove=(\"headers\", \"footers\", \"quotes\"))[\"data\"]\n",
    "# # only use 100 documents out of 18846 for faster iteration\n",
    "# docs = docs[:100]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "len(docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "print(\"The first document contains the following text:\\n\\n\", docs[1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# write the list to a file to be easily passed to the training job\n",
    "\n",
    "with open(training_file_path, \"w\") as f:\n",
    "    for line in docs:\n",
    "        f.write(\n",
    "            line.replace(\"\\n\", \"\\\\n\") + \"\\n\"\n",
    "        )  # preserve newlines symbols (\\n) and write one document per line"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### SageMaker Python SDK Local Training\n",
    "To represent our training, we use the Estimator class, which needs to be configured in five steps. \n",
    "1. IAM role - our AWS execution role ---> make sure to add the \"AmazonEC2ContainerRegistryFullAccess\" policy to the execution role, otherwise the image push will fail\n",
    "2. train_instance_count - number of instances to use for training.\n",
    "3. train_instance_type - type of instance to use for training. For training locally, we specify `local` or `local_gpu`.\n",
    "4. image_name - our custom PyTorch Docker image we created.\n",
    "5. hyperparameters - hyperparameters we want to pass.\n",
    "\n",
    "Let's start with setting up our IAM role. We make use of a helper function within the Python SDK. This function throws an exception if runs outside a SageMaker notebook instance, as it gets metadata from the notebook instance. If running outside, you must provide an IAM role with proper access stated above in [Permissions](#Permissions)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from sagemaker import get_execution_role\n",
    "\n",
    "# role to be used for batch submission\n",
    "role = get_execution_role()\n",
    "print(\"role\", role)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Check the local SageMaker setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Lets set up our SageMaker notebook instance for local mode.\n",
    "!/bin/bash ./utils/setup.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# CHECK if GPUs are available and set the corresponding \"instance_type\"\n",
    "import os\n",
    "import subprocess\n",
    "\n",
    "instance_type_local = \"local\"\n",
    "\n",
    "if subprocess.call(\"nvidia-smi\") == 0:\n",
    "    ## Set type to GPU if one is present\n",
    "    instance_type_local = \"local_gpu\"\n",
    "\n",
    "print(\"Instance type = \" + instance_type_local)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Fit (+ Deploy, Predict)\n",
    "\n",
    "Now that the rest of our estimator is configured, we can call `fit()` with the path to our local 20 Newsgroups dataset prefixed with `file://`. This invokes our PyTorch container with 'train' and passes in our hyperparameters and other metadata as json files in /opt/ml/input/config within the container to our program entry point defined in the Dockerfile.\n",
    "\n",
    "After our training has succeeded, our training algorithm outputs our trained model within the /opt/ml/model directory, which is used to handle predictions.\n",
    "\n",
    "If the container used for the training is also used for the inference (NOT IN THIS CASE!) we can then conveniently call `deploy()` with an instance_count and instance_type, which is 1 and `local`. This invokes our PyTorch container with 'serve', which setups our container to handle prediction requests as defined [here](https://github.com/aws/sagemaker-pytorch-container/blob/master/src/sagemaker_pytorch_container/serving.py#L103). What is returned is a predictor, which is used to make inferences against our trained model.\n",
    "\n",
    "After our prediction, we can delete our endpoint.\n",
    "\n",
    "We recommend always testing and training your training algorithm locally first, as it provides quicker iterations and better debuggability."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Local training: Fit"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.estimator import Estimator\n",
    "\n",
    "# define hyperparameters\n",
    "hyperparameters = {\"language\": \"english\"}\n",
    "\n",
    "# prepare training job\n",
    "estimator = Estimator(\n",
    "    role=role,\n",
    "    train_instance_count=1,\n",
    "    train_instance_type=instance_type_local,\n",
    "    image_uri=training_algorithm_name + \":latest\",\n",
    "    hyperparameters=hyperparameters,\n",
    ")\n",
    "\n",
    "# launch training job\n",
    "print(f\"file://{training_file_path}\")\n",
    "estimator.fit(f\"file://{training_file_path}\")\n",
    "# estimator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check where the fitted model has been stored after fit\n",
    "estimator.model_data"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### A sleek alternative to \"Bring your own container\": bring your own requirements.txt\n",
    "\n",
    "Instead of extending the PyTorch container through a Dockerfile, in many cases we can simply import additional libraries to make our model work. \n",
    "In these cases, adding the requirements.txt inside the same folder as the code will install the required dependencies at runtime. \n",
    "The source directory must be specified in the `source_dir` argument, and the entry point in the `entry_point` argument, when creating PyTorch estimator. Further documentation [here](https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#using-third-party-libraries)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "\n",
    "my_session = boto3.session.Session()\n",
    "my_region = my_session.region_name\n",
    "my_image_uri = (\n",
    "    \"763104351884.dkr.ecr.\"\n",
    "    + my_region\n",
    "    + \".amazonaws.com/pytorch-training:1.12.1-gpu-py38-cu113-ubuntu20.04-sagemaker\"\n",
    ")\n",
    "print(\n",
    "    \"my_image_uri:\",\n",
    "    my_image_uri,\n",
    "    \"\\nrole:\",\n",
    "    role,\n",
    ")\n",
    "from sagemaker.estimator import Estimator\n",
    "\n",
    "# define hyperparameters\n",
    "hyperparameters = {\"language\": \"english\"}\n",
    "\n",
    "# prepare training job\n",
    "estimator = Estimator(\n",
    "    role=role,\n",
    "    train_instance_count=1,\n",
    "    train_instance_type=instance_type_local,\n",
    "    image_uri=my_image_uri,\n",
    "    hyperparameters=hyperparameters,\n",
    "    source_dir=\"container/bert-topic-byoreq\",  # this directory contains the entry point code and the requirements.txt\n",
    "    entry_point=\"bert-topic.py\",  # this argument is used to override internal container entrypoint, if needed!\n",
    ")\n",
    "\n",
    "# launch training job\n",
    "print(f\"file://{training_file_path}\")\n",
    "estimator.fit(f\"file://{training_file_path}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "estimator.model_data"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Training your Algorithm in Amazon SageMaker in batch mode\n",
    "Once you have your container packaged and/or validated, you can use it to train models in batch mode. Let's do that with the algorithm we made above.\n",
    "\n",
    "#### Set up the environment\n",
    "Here we specify the bucket to use (the role has been defined earlier)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# S3 prefix\n",
    "prefix = \"DEMO-pytorch-bert-topic\""
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create the session\n",
    "\n",
    "The session remembers our connection parameters to SageMaker. We use it to perform all of our SageMaker operations."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker as sage\n",
    "\n",
    "sess = sage.Session()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Upload the data for training\n",
    "\n",
    "We will use the tools provided by the SageMaker Python SDK to upload the data to a default bucket."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_location = sess.upload_data(training_file_path, key_prefix=prefix)\n",
    "data_location"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Training on SageMaker in batch mode\n",
    "Training a model on SageMaker with the Python SDK is done in a way that is similar to the way we trained it locally. This is done by changing our train_instance_type from `local` to one of our [supported EC2 instance types](https://aws.amazon.com/sagemaker/pricing/instance-types/).\n",
    "\n",
    "In addition, we must now specify the ECR image URL, which we just pushed above.\n",
    "\n",
    "Finally, our local training dataset has to be in Amazon S3 and the S3 URL to our dataset is passed into the `fit()` call.\n",
    "\n",
    "Let's first fetch our ECR image url that corresponds to the image we just built and pushed.\n",
    "\n",
    "Also in this case, whenever useful, one can use the leverage a pre-existing base image with the \"bring your own requirements.txt\" option (not shown below)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "\n",
    "client = boto3.client(\"sts\")\n",
    "account = client.get_caller_identity()[\"Account\"]\n",
    "\n",
    "my_session = boto3.session.Session()\n",
    "region = my_session.region_name\n",
    "\n",
    "ecr_image = f\"{account}.dkr.ecr.{region}.amazonaws.com/{training_algorithm_name}:latest\"\n",
    "\n",
    "print(\"using ecr_image:\", ecr_image)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After training the mode, the output will be packaged into a tarred file `model.tar.gz` and copied to the S3 bucket in use for the training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker.estimator import Estimator\n",
    "\n",
    "hyperparameters = {\"language\": \"english\"}\n",
    "\n",
    "# instance_type = \"ml.g4dn.xlarge\"\n",
    "instance_type = \"ml.p2.xlarge\"\n",
    "\n",
    "estimator = Estimator(\n",
    "    role=role,\n",
    "    train_instance_count=1,\n",
    "    train_instance_type=instance_type,\n",
    "    image_uri=ecr_image,\n",
    "    hyperparameters=hyperparameters,\n",
    ")\n",
    "\n",
    "estimator.fit(data_location)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# check where the fitted model has been stored after fit\n",
    "estimator.model_data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! docker system df\n",
    "! docker image ls\n",
    "# remove downloaded docker images to avoid \"no space left on device\" message\n",
    "! docker system prune -f\n",
    "! docker image prune --all --force"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Part 3: Packaging and Uploading your Algorithm for inference\n",
    "\n",
    "### The inference Dockerfile\n",
    "\n",
    "Also in this case the Dockerfile describes the image that we want to build. \n",
    "\n",
    "We start from the SageMaker PyTorch image as the base, but the *inference* one. \n",
    "\n",
    "So the SageMaker PyTorch ECR image in this case would be:\n",
    "* FROM 763104351884.dkr.ecr.{region}.amazonaws.com/pytorch-inference:1.12.1-gpu-py38-cu113-ubuntu20.04-sagemaker\n",
    "\n",
    "Next, we install the required additional libraries (we make sure that to install numba==0.53.1 and nvgpu together with berttopic, otherwise we would get an error!) and add the code that implements our specific algorithm to the container, and set up the right environment for it to run under.\n",
    "\n",
    "Finally, we need to specify two environment variables.\n",
    "1. SAGEMAKER_SUBMIT_DIRECTORY - the directory within the container containing our Python script for training and inference.\n",
    "2. SAGEMAKER_PROGRAM - the Python script that should be invoked for training and inference.\n",
    "\n",
    "Let's look at the Dockerfile for this example."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat container/Dockerfile-inference"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Writing your own inference script (bert-topic-inference.py)\n",
    "\n",
    "Given the use of a pre-packaged SageMaker PyTorch container, the only requirement to write an inference script is that it has to define the following template functions:\n",
    "- `model_fn()` reading the content of an existing model weights directory saved as a `tar.gz` in s3. We will use it to load the trained Model.\n",
    "- `input_fn()` used here simply to format the data receives from a request made to the endpoint.\n",
    "- `predict_fn()` calling the output of `model_fn()` to run inference on the output of `input_fn()`.\n",
    "\n",
    "Optionally a `output_fn()` can be created for inference formatting, using the output of `predict_fn()`. Here it is especially useful since the BERTopic inference output is not really standard, featuring a tuple containing 2 lists.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pygmentize container/bert-topic/bert-topic-inference.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The name of our algorithm -- i.e. the name of the inference container\n",
    "inference_algorithm_name = \"pytorch-inference-bert-topic-example\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# building and pushing the container\n",
    "! cd container && sh build_and_push.sh {inference_algorithm_name} Dockerfile-inference"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inference from Containerized SageMaker Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pprint\n",
    "import boto3\n",
    "\n",
    "pp = pprint.PrettyPrinter(indent=1)\n",
    "\n",
    "sm_boto3 = boto3.client(\"sagemaker\")\n",
    "\n",
    "region = boto3.Session().region_name\n",
    "account_id = boto3.client(\"sts\").get_caller_identity()[\"Account\"]\n",
    "\n",
    "image_uri_inference = (\n",
    "    f\"{account_id}.dkr.ecr.{region}.amazonaws.com/{inference_algorithm_name}:latest\"\n",
    ")\n",
    "image_uri_inference"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Prepare session to run container locally"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sagemaker\n",
    "from sagemaker.local import LocalSession\n",
    "\n",
    "session_local = LocalSession()\n",
    "session_local.config = {instance_type_local: {\"local_code\": True}}\n",
    "print(type(session_local))\n",
    "\n",
    "from sagemaker import get_execution_role\n",
    "\n",
    "role = get_execution_role()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Check for locally running Docker containers (and stop them if needed)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! docker ps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# if model doesn't deploy locally, kill other running containers\n",
    "# docker stop $(docker ps -a -q)\n",
    "# docker rm $(docker ps -a -q)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Deploy container locally to create a local endpoint \n",
    "##### (WARNING: cells stdout will start behaving fuzzy, showing container log info, till the local container is active...)\n",
    "\n",
    "Disclaimer: if the inference container is missing the default `behavior` is possible to use local files (`source_dir` and `entry_point` arguments)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker import Model\n",
    "\n",
    "model_data = estimator.model_data\n",
    "\n",
    "estimator = Model(\n",
    "    image_uri=image_uri_inference,\n",
    "    model_data=model_data,\n",
    "    role=role,\n",
    "    source_dir=\"container/bert-topic\",  # this argument is used to override internal container entrypoint, if needed!\n",
    "    entry_point=\"bert-topic-inference.py\",  # this argument is used to override internal container entrypoint, if needed!\n",
    "    sagemaker_session=session_local,  # local session\n",
    "    #                   predictor_cls=None,\n",
    "    #                   env=None,\n",
    "    #                   name=None,\n",
    "    #                   vpc_config=None,\n",
    "    #                   enable_network_isolation=False,\n",
    "    #                   model_kms_key=None,\n",
    "    #                   image_config=None,\n",
    "    #                   code_location=None,\n",
    "    #                   container_log_level=20,\n",
    "    #                   dependencies=None,\n",
    "    #                   git_config=None\n",
    ")\n",
    "\n",
    "predictor = estimator.deploy(1, instance_type_local)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Invoke locally deployed endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "sagemaker_session = LocalSession()\n",
    "sagemaker_session.config = {instance_type_local: {\"local_code\": True}}\n",
    "\n",
    "sm_client = sagemaker_session.sagemaker_runtime_client\n",
    "response = sm_client.invoke_endpoint(\n",
    "    EndpointName=\"local-endpoint\",\n",
    "    ContentType=\"application/json\",\n",
    "    Body=json.dumps([\"Do you like Toyota Corolla?\", \"I have a Nissan\"]),\n",
    ")\n",
    "\n",
    "r = response[\"Body\"]\n",
    "print(\"RESULT r.read().decode():\", r.read().decode())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%sh\n",
    "# check if docker containers are running and kill them if needed\n",
    "docker ps\n",
    "docker stop $(docker ps -a -q)\n",
    "docker rm $(docker ps -a -q)\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### A sleek alternative to \"Bring your own container\": bring your own requirements.txt\n",
    "\n",
    "Also in this case, instead of extending the PyTorch container through a Dockerfile, we can simply import additional libraries to make our model work, if needed. \n",
    "The process is the same, i.e. adding the requirements.txt inside the same folder as the code will install the required dependencies at runtime. \n",
    "The source directory must be specified in the `source_dir` argument, and the entry point in the `entry_point` argument, when creating PyTorch estimator. Further documentation [here](https://sagemaker.readthedocs.io/en/stable/frameworks/pytorch/using_pytorch.html#using-third-party-libraries)\n",
    "\n",
    "*CAVEAT*: Installing libraries on the fly through requirements.txt implies additional overhead. This is negligible in case of \"always on\" single model endpoint deployment, but could have an impact in case of serverless or multi-model deployment. \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "\n",
    "my_session = boto3.session.Session()\n",
    "my_region = my_session.region_name\n",
    "my_image_uri_inference = (\n",
    "    \"763104351884.dkr.ecr.\"\n",
    "    + my_region\n",
    "    + \".amazonaws.com/pytorch-inference:1.12.1-gpu-py38-cu113-ubuntu20.04-sagemaker\"\n",
    ")\n",
    "print(\"my_image_uri_inference:\", my_image_uri_inference)\n",
    "from sagemaker.estimator import Estimator\n",
    "\n",
    "# define hyperparameters\n",
    "hyperparameters = {\"language\": \"english\"}\n",
    "\n",
    "estimator = Model(\n",
    "    image_uri=my_image_uri_inference,\n",
    "    model_data=model_data,\n",
    "    role=role,\n",
    "    source_dir=\"container/bert-topic-byoreq\",  # this directory contains the entrypoint code and the requirements.txt\n",
    "    entry_point=\"bert-topic-inference.py\",  # this argument is used to override internal container entrypoint, if needed!\n",
    "    sagemaker_session=session_local,  # local session\n",
    "    #                   predictor_cls=None,\n",
    "    #                   env=None,\n",
    "    #                   name=None,\n",
    "    #                   vpc_config=None,\n",
    "    #                   enable_network_isolation=False,\n",
    "    #                   model_kms_key=None,\n",
    "    #                   image_config=None,\n",
    "    #                   code_location=None,\n",
    "    #                   container_log_level=20,\n",
    "    #                   dependencies=None,\n",
    "    #                   git_config=None\n",
    ")\n",
    "\n",
    "predictor = estimator.deploy(1, instance_type_local)  #  (ignore the nvgpu errors)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Invoke locally deployed endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "sagemaker_session = LocalSession()\n",
    "sagemaker_session.config = {instance_type_local: {\"local_code\": True}}\n",
    "\n",
    "sm_client = sagemaker_session.sagemaker_runtime_client\n",
    "response = sm_client.invoke_endpoint(\n",
    "    EndpointName=\"local-endpoint\",\n",
    "    ContentType=\"application/json\",\n",
    "    Body=json.dumps([\"Do you like Toyota Corolla?\", \"I have a Nissan\"]),\n",
    ")\n",
    "\n",
    "r = response[\"Body\"]\n",
    "print(\"RESULT r.read().decode():\", r.read().decode())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%sh\n",
    "# check if docker containers are running and kill them if needed\n",
    "docker ps\n",
    "docker stop $(docker ps -a -q)\n",
    "docker rm $(docker ps -a -q)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Deploy container remotely to create a managed Amazon SageMaker endpoint for real-time inference"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sagemaker import Model\n",
    "\n",
    "import sagemaker as sage\n",
    "\n",
    "sess = sage.Session()\n",
    "\n",
    "# instance_type = \"ml.m5.xlarge\" # no GPU, will trigger an error\n",
    "# instance_type = \"ml.g4dn.xlarge\"\n",
    "instance_type = \"ml.p2.xlarge\"\n",
    "model_data = estimator.model_data\n",
    "\n",
    "endpoint_deployment_name = \"my-BERTopic-test-endpoint\"\n",
    "\n",
    "estimator = Model(\n",
    "    image_uri=image_uri_inference,\n",
    "    model_data=model_data,\n",
    "    role=role,\n",
    "    source_dir=\"container/bert-topic\",\n",
    "    entry_point=\"bert-topic-inference.py\",\n",
    "    sagemaker_session=sess,  # not local session anymore\n",
    "    #                   predictor_cls=None,\n",
    "    #                   env=None,\n",
    "    #                   name=None,\n",
    "    #                   vpc_config=None,\n",
    "    #                   enable_network_isolation=False,\n",
    "    #                   model_kms_key=None,\n",
    "    #                   image_config=None,\n",
    "    #                   code_location=None,\n",
    "    #                   container_log_level=20,\n",
    "    #                   dependencies=None,\n",
    "    #                   git_config=None\n",
    ")\n",
    "\n",
    "# deploy the model\n",
    "predictor = estimator.deploy(1, instance_type, endpoint_name=endpoint_deployment_name)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Manually check the deployed endpoint name, it will have the name \n",
    "- my-BERTopic-test-endpoint"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Invoke remotely deployed endpoint"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm_client = sess.sagemaker_runtime_client\n",
    "endpoint_name = \"my-BERTopic-test-endpoint\"\n",
    "response = sm_client.invoke_endpoint(\n",
    "    EndpointName=endpoint_name,\n",
    "    ContentType=\"application/json\",\n",
    "    Body=json.dumps([\"Do you like Toyota Corolla?\", \"I have a Nissan\"]),\n",
    ")\n",
    "\n",
    "r = response[\"Body\"]\n",
    "print(\"RESULT r.read().decode():\", r.read().decode())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Optional cleanup of the create endpoint\n",
    "The created endpoint can be deleted with the code below.\n",
    "\n",
    "This part represent the end of the notebook."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "\n",
    "client = boto3.client(\"sagemaker\")\n",
    "response = client.describe_endpoint_config(EndpointConfigName=endpoint_deployment_name)\n",
    "model_name = response[\"ProductionVariants\"][0][\"ModelName\"]\n",
    "client.delete_endpoint(EndpointName=endpoint_deployment_name)\n",
    "client.delete_endpoint_config(EndpointConfigName=endpoint_deployment_name)\n",
    "client.delete_model(ModelName=model_name)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reference\n",
    "- [How Amazon SageMaker interacts with your Docker container for training](https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-training-algo.html)\n",
    "- [How Amazon SageMaker interacts with your Docker container for inference](https://docs.aws.amazon.com/sagemaker/latest/dg/your-algorithms-inference-code.html)\n",
    "- [BERTtopic model](https://maartengr.github.io/BERTopic/index.html)\n",
    "- [20 Newsgroups dataset](http://qwone.com/~jason/20Newsgroups/)\n",
    "- [SageMaker Python SDK](https://github.com/aws/sagemaker-python-sdk)\n",
    "- [Dockerfile](https://docs.docker.com/engine/reference/builder/)\n",
    "- [scikit-bring-your-own](https://github.com/awslabs/amazon-sagemaker-examples/blob/master/advanced_functionality/scikit_bring_your_own/scikit_bring_your_own.ipynb)\n",
    "- [SageMaker PyTorch container](https://github.com/aws/sagemaker-pytorch-container)\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook CI Test Results\n",
    "\n",
    "This notebook was tested in multiple regions. The test results are as follows, except for us-west-2 which is shown at the top of the notebook.\n",
    "\n",
    "![This us-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-1/advanced_functionality|pytorch_extend_container_train_deploy_bertopic|BERTtopic_extending_container.ipynb)\n",
    "\n",
    "![This us-east-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-east-2/advanced_functionality|pytorch_extend_container_train_deploy_bertopic|BERTtopic_extending_container.ipynb)\n",
    "\n",
    "![This us-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/us-west-1/advanced_functionality|pytorch_extend_container_train_deploy_bertopic|BERTtopic_extending_container.ipynb)\n",
    "\n",
    "![This ca-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ca-central-1/advanced_functionality|pytorch_extend_container_train_deploy_bertopic|BERTtopic_extending_container.ipynb)\n",
    "\n",
    "![This sa-east-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/sa-east-1/advanced_functionality|pytorch_extend_container_train_deploy_bertopic|BERTtopic_extending_container.ipynb)\n",
    "\n",
    "![This eu-west-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-1/advanced_functionality|pytorch_extend_container_train_deploy_bertopic|BERTtopic_extending_container.ipynb)\n",
    "\n",
    "![This eu-west-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-2/advanced_functionality|pytorch_extend_container_train_deploy_bertopic|BERTtopic_extending_container.ipynb)\n",
    "\n",
    "![This eu-west-3 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-west-3/advanced_functionality|pytorch_extend_container_train_deploy_bertopic|BERTtopic_extending_container.ipynb)\n",
    "\n",
    "![This eu-central-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-central-1/advanced_functionality|pytorch_extend_container_train_deploy_bertopic|BERTtopic_extending_container.ipynb)\n",
    "\n",
    "![This eu-north-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/eu-north-1/advanced_functionality|pytorch_extend_container_train_deploy_bertopic|BERTtopic_extending_container.ipynb)\n",
    "\n",
    "![This ap-southeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-1/advanced_functionality|pytorch_extend_container_train_deploy_bertopic|BERTtopic_extending_container.ipynb)\n",
    "\n",
    "![This ap-southeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-southeast-2/advanced_functionality|pytorch_extend_container_train_deploy_bertopic|BERTtopic_extending_container.ipynb)\n",
    "\n",
    "![This ap-northeast-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-1/advanced_functionality|pytorch_extend_container_train_deploy_bertopic|BERTtopic_extending_container.ipynb)\n",
    "\n",
    "![This ap-northeast-2 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-northeast-2/advanced_functionality|pytorch_extend_container_train_deploy_bertopic|BERTtopic_extending_container.ipynb)\n",
    "\n",
    "![This ap-south-1 badge failed to load. Check your device's internet connectivity, otherwise the service is currently unavailable](https://prod.us-west-2.tcx-beacon.docs.aws.dev/sagemaker-nb/ap-south-1/advanced_functionality|pytorch_extend_container_train_deploy_bertopic|BERTtopic_extending_container.ipynb)\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_p39",
   "language": "python",
   "name": "conda_pytorch_p39"
  },
  "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.9.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}