{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Custom SageMaker Studio Kernel Demo\n",
    "First, open SageMaker Studio (**this notebook should run inside Studio**).\n",
    "\n",
    "### 1. Get your the IAM Role associated with this user inside Studio:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sagemaker\n",
    "import boto3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sm_client = boto3.client(\"sagemaker\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sess = sagemaker.session.Session()\n",
    "sess.get_caller_identity_arn()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Go to the IAM console and choose `Roles`. Find the role above.\n",
    "\n",
    "Add a new policy to the role, click on `Attach Policies`, then `Create Policy`.\n",
    "\n",
    "Choose the JSON editor, and substitute everything by:\n",
    "```\n",
    "{\n",
    "    \"Version\": \"2012-10-17\",\n",
    "    \"Statement\": [\n",
    "        {\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": [\n",
    "                \"ecr:CreateRepository\",\n",
    "                \"ecr:BatchGetImage\",\n",
    "                \"ecr:CompleteLayerUpload\",\n",
    "                \"ecr:DescribeImages\",\n",
    "                \"ecr:DescribeRepositories\",\n",
    "                \"ecr:UploadLayerPart\",\n",
    "                \"ecr:ListImages\",\n",
    "                \"ecr:InitiateLayerUpload\",\n",
    "                \"ecr:BatchCheckLayerAvailability\",\n",
    "                \"ecr:GetDownloadUrlForLayer\",\n",
    "                \"ecr:PutImage\"\n",
    "            ],\n",
    "            \"Resource\": \"arn:aws:ecr:*:*:repository/smstudio*\"\n",
    "        },\n",
    "        {\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": \"ecr:GetAuthorizationToken\",\n",
    "            \"Resource\": \"*\"\n",
    "        },\n",
    "        {\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": [\n",
    "                \"codebuild:DeleteProject\",\n",
    "                \"codebuild:CreateProject\",\n",
    "                \"codebuild:BatchGetBuilds\",\n",
    "                \"codebuild:StartBuild\"\n",
    "            ],\n",
    "            \"Resource\": \"arn:aws:codebuild:*:*:project/sagemaker-studio*\"\n",
    "        },\n",
    "        {\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": \"iam:PassRole\",\n",
    "            \"Resource\": \"arn:aws:iam::*:role/*\",\n",
    "            \"Condition\": {\n",
    "                \"StringLikeIfExists\": {\n",
    "                    \"iam:PassedToService\": \"codebuild.amazonaws.com\"\n",
    "                }\n",
    "            }\n",
    "        },\n",
    "        {\n",
    "            \"Effect\": \"Allow\",\n",
    "            \"Action\": \"sagemaker:UpdateDomain\",\n",
    "            \"Resource\": \"*\"\n",
    "        }\n",
    "    ]\n",
    "}\n",
    "```\n",
    "\n",
    "Name the IAM policy `SageMakerStudioBuildCustomKernel`.\n",
    "\n",
    "### 3. Go back to the IAM role and attach this new created policy.\n",
    "\n",
    "Still inside the IAM Role, go to the tab `Trust relationships`, click on `Edit trust relationship`.\n",
    "\n",
    "Add the following:\n",
    "```\n",
    "{\n",
    "  \"Version\": \"2012-10-17\",\n",
    "  \"Statement\": [\n",
    "    {\n",
    "      ...\n",
    "    },\n",
    "    \n",
    "    ...\n",
    "    \n",
    "    {\n",
    "      \"Effect\": \"Allow\",\n",
    "      \"Principal\": {\n",
    "        \"Service\": \"codebuild.amazonaws.com\"\n",
    "      },\n",
    "      \"Action\": \"sts:AssumeRole\"\n",
    "    }\n",
    "  ]\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4. Create Dockerfile"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile Dockerfile\n",
    "FROM tensorflow/tensorflow:2.3.0\n",
    "RUN apt-get update \n",
    "RUN apt-get install -y git\n",
    "RUN pip install --upgrade pip\n",
    "RUN pip install ipykernel && \\\n",
    "    python -m ipykernel install --sys-prefix && \\\n",
    "    pip install --quiet --no-cache-dir \\\n",
    "    'boto3>1.0<2.0' \\\n",
    "    'sagemaker>2.0<3.0'\n",
    "    \n",
    "WORKDIR /root\n",
    "COPY train.py train.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 5. Create training script"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%writefile train.py\n",
    "import tensorflow as tf\n",
    "import os \n",
    "\n",
    "mnist = tf.keras.datasets.mnist\n",
    "(x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
    "x_train, x_test = x_train / 255.0, x_test / 255.0\n",
    "\n",
    "model = tf.keras.models.Sequential([\n",
    "  tf.keras.layers.Flatten(input_shape=(28, 28)),\n",
    "  tf.keras.layers.Dense(128, activation='relu'),\n",
    "  tf.keras.layers.Dropout(0.2),\n",
    "  tf.keras.layers.Dense(10, activation='softmax')\n",
    "])\n",
    "model.compile(optimizer='adam',\n",
    "              loss='sparse_categorical_crossentropy',\n",
    "              metrics=['accuracy'])\n",
    "model.fit(x_train, y_train, epochs=1)\n",
    "model.evaluate(x_test, y_test)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Create some environment variables to help"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def set_env_var(key, value):\n",
    "    os.environ[key]=value\n",
    "    print(f\"Environment variable:  {key}={value}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_env_var(\"APP_CONF_NAME\", \"custom-tf2\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_env_var(\"KERNEL_IMAGE_NAME\", \"tf2kernel\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_env_var(\"AWS_ACCOUNT_ID\", sess.account_id())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_env_var(\"AWS_REGION\", sess.boto_region_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "domain_id = sm_client.list_domains()['Domains'][0]['DomainId']\n",
    "set_env_var(\"STUDIO_DOMAIN_ID\", domain_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "set_env_var(\"ROLE_ARN\", sess.get_caller_identity_arn())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 7. Create image configuration of file system for Studio domain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import json, os"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('app-image-config-input.json', 'w') as f:\n",
    "    data = {\n",
    "                \"AppImageConfigName\": os.environ[\"APP_CONF_NAME\"],\n",
    "                \"KernelGatewayImageConfig\": {\n",
    "                    \"KernelSpecs\": [\n",
    "                        {\n",
    "                            \"Name\": \"python3\",\n",
    "                            \"DisplayName\": \"Python 3\"\n",
    "                        }\n",
    "                    ],\n",
    "                    \"FileSystemConfig\": {\n",
    "                        \"MountPath\": \"/root/data\",\n",
    "                        \"DefaultUid\": 0,\n",
    "                        \"DefaultGid\": 0\n",
    "                    }\n",
    "                }\n",
    "            }\n",
    "    json.dump(data, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 8. Create App configuration for Studio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('default-user-settings.json', 'w') as f:\n",
    "    data =  {\n",
    "              \"DefaultUserSettings\": {\n",
    "                \"KernelGatewayAppSettings\": {\n",
    "                  \"CustomImages\": [\n",
    "                      {\n",
    "                               \"ImageName\": os.environ[\"KERNEL_IMAGE_NAME\"],\n",
    "                               \"AppImageConfigName\": os.environ[\"APP_CONF_NAME\"]\n",
    "                            }\n",
    "                        ]\n",
    "                    }\n",
    "                }\n",
    "            }\n",
    "    json.dump(data, f)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 9. Install the tool Sagemaker Studio Image Build (for building Docker images within Studio)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install sagemaker-studio-image-build"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 10. Build Docker image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!sm-docker build . --repository smstudio-custom:$KERNEL_IMAGE_NAME"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 11. Go to the ECR console and check if the ECR repository called `smstudio-custom` is there and the image tag `tf2kernel` is also there\n",
    "\n",
    "Click here:\n",
    "[https://console.aws.amazon.com/ecr/](https://us-west-1.console.aws.amazon.com/ecr/repositories/private/795875386142/smstudio-custom?region=us-west-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 12. Publish image to Studio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!aws --region ${AWS_REGION} sagemaker create-image --image-name ${KERNEL_IMAGE_NAME} --role-arn ${ROLE_ARN}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!aws --region ${AWS_REGION} sagemaker create-image-version --image-name ${KERNEL_IMAGE_NAME} --base-image \"${AWS_ACCOUNT_ID}.dkr.ecr.${AWS_REGION}.amazonaws.com/smstudio-custom:${KERNEL_IMAGE_NAME}\"  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Create AppImageConfig for this image\n",
    "!aws --region ${AWS_REGION} sagemaker create-app-image-config --cli-input-json file://app-image-config-input.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Update the Domain, providing the Image and AppImageConfig\n",
    "!aws --region ${AWS_REGION} sagemaker update-domain --domain-id ${STUDIO_DOMAIN_ID} --cli-input-json file://default-user-settings.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 13. Test it!\n",
    "\n",
    " Click on the top in the `File` > `New Notebook`.\n",
    "- Select the new Kernel called `tf2kernel`.\n",
    "- Run:\n",
    "```\n",
    "import tensorflow as tf\n",
    "tf.__version__\n",
    "```\n",
    "\n",
    "It should output:\n",
    "`2.3.0`\n",
    "\n",
    "Should be like this:\n",
    "\n",
    "![custom-kernel.png](custom-kernel.png)\n",
    "\n",
    "Run: `!cat /root/train.py`\n",
    "It should output: our training script that we saved in the docker image.\n",
    "\n",
    "Should be like this:\n",
    "\n",
    "![script-saved-in-container.png](script-saved-in-container.png)"
   ]
  }
 ],
 "metadata": {
  "instance_type": "ml.t3.medium",
  "kernelspec": {
   "display_name": "Python 3 (Data Science)",
   "language": "python",
   "name": "python3__SAGEMAKER_INTERNAL__arn:aws:sagemaker:us-west-1:742091327244:image/datascience-1.0"
  },
  "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.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}