integrations-and-supported-tools/sagemaker/custom-docker-container/Using-Neptune-with-custom-container.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "abc6583e-be23-457f-9f35-d5e790aa9636",
   "metadata": {},
   "source": [
    "![Neptune + Amazon Sagemaker](https://neptune.ai/wp-content/uploads/2023/09/sagemaker.svg)\n",
    "\n",
    "# Using Neptune together with Amazon SageMaker training jobs \n",
    "\n",
    "<div class=\"alert alert-info\">You can run this part of the notebook either locally or from a SageMaker notebook. It would need additional dependencies such as AWS CLI tools and Docker.</div>\n",
    "\n",
    "This tutorial uses some code (with adaptations) from the [official AWS tutorial](https://github.com/aws/amazon-sagemaker-examples/tree/main/advanced_functionality/scikit_bring_your_own). We'll show how to add Neptune logging to a custom training job in SageMaker. For this, we are going to create a Docker container with pre-installed Neptune, and adapt Amazon's code by adding Netune logging to it."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6121b6c4-72ee-43ac-bef0-ad835de5bfbd",
   "metadata": {},
   "source": [
    "## Docker container\n",
    "\n",
    "Our container is a simplified version of the container in AWS's tutorial. We are using `python` base image and added `neptune` and `neptune-sklearn` as additional dependencies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29a9242d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "FROM python:3-slim-buster\n",
      "\n",
      "RUN pip --quiet --no-cache-dir install \\\n",
      "    numpy scipy scikit-learn pandas \\\n",
      "    neptune-sklearn\n",
      "\n",
      "ENV PYTHONUNBUFFERED=TRUE\n",
      "ENV PYTHONDONTWRITEBYTECODE=TRUE\n",
      "ENV PATH=\"/opt/program:${PATH}\"\n",
      "\n",
      "# Set up the program in the image\n",
      "COPY train /opt/program/train\n",
      "WORKDIR /opt/program\n"
     ]
    }
   ],
   "source": [
    "!cat Dockerfile"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "cb87bfb5-8313-4f4a-9399-b5acfaba5df9",
   "metadata": {},
   "source": [
    "### Training script\n",
    "\n",
    "The training script can be found in `decision_trees/train`. As compared to the script from AWS's tutorial, we added a few lines of our code:\n",
    "\n",
    "```diff\n",
    "[...]\n",
    "\n",
    "+ import neptune\n",
    "+ import neptune.integrations.sklearn as npt_utils\n",
    "\n",
    "[...]\n",
    "\n",
    "def train():\n",
    "+    run = neptune.init_run(\n",
    "+         tags=[\"sagemaker\"],\n",
    "+         source_files=[\"train\"],\n",
    "+    )\n",
    "     [...]\n",
    "     \n",
    "     # Now use scikit-learn's decision tree classifier to train the model.\n",
    "     clf = tree.DecisionTreeClassifier(max_leaf_nodes=max_leaf_nodes)\n",
    "     clf = clf.fit(train_X, train_y)\n",
    "\n",
    "+    run[\"cls_summary\"] = npt_utils.create_classifier_summary(\n",
    "+       clf, train_X, train_X, train_y, train_y\n",
    "+    )\n",
    "     [...]\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1351effd-f554-4621-b552-02bd273abb5e",
   "metadata": {},
   "source": [
    "### Build the Docker image and push it to ECR\n",
    "\n",
    "Next, we need to build the container. The following Bash script assumes that you have [AWS CLI](https://aws.amazon.com/cli/) installed on your machine. It again uses the code from the AWS tutorial. It is lengthy, because it automatically creates an ECR repository for us and pushes the image to it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6dfbd2e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "#!/bin/bash\n",
      "set -e\n",
      "\n",
      "algorithm_name=\"neptune-sagemaker-demo\"\n",
      "account=$(aws sts get-caller-identity --query Account --output text)\n",
      "chmod +x train\n",
      "\n",
      "# Get the region defined in the current configuration (default to us-west-2 if none defined)\n",
      "region=$(aws configure get region)\n",
      "region=\"${region:-us-west-2}\"\n",
      "\n",
      "full_name=\"${account}.dkr.ecr.${region}.amazonaws.com/${algorithm_name}:latest\"\n",
      "echo \"Image: ${full_name}\"\n",
      "\n",
      "# If the repository doesn't exist in ECR, create it.\n",
      "if ! aws ecr describe-repositories --repository-names \"${algorithm_name}\" > /dev/null\n",
      "then\n",
      "    aws ecr create-repository --repository-name \"${algorithm_name}\" > /dev/null\n",
      "fi\n",
      "\n",
      "# Get the login command from ECR and execute it directly\n",
      "aws ecr get-login-password --region \"${region}\" | docker login --username AWS --password-stdin \"${full_name}\"\n",
      "\n",
      "# Build the docker image locally with the image name and then push it to ECR\n",
      "# with the full name.\n",
      "docker build -t ${algorithm_name} .\n",
      "docker tag \"${algorithm_name}\" \"${full_name}\"\n",
      "docker push \"${full_name}\"\n",
      "\n",
      "echo \"Success!\"\n"
     ]
    }
   ],
   "source": [
    "!cat build_and_push.sh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e054e735",
   "metadata": {},
   "outputs": [],
   "source": [
    "!bash ./build_and_push.sh"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "id": "e35d838f-b0e0-4491-9d71-7b52799e47a6",
   "metadata": {},
   "source": [
    "## Start the training job\n",
    "\n",
    "<div class=\"alert alert-info\">This part of the notebook is to be run from an Amazon SageMaker notebook.</div>"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4d635bbe-d2fe-453b-be68-96d301d3e8ee",
   "metadata": {},
   "source": [
    "### Obtaining the Neptune token from AWS Secrets\n",
    "\n",
    "If you store a Neptune API token and project name in AWS Secrets, you can read them using the following code. If you do that, you can use your token and project name in the place of `NEPTUNE_API_TOKEN` and `NEPTUNE_PROJECT` in the next cell. Alternatively, you can add the code below to the `decision_trees/train` script before building the Docker image and read the secrets from the `secrets` dictionary instead of `os.envir`.\n",
    "\n",
    "If you want to read the secrets form AWS Secrets, make sure that your SageMaker Notebook has a role that allows for the access to the secrets, in particular the `secretsmanager:GetSecretValue` permission for the appropriate secret."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bba763f2-07ac-4290-8985-d7ddd9f4e4d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "from botocore.exceptions import ClientError\n",
    "import json\n",
    "\n",
    "secret_name = \"AmazonSageMaker-tw-neptune-v2\"  # replace with yours\n",
    "region_name = \"us-east-1\"  # replace with yours\n",
    "\n",
    "# Create a Secrets Manager client\n",
    "session = boto3.session.Session()\n",
    "client = session.client(service_name=\"secretsmanager\", region_name=region_name)\n",
    "\n",
    "get_secret_value_response = client.get_secret_value(SecretId=secret_name)\n",
    "\n",
    "# Decrypts secret using the associated KMS key.\n",
    "secrets = json.loads(get_secret_value_response[\"SecretString\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4a00286-7b46-4d66-bd49-4d3fc57ddc47",
   "metadata": {},
   "source": [
    "In the example below, we are going to use the anonymous token `neptune.ANONYMOUS_API_TOKEN`, so we first need to install the Neptune client library to obtain the token."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ddde67c2",
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install -U neptune"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16af2a91-fcf0-4961-8a42-09d2d642bc7d",
   "metadata": {},
   "source": [
    "## Training data\n",
    "\n",
    "We are going to use the Iris dataset. Below, we are downloading it from the official SageMaker repository of sample datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f9ec27ba-2db5-4292-8911-ec073b44534f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import boto3\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import os\n",
    "\n",
    "os.makedirs(\"./data\", exist_ok=True)\n",
    "\n",
    "s3_client = boto3.client(\"s3\")\n",
    "s3_client.download_file(\n",
    "    f\"sagemaker-sample-files\", \"datasets/tabular/iris/iris.data\", \"./data/iris.csv\"\n",
    ")\n",
    "\n",
    "df_iris = pd.read_csv(\"./data/iris.csv\", header=None)\n",
    "df_iris[4] = df_iris[4].map({\"Iris-setosa\": 0, \"Iris-versicolor\": 1, \"Iris-virginica\": 2})\n",
    "iris = df_iris[[4, 0, 1, 2, 3]].to_numpy()\n",
    "np.savetxt(\"./data/iris.csv\", iris, delimiter=\",\", fmt=\"%1.1f, %1.3f, %1.3f, %1.3f, %1.3f\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34980560-2b60-456d-800a-89790dd67663",
   "metadata": {},
   "source": [
    "## Training the machine learning model\n",
    "\n",
    "The code below was taken from the official AWS tutorial. The only difference is that we are passing the `NEPTUNE_API_TOKEN` and `NEPTUNE_PROJECT` as environment variables to the [estimator](https://sagemaker.readthedocs.io/en/stable/api/training/estimators.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2db1de88-4983-42d5-9762-aa625243e9c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from sagemaker import get_execution_role\n",
    "import sagemaker as sage\n",
    "import neptune\n",
    "\n",
    "# S3 prefix\n",
    "s3_prefix = \"neptune-sagemaker-demo-data\"\n",
    "\n",
    "role = get_execution_role()\n",
    "\n",
    "sess = sage.Session()\n",
    "\n",
    "WORK_DIRECTORY = \"data\"\n",
    "\n",
    "data_location = sess.upload_data(WORK_DIRECTORY, key_prefix=s3_prefix)\n",
    "\n",
    "account = sess.boto_session.client(\"sts\").get_caller_identity()[\"Account\"]\n",
    "region = sess.boto_session.region_name\n",
    "image = \"{}.dkr.ecr.{}.amazonaws.com/neptune-sagemaker-demo:latest\".format(account, region)\n",
    "\n",
    "tree = sage.estimator.Estimator(\n",
    "    image_uri=image,\n",
    "    role=role,\n",
    "    instance_count=1,\n",
    "    instance_type=\"ml.m5.large\",\n",
    "    output_path=\"s3://{}/output\".format(sess.default_bucket()),\n",
    "    sagemaker_session=sess,\n",
    "    environment={\n",
    "        \"NEPTUNE_API_TOKEN\": neptune.ANONYMOUS_API_TOKEN,\n",
    "        \"NEPTUNE_PROJECT\": \"common/showroom\",\n",
    "    },\n",
    ")\n",
    "\n",
    "tree.fit(data_location)"
   ]
  }
 ],
 "metadata": {
  "keep_output": true,
  "kernelspec": {
   "display_name": "py38",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.15"
  },
  "vscode": {
   "interpreter": {
    "hash": "a9715cf0b0024f6e1c62cb31a4f1f43970eb41991212681878768b4bfe53050a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}